Update app.py

app.py

@@ -1,5 +1,8 @@
 import torch
 import gradio as gr
+from PIL import Image
+import requests
+import io
 from transformers import pipeline
 
 CAPTION_MODELS = {
@@ -10,7 +13,14 @@ CAPTION_MODELS = {
 }
 
 # Simple caption creation
-def caption_image(model_choice, image_path):
+def caption_image(model_choice, image_input):
+    if isinstance(image_input, str):  # Hopefully a URL
+        image_path = image_input
+    else:  # Upload a file
+        image = Image.open(io.BytesIO(image_input))
+        image.save('temp_image_file.jpg')
+        image_path = 'temp_image_file.jpg'
+    
     captioner = pipeline(task="image-to-text",
                          model=CAPTION_MODELS[model_choice],
                          max_new_tokens=30,
@@ -23,5 +33,5 @@ def launch(model_choice, input):
     return caption_image(model_choice, input)
 
 model_dropdown = gr.Dropdown(choices=list(CAPTION_MODELS.keys()), label='Model Choice')
-iface = gr.Interface(launch, inputs=[model_dropdown, "text"], outputs="text")
-iface.launch()
+iface = gr.Interface(launch, inputs=[model_dropdown, gr.Data(type="file", label="Upload Image or Enter URL")], outputs="text")
+iface.launch()