Update app.py

app.py

@@ -10,20 +10,26 @@ CAPTION_MODELS = {
 }
 
 # Simple caption creation
-def caption_image(model_choice, image_input):
+def caption_image(model_choice, image_input, url_input):
+    if image_input is not None:
+        input_data = image_input
+    else:
+        input_data = url_input
+
     captioner = pipeline(task="image-to-text",
                          model=CAPTION_MODELS[model_choice],
                          max_new_tokens=30,
                          device_map="cpu", use_fast=True
                          )
-    caption = captioner(image_input)[0]['generated_text']
+    caption = captioner(input_data)[0]['generated_text']
     return str(caption).strip()
 
-def launch(model_choice, image_input):
-    return caption_image(model_choice, image_input)
+def launch(model_choice, image_input, url_input):
+    return caption_image(model_choice, image_input, url_input)
 
-model_dropdown = gr.Dropdown(choices=list(CAPTION_MODELS.keys()), label='Model Choice')
-image_input = gr.Image(type="pil", label="Input Image or URL")
+model_dropdown = gr.Dropdown(choices=list(CAPTION_MODELS.keys()), label='Select Caption Model')
+image_input = gr.Image(type="pil", label="Input Image")
+url_input = gr.Text(label="Input URL")
 
-iface = gr.Interface(launch, inputs=[model_dropdown, image_input], outputs="text")
+iface = gr.Interface(launch, inputs=[model_dropdown, image_input, url_input], outputs="text")
 iface.launch()