Update app.py

app.py

@@ -11,6 +11,48 @@ import numpy as np
 MODEL_PATH = "Qwen/Qwen3-VL-2B-Instruct"
 CPU_DEVICE = "cpu"
 
+
+# --- DETAILED PROMPT TEMPLATE ---
+DETAILED_ANALYSIS_PROMPT = """
+Analyze the provided webpage screenshot and generate a detailed markdown document. Break down the page into logical sections. For each section, identify and list all significant elements such as headers, navigation links, buttons, images, forms, and footer content. Your analysis should be structured as follows:
+
+# [Page Title or a General Description]
+
+## 1. Overall Summary
+Provide a brief, high-level overview of the webpage's purpose, primary content, and general layout.
+
+## 2. Header / Navigation Bar
+- **Logo:** Describe the logo if present.
+- **Navigation Links:** List all visible navigation links (e.g., Home, About, Services, Contact).
+- **Buttons / CTAs:** List any buttons, such as "Sign In," "Register," or "Get Started."
+- **Search Bar:** Note if a search bar is present.
+
+## 3. Main Content / Body
+This is the primary section of the page. Break it down logically.
+- **Hero Section (if applicable):** Describe the main banner, headline, subheading, and primary call-to-action (CTA).
+- **Key Sections:** Identify and describe each subsequent section. For each, list:
+  - **Headings and Subheadings.**
+  - **Text/Paragraphs:** Briefly summarize the content.
+  - **Images:** Describe the images and their purpose.
+  - **Buttons and Links:** List all interactive elements and their text.
+  - **Forms:** Describe any input fields, labels, and submission buttons.
+
+## 4. Sidebar (if present)
+Describe the content of the sidebar, including any navigation, filters, or advertisements.
+
+## 5. Footer
+- **Footer Links:** Create a comprehensive list of all links, often categorized under headings like "Company," "Resources," or "Legal."
+- **Social Media Icons:** List any social media platforms linked.
+- **Contact Information:** Note any address, phone number, or email.
+- **Copyright Information:** State the copyright notice.
+
+## 6. General UI/UX Observations
+- **Color Scheme:** Describe the primary colors used.
+- **Typography:** Comment on the font styles (headings vs. body).
+- **Layout:** Describe the overall structure (e.g., single-column, grid-based, etc.).
+"""
+
+
 # --- Model and Processor Loading ---
 print("Loading model and processor...")
 try:
@@ -112,7 +154,7 @@ def process_and_generate(image_input, text_prompt, processing_size=512):
         if model.config.pad_token_id is None:
             model.config.pad_token_id = model.config.eos_token_id
 
-        generated_ids = model.generate(**inputs, max_new_tokens=1024, do_sample=True, top_p=0.8, temperature=0.7)
+        generated_ids = model.generate(**inputs, max_new_tokens=2048, do_sample=True, top_p=0.8, temperature=0.7) # Increased tokens for detailed prompt
 
         generated_ids_trimmed = [
             out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
@@ -139,12 +181,18 @@ with gr.Blocks() as demo:
         """
     )
 
-    with gr.Accordion("Performance Controls", open=True):
-        full_page_checkbox = gr.Checkbox(
-            value=True,
-            label="Enable Full Height Page Capture",
-            info="If checked, captures the entire scrollable webpage. If unchecked, captures only the visible part."
-        )
+    with gr.Accordion("Controls", open=True):
+        with gr.Row():
+            use_template_checkbox = gr.Checkbox(
+                value=True,
+                label="Use Detailed Analysis Template",
+                info="If checked, uses a comprehensive prompt to dissect the page into sections. If unchecked, uses the prompt in the textbox below."
+            )
+            full_page_checkbox = gr.Checkbox(
+                value=True,
+                label="Enable Full Height Page Capture",
+                info="If checked, captures the entire scrollable webpage. If unchecked, captures only the visible part."
+            )
         max_dim_slider = gr.Slider(
             minimum=512,
             maximum=2048,
@@ -167,10 +215,10 @@ with gr.Blocks() as demo:
         screenshot_button = gr.Button("Capture Screenshot")
 
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             image_output = gr.Image(type="numpy", label="Screenshot")
-        with gr.Column():
-            text_prompt = gr.Textbox(label="Prompt", placeholder="e.g., Describe this webpage in detail as a markdown document.", value="Describe this webpage in detail as a markdown document.")
+        with gr.Column(scale=1):
+            text_prompt = gr.Textbox(label="Custom Prompt", placeholder="e.g., Describe this webpage in detail as a markdown document.", value="Describe this page's color scheme.")
             submit_button = gr.Button("Generate Markdown")
 
     with gr.Row():
@@ -183,16 +231,21 @@ with gr.Blocks() as demo:
             return screenshot_path
         else:
             raise gr.Error(screenshot_path)
-            
+
     # --- Function to handle the loading UI and processing ---
-    def generate_markdown_with_loading(image, prompt, size):
+    def generate_markdown_with_loading(image, user_prompt, processing_size, use_template):
         # Return a dictionary of updates to show the loading state
         yield {
-            output_text: "Processing, please wait...",
+            output_text: "## Processing, please wait... ⏳",
             submit_button: gr.update(interactive=False)
         }
+        
+        # Determine which prompt to use
+        final_prompt = DETAILED_ANALYSIS_PROMPT if use_template else user_prompt
+        
         # Process the data
-        result = process_and_generate(image, prompt, size)
+        result = process_and_generate(image, final_prompt, processing_size)
+        
         # Return a dictionary of updates with the final result
         yield {
             output_text: result,
@@ -208,7 +261,7 @@ with gr.Blocks() as demo:
 
     submit_button.click(
         fn=generate_markdown_with_loading,
-        inputs=[image_output, text_prompt, processing_size_slider],
+        inputs=[image_output, text_prompt, processing_size_slider, use_template_checkbox],
         outputs=[output_text, submit_button]
     )