Update server.py

server.py

@@ -1,21 +1,23 @@
-# perchance_server.py
+# perchance_server_with_pyvirtualdisplay.py
 """
 Perchance Image-Generation Server v2.0
 
-Changes from v1:
-  • Auto-detects 'invalid_key' and re-fetches userKey via zendriver
-    without manual intervention.
-  • Only ONE browser launch at startup (no duplicate fetch).
-  • Uses FastAPI *lifespan* context-manager → zero deprecation warnings.
-  • Coordinated key refresh: one refresh at a time; other workers wait.
-  • SSE events: key_refreshing / key_refreshed / key_refresh_failed
-    so the client knows exactly what's happening.
-  • Cleaner separation of concerns, better error handling.
-"""
+This variant adds optional pyvirtualdisplay support so the server can be
+hosted on headless environments (Hugging Face Spaces etc.) while keeping
+all original behaviour unchanged.
 
-# ═══════════════════════════════════════════════════════════════
-#  IMPORTS
-# ═══════════════════════════════════════════════════════════════
+Behaviour:
+ - If ZD_HEADLESS is True, zendriver will run headless as before.
+ - If ZD_HEADLESS is False and a DISPLAY is not present, we attempt to
+   start a pyvirtualdisplay.Display (Xvfb) automatically before launching
+   browsers. If pyvirtualdisplay is not installed or starting Xvfb fails,
+   we log a warning and continue.
+
+To run on Hugging Face Spaces, add `pyvirtualdisplay` to requirements.txt
+and ensure `xvfb` is available in the runtime (HF Spaces typically provide it).
+
+All original defaults and constants are preserved from the original file.
+"""
 
 import asyncio
 import base64
@@ -41,6 +43,13 @@ from sse_starlette.sse import EventSourceResponse
 import zendriver as zd
 from zendriver import cdp
 
+# Try to import pyvirtualdisplay (optional)
+try:
+    from pyvirtualdisplay import Display
+    _HAS_PYVIRTUALDISPLAY = True
+except Exception:
+    Display = None
+    _HAS_PYVIRTUALDISPLAY = False
 
 # ═══════════════════════════════════════════════════════════════
 #  CONFIGURATION
@@ -92,9 +101,6 @@ log = logging.getLogger("perchance")
 
 # ═══════════════════════════════════════════════════════════════
 #  GLOBAL STATE
-#
-#  Asyncio primitives (Lock, Event, Queue) are created inside
-#  the lifespan handler so they live in uvicorn's event loop.
 # ═══════════════════════════════════════════════════════════════
 
 USER_KEY: Optional[str] = None
@@ -114,6 +120,9 @@ TASK_QUEUES: Dict[str, asyncio.Queue]   = {}         # SSE event queues
 EXECUTOR = ThreadPoolExecutor(max_workers=EXECUTOR_THREADS)
 SCRAPER  = cloudscraper.create_scraper()
 
+# pyvirtualdisplay handle (optional)
+VDISPLAY: Optional[Display] = None
+
 
 # ═══════════════════════════════════════════════════════════════
 #  SMALL HELPERS
@@ -141,6 +150,55 @@ def _stamp() -> str:
     return datetime.utcnow().strftime("%Y%m%dT%H%M%S")
 
 
+# ═══════════════════════════════════════════════════════════════
+#  Virtual display helpers (pyvirtualdisplay)
+# ═══════════════════════════════════════════════════════════════
+
+def _start_virtual_display_if_needed(headless: bool):
+    """
+    Start pyvirtualdisplay.Display() if we're running non-headless in an
+    environment without DISPLAY.
+    This function is synchronous and safe to be run in a thread executor.
+    """
+    global VDISPLAY
+
+    if headless:
+        log.info("ZD_HEADLESS=True → not starting virtual display")
+        return
+
+    if os.environ.get("DISPLAY"):
+        log.info("DISPLAY already set: %s", os.environ.get("DISPLAY"))
+        return
+
+    if not _HAS_PYVIRTUALDISPLAY or Display is None:
+        log.warning(
+            "pyvirtualdisplay not installed — cannot create virtual DISPLAY. "
+            "Install pyvirtualdisplay in your environment to enable Xvfb.")
+        return
+
+    try:
+        VDISPLAY = Display(visible=0, size=(1280, 720))
+        VDISPLAY.start()
+        # pyvirtualdisplay sets DISPLAY env itself; log for visibility
+        log.info("Started virtual display via pyvirtualdisplay (DISPLAY=%s)", os.environ.get("DISPLAY"))
+    except Exception as exc:
+        VDISPLAY = None
+        log.exception("Failed to start virtual display: %s", exc)
+
+
+def _stop_virtual_display_if_needed():
+    global VDISPLAY
+    if VDISPLAY is None:
+        return
+    try:
+        VDISPLAY.stop()
+        log.info("Stopped virtual display")
+    except Exception:
+        log.exception("Error while stopping virtual display")
+    finally:
+        VDISPLAY = None
+
+
 # ═══════════════════════════════════════════════════════════════
 #  PERCHANCE HTTP CLIENT  (blocking – runs in ThreadPoolExecutor)
 # ═══════════════════════════════════════════════════════════════
@@ -464,6 +522,15 @@ async def fetch_key_via_browser(
         timeout, headless,
     )
 
+    # If we're in non-headless mode on a display-less host, ensure a virtual
+    # DISPLAY is started first. This call is synchronous so we run it in the
+    # event loop's default executor when called from async code.
+    loop = asyncio.get_running_loop()
+    try:
+        await loop.run_in_executor(None, partial(_start_virtual_display_if_needed, headless))
+    except Exception:
+        log.exception("Error while attempting to start virtual display")
+
     try:
         browser = await zd.start(headless=headless)
     except Exception as exc:
@@ -967,6 +1034,15 @@ async def lifespan(app: FastAPI):
     _key_refresh_lock = asyncio.Lock()
     JOB_QUEUE         = asyncio.Queue(maxsize=MAX_QUEUE_SIZE)
 
+    # If running non-headless and no DISPLAY is present, attempt to start
+    # a virtual display (pyvirtualdisplay/Xvfb). This runs in a thread
+    # executor because pyvirtualdisplay is synchronous.
+    loop = asyncio.get_running_loop()
+    try:
+        await loop.run_in_executor(None, partial(_start_virtual_display_if_needed, ZD_HEADLESS))
+    except Exception:
+        log.exception("Failed to ensure virtual display at startup")
+
     # ── initial key fetch (skip if already set from __main__) ──
     if USER_KEY:
         log.info("Using pre-fetched userKey (len=%d)", len(USER_KEY))
@@ -1014,12 +1090,19 @@ async def lifespan(app: FastAPI):
     except Exception:
         pass
     EXECUTOR.shutdown(wait=True)
+
+    # Stop virtual display if we started one
+    try:
+        await loop.run_in_executor(None, _stop_virtual_display_if_needed)
+    except Exception:
+        log.exception("Failed to stop virtual display cleanly")
+
     log.info("Shutdown complete")
 
 
 # ── app ──
 app = FastAPI(
-    title="Perchance Image Generation Server v2",
+    title="Perchance Image Generation Server v2 (pyvirtualdisplay)",
     lifespan=lifespan,
 )
 app.add_middleware(
@@ -1234,34 +1317,19 @@ async def get_output(filename: str):
 
 
 # ═══════════════════════════════════════════════════════════════
-#  MAIN
+#  MAIN (when run directly)
 # ═══════════════════════════════════════════════════════════════
 
-def _run_uvicorn():
+if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
-
 
-if __name__ == "__main__":
-    skip = os.environ.get("NO_INITIAL_FETCH", "") in ("1", "true", "True")
+    # Allow overriding via environment variables
+    ZD_HEADLESS = os.environ.get("ZD_HEADLESS", str(ZD_HEADLESS)) in ("1", "true", "True")
 
-    if not skip:
-        log.info("Pre-startup: fetching userKey via browser (timeout %ds) …", ZD_TIMEOUT)
-        try:
-            found = asyncio.run(
-                fetch_key_via_browser(timeout=ZD_TIMEOUT, headless=ZD_HEADLESS)
-            )
-            if found:
-                USER_KEY = found
-                log.info("Pre-startup fetch OK (len=%d)", len(found))
-            else:
-                log.warning(
-                    "Pre-startup fetch returned nothing. "
-                    "Server will start; set key via /set_user_key."
-                )
-        except Exception as exc:
-            log.exception("Pre-startup fetch error: %s", exc)
-    else:
-        log.info("NO_INITIAL_FETCH=1 → skipping pre-startup browser fetch")
+    # If running locally and not headless, start virtual display if needed
+    try:
+        _start_virtual_display_if_needed(ZD_HEADLESS)
+    except Exception:
+        log.exception("Failed to ensure virtual display in __main__")
 
-    _run_uvicorn()
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))