Update app.py

app.py

@@ -1,19 +1,29 @@
 from PIL import Image
 import requests
+import torch
 import gradio as gr
 
-from transformers import BlipProcessor, BlipForConditionalGeneration
+from transformers import pipeline
 
-model_id = "Salesforce/blip-image-captioning-base"
+CAPTION_MODELS = {
+    'blip-base': 'Salesforce/blip-image-captioning-base',
+    'blip-large': 'Salesforce/blip-image-captioning-large',
+    'vit-gpt2-coco-en': 'ydshieh/vit-gpt2-coco-en',
+}
 
-model = BlipForConditionalGeneration.from_pretrained(model_id)
-processor = BlipProcessor.from_pretrained(model_id)
+captioner = pipeline(task="image-to-text",
+                        model=CAPTION_MODELS['blip-base'],
+                        max_new_tokens=30,
+                        device_map="auto", use_fast=True
+                        )
+
+# Simple caption creation
+def caption_image(captioner, image_path):
+    caption = captioner(image_path)[0]['generated_text']
+    return str(caption).strip()
 
 def launch(input):
-    image = Image.open(requests.get(input, stream=True).raw).convert('RGB')
-    inputs = processor(image, return_tensors="pt")
-    out = model.generate(**inputs)
-    return processor.decode(out[0], skip_special_tokens=True)
+    return caption_image(captioner, input)
 
 iface = gr.Interface(launch, inputs="text", outputs="text")
-iface.launch()
+iface.launch()