Update app.py

app.py

@@ -1,55 +1,48 @@
 import requests 
-from PIL import Image, UnidentifiedImageError
+from PIL import Image
 from transformers import AutoTokenizer, AutoImageProcessor, VisionEncoderDecoderModel
 import gradio as gr
 import os
+from concurrent.futures import ThreadPoolExecutor
 
 # Load the model, tokenizer, and image processor with error handling
 def load_model_and_components(model_name):
-    try:
-        model = VisionEncoderDecoderModel.from_pretrained(model_name)
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        image_processor = AutoImageProcessor.from_pretrained(model_name)
-        return model, tokenizer, image_processor
-    except Exception as e:
-        raise RuntimeError(f"Error loading model components: {e}")
+    model = VisionEncoderDecoderModel.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    image_processor = AutoImageProcessor.from_pretrained(model_name)
+    return model, tokenizer, image_processor
 
-# Preload both models
+# Preload both models in parallel
 def preload_models():
     models = {}
-    models["laicsiifes/swin-distilbertimbau"] = load_model_and_components("laicsiifes/swin-distilbertimbau")
-    models["laicsiifes/swin-gportuguese-2"] = load_model_and_components("laicsiifes/swin-gportuguese-2")
+    model_names = ["laicsiifes/swin-distilbertimbau", "laicsiifes/swin-gportuguese-2"]
+    with ThreadPoolExecutor() as executor:
+        results = executor.map(load_model_and_components, model_names)
+    for name, result in zip(model_names, results):
+        models[name] = result
     return models
 
 models = preload_models()
-current_model_name = "laicsiifes/swin-distilbertimbau"
-model, tokenizer, image_processor = models[current_model_name]
 
 # Function to process the image and generate a caption
 def generate_caption(image, model_name):
-    try:
-        model, tokenizer, image_processor = models[model_name]
-        pixel_values = image_processor(image, return_tensors="pt").pixel_values
-        generated_ids = model.generate(pixel_values)
-        caption = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        return caption
-    except Exception:
-        return "Please upload a valid image."
+    model, tokenizer, image_processor = models[model_name]
+    pixel_values = image_processor(image, return_tensors="pt").pixel_values
+    generated_ids = model.generate(pixel_values, max_length=30, num_beams=2)
+    caption = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return caption
 
 # Predefined images for selection
 image_folder = "images"
 predefined_images_paths = [
-    os.path.join(image_folder, fname) for fname in os.listdir(image_folder) if fname.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))
+    os.path.join(image_folder, fname) for fname in os.listdir(image_folder) if fname.lower().endswith(('.png', '.jpg', '.jpeg'))
 ]
 
 # Gradio app
 def app(image=None, model_name="laicsiifes/swin-distilbertimbau"):
-    try:
-        if image is None:
-            return "Please upload a valid image."
-        return generate_caption(image, model_name)
-    except Exception:
+    if image is None:
         return "Please upload a valid image."
+    return generate_caption(image, model_name)
 
 # Define UI
 with gr.Blocks() as interface:
@@ -62,27 +55,34 @@ with gr.Blocks() as interface:
     """)
     with gr.Row():
         with gr.Column():
-            model_selector = gr.Dropdown(choices=list(models.keys()), 
-                                         value="laicsiifes/swin-distilbertimbau", 
-                                         label="Select Model")
-            image_display = gr.Image(type="pil", label="Image Preview", interactive=False)
+            model_selector = gr.Dropdown(
+                choices=list(models.keys()),
+                value="laicsiifes/swin-distilbertimbau",
+                label="Select Model"
+            )
+    with gr.Row():
+        with gr.Column():
             upload_button = gr.File(label="Upload an Image", file_types=["image"], type="filepath")
             examples = gr.Examples(predefined_images_paths, inputs=[upload_button], label="Examples")
-
+            image_display = gr.Image(type="pil", label="Image Preview", interactive=False)
+            generate_button = gr.Button("Generate")
         with gr.Column():
             output_text = gr.Textbox(label="Generated Caption")
-            
+
     # Define logic
-    def handle_uploaded_image(image, selected_model):
-        try:
-            if image is None:
-                return None, "Please upload a valid image."
-            pil_image = Image.open(image).convert("RGB")
-            return pil_image, generate_caption(pil_image, selected_model)
-        except Exception:
-            return None, "Please upload a valid image."
+    def handle_uploaded_image(image):
+        if image is None:
+            return None
+        pil_image = Image.open(image).convert("RGB")
+        return pil_image
+
+    def handle_generate_button(image, selected_model):
+        if image is None:
+            return "Please upload an image to generate a caption."
+        return generate_caption(image, selected_model)
 
     model_selector.change(fn=lambda _: (None, None, None), inputs=[model_selector], outputs=[image_display, upload_button, output_text])
-    upload_button.change(fn=handle_uploaded_image, inputs=[upload_button, model_selector], outputs=[image_display, output_text])
+    upload_button.change(fn=handle_uploaded_image, inputs=upload_button, outputs=image_display)
+    generate_button.click(fn=handle_generate_button, inputs=[image_display, model_selector], outputs=output_text)
 
-    interface.launch(share=False)
+    interface.launch(share=False)