Update app.py

app.py

@@ -20,7 +20,7 @@ model.to(device)
 TITLE = f"# [{model_name}](https://huggingface.co/{model_name})"
 
 
-def process_image(image, num_beams=5, min_p=0.0, top_p=1.0):
+def process_image(image, task="<CAPTION>", num_beams=5, min_p=0.0, top_p=1.0):
     """
     Process a single image to generate a caption.
     Supports image input as file path, numpy array, or PIL Image.
@@ -37,7 +37,7 @@ def process_image(image, num_beams=5, min_p=0.0, top_p=1.0):
 
         # Prepare inputs for the model
         inputs = processor(
-            text="<CAPTION>",
+            text=task,
             images=image,
             return_tensors="pt"
         )
@@ -85,6 +85,11 @@ with gr.Blocks(css=css) as demo:
 
                 submit_btn = gr.Button(value="Submit")
 
+                task_dropdown = gr.Dropdown(
+                    ["<CAPTION>", "<DETAILED_CAPTION>", "<MORE_DETAILED_CAPTION>"],
+                    value="<CAPTION>",
+                    label="Captioning Mode",
+                ),
                 num_beams_slider = gr.Slider(
                     minimum=1,
                     maximum=5,
@@ -109,16 +114,16 @@ with gr.Blocks(css=css) as demo:
 
         gr.Examples(
             [
-                ["eval_img_1.jpg", 5, 0.0, 1.0],
-                ["eval_img_2.jpg", 5, 0.0, 1.0],
-                ["eval_img_3.jpg", 5, 0.0, 1.0],
-                ["eval_img_4.jpg", 5, 0.0, 1.0],
-                ["eval_img_5.jpg", 5, 0.0, 1.0],
-                ["eval_img_6.jpg", 5, 0.0, 1.0],
-                ["eval_img_7.png", 5, 0.0, 1.0],
-                ["eval_img_8.jpg", 5, 0.0, 1.0],
+                ["eval_img_1.jpg", "<CAPTION>", 5, 0.0, 1.0],
+                ["eval_img_2.jpg", "<CAPTION>", 5, 0.0, 1.0],
+                ["eval_img_3.jpg", "<CAPTION>", 5, 0.0, 1.0],
+                ["eval_img_4.jpg", "<CAPTION>", 5, 0.0, 1.0],
+                ["eval_img_5.jpg", "<CAPTION>", 5, 0.0, 1.0],
+                ["eval_img_6.jpg", "<CAPTION>", 5, 0.0, 1.0],
+                ["eval_img_7.png", "<CAPTION>", 5, 0.0, 1.0],
+                ["eval_img_8.jpg", "<CAPTION>", 5, 0.0, 1.0],
             ],
-            inputs=[input_img, num_beams_slider, min_p_slider, top_p_slider],
+            inputs=[input_img, task_dropdown, num_beams_slider, min_p_slider, top_p_slider],
             outputs=[output_text],
             fn=process_image,
             label="Try captioning on below examples",
@@ -126,7 +131,7 @@ with gr.Blocks(css=css) as demo:
 
         submit_btn.click(
             process_image,
-            [input_img, num_beams_slider, min_p_slider, top_p_slider],
+            [input_img, task_dropdown, num_beams_slider, min_p_slider, top_p_slider],
             [output_text]
         )