Update app.py

app.py

@@ -3,35 +3,41 @@ from selenium import webdriver
 from selenium.common.exceptions import WebDriverException
 from PIL import Image
 from io import BytesIO
-import time  # <-- import time
+import time
 
 def take_screenshot(url):
     options = webdriver.ChromeOptions()
-    options.add_argument('--headless')
-    options.add_argument('--no-sandbox')
-    options.add_argument('--disable-dev-shm-usage')
+    options.add_argument("--headless")
+    options.add_argument("--no-sandbox")
+    options.add_argument("--disable-dev-shm-usage")
 
+    wd = None
     try:
         wd = webdriver.Chrome(options=options)
-        wd.set_window_size(1080, 720)  # Adjust the window size here
+        wd.set_window_size(1280, 720)
         wd.get(url)
-        wd.implicitly_wait(10)
-        time.sleep(2)  # <-- Add delay before screenshot
+        time.sleep(2)
         screenshot = wd.get_screenshot_as_png()
+
+        if len(screenshot) < 1000:
+            print("⚠ Screenshot too small, likely error.")
+            return None
+
+        return Image.open(BytesIO(screenshot))
+
     except WebDriverException as e:
-        return Image.new('RGB', (1, 1))
+        print("WebDriverException:", e)
+        return None
     finally:
         if wd:
             wd.quit()
 
-    return Image.open(BytesIO(screenshot))
-
 iface = gr.Interface(
     fn=take_screenshot,
-    inputs=gr.Textbox(label="Website URL", value="https://github.com/aiclubiitropar"),
-    outputs=gr.Image(type="pil", label="Screenshot", height=360, width=540),
+    inputs=gr.Textbox(label="Website URL", value="https://github.com"),
+    outputs=gr.Image(type="pil", label="Screenshot"),
     title="Website Screenshot",
-    description="Take a screenshot of a website."
+    description="Takes a screenshot of the website and displays it here."
 )
 
 iface.launch()