Update app.py

app.py

@@ -1,45 +1,45 @@
-import os
+import torch
 import requests
 from io import BytesIO
-
 from PIL import Image
 from transformers import AutoProcessor, AutoModelForVision2Seq
+import gradio as gr
 
-def generate_caption(image):
-  # Load pre-trained models & processors
-  model = AutoModelForVision2Seq.from_pretrained("microsoft/kosmos-2-patch14-224")
-  processor = AutoProcessor.from_pretrained("microsoft/kosmos-2-patch14-224")
-
-  prompt = "<grounding>An image of"
+def load_models():
+    # Load pre-trained models
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = AutoModelForVision2Seq.from_pretrained("microsoft/kosmos-2-patch14-224").to(device)
+    processor = AutoProcessor.from_pretrained("microsoft/kosmos-2-patch14-224")
+    return model, processor
 
-  # Open the uploaded image file
-  img = Image.open(BytesIO(image))
 
-  # Save the image locally and open it again to avoid potential issues with reusing the same PIL object
-  img.save("temp_image.jpg")
-  img = Image.open("temp_image.jpg")
+def generate_description(image):
+    model, processor = load_models()  
+    prompt = "<grounding>An image of"
+   
+    inputs = processor(text=prompt, images=image, padding='max_length', truncation=True, return_tensors="pt")
 
-  inputs = processor(text=prompt, images=img, return_tensors="pt")
-  
-  # Generate caption
-  generated_ids = model.generate(**inputs, max_new_tokens=128)
-  generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-  
-  # Process the generated caption
-  processed_text, _ = processor.post_process_generation(generated_text)
-  
-  return processed_text
+    # Move tensors to GPU if available
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
-import gradio as gr
+    # Generate description
+    generated_ids = model.generate(**inputs, max_new_tokens=128)
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
 
-title = 'Image Caption Generator'
-description = 'Generate descriptive captions for images.'
-examples = [["PRO-b0fe1914d67344d98e120a19cd1aadf1.jpg"]]
-article = '<p style="margin:auto;max-width:600px;">This tool generates descriptive captions for given images.</p>'
+    return generated_text
+  
 
-interface = gr.Interface(fn=generate_caption, 
-                        inputs=gr.Image(),
-                        outputs=gr.Textbox(),
-                        title=title, description=description, examples=examples, article=article)
-                        
-interface.launch()
+if __name__ == '__main__':
+    interface = gr.Interface(
+        generate_description, 
+        ["image"], 
+        "text",
+        capture_session=True,
+        allow_recording=False,
+        title="GPT-based Visual Storytelling",
+        description="Upload an image to get a detailed caption generated by our powerful AI!",
+        examples=[
+            ['PRO-b0fe1914d67344d98e120a19cd1aadf1.jpg']
+        ],
+    )
+    interface.launch()