Spaces:
Paused
Paused
| """HF Space: Context-Scaling — Structured-Prompt Studio. | |
| sdk: gradio. Gradio owns the launch (its listener is what HF's | |
| frontend expects to see); we attach FastAPI routes onto Gradio's | |
| underlying app and a middleware serves the SPA at `/`. | |
| GET / → static/index.html (SPA) (middleware) | |
| GET /static/* → static assets (mount) | |
| POST /api/generate_image → proxy to DiT Inference Endpoint (route) | |
| POST /api/generate_sp → proxy to PE Inference Endpoint (route) | |
| """ | |
| import os | |
| os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") | |
| import base64 | |
| import io | |
| import json | |
| import re | |
| import threading | |
| import time | |
| import traceback | |
| from pathlib import Path | |
| import asyncio | |
| import httpx | |
| import gradio as gr | |
| async def _sleep_async(seconds: float): | |
| await asyncio.sleep(seconds) | |
| from fastapi.responses import FileResponse, JSONResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from pydantic import BaseModel, Field | |
| DIT_BASE_MODEL_ID = os.environ.get("DIT_BASE_MODEL_ID", "Qwen/Qwen-Image") | |
| # SP overlay ships as sharded safetensors (each ~4.5 GB, under the LFS | |
| # per-file limit corp proxies choke on). Two glob patterns pick out the | |
| # DiT vs text-encoder shard sets; the module prefix has already been | |
| # stripped inside each shard. | |
| # | |
| # NOTE: the EvalKit serve_qwenimage.py only loads `dit_model.*` from the | |
| # training checkpoint and leaves text_encoder as the base pipeline's | |
| # original — so we do the same here. Set DIT_SP_TEXT_ENCODER_GLOB to a | |
| # non-empty value in Space secrets if you want to overlay text_encoder too. | |
| DIT_SP_REPO = os.environ.get("DIT_SP_REPO", "heheyas/Qwen-Image-SP") | |
| DIT_SP_TRANSFORMER_GLOB = os.environ.get("DIT_SP_TRANSFORMER_GLOB", "dit_model-*.safetensors") | |
| DIT_SP_TEXT_ENCODER_GLOB = os.environ.get("DIT_SP_TEXT_ENCODER_GLOB", "") | |
| STATIC_DIR = Path(__file__).resolve().parent / "static" | |
| # --------------------------------------------------------------------------- | |
| # PE (prompt-expansion) config — reads Space secrets set via HF API: | |
| # PE_ENDPOINT_URL e.g. https://xxx.us-east-1.aws.endpoints.huggingface.cloud | |
| # PE_ENDPOINT_TOKEN HF token with read access to the private endpoint | |
| # PE_MODEL_ID served-model id shown in /v1/models (default: our repo) | |
| # PE_MAX_TOKENS completion cap; SP + CoT can be big (default: 16384) | |
| # --------------------------------------------------------------------------- | |
| PE_ENDPOINT_URL = os.environ.get("PE_ENDPOINT_URL", "").rstrip("/") | |
| PE_ENDPOINT_TOKEN = os.environ.get("PE_ENDPOINT_TOKEN", "") | |
| PE_MODEL_ID = os.environ.get("PE_MODEL_ID", "heheyas/SP-PE-Qwen3.5-35B-A3B") | |
| PE_MAX_TOKENS = int(os.environ.get("PE_MAX_TOKENS", "16384")) | |
| # Sampling defaults match EvalKit's `qwen35-397b` preset used by | |
| # rewrite_and_generate.py (the eval script users benchmark this ckpt with). | |
| PE_TEMPERATURE = float(os.environ.get("PE_TEMPERATURE", "0.6")) | |
| PE_TOP_P = float(os.environ.get("PE_TOP_P", "0.95")) | |
| PE_TOP_K = int(os.environ.get("PE_TOP_K", "20")) | |
| PE_MIN_P = float(os.environ.get("PE_MIN_P", "0.0")) | |
| PE_PRESENCE_PENALTY = float(os.environ.get("PE_PRESENCE_PENALTY", "0.0")) | |
| PE_REPETITION_PENALTY = float(os.environ.get("PE_REPETITION_PENALTY", "1.0")) | |
| # Qwen3 CoT — required for the `<think>...</think>` wrapper my | |
| # _extract_sp_json parser looks for. | |
| PE_ENABLE_THINKING = os.environ.get("PE_ENABLE_THINKING", "1") != "0" | |
| PE_SYSTEM_PROMPT_PATH = Path(__file__).resolve().parent / "system_prompts" / "pe.txt" | |
| PE_SYSTEM_PROMPT = ( | |
| PE_SYSTEM_PROMPT_PATH.read_text(encoding="utf-8") | |
| if PE_SYSTEM_PROMPT_PATH.is_file() else "" | |
| ) | |
| def _split_dupe_id_dict(pairs): | |
| """If a JSON object's key list has repeated `id`, split into | |
| multiple dicts (each starting at an `id` key). Fine-tuned SP model | |
| sometimes drops the `},{` separator between elements and dumps them | |
| all as one dict with duplicated keys — the default json.loads | |
| behaviour (last-wins) collapses N elements into 1. | |
| """ | |
| if not any(k == "id" for k, _ in pairs): | |
| return dict(pairs) | |
| id_indices = [i for i, (k, _) in enumerate(pairs) if k == "id"] | |
| if len(id_indices) <= 1: | |
| return dict(pairs) | |
| # Multi-id → split | |
| out = [] | |
| id_indices.append(len(pairs)) # sentinel | |
| for a, b in zip(id_indices, id_indices[1:]): | |
| out.append(dict(pairs[a:b])) | |
| return out # list of dicts | |
| def _walk_normalize(obj): | |
| """After json.loads with object_pairs_hook=_split_dupe_id_dict, any | |
| 'elements' value might be a list containing an inner list (from the | |
| split). Flatten that so `elements` is a proper flat list of dicts.""" | |
| if isinstance(obj, dict): | |
| for k, v in list(obj.items()): | |
| if isinstance(v, list): | |
| flat = [] | |
| for item in v: | |
| if isinstance(item, list): | |
| flat.extend(item) | |
| else: | |
| flat.append(item) | |
| obj[k] = [_walk_normalize(x) for x in flat] | |
| else: | |
| obj[k] = _walk_normalize(v) | |
| return obj | |
| def _extract_sp_json(reply: str): | |
| """Pull the SP JSON out of the model's reply. | |
| The PE model wraps its reasoning in <think>...</think> and emits the | |
| JSON blueprint *after* the closing tag. Fall back to first-{ to last-} | |
| if the tag is missing (e.g. the model skipped CoT for a short prompt). | |
| Two-stage repair: | |
| 1. `_split_dupe_id_dict` via object_pairs_hook — handles the SP | |
| model's habit of dumping all elements into one dict with repeated | |
| `id` keys (missing `},{` between elements). | |
| 2. `json_repair` as a fallback for other malformations (trailing | |
| commas, unclosed quotes, single quotes, unquoted keys, etc.). | |
| """ | |
| if "</think>" in reply: | |
| _, after = reply.rsplit("</think>", 1) | |
| candidate = after.strip() | |
| else: | |
| candidate = reply | |
| m = re.search(r"\{[\s\S]*\}", candidate) | |
| if not m: | |
| return None | |
| text = m.group(0) | |
| # Stage 1: straight json.loads with our pairs hook. | |
| try: | |
| parsed = json.loads(text, object_pairs_hook=_split_dupe_id_dict) | |
| return _walk_normalize(parsed) if isinstance(parsed, dict) else parsed | |
| except json.JSONDecodeError: | |
| pass | |
| # Stage 2: json_repair rewrites the text into strict JSON, then we | |
| # re-run the pairs hook so the split-dupes step still applies. | |
| try: | |
| from json_repair import repair_json | |
| fixed = repair_json(text) | |
| parsed = json.loads(fixed, object_pairs_hook=_split_dupe_id_dict) | |
| return _walk_normalize(parsed) if isinstance(parsed, dict) else parsed | |
| except Exception: | |
| return None | |
| def _to_compact_single_quote(prompt: str) -> str: | |
| """Match EvalKit's inference.generate.compact_single_quote_json: | |
| parse the incoming string as JSON, escape any inner quotes with | |
| placeholders, dump compact (no spaces) into standard JSON, then | |
| swap structural " → ' and restore inner quotes to \\' and ". | |
| If the prompt isn't parseable JSON (raw NL text OR already in | |
| single-quote format), return as-is — the model was trained to | |
| accept both. | |
| """ | |
| try: | |
| parsed = json.loads(prompt) | |
| except (json.JSONDecodeError, TypeError): | |
| return prompt | |
| PH_S, PH_D = "@@SP_SINGLE_QUOTE@@", "@@SP_DOUBLE_QUOTE@@" | |
| def protect(o): | |
| if isinstance(o, dict): | |
| return {(protect(k) if isinstance(k, str) else k): protect(v) | |
| for k, v in o.items()} | |
| if isinstance(o, list): | |
| return [protect(x) for x in o] | |
| if isinstance(o, str): | |
| return o.replace("'", PH_S).replace('"', PH_D) | |
| return o | |
| dumped = json.dumps(protect(parsed), separators=(",", ":"), ensure_ascii=False) | |
| return dumped.replace('"', "'").replace(PH_S, "\\'").replace(PH_D, '"') | |
| DIT_ENDPOINT_URL = os.environ.get("DIT_ENDPOINT_URL", "").rstrip("/") | |
| DIT_ENDPOINT_TOKEN = os.environ.get("DIT_ENDPOINT_TOKEN", "") or \ | |
| os.environ.get("HF_TOKEN", "") | |
| # Endpoint init downloads ~90G + builds pipe → wall time up to ~10min | |
| # on cold start. Keep request timeout generous. | |
| DIT_REQUEST_TIMEOUT = float(os.environ.get("DIT_REQUEST_TIMEOUT", "1200")) | |
| def _run_dit(prompt: str, height: int, width: int, num_steps: int, | |
| cfg_scale: float, seed: int, negative_prompt: str = "") -> bytes: | |
| """Proxy DiT to the dedicated HF Inference Endpoint (heheyas/sp-dit-l40s). | |
| Uses the same 'inputs' contract as HF handler.py convention. Retries a | |
| handful of times while the endpoint scales up from zero.""" | |
| if not DIT_ENDPOINT_URL: | |
| raise RuntimeError("DIT_ENDPOINT_URL not configured " | |
| "(set it as a Space secret).") | |
| body = { | |
| "inputs": { | |
| "prompt": prompt, | |
| "negative_prompt": negative_prompt or "", | |
| "width": int(width), | |
| "height": int(height), | |
| "num_steps": int(num_steps), | |
| "seed": int(seed), | |
| "cfg_scale": float(cfg_scale), | |
| } | |
| } | |
| headers = {"Content-Type": "application/json"} | |
| if DIT_ENDPOINT_TOKEN: | |
| headers["Authorization"] = f"Bearer {DIT_ENDPOINT_TOKEN}" | |
| # Cold-start retry: endpoint returns 503 while scaling up; POST once | |
| # per 15s until warm or ~5min elapsed. | |
| last_status, last_body = None, None | |
| for attempt in range(20): | |
| try: | |
| with httpx.Client(timeout=DIT_REQUEST_TIMEOUT) as client: | |
| resp = client.post(DIT_ENDPOINT_URL, json=body, headers=headers) | |
| except httpx.HTTPError as e: | |
| print(f"[DiT] proxy attempt {attempt+1}: transport error {e!r}", | |
| flush=True) | |
| time.sleep(15) | |
| continue | |
| if resp.status_code == 200: | |
| data = resp.json() | |
| if "error" in data: | |
| raise RuntimeError(f"DiT endpoint returned error: {data['error']}") | |
| b64 = data.get("image_base64") | |
| if not b64: | |
| raise RuntimeError(f"DiT endpoint response missing image_base64: " | |
| f"{str(data)[:200]}") | |
| return base64.b64decode(b64) | |
| last_status, last_body = resp.status_code, resp.text[:200] | |
| if resp.status_code in (502, 503, 504): | |
| print(f"[DiT] proxy attempt {attempt+1}: HTTP {resp.status_code} " | |
| f"(cold-starting?); retry in 15s", flush=True) | |
| time.sleep(15) | |
| continue | |
| # Non-retryable | |
| raise RuntimeError(f"DiT endpoint HTTP {resp.status_code}: {last_body}") | |
| raise RuntimeError(f"DiT endpoint never returned 200 after 20 retries. " | |
| f"Last: HTTP {last_status} {last_body}") | |
| def _run_dit_pil(prompt, height, width, num_steps, cfg_scale, seed): | |
| from PIL import Image | |
| png = _run_dit(prompt, int(height), int(width), int(num_steps), | |
| float(cfg_scale), int(seed)) | |
| return Image.open(io.BytesIO(png)) | |
| # --------------------------------------------------------------------------- | |
| # Gradio Blocks — minimal smoke UI; the real app lives at /. | |
| # --------------------------------------------------------------------------- | |
| with gr.Blocks(title="Context-Scaling — DiT smoke") as demo: | |
| gr.Markdown("Fallback DiT smoke test. The main app is at [`/`](/).") | |
| with gr.Row(): | |
| with gr.Column(): | |
| _p = gr.Textbox(label="Prompt", lines=2, | |
| value="a red apple on a wooden table") | |
| _h = gr.Slider(512, 1536, value=1024, step=64, label="height") | |
| _w = gr.Slider(512, 1536, value=1024, step=64, label="width") | |
| _s = gr.Slider(10, 50, value=20, step=1, label="steps") | |
| _c = gr.Slider(1.0, 10.0, value=4.0, step=0.1, label="cfg") | |
| _sd = gr.Slider(0, 2**31 - 1, value=42, step=1, label="seed") | |
| _btn = gr.Button("Generate", variant="primary") | |
| _img = gr.Image(label="output", type="pil") | |
| _btn.click(_run_dit_pil, [_p, _h, _w, _s, _c, _sd], _img) | |
| # --------------------------------------------------------------------------- | |
| # Launch Gradio (non-blocking), then attach FastAPI routes + a middleware | |
| # that intercepts GET / and returns the SPA (Gradio otherwise owns /). | |
| # --------------------------------------------------------------------------- | |
| def _attach_fastapi_routes(app): | |
| """Add SPA middleware + /api/* routes + /static/* mount to gradio's app.""" | |
| class ImageRequest(BaseModel): | |
| prompt: str | |
| height: int = Field(1024, ge=64, le=2048) | |
| width: int = Field(1024, ge=64, le=2048) | |
| num_steps: int = Field(25, ge=1, le=100) | |
| seed: int = 42 | |
| cfg_scale: float = Field(4.0, ge=0.0, le=20.0) | |
| negative_prompt: str = "" | |
| class SPRequest(BaseModel): | |
| user_prompt: str | |
| width: int = 1024 | |
| height: int = 1024 | |
| # Can't add middleware after launch (FastAPI freezes it), so prepend | |
| # a Route for GET / that takes precedence over gradio's own /. | |
| from starlette.routing import Route | |
| async def _serve_spa(request): | |
| index = STATIC_DIR / "index.html" | |
| if index.is_file(): | |
| return FileResponse(str(index)) | |
| return JSONResponse(status_code=500, content={"error": f"missing {index}"}) | |
| app.router.routes.insert(0, Route("/", _serve_spa, methods=["GET", "HEAD"])) | |
| if STATIC_DIR.is_dir(): | |
| app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static") | |
| def health(): | |
| return { | |
| "status": "ok", | |
| "dit_endpoint_configured": bool(DIT_ENDPOINT_URL), | |
| "pe_endpoint_configured": bool(PE_ENDPOINT_URL and PE_ENDPOINT_TOKEN), | |
| } | |
| def generate_image(req: ImageRequest): | |
| try: | |
| # DiT was trained on compact single-quote SP strings; if the | |
| # incoming prompt is standard JSON (whether from the SPA or a | |
| # direct API caller), normalize it. Non-JSON strings and | |
| # already-single-quote SP are passed through untouched. | |
| dit_prompt = _to_compact_single_quote(req.prompt) | |
| png = _run_dit( | |
| prompt=dit_prompt, height=req.height, width=req.width, | |
| num_steps=req.num_steps, cfg_scale=req.cfg_scale, | |
| seed=req.seed, negative_prompt=req.negative_prompt, | |
| ) | |
| return {"image_base64": base64.b64encode(png).decode("ascii")} | |
| except Exception as e: | |
| traceback.print_exc() | |
| return JSONResponse(status_code=500, content={"error": f"DiT failed: {e}"}) | |
| # ── Async job pattern to bypass HF's edge 10-min HTTP keep-alive cap. | |
| # POST /api/generate_image_async → {task_id} | |
| # GET /api/task/{id} → {status, image_base64|error} | |
| import uuid, threading as _t | |
| _TASKS: dict[str, dict] = {} | |
| _TASKS_LOCK = _t.Lock() | |
| def _run_dit_task(task_id: str, req: ImageRequest): | |
| try: | |
| dit_prompt = _to_compact_single_quote(req.prompt) | |
| png = _run_dit( | |
| prompt=dit_prompt, height=req.height, width=req.width, | |
| num_steps=req.num_steps, cfg_scale=req.cfg_scale, | |
| seed=req.seed, negative_prompt=req.negative_prompt, | |
| ) | |
| with _TASKS_LOCK: | |
| _TASKS[task_id] = {"status": "done", | |
| "image_base64": base64.b64encode(png).decode("ascii")} | |
| except Exception as e: | |
| traceback.print_exc() | |
| with _TASKS_LOCK: | |
| _TASKS[task_id] = {"status": "failed", "error": f"DiT failed: {e}"} | |
| def generate_image_async(req: ImageRequest): | |
| task_id = uuid.uuid4().hex | |
| with _TASKS_LOCK: | |
| _TASKS[task_id] = {"status": "running"} | |
| _t.Thread(target=_run_dit_task, args=(task_id, req), daemon=True).start() | |
| return {"task_id": task_id} | |
| def get_task(task_id: str): | |
| with _TASKS_LOCK: | |
| result = _TASKS.get(task_id) | |
| if result is None: | |
| return JSONResponse(status_code=404, content={"error": "unknown task_id"}) | |
| return result | |
| def generate_sp(req: SPRequest): | |
| # Config sanity | |
| if not PE_ENDPOINT_URL: | |
| return JSONResponse(status_code=501, content={ | |
| "error": "PE endpoint is not configured " | |
| "(set PE_ENDPOINT_URL as a Space secret)."}) | |
| if not PE_ENDPOINT_TOKEN: | |
| return JSONResponse(status_code=501, content={ | |
| "error": "PE endpoint token is not configured " | |
| "(set PE_ENDPOINT_TOKEN as a Space secret)."}) | |
| if not PE_SYSTEM_PROMPT: | |
| return JSONResponse(status_code=500, content={ | |
| "error": f"missing system prompt at {PE_SYSTEM_PROMPT_PATH}"}) | |
| # Match EvalKit's L5 detail-progression eval script exactly: | |
| # --input_template "<prompt>" | |
| # i.e. the raw user prompt goes in, no width/height suffix; the | |
| # PE model picks its own aspect ratio via CoT Stage A. | |
| user_message = req.user_prompt.strip() | |
| payload = { | |
| "model": PE_MODEL_ID, | |
| "messages": [ | |
| {"role": "system", "content": PE_SYSTEM_PROMPT}, | |
| {"role": "user", "content": user_message}, | |
| ], | |
| "temperature": PE_TEMPERATURE, | |
| "top_p": PE_TOP_P, | |
| "top_k": PE_TOP_K, | |
| "min_p": PE_MIN_P, | |
| "presence_penalty": PE_PRESENCE_PENALTY, | |
| "repetition_penalty": PE_REPETITION_PENALTY, | |
| "max_tokens": PE_MAX_TOKENS, | |
| # Qwen3 opt-in CoT — matches eval script's --reasoning flag. | |
| "chat_template_kwargs": {"enable_thinking": PE_ENABLE_THINKING}, | |
| } | |
| headers = { | |
| "Authorization": f"Bearer {PE_ENDPOINT_TOKEN}", | |
| "Content-Type": "application/json", | |
| } | |
| # Timeout budget: single generation ~1min on A100; the retry loop | |
| # below covers scale-from-zero cold starts (503 for ~5min). | |
| import time as _time | |
| r = None | |
| last_status = None | |
| # Up to ~8 min total wall time: 15 tries × 30s sleep between 503s. | |
| for attempt in range(15): | |
| try: | |
| with httpx.Client(timeout=httpx.Timeout(300.0, connect=30.0)) as client: | |
| r = client.post( | |
| f"{PE_ENDPOINT_URL}/v1/chat/completions", | |
| headers=headers, json=payload, | |
| ) | |
| except httpx.TimeoutException as e: | |
| # A real 300s inference timeout — don't retry, surface it. | |
| return JSONResponse(status_code=504, content={ | |
| "error": f"PE endpoint timed out mid-generation: {e}"}) | |
| except Exception as e: | |
| traceback.print_exc() | |
| return JSONResponse(status_code=502, content={ | |
| "error": f"PE endpoint request failed: {type(e).__name__}: {e}"}) | |
| # 503 = scaled-to-zero cold start; retry until endpoint is up. | |
| if r.status_code == 503: | |
| last_status = 503 | |
| print(f"[PE] endpoint 503 (cold start), retry {attempt+1}/15 in 30s", | |
| flush=True) | |
| _time.sleep(30) | |
| continue | |
| break | |
| if r is None or r.status_code >= 400: | |
| code = r.status_code if r is not None else last_status | |
| body = (r.text[:400] if r is not None else "no response after retries") | |
| return JSONResponse(status_code=502, content={ | |
| "error": f"PE endpoint HTTP {code}: {body}"}) | |
| try: | |
| body = r.json() | |
| reply = body["choices"][0]["message"]["content"] | |
| except Exception as e: | |
| return JSONResponse(status_code=502, content={ | |
| "error": f"PE endpoint returned unparseable body: {e}", | |
| "raw": r.text[:400], | |
| }) | |
| sp = _extract_sp_json(reply) | |
| if sp is None: | |
| # Return the raw text so the UI can at least display something; | |
| # the SP editor's textarea mode will show it verbatim. | |
| return {"sp": reply} | |
| return {"sp": sp} | |
| async def generate_sp_stream(req: SPRequest): | |
| """Server-sent-events streaming variant of /api/generate_sp. | |
| Streams the vLLM endpoint's raw deltas (raw model text — CoT + | |
| JSON) back to the browser as SSE `data:` frames, so the SPA can | |
| render tokens incrementally. Emits a final `event: done` frame | |
| with the extracted SP JSON (or raw text on parse failure). | |
| """ | |
| from starlette.responses import StreamingResponse | |
| import json as _json | |
| import time as _time | |
| if not PE_ENDPOINT_URL or not PE_ENDPOINT_TOKEN or not PE_SYSTEM_PROMPT: | |
| async def _err_only(): | |
| msg = "PE endpoint or system prompt not configured" | |
| yield f"event: error\ndata: {_json.dumps({'error': msg})}\n\n" | |
| return StreamingResponse(_err_only(), media_type="text/event-stream") | |
| payload = { | |
| "model": PE_MODEL_ID, | |
| "messages": [ | |
| {"role": "system", "content": PE_SYSTEM_PROMPT}, | |
| {"role": "user", "content": req.user_prompt.strip()}, | |
| ], | |
| "temperature": PE_TEMPERATURE, "top_p": PE_TOP_P, | |
| "top_k": PE_TOP_K, "min_p": PE_MIN_P, | |
| "presence_penalty": PE_PRESENCE_PENALTY, | |
| "repetition_penalty": PE_REPETITION_PENALTY, | |
| "max_tokens": PE_MAX_TOKENS, | |
| "chat_template_kwargs": {"enable_thinking": PE_ENABLE_THINKING}, | |
| # vLLM OpenAI-compatible streaming | |
| "stream": True, | |
| } | |
| headers = { | |
| "Authorization": f"Bearer {PE_ENDPOINT_TOKEN}", | |
| "Content-Type": "application/json", | |
| "Accept": "text/event-stream", | |
| } | |
| async def _relay(): | |
| # Cold-start retry: hold the SSE connection open, retry the | |
| # POST until endpoint responds non-503. | |
| full_reply = [] | |
| import httpx as _httpx | |
| for attempt in range(15): | |
| try: | |
| async with _httpx.AsyncClient( | |
| timeout=_httpx.Timeout(600.0, connect=30.0) | |
| ) as client: | |
| async with client.stream( | |
| "POST", | |
| f"{PE_ENDPOINT_URL}/v1/chat/completions", | |
| headers=headers, json=payload, | |
| ) as resp: | |
| if resp.status_code == 503: | |
| yield (f"event: waking\ndata: " | |
| f"{_json.dumps({'attempt': attempt+1, 'msg': 'endpoint cold-starting, retry in 30s'})}\n\n") | |
| await _sleep_async(30) | |
| continue | |
| if resp.status_code >= 400: | |
| body = (await resp.aread()).decode( | |
| "utf-8", errors="ignore")[:400] | |
| yield (f"event: error\ndata: " | |
| f"{_json.dumps({'error': f'HTTP {resp.status_code}: {body}'})}\n\n") | |
| return | |
| async for raw in resp.aiter_lines(): | |
| if not raw or not raw.startswith("data: "): | |
| continue | |
| data = raw[6:] | |
| if data.strip() == "[DONE]": | |
| break | |
| try: | |
| chunk = _json.loads(data) | |
| delta = (chunk.get("choices") or [{}])[0].get("delta", {}) | |
| content = delta.get("content") or "" | |
| except Exception: | |
| content = "" | |
| if content: | |
| full_reply.append(content) | |
| yield (f"event: token\ndata: " | |
| f"{_json.dumps({'text': content})}\n\n") | |
| break # success — leave retry loop | |
| except _httpx.TimeoutException: | |
| yield (f"event: error\ndata: " | |
| f"{_json.dumps({'error': 'endpoint stream timeout'})}\n\n") | |
| return | |
| except Exception as e: | |
| yield (f"event: error\ndata: " | |
| f"{_json.dumps({'error': f'{type(e).__name__}: {e}'})}\n\n") | |
| return | |
| else: | |
| yield (f"event: error\ndata: " | |
| f"{_json.dumps({'error': 'endpoint 503 after 15 retries'})}\n\n") | |
| return | |
| # After stream completes, extract JSON and emit `done` frame. | |
| reply_text = "".join(full_reply) | |
| sp = _extract_sp_json(reply_text) | |
| done_payload = {"sp": sp} if sp is not None else {"sp": reply_text} | |
| yield f"event: done\ndata: {_json.dumps(done_payload)}\n\n" | |
| return StreamingResponse(_relay(), media_type="text/event-stream") | |
| # Non-blocking launch. `demo.app` becomes the underlying FastAPI once the | |
| # background server is up. | |
| demo.queue().launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| ssr_mode=False, | |
| prevent_thread_lock=True, | |
| ) | |
| # demo.app is created inside launch(); attach our extras once it exists. | |
| if demo.app is not None: | |
| _attach_fastapi_routes(demo.app) | |
| # Keep the main thread alive; gradio's server runs in a daemon thread. | |
| threading.Event().wait() | |