app.py

"""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)