Update app.py

app.py

@@ -1,8 +1,5 @@
 import torch
 import gradio as gr
-from PIL import Image
-import requests
-import io
 from transformers import pipeline
 
 CAPTION_MODELS = {
@@ -14,24 +11,19 @@ CAPTION_MODELS = {
 
 # Simple caption creation
 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,
                          device_map="cpu", use_fast=True
                          )
-    caption = captioner(image_path)[0]['generated_text']
+    caption = captioner(image_input)[0]['generated_text']
     return str(caption).strip()
 
-def launch(model_choice, input):
-    return caption_image(model_choice, input)
+def launch(model_choice, image_input):
+    return caption_image(model_choice, image_input)
 
 model_dropdown = gr.Dropdown(choices=list(CAPTION_MODELS.keys()), label='Model Choice')
-iface = gr.Interface(launch, inputs=[model_dropdown, gr.Data(type="file", label="Upload Image or Enter URL")], outputs="text")
+image_input = gr.Image(type="pil", label="Input Image or URL")
+
+iface = gr.Interface(launch, inputs=[model_dropdown, image_input], outputs="text")
 iface.launch()