"""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. Free text is not available here. Encoding it needs umt5-xxl (11 GB) on top of the 18 GB this Space already pulls at startup, so the prompt selectors are the project's 96-prompt bank -- which is the conditioning every published number refers to. Run the GitHub repo locally for free text. """ 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" 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 is ~0.3 s; this only trips when something is wrong # ---------------------------------------------------------------- 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; so is umt5-xxl, which only free text needs. """ 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"]) 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 log_worker_tracebacks(): """`Engine._run` swallows the traceback, keeping only `type: message`. That is right for a browser demo, where the message goes to a status bar, but here the failure is a thread inside a forked GPU worker and the traceback is the only way to find out where it happened. Log it on the way past. """ def wrap(cls, name): inner = getattr(cls, name) def outer(self, *a, **k): try: 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") wrap(StreamingVAEDecoder, "decode") log_worker_tracebacks() 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)) 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) 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_at, seconds, seed, *args, **kwargs): # Worst case is the whole clip generated at the slowest observed rate, plus the # world decode at the start. Declared tight on purpose: ZeroGPU compares the # request against the visitor's remaining quota, not the actual runtime. return int(min(200, 45 + float(seconds) * 1.2)) @spaces.GPU(duration=_duration) def run(world: str, steer_to: str, steer_at: float, 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_at: seconds into the clip at which to swap the conditioning. 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) steer_frame = int(float(steer_at) * FPS) steer_idx = PROMPTS.get(steer_to) if steer_to and steer_to != NO_STEER else 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, which is by far # the largest allocation in a run and leaves the caching allocator holding blocks # the block loop cannot reuse (reserved 30.5 G against 13.5 G allocated). Growing # past that OOMs, and a ZeroGPU worker cannot even format the OOM message -- # PyTorch asks NVML which processes hold memory, NVML is not available there, and # the run dies on an internal assert instead. Hand the cache back first. torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() frames, pending_steer, last_push = [], steer_idx 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(idx=steer_idx) 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: 1180px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ INTRO = f"""# LiveWan **Streaming text-to-video you can steer while it runs.** A 1.3B student distilled from Wan2.1-T2V-14B that generates video continuously instead of as a fixed clip: 750 ms of 640x368 at a time, extended block by block. Pick a world to open on, and optionally a prompt to swap to partway through. The swap keeps the K/V cache, so the scene *continues* rather than cutting. Checkpoint step {INFO['step']}. [Code]({GITHUB}) · [Weights](https://huggingface.co/{LIVEWAN_REPO}) """ NOTES = f""" - The live preview runs as fast as the GPU produces frames, which is faster than real time. The **mp4 below it is the honest 16 fps playback** — watch that one. - Prompts come from the project's 96-prompt bank, the conditioning every published number refers to. Free text needs umt5-xxl (11 GB) and is not loaded here; the [local demo]({GITHUB}) has it. - One stream at a time: the engine holds a single K/V cache, so requests queue. - A stream cannot exceed 1024 latent frames (~4.3 min) — that is where `WanModel`'s RoPE tables end. It stops itself and says so. """ with gr.Blocks(title="LiveWan") as demo: with gr.Column(elem_id="col-container"): gr.Markdown(INTRO) with gr.Row(): world = gr.Dropdown(WORLD_LABELS, value=WORLD_LABELS[2], label="Open on world", scale=1) steer_to = gr.Dropdown([NO_STEER] + PROMPT_LABELS, value=NO_STEER, label="Steer to", scale=1) with gr.Row(): steer_at = gr.Slider(1, 25, value=6, step=0.5, label="Steer at (seconds in)", scale=1) seconds = gr.Slider(5, 30, value=15, step=1, label="Generate (seconds of video)", scale=1) seed = gr.Number(value=0, precision=0, label="Seed", scale=0) run_btn = gr.Button("Stream", variant="primary") preview = gr.Image(label="Live preview (faster than real time)", height=368) stats = gr.Markdown() video = gr.Video(label="The clip, at 16 fps", autoplay=True) gr.Examples( examples=[ [WORLD_LABELS[2], NO_STEER, 6, 20, 0], [WORLD_LABELS[2], PROMPT_LABELS[60], 6, 18, 0], [WORLD_LABELS[0], PROMPT_LABELS[3], 8, 18, 0], [WORLD_LABELS[1], PROMPT_LABELS[45], 7, 18, 0], ], inputs=[world, steer_to, steer_at, seconds, seed], label="Try one", ) gr.Markdown(NOTES) run_btn.click(run, inputs=[world, steer_to, steer_at, seconds, seed], outputs=[preview, stats, video], concurrency_limit=1) demo.queue(max_size=12).launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)