Spaces:

aavetis
/

ugen-image-captioning

Running on Zero

aavetis commited on Jan 10

Commit

b9c1647

•

1 Parent(s): bf23625

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -11,7 +11,7 @@ processor = gen_model.VLMProcessor.from_pretrained("unum-cloud/uform-gen")
 @spaces.GPU
 def generate_caption(image, prompt):
     # Process the image and the prompt
-    inputs = processor(texts=[prompt], images=[image], return_tensors="pt")
     # Generate the output
     with torch.inference_mode():

 @spaces.GPU
 def generate_caption(image, prompt):
     # Process the image and the prompt
+    inputs = processor(texts=[prompt], images=[image], return_tensors="pt").to('cuda')
     # Generate the output
     with torch.inference_mode():