"""LiveWan on ZeroGPU: a streaming, steerable text-to-video demo. This Space drives the project's own serving engine (`wanstreamer.serve.engine.Engine`) rather than reimplementing the streaming maths. The engine opens a cached world, extends it block by block with the distilled 1.3B student, and decodes each block through a VAE whose causal-conv cache is kept alive across calls so the blocks join without a seam. That is the same code path `livewan-serve` runs locally. What ZeroGPU changes, and why the UI looks the way it does: a GPU worker is forked per request and cannot be steered from outside while it runs, so the demo takes the steer as a *schedule* -- "start in this world, swap the conditioning to this prompt at t = N seconds" -- instead of a live button. The swap itself is exactly the live one: `Engine.steer` replaces the cross-attention conditioning and leaves the K/V cache in place, so the scene continues rather than cutting. Steers can come from the project's 96-prompt bank or from free text, which umt5-xxl encodes here. The bank is the conditioning every published number refers to; free text is not numerically comparable to it, because umt5 embeddings vary slightly by hardware, so the same string encoded on the training box is not the same tensor. """ import os # Deliberately NOT `expandable_segments:True`. That is the usual fix for allocator # trouble under transient spikes, but here it *causes* it: expandable segments grow # through the CUDA VMM path, and the first growth inside a ZeroGPU worker aborts with # NVML_SUCCESS == r INTERNAL ASSERT FAILED ... CUDACachingAllocator.cpp # on an allocation of ~49 MB, while the same allocation succeeds with the default # allocator. Set it explicitly so a platform default cannot turn it back on. os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:False" # The SDPA backend order that wan21_patches/modules/attention.py pins (cuDNN first) # is left alone. Forcing EFFICIENT_ATTENTION here was tried, on the theory that this # sm_120 / torch 2.11 runtime was falling through to the MATH backend; it changed the # peak not at all (44.2 GB either way), so the backend is not the problem and the # project's own measured order stands. import queue import shutil import subprocess import sys import tempfile import time from pathlib import Path import spaces # must precede torch: it patches torch.cuda.* for module-scope loading import torch import cv2 import gradio as gr import imageio.v2 as imageio import numpy as np from huggingface_hub import snapshot_download APP = Path(__file__).resolve().parent ASSETS = APP / "assets" BASE_DIR = APP / "wan21_13b" WAN_REPO = APP / "wan21_repo" WORLDS_DIR = APP / "generated_worlds" LIVEWAN_REPO = "JonathanColetti/LiveWan" BASE_REPO = "Wan-AI/Wan2.1-T2V-1.3B" GITHUB = "https://github.com/JonathanColetti/LiveWan" FPS = 16 BLOCK_SECONDS = 0.75 # 3 latent frames -> 12 pixel frames at 16 fps NO_STEER = "don't steer, stay on the opening prompt" STALL_SECONDS = 45 # a block takes ~1.0 s here; this only trips on a real fault # ---------------------------------------------------------------- bootstrap def fetch_weights(): """Pull the student, the prompt bank, the four worlds and the Wan2.1 base. The optimiser shards (17.7 GB) are for resuming training and are skipped. """ snapshot_download( LIVEWAN_REPO, local_dir=str(ASSETS), allow_patterns=["checkpoints/t14b_b64/latest.pt", "data/prompts.pt", "out/world_p*.pt"]) # Wan2.1_VAE.pth decodes every block. The base transformer is the scaffold the # student's weights are loaded into (see Engine._load_student). snapshot_download( BASE_REPO, local_dir=str(BASE_DIR), allow_patterns=["Wan2.1_VAE.pth", "config.json", "diffusion_pytorch_model.safetensors", # umt5-xxl encodes free-text steers; the bank covers the # rest. 11 GB, and the reason this Space downloads 29 GB. "models_t5_umt5-xxl-enc-bf16.pth", "google/umt5-xxl/*"]) # Engine.load() prefers a local `umt5-tokenizer/` and otherwise resolves the hub # id, which would mean a network fetch inside every cold GPU worker. tok = BASE_DIR / "umt5-tokenizer" if not tok.exists(): tok.symlink_to(BASE_DIR / "google/umt5-xxl", target_is_directory=True) def install_wan_reference_code(): """Clone Wan2.1 and apply the project's two patches, then put it on sys.path. `attention.py` replaces upstream's `assert FLASH_ATTN_2_AVAILABLE` with an SDPA fallback that keeps q_lens/k_lens; `configs/__init__.py` adds the 640x368 size entries this project streams at. Both are the same files setup.sh copies. `wan/__init__.py` is emptied on purpose. Upstream's eagerly imports the T2V/I2V/ VACE pipelines, which drag in dashscope, xfuser and `torch.cuda.amp` wrappers this demo never calls; only `wan.configs` and `wan.modules` are needed. """ if not WAN_REPO.exists(): subprocess.run(["git", "clone", "-q", "--depth", "1", "https://github.com/Wan-Video/Wan2.1", str(WAN_REPO)], check=True) shutil.copy(APP / "wan21_patches/modules/attention.py", WAN_REPO / "wan/modules/attention.py") shutil.copy(APP / "wan21_patches/configs/__init__.py", WAN_REPO / "wan/configs/__init__.py") (WAN_REPO / "wan/__init__.py").write_text( "# emptied by the LiveWan Space: only wan.configs and wan.modules are used\n") if str(WAN_REPO) not in sys.path: sys.path.insert(0, str(WAN_REPO)) def ensure_cuda_amp_shim(): """`wan.modules.model` imports `torch.cuda.amp`, removed in some torch builds.""" try: import torch.cuda.amp # noqa: F401 except ImportError: import types shim = types.ModuleType("torch.cuda.amp") shim.autocast = lambda *a, **k: torch.amp.autocast("cuda", *a, **k) shim.custom_fwd = torch.amp.custom_fwd shim.custom_bwd = torch.amp.custom_bwd sys.modules["torch.cuda.amp"] = shim torch.cuda.amp = shim print("[boot] downloading weights", flush=True) fetch_weights() install_wan_reference_code() ensure_cuda_amp_shim() from wanstreamer.serve.engine import Engine # noqa: E402 (needs sys.path above) from wanstreamer.serve.streamdecode import StreamingVAEDecoder # noqa: E402 from wanstreamer.stream import FewStepStreamer # noqa: E402 def harden_worker(): """Run block generation under `no_grad`, and log what a failure saw. `FewStepStreamer.generate_block` carries no `torch.no_grad()` of its own -- unlike `set_world`, `stream` and the VAE's `decode`, which all do. Locally that costs nothing, because `Engine._load_student` calls `requires_grad_(False)` on every parameter, so no graph is ever built. Here the weights are loaded in the main process, packed to disk by ZeroGPU and streamed into a forked worker, and autograd ends up live again: all 30 layers of each of the three forwards a block does are retained, and one block peaks at 44.2 GB against a 47.4 GB slice. The no-grad forwards (the world prime, the decode) were the ones that worked. `Engine._run` also swallows tracebacks, keeping only `type: message`, which is right for a status bar and useless for a thread inside a forked worker. """ def wrap(cls, name, no_grad=False): inner = getattr(cls, name) def outer(self, *a, **k): try: if no_grad: with torch.no_grad(): return inner(self, *a, **k) return inner(self, *a, **k) except Exception: import traceback free, total = torch.cuda.mem_get_info() print(f"[fail] {name}: allocated=" f"{torch.cuda.memory_allocated() / 2**30:.1f}G peak=" f"{torch.cuda.max_memory_allocated() / 2**30:.1f}G reserved=" f"{torch.cuda.memory_reserved() / 2**30:.1f}G " f"device_free={free / 2**30:.1f}G/{total / 2**30:.1f}G", flush=True) traceback.print_exc() raise setattr(cls, name, outer) wrap(FewStepStreamer, "generate_block", no_grad=True) wrap(StreamingVAEDecoder, "decode") harden_worker() print(f"[boot] PYTORCH_CUDA_ALLOC_CONF=" f"{os.environ.get('PYTORCH_CUDA_ALLOC_CONF')!r}", flush=True) print("[boot] loading the engine", flush=True) engine = Engine( assets=ASSETS, wan_repo=WAN_REPO, base_dir=BASE_DIR, weights=ASSETS / "checkpoints/t14b_b64/latest.pt", device="cuda", worlds_dir=WORLDS_DIR, allow_worldgen=False, # the 5.7 GB base is the load scaffold, not a second model compile_vae=False, # torch.compile cannot run in a ZeroGPU worker ) engine.load(progress=lambda m: print(f"[boot] {m}", flush=True)) # The engine loads umt5-xxl lazily, on the first free-text prompt. That is right for # a long-lived local server and wrong here: a lazy load happens inside a forked GPU # worker and every cold worker pays it. Load it now, at module scope, so ZeroGPU # packs it to disk with the rest and streams it in. print("[boot] loading umt5-xxl for free text", flush=True) engine.encoder.load() INFO = engine.info() PROMPTS = {f"{p['idx']:>2} · {p['text']}": p["idx"] for p in INFO["prompts"]} WORLDS = {f"world {w['idx']} · {w['prompt'][:70]}": w["idx"] for w in INFO["worlds"]} WORLD_LABELS = list(WORLDS) PROMPT_LABELS = list(PROMPTS) W = {i: lbl for lbl, i in WORLDS.items()} # by world id, for the default # Themes are what the local UI uses to warn about a steer that will smear. The four # shipped worlds are named by their own bank prompt, so the world id is a prompt id. from wanstreamer.serve.conditioning import theme_of # noqa: E402 PROMPT_THEME = {p["idx"]: p["theme"] for p in INFO["prompts"]} WORLD_THEME = {w["idx"]: theme_of(w["idx"]) for w in INFO["worlds"]} print(f"[boot] ready: step {INFO['step']}, {len(PROMPTS)} prompts, " f"{len(WORLDS)} worlds", flush=True) # ------------------------------------------------------------------- render def hud(stats, wall): """The numbers the browser demo puts along the bottom of the stream.""" total = stats.get("total_s") or 0.0 rt = (BLOCK_SECONDS / total) if total else 0.0 return ( f"**{stats['seconds']:.2f} s** of video · {stats['blocks']} blocks · " f"latent frames **{stats['latent_frames']}/{stats['latent_frames_max']}**\n\n" f"block **{total * 1000:.0f} ms** " f"(generate {stats.get('gen_s', 0) * 1000:.0f} ms, " f"decode {stats.get('decode_s', 0) * 1000:.0f} ms) · " f"**{rt:.2f}x real time** · K/V cache **{stats['kv_mb']:.0f} MB** · " f"{wall:.1f} s on the GPU" ) def write_mp4(frames): path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name with imageio.get_writer(path, fps=FPS, codec="libx264", quality=8, macro_block_size=1, ffmpeg_params=["-pix_fmt", "yuv420p"]) as w: for f in frames: w.append_data(f) return path def _duration(world, steer_to, steer_text, steer_at, crossfade, seconds, seed, *args, **kwargs): # Measured: ~12 s to open a world (load, prime the cache, decode its 81 frames) # and ~1.02 s per block, a block being 0.75 s of video. Declared tight on purpose: # ZeroGPU compares the request against the visitor's remaining quota, not the # actual runtime, so a padded number locks people out for no reason. return int(min(200, 30 + float(seconds) * 1.8)) @spaces.GPU(duration=_duration) def run(world: str, steer_to: str, steer_text: str, steer_at: float, crossfade: int, seconds: float, seed: int): """Stream video from a cached world, optionally swapping the prompt mid-stream. Args: world: which of the four shipped worlds to open the stream on. steer_to: a prompt from the 96-prompt bank to swap to, or the no-steer option. steer_text: free text to swap to instead; takes precedence over steer_to. steer_at: seconds into the clip at which to swap the conditioning. crossfade: blocks to interpolate the conditioning over; 0 swaps instantly. seconds: how much video to generate, at 16 fps. seed: RNG seed for the block sampler. Yields: (preview frame, HUD line, finished mp4) — the mp4 only on the last yield. """ total_frames = int(float(seconds) * FPS) # A swap at or past the end of the clip would never fire and the run would look # like the control did nothing. Keep it at least two seconds from the end so the # result of the swap is actually visible. steer_frame = int(min(float(steer_at), max(1.0, float(seconds) - 2.0)) * FPS) text = (steer_text or "").strip() if text: steer_with = {"text": text} elif steer_to and steer_to != NO_STEER: steer_with = {"idx": PROMPTS[steer_to]} else: steer_with = None t0 = time.perf_counter() engine.start(world=WORLDS[world], seed=int(seed)) # Opening a world decodes all 81 of its pixel frames in one call, the largest # allocation in a run, and leaves the allocator holding ~17 GB the block loop has # no use for (reserved 30.5 G against 13.5 G allocated). Handing it back is not # what fixed the OOM this Space started with -- that was autograd, see # `harden_worker` -- but there is no reason to sit on it either. torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() frames, pending_steer, last_push = [], steer_with is not None, 0.0 # The worker runs in a thread and reports failure by setting status.state rather # than raising here, so poll it between frames: a long blocking get() would spend # the whole GPU reservation waiting on a stream that is already dead. ended = "" last_frame_at = time.perf_counter() try: while len(frames) < total_frames: try: jpg = engine.frames.get(timeout=0.5) except queue.Empty: state = engine.status.state if state == "error": raise gr.Error(f"engine error: {engine.status.error}") if state != "streaming": # The worker ends the stream itself at the 1024-latent-frame # RoPE ceiling and puts the reason in `detail`. ended = engine.status.detail break if time.perf_counter() - last_frame_at > STALL_SECONDS: raise gr.Error( f"no block in {STALL_SECONDS} s (state={state})") continue last_frame_at = time.perf_counter() frames.append(cv2.imdecode(np.frombuffer(jpg, np.uint8), cv2.IMREAD_COLOR)[:, :, ::-1]) if pending_steer and len(frames) >= steer_frame: engine.steer(crossfade=int(crossfade), **steer_with) pending_steer = False now = time.perf_counter() if now - last_push > 0.12: last_push = now yield frames[-1], hud(engine.stats(), now - t0), None finally: stats = engine.stats() print(f"[run] {len(frames)} frames, {stats['blocks']} blocks, " f"state={engine.status.state!r} detail={engine.status.detail!r} " f"error={engine.status.error!r}", flush=True) engine.stop() if not frames: raise gr.Error("no frames were produced") line = hud(stats, time.perf_counter() - t0) if ended: line += f"\n\n_{ended}_" yield frames[-1], line, write_mp4(frames) # ----------------------------------------------------------------------- ui CSS = """ #col-container { max-width: 1150px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ INTRO = f"""# LiveWan Streaming text-to-video you can steer while it runs. Swap the prompt partway through and the scene continues instead of cutting. Running on shared ZeroGPU hardware, so about 0.7x real time here; a dedicated GPU does about 2.7x. [Code]({GITHUB}) · [Weights](https://huggingface.co/{LIVEWAN_REPO}) """ def steer_warning(world, steer_to, steer_text): """The same warning the local UI shows before a steer that will smear. Cross-theme steers ask the model to reconcile a K/V cache full of one kind of scene with a prompt describing another, and it smears rather than resolving. Free text gets no theme, so it cannot be checked. """ if (steer_text or "").strip() or not steer_to or steer_to == NO_STEER: return gr.update(visible=False) wt = WORLD_THEME.get(WORLDS[world]) pt = PROMPT_THEME.get(PROMPTS[steer_to]) if not wt or not pt or wt == pt: return gr.update(visible=False) return gr.update(visible=True, value=( f"⚠️ {wt.lower()} to {pt.lower()} is a jump across themes, so expect the " f"picture to smear rather than resolve. Raise the crossfade to soften it, " f"or steer somewhere closer to where you started.")) def bound_swap(seconds, steer_at): """A swap must land far enough before the end to be worth watching.""" hi = max(1.0, float(seconds) - 2.0) return gr.update(maximum=hi, value=min(float(steer_at), hi)) with gr.Blocks(title="LiveWan") as demo: with gr.Column(elem_id="col-container"): gr.Markdown(INTRO) with gr.Row(equal_height=False): with gr.Column(scale=2): world = gr.Dropdown(WORLD_LABELS, value=W[60], label="Open on world") steer_to = gr.Dropdown([NO_STEER] + PROMPT_LABELS, value=NO_STEER, label="Steer to", info="From the prompt bank") steer_text = gr.Textbox( label="Or steer to your own text", lines=2, placeholder="A lighthouse in a storm", info="Encoded with umt5-xxl. Overrides the selection above.") warning = gr.Markdown(visible=False) seconds = gr.Slider(5, 30, value=15, step=1, label="Generate (seconds of video)") steer_at = gr.Slider(1, 13, value=6, step=0.5, label="Swap at (seconds in)") # The local UI defaults this to 0, which is right when a human is # watching and can react. For a one-click demo it is the wrong # default: at 0 a waterfall steered to an aurora has collapsed into # blocking by 14 s, and at 8 the scene is still standing. crossfade = gr.Slider(0, 12, value=6, step=1, label="Crossfade (blocks)", info="Blending the swap over a few blocks " "holds the scene together. 0 swaps " "instantly, which is harsher.") seed = gr.Number(value=0, precision=0, label="Seed") run_btn = gr.Button("Stream", variant="primary", size="lg") with gr.Column(scale=3): # Both at the stream's native 640x368: the live view is the point of # the demo, so it is not the small one. preview = gr.Image(label="Live, as each block lands", height=368) stats = gr.Markdown() video = gr.Video(label="The clip, at 16 fps", autoplay=True, height=368) # api_name=False keeps these out of the API tab and the MCP tool list; they are # UI wiring, not endpoints. `run` is the only thing worth calling from outside. for control in (world, steer_to, steer_text): control.change(steer_warning, [world, steer_to, steer_text], warning, api_name=False) seconds.change(bound_swap, [seconds, steer_at], steer_at, api_name=False) run_btn.click(run, inputs=[world, steer_to, steer_text, steer_at, crossfade, seconds, seed], outputs=[preview, stats, video], concurrency_limit=1) demo.queue(max_size=12).launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)