Update app.py

app.py

@@ -87,7 +87,7 @@ def load_model(model_name):
         yield f"Status: Error loading model: {e}", gr.update(interactive=True), None, None
 
 # --- Playwright Screenshot Function ---
-def take_screenshot(url, max_dimension=1024, full_page_capture=True):
+def take_screenshot(url, max_dimension=1024, full_page_capture=True, timeout=0):
     if not url.startswith('http://') and not url.startswith('https://'):
         url = 'http://' + url
     try:
@@ -95,6 +95,7 @@ def take_screenshot(url, max_dimension=1024, full_page_capture=True):
             browser = p.chromium.launch(headless=True)
             page = browser.new_page()
             page.goto(url, wait_until='networkidle')
+            time.sleep(timeout)
             screenshot_path = f"screenshot_{int(time.time())}.png"
             page.screenshot(path=screenshot_path, full_page=full_page_capture)
             browser.close()
@@ -165,8 +166,10 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         url_input = gr.Textbox(label="Website URL", placeholder="e.g., www.google.com")
-        screenshot_button = gr.Button("Capture Screenshot")
-
+        with gr.Column():
+            screenshot_button = gr.Button("Capture Screenshot")
+            timeout_slider = gr.Slider(1, 1, 10, step=1, label="Page Load Timeout")
+    
     with gr.Row():
         image_output = gr.Image(type="numpy", label="Screenshot", scale=1)
         with gr.Column(scale=1):
@@ -194,7 +197,7 @@ with gr.Blocks() as demo:
     
     screenshot_button.click(
         fn=update_image,
-        inputs=[url_input, max_dim_slider, full_page_checkbox],
+        inputs=[url_input, max_dim_slider, full_page_checkbox, timeout_slider],
         outputs=image_output
     )