Update app/main.py

app/main.py

@@ -1,22 +1,32 @@
 from fastapi import FastAPI, Query
 from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
 import base64
 from playwright.async_api import async_playwright
-import io
 
 app = FastAPI()
 
+# Optional: Allow access from any frontend (useful for Hugging Face Gradio frontends)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
 @app.get("/screenshot")
 async def screenshot(url: str = Query(..., description="Web page URL to capture")):
-    async with async_playwright() as p:
-        browser = await p.chromium.launch()
-        page = await browser.new_page()
-        await page.goto(url)
-        buffer = await page.screenshot(full_page=True)
-        await browser.close()
+    try:
+        async with async_playwright() as p:
+            browser = await p.chromium.launch()
+            page = await browser.new_page()
+            await page.goto(url, timeout=30000)  # 30s timeout
+            buffer = await page.screenshot(full_page=True)
+            await browser.close()
 
-    # Encode the screenshot to base64
-    b64_image = base64.b64encode(buffer).decode("utf-8")
-    data_url = f"data:image/png;base64,{b64_image}"
+        b64_image = base64.b64encode(buffer).decode("utf-8")
+        data_url = f"data:image/png;base64,{b64_image}"
+        return JSONResponse(content={"image_base64": data_url})
 
-    return JSONResponse(content={"image_base64": data_url})
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})