Update app.py

app.py

@@ -24,14 +24,6 @@ def preload_models():
 
 models = preload_models()
 
-# Function to process the image and generate a caption
-def generate_caption(image, model_name):
-    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
-
 # Predefined images for selection
 image_folder = "images"
 predefined_images = [
@@ -40,17 +32,22 @@ predefined_images = [
     if  fname.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.ppm')) 
 ]
 
-# Define components logic
-def handle_uploaded_image(image):
+# Function to preprocess the image to RGB format
+def preprocess_image(image):
     if image is None:
         return None, None
     pil_image = image.convert("RGB")
     return pil_image, None
 
-def handle_generate_button(image, selected_model):
+# Function to process the image and generate a caption
+def generate_caption(image, selected_model):
     if image is None:
         return "Please upload an image to generate a caption."
-    return generate_caption(image, selected_model)
+    model, tokenizer, image_processor = models[selected_model]
+    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
 
 # Define UI
 with gr.Blocks(theme=gr.themes.Citrus(primary_hue="blue", secondary_hue="orange")) as interface:
@@ -84,7 +81,7 @@ with gr.Blocks(theme=gr.themes.Citrus(primary_hue="blue", secondary_hue="orange"
     with gr.Row(variant='panel'):
             examples = gr.Examples(
                 examples=predefined_images,
-                fn=handle_uploaded_image,
+                fn=preprocess_image,
                 inputs=[image_display],
                 outputs=[image_display, output_text],
                 label="Examples"
@@ -93,9 +90,9 @@ with gr.Blocks(theme=gr.themes.Citrus(primary_hue="blue", secondary_hue="orange"
     # Define actions
     model_selector.change(fn=lambda: (None, None), outputs=[image_display, output_text])
 
-    image_display.upload(fn=handle_uploaded_image, inputs=[image_display], outputs=[image_display, output_text])
+    image_display.upload(fn=preprocess_image, inputs=[image_display], outputs=[image_display, output_text])
     image_display.clear(fn=lambda: None, outputs=[output_text])
 
-    generate_button.click(fn=handle_generate_button, inputs=[image_display, model_selector], outputs=output_text)
+    generate_button.click(fn=generate_caption, inputs=[image_display, model_selector], outputs=output_text)
 
     interface.launch(share=False)