add sliders

app.py

@@ -131,6 +131,7 @@ def plot_bbox(image, data, use_quad_boxes=False):
             plt.text(x1, y1, label, color='white', fontsize=8, bbox=dict(facecolor='red', alpha=0.5))
 
     ax.axis('off')
+    
     return fig
 
 def draw_ocr_bboxes(image, prediction):
@@ -145,6 +146,7 @@ def draw_ocr_bboxes(image, prediction):
                   "{}".format(label),
                   align="right",
                   fill=color)
+        
     return image
 
 def draw_bounding_boxes(image, quad_boxes, labels, color=(0, 255, 0), thickness=2):
@@ -161,12 +163,7 @@ def draw_bounding_boxes(image, quad_boxes, labels, color=(0, 255, 0), thickness=
 
 def process_image(image, task):
     prompt = TASK_PROMPTS[task]
-    # # Print the inputs for debugging
-    # print(f"\n--- Processing Task: {task} ---")
-    # print(f"Prompt: {prompt}")
     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
-    # # Print the input tensors for debugging
-    # print(f"Model Input: {inputs}")
     generated_ids = model.generate(
         **inputs,
         max_new_tokens=1024,
@@ -175,24 +172,38 @@ def process_image(image, task):
     )
 
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    # # Print the raw generated output for debugging
-    # print(f"Raw Model Output: {generated_text}")
     parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
-    # # Print the parsed answer for debugging
-    # print(f"Parsed Answer: {parsed_answer}")
+
     return parsed_answer
 
 
-def main_process(image, task):
-    result = process_image(image, task)
+def main_process(image, task, top_k, top_p, repetition_penalty, num_beams, max_tokens):
+    prompt = TASK_PROMPTS[task]
+    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
+    generated_ids = model.generate(
+        **inputs,
+        max_new_tokens=max_tokens,
+        num_beams=num_beams,
+        do_sample=True,
+        top_k=top_k,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty
+    )
+
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
+    return parsed_answer
 
+def process_and_update(image, task, top_k, top_p, repetition_penalty, num_beams, max_tokens):
+    if image is None:
+        return None, gr.update(visible=False), "Please upload an image first.", gr.update(visible=True)
+    result = main_process(image, task, top_k, top_p, repetition_penalty, num_beams, max_tokens)
+    
     if task in IMAGE_TASKS:
         if task == "📸✍🏻OCR with Region":
             fig = plot_bbox(image, result.get('<OCR_WITH_REGION>', {}), use_quad_boxes=True)
             output_image = fig_to_pil(fig)
             text_output = result.get('<OCR_WITH_REGION>', {}).get('recognized_text', 'No text found')
-            # # Debugging: Print the recognized text
-            # print(f"Recognized Text: {text_output}")
             return output_image, gr.update(visible=True), text_output, gr.update(visible=False)
         else:
             fig = plot_bbox(image, result.get(TASK_PROMPTS[task], {}))
@@ -201,7 +212,6 @@ def main_process(image, task):
     else:
         return None, gr.update(visible=False), str(result), gr.update(visible=True)
 
-
 def reset_outputs():
     return None, gr.update(visible=False), None, gr.update(visible=True)
 
@@ -224,20 +234,21 @@ with gr.Blocks(title="Tonic's 🙏🏻PLeIAs/📸📈✍🏻Florence-PDF") as if
                 image_input = gr.Image(type="pil", label="Input Image")
                 task_dropdown = gr.Dropdown(list(TASK_PROMPTS.keys()), label="Task", value="✍🏻Caption")            
                 with gr.Row():
-                    submit_button = gr.Button("Process")
-                    reset_button = gr.Button("Reset")
+                    submit_button = gr.Button("📸📈✍🏻Process")
+                    reset_button = gr.Button("♻️Reset")
+                with gr.Accordion("🧪Advanced Settings", open=False):
+                    top_k = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k")
+                    top_p = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.01, label="Top-p")
+                    repetition_penalty = gr.Slider(minimum=1.0, maximum=2.0, value=1.0, step=0.01, label="Repetition Penalty")
+                    num_beams = gr.Slider(minimum=1, maximum=6, value=3, step=1, label="Number of Beams")
+                    max_tokens = gr.Slider(minimum=1, maximum=1024, value=1000, step=1, label="Max Tokens")
             with gr.Column(scale=1):    
                 output_image = gr.Image(label="🙏🏻PLeIAs/📸📈✍🏻Florence-PDF", visible=False)
                 output_text = gr.Textbox(label="🙏🏻PLeIAs/📸📈✍🏻Florence-PDF", visible=True)
-
-    def process_and_update(image, task):
-        if image is None:
-            return None, gr.update(visible=False), "Please upload an image first.", gr.update(visible=True)
-        return main_process(image, task)
     
     submit_button.click(
         fn=process_and_update,
-        inputs=[image_input, task_dropdown],
+        inputs=[image_input, task_dropdown, top_k, top_p, repetition_penalty, num_beams, max_tokens],
         outputs=[output_image, output_image, output_text, output_text]
     )