"""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 `...` 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 ... 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 "" in reply:
_, after = reply.rsplit("", 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")
@app.get("/health")
def health():
return {
"status": "ok",
"dit_endpoint_configured": bool(DIT_ENDPOINT_URL),
"pe_endpoint_configured": bool(PE_ENDPOINT_URL and PE_ENDPOINT_TOKEN),
}
@app.post("/api/generate_image")
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}"}
@app.post("/api/generate_image_async")
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}
@app.get("/api/task/{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
@app.post("/api/generate_sp")
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 ""
# 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}
@app.post("/api/generate_sp/stream")
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()