Update app.py

app.py

@@ -21,7 +21,7 @@ def get_media_type(image_name):
     else:
         return None  # Extend this function based on the image formats you expect to handle
 
-def describe_image(image_path, api_key, model):
+def describe_image(image_path, api_key, model, prompt):
     """Send the image to Claude for description."""
     image_base64 = image_to_base64(image_path)
     media_type = get_media_type(image_path)
@@ -44,7 +44,7 @@ def describe_image(image_path, api_key, model):
                     },
                     {
                         "type": "text",
-                        "text": "Describe this image."
+                        "text": prompt
                     }
                 ],
             }
@@ -52,10 +52,10 @@ def describe_image(image_path, api_key, model):
     )
     return message.content
 
-def main(image_path, api_key, model_a, model_b):
+def main(image_path, api_key, model_a, model_b, prompt):
     if api_key:
-        description_a = describe_image(image_path, api_key, model_a)
-        description_b = describe_image(image_path, api_key, model_b)
+        description_a = describe_image(image_path, api_key, model_a, prompt)
+        description_b = describe_image(image_path, api_key, model_b, prompt)
         
         if isinstance(description_a, list) and len(description_a) > 0:
             description_a = description_a[0].text
@@ -81,20 +81,21 @@ with gr.Blocks() as iface:
         with gr.Column():
             image_input = gr.Image(type="filepath", label="Upload Image")
             api_key_input = gr.Textbox(type="password", label="Enter your Claude API Key")
-            run_button = gr.Button("Run")
         
         with gr.Column():
-            with gr.Row():
-                model_a_dropdown = gr.Dropdown(choices=model_options, label="Model A")
-                model_b_dropdown = gr.Dropdown(choices=model_options, label="Model B")
+            model_a_dropdown = gr.Dropdown(choices=model_options, label="Model A")
+            model_b_dropdown = gr.Dropdown(choices=model_options, label="Model B")
             
             with gr.Row():
                 output_a = gr.Textbox(label="Description from Model A")
                 output_b = gr.Textbox(label="Description from Model B")
     
+    prompt_input = gr.Textbox(label="Custom Prompt", value="Describe this image.")
+    run_button = gr.Button("Run")
+    
     run_button.click(
         fn=main,
-        inputs=[image_input, api_key_input, model_a_dropdown, model_b_dropdown],
+        inputs=[image_input, api_key_input, model_a_dropdown, model_b_dropdown, prompt_input],
         outputs=[output_a, output_b]
     )