noamrot commited on
Commit
060b48a
·
verified ·
1 Parent(s): d3c8f24

added use of ZERO so it runs on GPU

Browse files
Files changed (1) hide show
  1. app.py +6 -2
app.py CHANGED
@@ -1,16 +1,20 @@
1
  import gradio as gr
2
  from PIL import Image
 
3
  import torch
4
  from transformers import BlipProcessor, BlipForConditionalGeneration
5
 
6
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
7
 
 
8
  processor = BlipProcessor.from_pretrained("noamrot/FuseCap_Image_Captioning")
9
  model = BlipForConditionalGeneration.from_pretrained("noamrot/FuseCap_Image_Captioning").to(device)
10
 
 
11
  def inference(raw_image):
 
 
12
  text = "a picture of "
13
- inputs = processor(raw_image, text, return_tensors="pt").to(device)
14
  out = model.generate(**inputs)
15
  caption = processor.decode(out[0], skip_special_tokens=True)
16
  return caption
 
1
  import gradio as gr
2
  from PIL import Image
3
+ import spaces
4
  import torch
5
  from transformers import BlipProcessor, BlipForConditionalGeneration
6
 
 
7
 
8
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
9
  processor = BlipProcessor.from_pretrained("noamrot/FuseCap_Image_Captioning")
10
  model = BlipForConditionalGeneration.from_pretrained("noamrot/FuseCap_Image_Captioning").to(device)
11
 
12
+ @spaces.GPU(duration=15)
13
  def inference(raw_image):
14
+ cur_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
15
+ model = model.to(device)
16
  text = "a picture of "
17
+ inputs = processor(raw_image, text, return_tensors="pt").to(cur_device)
18
  out = model.generate(**inputs)
19
  caption = processor.decode(out[0], skip_special_tokens=True)
20
  return caption