Update app.py

app.py

@@ -6,22 +6,22 @@ import gradio as gr
 import PIL.Image
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers import BlipProcessor, BlipForConditionalGeneration
 
-DESCRIPTION = "# Image Captioning with GIT"
+DESCRIPTION = "# Image Captioning with LongCap"
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-model_id = "microsoft/git-large-coco"
-processor = AutoProcessor.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
+model_id = "unography/blip-long-cap"
+processor = BlipProcessor.from_pretrained(model_id)
+model = BlipForConditionalGeneration.from_pretrained(model_id).to(device)
 
 
 @spaces.GPU
 def run(image: PIL.Image.Image) -> str:
     inputs = processor(images=image, return_tensors="pt").to(device)
-    generated_ids = model.generate(pixel_values=inputs.pixel_values, num_beams=3, max_length=20, min_length=5)
-    generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    out = model.generate(pixel_values=pixel_values, max_length=300)
+    generated_caption = processor.decode(out[0], skip_special_tokens=True)
     return generated_caption