"""JoyAI-Video-Edit — streaming v2v editing on HF ZeroGPU. Port of the standalone uvicorn server (xvideo/serving/serve_joyomni_streaming.py) to ZeroGPU's fork model, without modifying the code it vendors: * the main web process holds zero CUDA state — a thin byte pipe between the browser WebSocket and two multiprocessing.Queues; * each session runs ONE @spaces.GPU generator, acquired through the ladder in ws(); the entire server-side machine (gate, PE, streaming session, output pump) runs unchanged inside the fork. Hardware: ZeroGPU, JOYOMNI_GPU_SIZE env picks the tier — xlarge (default, full RTX Pro 6000 Blackwell, 96GB) or large (half-card MIG, 48GB + JOYOMNI_LOW_VRAM=1); both are sm_120, matching the joyomni_ops wheel in wheels/. """ from __future__ import annotations import os import shutil from pathlib import Path # ZeroGPU's libstdc++ predates GCC 13; preload ours RTLD_GLOBAL before any native import. def _preload_libstdcxx() -> None: import ctypes cand = Path(__file__).resolve().parent / "libs" / "libstdc++.so.6" if not cand.is_file(): return try: ctypes.CDLL(str(cand), mode=ctypes.RTLD_GLOBAL) print(f"[boot] preloaded {cand}", flush=True) except OSError as e: print(f"[boot] WARNING: could not preload libstdc++: {e}", flush=True) _preload_libstdcxx() # Weights persist on /data; compile/JIT caches stay ephemeral (forks never compile). _DATA = Path("/data") if Path("/data").is_dir() and os.access("/data", os.W_OK) else (Path(__file__).resolve().parent / "deps") shutil.rmtree(_DATA / "cache", ignore_errors=True) # legacy compile-cache dir os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") os.environ.setdefault("GRADIO_SSR_MODE", "false") # Kernel path: FP8 + CUDA graph on plain cuDNN SDPA — on the Space GPU (RTX PRO 6000 # Blackwell) cuDNN is the fastest attention at the serving shapes. os.environ.setdefault("JOYOMNI_FP8_IMG", "1") os.environ.setdefault("JOYOMNI_FP8_TXT", "1") os.environ.setdefault("JOYOMNI_CUDA_GRAPH", "1") # Module-scope load runs under CUDA emulation (no real GPU) — skip the load-time # warmup; each session fork warms up as it runs. The standalone server leaves this unset. os.environ.setdefault("JOYOMNI_SKIP_LOAD_WARMUP", "1") # Do NOT set JOYOMNI_SKIP_VAE_COMPILE_WARMUP here: skipping the load-time VAE # warmups looks like a free ~6 min boot saving, but ZeroGPU's tensor packing hangs # without them (the warmup forwards materialize state packing depends on). # Empirically verified 2026-08-18 over five build cycles. import sys DEPLOY_ROOT = Path(__file__).resolve().parent if str(DEPLOY_ROOT) not in sys.path: sys.path.insert(0, str(DEPLOY_ROOT)) import asyncio import base64 import json import queue import re import tempfile import threading import time import traceback from multiprocessing import Queue as MPQueue from types import SimpleNamespace import warnings warnings.filterwarnings("ignore", message="ZeroGPU: Cannot get Gradio app Queue instance") import spaces CKPT_ROOT = Path(os.environ.get("JOYOMNI_CKPT_ROOT", str(_DATA / "checkpoints"))) # 60s = 120s billed on xlarge. SESSION_DURATION = int(os.environ.get("JOYOMNI_SPACES_DURATION", "60")) # xlarge = full RTX Pro 6000 (96GB); large = half-card MIG slice (48GB, needs # JOYOMNI_LOW_VRAM=1) — the larger pool when the xlarge tier is saturated. GPU_SIZE = os.environ.get("JOYOMNI_GPU_SIZE", "xlarge") # Prebuilt sm_120 wheel: pip runs before repo files exist on Spaces — install here. def _ensure_local_wheels() -> None: import importlib.util import subprocess wheel_dir = DEPLOY_ROOT / "wheels" specs = [ ("joyomni_ops", "joyomni_ops-0.1.0-cp310-cp310-linux_x86_64.whl"), ] for mod, whl in specs: if importlib.util.find_spec(mod) is not None: continue # already installed (warm container) path = wheel_dir / whl if not path.is_file(): print(f"[boot] WARNING: wheel missing, cannot install {mod}: {path}", flush=True) continue print(f"[boot] installing {mod} from {whl} ...", flush=True) subprocess.check_call( [sys.executable, "-m", "pip", "install", "--no-deps", "--no-index", str(path)] ) _ensure_local_wheels() # Pull weights to /data on first boot; later boots skip. ONNX detectors optional. def _resolve_weights() -> dict: CKPT_ROOT.mkdir(parents=True, exist_ok=True) jve_local = CKPT_ROOT / "JoyAI-Video-Edit" mimo_local = CKPT_ROOT / "MiMo-VL-7B-RL-2508" # Skip snapshot_download when the key files are already on /data (~30s less # cold-start). Either DIT (0811 preferred, 0804 fallback) satisfies readiness. _dit_0811 = jve_local / "dit" / "joyai_video_edit_dit_0811.pth" _dit_0804 = jve_local / "dit" / "joyai_video_edit_dit_0804.pth" jve_ready = (_dit_0811.is_file() or _dit_0804.is_file()) and \ (jve_local / "vae" / "diffusion_pytorch_model.safetensors").is_file() mimo_ready = (mimo_local / "config.json").is_file() and \ any(mimo_local.glob("*.safetensors")) if jve_ready and mimo_ready: print("[boot] weights already in /data — skipping hub check entirely", flush=True) else: from huggingface_hub import snapshot_download snapshot_download("jdopensource/JoyAI-Video-Edit", local_dir=str(jve_local)) snapshot_download("XiaomiMiMo/MiMo-VL-7B-RL-2508", local_dir=str(mimo_local)) # Ensure 0811 specifically via a targeted single-file fetch; keep 0804 if absent. if not _dit_0811.is_file(): try: from huggingface_hub import hf_hub_download hf_hub_download("jdopensource/JoyAI-Video-Edit", "dit/joyai_video_edit_dit_0811.pth", local_dir=str(jve_local)) except Exception as e: # noqa: BLE001 print(f"[boot] 0811 DIT not fetched ({type(e).__name__}: {e}); using 0804", flush=True) _dit = _dit_0811 if _dit_0811.is_file() else _dit_0804 return { "dit": str(_dit), "vae": str(jve_local / "vae"), "text_encoder": str(mimo_local), "face_onnx": str(DEPLOY_ROOT / "detectors" / "face_detection_yunet_2023mar.onnx"), "person_onnx": str(DEPLOY_ROOT / "detectors" / "yolov8n.onnx"), } # ZeroGPU requires models on cuda at module scope (packed once, mapped into forks). _RUNTIME = None _PATHS: dict | None = None _BUILD_ERROR: str | None = None def _build_runtime(): from xvideo.inductor_autotune_fix import install as _install_autotune_fix from xvideo.serving.joyomni_streaming import JoyOmniRuntime _install_autotune_fix() paths = _resolve_weights() rt = JoyOmniRuntime.load( paths["dit"], vae_ckpt=paths["vae"], text_encoder_ckpt=paths["text_encoder"], device="cuda", vae_encode_device="cuda", vae_decode_device="cuda", vae_pseudo_device="cuda", postprocess_device="cuda", seed=42, warmup_height=480, warmup_width=840, ) return rt, paths try: print("[boot] building JoyAI-Video-Edit runtime (first boot downloads ~63GB to /data)...", flush=True) _t0 = time.perf_counter() _RUNTIME, _PATHS = _build_runtime() print(f"[boot] runtime ready in {time.perf_counter() - _t0:.1f}s", flush=True) except Exception as e: # noqa: BLE001 _BUILD_ERROR = f"{type(e).__name__}: {e}\n{traceback.format_exc()}" print("[boot] runtime build FAILED:\n" + _BUILD_ERROR, flush=True) # GPU session: full pipeline inside the fork; WS replaced by # in_q {"kind":"text"|"bytes"|"close","data":...} / out_q ("json",dict)|("bin",bytes) def make_gpu_session(in_q, out_q, paths, duration, dead): """Build the @spaces.GPU generator for one session. in_q/out_q captured by CLOSURE — inherited through the fork (MPQueues are not picklable, so they cannot be decorated-function args).""" sched_t0 = time.monotonic() # ZeroGPU's duration window opens here — queue wait included @spaces.GPU(duration=duration, size=GPU_SIZE) def gpu_session(): # Queue slots can't be cancelled: if the client left, release without touching the engine. if dead.is_set(): print("[ws] fork: client left during queue wait — GPU released", flush=True) return yield from _run_session_in_fork(in_q, out_q, paths, duration, sched_t0) return gpu_session def _run_session_in_fork(in_q, out_q, paths, duration, sched_t0): """Inside the GPU fork: drive the vendored streaming engine on a worker thread and yield a sentinel per tick so the @spaces.GPU generator keeps the GPU lease alive for the session's lifetime.""" from xvideo.serving.zerogpu_engine import run_session_blocking done = threading.Event() def _drive(): try: run_session_blocking(_RUNTIME, paths, in_q, out_q) # reports its own errors finally: done.set() threading.Thread(target=_drive, daemon=True).start() while not done.wait(timeout=1.0): left = duration - (time.monotonic() - sched_t0) try: out_q.put_nowait(("json", {"type": "lease_left", "left_s": max(0, round(left))})) except Exception: # noqa: BLE001 pass yield "tick" yield "done" @spaces.GPU(duration=30, size=GPU_SIZE) def gpu_probe(): """Hardware introspection + micro-benchmarks inside the fork. Answers, for a few GPU-seconds: which physical GPU ZeroGPU attached, at what clocks/TGP, and the achievable bf16 / fp8 GEMM throughput + HBM bandwidth. Used to compare a dev box against this Space apples-to-apples.""" import subprocess import time as _t import torch out = {"device": torch.cuda.get_device_name(0), "capability": list(torch.cuda.get_device_capability(0)), "torch": torch.__version__} try: out["cpu_count"] = os.cpu_count() out["cpu_affinity"] = len(os.sched_getaffinity(0)) out["loadavg"] = [round(x, 1) for x in os.getloadavg()] model = "" for line in open("/proc/cpuinfo"): if line.startswith("model name"): model = line.split(":", 1)[1].strip() break out["cpu_model"] = model t0 = _t.perf_counter() x = 0 for i in range(3_000_000): x += i * i out["cpu_1t_spin_ms"] = round((_t.perf_counter() - t0) * 1000, 1) except Exception as e: # noqa: BLE001 out["cpu_err"] = repr(e) try: free, total = torch.cuda.mem_get_info() out["mem_total_gb"] = round(total / 2**30, 1) out["mem_free_gb"] = round(free / 2**30, 1) except Exception as e: # noqa: BLE001 out["mem_err"] = repr(e) try: smi = subprocess.run( ["nvidia-smi", "--query-gpu=name,clocks.sm,clocks.max.sm,power.limit,power.draw,temperature.gpu", "--format=csv,noheader"], capture_output=True, text=True, timeout=10) out["smi"] = smi.stdout.strip() or smi.stderr.strip() except Exception as e: # noqa: BLE001 out["smi"] = repr(e) dev = torch.device("cuda:0") a = torch.randn(8192, 8192, device=dev, dtype=torch.bfloat16) b = torch.randn(8192, 8192, device=dev, dtype=torch.bfloat16) for _ in range(3): a @ b torch.cuda.synchronize() t0 = _t.perf_counter() n = 20 for _ in range(n): a @ b torch.cuda.synchronize() dt = _t.perf_counter() - t0 out["bf16_gemm_tflops"] = round(n * 2 * 8192**3 / dt / 1e12, 1) try: af = (a.float() / a.float().abs().amax()).to(torch.float8_e4m3fn) bf = (b.float() / b.float().abs().amax()).t().contiguous().t().to(torch.float8_e4m3fn) sa = torch.ones(1, device=dev) sb = torch.ones(1, device=dev) for _ in range(3): torch._scaled_mm(af, bf, scale_a=sa, scale_b=sb, out_dtype=torch.bfloat16) torch.cuda.synchronize() t0 = _t.perf_counter() for _ in range(n): torch._scaled_mm(af, bf, scale_a=sa, scale_b=sb, out_dtype=torch.bfloat16) torch.cuda.synchronize() dt = _t.perf_counter() - t0 out["fp8_scaled_mm_tflops"] = round(n * 2 * 8192**3 / dt / 1e12, 1) except Exception as e: # noqa: BLE001 out["fp8_err"] = repr(e) try: big = torch.empty(2**30, device=dev, dtype=torch.uint8) dst = torch.empty_like(big) for _ in range(2): dst.copy_(big) torch.cuda.synchronize() t0 = _t.perf_counter() for _ in range(8): dst.copy_(big) torch.cuda.synchronize() dt = _t.perf_counter() - t0 out["hbm_copy_gbps"] = round(8 * 2 * big.numel() / dt / 1e9, 0) except Exception as e: # noqa: BLE001 out["bw_err"] = repr(e) try: from xvideo.models.vae import vae_compile as _vc vae_d = _RUNTIME.decode_vae vae_e = _RUNTIME.pipeline.vae lat_c = int(getattr(vae_d, "latent_channels", 16) or 16) z = torch.zeros(1, lat_c, 2, 60, 105, device=dev, dtype=torch.bfloat16) z = _vc.prep_input(z) with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): for _ in range(3): vae_d.decode(z, return_dict=False)[0] torch.cuda.synchronize() t0 = _t.perf_counter() for _ in range(15): vae_d.decode(z, return_dict=False)[0] torch.cuda.synchronize() out["vae_decode_direct_ms"] = round((_t.perf_counter() - t0) / 15 * 1000, 2) x = torch.zeros(1, 3, 9, 480, 840, device=dev, dtype=torch.bfloat16) x = _vc.prep_input(x) with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): for _ in range(3): vae_e.encode(x) torch.cuda.synchronize() t0 = _t.perf_counter() for _ in range(10): vae_e.encode(x) torch.cuda.synchronize() out["vae_encode_direct_ms"] = round((_t.perf_counter() - t0) / 10 * 1000, 2) except Exception as e: # noqa: BLE001 out["vae_probe_err"] = repr(e) try: import torch.nn.functional as F torch.backends.cudnn.benchmark = True for tag, (shape, cout) in {"conv_a": ((1, 384, 2, 60, 105), 384), "conv_b": ((1, 128, 8, 120, 210), 128)}.items(): x = torch.randn(*shape, device=dev, dtype=torch.bfloat16).to(memory_format=torch.channels_last_3d) w = torch.randn(cout, shape[1], 3, 3, 3, device=dev, dtype=torch.bfloat16).to(memory_format=torch.channels_last_3d) for _ in range(8): F.conv3d(x, w, padding=1) torch.cuda.synchronize() t0 = _t.perf_counter() for _ in range(30): F.conv3d(x, w, padding=1) torch.cuda.synchronize() out[tag + "_ms"] = round((_t.perf_counter() - t0) / 30 * 1000, 3) except Exception as e: # noqa: BLE001 out["conv_err"] = repr(e) try: tiny = torch.ones(16, device=dev) for _ in range(200): tiny.add_(1) torch.cuda.synchronize() n2 = 2000 t0 = _t.perf_counter() for _ in range(n2): tiny.add_(1) torch.cuda.synchronize() out["launch_us_per_kernel"] = round((_t.perf_counter() - t0) / n2 * 1e6, 2) g = torch.cuda.CUDAGraph() s = torch.cuda.Stream() s.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(s): for _ in range(3): tiny.add_(1) torch.cuda.current_stream().wait_stream(s) with torch.cuda.graph(g): for _ in range(500): tiny.add_(1) for _ in range(3): g.replay() torch.cuda.synchronize() t0 = _t.perf_counter() for _ in range(20): g.replay() torch.cuda.synchronize() out["graph_replay_us_per_kernel"] = round((_t.perf_counter() - t0) / (20 * 500) * 1e6, 3) except Exception as e: # noqa: BLE001 out["launch_err"] = repr(e) return out # Main process: gr.Server (FastAPI) + thin byte-pipe WebSocket. from gradio import Server # noqa: E402 from fastapi import Request, WebSocket, WebSocketDisconnect # noqa: E402 from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, Response # noqa: E402 from starlette.datastructures import Headers # noqa: E402 # Mirrors zerogpu_engine._default_args: sessions record here (owner's decision). _RECORD_DIR = os.environ.get("JOYOMNI_RECORD_DIR", "/data/recordings" if Path("/data").is_dir() else "") or None _INDEX_HTML = (DEPLOY_ROOT / "static" / "index.html").read_text(encoding="utf-8") # JOYOMNI_MAINTENANCE=1 serves this instead of the live demo (delete the variable to restore). _MAINTENANCE_HTML = """ JoyAI Video Edit
\U0001f6e0️

Down for maintenance — back soon

The live demo is temporarily offline while we work on it.

Built by the JoyAI Video Edit team · runs natively at 1248×720 · prompt enhancement by Claude Fable 5
""" app = Server(title="JoyAI-Video-Edit (ZeroGPU)") @app.get("/", response_class=HTMLResponse) async def index(): if os.environ.get("JOYOMNI_MAINTENANCE"): return HTMLResponse(_MAINTENANCE_HTML) # __SERVER_DEFAULTS__ filled from serve's argparse defaults. from xvideo.serving.serve_joyomni_streaming import build_parser _sd = build_parser().parse_args([]) defaults = {k: getattr(_sd, k) for k in ( "width", "height", "num_inference_steps", "output_quality", "seed", "fps", "online_gate", "kv_reset_frames", "static_diff_thresh", "freeze_kv_on_static", "profile_timings", "max_temporal_ids")} defaults.update({"use_pe": bool(os.environ.get("OPENAI_API_KEY")), "pe_available": bool(os.environ.get("OPENAI_API_KEY")), "record_enabled": _RECORD_DIR is not None}) return _INDEX_HTML.replace("__SERVER_DEFAULTS__", json.dumps(defaults)) @app.get("/ref-images") async def ref_images(): # Reference presets: reuse serve's loader (reads rv2v_reference/, returns name->data-url). from xvideo.serving.serve_joyomni_streaming import _ref_images_cached return _ref_images_cached() # Team gallery (Space-only): the client shows its strip when /team.json exists. _TEAM_DIR = DEPLOY_ROOT / "teams" _TEAM_ORDER = ("us", "wenxun", "wenxun-x", "xinran", "yicheng") @app.get("/team.json") async def team_manifest(): items = [{"name": n, "thumb": f"/team/thumb/{n}", "full": f"/team/full/{n}"} for n in _TEAM_ORDER if (_TEAM_DIR / f"{n}.png").is_file() and (_TEAM_DIR / "thumb" / f"{n}.webp").is_file()] return JSONResponse(items, status_code=200 if items else 404) @app.get("/team/{kind}/{name}") async def team_image(kind: str, name: str): if name not in _TEAM_ORDER or kind not in ("thumb", "full"): return JSONResponse({"error": "not found"}, status_code=404) path = _TEAM_DIR / "thumb" / f"{name}.webp" if kind == "thumb" else _TEAM_DIR / f"{name}.png" if not path.is_file(): return JSONResponse({"error": "not found"}, status_code=404) return FileResponse(str(path), media_type="image/webp" if kind == "thumb" else "image/png", headers={"Cache-Control": "public, max-age=86400"}) @app.get("/health") async def health(): return {"ok": _BUILD_ERROR is None, "build_error": _BUILD_ERROR} @app.get("/download_last") async def download_last(rec: str = ""): if _RECORD_DIR is None: return JSONResponse({"error": "Recording is not enabled."}, status_code=404) if not re.fullmatch(r"\d+_\d+", rec): return JSONResponse({"error": "Missing or invalid rec id."}, status_code=400) base = Path(_RECORD_DIR) / rec if not base.is_dir(): return JSONResponse({"error": "No such recording."}, status_code=404) segments = sorted(base.glob("output_*.mp4")) if not segments: return JSONResponse({"error": "Recording file has not been generated yet. Try again later."}, status_code=404) download_name = f"joyomni_{base.name}.mp4" if len(segments) == 1: return FileResponse(str(segments[0]), media_type="video/mp4", filename=download_name) try: import imageio_ffmpeg ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe() except Exception as exc: # noqa: BLE001 return JSONResponse({"error": f"ffmpeg unavailable: {exc!r}"}, status_code=500) list_path = out_path = None try: fd_list, list_path = tempfile.mkstemp(suffix=".txt", prefix="rv2v_cat_") with os.fdopen(fd_list, "w", encoding="utf-8") as f: for seg in segments: f.write(f"file '{seg.as_posix()}'\n") fd_out, out_path = tempfile.mkstemp(suffix=".mp4", prefix="rv2v_dl_") os.close(fd_out) proc = await asyncio.create_subprocess_exec( ffmpeg_exe, "-y", "-f", "concat", "-safe", "0", "-i", list_path, "-c", "copy", "-movflags", "+faststart", out_path, stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.PIPE, ) _, stderr = await proc.communicate() if proc.returncode != 0: tail = (stderr or b"").decode("utf-8", "replace") return JSONResponse({"error": f"ffmpeg failed: {tail}"}, status_code=500) with open(out_path, "rb") as f: data = f.read() return Response(content=data, media_type="video/mp4", headers={"Content-Disposition": f'attachment; filename="{download_name}"'}) except Exception as exc: # noqa: BLE001 return JSONResponse({"error": f"download encode error: {exc!r}"}, status_code=500) finally: for p in (list_path, out_path): if p: try: os.unlink(p) except OSError: pass # Debug endpoints burn quota: gated by the JOYOMNI_DEBUG_KEY secret (?key=...); no secret = disabled. _DEBUG_KEY = os.environ.get("JOYOMNI_DEBUG_KEY", "") def _debug_denied(request: Request): if not _DEBUG_KEY or request.query_params.get("key") != _DEBUG_KEY: return JSONResponse({"error": "forbidden"}, status_code=403) return None @spaces.GPU(duration=120, size=GPU_SIZE) def engine_bench(fresh_cache: bool = False, pace: float = 12.0, chunks: int = 16): """Engine-only throughput probe inside the fork: no session wrapper, no queues — the exact mirror of the dev-box bench. Separates 'fork environment' from 'session machinery' when chasing Space-vs-local speed gaps. Samples SM clocks while running to expose power-management effects.""" import subprocess as _sp import threading as _th import time as _t import numpy as np from PIL import Image as _Image if fresh_cache: import torch._inductor.config # noqa: F401 os.environ["TORCHINDUCTOR_FORCE_DISABLE_CACHES"] = "1" from xvideo.serving.joyomni_streaming import StreamingSettings clocks = [] stop_clk = _th.Event() def _clk(): while not stop_clk.is_set(): try: r = _sp.run(["nvidia-smi", "--query-gpu=clocks.sm,power.draw", "--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=5) clocks.append(r.stdout.strip().split("\n")[0]) except Exception: # noqa: BLE001 pass stop_clk.wait(0.5) _th.Thread(target=_clk, daemon=True).start() settings = StreamingSettings(height=480, width=840, num_inference_steps=2, seed=42, max_temporal_ids=8, profile_timings=True, output_codec="mjpeg") sess = _RUNTIME.create_v2v_session( "Turn the scene into an oil painting with warm colors.", settings=settings) rng = np.random.default_rng(7) base = rng.integers(0, 256, size=(480, 840, 3), dtype=np.uint8) recs = [] t0 = _t.time() i = 0 try: while len(recs) < chunks and i < 400 and _t.time() - t0 < 95: if pace > 0: target = t0 + i / pace now = _t.time() if target > now: _t.sleep(target - now) frame = _Image.fromarray(np.roll(base, (i * 7) % 840, axis=1), mode="RGB") for r in sess.push_frame(frame, {"seq": i + 1, "t_capture_ms": 0.0}): p = r.profile or {} recs.append({k: round(float(p[k]), 4) for k in ("dit_denoise_s", "vae_encode_s", "reference_prepare_s", "vae_decode_s", "jpeg_encode_s", "frames_to_tensor_s", "total_server_chunk_s") if p.get(k) is not None}) i += 1 deadline = _t.time() + 20.0 while len(recs) < chunks and _t.time() < deadline: r = sess.wait_async_result(timeout=0.5) if r is not None: p = r.profile or {} recs.append({k: round(float(p[k]), 4) for k in ("dit_denoise_s", "vae_encode_s", "reference_prepare_s", "vae_decode_s", "jpeg_encode_s", "frames_to_tensor_s", "total_server_chunk_s") if p.get(k) is not None}) finally: stop_clk.set() try: sess.close() except Exception: # noqa: BLE001 pass st = recs[5:] or recs def _m(k): v = [c[k] for c in st if k in c] return round(sum(v) / len(v), 4) if v else None return {"chunks": len(recs), "fresh_cache": bool(fresh_cache), "pace": pace, "dit": _m("dit_denoise_s"), "enc": _m("vae_encode_s"), "ref": _m("reference_prepare_s"), "dec": _m("vae_decode_s"), "jpg": _m("jpeg_encode_s"), "f2t": _m("frames_to_tensor_s"), "total": _m("total_server_chunk_s"), "recs": recs, "clocks": clocks[-60:]} @app.get("/gpu-probe") async def do_gpu_probe(request: Request): denied = _debug_denied(request) if denied is not None: return denied loop = asyncio.get_event_loop() try: return {"result": await loop.run_in_executor(None, gpu_probe)} except Exception as e: # noqa: BLE001 return JSONResponse({"error": str(e)}, status_code=500) @app.get("/engine-bench") async def do_engine_bench(request: Request, fresh_cache: int = 0, pace: float = 12.0, chunks: int = 16): denied = _debug_denied(request) if denied is not None: return denied loop = asyncio.get_event_loop() try: return {"result": await loop.run_in_executor( None, lambda: engine_bench(bool(fresh_cache), float(pace), int(chunks)))} except Exception as e: # noqa: BLE001 return JSONResponse({"error": str(e)}, status_code=500) def _token_user(tok: str) -> bool: try: pl = tok.split(".")[1] claims = json.loads(base64.urlsafe_b64decode(pl + "=" * (-len(pl) % 4))) return claims.get("user") is not None except Exception: # noqa: BLE001 return False # sid -> (headers, expiry). When embedded on huggingface.co, the parent page mints # user-scoped quota headers via the "zerogpu-headers" postMessage handshake (same # protocol gradio's frontend speaks). The frontend relays them here over a plain GET # (WS upgrades can't carry custom headers) and /ws redeems the one-time ticket, so # the visitor tier bills the logged-in user instead of the anonymous IP pool. _MINTED: dict[str, tuple[dict, float]] = {} @app.get("/mint/{sid}") async def mint(sid: str, request: Request): picked = {k.lower(): v for k, v in request.headers.items() if k.lower() == "x-ip-token" or k.lower().startswith("x-zerogpu")} user = _token_user(picked.get("x-ip-token", "")) if picked: now = time.time() for k in [k for k, (_, exp) in _MINTED.items() if exp < now]: _MINTED.pop(k, None) if len(_MINTED) < 512: _MINTED[sid[:64]] = (picked, now + 120.0) print(f"[mint] stashed={sorted(picked)} user={'yes' if user else 'no'}", flush=True) return {"stashed": bool(picked), "user": user, "names": sorted(picked)} @app.websocket("/ws") async def ws(websocket: WebSocket): await websocket.accept() # Maintenance covers the GPU path too: with the app-pool rungs in the # ladder, a stale tab or a raw WS client could otherwise still bill the # owner while the front page says "closed". if os.environ.get("JOYOMNI_MAINTENANCE"): await websocket.send_json({"type": "error", "code": "maintenance", "message": "under maintenance"}) await websocket.close() return if _BUILD_ERROR is not None: await websocket.send_json({"type": "error", "message": "model failed to load"}) await websocket.close() return # 512 slots never fill (frame = 2 slots, frontend caps 32 unacked); a dropped # frame is never acked and would permanently strangle the credit gate. in_q: "MPQueue" = MPQueue(maxsize=512) out_q: "MPQueue" = MPQueue(maxsize=64) stop = threading.Event() gone = threading.Event() _hdrs = websocket.headers _src = "header" _sid = websocket.query_params.get("mint") if _sid and (_m := _MINTED.pop(_sid, None)) is not None and _m[1] > time.time(): _hdrs = Headers({**dict(websocket.headers), **_m[0]}) _src = "minted" _tok = _hdrs.get("x-ip-token") _user = bool(_tok) and _token_user(_tok) print(f"[ws] x-ip-token: {'yes' if _tok else 'no'} user={'yes' if _user else 'no'} src={_src}", flush=True) def _worker(): from gradio.context import LocalContext # Acquisition ladder: visitor quota first (user token or anonymous IP # pool), then the shared app pool, each at SESSION_DURATION/30s. # Quota refusals are instant (checked at schedule time), so a step down # costs one round-trip; "No GPU was available" means the physical pool # is busy — no rung or pool fixes that, so the ladder aborts there. # Without an x-ip-token the "visitor" identity would silently bill the # app anyway (spaces falls back to the app token), so skip those rungs. _billx = 2 if GPU_SIZE == "xlarge" else 1 tiers = [(pool, req, d) for pool, req in (("visitor", SimpleNamespace(headers=_hdrs)), ("app", None)) if pool != "visitor" or _tok for d in dict.fromkeys((SESSION_DURATION, 30))] got_gpu = False granted = None def _put(payload): try: out_q.put(("json", payload)) except Exception: # noqa: BLE001 pass def _err(code, msg): payload = {"type": "error", "code": code, "message": msg} m = re.search(r"(?:re)?try\s+again\s+in\s+([0-9:]+)", msg, re.I) if m: payload["retry_in"] = m.group(1) _put(payload) last = "" try: for pool, req, dur in tiers: if stop.is_set() or gone.is_set(): return _put({"type": "gpu_wait", "pool": pool, "seconds": dur, "gpu_size": GPU_SIZE, "billing_x": _billx}) try: LocalContext.request.set(req) gen = make_gpu_session(in_q, out_q, _PATHS, dur, gone) for _ in gen(): if not got_gpu: got_gpu = True granted = (pool, dur) print(f"[ws] lease acquired: {pool}@{dur}s " f"(billed x{_billx})", flush=True) _put({"type": "gpu_ready", "pool": pool, "seconds": dur, "gpu_size": GPU_SIZE, "billing_x": _billx}) if stop.is_set(): break return # session ended inside its lease (client stop/close) except Exception as exc: # noqa: BLE001 msg = (str(exc) or exc.__class__.__name__).strip() low = msg.lower() if got_gpu: # Ran, then died. Lease expiry surfaces as "GPU task # aborted" — a normal end, not an error. print("[ws] session ended:\n" + traceback.format_exc(), flush=True) if "aborted" in low: _put({"type": "lease_expired", "pool": granted[0], "seconds": granted[1], "billing_x": _billx}) else: _err("gpu_error", msg) return print(f"[ws] {pool}@{dur}s failed: {msg!r}", flush=True) last = msg if "no gpu was available" in low: _err("gpu_busy", msg) return if not any(k in low for k in ("quota", "exceeded", "limit", "credits")): _err("gpu_error", msg) return # quota refusal -> next rung _err("gpu_quota", last or "quota exceeded") finally: stop.set() threading.Thread(target=_worker, daemon=True).start() loop = asyncio.get_event_loop() async def send_loop(): while not stop.is_set() or not _q_empty(out_q): try: kind, payload = await loop.run_in_executor(None, lambda: out_q.get(timeout=0.1)) except queue.Empty: continue except (OSError, ValueError): break try: if kind == "json": await websocket.send_json(payload) else: await websocket.send_bytes(payload) except (WebSocketDisconnect, RuntimeError): break try: await websocket.send_json({"type": "session_timeout", "message": "GPU lease ended; reconnect to re-queue"}) except Exception: # noqa: BLE001 pass async def recv_loop(): while not stop.is_set(): try: msg = await websocket.receive() except (WebSocketDisconnect, RuntimeError): break if msg.get("type") == "websocket.disconnect": break if msg.get("text") is not None: # Answer pings here in the main process: during the GPU-queue wait # nothing drains in_q, and a queued ping would read as fake 10s+ RTT. if '"ping"' in msg["text"]: try: p = json.loads(msg["text"]) if p.get("type") == "ping": await websocket.send_json({"type": "pong", "t": p.get("t")}) continue except (ValueError, KeyError): pass _q_put(in_q, {"kind": "text", "data": msg["text"]}) elif msg.get("bytes") is not None: _q_put(in_q, {"kind": "bytes", "data": msg["bytes"]}) gone.set() _q_put(in_q, {"kind": "close"}) try: await asyncio.gather(send_loop(), recv_loop()) except (WebSocketDisconnect, RuntimeError): pass finally: stop.set() gone.set() _q_put(in_q, {"kind": "close"}) def _q_put(q, item) -> None: try: q.put_nowait(item) except queue.Full: pass # inbound frames are droppable under backpressure def _q_empty(q) -> bool: try: return q.empty() except (OSError, ValueError, NotImplementedError): return True # Satisfy ZeroGPU's startup scan (no Gradio Blocks events to walk with gr.Server). # NOTE: do NOT create/mount any gr.Blocks here — its mere existence breaks the # FastAPI-demo startup path (the spaces launch hook never fires -> the platform # reports "No @spaces.GPU function detected"). Verified over four build cycles. spaces.GPU(lambda: None) demo = app if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)