Update app.py

app.py

@@ -1,6 +1,6 @@
+import torch
 import gradio as gr
 from transformers import pipeline
-import torch
 
 CAPTION_MODELS = {
     'blip-base': 'Salesforce/blip-image-captioning-base',
@@ -14,11 +14,14 @@ CAPTION_MODELS = {
 # Create a dictionary to store loaded models
 loaded_models = {}
 
-# Modify caption_image to accept and process lists of images
-def caption_image(model_choice, images_input, urls_input, load_in_8bit, device):
-    input_data = images_input if all(i is not None for i in images_input) else urls_input
+# Simple caption creation
+def caption_image(model_choice, image_input, url_input, load_in_8bit, device):
+    if image_input is not None:
+        input_data = image_input
+    else:
+        input_data = url_input
 
-    model_key = (model_choice, load_in_8bit, device)  # Update the model key to include the device
+    model_key = (model_choice, load_in_8bit)  # Create a tuple to represent the unique combination of model and 8bit loading
 
     # Check if the model is already loaded
     if model_key in loaded_models:
@@ -29,7 +32,7 @@ def caption_image(model_choice, images_input, urls_input, load_in_8bit, device):
         captioner = pipeline(task="image-to-text",
                             model=CAPTION_MODELS[model_choice],
                             max_new_tokens=30,
-                            device=device, # Set the device as selected by the user
+                            device=device, # Set the device as selected
                             model_kwargs=model_kwargs, 
                             torch_dtype=dtype,  # Set the floating point
                             use_fast=True
@@ -37,16 +40,14 @@ def caption_image(model_choice, images_input, urls_input, load_in_8bit, device):
         # Store the loaded model
         loaded_models[model_key] = captioner
 
-    captions = captioner(input_data)  # Run the model on the batch of images
-    results = [str(caption['generated_text']).strip() for caption in captions]  # Extract the captions from the outputs
-
-    return results
+    caption = captioner(input_data)
+    return [str(c['generated_text']).strip() for c in caption]
 
 model_dropdown = gr.Dropdown(choices=list(CAPTION_MODELS.keys()), label='Select Caption Model')
-image_input = gr.Image(type="pil", label="Input Image")  # Now takes multiple images
-url_input = gr.Text(label="Input URL")  # Now takes multiple URLs
+image_input = gr.Image(type="pil", label="Input Image", multiple=True)  # Enable multiple inputs
+url_input = gr.Text(label="Input URL")
 load_in_8bit = gr.Checkbox(label="Load model in 8bit")
-device = gr.Radio(choices=['cpu', 'cuda'], label='Device')  # Radio button for device selection
+device = gr.Radio(['cpu', 'cuda'], label='Select device', default='cpu')
 
-iface = gr.Interface(caption_image, inputs=[model_dropdown, image_input, url_input, load_in_8bit, device], outputs=gr.outputs.Textbox(type="auto", label="Caption"))
+iface = gr.Interface(caption_image, inputs=[model_dropdown, image_input, url_input, load_in_8bit, device], outputs=gr.interfaces.outputs.Textbox(type="text", label="Caption"))
 iface.launch()