Update app.py

app.py

@@ -1,64 +1,50 @@
 import gradio as gr
-import requests
-import concurrent.futures
-import time
-import os
-from selenium import webdriver
-from selenium.webdriver.chrome.options import Options
+import asyncio
+from playwright.async_api import async_playwright
+import base64
 from PIL import Image
 import io
 
-def capture_screenshot(url):
-    chrome_options = Options()
-    chrome_options.add_argument("--headless")
-    chrome_options.add_argument("--no-sandbox")
-    chrome_options.add_argument("--disable-dev-shm-usage")
-    
-    driver = webdriver.Chrome(options=chrome_options)
-    driver.get(url)
-    time.sleep(2)  # Wait for the page to load
-    screenshot = driver.get_screenshot_as_png()
-    driver.quit()
-    
-    return Image.open(io.BytesIO(screenshot))
+async def capture_screenshot(url):
+    async with async_playwright() as p:
+        browser = await p.chromium.launch()
+        page = await browser.new_page()
+        await page.goto(url)
+        screenshot = await page.screenshot()
+        await browser.close()
+        return Image.open(io.BytesIO(screenshot))
 
-def open_link(url, index):
-    start_time = time.time()
+async def open_link(url, index):
     try:
-        response = requests.get(url, timeout=10)
-        elapsed_time = time.time() - start_time
-        status = f"Request {index+1}: Status {response.status_code}, Time: {elapsed_time:.2f}s, Size: {len(response.content)} bytes"
-        
-        screenshot = capture_screenshot(url)
-        return status, screenshot
+        screenshot = await capture_screenshot(url)
+        return f"Request {index+1}: Successfully captured screenshot", screenshot
     except Exception as e:
         return f"Request {index+1}: Error: {str(e)}", None
 
-def open_links(link, count):
-    results = []
-    screenshots = []
-    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
-        futures = [executor.submit(open_link, link, i) for i in range(count)]
-        for future in concurrent.futures.as_completed(futures):
-            status, screenshot = future.result()
-            results.append(status)
-            if screenshot:
-                screenshots.append(screenshot)
+async def open_links(link, count):
+    tasks = [open_link(link, i) for i in range(count)]
+    results = await asyncio.gather(*tasks)
     
-    return "\n".join(results), screenshots[0] if screenshots else None
+    statuses = [result[0] for result in results]
+    screenshots = [result[1] for result in results if result[1] is not None]
+    
+    return "\n".join(statuses), screenshots[0] if screenshots else None
+
+def gradio_open_links(link, count):
+    return asyncio.run(open_links(link, count))
 
 iface = gr.Interface(
-    fn=open_links,
+    fn=gradio_open_links,
     inputs=[
         gr.Textbox(label="Enter link", placeholder="https://example.com"),
-        gr.Slider(minimum=1, maximum=10, step=1, value=1, label="Number of times to open")
+        gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Number of times to open")
     ],
     outputs=[
         gr.Textbox(label="Results"),
         gr.Image(label="Screenshot")
     ],
     title="Link Opener with Screenshot",
-    description="Enter a link and choose how many times to open it. The app will attempt to open the link, return the status, response time, content size for each attempt, and display a screenshot of the first successful attempt.",
+    description="Enter a link and choose how many times to open it. The app will attempt to open the link, capture a screenshot, and display the first successful screenshot.",
     theme="huggingface",
     examples=[["https://www.example.com", 1]]
 )