nielsr HF staff commited on
Commit
a808d84
1 Parent(s): 329d18e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -2
app.py CHANGED
@@ -10,10 +10,15 @@ git_model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-coco")
10
  blip_processor = AutoProcessor.from_pretrained("Salesfoce/blip-image-captioning-base")
11
  blip_model = BlipForConditionalGeneration.from_pretrained("Salesfoce/blip-image-captioning-base")
12
 
 
 
 
 
 
13
  def generate_caption(processor, model, image):
14
- inputs = processor(image=image, return_tensors="pt")
15
 
16
- generated_ids = model.generate(pixel_values=pixel_values, max_length=50)
17
 
18
  generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
19
 
 
10
  blip_processor = AutoProcessor.from_pretrained("Salesfoce/blip-image-captioning-base")
11
  blip_model = BlipForConditionalGeneration.from_pretrained("Salesfoce/blip-image-captioning-base")
12
 
13
+ device = "cuda" if torch.cuda.is_available() else "cpu"
14
+
15
+ git_model.to(device)
16
+ blip_model.to(device)
17
+
18
  def generate_caption(processor, model, image):
19
+ inputs = processor(image=image, return_tensors="pt").to(device)
20
 
21
+ generated_ids = model.generate(pixel_values=inputs.pixel_values, max_length=50)
22
 
23
  generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
24