AdarshJi/CSSSC
0
1# perchance_server_with_pyvirtualdisplay.py2"""3Perchance Image-Generation Server v2.04 5This variant adds optional pyvirtualdisplay support so the server can be6hosted on headless environments (Hugging Face Spaces etc.) while keeping7all original behaviour unchanged.8 9Behaviour:10 - If ZD_HEADLESS is True, zendriver will run headless as before.11 - If ZD_HEADLESS is False and USE_VIRTUAL_DISPLAY is True and a DISPLAY12 is not present, we attempt to start a pyvirtualdisplay.Display (Xvfb)13 automatically before launching browsers. If pyvirtualdisplay is not14 installed or starting Xvfb fails, we log a warning and continue.15 - If USE_VIRTUAL_DISPLAY is False we will NOT attempt to start a virtual16 display — you must provide a DISPLAY yourself (or set ZD_HEADLESS=True)17 if running on a headless host.18 19To run on Hugging Face Spaces, add `pyvirtualdisplay` to requirements.txt20and ensure `xvfb` is available in the runtime (HF Spaces typically provide it).21 22All original defaults and constants are preserved from the original file.23"""24 25import asyncio26import base6427import json28import logging29import os30import random31import string32import time33import uuid34from concurrent.futures import ThreadPoolExecutor35from contextlib import asynccontextmanager36from datetime import datetime37from functools import partial38from pathlib import Path39from typing import Any, Dict, List, Optional40 41import cloudscraper42from fastapi import FastAPI, HTTPException, Request43from fastapi.middleware.cors import CORSMiddleware44from fastapi.responses import FileResponse45from sse_starlette.sse import EventSourceResponse46import zendriver as zd47from zendriver import cdp48 49# Try to import pyvirtualdisplay (optional)50try:51 from pyvirtualdisplay import Display52 _HAS_PYVIRTUALDISPLAY = True53except Exception:54 Display = None55 _HAS_PYVIRTUALDISPLAY = False56 57# ═══════════════════════════════════════════════════════════════58# CONFIGURATION59# ═══════════════════════════════════════════════════════════════60 61# --- Perchance API ---62BASE_URL = "https://image-generation.perchance.org"63API_GENERATE = "/api/generate"64API_DOWNLOAD = "/api/downloadTemporaryImage"65API_AWAIT = "/api/awaitExistingGenerationRequest"66API_ACCESS_CODE = "/api/getAccessCodeForAdPoweredStuff"67 68# --- Browser automation (zendriver) ---69TARGET_URL = "https://perchance.org/ai-text-to-image-generator"70IMAGE_GEN_ORIGIN = "https://image-generation.perchance.org"71ZD_TIMEOUT = 90 # seconds for key-fetch attempt72ZD_HEADLESS = False # True → hide browser window73CLICK_INTERVAL = 0.3574CLICK_JITTER = 8.075KEY_PREFIX = "userKey"76 77# --- Virtual display toggle (new) ---78# If True the server will attempt to auto-start a pyvirtualdisplay (Xvfb)79# when no DISPLAY is present and ZD_HEADLESS is False. If False the server80# will not try to start a virtual display and you must provide a DISPLAY81# or set ZD_HEADLESS=True.82USE_VIRTUAL_DISPLAY = True83 84# --- HTTP / generation ---85HTTP_TIMEOUT = 3086MAX_DOWNLOAD_WAIT = 18087BACKOFF_INIT = 0.788MAX_GEN_RETRIES = 6 # retries inside generate_one()89 90# --- Key-refresh policy ---91MAX_KEY_RETRIES = 3 # per-image retries when key is invalid92KEY_REFRESH_COOLDOWN = 30 # min seconds between two refreshes93MAX_REFRESH_FAILURES = 5 # consecutive failures → stop auto-refresh94 95# --- Server ---96WORKER_COUNT = 397MAX_QUEUE_SIZE = 100098EXECUTOR_THREADS = 1699OUTPUT_DIR = Path("outputs")100OUTPUT_DIR.mkdir(exist_ok=True, parents=True)101 102 103# ═══════════════════════════════════════════════════════════════104# LOGGING105# ═══════════════════════════════════════════════════════════════106 107LOG_FMT = "%(asctime)s | %(levelname)-7s | %(message)s"108logging.basicConfig(level=logging.INFO, format=LOG_FMT)109log = logging.getLogger("perchance")110 111 112# ═══════════════════════════════════════════════════════════════113# GLOBAL STATE114# ═══════════════════════════════════════════════════════════════115 116USER_KEY: Optional[str] = None117 118# -- set in lifespan --119_key_lock: Optional[asyncio.Lock] = None # guard USER_KEY reads/writes120_key_valid: Optional[asyncio.Event] = None # cleared while refreshing121_key_refresh_lock: Optional[asyncio.Lock] = None # one refresh at a time122_key_last_ts: float = 0.0 # last successful refresh123_key_fail_count: int = 0 # consecutive refresh failures124 125JOB_QUEUE: Optional[asyncio.Queue] = None126 127TASKS: Dict[str, Dict[str, Any]] = {}128TASK_QUEUES: Dict[str, asyncio.Queue] = {} # SSE event queues129 130EXECUTOR = ThreadPoolExecutor(max_workers=EXECUTOR_THREADS)131SCRAPER = cloudscraper.create_scraper()132 133# pyvirtualdisplay handle (optional)134VDISPLAY: Optional[Display] = None135 136 137# ═══════════════════════════════════════════════════════════════138# SMALL HELPERS139# ═══════════════════════════════════════════════════════════════140 141def _safe(s: str) -> str:142 """Sanitise string for filenames."""143 ok = set(string.ascii_letters + string.digits + "-_.()")144 return "".join(c if c in ok else "_" for c in s)[:120]145 146 147def _sid() -> str:148 return "".join(random.choices(string.ascii_lowercase + string.digits, k=8))149 150 151def _now() -> str:152 return datetime.utcnow().isoformat(timespec="milliseconds") + "Z"153 154 155def _reqid() -> str:156 return f"{time.time():.6f}-{_sid()}"157 158 159def _stamp() -> str:160 return datetime.utcnow().strftime("%Y%m%dT%H%M%S")161 162 163# ═══════════════════════════════════════════════════════════════164# Virtual display helpers (pyvirtualdisplay)165# ═══════════════════════════════════════════════════════════════166 167def _start_virtual_display_if_needed(headless: bool):168 """169 Start pyvirtualdisplay.Display() if we're running non-headless in an170 environment without DISPLAY and USE_VIRTUAL_DISPLAY is True.171 This function is synchronous and safe to be run in a thread executor.172 """173 global VDISPLAY174 175 if headless:176 log.info("ZD_HEADLESS=True → not starting virtual display")177 return178 179 if not USE_VIRTUAL_DISPLAY:180 log.info("USE_VIRTUAL_DISPLAY=False → not starting virtual display; expecting manual DISPLAY or headless mode.")181 return182 183 if os.environ.get("DISPLAY"):184 log.info("DISPLAY already set: %s", os.environ.get("DISPLAY"))185 return186 187 if not _HAS_PYVIRTUALDISPLAY or Display is None:188 log.warning(189 "pyvirtualdisplay not installed — cannot create virtual DISPLAY. "190 "Install pyvirtualdisplay in your environment to enable Xvfb.")191 return192 193 try:194 VDISPLAY = Display(visible=0, size=(1280, 720))195 VDISPLAY.start()196 # pyvirtualdisplay sets DISPLAY env itself; log for visibility197 log.info("Started virtual display via pyvirtualdisplay (DISPLAY=%s)", os.environ.get("DISPLAY"))198 except Exception as exc:199 VDISPLAY = None200 log.exception("Failed to start virtual display: %s", exc)201 202 203def _stop_virtual_display_if_needed():204 global VDISPLAY205 if VDISPLAY is None:206 return207 try:208 VDISPLAY.stop()209 log.info("Stopped virtual display")210 except Exception:211 log.exception("Error while stopping virtual display")212 finally:213 VDISPLAY = None214 215 216# ═══════════════════════════════════════════════════════════════217# PERCHANCE HTTP CLIENT (blocking – runs in ThreadPoolExecutor)218# ═══════════════════════════════════════════════════════════════219 220class PerchanceClient:221 """All blocking HTTP work against the Perchance API."""222 223 def __init__(self):224 self.base = BASE_URL.rstrip("/")225 self.s = SCRAPER226 self.h = {227 "Accept": "*/*",228 "Content-Type": "application/json;charset=UTF-8",229 "Origin": IMAGE_GEN_ORIGIN,230 "Referer": f"{IMAGE_GEN_ORIGIN}/embed",231 "User-Agent": (232 "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "233 "AppleWebKit/537.36 (KHTML, like Gecko) "234 "Chrome/131.0.0.0 Safari/537.36"235 ),236 }237 238 # ---- low-level helpers ----239 240 def get_ad_code(self) -> str:241 try:242 r = self.s.get(243 f"{self.base}{API_ACCESS_CODE}",244 timeout=HTTP_TIMEOUT, headers=self.h,245 )246 r.raise_for_status()247 return r.text.strip()248 except Exception:249 return ""250 251 def _post(self, body: dict, params: dict) -> dict:252 try:253 r = self.s.post(254 f"{self.base}{API_GENERATE}",255 json=body, params=params,256 timeout=HTTP_TIMEOUT, headers=self.h,257 )258 r.raise_for_status()259 try:260 return r.json()261 except Exception:262 return {"status": "invalid_json", "raw": r.text}263 except Exception as exc:264 return {"status": "fetch_failure", "error": str(exc)}265 266 def _await_prev(self, key: str):267 try:268 self.s.get(269 f"{self.base}{API_AWAIT}",270 params={"userKey": key, "__cacheBust": random.random()},271 timeout=20, headers=self.h,272 )273 except Exception:274 pass275 276 # ---- generate one image ----277 278 def generate_one(279 self, *,280 prompt: str,281 negative_prompt: str = "",282 seed: int = -1,283 resolution: str = "512x768",284 guidance_scale: float = 7.0,285 channel: str = "ai-text-to-image-generator",286 sub_channel: str = "private",287 user_key: str = "",288 ad_access_code: str = "",289 request_id: str = "",290 ) -> dict:291 """292 Returns ONE of:293 {"imageId": ..., "seed": ...}294 {"inline": ..., "seed": ...}295 {"error": "invalid_key"} ← caller must refresh key296 {"error": "<other>", ...}297 """298 request_id = request_id or _reqid()299 params = {300 "userKey": user_key,301 "requestId": request_id,302 "adAccessCode": ad_access_code,303 "__cacheBust": random.random(),304 }305 body = {306 "prompt": prompt,307 "negativePrompt": negative_prompt,308 "seed": seed,309 "resolution": resolution,310 "guidanceScale": guidance_scale,311 "channel": channel,312 "subChannel": sub_channel,313 "userKey": user_key,314 "adAccessCode": ad_access_code,315 "requestId": request_id,316 }317 318 ad_refreshed = False319 320 for att in range(1, MAX_GEN_RETRIES + 1):321 res = self._post(body, params)322 st = res.get("status")323 324 # ---- success ----325 if st == "success":326 iid = res.get("imageId")327 urls = res.get("imageDataUrls")328 if iid:329 log.info("Got imageId: %s", iid)330 return {"imageId": iid, "seed": res.get("seed")}331 if urls:332 return {"inline": urls[0], "seed": res.get("seed")}333 log.error("success but empty payload: %s", str(res)[:300])334 return {"error": "empty_success", "raw": res}335 336 # ---- invalid key → return immediately (do NOT retry here) ----337 if st == "invalid_key":338 log.warning("Server says invalid_key")339 return {"error": "invalid_key"}340 341 # ---- previous request still running ----342 if st == "waiting_for_prev_request_to_finish":343 log.info("Waiting for prev request to finish …")344 self._await_prev(user_key)345 time.sleep(0.3 + random.random() * 0.3)346 continue347 348 # ---- ad access code expired ----349 if st == "invalid_ad_access_code" and not ad_refreshed:350 code = self.get_ad_code()351 if code:352 ad_access_code = code353 params["adAccessCode"] = code354 body["adAccessCode"] = code355 ad_refreshed = True356 log.info("Refreshed ad code → retry")357 time.sleep(0.8)358 continue359 return {"error": "invalid_ad_access_code"}360 361 # ---- transient gen failure ----362 if st == "gen_failure" and res.get("type") == 1:363 log.warning("gen_failure type 1 → retry after 2.5 s")364 time.sleep(2.5)365 continue366 367 # ---- network / stale ----368 if st in (None, "fetch_failure", "invalid_json", "stale_request"):369 log.info("Transient error (status=%s) attempt %d/%d", st, att, MAX_GEN_RETRIES)370 time.sleep(1.0)371 continue372 373 # ---- anything else ----374 log.error("Unhandled status '%s': %s", st, str(res)[:300])375 return {"error": f"unhandled_{st}", "raw": res}376 377 return {"error": "max_retries_exceeded"}378 379 # ---- download ----380 381 def download_image(self, image_id: str, prefix: str = "img") -> str:382 """Poll until the image is ready, save to OUTPUT_DIR, return path."""383 url = f"{self.base}{API_DOWNLOAD}?imageId={image_id}"384 t0 = time.time()385 bk = BACKOFF_INIT386 387 while True:388 elapsed = time.time() - t0389 if elapsed >= MAX_DOWNLOAD_WAIT:390 raise TimeoutError(391 f"Download timed out ({elapsed:.0f}s) for {image_id}"392 )393 try:394 r = self.s.get(url, timeout=HTTP_TIMEOUT,395 headers=self.h, stream=True)396 if r.status_code == 200:397 ct = r.headers.get("Content-Type", "")398 ext = (399 ".png" if "png" in ct else400 ".webp" if "webp" in ct else ".jpg"401 )402 fn = _safe(f"{prefix}_{image_id[:12]}{ext}")403 fp = str(OUTPUT_DIR / fn)404 with open(fp, "wb") as f:405 for chunk in r.iter_content(8192):406 if chunk:407 f.write(chunk)408 log.info("Saved → %s", fp)409 return fp410 except Exception:411 pass412 413 time.sleep(bk)414 bk = min(bk * 1.8, 8.0)415 416 417CLIENT = PerchanceClient()418 419 420# ═══════════════════════════════════════════════════════════════421# ZENDRIVER – browser automation to extract userKey422# ═══════════════════════════════════════════════════════════════423 424async def _cdp_mouse(tab, typ, x, y, **kw):425 await tab.send(426 cdp.input_.dispatch_mouse_event(427 type_=typ, x=float(x), y=float(y), **kw,428 )429 )430 431 432async def _viewport_center(tab):433 try:434 v = await tab.evaluate(435 "(()=>({w:innerWidth,h:innerHeight}))()",436 await_promise=False, return_by_value=True,437 )438 return (v["w"] / 2.0, v["h"] / 2.0)439 except Exception:440 return (600.0, 400.0)441 442 443async def _ls_get(tab, key):444 try:445 return await tab.evaluate(446 f"localStorage&&localStorage.getItem({json.dumps(key)})",447 await_promise=True, return_by_value=True,448 )449 except Exception:450 return None451 452 453async def _clicker_loop(tab, stop: asyncio.Event):454 """Simulate steady centre-clicks on *tab* until *stop* is set."""455 try:456 await tab.evaluate(457 "window.focus&&window.focus()",458 await_promise=False, return_by_value=False,459 )460 except Exception:461 pass462 463 centre = await _viewport_center(tab)464 centre_upd = time.time()465 466 while not stop.is_set():467 if time.time() - centre_upd > 2.5:468 centre = await _viewport_center(tab)469 centre_upd = time.time()470 471 jx = random.uniform(-CLICK_JITTER, CLICK_JITTER)472 jy = random.uniform(-CLICK_JITTER, CLICK_JITTER)473 cx, cy = centre[0] + jx, centre[1] + jy474 475 try:476 await _cdp_mouse(tab, "mouseMoved", cx, cy, pointer_type="mouse")477 await asyncio.sleep(random.uniform(0.02, 0.08))478 await _cdp_mouse(479 tab, "mousePressed", cx, cy,480 button=cdp.input_.MouseButton.LEFT,481 click_count=1, buttons=1,482 )483 await asyncio.sleep(random.uniform(0.03, 0.12))484 await _cdp_mouse(485 tab, "mouseReleased", cx, cy,486 button=cdp.input_.MouseButton.LEFT,487 click_count=1, buttons=0,488 )489 except Exception:490 pass491 492 # interruptible sleep493 try:494 await asyncio.wait_for(495 stop.wait(),496 timeout=CLICK_INTERVAL * random.uniform(0.85, 1.15),497 )498 break499 except asyncio.TimeoutError:500 pass501 502 503async def _poll_for_key(tab, stop: asyncio.Event, max_sec: int):504 """Poll localStorage every 250 ms for a userKey entry."""505 t0 = time.time()506 while not stop.is_set() and (time.time() - t0) < max_sec:507 val = await _ls_get(tab, f"{KEY_PREFIX}-0")508 if val:509 return val510 try:511 keys = await tab.evaluate(512 "Object.keys(localStorage||{}).filter(k=>k.includes('userKey'))",513 await_promise=False, return_by_value=True,514 )515 for k in (keys or []):516 v = await _ls_get(tab, k)517 if v:518 return v519 except Exception:520 pass521 await asyncio.sleep(0.25)522 return None523 524 525async def fetch_key_via_browser(526 timeout: int = ZD_TIMEOUT,527 headless: bool = ZD_HEADLESS,528) -> Optional[str]:529 """530 Launch Chrome → navigate to Perchance → click to trigger531 ad/verification → read userKey from localStorage → close browser.532 Returns the key string or None.533 """534 log.info(535 "Launching browser for userKey (timeout=%ds, headless=%s)",536 timeout, headless,537 )538 539 # If we're in non-headless mode on a display-less host, ensure a virtual540 # DISPLAY is started first. This call is synchronous so we run it in the541 # event loop's default executor when called from async code.542 loop = asyncio.get_running_loop()543 try:544 await loop.run_in_executor(None, partial(_start_virtual_display_if_needed, headless))545 except Exception:546 log.exception("Error while attempting to start virtual display")547 548 try:549 browser = await zd.start(headless=headless)550 except Exception as exc:551 log.exception("Browser start failed: %s", exc)552 return None553 554 stop = asyncio.Event()555 result = None556 557 try:558 page_tab = await browser.get(TARGET_URL)559 log.info("Opened %s", TARGET_URL)560 await asyncio.sleep(2.0)561 562 origin_tab = await browser.get(IMAGE_GEN_ORIGIN, new_tab=True)563 log.info("Opened %s", IMAGE_GEN_ORIGIN)564 await asyncio.sleep(1.0)565 566 await page_tab.bring_to_front()567 await asyncio.sleep(0.5)568 569 clicker = asyncio.create_task(_clicker_loop(page_tab, stop))570 poller = asyncio.create_task(_poll_for_key(origin_tab, stop, timeout))571 572 try:573 done, _ = await asyncio.wait({poller}, timeout=timeout)574 if poller in done:575 result = poller.result()576 finally:577 stop.set()578 if not clicker.done():579 clicker.cancel()580 try:581 await clicker582 except asyncio.CancelledError:583 pass584 585 for t in (origin_tab, page_tab):586 try:587 await t.close()588 except Exception:589 pass590 finally:591 try:592 await browser.stop()593 except Exception:594 pass595 596 if result:597 log.info("Fetched userKey (len=%d)", len(result))598 else:599 log.warning("Could not fetch userKey within %ds", timeout)600 return result601 602 603# ═══════════════════════════════════════════════════════════════604# KEY MANAGEMENT – coordinated refresh across workers605# ═══════════════════════════════════════════════════════════════606 607async def _broadcast(event: dict):608 """Push an event into every active task's SSE queue."""609 for tid, q in TASK_QUEUES.items():610 task = TASKS.get(tid)611 if task and task["status"] in ("queued", "running"):612 try:613 q.put_nowait(event)614 except asyncio.QueueFull:615 pass616 617 618async def refresh_user_key() -> Optional[str]:619 """620 Coordinate a single key refresh. If another coroutine is already621 refreshing, we simply wait for it to finish and return the new key.622 623 Returns the new key string, or None on failure.624 """625 global USER_KEY, _key_last_ts, _key_fail_count626 627 async with _key_refresh_lock:628 # ── double-check: maybe another coroutine just refreshed ──629 age = time.time() - _key_last_ts630 if age < KEY_REFRESH_COOLDOWN and USER_KEY:631 log.info(632 "Key was refreshed %.1fs ago → reusing existing key", age,633 )634 return USER_KEY635 636 # ── too many consecutive failures? ──637 if _key_fail_count >= MAX_REFRESH_FAILURES:638 log.error(639 "Key refresh disabled: %d consecutive failures. "640 "Set key manually via POST /set_user_key",641 _key_fail_count,642 )643 await _broadcast({644 "type": "key_refresh_failed",645 "time": _now(),646 "message": (647 f"Auto-refresh disabled after {_key_fail_count} failures. "648 "Please set userKey manually via /set_user_key"649 ),650 })651 return None652 653 # ── signal "key is being refreshed" ──654 _key_valid.clear()655 log.info("Starting userKey refresh via browser …")656 657 await _broadcast({658 "type": "key_refreshing",659 "time": _now(),660 "message": "UserKey expired — refreshing via browser automation …",661 })662 663 try:664 new_key = await fetch_key_via_browser(665 timeout=ZD_TIMEOUT, headless=ZD_HEADLESS,666 )667 668 if new_key:669 async with _key_lock:670 USER_KEY = new_key671 _key_last_ts = time.time()672 _key_fail_count = 0673 674 log.info("UserKey refreshed OK (len=%d)", len(new_key))675 await _broadcast({676 "type": "key_refreshed",677 "time": _now(),678 "message": "UserKey refreshed – resuming generation.",679 })680 return new_key681 682 # fetch returned None683 _key_fail_count += 1684 log.error(685 "Key refresh returned nothing (failure #%d/%d)",686 _key_fail_count, MAX_REFRESH_FAILURES,687 )688 await _broadcast({689 "type": "key_refresh_failed",690 "time": _now(),691 "message": (692 f"Key refresh failed (attempt {_key_fail_count}"693 f"/{MAX_REFRESH_FAILURES})"694 ),695 })696 return None697 698 except Exception as exc:699 _key_fail_count += 1700 log.exception(701 "Key refresh error (failure #%d/%d): %s",702 _key_fail_count, MAX_REFRESH_FAILURES, exc,703 )704 await _broadcast({705 "type": "key_refresh_failed",706 "time": _now(),707 "message": f"Key refresh error: {exc}",708 })709 return None710 711 finally:712 # ALWAYS unblock waiters, even on failure713 _key_valid.set()714 715 716# ═══════════════════════════════════════════════════════════════717# TASK MODEL718# ═══════════════════════════════════════════════════════════════719 720def create_task(721 prompts: List[str],722 count: int,723 resolution: str,724 guidance: float,725 negative: str,726 sub_channel: str,727) -> dict:728 tid = str(uuid.uuid4())729 task = {730 "id": tid,731 "prompts": prompts,732 "count": count,733 "resolution": resolution,734 "guidance": guidance,735 "negative": negative,736 "sub_channel": sub_channel,737 "created_at": _now(),738 "status": "queued", # queued → running → done / failed739 "total_images": len(prompts) * count,740 "completed": 0,741 "results": [],742 "error": None,743 }744 TASKS[tid] = task745 TASK_QUEUES[tid] = asyncio.Queue()746 return task747 748 749# ═══════════════════════════════════════════════════════════════750# WORKER — image generation + key-refresh retry loop751# ═══════════════════════════════════════════════════════════════752 753async def _save_inline(data_url: str, prompt: str) -> str:754 """Decode base-64 data URL → file. Returns path."""755 loop = asyncio.get_running_loop()756 header, b64 = (data_url.split(",", 1) + [""])[:2] if "," in data_url else ("", data_url)757 ext = ".png" if "png" in header else ".jpg"758 fn = _safe(f"{prompt[:30]}_{_stamp()}_{_sid()}{ext}")759 fp = OUTPUT_DIR / fn760 raw = base64.b64decode(b64)761 await loop.run_in_executor(EXECUTOR, fp.write_bytes, raw)762 log.info("Saved inline → %s", fp)763 return str(fp)764 765 766async def _download(image_id: str, prompt: str) -> str:767 """Download via PerchanceClient (blocking, in executor)."""768 loop = asyncio.get_running_loop()769 prefix = f"{_safe(prompt[:30])}_{_stamp()}_{_sid()}"770 return await loop.run_in_executor(771 EXECUTOR,772 partial(CLIENT.download_image, image_id, prefix),773 )774 775 776async def _generate_single(777 prompt: str,778 task: dict,779 idx: int,780 queue: asyncio.Queue,781 ad_code: str,782) -> Optional[str]:783 """784 Generate + save one image.785 786 On 'invalid_key', triggers a coordinated key refresh and retries787 up to MAX_KEY_RETRIES times. Returns the saved filepath or None.788 """789 loop = asyncio.get_running_loop()790 tid = task["id"]791 792 for key_try in range(1, MAX_KEY_RETRIES + 1):793 794 # ── wait if a refresh is in progress ──795 await _key_valid.wait()796 797 # ── read current key ──798 async with _key_lock:799 active_key = USER_KEY800 801 if not active_key:802 await queue.put({803 "type": "error",804 "time": _now(),805 "task_id": tid,806 "message": "No userKey available. Set via /set_user_key",807 })808 return None809 810 # ── blocking generation in thread-pool ──811 result = await loop.run_in_executor(812 EXECUTOR,813 partial(814 CLIENT.generate_one,815 prompt=prompt,816 negative_prompt=task["negative"],817 seed=-1,818 resolution=task["resolution"],819 guidance_scale=task["guidance"],820 channel="ai-text-to-image-generator",821 sub_channel=task["sub_channel"],822 user_key=active_key,823 ad_access_code=ad_code,824 request_id=_reqid(),825 ),826 )827 828 # ── invalid_key → refresh + retry ──829 if result.get("error") == "invalid_key":830 log.warning(831 "invalid_key for task %s (key_try %d/%d) → refreshing",832 tid, key_try, MAX_KEY_RETRIES,833 )834 await queue.put({835 "type": "key_invalid",836 "time": _now(),837 "task_id": tid,838 "attempt": key_try,839 "max_attempts": MAX_KEY_RETRIES,840 "message": "UserKey invalid — refreshing …",841 })842 843 new_key = await refresh_user_key()844 if new_key:845 # also refresh ad code with fresh key846 ad_code = await loop.run_in_executor(847 EXECUTOR, CLIENT.get_ad_code,848 )849 continue # ← retry generation850 else:851 await queue.put({852 "type": "error",853 "time": _now(),854 "task_id": tid,855 "message": "Could not refresh userKey — aborting image",856 })857 return None858 859 # ── other errors ──860 if result.get("error"):861 log.warning(862 "Gen error task=%s prompt='%.40s': %s",863 tid, prompt, result,864 )865 await queue.put({866 "type": "gen_error",867 "time": _now(),868 "task_id": tid,869 "prompt": prompt,870 "index": idx,871 "error": result,872 })873 return None874 875 # ── success → save ──876 try:877 if result.get("inline"):878 fp = await _save_inline(result["inline"], prompt)879 elif result.get("imageId"):880 fp = await _download(result["imageId"], prompt)881 else:882 log.error("Unexpected result: %s", result)883 return None884 885 seed = result.get("seed")886 task["completed"] += 1887 task["results"].append({888 "prompt": prompt,889 "index": idx,890 "path": fp,891 "seed": seed,892 })893 await queue.put({894 "type": "image_ready",895 "time": _now(),896 "task_id": tid,897 "prompt": prompt,898 "index": idx,899 "path": fp,900 "seed": seed,901 "completed": task["completed"],902 "total": task["total_images"],903 })904 return fp905 906 except Exception as exc:907 log.exception("Save/download error task=%s: %s", tid, exc)908 await queue.put({909 "type": "download_error",910 "time": _now(),911 "task_id": tid,912 "prompt": prompt,913 "index": idx,914 "error": str(exc),915 })916 return None917 918 # exhausted key retries919 log.error("Exhausted key retries for task %s prompt='%.40s'", tid, prompt)920 return None921 922 923async def worker_loop(worker_id: int, semaphore: asyncio.Semaphore):924 """Long-running coroutine: pull jobs → generate images."""925 log.info("Worker %d started", worker_id)926 loop = asyncio.get_running_loop()927 928 while True:929 job = await JOB_QUEUE.get()930 931 # shutdown sentinel932 if job is None:933 log.info("Worker %d shutting down", worker_id)934 JOB_QUEUE.task_done()935 break936 937 task = job["task"]938 tid = task["id"]939 queue = TASK_QUEUES.get(tid)940 941 log.info(942 "Worker %d → task %s (%d images)",943 worker_id, tid, task["total_images"],944 )945 task["status"] = "running"946 if queue:947 await queue.put({948 "type": "started",949 "time": _now(),950 "task_id": tid,951 "total_images": task["total_images"],952 })953 954 # fetch ad code once per task955 ad_code = await loop.run_in_executor(EXECUTOR, CLIENT.get_ad_code)956 957 # heartbeat coroutine958 async def _heartbeat():959 while task["status"] == "running":960 await asyncio.sleep(5.0)961 if queue and task["status"] == "running":962 try:963 queue.put_nowait({964 "type": "heartbeat",965 "time": _now(),966 "task_id": tid,967 "completed": task["completed"],968 "total": task["total_images"],969 })970 except asyncio.QueueFull:971 pass972 973 hb = asyncio.create_task(_heartbeat())974 975 try:976 for prompt in task["prompts"]:977 for i in range(task["count"]):978 async with semaphore:979 await _generate_single(980 prompt, task, i, queue, ad_code,981 )982 if task["status"] == "failed":983 break984 if task["status"] == "failed":985 break986 987 # decide final status988 if task["status"] != "failed":989 if task["completed"] == 0 and task["total_images"] > 0:990 task["status"] = "failed"991 task["error"] = "No images generated successfully"992 else:993 task["status"] = "done"994 995 if queue:996 await queue.put({997 "type": task["status"], # "done" or "failed"998 "time": _now(),999 "task_id": tid,1000 "completed": task["completed"],1001 "total": task["total_images"],1002 "error": task.get("error"),1003 })1004 1005 except Exception as exc:1006 log.exception("Worker %d task %s crashed: %s", worker_id, tid, exc)1007 task["status"] = "failed"1008 task["error"] = str(exc)1009 if queue:1010 await queue.put({1011 "type": "failed",1012 "time": _now(),1013 "task_id": tid,1014 "error": str(exc),1015 })1016 1017 finally:1018 hb.cancel()1019 try:1020 await hb1021 except asyncio.CancelledError:1022 pass1023 1024 if queue:1025 await queue.put({"type": "eof", "time": _now(), "task_id": tid})1026 1027 JOB_QUEUE.task_done()1028 log.info(1029 "Worker %d task %s finished (%s, %d/%d)",1030 worker_id, tid, task["status"],1031 task["completed"], task["total_images"],1032 )1033 1034 1035# ═══════════════════════════════════════════════════════════════1036# FASTAPI – lifespan + app + endpoints1037# ═══════════════════════════════════════════════════════════════1038 1039@asynccontextmanager1040async def lifespan(app: FastAPI):1041 global USER_KEY, _key_lock, _key_valid, _key_refresh_lock1042 global _key_last_ts, _key_fail_count, JOB_QUEUE1043 1044 # ── create asyncio primitives in uvicorn's loop ──1045 _key_lock = asyncio.Lock()1046 _key_valid = asyncio.Event()1047 _key_valid.set() # assume usable initially1048 _key_refresh_lock = asyncio.Lock()1049 JOB_QUEUE = asyncio.Queue(maxsize=MAX_QUEUE_SIZE)1050 1051 # If configured to start virtual display, attempt to do so when needed.1052 loop = asyncio.get_running_loop()1053 try:1054 await loop.run_in_executor(None, partial(_start_virtual_display_if_needed, ZD_HEADLESS))1055 except Exception:1056 log.exception("Failed to ensure virtual display at startup")1057 1058 # ── initial key fetch (skip if already set from __main__) ──1059 if USER_KEY:1060 log.info("Using pre-fetched userKey (len=%d)", len(USER_KEY))1061 _key_last_ts = time.time()1062 else:1063 skip = os.environ.get("NO_INITIAL_FETCH", "") in ("1", "true", "True")1064 if not skip:1065 try:1066 key = await fetch_key_via_browser(1067 timeout=ZD_TIMEOUT, headless=ZD_HEADLESS,1068 )1069 if key:1070 USER_KEY = key1071 _key_last_ts = time.time()1072 log.info("Fetched userKey at startup (len=%d)", len(key))1073 else:1074 log.warning(1075 "Startup key fetch failed. "1076 "Use /set_user_key or /fetch_user_key."1077 )1078 except Exception as exc:1079 log.exception("Startup key fetch error: %s", exc)1080 else:1081 log.info("NO_INITIAL_FETCH=1 → skipping browser key fetch")1082 1083 # ── launch workers ──1084 sem = asyncio.Semaphore(WORKER_COUNT)1085 workers = [1086 asyncio.create_task(worker_loop(i + 1, sem))1087 for i in range(WORKER_COUNT)1088 ]1089 log.info("Launched %d workers", WORKER_COUNT)1090 1091 # ---------- server is running ----------1092 yield1093 # ---------- shutdown begins ------------1094 1095 log.info("Shutdown: sending stop sentinels to workers …")1096 for _ in range(WORKER_COUNT):1097 await JOB_QUEUE.put(None)1098 await asyncio.gather(*workers, return_exceptions=True)1099 1100 try:1101 SCRAPER.close()1102 except Exception:1103 pass1104 EXECUTOR.shutdown(wait=True)1105 1106 # Stop virtual display if we started one1107 try:1108 await loop.run_in_executor(None, _stop_virtual_display_if_needed)1109 except Exception:1110 log.exception("Failed to stop virtual display cleanly")1111 1112 log.info("Shutdown complete")1113 1114 1115# ── app ──1116app = FastAPI(1117 title="Perchance Image Generation Server v2 (pyvirtualdisplay)",1118 lifespan=lifespan,1119)1120app.add_middleware(1121 CORSMiddleware,1122 allow_origins=["*"],1123 allow_methods=["*"],1124 allow_headers=["*"],1125)1126 1127 1128# ───────────── endpoints ─────────────1129 1130@app.get("/health")1131async def health():1132 async with _key_lock:1133 has_key = USER_KEY is not None1134 return {1135 "status": "ok",1136 "has_user_key": has_key,1137 "queue_size": JOB_QUEUE.qsize() if JOB_QUEUE else 0,1138 "active_tasks": sum(1139 1 for t in TASKS.values() if t["status"] in ("queued", "running")1140 ),1141 }1142 1143 1144@app.get("/user_key")1145async def user_key_info():1146 async with _key_lock:1147 has = USER_KEY is not None1148 ln = len(USER_KEY) if has else 01149 return {"has_user_key": has, "key_length": ln}1150 1151 1152@app.post("/set_user_key")1153async def set_user_key(payload: Dict[str, str]):1154 global USER_KEY, _key_last_ts, _key_fail_count1155 key = payload.get("userKey", "").strip()1156 if not key:1157 raise HTTPException(400, "userKey required")1158 async with _key_lock:1159 USER_KEY = key1160 _key_last_ts = time.time()1161 _key_fail_count = 01162 _key_valid.set() # unblock any waiting workers1163 log.info("userKey set via API (len=%d)", len(key))1164 return {"status": "ok", "key_length": len(key)}1165 1166 1167@app.post("/fetch_user_key")1168async def fetch_user_key_endpoint():1169 """Trigger a background browser-based key fetch."""1170 global _key_fail_count1171 1172 async def _bg():1173 global _key_fail_count1174 _key_fail_count = 0 # reset so refresh is allowed1175 await refresh_user_key()1176 1177 asyncio.create_task(_bg())1178 return {"status": "started", "note": "Browser key fetch running in background"}1179 1180 1181@app.post("/generate")1182async def submit_job(payload: Dict[str, Any]):1183 """1184 POST /generate1185 Body:1186 {1187 "prompts": ["a cat in space", "sunset over mountains"],1188 "count": 2,1189 "resolution": "512x768",1190 "guidance": 7.0,1191 "negative": "",1192 "subChannel": "private"1193 }1194 Returns:1195 { "task_id": "...", "stream_url": "/stream/...", "queue_position": N }1196 """1197 prompts = payload.get("prompts") or payload.get("prompt") or []1198 if isinstance(prompts, str):1199 prompts = [prompts]1200 if not isinstance(prompts, list) or not prompts: