LiveWan streaming demo on ZeroGPU

.gitignore

@@ -0,0 +1,6 @@
+assets/
+wan21_13b/
+wan21_repo/
+generated_worlds/
+__pycache__/
+*.pyc

README.md

@@ -1,13 +1,55 @@
 ---
 title: LiveWan
-emoji: 📚
-colorFrom: green
-colorTo: red
+emoji: 🎞️
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
 sdk_version: 6.24.0
-python_version: '3.12'
 app_file: app.py
+python_version: "3.10"
+short_description: Streaming text-to-video you can steer mid-stream
+startup_duration_timeout: 1h
 pinned: false
+license: apache-2.0
+models:
+  - JonathanColetti/LiveWan
+  - Wan-AI/Wan2.1-T2V-1.3B
+tags:
+  - text-to-video
+  - streaming
+  - real-time
+  - wan2.1
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# LiveWan
+
+Streaming, steerable text-to-video from a 1.3B student distilled out of
+Wan2.1-T2V-14B with SF-DMD. It generates video continuously rather than as a fixed
+clip — 750 ms of 640×368 at a time, extended block by block — and the text
+conditioning can be swapped mid-stream without clearing the K/V cache, so the scene
+continues instead of cutting.
+
+- Code: https://github.com/JonathanColetti/LiveWan
+- Weights and data: https://huggingface.co/JonathanColetti/LiveWan
+
+## What this Space runs
+
+`app.py` drives the project's own serving engine, `wanstreamer.serve.engine.Engine`
+— the same code path `livewan-serve` runs locally. `wanstreamer/` and
+`wan21_patches/` here are copies of the GitHub repo's; the Wan2.1 reference code is
+cloned and patched at startup exactly as `setup.sh` does.
+
+Two differences from running it locally, both forced by ZeroGPU:
+
+- **Steering is a schedule, not a button.** A GPU worker is forked per request and
+  cannot be steered from outside while it runs, so the demo takes the swap up front
+  ("switch to this prompt at t = N seconds"). The swap itself is the live one:
+  `Engine.steer` replaces the cross-attention conditioning and leaves the cache in
+  place.
+- **No free text.** Encoding it needs umt5-xxl, 11 GB on top of the 18 GB this
+  Space already pulls at startup, so the 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.
+
+`torch.compile` of the VAE decoder is also off, since it cannot run in a ZeroGPU
+worker. That costs roughly 50 ms per block.

app.py

@@ -0,0 +1,308 @@
+"""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
+
+# Set before torch is imported anywhere: the decoder allocates and frees a
+# pixel-space buffer every block, which is the pattern expandable segments exist for.
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+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"
+
+# ---------------------------------------------------------------- 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)
+
+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))
+    frames, pending_steer, last_push = [], steer_idx is not None, 0.0
+
+    try:
+        while len(frames) < total_frames:
+            if engine.status.state == "error":
+                raise gr.Error(f"engine error: {engine.status.error}")
+            try:
+                jpg = engine.frames.get(timeout=60)
+            except queue.Empty:
+                # The worker sets state to idle and puts the reason in `detail` when
+                # a stream ends on its own (the 1024-latent-frame RoPE ceiling).
+                if engine.status.state != "streaming":
+                    break
+                raise gr.Error("the stream stalled waiting for a block")
+
+            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:
+        engine.stop()
+
+    if not frames:
+        raise gr.Error("no frames were produced")
+    note = engine.status.detail
+    line = hud(engine.stats(), time.perf_counter() - t0)
+    if note:
+        line += f"\n\n_{note}_"
+    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(theme=gr.themes.Citrus(), css=CSS, 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, show_download_button=False)
+        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(mcp_server=True)

requirements.txt

@@ -0,0 +1,19 @@
+# torch, gradio, spaces and huggingface_hub come with the ZeroGPU Gradio image.
+torchvision
+numpy
+einops
+safetensors
+transformers
+tokenizers
+sentencepiece
+accelerate
+diffusers
+# the Wan2.1 reference package's own needs
+easydict
+ftfy
+regex
+tqdm
+# this Space: jpeg frames out of the engine, mp4 in
+opencv-python-headless
+imageio
+imageio-ffmpeg

wan21_patches/configs/__init__.py

@@ -0,0 +1,58 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import copy
+import os
+
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+from .wan_i2v_14B import i2v_14B
+from .wan_t2v_1_3B import t2v_1_3B
+from .wan_t2v_14B import t2v_14B
+
+# the config of t2i_14B is the same as t2v_14B
+t2i_14B = copy.deepcopy(t2v_14B)
+t2i_14B.__name__ = 'Config: Wan T2I 14B'
+
+# the config of flf2v_14B is the same as i2v_14B
+flf2v_14B = copy.deepcopy(i2v_14B)
+flf2v_14B.__name__ = 'Config: Wan FLF2V 14B'
+flf2v_14B.sample_neg_prompt = "镜头切换，" + flf2v_14B.sample_neg_prompt
+
+WAN_CONFIGS = {
+    't2v-14B': t2v_14B,
+    't2v-1.3B': t2v_1_3B,
+    'i2v-14B': i2v_14B,
+    't2i-14B': t2i_14B,
+    'flf2v-14B': flf2v_14B,
+    'vace-1.3B': t2v_1_3B,
+    'vace-14B': t2v_14B,
+}
+
+SIZE_CONFIGS = {
+    '720*1280': (720, 1280),
+    '1280*720': (1280, 720),
+    '480*832': (480, 832),
+    '832*480': (832, 480),
+    '1024*1024': (1024, 1024),
+    # Wan-Streamer v0.2 target resolution (landscape). Legal for t2v-1.3B:
+    # vae_stride=(4,8,8), patch_size=(1,2,2) -> 640/8/2 = 40, 368/8/2 = 23,
+    # giving a 23x40 = 920-token latent frame.
+    '640*368': (640, 368),
+}
+
+MAX_AREA_CONFIGS = {
+    '720*1280': 720 * 1280,
+    '1280*720': 1280 * 720,
+    '480*832': 480 * 832,
+    '832*480': 832 * 480,
+    '640*368': 640 * 368,
+}
+
+SUPPORTED_SIZES = {
+    't2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
+    't2v-1.3B': ('480*832', '832*480', '640*368'),
+    'i2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
+    'flf2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
+    't2i-14B': tuple(SIZE_CONFIGS.keys()),
+    'vace-1.3B': ('480*832', '832*480'),
+    'vace-14B': ('720*1280', '1280*720', '480*832', '832*480')
+}

wan21_patches/modules/attention.py

@@ -0,0 +1,380 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+#
+# MODIFIED for Blackwell (sm_100) / no-flash-attn operation.
+#
+# Why: every attention call site in this repo (model.py:149/179/220/222,
+# clip.py:85/197, streaming_blocks.py) calls `flash_attention`, NOT `attention`.
+# Upstream `flash_attention` ends in `assert FLASH_ATTN_2_AVAILABLE`, so with no
+# flash_attn wheel (none exists for sm_100) every forward pass raised a bare
+# AssertionError. The SDPA fallback that upstream put in `attention()` had zero
+# callers and was dead code.
+#
+# Fix: `flash_attention` now dispatches to a correct SDPA implementation when no
+# flash-attn is installed. Unlike upstream's `attention()`, this fallback does NOT
+# discard q_lens/k_lens -- it materialises an explicit boolean mask whenever the
+# lengths imply real padding. Silently attending over padding was the specific
+# "runs fine, quietly wrong" failure mode we had to rule out.
+import os
+import warnings
+from contextlib import nullcontext
+
+import torch
+import torch.nn.functional as F
+
+try:
+    import flash_attn_interface
+    FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+    FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_2_AVAILABLE = False
+
+__all__ = [
+    'flash_attention',
+    'attention',
+    'sdpa_attention',
+    'sdpa_backend_ctx',
+]
+
+# --- SDPA backend selection -------------------------------------------------
+# cuDNN attention is not an alternative to SDPA; it is one of SDPA's backends.
+# Measured on this machine (B200 / sm_100, torch 2.13.0+cu130, cuDNN 9.2.0),
+# bf16, 12 heads x 128 dim, time per call:
+#
+#   shape (Lq x Lk)         FLASH    MEM_EFF   CUDNN    MATH
+#   7800 x 7800  (480x832)  0.967    2.258     0.281    OOM-ish (1.4 GiB)
+#   4600 x 4600  (640x368)  0.353    0.820     0.117    4.955
+#   7800 x  512  (cross)    0.074    0.161     0.030    0.922
+#   1560 x 6240  (stream)   0.316    0.579     0.079    2.420
+#   1560 x 6240 + bool mask   n/a    0.747     0.212    2.850
+#
+# So cuDNN is 3-4x faster than the flash backend here, and it is the only fused
+# backend that accepts an explicit attn_mask (flash rejects non-null masks, which
+# would otherwise drop the padded path to MATH at ~13x the cost).
+#
+# torch's default dispatch already prefers cuDNN at these shapes, but that is an
+# implicit heuristic that varies by version and shape. Pin the order explicitly
+# so the fast path is guaranteed rather than incidental. All backends stay
+# enabled, so an unsupported shape degrades instead of raising.
+#
+# NOTE: because cuDNN was already the default, this pinning is a robustness
+# measure, not a speedup -- do not book it as one.
+_BACKEND_ORDER = ['CUDNN_ATTENTION', 'FLASH_ATTENTION', 'EFFICIENT_ATTENTION', 'MATH']
+
+
+def _build_backend_ctx():
+    """Return a callable giving a context manager that pins SDPA backend order."""
+    try:
+        from torch.nn.attention import SDPBackend, sdpa_kernel
+    except ImportError:
+        return lambda: nullcontext()
+
+    # WAN_SDPA_BACKEND forces a single backend, for benchmarking/debugging only.
+    # It disables all fallbacks, so e.g. WAN_SDPA_BACKEND=FLASH_ATTENTION will
+    # raise "No available kernel" on any masked call. That is intended.
+    override = os.environ.get('WAN_SDPA_BACKEND', '').strip().upper()
+    order = [override] if override else _BACKEND_ORDER
+    backends = [getattr(SDPBackend, n) for n in order if hasattr(SDPBackend, n)]
+    if not backends:
+        return lambda: nullcontext()
+
+    # set_priority=True keeps every listed backend enabled but fixes the order.
+    try:
+        with sdpa_kernel(backends, set_priority=True):
+            pass
+    except TypeError:
+        return lambda: sdpa_kernel(backends)
+    return lambda: sdpa_kernel(backends, set_priority=True)
+
+
+sdpa_backend_ctx = _build_backend_ctx()
+
+
+def _needs_mask(lens, length):
+    """True if `lens` describes anything shorter than the padded `length`."""
+    if lens is None:
+        return False
+    return bool(torch.as_tensor(lens).min().item() < length)
+
+
+def sdpa_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    dtype=torch.bfloat16,
+):
+    """Correct SDPA replacement for flash_attn_varlen_func.
+
+    Shapes match the flash-attn convention:
+        q: [B, Lq, Nq, C1]   k: [B, Lk, Nk, C1]   v: [B, Lk, Nk, C2]
+    Returns [B, Lq, Nq, C2] in q's original dtype.
+
+    Semantics deliberately mirror flash-attn's varlen kernel:
+
+    * `q_lens` / `k_lens` are HONOURED. Keys at index >= k_lens[b] are masked out;
+      query rows at index >= q_lens[b] are zeroed on output. Upstream's
+      `attention()` set attn_mask=None here and only warned, which computes
+      attention over padding and is silently wrong.
+    * `causal=True` uses BOTTOM-RIGHT alignment when Lq != Lk, matching
+      flash-attn (query i sees keys 0 .. i + Lk - Lq). Note torch's
+      `is_causal=True` is TOP-LEFT aligned instead, so the two disagree for
+      Lq != Lk -- that discrepancy is exactly the streaming-KV-cache bug this
+      project had to fix. For block-causal streaming where the K/V cache holds
+      only past+current frames, the correct call is `causal=False`: causality is
+      already enforced by cache contents, and a mask would wrongly serialise
+      tokens *within* a frame.
+    """
+    assert tuple(window_size) == (-1, -1), (
+        f'sliding-window attention (window_size={window_size}) has no SDPA '
+        'fallback; it would be silently ignored.')
+
+    b, lq, nq, _ = q.shape
+    lk, nk = k.shape[1], k.shape[2]
+    out_dtype = q.dtype
+
+    q = q.to(dtype)
+    k = k.to(dtype)
+    v = v.to(dtype)
+
+    if q_scale is not None:
+        q = q * q_scale
+
+    # [B, L, N, C] -> [B, N, L, C]
+    q = q.transpose(1, 2)
+    k = k.transpose(1, 2)
+    v = v.transpose(1, 2)
+
+    # grouped-query attention: replicate k/v heads to match q heads
+    if nq != nk:
+        assert nq % nk == 0, f'Nq ({nq}) must be divisible by Nk ({nk})'
+        rep = nq // nk
+        k = k.repeat_interleave(rep, dim=1)
+        v = v.repeat_interleave(rep, dim=1)
+
+    mask_k = _needs_mask(k_lens, lk)
+    mask_q = _needs_mask(q_lens, lq)
+
+    if not mask_k and (not causal or lq == lk):
+        # Fast path: no padding, and either no mask or a square causal mask
+        # (square => top-left and bottom-right alignment coincide).
+        with sdpa_backend_ctx():
+            out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=None, is_causal=causal,
+                dropout_p=dropout_p, scale=softmax_scale)
+    elif (mask_k and not causal
+          and int(torch.as_tensor(k_lens).min()) > 0):
+        # Key-padding only. The mask does not depend on the query index, so a
+        # [B, 1, 1, Lk] mask is exactly equivalent to the [B, 1, Lq, Lk] one
+        # built below and SDPA broadcasts it. Worth a special case: the block-
+        # causal trainer batches blocks with differing prefix lengths, where the
+        # dense form would be a 244 MiB bool tensor per attention call.
+        kl = torch.as_tensor(k_lens, device=q.device).reshape(b, 1, 1, 1)
+        keep = torch.arange(lk, device=q.device).reshape(1, 1, 1, lk) < kl
+        with sdpa_backend_ctx():
+            out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=keep, is_causal=False,
+                dropout_p=dropout_p, scale=softmax_scale)
+    else:
+        # Build an explicit boolean keep-mask [B, 1, Lq, Lk].
+        keep = torch.ones(b, 1, lq, lk, dtype=torch.bool, device=q.device)
+
+        if mask_k:
+            kl = torch.as_tensor(k_lens, device=q.device).reshape(b, 1, 1, 1)
+            key_idx = torch.arange(lk, device=q.device).reshape(1, 1, 1, lk)
+            keep &= key_idx < kl
+
+        if causal:
+            # bottom-right aligned: key j visible to query i iff j <= i + (lk - lq)
+            qi = torch.arange(lq, device=q.device).reshape(1, 1, lq, 1)
+            kj = torch.arange(lk, device=q.device).reshape(1, 1, 1, lk)
+            keep &= kj <= qi + (lk - lq)
+
+        # A row with nothing visible would give NaN from softmax. Let such rows
+        # attend to key 0, then zero the result below.
+        dead = ~keep.any(dim=-1, keepdim=True)
+        if dead.any():
+            keep = keep | (dead & (torch.arange(lk, device=q.device) == 0).reshape(1, 1, 1, lk))
+
+        with sdpa_backend_ctx():
+            out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=keep, is_causal=False,
+                dropout_p=dropout_p, scale=softmax_scale)
+        if dead.any():
+            out = out.masked_fill(dead, 0.0)
+
+    out = out.transpose(1, 2).contiguous()  # -> [B, Lq, Nq, C2]
+
+    if mask_q:
+        ql = torch.as_tensor(q_lens, device=out.device).reshape(b, 1, 1, 1)
+        qi = torch.arange(lq, device=out.device).reshape(1, lq, 1, 1)
+        out = out * (qi < ql)
+
+    return out.to(out_dtype)
+
+
+def flash_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    version=None,
+):
+    """
+    q:              [B, Lq, Nq, C1].
+    k:              [B, Lk, Nk, C1].
+    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+    q_lens:         [B].
+    k_lens:         [B].
+    dropout_p:      float. Dropout probability.
+    softmax_scale:  float. The scaling of QK^T before applying softmax.
+    causal:         bool. Whether to apply causal attention mask.
+    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
+    deterministic:  bool. If True, slightly slower and uses more memory.
+    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+    """
+    half_dtypes = (torch.float16, torch.bfloat16)
+    assert dtype in half_dtypes
+    assert q.device.type == 'cuda' and q.size(-1) <= 256
+
+    if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE):
+        # No flash-attn (e.g. Blackwell sm_100): use the correct SDPA path.
+        return sdpa_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            dtype=dtype,
+        )
+
+    # params
+    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # preprocess query
+    if q_lens is None:
+        q = half(q.flatten(0, 1))
+        q_lens = torch.tensor(
+            [lq] * b, dtype=torch.int32).to(
+                device=q.device, non_blocking=True)
+    else:
+        q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+    # preprocess key, value
+    if k_lens is None:
+        k = half(k.flatten(0, 1))
+        v = half(v.flatten(0, 1))
+        k_lens = torch.tensor(
+            [lk] * b, dtype=torch.int32).to(
+                device=k.device, non_blocking=True)
+    else:
+        k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+        v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+    q = q.to(v.dtype)
+    k = k.to(v.dtype)
+
+    if q_scale is not None:
+        q = q * q_scale
+
+    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+        warnings.warn(
+            'Flash attention 3 is not available, use flash attention 2 instead.'
+        )
+
+    # apply attention
+    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+        # Note: dropout_p, window_size are not supported in FA3 now.
+        x = flash_attn_interface.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            seqused_q=None,
+            seqused_k=None,
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            deterministic=deterministic)[0].unflatten(0, (b, lq))
+    else:
+        assert FLASH_ATTN_2_AVAILABLE
+        x = flash_attn.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic).unflatten(0, (b, lq))
+
+    # output
+    return x.type(out_dtype)
+
+
+def attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    fa_version=None,
+):
+    return flash_attention(
+        q=q,
+        k=k,
+        v=v,
+        q_lens=q_lens,
+        k_lens=k_lens,
+        dropout_p=dropout_p,
+        softmax_scale=softmax_scale,
+        q_scale=q_scale,
+        causal=causal,
+        window_size=window_size,
+        deterministic=deterministic,
+        dtype=dtype,
+        version=fa_version,
+    )

wanstreamer/__init__.py

@@ -0,0 +1,14 @@
+"""Corrected block-causal streaming implementation for Wan2.1.
+
+See PROGRESS.md for the measurements behind each design decision.
+"""
+from .rope import RopeTable, apply_rope
+from .kvcache import StreamingKVCache
+from .core import (timestep_to_train_scale, ModulationCache, block_forward,
+                   frame_forward, sequence_forward, make_rope_table, make_cache,
+                   latent_geometry)
+
+__all__ = ['RopeTable', 'apply_rope', 'StreamingKVCache',
+           'timestep_to_train_scale', 'ModulationCache', 'block_forward',
+           'frame_forward', 'sequence_forward', 'make_rope_table', 'make_cache',
+           'latent_geometry']

wanstreamer/blockcausal.py

@@ -0,0 +1,423 @@
+"""Blockcausal Wan forward shared by training and streaming inference.
+
+One implementation serves both regimes, which is the point: what the student is
+trained under is literally the code that runs at deployment.
+
+    inference   BufferKV  -> preallocated StreamingKVCache, no_grad, batch 1
+    training    TrainKV   -> plain per-layer tensors, autograd-safe, batched
+
+The contract in both cases is the papers' "KV construction": the past is a
+*clean* (t=0) key/value prefix laid down by an earlier pass, and the current
+block is the only thing carrying noise. Block-causality is therefore structural
+-- the key set only ever holds past + current -- so no attention mask is needed
+over it, and adding one would wrongly serialise tokens within a block.
+
+Two ways to run the noisy blocks of a clip, both against the same clean prefix:
+
+  block_forward           one block, sequential. What inference does.
+  parallel_blocks_forward every block of the clip as a batch row, one forward.
+                          Each row carries its own noise level and its own
+                          absolute RoPE offset, and reads the prefix slice its
+                          start index implies. Used for training.
+
+PER-FRAME TIMESTEP CONDITIONING (v0.2). Upstream Wan folds one shared `e0` into
+every token, so a sequence mixing a clean past with a noisy present is
+inexpressible in a single forward. `time_embed` here accepts a per-latent-frame
+vector, giving modulation [F, 6, dim] applied to x viewed as [B, F, S, dim].
+With the pre-computed clean K/V above it is not *required*. The clean prefix
+is laid down by its own t=0 pass -- but it is what makes a block of mixed noise
+levels expressible at all, and a uniform scalar collapses to upstream's
+behaviour exactly, so the original weights load and run unchanged.
+"""
+import torch
+import torch.utils.checkpoint as ckpt
+
+from wan.modules.attention import flash_attention
+from wan.modules.model import sinusoidal_embedding_1d
+
+
+def time_embed(model, t, device, time_scale=1000.0):
+    """Flow fraction(s) in [0,1] -> (e [N, dim], e0 [N, 6, dim]).
+
+    N == 1              uniform timestep (upstream behaviour)
+    N == batch rows     one timestep per row
+    N == latent frames  per-frame timestep conditioning (batch must be 1)
+    `time_scale` maps the internal [0,1] fraction onto the checkpoint's own
+    convention: 1000 for stock Wan, 1.0 for a model fine-tuned on torch.rand().
+    """
+    if not torch.is_tensor(t):
+        t = torch.as_tensor([float(t)], device=device)
+    t = t.to(device=device, dtype=torch.float32).reshape(-1)
+    # The time embedding runs outside autocast so its precision is the module's
+    # own. Cast the input to match, rather than assuming fp32 weights: the
+    # frozen teacher / critic base is held in bf16 to fit three networks on one
+    # 40 GB card, and upstream's hard fp32 assumption would fail on it.
+    wd = next(model.time_embedding.parameters()).dtype
+    with torch.amp.autocast('cuda', enabled=False):
+        e = model.time_embedding(
+            sinusoidal_embedding_1d(model.freq_dim, t * time_scale).to(wd))
+        e0 = model.time_projection(e).unflatten(1, (6, model.dim))
+    return e, e0
+
+
+class Modulation:
+    """Per-layer AdaLN chunks for one set of timesteps.
+
+    `(block.modulation + e0).chunk(6)` is identical for every token of a frame,
+    so it is computed once per denoising step rather than per layer x frame.
+    Each chunk is [N, 1, dim].
+    """
+
+    def __init__(self, blocks, e0):
+        with torch.amp.autocast('cuda', enabled=False):
+            self.chunks = [(blk.modulation + e0).chunk(6, dim=1) for blk in blocks]
+
+    def __getitem__(self, i):
+        return self.chunks[i]
+
+
+def _bcast(g, x, per_frame, n_frames):
+    """Broadcast a modulation chunk [N,1,dim] against tokens x [B,L,dim].
+
+    per_frame=False: N is 1 or B, which already broadcasts.
+    per_frame=True:  N is the number of latent frames; tokens are frame-major,
+                     so view x as [B, F, S, dim] and give g a leading axis.
+    """
+    if not per_frame:
+        return g
+    return g.unsqueeze(0)              # [1, F, 1, dim] against [B, F, S, dim]
+
+
+def _as_frames(x, per_frame, n_frames):
+    if not per_frame:
+        return x
+    b, l, d = x.shape
+    return x.view(b, n_frames, l // n_frames, d)
+
+
+def _modulate(x, norm, shift, scale, per_frame, n_frames):
+    y = _as_frames(norm(x).float(), per_frame, n_frames)
+    y = y * (1 + _bcast(scale, x, per_frame, n_frames)) \
+        + _bcast(shift, x, per_frame, n_frames)
+    return y.reshape(x.shape[0], x.shape[1], x.shape[2]) if per_frame else y
+
+
+def _gate(x, y, g, per_frame, n_frames):
+    """x + y * g, in float32, with per-frame broadcasting."""
+    with torch.amp.autocast('cuda', dtype=torch.float32):
+        if not per_frame:
+            return x + y * g
+        b, l, d = x.shape
+        out = _as_frames(x, True, n_frames) + \
+            _as_frames(y, True, n_frames) * g.unsqueeze(0)
+        return out.reshape(b, l, d)
+
+
+def _head_from(head, x, e, per_frame, n_frames):
+    """Split out from `_head` so the FSDP shard unit can enter it with the head
+    module it owns rather than reaching through the model (see wanstreamer.fsdp)."""
+    with torch.amp.autocast('cuda', enabled=False):
+        m = (head.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+        y = _modulate(x, head.norm, m[0], m[1], per_frame, n_frames)
+        return head.head(y.to(head.head.weight.dtype))
+
+
+def _head(model, x, e, per_frame, n_frames):
+    return _head_from(model.head, x, e, per_frame, n_frames)
+
+
+def apply_rope(x, tbl):
+    """x [B, L, n, d] real -> rotated, by tbl [L, 1, c] (shared) or [B, L, 1, c].
+
+    float32 complex rather than upstream's float64; verified equivalent within
+    bf16 tolerance by tests/test_streaming_core.py.
+    """
+    b, l, n, d = x.shape
+    xc = torch.view_as_complex(x.float().reshape(b, l, n, d // 2, 2))
+    if tbl.dim() == 3:
+        tbl = tbl.unsqueeze(0)
+    return torch.view_as_real(xc * tbl).flatten(3)
+
+
+# ------------------------------------------------------------------- K/V stores
+class BufferKV:
+    """Adapter over the preallocated StreamingKVCache used at inference."""
+
+    def __init__(self, cache):
+        self.cache = cache
+
+    def context(self, layer, k, v, **_):
+        self.cache.write(layer, k, v)
+        ck, cv = self.cache.context(layer, k.shape[1])
+        return ck, cv, None
+
+
+class BufferPrefixKV:
+    """Read-only slice of a StreamingKVCache, in TrainKV's calling convention.
+
+    The self-forcing trainer rolls out through the preallocated cache and then
+    has to redo one recorded denoising step *with gradient*, against the prefix
+    as it stood at that point. The buffer only ever grows during a rollout (no
+    eviction is used in training), so `buffer[:upto]` is exactly that prefix,
+    and reading it as a view keeps it detached for free.
+    """
+
+    def __init__(self, cache):
+        self.cache = cache
+
+    def context(self, layer, k, v, upto=None, rows=1, k_lens=None):
+        if not upto:
+            return k, v, None
+        pk = self.cache.k[layer, :upto].unsqueeze(0)
+        pv = self.cache.v[layer, :upto].unsqueeze(0)
+        if rows > 1:
+            pk, pv = pk.expand(rows, -1, -1, -1), pv.expand(rows, -1, -1, -1)
+        return torch.cat([k, pk], 1), torch.cat([v, pv], 1), k_lens
+
+
+class TrainKV:
+    """Clean per-layer K/V prefix for a whole clip: plain tensors, autograd-safe.
+
+    Built once per clip by `build_clean_kv` under no_grad; every noisy block then
+    reads the slice its start index implies, so each block trains against exactly
+    the prefix it would see at deployment.
+    """
+
+    def __init__(self, num_layers):
+        self.k = [None] * num_layers
+        self.v = [None] * num_layers
+
+    def append(self, layer, k, v):
+        if self.k[layer] is None:
+            self.k[layer], self.v[layer] = k, v
+        else:
+            self.k[layer] = torch.cat([self.k[layer], k], dim=1)
+            self.v[layer] = torch.cat([self.v[layer], v], dim=1)
+
+    @property
+    def tokens(self):
+        return 0 if self.k[0] is None else self.k[0].shape[1]
+
+    def context(self, layer, k, v, upto=None, rows=1, k_lens=None):
+        """Keys are laid out [current block ; clean prefix] -- current FIRST.
+
+        That ordering is what makes the batched-block path expressible with a
+        key-*length* mask: row i must see its own block plus prefix[:start_i*S],
+        and only in this order are those two runs contiguous from index 0.
+        Attention is permutation-invariant over keys, so the sequential path
+        (which passes no mask at all) is unaffected.
+        """
+        pk, pv = self.k[layer], self.v[layer]
+        if pk is None or upto == 0:
+            return k, v, None
+        pk, pv = pk[:, :upto], pv[:, :upto]
+        if rows > 1:
+            pk = pk.expand(rows, -1, -1, -1)
+            pv = pv.expand(rows, -1, -1, -1)
+        return torch.cat([k, pk], 1), torch.cat([v, pv], 1), k_lens
+
+
+# ------------------------------------------------------------------- the layer
+def _layer(blk, x, ec, tbl, kv_ctx, ctx, ctx_lens, dtype, per_frame, n_frames):
+    """One WanAttentionBlock in block-causal mode. kv_ctx(k,v) -> (K, V, k_lens)."""
+    sa_in = _modulate(x, blk.norm1, ec[0], ec[1], per_frame, n_frames)
+    b, s = sa_in.shape[0], sa_in.shape[1]
+    n, d = blk.num_heads, blk.dim // blk.num_heads
+    sa = blk.self_attn
+
+    q = apply_rope(sa.norm_q(sa.q(sa_in)).view(b, s, n, d), tbl).to(dtype)
+    k = apply_rope(sa.norm_k(sa.k(sa_in)).view(b, s, n, d), tbl).to(dtype)
+    v = sa.v(sa_in).view(b, s, n, d).to(dtype)
+
+    ck, cv, k_lens = kv_ctx(k, v)
+    y = flash_attention(q=q, k=ck, v=cv, k_lens=k_lens,
+                        window_size=(-1, -1), causal=False)
+    x = _gate(x, sa.o(y.flatten(2)), ec[2], per_frame, n_frames)
+    x = x + blk.cross_attn(blk.norm3(x), ctx, ctx_lens)
+    yf = blk.ffn(_modulate(x, blk.norm2, ec[3], ec[4], per_frame, n_frames))
+    return _gate(x, yf, ec[5], per_frame, n_frames)
+
+
+def _run(model, z, t, tbl, ctx, ctx_lens, kv_ctx_factory, dtype, time_scale,
+         per_frame, grad_checkpoint, emb=None):
+    """Shared body: patch-embed -> 30 block-causal layers -> head -> unpatchify.
+
+    `emb` supplies a precomputed (e, e0). The CUDA-graph path needs it because
+    `sinusoidal_embedding_1d` builds its frequency vector on the CPU and copies
+    it to the device, which cannot be captured -- and the time embedding is two
+    small linears, so hoisting it out of the graph costs nothing.
+    """
+    B, _, F = z.shape[0], z.shape[1], z.shape[2]
+    e, e0 = emb if emb is not None else time_embed(model, t, z.device, time_scale)
+    nf = F if per_frame else 1
+    # `wanstreamer.fsdp.shard_model` attaches these: per-layer nn.Modules that
+    # are the FSDP shard units. They must be ENTERED, because FSDP2 all-gathers
+    # a module's parameters from a pre-forward hook on that module -- reaching
+    # into `blk.self_attn.q` from outside, as `_layer` does, would run on
+    # sharded parameters without ever erroring. Absent them nothing changes.
+    layers = getattr(model, 'causal_layers', None)
+    head_mod = getattr(model, 'causal_head', None)
+    mod = Modulation(model.blocks, e0) if layers is None else None
+    with torch.amp.autocast('cuda', dtype=dtype):
+        x = model.patch_embedding(z.to(dtype))
+        gf, gh, gw = (int(s) for s in x.shape[2:])
+        x = x.flatten(2).transpose(1, 2)
+        for li in range(len(model.blocks)):
+            kv_ctx = kv_ctx_factory(li)
+            if layers is None:
+                fn = _layer
+                args = (model.blocks[li], x, mod[li], tbl, kv_ctx, ctx,
+                        ctx_lens, dtype, per_frame, nf)
+            else:
+                # The sharded path passes e0 rather than a precomputed
+                # modulation chunk: `blk.modulation` is a plain Parameter read
+                # outside any forward by `Modulation`, so under FSDP it would
+                # still be a shard at that point. CausalBlock folds that one add
+                # into the layer, where the gather has already happened. It is
+                # the same arithmetic and the same number of evaluations -- each
+                # block's chunks are built exactly once per forward either way.
+                fn = layers[li]
+                args = (x, e0, tbl, kv_ctx, ctx, ctx_lens, dtype, per_frame, nf)
+            if grad_checkpoint and torch.is_grad_enabled():
+                x = ckpt.checkpoint(fn, *args, use_reentrant=False)
+            else:
+                x = fn(*args)
+        h = head_mod(x, e, per_frame, nf) if head_mod is not None \
+            else _head(model, x, e, per_frame, nf)
+        return _unpatchify(model, h, gf, gh, gw)
+
+
+def _unpatchify(model, x, gf, gh, gw):
+    """[B, L, out_dim*prod(patch)] -> [B, C, F, H, W].
+
+    Upstream's `WanModel.unpatchify` takes the grid as a *tensor* and calls
+    `.tolist()` on it, which builds a CPU tensor and syncs -- neither is legal
+    inside a CUDA graph capture. The grid is statically known here (it comes
+    from the patch-embedding output shape), so do the same reshape with Python
+    ints. Equivalence to upstream is asserted in scripts/verify_blockcausal.py.
+    """
+    p0, p1, p2 = model.patch_size
+    c = model.out_dim
+    b = x.shape[0]
+    u = x[:, :gf * gh * gw].view(b, gf, gh, gw, p0, p1, p2, c)
+    u = u.permute(0, 7, 1, 4, 2, 5, 3, 6)          # b c f p0 h p1 w p2
+    return u.reshape(b, c, gf * p0, gh * p1, gw * p2)
+
+
+# --------------------------------------------------------------- entry points
+def block_forward(model, z, t, t_start, rope, ctx, ctx_lens, kv=None,
+                  collect=None, dtype=torch.bfloat16, time_scale=1000.0,
+                  prefix_upto=None, per_frame=False, grad_checkpoint=False,
+                  tbl=None, emb=None):
+    """Velocity for ONE block of latent frames against the clean prefix.
+
+    z        [B, C, F, H, W] noisy latents of the current block
+    t        float, or [F] flow fractions in [0,1] when per_frame
+    t_start  absolute latent-frame index of z[:, :, 0] (drives temporal RoPE)
+    kv       BufferKV (inference) | TrainKV (training) | None (self-attention only)
+    collect  list receiving this block's per-layer (k, v), or None
+    tbl      precomputed RoPE table; overrides t_start. The CUDA-graph path
+             supplies it from a static buffer, since it is the one input that
+             changes with the absolute frame index.
+    """
+    if tbl is None:
+        tbl = rope.span(t_start, z.shape[2])
+
+    def factory(li):
+        def kv_ctx(k, v):
+            if collect is not None:
+                collect.append((k, v))
+            if kv is None:
+                return k, v, None
+            if prefix_upto is None:
+                return kv.context(li, k, v)
+            return kv.context(li, k, v, upto=prefix_upto)
+        return kv_ctx
+
+    return _run(model, z, t, tbl, ctx, ctx_lens, factory, dtype, time_scale,
+                per_frame, grad_checkpoint, emb=emb)
+
+
+@torch.no_grad()
+def build_clean_kv(model, z, rope, ctx, ctx_lens, world_frames, block_frames,
+                   dtype=torch.bfloat16, time_scale=1000.0, rope_gap=0):
+    """Clean (t=0) K/V for a whole clip, laid down exactly as deployment lays it
+    down: the world primed as one bidirectional block, then each event block
+    committed after being re-run at t=0. z: [B, C, F, H, W] clean latents.
+
+    `rope_gap` shifts the *event* frames' temporal indices further from the
+    world's -- see `event_rope_index`.
+    """
+    kv = TrainKV(len(model.blocks))
+    F, S = z.shape[2], rope.seq
+    spans, f = ([(0, world_frames)] if world_frames else []), world_frames
+    while f < F:
+        n = min(block_frames, F - f)
+        spans.append((f, n))
+        f += n
+    for t0, n in spans:
+        collect = []
+        idx = t0 if t0 < world_frames or t0 == 0 else t0 + rope_gap
+        block_forward(model, z[:, :, t0:t0 + n], 0.0, idx, rope, ctx, ctx_lens,
+                      kv=kv, collect=collect, dtype=dtype,
+                      time_scale=time_scale, prefix_upto=t0 * S)
+        for li, (k, v) in enumerate(collect):
+            kv.append(li, k, v)
+    return kv
+
+
+def event_rope_index(frame, world_frames, rope_gap):
+    """Absolute temporal index for RoPE, with the world pushed into the past.
+
+    RoPE is relative: attention depends only on index differences. In a long
+    stream the event window slides but the world block is *pinned*, so by unit
+    60 the current query sits ~180 latent frames from the world -- a relative
+    distance training never showed the model if a clip is only 21 frames long,
+    and rotary attention degrades badly off-distribution.
+
+    Deployment produces that geometry for free (evicted events leave a real gap
+    in the cache). Training has to simulate it, which costs nothing: keep the
+    world at 0..world_frames and shift every event frame by a random gap. The
+    content stays contiguous; only the positional distance changes, which is
+    exactly the axis that needs covering.
+    """
+    return frame if frame < world_frames else frame + rope_gap
+
+
+def parallel_blocks_forward(model, z_noisy, t_rows, starts, rope, ctx, ctx_lens,
+                            kv, block_frames, dtype=torch.bfloat16,
+                            time_scale=1000.0, grad_checkpoint=True):
+    """Every event block of a clip in ONE forward, as batch rows.
+
+    Row i holds block `starts[i]`, carries its own noise level `t_rows[i]` and
+    its own absolute RoPE offset, and attends over
+    [clean prefix[:starts[i]*S] ; its own noisy tokens]. Because the visible key
+    set depends only on the row -- never on the query index -- this costs a
+    [nb, 1, 1, Lk] key-padding mask instead of a dense [Lq, Lk] one (see
+    sdpa_attention's key-padding fast path).
+
+    z_noisy [1, C, F, H, W] full-clip noisy latents; returns [nb, C, b, H, W].
+    """
+    S, nb = rope.seq, len(starts)
+    blocks = torch.cat([z_noisy[:, :, s:s + block_frames] for s in starts], 0)
+    tbl = torch.stack([rope.span(s, block_frames).squeeze(1) for s in starts]
+                      ).unsqueeze(2)                       # [nb, b*S, 1, c]
+    max_prefix = max(starts) * S
+    # keys are [own block ; clean prefix], so row i keeps the first
+    # block_frames*S + starts[i]*S of them (see TrainKV.context)
+    k_lens = torch.tensor([block_frames * S + s * S for s in starts],
+                          device=z_noisy.device, dtype=torch.long)
+    ctx_b = ctx.expand(nb, -1, -1) if ctx.shape[0] == 1 else ctx
+    cl = ctx_lens.expand(nb) if ctx_lens.numel() == 1 else ctx_lens
+
+    def factory(li):
+        return lambda k, v: kv.context(li, k, v, upto=max_prefix, rows=nb,
+                                       k_lens=k_lens)
+
+    return _run(model, blocks, t_rows, tbl, ctx_b, cl, factory, dtype,
+                time_scale, False, grad_checkpoint)
+
+
+def block_starts(num_frames, world_frames, block_frames):
+    return [s for s in range(world_frames, num_frames, block_frames)
+            if s + block_frames <= num_frames]

wanstreamer/core.py

@@ -0,0 +1,151 @@
+"""Block-causal streaming forward for Wan2.1, corrected and de-overheaded.
+
+Corrections relative to streaming_blocks.py, each backed by a measurement
+recorded in PROGRESS.md:
+
+  1. TIMESTEP SCALE. final.pt was fine-tuned with t in [0,1]
+     (train_streaming.py:152 draws torch.rand()), not the scheduler's [0,1000].
+     Callers must pass t in [0,1]; `timestep_to_train_scale` does the conversion.
+     Measured: normalised flow error 0.499 -> 0.168.
+  2. TEMPORAL RoPE. The old path gave every frame temporal index 0. Here each
+     latent frame is rotated at its true absolute index, so cached keys carry
+     real temporal position. Measured: 0.180 -> 0.107 (uniform, t=0.5).
+  3. NO CAUSAL MASK. The K/V cache holds only past+current frames, so full
+     attention over it is already block-causal. The old code passed causal=True
+     with Lq != Lk, which imposes a spurious raster ordering *within* a frame.
+     Measured: slightly better error AND 1.56x faster (no mask materialisation).
+
+Overhead reductions, targeting the profile
+"""
+import torch
+
+from wan.modules.attention import flash_attention
+
+from .rope import RopeTable, apply_rope
+from .kvcache import StreamingKVCache
+
+
+def timestep_to_train_scale(t_val, num_train_timesteps=1000):
+    """Scheduler timestep (0..1000) -> the t in [0,1] final.pt was trained on."""
+    return float(t_val) / float(num_train_timesteps)
+
+
+class ModulationCache:
+    """Per-block AdaLN modulation chunks for one timestep.
+
+    `(block.modulation + e0).chunk(6)` is identical for every frame at a given
+    denoising step, but the old code recomputed it per block *per frame*.
+    """
+
+    def __init__(self, blocks, e0):
+        self.chunks = []
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            for blk in blocks:
+                self.chunks.append((blk.modulation + e0).chunk(6, dim=1))
+
+    def __getitem__(self, i):
+        return self.chunks[i]
+
+
+def block_forward(blk, x, mod, rope_tbl, t_index, cache, layer, ctx, ctx_lens):
+    """One WanAttentionBlock in block-causal streaming mode.
+
+    x: [B, S, dim] tokens of the CURRENT latent frame only.
+    Returns the updated x. K/V for this frame is written (uncommitted) to `cache`.
+    """
+    ec = mod
+    sa_in = blk.norm1(x).float() * (1 + ec[1]) + ec[0]
+
+    b, s = sa_in.shape[0], sa_in.shape[1]
+    n = blk.num_heads
+    d = blk.dim // n
+    sa = blk.self_attn
+
+    q = sa.norm_q(sa.q(sa_in)).view(b, s, n, d)
+    k = sa.norm_k(sa.k(sa_in)).view(b, s, n, d)
+    v = sa.v(sa_in).view(b, s, n, d)
+
+    tbl = rope_tbl.frame(t_index)
+    q = apply_rope(q, tbl)
+    k = apply_rope(k, tbl)
+
+    # Write current chunk, then attend over past+current as one view.
+    cache.write(layer, k.to(cache.k.dtype), v.to(cache.v.dtype))
+    ctx_k, ctx_v = cache.context(layer, s)
+
+    y = flash_attention(q=q.to(ctx_k.dtype), k=ctx_k, v=ctx_v,
+                        window_size=(-1, -1), causal=False)
+    y = sa.o(y.flatten(2))
+
+    with torch.amp.autocast('cuda', dtype=torch.float32):
+        x = x + y * ec[2]
+    x = x + blk.cross_attn(blk.norm3(x), ctx, ctx_lens)
+    y = blk.ffn(blk.norm2(x).float() * (1 + ec[4]) + ec[3])
+    with torch.amp.autocast('cuda', dtype=torch.float32):
+        x = x + y * ec[5]
+    return x
+
+
+@torch.no_grad()
+def frame_forward(model, latent_frame, e, mod, rope_tbl, t_index, cache,
+                  ctx, ctx_lens):
+    """Denoise ONE latent frame against the cached clean past.
+
+    latent_frame: [B, C, 1, H, W]. Returns predicted velocity [C, 1, H, W].
+    Does not commit the frame's K/V -- the caller commits once the chunk is final.
+    """
+    x = model.patch_embedding(latent_frame)
+    grid = torch.stack([torch.tensor(x.shape[2:], dtype=torch.long, device=x.device)
+                        for _ in range(x.shape[0])])
+    x = x.flatten(2).transpose(1, 2)
+
+    for i, blk in enumerate(model.blocks):
+        x = block_forward(blk, x, mod[i], rope_tbl, t_index, cache, i, ctx, ctx_lens)
+
+    return model.unpatchify(model.head(x, e), grid)[0]
+
+
+@torch.no_grad()
+def sequence_forward(model, latents, e, mod, rope_tbl, cache, ctx, ctx_lens,
+                     start_index=0, commit=True):
+    """Run a run of latent frames causally, committing each as it completes.
+
+    Used to (a) prime the cache from clean context frames and (b) evaluate the
+    model over a whole sequence for diagnostics. latents: [B, C, F, H, W].
+    """
+    outs = []
+    for f in range(latents.shape[2]):
+        out = frame_forward(model, latents[:, :, f:f + 1], e, mod, rope_tbl,
+                            start_index + f, cache, ctx, ctx_lens)
+        outs.append(out)
+        if commit:
+            cache.commit(rope_tbl.seq)
+    return torch.cat(outs, dim=1)
+
+
+def make_rope_table(model, h_patches, w_patches, max_frames, device):
+    return RopeTable(model.freqs.to(device), h_patches, w_patches, max_frames, device)
+
+
+def make_cache(model, tokens_per_frame, max_frames, device, dtype=torch.bfloat16):
+    n = model.num_heads
+    d = model.dim // n
+    return StreamingKVCache(
+        num_layers=len(model.blocks),
+        max_tokens=tokens_per_frame * max_frames,
+        num_heads=n, head_dim=d, device=device, dtype=dtype,
+        scratch_tokens=tokens_per_frame)
+
+
+def latent_geometry(width, height, vae_stride=(4, 8, 8), patch_size=(1, 2, 2)):
+    """Pixel (W,H) -> latent (H_lat, W_lat) and patch grid (H_p, W_p).
+
+    Mirrors WanModel/run_demo: target_shape[2] = size[1]//vae_stride[1] (height),
+    target_shape[3] = size[0]//vae_stride[2] (width).
+    """
+    h_lat = height // vae_stride[1]
+    w_lat = width // vae_stride[2]
+    if h_lat % patch_size[1] or w_lat % patch_size[2]:
+        raise ValueError(f'{width}x{height} -> latent {h_lat}x{w_lat} not '
+                         f'divisible by patch size {patch_size[1:]}')
+    return h_lat, w_lat, h_lat // patch_size[1], w_lat // patch_size[2]

wanstreamer/data.py

@@ -0,0 +1,65 @@
+"""Teacher-latent dataset plus the shared flow-matching conventions."""
+import glob
+import os
+
+import torch
+from torch.utils.data import Dataset
+
+
+class TeacherLatents(Dataset):
+    """Teacher ODE samples: clean latents [16, F, H, W] + their prompt index."""
+
+    def __init__(self, root, prompts_path, frames=None):
+        self.files = sorted(glob.glob(os.path.join(root, '*.pt')))
+        if not self.files:
+            raise RuntimeError(f'no .pt latents under {root}')
+        blob = torch.load(prompts_path, map_location='cpu')
+        self.pos = blob['pos']            # [P, 512, 4096] fp16
+        self.neg = blob['neg']            # [1, 512, 4096] fp16
+        self.prompts = blob['prompts']
+        self.frames = frames
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, i):
+        d = torch.load(self.files[i], map_location='cpu', weights_only=False)
+        lat = d['latents'].float()
+        if self.frames is not None:
+            lat = lat[:, :self.frames]
+        pi = int(d['prompt_idx'])
+        return lat, self.pos[pi].float(), pi
+
+
+def shifted_uniform_t(n, shift=5.0, device='cpu', generator=None):
+    """Flow fractions distributed as Wan's inference sigma schedule.
+
+    Wan samples with shift=5, which concentrates steps near t=1. Training with
+    plain U(0,1) would under-serve exactly the region the student spends most of
+    its few steps in, so draw from the same reparameterisation the sampler uses:
+        t = shift*u / (1 + (shift-1)*u),  u ~ U(0,1).
+    """
+    u = torch.rand(n, device=device, generator=generator)
+    return shift * u / (1 + (shift - 1) * u)
+
+
+def add_noise(z0, t, noise=None):
+    """Wan's rectified-flow convention, matching train_streaming.py:
+
+        z_t    = (1 - t) * z0 + t * noise
+        target = noise - z0        (velocity the model predicts)
+
+    t broadcasts over [B] or [B, F]; z0 is [B, C, F, H, W].
+    """
+    if noise is None:
+        noise = torch.randn_like(z0)
+    while t.dim() < z0.dim():
+        t = t.unsqueeze(-1)
+    return (1 - t) * z0 + t * noise, noise - z0, noise
+
+
+def x0_from_velocity(z_t, v, t):
+    """Invert the flow parameterisation: z0 = z_t - t * v."""
+    while t.dim() < z_t.dim():
+        t = t.unsqueeze(-1)
+    return z_t - t * v

wanstreamer/dmd.py

@@ -0,0 +1,90 @@
+"""Distribution Matching Distillation for the blockcausal student.
+
+The generator loss is the DMD2 form. Given a generated clean latent `x0`, noise
+it to some level t and ask two score networks where they think it came from:
+
+    x0_real = teacher(x_t, t)      with CFG   -- the distribution we want
+    x0_fake = critic(x_t, t)                  -- the distribution we have
+    grad    = (x0_fake - x0_real) / normalizer
+    L_gen   = 0.5 * || x0 - stopgrad(x0 - grad) ||^2
+
+so d L_gen / d x0 == grad, which is the KL gradient between the two
+distributions. The normalizer is DMD2 scale invariant one. Which is mean |x0 - x0_real|
+. Without which the loss magnitude swings by orders of magnitude across t and
+the generator LR cannot be set at all.
+
+The critic is trained in the ordinary way, a flow-matching loss on the student's
+own samples, so it tracks a moving target. That is the whole reason DMD needs
+two time scales: the critic must stay ahead of the generator.
+
+Why the teacher is allowed to be bidirectional: it only ever scores a finished
+clip, never generates one. Causality is a constraint on the students sampling
+procedure, not on the reference distribution.
+"""
+import torch
+import torch.nn.functional as tnnF
+
+from . import blockcausal as bc
+from .data import add_noise
+
+
+def bidirectional_velocity(model, z, t, rope, ctx, ctx_lens, dtype,
+                           time_scale=1000.0, grad_checkpoint=False):
+    """Full bidirectional forward over a whole clip.
+
+    `block_forward` with no K/V prefix and the entire clip as one block *is* the
+    bidirectional forward -- attention runs over the clip and nothing else -- so
+    teacher, critic and student all go through one code path and one RoPE
+    implementation.
+    """
+    return bc.block_forward(model, z, t, 0, rope, ctx, ctx_lens, kv=None,
+                            dtype=dtype, time_scale=time_scale,
+                            grad_checkpoint=grad_checkpoint)
+
+
+def cfg_velocity(model, z, t, rope, ctx_pos, ctx_neg, ctx_lens, dtype,
+                 guidance, time_scale=1000.0):
+    """Teacher velocity with classifier-free guidance, cond and uncond batched."""
+    if not guidance or guidance == 1.0:
+        return bidirectional_velocity(model, z, t, rope, ctx_pos, ctx_lens,
+                                      dtype, time_scale)
+    z2 = torch.cat([z, z], 0)
+    ctx2 = torch.cat([ctx_pos, ctx_neg], 0)
+    lens2 = ctx_lens.expand(2)
+    v = bidirectional_velocity(model, z2, t, rope, ctx2, lens2, dtype, time_scale)
+    vc, vu = v[0:1], v[1:2]
+    return vu + guidance * (vc - vu)
+
+
+@torch.no_grad()
+def dmd_gradient(x0, t, z_t, v_real, v_fake):
+    """DMD2's KL gradient wrt the generated clip, for the whole clip at once.
+
+    One normalizer for the clip rather than one per block: per-block
+    normalisation would rescale each block's gradient by its own error and so
+    quietly up-weight whichever block is currently worst.
+    """
+    while t.dim() < z_t.dim():
+        t = t.unsqueeze(-1)
+    x0_real = (z_t - t * v_real).float()
+    x0_fake = (z_t - t * v_fake).float()
+    normalizer = (x0.float() - x0_real).abs().mean().clamp_min(1e-4)
+    grad = torch.nan_to_num((x0_fake - x0_real) / normalizer)
+    return grad, {'dmd_grad_norm': float(grad.norm()),
+                  'dmd_normalizer': float(normalizer),
+                  'dmd_real_fake_gap': float((x0_real - x0_fake).abs().mean())}
+
+
+def dmd_surrogate(x0_hat, grad, scale=1.0):
+    """Scalar whose gradient wrt x0_hat is exactly `grad` (times scale)."""
+    return scale * 0.5 * tnnF.mse_loss(x0_hat.float(),
+                                       (x0_hat.detach() - grad).float())
+
+
+def critic_loss(model, x0, t, rope, ctx, ctx_lens, dtype, time_scale=1000.0,
+                grad_checkpoint=True, noise=None):
+    """Flow-matching loss for the fake-score network on student samples."""
+    z_t, target, _ = add_noise(x0, t, noise)
+    v = bidirectional_velocity(model, z_t, t, rope, ctx, ctx_lens, dtype,
+                               time_scale, grad_checkpoint=grad_checkpoint)
+    return tnnF.mse_loss(v.float(), target.float())

wanstreamer/fsdp.py

@@ -0,0 +1,213 @@
+"""FSDP2 sharding for the block-causal student, and why it needs help.
+
+Both trainers here are data-parallel today: Stage 1 wraps its step in a module
+so DDP's hooks fire, and DMD all-reduces gradients by hand. Both keep a full
+fp32 copy of the model, its gradients and its AdamW moments on every GPU --
+16 bytes per parameter, 21 GB at 1.3B, and 224 GB at 14B, which is why nothing
+larger than 1.3B trains on this box at all. Sharding those three tensor sets
+across N ranks is the only thing that changes that, so it has to work before a
+larger base is worth discussing.
+
+THE TRAP THIS MODULE EXISTS FOR. FSDP2 all-gathers a sharded module's
+parameters in a pre-forward hook installed on *that module*. `blockcausal._run`
+never calls a `WanAttentionBlock` -- it passes the block to `_layer`, which
+reaches into `blk.self_attn.q`, `blk.norm1`, `blk.ffn` and so on. Shard at the
+block and that hook never fires, so the layer runs against parameters that are
+still 1/N shards. It is structurally the same mistake DDP already invites here
+("DDP does not sync gradients if you call the bare model
+through helper functions"), and it is why this port is worth debugging at 1.3B.
+
+Measured, on this codebase, the mis-shard is LOUD rather than silent: the first
+parameter touched outside a forward is `blk.modulation`, and DTensor refuses
+the mixed operand
+
+    RuntimeError: aten.add.Tensor got mixed torch.Tensor and DTensor
+
+(scripts/verify_fsdp.py check 3 asserts exactly this). That guard is DTensor's,
+not ours, and it only holds while every parameter stays a DTensor; it is not a
+reason to leave the shard boundary in the wrong place. The `stragglers` check
+in `shard_model` is the part that does not depend on someone else's invariant.
+
+`CausalBlock` and `CausalHead` fix it by being real modules whose forward *is*
+the computation, so the shard boundary and the call boundary coincide.
+`scripts/verify_fsdp.py` is the control: it asserts a sharded forward matches
+the single-GPU forward, and that a deliberately mis-sharded one does not.
+
+What is sharded: the 30 transformer blocks (98% of parameters) plus, when they
+are big enough to matter, the head and the embeddings -- all of which are
+entered through their own `__call__` and so need no wrapper.
+"""
+import torch
+from torch.distributed.checkpoint.state_dict import (
+    get_model_state_dict, StateDictOptions
+)
+
+from . import blockcausal as bc
+
+
+class CausalBlock(torch.nn.Module):
+    """One WanAttentionBlock as an FSDP shard unit.
+
+    Holds the block as a child so `fully_shard(self)` shards the block's
+    parameters, and runs `_layer` in its own forward so entering the shard is
+    the same act as entering the computation.
+    """
+
+    def __init__(self, blk):
+        super().__init__()
+        self.blk = blk
+
+    def forward(self, x, e0, tbl, kv_ctx, ctx, ctx_lens, dtype, per_frame,
+                n_frames):
+        with torch.amp.autocast('cuda', enabled=False):
+            ec = (self.blk.modulation + e0).chunk(6, dim=1)
+        return bc._layer(self.blk, x, ec, tbl, kv_ctx, ctx, ctx_lens, dtype,
+                         per_frame, n_frames)
+
+
+class CausalHead(torch.nn.Module):
+    """The output head as an FSDP shard unit, for the same reason: `_head`
+    reaches through `model.head` to `model.head.head` and reads
+    `model.head.modulation` directly."""
+
+    def __init__(self, head):
+        super().__init__()
+        self.head = head
+
+    def forward(self, x, e, per_frame, n_frames):
+        return bc._head_from(self.head, x, e, per_frame, n_frames)
+
+
+def attach_causal_modules(model):
+    """Install the per-layer / head modules `blockcausal._run` will enter, and
+    make each of them the *only* registered path to its parameters.
+
+    Re-registering is the subtle half. `CausalBlock(blk)` holds `blk` as a
+    child, but `model.blocks[i]` still holds it too, so every block parameter
+    now has two names in the module tree. FSDP2 shards by walking that tree and
+    skipping what a child FSDP module already owns; reached by its second name
+    a parameter looks unmanaged, and sharding it again dies with "Cannot
+    concatenate overlapping meshes". So `blocks` and `head` are demoted to
+    plain attributes -- still there for `len(model.blocks)` and for the
+    unsharded code path, invisible to `named_parameters()`.
+
+    Separated from `shard_model` so the correctness control can run the wrapped
+    path *without* FSDP and show that wrapping alone changes nothing.
+    """
+    if getattr(model, 'causal_layers', None) is not None:
+        return model
+    blocks = list(model.blocks)
+    head = model.head
+    layers = torch.nn.ModuleList(CausalBlock(b) for b in blocks)
+    chead = CausalHead(head)
+    del model.blocks, model.head          # drop from _modules
+    model.causal_layers = layers
+    model.causal_head = chead
+    # object.__setattr__, because nn.Module.__setattr__ would re-register an
+    # nn.Module value and undo exactly what the `del` above achieved.
+    object.__setattr__(model, 'blocks', blocks)
+    object.__setattr__(model, 'head', head)
+    return model
+
+
+def shard_model(model, mesh=None, reshard_after_forward=True, mp_policy=None,
+                ignored_params=None):
+    """Shard a WanModel across `mesh` with FSDP2, in place. Returns the model.
+
+    Call BEFORE moving to device and before building the optimizer: FSDP2
+    replaces `.weight` with a DTensor, and an optimizer built over the
+    unsharded parameters would hold stale references.
+
+    `reshard_after_forward=False` keeps parameters gathered between forward and
+    backward. That trades memory for collectives, and it matters here more than
+    in ordinary training: one DMD iteration runs the student a few dozen times
+    inside its rollout, so resharding after each one pays 30 all gathers per
+    forward to reclaim memory the (no_grad) rollout never needed back.
+
+    There is deliberately no root `fully_shard(model)`. The root's hook fires
+    on `model.forward`, and nothing here ever calls it. `block_forward` drives
+    the model from outside. A parameter caught only by the root group would
+    therefore stay sharded through every forward, silently. Everything is
+    sharded at a module that is genuinely entered instead, and the assertion at
+    the end refuses to return a model where that did not hold.
+    """
+    from torch.distributed.fsdp import fully_shard
+
+    # An explicit, NAMED mesh. `fully_shard(mesh=None)` synthesises one per
+    # call with `mesh_dim_names=None`, and composing shards over children then
+    # dies in `_init_sharded_param` trying to concatenate those names.
+    if mesh is None:
+        from torch.distributed.device_mesh import init_device_mesh
+        import torch.distributed as dist
+        mesh = init_device_mesh('cuda', (dist.get_world_size(),),
+                                mesh_dim_names=('dp',))
+    kw = {'mesh': mesh, 'reshard_after_forward': reshard_after_forward}
+    if mp_policy is not None:
+        kw['mp_policy'] = mp_policy
+    # `ignored_params` is how the LoRA critic survives sharding its own base.
+    # The frozen teacher is bf16 with no gradients; the adapter living inside
+    # the same blocks is fp32 and trainable, and it is 32 M parameters (small
+    # enough that replicating it and all reducing by hand is cheaper than
+    # sharding, and it keeps mixed dtypes out of a single FSDP parameter group.
+    if ignored_params:
+        kw['ignored_params'] = set(ignored_params)
+
+    attach_causal_modules(model)
+    units = list(model.causal_layers) + [model.causal_head]
+    # These four are invoked through their own __call__ already (patch_embedding
+    # in `_run`, text_embedding in the trainers, time_embedding and
+    # time_projection in `time_embed`), so they need no wrapper.
+    units += [model.patch_embedding, model.text_embedding,
+              model.time_embedding, model.time_projection]
+    for m in units:
+        fully_shard(m, **kw)
+
+    ignore = set(id(p) for p in (ignored_params or ()))
+    stragglers = [n for n, p in model.named_parameters()
+                  if id(p) not in ignore
+                  and not isinstance(p, torch.distributed.tensor.DTensor)]
+    if stragglers:
+        raise RuntimeError(
+            f'{len(stragglers)} parameter(s) were not sharded and are not '
+            f'reachable from any entered module -- they would train '
+            f'unsynchronised: {stragglers[:8]}')
+    return model
+
+
+def set_grad_sync(model, flag):
+    """FSDP2's equivalent of `DDP.no_sync()`, over every shard unit.
+
+    With gradient accumulation, leaving this on costs one reduce-scatter per
+    micro-step instead of one per optimizer step. It is not a correctness issue
+    -- reducing each micro-step and summing equals summing then reducing but
+    at `--accum 4` it is four times the collectives for the same gradient.
+    """
+    for m in model.modules():
+        if hasattr(m, 'set_requires_gradient_sync'):
+            m.set_requires_gradient_sync(flag)
+
+
+def full_state_dict(model):
+    """Unsharded fp32 state dict on CPU, in the same key layout the existing
+    checkpoints use, so `demo.py --weights` keeps working unchanged.
+
+    Under FSDP2 `model.state_dict()` returns DTensors -- saving those would
+    produce a checkpoint that only reloads onto an identical mesh, and every
+    evaluation script here loads onto one GPU.
+    """
+
+    sd = get_model_state_dict(
+        model, options=StateDictOptions(full_state_dict=True, cpu_offload=True))
+    # `attach_causal_modules` renamed the shard units, so the keys now read
+    # `causal_layers.7.blk.*` and `causal_head.head.*`. Put them back under the
+    # names WanModel.load_state_dict expects, so demo.py and every other
+    # evaluation script keep loading these checkpoints on a single GPU.
+    out = {}
+    for k, v in sd.items():
+        if k.startswith('causal_layers.'):
+            i, rest = k[len('causal_layers.'):].split('.blk.', 1)
+            k = f'blocks.{i}.{rest}'
+        elif k.startswith('causal_head.head.'):
+            k = 'head.' + k[len('causal_head.head.'):]
+        out[k] = v
+    return out

wanstreamer/graphrunner.py

@@ -0,0 +1,97 @@
+"""CUDA-graph capture for the streaming forward.
+
+Why this and not a smaller optimisation: the prior profiling of this pipeline
+found ~7,300 kernel launches and only ~17% arithmetic in a single frame-forward,
+and the sweep here reproduces the signature exactly -- a 1-latent-frame unit and
+a 3-latent-frame unit cost nearly the same per forward (119 ms vs 168 ms) even
+though the second does 3x the work. The bottleneck is launch overhead, so the
+fix has to remove launches, not work. A CUDA graph replays the whole forward as
+one submission.
+
+The capture is only valid because streaming reaches a genuine steady state:
+
+  * the event K/V window is bounded, so once it is full `cache.length` stops
+    changing and every shape in the forward is constant;
+  * the cache is one preallocated buffer written in place at a constant offset,
+    so the graph can own those writes;
+  * only three things vary between calls -- the latents, the timestep, and the
+    RoPE table slice for the current absolute frame index -- and all three are
+    copied into static input buffers before replay.
+
+So: run eagerly until the window fills, capture once, then replay. Anything that
+would change a shape (a different block size, the window not yet full) falls
+back to eager automatically.
+"""
+import torch
+
+
+class GraphedForward:
+    """Captures `fn(z, tbl, e, e0) -> v` for one fixed streaming configuration."""
+
+    def __init__(self, fn, z_shape, tbl_shape, e_shape, e0_shape, device,
+                 dtype=torch.float32, warmup=3):
+        self.fn = fn
+        self.device = device
+        self.z = torch.zeros(z_shape, device=device, dtype=dtype)
+        self.tbl = torch.zeros(tbl_shape, device=device, dtype=torch.complex64)
+        self.e = torch.zeros(e_shape, device=device, dtype=torch.float32)
+        self.e0 = torch.zeros(e0_shape, device=device, dtype=torch.float32)
+        self.graph = None
+        self.warmup = warmup
+        self.out = None
+        self.key = None
+
+    def capture(self):
+        # Warm up on a side stream first: cuDNN/cuBLAS pick algorithms and
+        # allocate workspaces on first call, and that must not happen during
+        # capture.
+        s = torch.cuda.Stream()
+        s.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(s):
+            for _ in range(self.warmup):
+                out = self.fn(self.z, self.tbl, self.e, self.e0)
+        torch.cuda.current_stream().wait_stream(s)
+        torch.cuda.synchronize()
+
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph):
+            self.out = self.fn(self.z, self.tbl, self.e, self.e0)
+        return self
+
+    def __call__(self, z, tbl, e, e0):
+        self.z.copy_(z)
+        self.tbl.copy_(tbl)
+        self.e.copy_(e)
+        self.e0.copy_(e0)
+        self.graph.replay()
+        return self.out
+
+
+class SteadyStateGraphs:
+    """One graph per distinct (cache length, block size) the stream settles into.
+
+    A stream with a bounded window visits at most a handful of cache lengths
+    before it saturates, so keying on the length is enough; the dict never grows
+    without bound. `enabled=False` makes every call fall through to eager, which
+    is what the correctness comparison in scripts/bench_graph.py uses.
+    """
+
+    def __init__(self, enabled=True, max_graphs=4):
+        self.enabled = enabled
+        self.graphs = {}
+        self.max_graphs = max_graphs
+        self.captures = 0
+        self.replays = 0
+
+    def get(self, key, make):
+        if not self.enabled:
+            return None
+        g = self.graphs.get(key)
+        if g is None:
+            if len(self.graphs) >= self.max_graphs:
+                return None
+            g = make().capture()
+            self.graphs[key] = g
+            self.captures += 1
+        self.replays += 1
+        return g

wanstreamer/kvcache.py

@@ -0,0 +1,105 @@
+"""Preallocated K/V cache for block-causal streaming.
+
+Replaces streaming_blocks.KVFrameCache, which stored a deque of per-frame tensors
+and rebuilt the context with `torch.cat` on every `get()`. The profile in
+PROGRESS.md §6 counted 302 `aten::cat` calls per single frame-forward.
+
+Here the cache is one preallocated buffer per layer and the context is a *view*,
+so steady-state streaming does no allocation and no copying beyond writing the
+new frame's K/V.
+
+Commit semantics matter for multi-step denoising. The current chunk's K/V changes
+at every denoising step, so it must not be committed until the chunk is final:
+
+    for step in steps:                     # chunk still noisy
+        cache.write(layer, k, v)           # scratch region at [len, len+S)
+        ctx_k, ctx_v = cache.context(layer, S)   # past + current, one view
+    cache.commit(S)                        # chunk finalised, becomes past
+
+Because the cache only ever holds past + current frames, full attention over that
+view IS block-causal -- bidirectional within the current chunk, unrestricted over
+the past, and no future keys exist to leak. No mask is needed; see
+tests/test_attention_fallback.py case 5.
+"""
+import torch
+
+
+class StreamingKVCache:
+    """Per-layer ring-free preallocated K/V store with an uncommitted scratch tail."""
+
+    def __init__(self, num_layers, max_tokens, num_heads, head_dim,
+                 device, dtype=torch.bfloat16, scratch_tokens=None):
+        self.num_layers = num_layers
+        self.max_tokens = max_tokens
+        self.scratch = scratch_tokens if scratch_tokens is not None else max_tokens
+        cap = max_tokens + self.scratch
+        self.k = torch.zeros(num_layers, cap, num_heads, head_dim,
+                             device=device, dtype=dtype)
+        self.v = torch.zeros(num_layers, cap, num_heads, head_dim,
+                             device=device, dtype=dtype)
+        self.length = 0          # committed tokens
+        self._pending = 0        # tokens written to scratch, not yet committed
+        # Monotone count of eviction events. `BufferPrefixKV` reconstructs a
+        # recorded step's prefix as `buffer[:upto]`, which is only that step's
+        # actual prefix if nothing has been evicted since it was recorded --
+        # eviction shifts surviving keys down and discards the oldest outright.
+        # The self-forcing trainer compares this counter against the value it
+        # stamped on each recorded block and refuses to take the gradient on a
+        # block whose prefix has since moved. See stream.trim_to_window.
+        self.evictions = 0
+
+    def reset(self):
+        self.length = 0
+        self._pending = 0
+        self.evictions = 0
+
+    def write(self, layer, k, v):
+        """Write the current (uncommitted) chunk's K/V at [length, length+S)."""
+        s = k.shape[1]
+        if self.length + s > self.k.shape[1]:
+            raise RuntimeError(
+                f'K/V cache overflow: {self.length} committed + {s} pending > '
+                f'capacity {self.k.shape[1]}. Raise max_tokens or evict.')
+        self.k[layer, self.length:self.length + s].copy_(k[0])
+        self.v[layer, self.length:self.length + s].copy_(v[0])
+        self._pending = s
+
+    def context(self, layer, s):
+        """View of past + current: ([1, length+s, n, d], same for v). No copy."""
+        end = self.length + s
+        return (self.k[layer, :end].unsqueeze(0),
+                self.v[layer, :end].unsqueeze(0))
+
+    def commit(self, s):
+        """Promote the pending chunk to committed history."""
+        self.length += s
+        self._pending = 0
+
+    def evict_front(self, s, protect=0):
+        """Drop `s` tokens from the front of the evictable region.
+
+        `protect` pins the first `protect` tokens permanently -- used to keep the
+        WORLD block resident while sliding a window over the event stream, which
+        is the behaviour the papers describe (W is persistent, events stream).
+        Eviction therefore removes the OLDEST EVENTS, not the world.
+        """
+        if s <= 0:
+            return
+        evictable = self.length - protect
+        s = min(s, max(0, evictable))
+        if s == 0:
+            return
+        keep = self.length - protect - s          # events surviving the eviction
+        if keep > 0:
+            src = protect + s
+            self.k[:, protect:protect + keep].copy_(self.k[:, src:src + keep])
+            self.v[:, protect:protect + keep].copy_(self.v[:, src:src + keep])
+        self.length = protect + keep
+        self.evictions += 1
+
+    @property
+    def num_tokens(self):
+        return self.length
+
+    def memory_bytes(self):
+        return self.k.numel() * self.k.element_size() * 2

wanstreamer/lora.py

@@ -0,0 +1,95 @@
+"""Minimal LoRA, used to make the DMD critic share the teacher's weights.
+
+DMD needs three networks: the causal student, a frozen *real* score (the
+original bidirectional Wan) and a trainable *fake* score that tracks the
+student's own output distribution. Three full 1.4B copies plus the student's
+AdamW state does not fit on a 40 GB card, and the fake score's correct
+initialisation is exactly the real score anyway -- so the fake score is the same
+frozen base with a low-rank adapter on top, toggled by a flag:
+
+    with lora_enabled(base, False): v_real = ...      # teacher
+    with lora_enabled(base, True):  v_fake = ...      # critic
+
+That costs ~32 M trainable parameters instead of 1.4 B, starts the critic at the
+right place by construction, and leaves the base weights bit-identical between
+the two roles.
+"""
+from contextlib import contextmanager
+
+import torch
+import torch.nn as nn
+
+
+class LoRALinear(nn.Module):
+    def __init__(self, base: nn.Linear, rank=32, alpha=None):
+        super().__init__()
+        self.base = base
+        self.base.weight.requires_grad_(False)
+        if self.base.bias is not None:
+            self.base.bias.requires_grad_(False)
+        self.rank = rank
+        self.scale = (alpha or rank) / rank
+        self.a = nn.Parameter(torch.zeros(rank, base.in_features))
+        self.b = nn.Parameter(torch.zeros(base.out_features, rank))
+        nn.init.normal_(self.a, std=1.0 / rank)      # b stays zero -> starts as identity
+        self.enabled = True
+
+    def forward(self, x):
+        y = self.base(x)
+        if not self.enabled:
+            return y
+        h = nn.functional.linear(x.to(self.a.dtype), self.a)
+        return y + nn.functional.linear(h, self.b).to(y.dtype) * self.scale
+
+
+TARGETS = ('self_attn.q', 'self_attn.k', 'self_attn.v', 'self_attn.o',
+           'ffn.0', 'ffn.2')
+
+
+def inject_lora(model, rank=32, alpha=None, targets=TARGETS):
+    """Wrap the targeted Linears of every transformer block. Returns the new
+    parameters, and freezes everything else in the model."""
+    for p in model.parameters():
+        p.requires_grad_(False)
+    n = 0
+    for blk in model.blocks:
+        for name in targets:
+            parent, _, leaf = name.rpartition('.')
+            mod = blk.get_submodule(parent) if parent else blk
+            lin = getattr(mod, leaf) if not leaf.isdigit() else mod[int(leaf)]
+            if isinstance(lin, LoRALinear):
+                continue
+            wrapped = LoRALinear(lin, rank, alpha).to(lin.weight.device)
+            wrapped.a.data = wrapped.a.data.float()
+            wrapped.b.data = wrapped.b.data.float()
+            if leaf.isdigit():
+                mod[int(leaf)] = wrapped
+            else:
+                setattr(mod, leaf, wrapped)
+            n += 1
+    params = [p for p in model.parameters() if p.requires_grad]
+    total = sum(p.numel() for p in params)
+    return params, n, total
+
+
+def set_lora(model, on):
+    for m in model.modules():
+        if isinstance(m, LoRALinear):
+            m.enabled = on
+
+
+@contextmanager
+def lora_enabled(model, on):
+    prev = [m.enabled for m in model.modules() if isinstance(m, LoRALinear)]
+    set_lora(model, on)
+    try:
+        yield
+    finally:
+        for m, p in zip((m for m in model.modules()
+                         if isinstance(m, LoRALinear)), prev):
+            m.enabled = p
+
+
+def lora_state_dict(model):
+    return {k: v for k, v in model.state_dict().items()
+            if k.endswith('.a') or k.endswith('.b')}

wanstreamer/metrics.py

@@ -0,0 +1,183 @@
+"""Objective quality statistics for a streamed video.
+
+There is no reference video to compare against -- the stream is a continuation
+that never existed -- so these are no-reference statistics chosen to catch the
+three failure modes actually observed in this project (PROGRESS.md §8):
+
+  collapse   the model loses signal and output greys out.
+             -> latent std ratio vs the world, and pixel contrast.
+  blur       output stays plausible but soft, losing high frequencies.
+             -> Laplacian variance (standard sharpness proxy), relative to the
+                world frames produced by the same VAE at the same resolution, so
+                the number is a ratio against a known-good reference.
+  freeze/flicker  motion either stops or becomes incoherent.
+             -> mean |frame_t - frame_{t-1}|; a healthy talking-head sits well
+                above 0 and well below the world's own inter-frame delta * ~2.
+
+Every statistic is reported for the WORLD segment and the GENERATED segment
+separately, plus their ratio. Ratios near 1.0 mean the stream matches the quality
+of the bidirectional teacher output it continues -- that is the target.
+
+`drift` measures whether the stream wanders away from the established scene:
+cosine similarity in latent space between each generated frame and the mean world
+latent. A monotone decline is drift; a flat line is a stable identity.
+"""
+import numpy as np
+
+
+def _lap_var(gray):
+    """Variance of the 3x3 Laplacian -- higher is sharper. gray: [H, W] float."""
+    lap = (-4.0 * gray[1:-1, 1:-1]
+           + gray[:-2, 1:-1] + gray[2:, 1:-1]
+           + gray[1:-1, :-2] + gray[1:-1, 2:])
+    return float(lap.var())
+
+
+def _blockiness(gray, period):
+    """Ratio of gradient energy ON a `period`-aligned grid to gradient energy off
+    it. ~1.0 means no grid structure; >1.2 is visible blocking.
+
+    This exists because Laplacian sharpness is gameable: the highest-sharpness arm
+    in out_sweep/round2 scored 0.38 while decoding to a blocky checkerboard, and a
+    periodic artifact grid raises high-frequency energy exactly like real detail
+    does. Wan's VAE has spatial stride 8 and the DiT patch is 2 latent cells, so
+    artifacts land on 8- and 16-pixel grids; measuring those periods separates
+    structure from detail.
+    """
+    dv = np.abs(np.diff(gray, axis=1))       # [H, W-1], boundary between x,x+1
+    dh = np.abs(np.diff(gray, axis=0))
+    out = []
+    for d, axis in ((dv, 1), (dh, 0)):
+        n = d.shape[axis]
+        on_idx = np.arange(period - 1, n, period)
+        if len(on_idx) < 2:
+            continue
+        mask = np.zeros(n, dtype=bool)
+        mask[on_idx] = True
+        on = d[:, mask] if axis == 1 else d[mask, :]
+        off = d[:, ~mask] if axis == 1 else d[~mask, :]
+        if off.size and off.mean() > 1e-8:
+            out.append(float(on.mean() / off.mean()))
+    return max(out) if out else float('nan')
+
+
+def video_stats(pix, split):
+    """pix: uint8 [N, H, W, 3] numpy. split: index where generated frames begin."""
+    g = pix.astype(np.float32).mean(axis=3) / 255.0        # [N, H, W] luma
+    sharp = np.array([_lap_var(g[i]) for i in range(g.shape[0])])
+    contrast = g.reshape(g.shape[0], -1).std(axis=1)
+    delta = np.abs(np.diff(g, axis=0)).mean(axis=(1, 2))   # length N-1
+
+    def seg(a, lo, hi):
+        s = a[lo:hi]
+        return float(s.mean()) if len(s) else float('nan')
+
+    n = g.shape[0]
+    out = {
+        'sharpness_world': seg(sharp, 0, split),
+        'sharpness_gen': seg(sharp, split, n),
+        'contrast_world': seg(contrast, 0, split),
+        'contrast_gen': seg(contrast, split, n),
+        # deltas are between frames i and i+1, so the generated segment's own
+        # deltas start at `split` (the seam frame is excluded as it spans both)
+        'interframe_world': seg(delta, 0, max(0, split - 1)),
+        'interframe_gen': seg(delta, split, n - 1),
+    }
+    # blockiness on the last 8 generated frames (artifacts accumulate, so the tail
+    # is where they show) vs the world's own baseline from the same VAE
+    for per in (8, 16):
+        wb = np.mean([_blockiness(g[i], per) for i in range(min(split, 8))])
+        gb = np.mean([_blockiness(g[i], per) for i in range(max(split, n - 8), n)])
+        out[f'block{per}_world'] = float(wb)
+        out[f'block{per}_gen'] = float(gb)
+    out['blockiness'] = float(max(out['block8_gen'] / max(out['block8_world'], 1e-6),
+                                 out['block16_gen'] / max(out['block16_world'], 1e-6)))
+
+    for key in ('sharpness', 'contrast', 'interframe'):
+        w, gg = out[f'{key}_world'], out[f'{key}_gen']
+        out[f'{key}_ratio'] = float(gg / w) if w else float('nan')
+    # trend over the generated segment: first vs last third, to expose decay
+    gs = sharp[split:]
+    if len(gs) >= 6:
+        third = len(gs) // 3
+        out['sharpness_first_third'] = float(gs[:third].mean())
+        out['sharpness_last_third'] = float(gs[-third:].mean())
+        out['sharpness_decay'] = float(gs[-third:].mean() / gs[:third].mean())
+    return out
+
+
+def channel_drift(world_lat, gen_lat, tail_frac=0.34):
+    """Per-channel first/second-moment departure from the world, in latent space.
+
+    Every no-reference statistic in this file, and the reference-based Frechet
+    one in scripts/fwd_score.py, failed to rank the colour/saturation artifact
+    (FINDINGS.md §5a: the Frechet metric put the most saturated run first and
+    the only clean run fifth). They fail for the same reason -- each pools away
+    the axis the artifact actually lives on.
+
+    That axis is not hidden. `FewStepStreamer._renorm`, the inference patch that
+    hides the artifact, works by matching the generated block's PER-CHANNEL mean
+    and std to the world's. So the artifact is, by construction, a per-channel
+    moment departure, and measuring it needs no network at all:
+
+        mu_shift[c] = |mean(gen[c]) - mean(world[c])| / std(world[c])
+        sd_ratio[c] =  std(gen[c])  / std(world[c])
+
+    Reported as the WORST channel, not the mean: a cast in one or two of the 16
+    latent channels is exactly what a magenta face is, and averaging over
+    channels would dilute it back into invisibility. Measured over the tail of
+    the rollout, where drift has accumulated.
+
+    Note this is only meaningful with the latent-norm patch OFF; with it on it
+    is being directly optimised, and reads ~0 by construction.
+    """
+    w = world_lat.float()
+    g = gen_lat.float()
+    n_tail = max(1, int(g.shape[1] * tail_frac))
+    t = g[:, -n_tail:]
+    wm = w.mean(dim=(1, 2, 3))
+    ws = w.std(dim=(1, 2, 3)).clamp_min(1e-6)
+    mu_shift = ((t.mean(dim=(1, 2, 3)) - wm).abs() / ws)
+    sd_ratio = (t.std(dim=(1, 2, 3)) / ws)
+    return {
+        'chan_mu_shift_worst': float(mu_shift.max()),
+        'chan_mu_shift_mean': float(mu_shift.mean()),
+        'chan_sd_ratio_worst': float(sd_ratio.max()),
+        'chan_sd_ratio_min': float(sd_ratio.min()),
+        'chan_worst_index': int(mu_shift.argmax()),
+    }
+
+
+def latent_stats(world_lat, gen_lat):
+    """world_lat/gen_lat: torch [C, F, H, W] float. Adds drift + collapse terms."""
+    w = world_lat.float()
+    g = gen_lat.float()
+    ref = w.mean(dim=1, keepdim=True).flatten()
+    ref = ref / (ref.norm() + 1e-8)
+    cos = []
+    for i in range(g.shape[1]):
+        f = g[:, i].flatten()
+        cos.append(float((f / (f.norm() + 1e-8) @ ref)))
+    per_frame_std = [float(g[:, i].std()) for i in range(g.shape[1])]
+    third = max(1, len(cos) // 3)
+    return {
+        'std_world': float(w.std()),
+        'std_gen': float(g.std()),
+        'std_ratio': float(g.std() / w.std()),
+        'std_first_frame': per_frame_std[0],
+        'std_last_frame': per_frame_std[-1],
+        'std_frame_decay': float(per_frame_std[-1] / per_frame_std[0]),
+        'world_cos_first_third': float(np.mean(cos[:third])),
+        'world_cos_last_third': float(np.mean(cos[-third:])),
+        'world_cos_drift': float(np.mean(cos[-third:]) - np.mean(cos[:third])),
+    }
+
+
+def summary_line(vs, ls):
+    return (f"sharp {vs['sharpness_ratio']:.2f}x  "
+            f"contrast {vs['contrast_ratio']:.2f}x  "
+            f"motion {vs['interframe_ratio']:.2f}x  "
+            f"blockiness {vs['blockiness']:.2f}x  "
+            f"latent-std {ls['std_ratio']:.2f}x  "
+            f"sharp-decay {vs.get('sharpness_decay', float('nan')):.2f}  "
+            f"world-cos-drift {ls['world_cos_drift']:+.3f}")

wanstreamer/pipeline.py

@@ -0,0 +1,305 @@
+"""Streaming generator: persistent world W + autoregressive event stream.
+
+Implements the modelling contract the two papers share, at 1.3B scale:
+
+    p(e_1..K | W, x_1..K) = prod_k p(e_k | W, x_<=k, e_<k)
+
+  * WORLD W -- the established scene/character. Primed ONCE as a single
+    bidirectional block, then its K/V is frozen in the cache (v0.2's "KV
+    construction").
+  * EVENT STREAM -- emitted one chunk at a time. A chunk attends bidirectionally
+    within itself and freely over the whole world plus all prior events, and to
+    nothing else.
+
+BLOCK-CAUSALITY IS STRUCTURAL. The cache only ever holds world + committed
+events, so full attention over [cache, chunk] already gives exactly "world
+bidirectional, events causal". No attention mask is required -- verified in
+tests/test_attention_fallback.py case 5, and a mask would wrongly serialise
+tokens *within* a chunk.
+
+CHUNK SIZE. final.pt was fine-tuned with --num-frames 4, so it only ever saw
+4-latent-frame clips. Generating one latent frame in isolation is out of
+distribution; `chunk_frames=4` matches training. One latent frame = 4 pixel
+frames = 160 ms at 25 FPS (Wan VAE temporal stride 4), so a 4-frame chunk covers
+640 ms of video and must be produced in under 640 ms to sustain 25 FPS.
+
+Timesteps are fed to the model in the [0,1] convention final.pt was fine-tuned on
+(train_streaming.py:152 draws torch.rand()); see PROGRESS.md §5.
+"""
+import time
+
+import torch
+
+from wan.modules.model import sinusoidal_embedding_1d
+from wan.modules.attention import flash_attention
+from wan.utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
+                                  get_sampling_sigmas, retrieve_timesteps)
+
+from .core import (ModulationCache, make_rope_table, make_cache, latent_geometry,
+                   timestep_to_train_scale)
+from .rope import apply_rope
+
+
+class StreamingGenerator:
+    def __init__(self, model, ctx_emb, width=640, height=368, max_frames=64,
+                 device='cuda', dtype=torch.bfloat16, window_frames=None,
+                 time_scale=1.0, cache_frames=None):
+        """window_frames: if set, evict oldest frames so the cache holds at most
+        this many (bounded memory + bounded per-chunk cost). None = unbounded.
+
+        time_scale converts the internal flow fraction t in [0,1] into the value
+        the WEIGHTS expect at their time embedding:
+
+            final.pt        -> 1.0     (fine-tuned on torch.rand(), i.e. [0,1])
+            original Wan    -> 1000.0  (trained on scheduler timesteps [0,1000])
+
+        This must match the checkpoint. Getting it backwards is not a small
+        degradation, it is total: measured normalised velocity error 0.064 vs
+        0.732 for the original weights under the two conventions (PROGRESS.md §5).
+        Everything internal to this class speaks the [0,1] fraction.
+        """
+        self.model = model
+        self.device = device
+        self.dtype = dtype
+        self.time_scale = float(time_scale)
+        self.h_lat, self.w_lat, self.hp, self.wp = latent_geometry(width, height)
+        self.tokens_per_frame = self.hp * self.wp
+        self.max_frames = max_frames
+        self.window_frames = window_frames
+        # RoPE must span the whole stream (absolute temporal indices keep growing),
+        # but the K/V cache only ever holds world + window, so sizing it by
+        # max_frames wastes memory quadratically in stream length: 30 layers x
+        # 6 KB/token means an unbounded 173 frame buffer at 1560 tok/frame is
+        # ~48 GiB, versus ~6 GiB for a 21 frame working set.
+        self.rope = make_rope_table(model, self.hp, self.wp, max_frames, device)
+        self.cache = make_cache(model, self.tokens_per_frame,
+                                cache_frames or max_frames, device, dtype)
+        self.ctx = ctx_emb
+        self.ctx_lens = torch.tensor([ctx_emb.shape[1]], device=device, dtype=torch.long)
+        self.n_world = 0
+        self.n_frames = 0        # total committed latent frames
+        self.ref_stats = None    # per-channel (mean, std) of the world latents
+
+    def _time_embed(self, t_frac):
+        """t_frac is the flow fraction in [0,1]; time_scale maps it to the
+        checkpoint's own convention."""
+        tv = (torch.ones(1, device=self.device, dtype=torch.float32)
+              * float(t_frac) * self.time_scale)
+        with torch.amp.autocast('cuda', enabled=False):
+            e = self.model.time_embedding(
+                sinusoidal_embedding_1d(self.model.freq_dim, tv).float())
+            e0 = self.model.time_projection(e).unflatten(1, (6, self.model.dim))
+        return e, e0
+
+    @torch.no_grad()
+    def chunk_forward(self, z, t_frac, t_start, use_cache=True, ctx=None):
+        """Run a chunk of latent frames against the cached past.
+
+        z: [1, C, N, H_lat, W_lat]. t_frac is the flow fraction in [0,1].
+        Writes the chunk's K/V into the cache scratch region (uncommitted).
+        Returns predicted velocity [C, N, H_lat, W_lat].
+        Attention is unmasked over [cache, chunk] == block-causal.
+        """
+        model = self.model
+        ctx = self.ctx if ctx is None else ctx
+        N = z.shape[2]
+        S = N * self.tokens_per_frame
+        e, e0 = self._time_embed(t_frac)
+        mod = ModulationCache(model.blocks, e0)
+
+        with torch.amp.autocast('cuda', dtype=self.dtype):
+            x = model.patch_embedding(z.to(self.dtype))
+            grid = torch.stack([torch.tensor(x.shape[2:], dtype=torch.long,
+                                             device=self.device)])
+            x = x.flatten(2).transpose(1, 2)
+            tbl = self.rope.span(t_start, N)
+
+            for li, blk in enumerate(model.blocks):
+                ec = mod[li]
+                sa_in = blk.norm1(x).float() * (1 + ec[1]) + ec[0]
+                n = blk.num_heads
+                d = blk.dim // n
+                sa = blk.self_attn
+                q = sa.norm_q(sa.q(sa_in)).view(1, S, n, d)
+                k = sa.norm_k(sa.k(sa_in)).view(1, S, n, d)
+                v = sa.v(sa_in).view(1, S, n, d)
+                q = apply_rope(q, tbl).to(self.dtype)
+                k = apply_rope(k, tbl).to(self.dtype)
+                v = v.to(self.dtype)
+
+                if use_cache:
+                    self.cache.write(li, k, v)
+                    ck, cv = self.cache.context(li, S)
+                else:
+                    ck, cv = k, v
+                y = flash_attention(q=q, k=ck, v=cv, window_size=(-1, -1),
+                                    causal=False)
+                y = sa.o(y.flatten(2))
+                with torch.amp.autocast('cuda', dtype=torch.float32):
+                    x = x + y * ec[2]
+                x = x + blk.cross_attn(blk.norm3(x), ctx, self.ctx_lens)
+                yf = blk.ffn(blk.norm2(x).float() * (1 + ec[4]) + ec[3])
+                with torch.amp.autocast('cuda', dtype=torch.float32):
+                    x = x + yf * ec[5]
+
+            return model.unpatchify(model.head(x, e), grid)[0]
+
+    def _commit(self, n_frames):
+        self.cache.commit(n_frames * self.tokens_per_frame)
+        self.n_frames += n_frames
+        if self.window_frames is not None:
+            # Persistent world + sliding event window: the world block is pinned,
+            # so eviction drops the OLDEST EVENTS. Bounds both memory and the
+            # per chunk attention cost, which is what makes throughput flat in
+            # stream length rather than degrading.
+            protect = self.n_world * self.tokens_per_frame
+            budget = (self.n_world + self.window_frames) * self.tokens_per_frame
+            excess = self.cache.num_tokens - budget
+            if excess > 0:
+                self.cache.evict_front(excess, protect=protect)
+
+    @torch.no_grad()
+    def set_world(self, world_latents, t_frac=0.0):
+        """Prime the cache from clean world latents, attended bidirectionally."""
+        self.cache.reset()
+        self.n_frames = 0
+        z = world_latents.unsqueeze(0).to(self.device, self.dtype)
+        F = z.shape[2]
+        assert F <= self.max_frames
+        self.chunk_forward(z, t_frac, t_start=0, use_cache=True)
+        self._commit(F)
+        self.n_world = F
+        # Per channel moments of the world act as the in dist reference
+        # for latent_norm (see _renorm). Computed over (F, H, W) per channel.
+        w = world_latents.float()
+        self.ref_stats = (w.mean(dim=(1, 2, 3), keepdim=True).unsqueeze(0),
+                          w.std(dim=(1, 2, 3), keepdim=True).unsqueeze(0))
+        return F
+
+    def _renorm(self, z, strength):
+        """Pull a finished chunk's per-channel moments back to the world's.
+
+        Autoregressive rollout with a model that was never trained on its own
+        outputs accumulates error, and here that error is measurably a SCALE
+        divergence: with the original bidirectional weights, generated latent std
+        grows ~2x over 20 chunks and reaches 3x the world's, at which point the
+        VAE's input range is exceeded and frames decode to flat grey
+        (contrast 0.018 vs the world's 0.194 -- out_sweep/round1).
+
+        Matching the first two moments per channel is the cheapest correction that
+        targets exactly that failure, and it is applied BEFORE the chunk's K/V is
+        recomputed and committed, so the cached history stays in-distribution and
+        the correction cannot accumulate. strength blends: 0 = off, 1 = full
+        moment match. It cannot fix a wrong direction, only a wrong scale, so if
+        content quality is the problem this will not rescue it.
+        """
+        if strength <= 0 or self.ref_stats is None:
+            return z
+        mu, sd = self.ref_stats
+        zm = z.mean(dim=(2, 3, 4), keepdim=True)
+        zs = z.std(dim=(2, 3, 4), keepdim=True)
+        z_n = (z - zm) / (zs + 1e-5) * sd.to(z.dtype) + mu.to(z.dtype)
+        return (1.0 - strength) * z + strength * z_n
+
+
+    @torch.no_grad()
+    def generate_chunk(self, chunk_frames=4, num_steps=3, generator=None,
+                       sampler='dpm', shift=5.0, guidance=None, ctx_neg=None,
+                       rope_start=None, t_max=1.0, anchor=None, latent_norm=0.0):
+        """Emit one event chunk autoregressively. Returns (latents, seconds).
+
+        t_max < 1.0 warm-starts the chunk instead of denoising from pure noise:
+            z_init = (1 - t_max) * anchor + t_max * noise
+        where `anchor` is the last clean latent frame, broadcast over the chunk.
+        This is SDEdit-style continuation. It is worth doing here because
+        diag/context_len_probe.py measured this model's velocity error at 0.33 for
+        t=0.8 versus 0.14 for t=0.5 -- denoising from t=1 integrates through the
+        region where it is least accurate. Trade-off: lower t_max means less
+        motion diversity, since the chunk starts closer to the previous frame.
+        """
+        t_start = self.n_frames if rope_start is None else rope_start
+        if t_start + chunk_frames > self.max_frames:
+            raise RuntimeError(f'temporal index {t_start + chunk_frames} exceeds '
+                               f'max_frames {self.max_frames}')
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+
+        noise = torch.randn(1, 16, chunk_frames, self.h_lat, self.w_lat,
+                            device=self.device, dtype=torch.float32,
+                            generator=generator)
+        if t_max < 1.0:
+            if anchor is None:
+                raise ValueError('t_max < 1.0 requires an anchor latent frame')
+            a = anchor.to(self.device, torch.float32).unsqueeze(0)
+            if a.shape[2] != chunk_frames:
+                a = a[:, :, -1:].expand(-1, -1, chunk_frames, -1, -1)
+            z = (1 - t_max) * a + t_max * noise
+        else:
+            z = noise
+
+        def predict(z_in, t_frac):
+            v = self.chunk_forward(z_in, t_frac, t_start).float().unsqueeze(0)
+            if guidance and ctx_neg is not None:
+                vu = self.chunk_forward(z_in, t_frac, t_start,
+                                        ctx=ctx_neg).float().unsqueeze(0)
+                v = vu + guidance * (v - vu)
+            return v
+
+        if sampler == 'euler':
+            ts = torch.linspace(t_max, 0.0, num_steps + 1)
+            for i in range(num_steps):
+                tc, tn = float(ts[i]), float(ts[i + 1])
+                z = z + (tn - tc) * predict(z, tc)
+        elif sampler == 'dpm':
+            sch = FlowDPMSolverMultistepScheduler(
+                num_train_timesteps=1000, shift=1, use_dynamic_shifting=False)
+            sig = get_sampling_sigmas(num_steps, shift)
+            if t_max < 1.0:                     # rescale sigmas into [0, t_max]
+                sig = sig * t_max               # keep it a numpy array
+            tsteps, _ = retrieve_timesteps(sch, device=self.device, sigmas=sig)
+            for tv in tsteps:
+                v = predict(z, timestep_to_train_scale(tv))
+                z = sch.step(v, tv, z, return_dict=False)[0]
+        else:
+            raise ValueError(sampler)
+
+        z = self._renorm(z, latent_norm)
+
+        # recompute the finished chunk's K/V at t=0 so history is clean, then commit
+        self.chunk_forward(z, 0.0, t_start)
+        self._commit(chunk_frames)
+
+        torch.cuda.synchronize()
+        return z[0], time.perf_counter() - t0
+
+    @torch.no_grad()
+    def stream(self, num_chunks, chunk_frames=4, num_steps=3, seed=0,
+               log_every=5, anchor=None, world_anchor=None, world_anchor_weight=0.0,
+               **kw):
+        """world_anchor_weight w > 0 blends a fixed WORLD reference frame into the
+        warm start anchor:  anchor = (1-w)·last_generated + w·world_ref.
+
+        Rationale: pinning the worlds K/V stabilises attention but the warm-start
+        anchor still forms a chain (each chunk starts from the previous chunks
+        output), so errors compound and long streams drift. v0.3 casts W as the
+        persistent carrier of scene and character identity, so letting W enter the
+        initialisation too is the faithful reading, not just a hack.
+        """
+        g = torch.Generator(device=self.device).manual_seed(seed)
+        lats, times = [], []
+        for u in range(num_chunks):
+            lat, dt = self.generate_chunk(chunk_frames=chunk_frames,
+                                          num_steps=num_steps, generator=g,
+                                          anchor=anchor, **kw)
+            anchor = lat[:, -1:]        # newest clean frame anchors the next chunk
+            if world_anchor_weight > 0 and world_anchor is not None:
+                w = world_anchor_weight
+                anchor = (1 - w) * anchor + w * world_anchor.to(anchor.device,
+                                                                anchor.dtype)
+            lats.append(lat)
+            times.append(dt)
+            if log_every and (u + 1) % log_every == 0:
+                print(f'  chunk {u+1}/{num_chunks}: {dt*1000:7.1f} ms '
+                      f'({dt/chunk_frames*1000:6.1f} ms/latent-frame, '
+                      f'cache {self.cache.num_tokens} tok)')
+        return torch.cat(lats, dim=1), times

wanstreamer/prompts.py

@@ -0,0 +1,150 @@
+"""Prompt bank for teacher data generation.
+
+Composition is deliberate rather than arbitrary. The deployment story is a
+persistent world W that a stream of events continues, so the bank is weighted
+towards content with a stable subject and continuous, non-cut motion -- that is
+the distribution the student has to stay in over a long rollout. Hard cuts,
+crowds and rapid scene changes are exactly what a block-causal model with a
+sliding K/V window cannot represent, so they are kept out.
+
+Roughly half is human subjects (the papers' interaction use case), the rest
+animals, nature and urban scenes so the model does not collapse onto one mode.
+"""
+
+PEOPLE = [
+    "A woman with long dark hair speaking directly to the camera, natural facial expressions, soft studio lighting, shallow depth of field",
+    "A man in a grey sweater talking to the camera and gesturing calmly, warm indoor lighting, blurred bookshelf background",
+    "A young woman smiling and nodding while listening, soft window light from the left, neutral background",
+    "An older man with a short white beard telling a story to the camera, warm lamp light, cozy living room",
+    "A woman in a yellow raincoat looking around, light rain falling, overcast daylight, city street behind her",
+    "A chef in a white coat carefully plating a dish, overhead kitchen lighting, steam rising",
+    "A barista steaming milk behind a counter, morning light through a large window, cafe interior",
+    "A violinist playing on a small stage, warm spotlight, dark background",
+    "A woman reading a book by a window, turning a page, soft afternoon light, dust motes in the air",
+    "A man drinking coffee from a mug and looking out of a rainy window, muted daylight",
+    "A painter working on a canvas, brush strokes visible, studio skylight, paint-spattered apron",
+    "A woman practicing yoga on a mat, slow controlled movement, morning light on a wooden floor",
+    "A doctor in a white coat speaking to the camera in a bright clinic, professional lighting",
+    "A teacher writing on a whiteboard and turning to speak, classroom, fluorescent light",
+    "A man repairing a bicycle wheel in a garage, focused expression, work lamp lighting",
+    "A woman arranging flowers in a vase, bright kitchen, sunlight on the counter",
+    "A street musician playing guitar on a sidewalk, passersby blurred behind, golden hour",
+    "A potter shaping clay on a spinning wheel, hands wet with slip, warm workshop light",
+    "A woman in a wool coat walking slowly through a park in autumn, leaves falling",
+    "A man in a suit adjusting his tie in front of a mirror, hotel room, warm lamps",
+    "A child blowing bubbles in a garden, sunlight through the bubbles, summer afternoon",
+    "A woman laughing while talking on the phone, cafe interior, bokeh lights behind",
+    "A dancer moving slowly in an empty studio, mirrors and barre, cool daylight",
+    "A fisherman mending a net on a dock, weathered hands, overcast coastal light",
+    "A woman scientist looking into a microscope and then up at the camera, laboratory lighting",
+    "A man playing chess alone, moving a piece, window light across the board",
+    "A woman with curly hair singing into a studio microphone, headphones on, dim red light",
+    "An elderly woman knitting in an armchair, fireplace glow, quiet living room",
+    "A carpenter sanding a wooden plank, sawdust in the air, workshop window light",
+    "A woman tying her running shoes on a park bench, early morning, mist in the background",
+    "A man walking a dog along a canal, evening light, calm water",
+    "A woman applying makeup in front of a lit mirror, dressing room, soft bulbs",
+    "A librarian pulling a book from a high shelf, tall wooden stacks, warm reading lamps",
+    "A woman in a lab coat writing notes on a clipboard, bright modern interior",
+    "A man kneading dough on a floured counter, bakery, morning sun through a window",
+    "A woman in a denim jacket leaning on a railing overlooking a city, wind in her hair",
+    "A surfer in a wetsuit walking up a beach carrying a board, late afternoon sun",
+    "A woman gardening, planting a seedling into dark soil, bright overcast daylight",
+    "A man playing piano in a dim room, hands on the keys, single warm lamp",
+    "A woman photographer adjusting a camera lens, city rooftop, dusk",
+]
+
+ANIMALS = [
+    "A ginger cat sitting on a windowsill, tail flicking, sunlight across its fur",
+    "A golden retriever running across a grassy field toward the camera, sunny day",
+    "A horse grazing in a misty meadow at dawn, breath visible in the cold air",
+    "A red fox stepping carefully through fresh snow in a pine forest",
+    "An owl slowly turning its head on a branch, dappled forest light",
+    "A pod of dolphins swimming just under a clear turquoise surface",
+    "A hummingbird hovering at a red flower, wings blurred, bright garden",
+    "A sheepdog herding sheep across a green hillside, overcast light",
+    "A tabby kitten batting at a dangling string, wooden floor, warm indoor light",
+    "A deer lifting its head alertly in a clearing, early morning fog",
+    "A parrot with green and red feathers preening on a branch, tropical foliage",
+    "A school of small silver fish turning together over a coral reef, shafts of sunlight",
+    "A bumblebee moving across lavender flowers, summer sunlight, shallow focus",
+    "An elephant slowly walking across a dry savannah, dust rising, late afternoon",
+    "A black cat stretching on a sunlit rug, slow movement, quiet room",
+    "A hawk perched on a fence post scanning the field, wind in the grass",
+]
+
+NATURE = [
+    "Waves rolling onto a rocky shore at sunset, spray catching the light",
+    "A mountain stream running over smooth stones, sunlight through overhanging leaves",
+    "Tall grass moving in the wind on a hillside, clouds drifting overhead",
+    "Snow falling slowly through a dense pine forest, quiet grey light",
+    "A waterfall in a green canyon, mist rising, shafts of sunlight",
+    "Autumn leaves drifting down onto a still pond, reflections rippling",
+    "A field of sunflowers swaying, bright blue sky with scattered clouds",
+    "Clouds moving over a desert mesa, long shadows, late afternoon",
+    "A lavender field in bloom, gentle wind, hazy summer sun",
+    "Fog rolling slowly through a valley of dark evergreens at dawn",
+    "Rain falling on the surface of a lake, concentric ripples, grey daylight",
+    "A campfire burning at dusk, sparks rising, forest silhouettes behind",
+    "Northern lights shifting slowly above a snowy ridge, deep blue night",
+    "A wheat field rippling in the wind under a dramatic evening sky",
+    "Palm fronds moving in a warm breeze against a bright tropical sky",
+    "Ice floating slowly down a wide grey river, bare winter trees on the bank",
+]
+
+URBAN = [
+    "A quiet city street at night, wet asphalt reflecting neon signs, light rain",
+    "Steam rising from a manhole on a cold morning, cars passing slowly",
+    "A cafe terrace in the afternoon, awning moving slightly, people seated",
+    "A subway train arriving at a platform, motion blur, fluorescent light",
+    "Traffic crossing a bridge at dusk, headlights and taillights, city skyline behind",
+    "A narrow European alley with laundry hanging, warm afternoon light",
+    "Rain on a bus window, blurred city lights beyond the glass",
+    "A market stall with fresh produce, vendor arranging fruit, bright daylight",
+    "A quiet bookshop interior, dust in a shaft of sunlight, shelves of books",
+    "A rooftop at golden hour overlooking a dense city, antennas and water tanks",
+    "An empty basketball court at dusk, net moving slightly in the wind",
+    "A ferry crossing a harbour, gulls following, overcast light",
+    "A neon-lit ramen shop at night seen from the street, steam in the doorway",
+    "An old tram turning a corner on cobblestones, autumn trees along the street",
+]
+
+OBJECTS = [
+    "Coffee being poured into a glass cup, swirling crema, morning kitchen light",
+    "A candle flame flickering in a dark room, wax slowly melting",
+    "Ink diffusing into a glass of clear water, soft studio lighting",
+    "A record spinning on a turntable, needle in the groove, warm lamp light",
+    "Dough rising in a bowl in time lapse, kitchen window light",
+    "Water droplets running down a cold glass bottle, dark background",
+    "A mechanical watch movement ticking, extreme close-up, jeweller's lighting",
+    "Paint colours swirling together on a wet canvas, overhead studio light",
+    "Steam curling from a bowl of soup on a wooden table, warm evening light",
+    "A kite flying against a blue sky, string taut, wind moving the tail",
+]
+
+ALL = PEOPLE + ANIMALS + NATURE + URBAN + OBJECTS
+
+
+from .prompts_ext import EXT_ALL, REPLACEMENTS, assert_no_minors  # noqa: E402
+
+
+def prompt_bank():
+    bank = list(ALL)
+    for i, text in REPLACEMENTS.items():
+        bank[i] = text
+    bank += EXT_ALL
+    assert_no_minors(bank)
+    return bank
+
+
+def categories():
+    from . import prompts_ext as _x
+    return {'people': PEOPLE + _x.PEOPLE_EXT, 'animals': ANIMALS + _x.ANIMALS_EXT,
+            'nature': NATURE + _x.NATURE_EXT, 'urban': URBAN + _x.URBAN_EXT,
+            'objects': OBJECTS + _x.OBJECTS_EXT}
+
+
+if __name__ == '__main__':
+    for name, lst in categories().items():
+        print(f'{name:8s} {len(lst)}')
+    print(f'{"TOTAL":8s} {len(prompt_bank())}  (v5 bank {len(ALL)} at indices 0-95)')

wanstreamer/prompts_ext.py

@@ -0,0 +1,1051 @@
+"""Prompt-bank extension: the diversity fix FINDINGS §7 item 0b calls a blocker.
+
+96 prompts / 432 clips is what v5 trained on, and it is the stated reason the
+rolling-K/V work improved world_p0 and failed to generalise to the other three
+worlds. At effective batch 64 the old bank is nine iterations per epoch. This
+module adds 904 prompts to reach 1000.
+
+TWO RULES THIS FILE FOLLOWS, AND WHY.
+
+1. **It never touches indices 0-95.** `prompts.py` keeps its original bank
+   verbatim and `prompt_bank()` appends this module afterwards, so prompt 0 is
+   still the talking head, 44 the owl, 60 the waterfall and 82 the court. Those
+   four are `out/world_p{0,44,60,82}.pt` and every 1:1 comparison in FINDINGS;
+   renumbering them would silently invalidate the entire before/after evaluation
+   and the teacher clips already generated for them.
+
+2. **Same composition discipline as the original.** Stable subject, continuous
+   non-cut motion, one scene per prompt. No hard cuts, no crowds, no rapid scene
+   changes -- a block-causal model with a sliding K/V window cannot represent
+   them, so training on them teaches nothing it can keep. Roughly half remains
+   human subjects for the interaction use case.
+
+What is deliberately WIDER than the original bank, since narrowness was the
+defect being fixed: many more trades, activities and interiors; the same kind of
+subject under different times of day, weather and light; more variation in
+framing distance, in how much of the frame the subject occupies, and in how fast
+the motion is. Content type, lighting and setting vary as independently as a
+hand-written list can manage, so the student does not learn that (say) a
+workshop implies warm light.
+"""
+
+# --------------------------------------------------------------------- people
+# Direct-to-camera and conversational. The papers' interaction case, and the
+# thing mid-stream text switching will eventually have to stay inside, so it is
+# the largest single block.
+TALKING = [
+    "A woman in a green blouse speaking to the camera and pausing to think, soft key light, plain grey backdrop",
+    "A man with round glasses explaining something to the camera, hands moving in front of him, warm office lighting",
+    "A young man in a hooded sweatshirt talking casually to the camera, bedroom behind him, string lights out of focus",
+    "A woman with short silver hair speaking calmly from an armchair, window light from the right, quiet room",
+    "A man with a full beard laughing mid-sentence while talking to the camera, kitchen background, morning light",
+    "A woman in a blazer presenting to the camera with confident posture, softbox lighting, dark backdrop",
+    "A man in a plaid shirt leaning forward as he speaks, workshop behind him, overhead work lights",
+    "A woman with braided hair nodding and answering, cafe table, afternoon light through a window",
+    "A man in a turtleneck speaking quietly to the camera, dim room, a single lamp to one side",
+    "A woman in hospital scrubs talking to the camera, corridor blurred behind her, cool ceiling light",
+    "A man in a denim jacket gesturing while telling a story, brick wall behind, late afternoon sun",
+    "A woman with freckles smiling and speaking, close framing, natural window light, white wall",
+    "An older man in a cardigan speaking slowly and thoughtfully, bookshelves behind, warm reading lamp",
+    "A woman in a striped shirt raising her eyebrows as she makes a point, soft even studio light",
+    "A man with a shaved head speaking directly to camera, arms folded, grey seamless background",
+    "A woman in a knitted jumper cupping a mug while talking, sofa, soft lamp light, evening",
+    "A young woman with glasses reading aloud and glancing up at the camera, desk lamp, dark room",
+    "A man in a chef's jacket describing something to the camera, stainless kitchen behind, bright light",
+    "A woman in a high-collared coat speaking outdoors, breath faintly visible, overcast winter light",
+    "A man seated at a desk turning from a monitor to speak to the camera, blue screen glow on his face",
+    "A woman with long red hair tilting her head while listening, soft rim light, dark background",
+    "A man in a linen shirt speaking on a balcony, sea out of focus behind him, bright midday sun",
+    "A woman in a lab coat explaining a result to the camera, whiteboard behind, fluorescent light",
+    "An older woman with reading glasses on a chain speaking warmly, conservatory, filtered daylight",
+    "A man in a mechanic's overall wiping his hands while talking, garage, harsh overhead lamp",
+    "A woman in a silk scarf speaking with animated expressions, hotel lobby, warm ambient light",
+    "A young man with curly hair pausing and starting again mid-sentence, plain wall, flat daylight",
+    "A woman in a puffer jacket talking to the camera in falling snow, grey afternoon light",
+    "A man in a suit jacket without a tie speaking calmly, glass office, city blurred behind",
+    "A woman with a nose ring smiling and speaking, record shop behind her, warm tungsten light",
+    "A man in a fleece speaking from inside a tent doorway, dawn light on the fabric",
+    "A woman in a black polo neck speaking to camera in profile then turning to face it, dramatic side light",
+    "A man in a rugby shirt talking while catching his breath, park behind him, bright overcast sky",
+    "A woman in a tailored coat speaking on a station platform, train lights behind, evening",
+    "An older man with a walking stick seated on a bench, speaking to the camera, dappled park light",
+    "A woman in a chef's apron tasting from a spoon then talking to the camera, warm kitchen light",
+    "A man with a moustache speaking while leaning against a doorframe, hallway light behind him",
+    "A woman in a summer dress speaking in a garden, wind moving her hair, hazy afternoon sun",
+    "A young woman in a sports top speaking after exercise, gym behind her, cool overhead light",
+    "A man in a wool hat speaking outdoors at dusk, streetlight coming on behind him",
+    "A woman in a paint-stained shirt speaking in a studio, canvases behind, north-facing skylight",
+    "A man in a collarless shirt speaking with his hands clasped, softly lit interior, shallow focus",
+    "A woman with a pixie cut speaking energetically, colourful mural behind her, bright daylight",
+    "A man in a heavy coat speaking on a bridge, river and grey sky behind him",
+    "A woman seated cross-legged on a floor cushion talking to the camera, low warm light",
+    "A man in a bow tie speaking with precise gestures, wood-panelled room, warm lamps",
+    "A woman in a raincoat speaking under an umbrella, rain visible, muted grey daylight",
+    "A young man in a football shirt speaking excitedly, garden fence behind him, evening sun",
+    "A woman in a velvet jacket speaking under stage lighting, dark auditorium behind",
+    "A man with grey stubble speaking quietly by a window, rain running down the glass",
+    "A woman in a cycling jersey speaking while holding a helmet, roadside, bright morning",
+    "A man in a fisherman's jumper speaking on a harbour wall, boats out of focus behind",
+    "A woman in a graduation gown speaking to the camera and smiling, stone building behind",
+    "A man in a dressing gown speaking with a mug in hand, morning kitchen, low sun",
+    "A woman in a fitted blazer speaking in front of a bookcase, warm practical lamps",
+    "A man in a hi-vis vest speaking on a building site, cranes out of focus, hard daylight",
+    "A woman with silver earrings speaking closely to the camera, deep shadow on one side",
+    "A man in a cricket jumper speaking on a lawn, long shadows, late golden light",
+    "A woman in a hijab speaking calmly to the camera, soft window light, neutral interior",
+    "A man in a beanie speaking in a stairwell, cool fluorescent light overhead",
+]
+
+# Trades and craft. Continuous, repetitive, hand-scale motion with a fixed
+# subject -- close to ideal for a sliding-window model, and almost absent from
+# the original bank.
+TRADES = [
+    "A blacksmith hammering glowing metal on an anvil, sparks scattering, dark forge",
+    "A glassblower turning a gather of molten glass, furnace glow on their face",
+    "A cobbler stitching the welt of a leather shoe, work lamp, cluttered bench",
+    "A watchmaker fitting a tiny gear with tweezers, loupe in one eye, bright task light",
+    "A tailor pinning fabric on a dress form, chalk marks, daylight through a tall window",
+    "A weaver passing a shuttle through a loom, coloured threads, warm workshop light",
+    "A luthier planing the top of a guitar body, curls of wood falling, window light",
+    "A bookbinder pressing a spine in a clamp, glue brush at hand, soft overhead light",
+    "A calligrapher drawing a long stroke with a broad nib, ink pooling, desk lamp",
+    "A printmaker rolling ink across a plate, hand press behind, bright studio",
+    "A stonemason chiselling a block of limestone, dust in the air, open yard light",
+    "A cooper hammering a hoop onto a barrel, wood shavings underfoot, dim workshop",
+    "A farrier shaping a horseshoe over a small anvil, stable doorway light behind",
+    "A sign painter drawing a gold letter with a fine brush, glass shopfront, street light behind",
+    "A neon bender heating a glass tube over a flame, dark workshop, orange glow",
+    "A ceramicist trimming a leather-hard bowl on a wheel, clay ribbons curling away",
+    "A jeweller soldering a ring with a small torch, bright pinpoint flame, dark bench",
+    "A leatherworker burnishing an edge with a wooden tool, warm bench lamp",
+    "A knife sharpener drawing a blade across a wet stone, steady rhythm, window light",
+    "An upholsterer stretching fabric over a chair frame and tacking it down, workshop light",
+    "A cabinetmaker cutting a dovetail with a fine saw, sawdust, clear north light",
+    "A basket weaver bending willow into a rim, hands working steadily, shaded porch",
+    "A candlemaker dipping wicks into wax, slow repeated motion, dim warm room",
+    "A cheesemaker turning a wheel on a wooden shelf, cool cellar light",
+    "A brewer stirring a mash tun with a long paddle, steam rising, industrial lighting",
+    "A cooper's apprentice sanding a stave, focused, dust catching a shaft of sunlight",
+    "A tattoo artist working on a forearm, gloved hands, bright directional lamp",
+    "A barber trimming a beard with clippers, mirror behind, warm barbershop light",
+    "A florist wiring a stem into an arrangement, cool shop light, buckets of flowers",
+    "A milliner shaping felt over a wooden hat block with steam, dim workroom",
+    "A glazier scoring a sheet of glass and snapping it cleanly, workshop daylight",
+    "A locksmith cutting a key on a machine, sparks and metal dust, harsh work light",
+    "A saddler stitching leather with two needles, saddle clamped upright, warm light",
+    "A restorer cleaning varnish from an oil painting with a swab, raking light",
+    "A model maker gluing a tiny part with tweezers, magnifier lamp, dark room",
+    "A clockmaker winding a movement and listening to it, quiet workshop, side light",
+    "A stained-glass artist leading a panel together, coloured light falling on the bench",
+    "A woodturner shaping a bowl on a lathe, shavings flying, focused work light",
+    "A shoemaker hammering a heel into place, last held between the knees, dim shop",
+    "An engraver cutting a line into a copper plate with a burin, bright angled lamp",
+    "A bell founder scraping a mould surface smooth, dusty foundry, high windows",
+    "A thatcher combing reed into place on a roof, bright overcast sky behind",
+    "A sailmaker feeding heavy canvas through an industrial machine, loft daylight",
+    "A wheelwright fitting a spoke into a wooden hub, open workshop door, bright outside",
+    "A gunsmith checkering a walnut stock with a fine file, bench lamp, dark room",
+    "A potter centring a lump of clay, wet hands, slip running, warm studio",
+    "A silversmith planishing a bowl with a small hammer, rhythmic ringing, dim bench",
+    "A bowyer scraping the belly of a longbow, long even strokes, window light",
+    "A glass engraver working a wheel against a goblet, water dripping, bright lamp",
+    "A tanner working oil into a hide with a cloth, dim tannery, high dusty light",
+    "A parchment maker scraping a stretched skin with a curved blade, cool daylight",
+    "A mosaic artist setting tesserae into wet mortar, bright overhead studio light",
+    "A gilder laying gold leaf onto a frame with a brush, still air, warm lamp",
+    "A blacksmith quenching a hot blade in oil, steam and flame flaring briefly",
+    "A carpenter marking a mortise with a gauge, pencil behind the ear, soft daylight",
+    "A pipe organ builder voicing a metal pipe, quiet church interior, tall windows",
+    "A rope maker twisting strands on a walk, long perspective, flat afternoon light",
+    "A knife maker grinding a bevel on a belt sander, sparks streaming, dark shop",
+    "A puppet maker articulating a wooden joint, tools on the bench, warm lamp",
+    "A book conservator flattening a page under weights, gloves, cool even light",
+    "A brass polisher buffing a lamp on a wheel, cloth blur, dim workshop",
+    "A dyer lifting fabric from an indigo vat, deep blue dripping, shaded courtyard",
+    "A furniture restorer rubbing wax into a tabletop, circular motion, warm light",
+    "A lens grinder polishing an optic on a spinning pad, water film, bright bench",
+    "A hat blocker brushing a felt brim, steam kettle beside, dim workroom",
+    "A screen printer pulling a squeegee across a screen, bright even studio light",
+    "A cobbler's wheel spinning as a sole is trimmed, dust and leather smell implied",
+    "A bronze caster pouring molten metal into a mould, intense orange glow, dark foundry",
+    "A chair caner threading rush through a seat frame, patient hands, porch light",
+    "A stone carver dusting chips from a half-finished relief, raking sunlight",
+]
+
+# Food preparation. Hands, steam, liquids and repetitive motion at a fixed
+# camera -- and a very different texture and colour statistic from the studio
+# talking heads above.
+KITCHEN = [
+    "A cook folding an omelette in a pan, gas flame beneath, warm kitchen light",
+    "A baker scoring a loaf before it goes into the oven, flour on the crust, morning light",
+    "A chef julienning carrots quickly and evenly, steel bench, bright overhead light",
+    "A pastry chef piping cream onto a tart, steady hands, cool even light",
+    "A cook tossing vegetables in a wok over a high flame, steam and smoke rising",
+    "A barista pouring a rosetta into a flat white, close framing, cafe window light",
+    "A person kneading pasta dough and folding it over, wooden board, soft daylight",
+    "A chef spooning sauce around a plate, tweezers at hand, dark restaurant pass",
+    "A cook cracking eggs into a bowl one after another, marble counter, bright light",
+    "A person rolling sushi on a bamboo mat, precise movements, cool clean light",
+    "A baker dusting a tray of pastries with icing sugar, warm bakery light",
+    "A chef searing a steak, smoke rising, flames occasionally licking up",
+    "A person stirring risotto slowly in a wide pan, steam, evening kitchen light",
+    "A cook grating cheese over a bowl of pasta, warm domestic light",
+    "A pastry chef tempering chocolate on a marble slab, spatula sweeping, cool light",
+    "A person pressing tortillas and laying them on a hot griddle, steam rising",
+    "A chef deboning a fish with a thin blade, clean stainless surface, bright light",
+    "A cook ladling soup into a bowl, steam curling up, warm evening light",
+    "A person peeling apples in one long ribbon, kitchen table, window light",
+    "A baker shaping baguettes on a floured bench, rhythmic, early morning light",
+    "A cook flipping pancakes on a griddle, batter bubbling, bright kitchen",
+    "A person whisking egg whites in a copper bowl, steady circular motion, soft light",
+    "A chef torching the top of a creme brulee, flame passing over sugar, dark counter",
+    "A cook pulling noodles by hand, stretching and folding, flour dust, bright light",
+    "A person pouring batter into a hot waffle iron, steam escaping the seam",
+    "A chef arranging thin slices of fish on a plate, tweezers, focused task light",
+    "A cook stirring a large pot of stew with a wooden spoon, steam, warm light",
+    "A person icing a cake with a palette knife, turntable rotating slowly, soft light",
+    "A tea master pouring hot water over leaves in a small pot, steam, quiet dim room",
+    "A cook shucking oysters on ice, knife twisting, cold blue-tinted light",
+    "A person rolling dumplings and pleating the edges, board of flour, bright light",
+    "A chef straining pasta into a colander, steam billowing up, kitchen light",
+    "A person grinding spices in a mortar and pestle, steady rhythm, warm light",
+    "A cook basting a roast with a spoon, oven door open, warm interior glow",
+    "A bartender stirring a drink in a mixing glass, ice ringing, dim bar light",
+    "A person pouring olive oil in a thin stream over a salad, bright daylight",
+    "A cook flipping a crepe in a wide pan, quick wrist motion, warm kitchen",
+    "A person spooning yoghurt into a bowl and adding berries, morning window light",
+    "A chef plating a dessert with a quenelle of ice cream, cold plate, dark background",
+    "A person pouring coffee from a stovetop pot into a small cup, morning light",
+]
+
+# Music and performance. Continuous articulated motion, strong rhythmic
+# structure, and a wide range of lighting from dark stages to bright rehearsal
+# rooms.
+MUSIC = [
+    "A cellist drawing a long bow stroke, eyes closed, warm rehearsal room light",
+    "A drummer playing a steady groove, sticks blurring, dim practice space",
+    "A pianist's hands moving across the keys, close framing, single warm lamp",
+    "A saxophonist playing under a single spotlight, smoke in the beam, dark club",
+    "A singer holding a long note at a microphone, eyes shut, deep red stage light",
+    "A guitarist bending a string on an electric guitar, amp glowing behind, dark stage",
+    "A harpist plucking with both hands, gilded frame catching warm light",
+    "A flautist playing in a bright rehearsal room, music stand, tall windows",
+    "A double bass player walking a line, upright instrument, dim jazz club",
+    "A violinist tuning and then beginning to play, warm practice room light",
+    "A DJ moving a fader and nodding to the beat, coloured lights sweeping behind",
+    "A conductor raising both hands and cueing an entry, dark hall, focused light",
+    "An accordionist playing on a stone step, warm late afternoon sun",
+    "A banjo player picking quickly on a wooden porch, bright overcast light",
+    "A trumpeter emptying a valve and lifting the horn to play, dim stage",
+    "A choir singer holding a folder and singing, stained glass light behind",
+    "A tabla player's hands moving rapidly on the drums, warm floor lamp",
+    "A koto player plucking strings, seated on a mat, soft diffuse daylight",
+    "A church organist's feet moving on the pedalboard, dim gallery light",
+    "A busker playing a fiddle in an underpass, hard directional light, tiled walls",
+    "A harmonica player cupping the instrument, eyes closed, warm bar light",
+    "A bagpiper playing on a windy hillside, tartan moving, grey daylight",
+    "A marimba player's mallets moving across the bars, bright rehearsal light",
+    "A sitar player sitting cross-legged, hands on the frets, dim warm room",
+    "A synth player adjusting a knob while a pad sustains, blue and purple glow",
+    "A bass guitarist locked in with the drums, stage haze, moving lights",
+    "A ballet dancer rising onto pointe and holding it, mirrored studio, cool light",
+    "A tap dancer's shoes striking a wooden floor, sharp rhythm, warm stage light",
+    "A flamenco dancer turning with a fan, red dress, dark stage, hard side light",
+    "A contemporary dancer rolling slowly across a floor, bare studio, north light",
+    "A puppeteer working a marionette on a small stage, warm footlights",
+    "A magician turning a coin over their knuckles, close framing, dark background",
+    "A mime moving slowly against an invisible wall, plain backdrop, flat light",
+    "A juggler keeping three clubs moving, park behind, golden hour light",
+    "A circus aerialist turning slowly on a hoop, dark tent, single beam",
+    "An actor delivering a monologue on a bare stage, single overhead light",
+    "A stand-up comic gesturing at a microphone, brick wall behind, warm spot",
+    "A street performer painting themselves still as a statue, bright square daylight",
+    "A beatboxer close to a microphone, hand cupped, dark room, rim light",
+    "A ukulele player strumming on a beach at sunset, warm backlight",
+]
+
+# Sport and physical movement. Fast, whole-body motion with a stable frame --
+# the opposite end of the motion-magnitude axis from the talking heads.
+SPORT = [
+    "A swimmer doing front crawl in a pool lane, underwater view, shafts of light",
+    "A boxer working a speed bag, rhythmic blur, dusty gym light",
+    "A climber reaching for a hold on an indoor wall, chalk dust, bright colour",
+    "A runner on a track rounding a bend, morning light, long shadows",
+    "A rower pulling on an indoor machine, steady breathing, cool gym light",
+    "A skateboarder rolling slowly across an empty park bowl, late afternoon sun",
+    "A tennis player serving, ball toss at the top of frame, bright daylight",
+    "A footballer juggling a ball on their instep, grass pitch, evening light",
+    "A weightlifter setting up over a barbell and taking a breath, harsh gym light",
+    "A gymnast holding a handstand on a beam, quiet hall, high windows",
+    "A yoga teacher moving slowly into a low lunge, warm studio, wooden floor",
+    "A cyclist climbing out of the saddle on a hill road, low sun behind",
+    "A kayaker paddling across flat water at dawn, mist on the surface",
+    "A surfer paddling out through a small wave, bright morning glare",
+    "A skier carving through soft snow, spray rising, blue shadows",
+    "A snowboarder traversing slowly across a slope, overcast flat light",
+    "A fencer lunging and recovering, white kit, cool hall lighting",
+    "A judo practitioner gripping a partner's jacket and shifting weight, dojo light",
+    "An archer drawing a bow and holding the anchor, field behind, soft daylight",
+    "A basketball player shooting free throws alone, empty gym, high windows",
+    "A dancer skipping rope quickly, dust on the floor, single hanging bulb",
+    "A martial artist practising a slow form in a park at sunrise, mist",
+    "A horse and rider cantering along a beach, spray from the hooves, low sun",
+    "A rock climber on a sea cliff reaching up, sea moving far below, bright light",
+    "A hiker stepping up onto a rock and pausing, wide valley behind, hazy light",
+    "A pole vaulter running in with the pole, stadium, bright floodlights",
+    "A table tennis player rallying, fast bat movement, flat indoor light",
+    "A golfer swinging through and holding the finish, fairway, low golden light",
+    "A rugby player passing along the line, wet grass, floodlit evening",
+    "A goalkeeper diving to save a low shot, spray of turf, bright daylight",
+    "A sprinter setting into blocks and breathing out, stadium, warm evening light",
+    "A cricketer bowling in from the crease, dusty pitch, hard midday sun",
+    "A diver walking to the end of a board and pausing, still pool below",
+    "A parkour athlete stepping along a low wall, quiet street, morning light",
+    "A wheelchair racer pushing along a track, hands blurring, bright daylight",
+    "A person on a climbing treadmill, steady hand movement, dim gym",
+    "An ice skater gliding across an outdoor rink, string lights above, dusk",
+    "A speed skater in a low crouch on a long oval, cool indoor light",
+    "A trail runner descending a rocky path, dust rising, dappled forest light",
+    "A stand-up paddleboarder crossing a still lake at dawn, reflections",
+]
+
+# Work, study and screens. Small motion against a fixed subject, and the screen
+# glow gives a light source the rest of the bank does not have.
+WORK = [
+    "A programmer typing quickly, code reflected in their glasses, dark room",
+    "An architect drawing on a large sheet with a scale rule, drafting lamp",
+    "An editor scrubbing through footage on a timeline, dual monitors, dim suite",
+    "An accountant turning pages of a ledger and marking a figure, desk lamp",
+    "A translator speaking into a headset in a booth, small light on the desk",
+    "A radio host adjusting a microphone arm and leaning in, dim studio, red light",
+    "An air traffic controller watching a screen and speaking, dark room, green glow",
+    "A watch repairer's hands under a magnifier, bright ring light",
+    "A designer sketching on a tablet with a stylus, bright screen, dark surround",
+    "A student highlighting a textbook at a library desk, green banker's lamp",
+    "A teacher marking papers at a kitchen table late at night, single lamp",
+    "A cartographer tracing a coastline on a large map, bright even light",
+    "A stenographer's hands moving on a small keyboard, courtroom, flat light",
+    "A tailor's assistant measuring a sleeve and noting it down, shop daylight",
+    "A curator adjusting the angle of a small object in a display case, gallery light",
+    "A sound engineer moving faders on a large console, meters flickering, dim room",
+    "A photographer reviewing frames on a camera back, studio strobes behind",
+    "A pilot running through a checklist in a cockpit, instrument glow, dusk outside",
+    "A ship's officer plotting a course on a chart table, red night lighting",
+    "A dispatcher speaking into a headset, wall of screens behind, cool light",
+    "A journalist typing on a laptop in a hotel room, window light, city outside",
+    "A clerk stamping documents one after another, post office, flat overhead light",
+    "A biologist labelling sample tubes in a rack, cold laboratory light",
+    "A weather forecaster gesturing at a map, studio lighting, blue backdrop",
+    "A pharmacist counting tablets into a tray, bright counter light",
+    "A bank teller counting notes through a machine, quiet interior, even light",
+    "A barista's colleague writing the specials on a chalkboard, morning shop light",
+    "A tattoo shop receptionist flipping through an appointment book, warm lamp",
+    "A watch dealer examining a dial with a loupe, dark velvet tray, bright spot",
+    "An auctioneer taking bids and gesturing, warm room light, wood panelling",
+    "A locksmith at a van workbench cutting a key, roadside, overcast light",
+    "A gardener at a potting bench filling trays with compost, greenhouse light",
+    "A vet writing notes after an examination, clinic, cool ceiling light",
+    "A driving instructor pointing at the road ahead from the passenger seat, daylight",
+    "A projectionist threading film through a projector, dim booth, warm lamp",
+    "A bookseller shelving new stock and straightening spines, warm shop light",
+    "A museum conservator brushing dust from a fossil, bright angled lamp",
+    "A drone pilot watching a controller screen in a field, bright overcast sky",
+    "A seamstress feeding fabric through an industrial machine, task light",
+    "A call centre worker adjusting a headset and smiling while speaking, flat light",
+]
+
+# Outdoors and travel. Wide scenes with a person in them, moving daylight and
+# weather -- deliberately the hardest lighting in the bank.
+OUTDOOR = [
+    "A hiker adjusting the straps of a rucksack on a ridge, wind, bright hazy light",
+    "A woman standing at the edge of a cliff looking out to sea, wind in her coat",
+    "A man rowing a small wooden boat across a still lake at dawn",
+    "A traveller waiting on a rural platform with a suitcase, low winter sun",
+    "A cyclist stopping to drink from a bottle on a country lane, dappled shade",
+    "A shepherd walking behind a flock along a stone-walled track, grey light",
+    "A woman reading on a picnic blanket, grass moving in the wind, bright afternoon",
+    "A man skimming a stone across flat water, ripples spreading, evening light",
+    "A woman photographing a mountain range from a viewpoint, cold clear light",
+    "A man chopping firewood outside a cabin, breath visible, low winter sun",
+    "A woman pitching a tent in a meadow, fabric catching the wind, golden hour",
+    "A man fishing from a riverbank, line casting out, mist on the water",
+    "A woman walking a coastal path with a dog, sea and sky behind, bright wind",
+    "A man refilling a water bottle from a mountain stream, cold clear light",
+    "A woman feeding chickens in a farmyard, morning light through a gate",
+    "A man on a ladder picking apples into a canvas bag, orchard, soft autumn light",
+    "A woman digging in an allotment bed, spade turning soil, flat overcast light",
+    "A man loading hay bales onto a trailer, dust in the air, hard summer sun",
+    "A woman leading a horse out of a stable, warm doorway light behind",
+    "A man scraping ice from a windscreen, breath visible, blue pre-dawn light",
+    "A woman shovelling snow from a path, steady rhythm, grey winter light",
+    "A man sitting on a rock eating a sandwich, wide valley behind, hazy light",
+    "A woman crossing a wooden footbridge over a stream, dappled forest light",
+    "A man launching a kayak from a shingle beach, small waves, overcast",
+    "A woman in waders casting a fly line, river bend, low golden light",
+    "A man tending a beehive in a veil and gloves, smoke drifting, bright meadow",
+    "A woman sketching a landscape in a small notebook, wind moving the pages",
+    "A man walking along a railway embankment at dusk, long shadows",
+    "A woman sitting on a harbour wall watching boats, gulls overhead, bright grey",
+    "A man tying a boat to a mooring ring, water slapping the stone, soft light",
+    "A woman running along a beach at low tide, wet sand reflecting the sky",
+    "A man setting up a telescope in a field at dusk, first stars appearing",
+    "A woman collecting shells at the tideline, low sun behind her",
+    "A man walking through tall reeds at the edge of a marsh, warm evening light",
+    "A woman climbing a stile in a hedgerow, sheep field beyond, bright overcast",
+    "A man carrying a canoe on his shoulders down to the water, dappled shade",
+    "A woman stopping to look at a map at a trail junction, forest light",
+    "A man raking leaves into a pile in a garden, autumn light, low sun",
+    "A woman hanging washing on a line, sheets moving in the wind, bright day",
+    "A man watering a vegetable garden with a can, evening light, long shadows",
+]
+
+# Quiet domestic interiors. Small motion, warm practical light, and the
+# strongest test of whether identity and fine detail hold over a long rollout.
+DOMESTIC = [
+    "A woman folding laundry on a bed, slow steady motion, afternoon window light",
+    "A man ironing a shirt, steam rising, kitchen light behind him",
+    "A woman watering houseplants on a windowsill, low sun through the leaves",
+    "A man tuning a radio dial, warm lamp, evening living room",
+    "A woman brushing her hair in front of a window, soft morning light",
+    "A man lighting a fire in a wood stove, flames catching, dim room",
+    "A woman writing a letter at a desk, fountain pen, single lamp, night",
+    "A man polishing shoes on a newspaper, circular motion, kitchen light",
+    "A woman sorting through a box of old photographs, warm attic light",
+    "A man making a bed, sheet billowing and settling, bright bedroom",
+    "A woman doing a jigsaw puzzle at a table, pieces fitting, lamp light",
+    "A man winding a mantel clock, quiet room, afternoon light",
+    "A woman painting her nails at a small table, focused, window light",
+    "A man sharpening a pencil with a small knife, shavings falling, desk lamp",
+    "A woman kneading dough at a kitchen table, flour dust in the light",
+    "A man repotting a plant, soil spilling onto newspaper, bright kitchen",
+    "A woman lighting candles on a table one by one, room darkening around",
+    "A man sewing a button onto a coat, thread pulled through, lamp light",
+    "A woman stacking books onto a shelf, dust motes in a sunbeam",
+    "A man cleaning a camera lens with a cloth, careful movements, window light",
+    "A woman drinking tea and watching rain on the window, muted grey light",
+    "A man cutting flowers and putting them in a jar, bright kitchen counter",
+    "A woman rocking slowly in a chair with a blanket, firelight on her face",
+    "A man scrolling through records in a crate, warm shop-like room light",
+    "A woman plaiting her own hair in front of a mirror, morning bathroom light",
+    "A man wiping down a kitchen counter, steady sweeps, evening light",
+    "A woman putting on earrings in front of a mirror, warm bulbs around the frame",
+    "A man setting a table, placing cutlery one piece at a time, low sun",
+    "A woman stretching on a mat in a living room, morning light on the floor",
+    "A man feeding a cat from a tin, cat winding around his legs, kitchen light",
+    "A woman opening curtains and light flooding into a room",
+    "A man carrying a mug through a hallway, warm light spilling from a doorway",
+    "A woman folding paper into a small shape, precise hands, desk lamp",
+    "A man restringing a guitar, winding the tuner, warm living room light",
+    "A woman shelling peas into a bowl, steady rhythm, kitchen window light",
+    "A man dusting a bookshelf with a cloth, sunbeam full of motes",
+    "A woman knitting while watching something off camera, lamp beside her",
+    "A man buttoning a coat by the front door, hallway light, about to leave",
+    "A woman putting a record on a turntable and lowering the needle, warm light",
+    "A man drying dishes and stacking them, window over the sink, evening light",
+]
+
+# Clinical, scientific and ceremonial. Cool and neutral lighting, and content
+# types with almost no representation in the original bank.
+INSTITUTIONAL = [
+    "A surgeon tying a suture knot, gloved hands, bright overhead theatre light",
+    "A nurse taking a blood pressure reading, cuff inflating, clinic light",
+    "A dentist examining with a small mirror, bright chair lamp, blue gloves",
+    "A physiotherapist guiding a shoulder through a slow range of motion",
+    "A radiographer positioning a patient's arm, cool scanner room light",
+    "A paramedic checking a monitor in the back of an ambulance, harsh strip light",
+    "A researcher pipetting into a row of wells, bright bench light, blue gloves",
+    "A chemist swirling a flask over a stirrer, coloured solution, fume hood light",
+    "A microscopist adjusting focus and looking up, dim room, illuminated stage",
+    "An astronomer adjusting a telescope mount in a dome, red safety light",
+    "A geologist splitting a rock with a hammer and examining the face, field light",
+    "A marine biologist lifting a sample net onto a deck, bright sea glare",
+    "An archaeologist brushing soil from a fragment in a trench, hard sun",
+    "A vet listening to a dog's chest with a stethoscope, calm clinic light",
+    "A lab technician loading a centrifuge and closing the lid, cool light",
+    "An engineer testing a circuit board with a probe, bench lamp, dark room",
+    "A metrologist reading a dial gauge on a granite table, clean bright light",
+    "A wind tunnel technician adjusting a model, cavernous space, work lights",
+    "A librarian scanning a barcode and stamping a date, quiet counter light",
+    "A judge writing a note at the bench, wood panelling, warm formal light",
+    "A priest lighting altar candles, tall stained glass windows behind",
+    "A monk sweeping a courtyard slowly, early light on stone",
+    "A bride adjusting a veil in front of a tall mirror, soft daylight",
+    "A tea ceremony host whisking matcha in a bowl, tatami room, diffuse light",
+    "A calligrapher preparing ink on a stone before writing, quiet dim room",
+    "A flag bearer holding a pole steady in the wind, bright overcast sky",
+    "A choir conductor lowering their hands at the end of a phrase, church light",
+    "A glassblower's assistant opening a furnace door, intense orange light spilling",
+    "A miner adjusting a headlamp underground, dust in the beam, black surround",
+    "A lighthouse keeper climbing a spiral stair, small windows, cold daylight",
+]
+
+# Older adults. Motion statistics, pacing and skin/hair detail that the rest of
+# the bank does not cover.
+#
+# NOTE: this bank contains no minors, by request. Nothing here depicts or
+# implies anyone under 18. See `assert_no_minors()` at the bottom of this file,
+# which enforces it as a test rather than a convention.
+ELDERS = [
+    "An older woman tending roses with secateurs, sunhat, bright garden light",
+    "An older man doing a crossword with a pen, reading glasses, lamp light",
+    "An elderly couple walking arm in arm down a lane, autumn trees, low sun",
+    "An older woman playing an upright piano from memory, warm parlour light",
+    "An older man feeding pigeons from a bench, city square, flat daylight",
+    "An elderly woman threading a needle carefully, bright window light",
+    "An older man polishing a vintage car bonnet, garage, hard work light",
+    "An elderly woman arranging photographs in an album, warm table lamp",
+    "An older man tying fishing flies at a bench, magnifier, focused lamp",
+    "An elderly woman watering a windowbox, morning light on the street below",
+    "An older man kneading bread dough slowly on a scrubbed table, morning light",
+    "An elderly woman winding wool into a ball from a skein, lamp beside her",
+    "An older man walking a dog slowly along a canal path, low autumn sun",
+    "An elderly woman pruning a bonsai with small shears, bright bench light",
+    "An older man carving a walking stick handle with a knife, porch light",
+    "An elderly woman practising scales on a violin, warm sitting room",
+    "An older man cleaning a pair of spectacles with a cloth, window light",
+    "An elderly woman stirring jam in a wide pan, steam rising, kitchen light",
+    "An older man setting up a chessboard piece by piece, warm lamp",
+    "An elderly woman doing slow tai chi in a park at dawn, mist on the grass",
+]
+
+PEOPLE_EXT = (TALKING + TRADES + KITCHEN + MUSIC + SPORT + WORK + OUTDOOR
+              + DOMESTIC + INSTITUTIONAL + ELDERS)
+
+# -------------------------------------------------------------------- animals
+ANIMALS_EXT = [
+    "A snow leopard picking its way along a rocky ledge, cold overcast light",
+    "A grey wolf trotting across a frozen lake, low winter sun behind",
+    "A brown bear turning over a rock in a shallow river, bright forest light",
+    "A lynx sitting still in deep snow, ears twitching, flat grey light",
+    "A stag standing in bracken and lifting its head, autumn morning mist",
+    "A herd of bison moving slowly through long grass, dust and low sun",
+    "A giraffe reaching up to strip leaves from a thorn tree, hard midday light",
+    "A zebra flicking its tail in dry grass, heat haze, bright savannah",
+    "A hippopotamus surfacing slowly in a muddy river, ears flicking water",
+    "A rhinoceros walking through dust at sunset, long shadows behind",
+    "A cheetah scanning the plain from a low mound, wind in the grass",
+    "A lioness lying in shade and yawning, dappled acacia light",
+    "A meerkat standing upright on a mound, alert, bright desert light",
+    "A sloth moving one arm slowly along a branch, humid green light",
+    "An orangutan swinging slowly between branches, dappled canopy light",
+    "A gorilla sitting and chewing a stem, deep forest shade",
+    "A lemur sunning itself on a branch, arms out, bright morning light",
+    "A red panda walking along a mossy log, soft overcast forest light",
+    "A koala shifting position in a eucalyptus fork, hazy blue-grey light",
+    "A kangaroo grazing then lifting its head, dry paddock, low golden sun",
+    "An otter rolling onto its back in a river, water beading on fur",
+    "A beaver swimming across a still pond at dusk, wake spreading behind",
+    "A badger emerging from a sett at twilight, low blue light",
+    "A hedgehog moving through leaf litter, torchlit, dark surround",
+    "A squirrel turning a nut over in its paws on a branch, autumn light",
+    "A hare sitting motionless in a ploughed field, low morning sun",
+    "A mole hill of soil shifting as something moves beneath, bright lawn",
+    "A field mouse climbing a grass stem, shallow focus, warm meadow light",
+    "A bat hanging and stretching one wing, dim cave light",
+    "A porcupine walking slowly across pine needles, dappled forest light",
+    "A raccoon washing something in shallow water, night, cool moonlight",
+    "A skunk foraging along a fence line at dusk, low warm light",
+    "A wild boar rooting in mud in a dark wood, shafts of light",
+    "A mountain goat standing on a narrow ledge, bright alpine light",
+    "An ibex silhouetted on a ridge at sunset, warm rim light",
+    "A camel walking across dunes, long shadow, hard low sun",
+    "A llama chewing and looking toward the camera, high altitude clear light",
+    "An alpaca in a misty field at dawn, breath visible",
+    "A donkey shaking its head in a dusty paddock, bright afternoon",
+    "A cow chewing cud in a barn doorway, bright light outside, dark within",
+    "A calf standing up unsteadily in straw, warm barn light",
+    "A pig rooting in mud in a farmyard, flat overcast light",
+    "A goat balancing on a fallen tree trunk, dappled woodland light",
+    "A ram standing in heather on a hillside, wind, grey light",
+    "A sheepdog lying alert with its head on its paws, warm kitchen light",
+    "A border collie crouching low and watching sheep, bright field",
+    "A labrador shaking water from its coat after a swim, spray flying",
+    "A beagle following a scent along a hedgerow, nose down, soft light",
+    "A greyhound running flat out across a field, motion blur behind",
+    "A terrier digging enthusiastically in sand, bright beach light",
+    "A husky pulling on a harness in snow, breath steaming, blue shadows",
+    "A dachshund trotting along a pavement, low camera, warm evening light",
+    "A pug asleep on a sofa and twitching, warm lamp light",
+    "A great dane standing and stretching, tall against a doorway, hall light",
+    "A dog shaking itself dry in slow motion, droplets flying, bright light",
+    "A puppy sniffing at a flower in a garden, soft morning light",
+    "A siamese cat stalking slowly across a wooden floor, low sun",
+    "A maine coon cat grooming its tail, warm window light",
+    "A cat leaping onto a windowsill and settling, afternoon light",
+    "A cat's paw batting at a water surface in a bowl, ripples, kitchen light",
+    "A cat sleeping curled in a patch of sun, slow breathing",
+    "A cat watching birds through glass, tail twitching, bright window",
+    "A rabbit washing its face with both paws, hutch straw, soft light",
+    "A guinea pig eating a piece of vegetable, warm indoor light",
+    "A ferret exploring along a shelf, quick movements, dim room",
+    "A hamster running in a wheel, warm cage light, dark surround",
+    "A horse tossing its mane in a paddock, dust and low sun",
+    "A foal running alongside a mare in a field, bright morning",
+    "A horse drinking from a trough, ripples, hard midday light",
+    "A shire horse pulling a cart along a lane, dust, warm afternoon",
+    "An eagle perched on a crag, feathers ruffling in the wind, cold light",
+    "A kestrel hovering over a field, wings adjusting, bright overcast",
+    "A barn owl flying slowly along a hedge at dusk, pale against dark",
+    "A heron standing motionless in shallow water then striking, grey light",
+    "A kingfisher perched on a reed above a stream, bright dappled light",
+    "A robin hopping along a fence rail, cold winter light",
+    "A blackbird pulling a worm from a lawn, bright overcast",
+    "A swan gliding across still water, reflection perfect beneath",
+    "A duck landing on a pond, feet skimming the surface, spray",
+    "A goose stretching its wings on a riverbank, morning mist",
+    "A pelican folding its wings after landing on a post, bright coastal light",
+    "A puffin standing on a cliff edge, wind ruffling, grey sea behind",
+    "A penguin waddling across ice then sliding onto its belly, flat white light",
+    "A flamingo standing on one leg and preening, shallow pink water",
+    "A peacock slowly raising its tail, iridescent, bright garden light",
+    "A rooster crowing on a fence post, early morning light",
+    "A hen scratching in a dust bath, bright farmyard",
+    "A pigeon strutting along a window ledge, city light behind",
+    "A crow examining something on tarmac, overcast urban light",
+    "A magpie hopping across a lawn, bright morning",
+    "A woodpecker drumming on a trunk, dappled forest light",
+    "A swallow perched on a wire, wind ruffling, bright sky behind",
+    "A stork standing in a nest on a chimney, warm evening light",
+    "A hummingbird feeding at a feeder, wings a blur, bright garden",
+    "A parrot cracking a seed in its beak, tropical shade",
+    "A budgerigar preening on a perch, warm indoor light",
+    "A canary hopping between perches, bright window light",
+    "A seagull hovering into the wind above a harbour, grey bright sky",
+    "A tern diving toward the sea and pulling up, bright glare",
+    "A turtle swimming slowly over a reef, sunlight rippling on its shell",
+    "A sea turtle surfacing to breathe, calm water, bright light",
+    "A shoal of sardines turning as one, shafts of light through blue water",
+    "A manta ray gliding past, silhouetted against the surface",
+    "A reef shark cruising along a drop-off, deep blue water",
+    "A clownfish moving among anemone tentacles, bright reef light",
+    "A seahorse gripping a strand of weed, gentle current, soft light",
+    "An octopus changing colour across a rock, shallow clear water",
+    "A cuttlefish hovering and rippling its fins, dim blue water",
+    "A jellyfish pulsing slowly in dark water, translucent, single beam",
+    "A starfish moving imperceptibly across a rock pool, bright shallow water",
+    "A crab scuttling sideways across wet sand, low sun",
+    "A lobster backing into a crevice, dim underwater light",
+    "A seal hauling itself onto a rock, spray, bright coastal light",
+    "A sea lion turning underwater in a shaft of sunlight",
+    "A whale's tail lifting and sliding under the surface, grey sea",
+    "A dolphin leaping clear of the water, spray catching the sun",
+    "A frog inflating its throat on a lily pad, warm pond light",
+    "A tree frog gripping a wet leaf, rain dripping, green light",
+    "A newt swimming through pondweed, dappled shallow light",
+    "A chameleon moving one foot at a time along a branch, bright foliage",
+    "A gecko clinging to a wall and blinking, warm lamp light",
+    "A snake moving slowly across warm sand, low sun, long shadow",
+    "A python coiled and slowly shifting, dim terrarium light",
+    "An iguana basking on a rock, bright tropical sun",
+    "A tortoise walking slowly across grass, bright afternoon",
+    "A crocodile lying motionless with one eye open, muddy bank, hard light",
+    "A monarch butterfly opening and closing its wings on a flower, bright sun",
+    "A dragonfly hovering over water then darting away, bright reflections",
+    "A praying mantis turning its head, shallow focus, warm light",
+    "A spider repairing a web strand by strand, backlit dew",
+    "A snail moving along a wet leaf, slow trail, grey light",
+    "A colony of ants carrying leaf pieces along a branch, dappled light",
+    "A beetle climbing a grass stem, bright meadow, shallow depth of field",
+    "A moth resting on a lit window at night, wings still",
+]
+
+# --------------------------------------------------------------------- nature
+NATURE_EXT = [
+    "A glacier calving slowly into a fjord, ice falling, flat grey light",
+    "Meltwater running in channels across blue glacier ice, bright midday sun",
+    "An iceberg drifting in still black water, pale under an overcast sky",
+    "Frost forming on a window pane, crystals spreading, cold dawn light",
+    "Icicles dripping steadily from a roof edge, bright winter sun",
+    "A frozen waterfall with water still running behind the ice, blue shade",
+    "Snow sliding off a pine branch in a clump, quiet forest",
+    "A blizzard blowing across an empty road, headlights faint, grey light",
+    "Wind-blown snow streaming off a mountain ridge, hard alpine light",
+    "A thaw stream cutting through old snow, bright spring light",
+    "Steam rising from a hot spring into cold air, blue morning light",
+    "A geyser building and then erupting, spray drifting, bright sky",
+    "Bubbling mud in a geothermal pool, sulphurous steam, flat light",
+    "Lava crusting over and cracking bright orange, dark night",
+    "Ash drifting from a volcanic vent, grey daylight, stark ground",
+    "Black sand beach with waves washing over it, high contrast light",
+    "A desert dune ridge with sand streaming off the crest, low sun",
+    "Heat shimmer over a salt flat, hard white light, distant mountains",
+    "A dust devil moving across dry ground, bright afternoon",
+    "A cactus standing against a deep blue desert sky, hard shadows",
+    "Rain falling on parched earth, dust puffing up, grey storm light",
+    "A flash flood running down a dry wash, muddy water, overcast",
+    "A canyon river seen from above, slow green water, deep shade",
+    "Sandstone walls glowing with reflected light in a slot canyon",
+    "A mesa at sunrise with mist in the valley below, warm light",
+    "Rolling fog spilling over a coastal ridge, late afternoon sun",
+    "A cloud inversion filling a valley at dawn, peaks above it",
+    "Cumulus clouds building over a plain, shadows moving across fields",
+    "A thunderstorm anvil lit from within, distant, dusk",
+    "Lightning branching across a night sky over open country",
+    "Rain sweeping visibly across a lake toward the camera, grey light",
+    "A rainbow forming over wet fields as the sun breaks through",
+    "Sunbeams breaking through cloud onto the sea, high contrast",
+    "A sunset burning orange behind silhouetted hills, clouds moving",
+    "The last light fading from a mountain face, cold blue shadow rising",
+    "Stars wheeling slowly above a dark treeline, deep night",
+    "The Milky Way over a still desert horizon, faint airglow",
+    "A full moon rising behind bare winter branches, cold light",
+    "Moonlight on rippling water, silver path, dark surround",
+    "Dawn light spreading across a frozen lake, pink to gold",
+    "A meadow of wildflowers moving in the wind, hazy summer sun",
+    "Poppies in a green field bending in a gust, bright overcast",
+    "A hillside of heather in flower, low cloud drifting over",
+    "Bluebells covering a woodland floor, dappled spring light",
+    "Long grass bending in waves across a hillside, low golden sun",
+    "A bamboo grove swaying, green filtered light, quiet",
+    "Aspen leaves trembling in a breeze, bright autumn yellow",
+    "Birch trunks white against dark forest, snow on the ground",
+    "A single oak in a field with cloud shadows passing over it",
+    "Moss and ferns on a damp forest floor, soft green light",
+    "Mushrooms on a fallen log, shallow focus, dim forest light",
+    "Mist hanging between conifers at first light, still air",
+    "Rain dripping from a leaf canopy, forest floor dark, green light",
+    "Sunlight moving across a forest floor as clouds pass",
+    "A stream running under exposed tree roots, dappled light",
+    "A beaver dam holding back a still pond, evening light",
+    "A reed bed moving in the wind at the edge of a lake, grey light",
+    "Water lilies on a still pond, insects skimming the surface",
+    "A mountain tarn perfectly reflecting the ridge above, still dawn",
+    "A river braiding across a gravel plain, seen wide, flat light",
+    "Rapids breaking white over boulders, spray, bright daylight",
+    "A waterfall plunging into a dark plunge pool, mist rising",
+    "A tidal estuary at low water, channels shining, wide grey sky",
+    "Waves breaking over a reef, spray blown back by offshore wind",
+    "A big swell rolling in under grey cloud, no horizon detail",
+    "Foam sliding back down a steep shingle beach, low sun",
+    "Rock pools reflecting the sky at low tide, bright light",
+    "Kelp moving in the surge in shallow clear water",
+    "Sea caves with light reflecting off the water onto the walls",
+    "A sea stack with waves washing around its base, evening light",
+    "Cliffs with seabirds wheeling below the top, bright overcast",
+    "A lighthouse beam sweeping through sea mist at night",
+    "A calm sea at dawn with no wind, pale gradient sky",
+    "Sunlight sparkling on choppy water, high contrast, hard light",
+    "A coral reef in shallow water, light rippling across it",
+    "Seagrass bending with the current, clear turquoise water",
+    "Rain dimpling the surface of a tropical lagoon, warm grey light",
+    "A mangrove root system in shallow water, dappled light",
+    "A salt marsh with channels filling on the tide, wide flat light",
+    "Sand ripples underwater with light patterns moving over them",
+    "A field of ripe barley moving in the wind, warm evening light",
+    "A vineyard on a slope with mist in the rows at dawn",
+    "An olive grove with silver leaves turning in the wind, hard sun",
+    "A tea plantation on terraced hills, low cloud, soft light",
+    "Rice paddies reflecting the sky, a slight breeze, flat light",
+    "A field of tulips in blocks of colour, bright overcast",
+    "An orchard in blossom with petals drifting down, soft sun",
+    "A hedgerow in autumn heavy with berries, low golden light",
+    "A stubble field after harvest with dust blowing, hard sun",
+    "A ploughed field with gulls following, grey winter light",
+    "A hay meadow with insects in the air, backlit summer light",
+    "A pumpkin field in October, low sun, long shadows",
+    "A cotton field with bolls moving in the wind, bright light",
+    "Sunflowers turning toward the light, blue sky, scattered cloud",
+    "A lavender row with bees working along it, hazy sun",
+    "A greenhouse interior with condensation running down the glass",
+    "Seedlings in trays with water droplets on the leaves, soft light",
+    "A compost heap steaming faintly in cold morning air",
+    "A garden pond with insects landing on the surface, dappled light",
+    "Ivy moving slightly on an old stone wall, overcast light",
+    "Cherry blossom against a grey sky, petals falling",
+    "Autumn leaves turning on a single maple, wind picking up",
+    "Bare branches against a white winter sky, snow beginning",
+    "Buds opening on a branch in early spring, soft light",
+    "A wildfire smoke haze turning the sun deep orange, dim light",
+    "Fresh snow untouched on a hillside at first light, blue shadows",
+    "A hail shower bouncing off a stone path, hard bright light",
+    "Wind flattening long grass in gusts across a headland",
+    "Dew on a spider web catching the low morning sun",
+    "Frost patterns on a fallen leaf, cold shallow light",
+    "A puddle freezing over, ice crystals spreading, dull light",
+    "Sea smoke rising off water on a freezing morning, low sun",
+    "A sandstorm approaching across flat desert, dimming the light",
+    "A meteor streaking across a dark star field, brief",
+    "Clouds racing across the moon, light coming and going",
+    "Bioluminescence glowing in breaking waves at night",
+    "A tide pool at dusk reflecting the last colour in the sky",
+    "Sun setting directly into the sea, path of light on the water",
+    "A high alpine glacier under hard blue sky, wind blowing snow",
+    "A crevasse field seen from above, deep blue shadow",
+    "A mountain hut with cloud streaming past, cold flat light",
+    "Scree slope with small stones sliding, hard midday sun",
+    "An alpine meadow with flowers below a snow peak, bright light",
+    "A rock face with water seeping and dripping, shaded",
+    "Wind-carved rock formations in a desert, low raking light",
+    "A limestone pavement with grass in the cracks, grey light",
+    "A peat bog with cotton grass moving in the wind, flat light",
+    "A chalk stream running clear over weed, bright dappled light",
+    "A weir with water pouring evenly over the lip, grey light",
+    "A canal reflecting overhanging trees, still, soft light",
+    "A flooded field with trees standing in water, grey calm light",
+    "A river in spate carrying branches, brown water, overcast",
+    "A dry riverbed with cracked mud, hard sun, heat shimmer",
+    "A spring bubbling up through sand, clear water, dappled light",
+    "A waterfall seen from behind through the falling curtain",
+    "Rain on a tin roof seen from beneath the eaves, grey light",
+    "A storm passing and light returning across a valley",
+]
+
+# ---------------------------------------------------------------------- urban
+URBAN_EXT = [
+    "A cobbled street in the rain at night, single streetlamp, reflections",
+    "A wide avenue at dawn with almost no traffic, low sun down the street",
+    "A crossing signal changing and a few people stepping off the kerb",
+    "A revolving door turning slowly in a lobby, bright exterior behind",
+    "An escalator moving in an empty station concourse, cool light",
+    "A lift lobby with floor indicators changing, polished stone, even light",
+    "A parking garage with a car's headlights sweeping the pillars",
+    "A loading bay with a shutter rolling up, harsh sodium light",
+    "A fire escape on a brick wall with steam drifting past, night",
+    "A back alley with a single door light and rain falling",
+    "A laundrette at night with machines turning, fluorescent light",
+    "A late-night diner window seen from the street, warm interior",
+    "A bar interior with bottles backlit, dim and warm, quiet",
+    "A cinema foyer with carpet and dim wall lights, empty",
+    "A theatre auditorium with house lights slowly dimming",
+    "A concert hall stage empty with a single work light",
+    "A museum gallery with light from a high window across a floor",
+    "An art gallery corridor with spotlit walls, polished concrete",
+    "A cathedral nave with light falling through clerestory windows",
+    "A mosque interior with patterned light on the carpet",
+    "A temple courtyard with incense smoke drifting, warm light",
+    "A synagogue interior with warm lamps and wooden benches",
+    "A university quad with long shadows in late afternoon",
+    "A school corridor empty between classes, fluorescent light",
+    "A hospital corridor with a trolley passing, cool even light",
+    "An office floor after hours with a few desk lamps still on",
+    "A server room with rows of blinking lights, cold blue glow",
+    "A control room with banks of monitors, dim surround",
+    "A print works with a press running, harsh industrial light",
+    "A warehouse aisle with a forklift moving away, high bay lights",
+    "A shipping container yard with a crane moving, hard daylight",
+    "A dockside with a ship's hull filling the frame, grey light",
+    "A fishing harbour at dawn with boats returning, soft light",
+    "A marina with halyards moving and reflections shifting, bright light",
+    "A canal lock filling slowly, water churning, overcast",
+    "A railway yard with wagons being shunted, flat grey light",
+    "A level crossing with barriers lowering, evening light",
+    "A tram stop with a tram arriving, wet road, night",
+    "A bus depot with buses idling, exhaust visible, cold morning",
+    "An airport gate with a plane pushing back beyond the glass",
+    "A runway seen through heat haze, aircraft taxiing, hard sun",
+    "A motorway bridge with traffic streaming beneath at dusk",
+    "A toll booth at night with headlights approaching, sodium light",
+    "A petrol station forecourt at night, bright canopy, dark surround",
+    "A car wash with brushes and water on the windscreen from inside",
+    "A multi-storey rooftop car park at sunset, city beyond",
+    "A construction site with a crane slewing slowly, bright sky",
+    "Scaffolding on a building facade with sheeting moving in the wind",
+    "A demolition site with dust settling, hard afternoon light",
+    "A newly poured concrete floor being smoothed, work lights",
+    "A road crew laying tarmac, steam rising, bright daylight",
+    "A window cleaner's cradle moving down a glass facade",
+    "A city square with pigeons and a fountain running, bright light",
+    "A fountain switching off and the water settling, evening",
+    "A park bandstand empty with leaves blowing across, autumn light",
+    "A municipal rose garden with a sprinkler turning, bright sun",
+    "A cemetery with long shadows between stones, low winter sun",
+    "A canal towpath with a narrowboat passing slowly, soft light",
+    "A pedestrian bridge with the river running below, grey light",
+    "A riverside walk at blue hour with lights coming on",
+    "A skyline seen across water with lights reflecting, dusk",
+    "A rooftop garden with plants moving in the wind, city behind",
+    "A balcony with washing drying and traffic sound implied, warm light",
+    "An apartment block facade with lights on in some windows, night",
+    "A stairwell with light falling from a high window, dust in the air",
+    "A hotel corridor with patterned carpet and wall sconces",
+    "A hotel room window with curtains moving and city beyond",
+    "A launderette window seen from outside at night, warm inside",
+    "A bakery window with trays being filled from behind, morning",
+    "A butcher's shop with a counter light and hanging scales",
+    "A greengrocer arranging crates on the pavement, bright morning",
+    "A flower stall with buckets of stems, awning shadow, bright light",
+    "A fish market slab with ice and a hose running, cold light",
+    "A spice market with sacks open, warm dusty light",
+    "A night market with hanging bulbs and steam, warm and dark",
+    "A food truck window with steam and a warm interior light, night",
+    "A vending machine glowing in an empty passage at night",
+    "A phone box lit from within on an empty street",
+    "A bus shelter with rain running down the glass, night",
+    "A pedestrian underpass with tiled walls and a strip light",
+    "A graffiti wall with paint still wet, bright daylight",
+    "A mural being painted from a cherry picker, bright sky",
+    "A shopfront shutter with peeling paint, hard afternoon sun",
+    "A barber's pole turning outside a shop, evening light",
+    "A neon sign flickering on above a doorway, dusk",
+    "A cinema marquee with letters lit, rain on the pavement",
+    "An arcade with machine lights and no one there, dark surround",
+    "A bowling alley lane with pins resetting, low warm light",
+    "A swimming pool hall empty with light rippling on the ceiling",
+    "An ice rink empty with the surface being resurfaced, cold light",
+    "A gym at night with one row of lights on, machines still",
+    "A boxing gym with a bag swinging slightly, dusty light",
+    "A dance studio empty with mirrors and a barre, north light",
+    "A recording studio live room empty, warm lamps, cables",
+    "A radio tower blinking on a hill at dusk, city lights below",
+    "A water tower against an evening sky, birds circling",
+    "A power station cooling tower with steam drifting, flat light",
+    "A wind turbine turning slowly against grey cloud",
+    "A solar farm with panels tracking, hard bright light",
+    "An electricity pylon line receding across fields, low sun",
+    "A railway tunnel mouth with a train emerging, hard contrast",
+    "A metro platform with a train's lights approaching in the tunnel",
+    "A ticket hall with departure boards updating, cool light",
+    "A cable car ascending above a city, cabin swaying slightly",
+    "A funicular climbing a steep slope, town below, bright light",
+    "A ferris wheel turning slowly at dusk, lights coming on",
+    "A carousel turning with lights and horses, evening",
+    "A fairground at night with lights moving, dark between",
+    "A pier stretching out to sea with lamps lit, blue hour",
+    "A promenade with waves breaking over the sea wall, grey light",
+    "A beach hut row with paint bright under a flat overcast sky",
+    "A seaside arcade front at night, saturated light, wet pavement",
+    "A harbour wall with waves bursting over it, storm light",
+    "A coastal road with mist blowing across, headlights faint",
+    "A mountain village street with snow banked at the sides",
+    "A hill town at dusk with lights coming on in the houses",
+    "A desert highway with heat shimmer and no traffic",
+    "A dirt road with dust hanging after a vehicle has passed",
+    "A tree-lined avenue with light flickering through the canopy",
+]
+
+# -------------------------------------------------------------------- objects
+OBJECTS_EXT = [
+    "Honey drizzling slowly from a dipper into a jar, warm backlight",
+    "Milk being poured into black coffee, clouds blooming, bright light",
+    "A tea bag steeping and colour spreading through hot water",
+    "Butter melting in a hot pan and sliding across it",
+    "Chocolate being poured into a mould, glossy, cool studio light",
+    "Caramel being pulled and folded on a marble slab, warm light",
+    "Cheese being sliced with a wire, clean edge, bright light",
+    "Bread being torn open with steam escaping, warm light",
+    "A soft-boiled egg being opened and yolk running, bright light",
+    "Batter being whisked until smooth, bowl rotating, soft light",
+    "Herbs being chopped finely, blade rocking, bright board",
+    "Citrus being zested with a fine grater, oil misting, close",
+    "Garlic being crushed under a knife blade, bright kitchen light",
+    "A jar of preserves being sealed and the lid popping, warm light",
+    "Popcorn popping in a covered pan, lid lifting slightly",
+    "Rice being rinsed in a bowl, water clouding, bright light",
+    "Noodles being lifted from broth with chopsticks, steam rising",
+    "Sparkling water being poured, bubbles rising, backlit glass",
+    "Wine being swirled in a glass, legs running down the side",
+    "Beer being poured with a head forming, warm bar light",
+    "Ice cracking as liquid is poured over it, close, dark background",
+    "A cocktail being strained into a chilled glass, dim bar light",
+    "A teapot pouring a thin steady stream, steam, soft light",
+    "A moka pot bubbling on a hob, coffee rising, warm light",
+    "An espresso extraction running into a cup, close, machine light",
+    "Steam wand frothing milk, vortex forming, cafe light",
+    "A kettle boiling with steam pouring from the spout, kitchen light",
+    "Oil shimmering in a hot pan just before smoking",
+    "A candle being lit and the flame settling, dark room",
+    "A match striking and flaring, close, black background",
+    "Incense smoke curling upward in still air, single beam of light",
+    "A wood fire collapsing and sparks rising, dark surround",
+    "Embers glowing and pulsing in a grate, dark room",
+    "A gas hob ring igniting, blue flame steadying",
+    "A blowtorch flame playing across metal, dark workshop",
+    "Molten metal being poured into a mould, intense glow",
+    "Sparks streaming from an angle grinder, dark background",
+    "A welding arc flaring and dimming, deep shadow around",
+    "Solder melting and flowing along a joint, bright close light",
+    "A soldering iron tip being cleaned on a sponge, steam",
+    "A drill bit cutting into wood, shavings spiralling out",
+    "A saw blade cutting through a board, dust rising, work light",
+    "A plane taking a long shaving off a board, curling away",
+    "Sandpaper moving across wood, dust and changing sheen",
+    "A chisel paring a shaving from end grain, bright bench light",
+    "A clock's second hand sweeping, close, warm light",
+    "A pendulum swinging in a case, brass catching lamp light",
+    "An hourglass with sand running, backlit, dark background",
+    "A metronome ticking back and forth, warm lamp",
+    "A music box cylinder turning and plucking the comb, close",
+    "A record needle settling into the groove, close, warm light",
+    "A cassette reel turning behind a window, warm light",
+    "A film projector reel turning, dust in the beam, dark booth",
+    "A camera shutter opening and closing in slow motion",
+    "A lens aperture blades closing down, close, studio light",
+    "A typewriter key striking the page, ink impression, lamp light",
+    "A fountain pen nib laying a wet line of ink, close",
+    "Ink being drawn into a pen from a bottle, warm light",
+    "A pencil being sharpened, shavings curling, bright light",
+    "A brush loading paint from a palette, studio light",
+    "Watercolour blooming into wet paper, bright even light",
+    "A palette knife spreading thick paint on canvas, raking light",
+    "Ink being rolled out on a slab with a brayer, even light",
+    "A screen print squeegee lifting to reveal a print, bright light",
+    "Clay being wedged on a bench, hands pressing, workshop light",
+    "A glaze being poured over a bisque bowl, cool studio light",
+    "A kiln peephole glowing orange, dark surround",
+    "Thread being wound onto a bobbin, close, task light",
+    "A needle pulling thread through fabric, close, lamp light",
+    "Scissors cutting a long clean line through fabric, bright light",
+    "A knot being tied in rope, hands working, flat light",
+    "A shoelace being threaded and pulled tight, close",
+    "Wool being spun onto a spindle, fibre drawing out, warm light",
+    "A loom shuttle passing and the beater knocking up, workshop light",
+    "Dye spreading through white fabric in a bath, overhead light",
+    "Fabric being steamed and creases falling out, backlit",
+    "A paper sheet being folded sharply, crisp edge, bright light",
+    "A book page turning in a draught, warm reading light",
+    "Dust motes drifting through a shaft of light in a still room",
+    "Rain running down a window pane, lights beyond out of focus",
+    "Condensation forming on a cold glass, drops beginning to run",
+    "Soap bubbles forming and popping in a sink, bright light",
+    "Water swirling down a drain, close, bright light",
+    "A drop falling into still water, crown and ripples, dark background",
+    "Oil and water separating in a jar, backlit, slow motion",
+    "Sand running through fingers onto a pile, low sun",
+    "A magnet dragging iron filings into lines, bright even light",
+    "A pendulum of a Newton's cradle clicking, warm desk light",
+    "A spinning top slowing and beginning to wobble, dark surface",
+    "A domino line falling in sequence, bright even light",
+    "A balloon slowly deflating, surface wrinkling, plain background",
+    "A feather falling slowly through still air, dark background",
+    "A soap film shifting colours across a wire loop, dark surround",
+    "Frost melting off a windscreen in the morning sun",
+]
+
+EXT_ALL = PEOPLE_EXT + ANIMALS_EXT + NATURE_EXT + URBAN_EXT + OBJECTS_EXT
+
+# One in-place substitution in the ORIGINAL bank. Index 20 was
+# "A child blowing bubbles in a garden" -- this bank is to contain no minors,
+# and deleting the entry instead of replacing it would renumber prompts 21-95
+# and silently break world_p44 (owl), world_p60 (waterfall) and world_p82
+# (court), which are referenced by index everywhere in FINDINGS and by
+# `demo.py --prompt-idx`. Substituting the text keeps every index where it was.
+#
+# Consequence to act on: the six teacher clips already generated for prompt 20
+# (`data/teacher/p0020_s0*.pt`) were made from the old text and must be deleted
+# so `gen_teacher.py` regenerates them. Nothing else is affected.
+REPLACEMENTS = {
+    20: "A woman blowing bubbles in a garden, sunlight through the bubbles, summer afternoon",
+}
+
+# Words that would indicate a minor in a prompt. Checked against the FINAL bank,
+# so it covers the original 96 as well as the extension -- a convention that is
+# only written down in a docstring is one nobody notices breaking.
+_MINOR_TERMS = (
+    'child', 'children', 'kid', 'kids', 'boy', 'girl', 'toddler', 'baby',
+    'babies', 'infant', 'teen', 'teenage', 'teenager', 'youngster', 'schoolboy',
+    'schoolgirl', 'schoolchild', 'pupil', 'playground', 'nursery', 'daughter',
+    'son', 'juvenile', 'minor', 'adolescent', 'youth', 'newborn', 'grandchild',
+)
+
+
+def assert_no_minors(prompts):
+    """Raise if any prompt reads as depicting a minor. Word-boundary matched, so
+    'person' does not trip 'son' and 'stonemason' does not trip it either."""
+    import re
+    pat = re.compile(r'\b(' + '|'.join(_MINOR_TERMS) + r')\b', re.I)
+    bad = [(i, p) for i, p in enumerate(prompts) if pat.search(p)]
+    if bad:
+        raise AssertionError(
+            f'{len(bad)} prompt(s) reference a minor; this bank must contain '
+            f'none. First: index {bad[0][0]} -- {bad[0][1]!r}')
+    return True

wanstreamer/rope.py

@@ -0,0 +1,77 @@
+"""Precomputed 3D RoPE with a temporal offset.
+
+Two problems with upstream `wan.modules.model.rope_apply`:
+
+  1. It reconstructs its frequency tensor on EVERY call -- a float64 cat/expand/
+     reshape over the whole sequence -- and the streaming path calls it twice per
+     block, 30 blocks per latent frame (60x per frame-forward). The profile in
+     PROGRESS.md §6 attributes a large share of the 55 ms frame-forward to this
+     kind of churn.
+  2. It has no temporal offset, so the streaming path (which patch-embeds one
+     latent frame at a time, giving grid f=1) always indexes freqs[0][:1] --
+     temporal position 0 for every frame. Measured in diag/rope_probe.py: using
+     the true absolute frame index instead cuts normalised flow error by up to
+     3.2x, because the pretrained backbone still expects real temporal RoPE.
+
+This module precomputes the per-frame frequency table once per (resolution,
+max_frames) and reduces application to one complex multiply.
+"""
+import torch
+
+
+class RopeTable:
+    """Per-temporal-index RoPE frequency tables for a fixed spatial grid.
+
+    freqs: the model's [1024, c] complex buffer (WanModel.freqs), c = head_dim/2.
+    Table t holds the flattened (1, h, w) grid at temporal position t, shaped
+    [h*w, 1, c] so it broadcasts over batch and heads.
+    """
+
+    def __init__(self, freqs, h, w, max_frames, device, dtype=torch.complex64):
+        c = freqs.shape[1]
+        # upstream split: temporal gets the remainder, height and width get c//3 each
+        f_t, f_h, f_w = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+        assert max_frames <= f_t.shape[0], (
+            f'max_frames={max_frames} exceeds rope_params table ({f_t.shape[0]})')
+        assert h <= f_h.shape[0] and w <= f_w.shape[0], (
+            f'grid {h}x{w} exceeds rope_params table '
+            f'({f_h.shape[0]}x{f_w.shape[0]}) -- raise rope_params(1024, ...)')
+
+        spatial = torch.cat([
+            f_h[:h].view(h, 1, -1).expand(h, w, -1),
+            f_w[:w].view(1, w, -1).expand(h, w, -1),
+        ], dim=-1).reshape(h * w, -1)                      # [S, c_h + c_w]
+
+        tables = []
+        for t in range(max_frames):
+            temporal = f_t[t].view(1, -1).expand(h * w, -1)  # [S, c_t]
+            tables.append(torch.cat([temporal, spatial], dim=-1))
+        # [max_frames, S, 1, c]
+        self.table = torch.stack(tables).unsqueeze(2).to(device=device, dtype=dtype)
+        self.h, self.w, self.seq = h, w, h * w
+        self.max_frames = max_frames
+
+    def frame(self, t_index):
+        """[S, 1, c] complex table for a single latent frame at temporal index t."""
+        if t_index >= self.max_frames:
+            raise IndexError(f'temporal index {t_index} >= max_frames {self.max_frames}')
+        return self.table[t_index]
+
+    def span(self, t_start, num_frames):
+        """[num_frames*S, 1, c] for a contiguous run of latent frames."""
+        end = t_start + num_frames
+        if end > self.max_frames:
+            raise IndexError(f'span [{t_start},{end}) exceeds max_frames {self.max_frames}')
+        return self.table[t_start:end].reshape(num_frames * self.seq, 1, -1)
+
+
+def apply_rope(x, table):
+    """x: [B, L, n, d] real -> [B, L, n, d] real, rotated by `table` [L, 1, c].
+
+    Done in float32 complex rather than upstream's float64. Verified equivalent
+    within bf16 tolerance by tests/test_streaming_core.py.
+    """
+    b, l, n, d = x.shape
+    xc = torch.view_as_complex(x.float().reshape(b, l, n, d // 2, 2))
+    out = torch.view_as_real(xc * table.unsqueeze(0))
+    return out.flatten(3).to(x.dtype) if x.dtype != torch.float32 else out.flatten(3)

wanstreamer/serve/__init__.py

@@ -0,0 +1,7 @@
+"""Browser demo layer: a serving shell around the project's inference core.
+
+Everything that touches the streaming maths lives in `wanstreamer` proper
+(`stream.FewStepStreamer`, `blockcausal`, `kvcache`, `rope`). This package only
+drives it: session management, JPEG framing, a websocket, world generation from
+text, and the UI.
+"""

wanstreamer/serve/auth.py

@@ -0,0 +1,56 @@
+"""Shared token gate for the demo server.
+
+Starlette's `BaseHTTPMiddleware` never sees websocket scopes, and the frame stream is
+a websocket, so this is a plain ASGI wrapper rather than a FastAPI middleware.
+
+A token supplied in the query string is echoed back as a cookie, so visiting
+`/?token=...` once is enough: the websocket handshake then carries it automatically
+and the URL can be shared without the query string surviving in every request.
+"""
+
+from urllib.parse import parse_qs
+
+COOKIE = "livewan_token"
+
+
+class TokenAuth:
+    def __init__(self, app, token):
+        self.app, self.token = app, token
+
+    def _supplied(self, scope):
+        headers = {k.lower(): v for k, v in (scope.get("headers") or [])}
+        token = parse_qs(scope.get("query_string", b"").decode()).get("token", [None])[0]
+        if token:
+            return token
+        auth = headers.get(b"authorization", b"").decode()
+        if auth.startswith("Bearer "):
+            return auth[7:]
+        for part in headers.get(b"cookie", b"").decode().split(";"):
+            k, _, v = part.strip().partition("=")
+            if k == COOKIE:
+                return v
+        return None
+
+    async def __call__(self, scope, receive, send):
+        if not self.token or scope["type"] not in ("http", "websocket"):
+            return await self.app(scope, receive, send)
+
+        if self._supplied(scope) != self.token:
+            if scope["type"] == "websocket":
+                return await send({"type": "websocket.close", "code": 1008})
+            await send({"type": "http.response.start", "status": 401,
+                        "headers": [(b"content-type", b"text/plain; charset=utf-8")]})
+            return await send({"type": "http.response.body",
+                               "body": b"unauthorized -- append ?token=<token> to the URL"})
+
+        if scope["type"] == "http":
+            async def _send(msg):
+                if msg["type"] == "http.response.start":
+                    msg["headers"] = list(msg.get("headers") or []) + [
+                        (b"set-cookie",
+                         f"{COOKIE}={self.token}; Path=/; SameSite=Lax; Max-Age=604800".encode())
+                    ]
+                await send(msg)
+
+            return await self.app(scope, receive, _send)
+        await self.app(scope, receive, send)

wanstreamer/serve/conditioning.py

@@ -0,0 +1,94 @@
+"""Prompt conditioning for the demo: the 96 prompt bank, plus free text encoding.
+
+Two sources of conditioning, deliberately kept distinct:
+
+* **The bank** (`data/prompts.pt`) -- 96 umt5-xxl embeddings that every clip in the
+  project was generated and trained under. Zero cost, and exactly the conditioning
+  the published numbers refer to.
+* **Free text** -- encoded here, on this machine, with umt5-xxl (11.4 GB, loaded
+  lazily on first use). Note that umt5 embeddings are mildly hardware-dependent, so
+  text encoded here is not numerically identical to what the same string would give
+  on the training box. It looks fine; it just isn't the *same* conditioning, so
+  free-text results are not strictly comparable to the bank's published metrics.
+
+Both paths end in the same shape: [1, 512, 4096], zero padded, which is what
+`WanModel` expects.
+"""
+
+from pathlib import Path
+
+import torch
+
+TEXT_LEN = 512
+
+THEMES = [
+    ("People", 0, 40),
+    ("Animals", 40, 56),
+    ("Nature", 56, 72),
+    ("City", 72, 86),
+    ("Objects", 86, 96),
+]
+
+
+def theme_of(idx):
+    for name, lo, hi in THEMES:
+        if lo <= idx < hi:
+            return name
+    return "Other"
+
+
+class PromptBank:
+    def __init__(self, path):
+        d = torch.load(path, map_location="cpu", weights_only=False)
+        self.texts = d["prompts"]
+        self.pos = d["pos"]  # [96, 512, 4096] fp16
+        self.neg = d["neg"]
+        self.neg_prompt = d["neg_prompt"]
+
+    def __len__(self):
+        return len(self.texts)
+
+    def embedding(self, idx):
+        return self.pos[idx : idx + 1].float()
+
+    def catalogue(self):
+        return [
+            {"idx": i, "text": t, "theme": theme_of(i)} for i, t in enumerate(self.texts)
+        ]
+
+
+class TextEncoder:
+    """Lazy umt5-xxl. 11.4 GB -- only loaded if someone actually types a prompt."""
+
+    def __init__(self, checkpoint, tokenizer_path, wan_repo, device="cuda"):
+        self.checkpoint = str(checkpoint)
+        self.tokenizer_path = str(tokenizer_path)
+        self.wan_repo = str(wan_repo)
+        self.device = device
+        self._model = None
+
+    @property
+    def loaded(self):
+        return self._model is not None
+
+    def load(self):
+        if self._model is not None:
+            return
+        from wan.modules.t5 import T5EncoderModel
+
+        self._model = T5EncoderModel(
+            text_len=TEXT_LEN, dtype=torch.bfloat16, device=self.device,
+            checkpoint_path=self.checkpoint, tokenizer_path=self.tokenizer_path)
+
+    @torch.no_grad()
+    def encode(self, text):
+        """-> [1, 512, 4096] float32, zero-padded exactly as the bank is."""
+        self.load()
+        ctx = self._model([text], self.device)[0]  # [L, 4096], L <= 512
+        out = torch.zeros(TEXT_LEN, ctx.shape[1], dtype=torch.float32, device=ctx.device)
+        out[: ctx.shape[0]] = ctx.float()
+        return out.unsqueeze(0).cpu()
+
+    def unload(self):
+        self._model = None
+        torch.cuda.empty_cache()

wanstreamer/serve/engine.py

@@ -0,0 +1,482 @@
+"""Session engine for the browser demo.
+
+This is a thin serving layer over the project's own inference core: the student is
+rolled out by `wanstreamer.stream.FewStepStreamer`, which is the code the checkpoint
+was distilled and measured under. Nothing here reimplements the streaming maths --
+the block loop, the block-causal K/V cache, the renoise sampler and `latent_norm`
+all live in the core, and this module only drives them and turns latents into JPEG.
+
+One GPU, one stream. A background worker generates blocks and pushes frames into a
+bounded queue. The queue is deliberately short (~1.5 s) because it *is* the steering
+latency (anything a viewer has already been sent cannot be steered any more).
+
+Two controls, doing genuinely different things:
+
+  steer  -- swap the cross-attention conditioning, keep the K/V cache. The scene
+            continues. Cheap, instant, no reload.
+  scene  -- tear the stream down and reopen from another world. A cut.
+"""
+
+import json
+import queue
+import threading
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import torch
+
+from ..stream import FewStepStreamer
+from .paths import WORLDS_DIR
+from .streamdecode import StreamingVAEDecoder
+from .conditioning import PromptBank, TextEncoder
+from .worldgen import WorldGenerator
+
+WORLDS = [0, 44, 60, 82]  # the four shipped caches; generated ones get ids from 1000
+GENERATED_BASE_ID = 1000
+FPS = 16
+
+# Defaults matching the published command in the model card. `shift=1.0` is the
+# uniform few-step spacing the student was distilled under (see
+# wanstreamer.stream.few_step_schedule); `window` is in LATENT FRAMES, not blocks.
+DEFAULTS = dict(block=3, steps=2, window=6, latent_norm=1.0, shift=1.0,
+                sampler="renoise")
+
+
+def _encode_jpeg(arr, quality=88):
+    import cv2
+
+    ok, buf = cv2.imencode(".jpg", arr[:, :, ::-1], [cv2.IMWRITE_JPEG_QUALITY, quality])
+    if not ok:
+        raise RuntimeError("jpeg encode failed")
+    return buf.tobytes()
+
+
+class StreamExhausted(RuntimeError):
+    pass
+
+
+@dataclass
+class Status:
+    state: str = "idle"  # idle | loading | generating | streaming | error
+    detail: str = ""
+    world: int | None = None
+    prompt: str = ""
+    prompt_source: str = ""  # "bank" | "text"
+    stats: dict = field(default_factory=dict)
+    error: str = ""
+    progress: float | None = None  # 0..1 while generating a world
+
+
+class Engine:
+    def __init__(self, assets, wan_repo, base_dir, weights=None, device="cuda",
+                 jpeg_quality=88, buffer_seconds=1.5, worlds_dir=None,
+                 allow_worldgen=True, size=(640, 368), compile_vae=True):
+        self.assets = Path(assets)
+        self.wan_repo = Path(wan_repo)
+        self.base_dir = Path(base_dir)
+        self.weights = Path(weights or self.assets / "checkpoints/t14b_b64/latest.pt")
+        self.worlds_dir = Path(worlds_dir or WORLDS_DIR)
+        self.allow_worldgen = allow_worldgen
+        self.size = size
+        self.device = device
+        self.jpeg_quality = jpeg_quality
+        self.compile_vae = compile_vae
+        self.maxframes = int(buffer_seconds * FPS)
+
+        self.worlds = {}
+        self.worldgen = None
+        self.status = Status()
+        self.lock = threading.Lock()
+        self._genlock = threading.Lock()
+        self.frames = queue.Queue(maxsize=self.maxframes)
+        self._worker = None
+        self._stop = threading.Event()
+        self._pending_prompt = None
+
+        self.model = self.decoder = self.streamer = None
+        self._cur_emb = None
+        self.bank = self.encoder = None
+        self.step = None
+        self.cfg = None
+        self._loaded = False
+        self._timings = {}
+        self._frames_emitted = 0
+        self._blocks = 0
+
+    # ------------------------------------------------------------------ load
+
+    def load(self, progress=None):
+        def say(msg):
+            self.status.state, self.status.detail = "loading", msg
+            if progress:
+                progress(msg)
+
+        say("reading the prompt bank")
+        self.bank = PromptBank(self.assets / "data/prompts.pt")
+
+        say(f"loading the student ({self.weights.name})")
+        self.model, self.cfg, self.step = self._load_student()
+
+        say("loading the Wan2.1 VAE")
+        torch.backends.cudnn.benchmark = True
+        self.decoder = StreamingVAEDecoder(
+            self.base_dir / "Wan2.1_VAE.pth", self.wan_repo, self.device
+        )
+
+        if self.compile_vae:
+            say("compiling the decoder (one-off, ~40 s)")
+            try:
+                import torch._dynamo as dynamo
+
+                dynamo.config.recompile_limit = 64
+                self.decoder.model.decoder = torch.compile(
+                    self.decoder.model.decoder, dynamic=False)
+                w = torch.load(self.assets / "out/world_p60.pt", map_location="cpu",
+                               weights_only=False)
+                self.decoder.decode(w["latents"][:, :6].to(self.device))
+                self.decoder.reset()
+            except Exception as e:  # an optimisation, never a requirement
+                self.status.detail = f"decoder compile skipped: {e}"
+
+        tok = self.base_dir / "umt5-tokenizer"
+        self.encoder = TextEncoder(
+            self.base_dir / "models_t5_umt5-xxl-enc-bf16.pth",
+            tok if tok.exists() else "google/umt5-xxl", self.wan_repo, self.device,
+        )
+        if self.allow_worldgen:
+            self.worldgen = WorldGenerator(
+                self.base_dir / "diffusion_pytorch_model.safetensors",
+                self.cfg, self.device,
+            )
+        self._index_worlds()
+        self._loaded = True
+        self.status.state, self.status.detail = "idle", "ready"
+
+    def _load_student(self):
+        """Build the stock WanModel and load the distilled weights into it.
+
+        Mirrors scripts/demo.py: a checkpoint that silently half-loaded would report
+        base-model quality as if it were trained, so every parameter is checked.
+        """
+        from wan.configs import WAN_CONFIGS
+        from wan.modules.model import WanModel
+        from safetensors.torch import load_file
+
+        cfg = WAN_CONFIGS["t2v-1.3B"]
+        m = WanModel(dim=cfg.dim, ffn_dim=cfg.ffn_dim, freq_dim=cfg.freq_dim,
+                     num_heads=cfg.num_heads, num_layers=cfg.num_layers,
+                     window_size=cfg.window_size, qk_norm=True,
+                     cross_attn_norm=True, eps=1e-6)
+        m.load_state_dict(
+            load_file(str(self.base_dir / "diffusion_pytorch_model.safetensors")),
+            strict=True)
+        sd = torch.load(self.weights, map_location="cpu", weights_only=False)
+        src = sd.get("model", sd)
+        res = m.load_state_dict(src, strict=False)
+        got = {n for n, _ in m.named_parameters()} - set(res.missing_keys)
+        if len(got) != len(list(m.named_parameters())) or res.unexpected_keys:
+            raise RuntimeError(
+                f"bad checkpoint load: {len(res.missing_keys)} missing, "
+                f"{len(res.unexpected_keys)} unexpected")
+        return m.to(self.device).eval().requires_grad_(False), cfg, sd.get("step")
+
+    # ---------------------------------------------------------------- worlds
+
+    def _index_worlds(self):
+        self.worlds = {
+            w: {"path": self.assets / f"out/world_p{w}.pt", "prompt": self.bank.texts[w],
+                "generated": False, "seconds": None}
+            for w in WORLDS
+        }
+        self.worlds_dir.mkdir(parents=True, exist_ok=True)
+        for meta_path in sorted(self.worlds_dir.glob("world_*.json")):
+            try:
+                meta = json.loads(meta_path.read_text())
+                pt = meta_path.with_suffix(".pt")
+                if pt.exists():
+                    self.worlds[int(meta["id"])] = {
+                        "path": pt, "prompt": meta.get("prompt", ""),
+                        "generated": True, "seconds": meta.get("seconds")}
+            except Exception:
+                continue  # a half-written world should not stop the server booting
+
+    def _next_world_id(self):
+        used = [i for i in self.worlds if i >= GENERATED_BASE_ID]
+        return max(used) + 1 if used else GENERATED_BASE_ID
+
+    @property
+    def worldgen_available(self):
+        return bool(self.worldgen and self.worldgen.available)
+
+    def generate_world(self, text=None, idx=None, steps=30, seed=0, guide=5.0):
+        """Make a new opening world from text (or a bank prompt) with the base model.
+
+        Slow (~52 s at 30 steps) and the only non-real-time step in the project.
+        """
+        if not self.worldgen_available:
+            raise RuntimeError(
+                "world generation is unavailable — the Wan2.1 base transformer "
+                "(diffusion_pytorch_model.safetensors) is not present")
+        if not self._genlock.acquire(blocking=False):
+            raise RuntimeError("already generating a world")
+        try:
+            self._stop_worker()
+            pos, label, source = self.resolve_prompt(idx, text)
+            self.status = Status(state="generating", detail="encoding the prompt",
+                                 prompt=label, prompt_source=source, progress=0.0)
+
+            def on_step(i, n):
+                self.status.progress = i / n
+                self.status.detail = f"denoising the world — step {i} of {n}"
+
+            lat, secs = self.worldgen.generate(
+                pos, self.bank.neg.float(), size=self.size, steps=steps, seed=seed,
+                guide=guide, progress=on_step)
+
+            self.status.detail, self.status.progress = "decoding", 1.0
+            self.decoder.reset()
+            pixels = self.decoder.decode(lat)
+            self.decoder.reset()
+
+            wid = self._next_world_id()
+            path = self.worlds_dir / f"world_{wid}.pt"
+            # keep the conditioning with the world: reopening needs no text encoder
+            torch.save({"latents": lat.cpu(), "pixels": pixels,
+                        "prompt_emb": pos.cpu(), "base_seconds": secs}, path)
+            path.with_suffix(".json").write_text(json.dumps(
+                {"id": wid, "prompt": label, "steps": steps, "seed": seed,
+                 "guide": guide, "seconds": round(secs, 1)}, indent=1))
+            self.worlds[wid] = {"path": path, "prompt": label, "generated": True,
+                                "seconds": round(secs, 1)}
+            self.status = Status(state="idle", detail="world ready", prompt=label,
+                                 prompt_source=source)
+            return wid
+        except Exception as e:
+            self.status = Status(state="error", error=f"{type(e).__name__}: {e}")
+            raise
+        finally:
+            self._genlock.release()
+
+    def delete_world(self, wid):
+        w = self.worlds.get(int(wid))
+        if not w or not w["generated"]:
+            raise ValueError("only generated worlds can be deleted")
+        Path(w["path"]).unlink(missing_ok=True)
+        Path(w["path"]).with_suffix(".json").unlink(missing_ok=True)
+        self.worlds.pop(int(wid), None)
+
+
+
+    def resolve_prompt(self, idx=None, text=None):
+        """-> (embedding [1, 512, 4096], label, source)"""
+        if text:
+            text = text.strip()
+            if not text:
+                raise ValueError("empty prompt")
+            return self.encoder.encode(text), text, "text"
+        if idx is None:
+            raise ValueError("need a prompt index or text")
+        idx = int(idx)
+        if not 0 <= idx < len(self.bank):
+            extra = ""
+            if idx >= GENERATED_BASE_ID:
+                extra = (f" — {idx} looks like a generated world id, not a prompt. "
+                         "Generated worlds carry their own conditioning: pass no "
+                         "prompt_idx/prompt_text and it will be used.")
+            raise ValueError(f"prompt index must be 0-{len(self.bank)-1}{extra}")
+        return self.bank.embedding(idx), self.bank.texts[idx], "bank"
+
+    def _set_text(self, emb):
+        """Push raw umt5 conditioning into the streamer, embedded as the core wants."""
+        emb = emb.to(self.device)
+        with torch.amp.autocast("cuda", enabled=False):
+            ctx = self.model.text_embedding(emb.float())
+        self.streamer.set_text(ctx)
+        self._cur_emb = emb   # kept so a crossfade can interpolate from it
+
+
+
+    def start(self, world, idx=None, text=None, seed=0, steps=None, block=None,
+              window=None, latent_norm=None, shift=None, sampler=None):
+        if not self._loaded:
+            raise RuntimeError("engine not loaded")
+        world = int(world)
+        entry = self.worlds.get(world)
+        if entry is None:
+            raise ValueError(f"unknown world {world}; have {sorted(self.worlds)}")
+
+        block = block or DEFAULTS["block"]
+        steps = steps or DEFAULTS["steps"]
+        window = window or DEFAULTS["window"]
+        shift = DEFAULTS["shift"] if shift is None else shift
+        sampler = sampler or DEFAULTS["sampler"]
+        latent_norm = DEFAULTS["latent_norm"] if latent_norm is None else latent_norm
+
+        with self.lock:
+            self._stop_worker()
+            self.status = Status(state="loading", detail="opening the world", world=world)
+            w = torch.load(entry["path"], map_location="cpu", weights_only=False)
+
+            emb = label = source = None
+            if idx is None and text is None:
+                if not entry["generated"]:
+                    idx = world
+                elif w.get("prompt_emb") is not None:
+                    emb, label, source = w["prompt_emb"].float(), entry["prompt"], "text"
+                else:
+                    text = entry["prompt"]
+            if emb is None:
+                emb, label, source = self.resolve_prompt(idx, text)
+            self.status.prompt, self.status.prompt_source = label, source
+
+            world_lat = w["latents"]
+            nw = world_lat.shape[1]
+            self.streamer = None
+            torch.cuda.empty_cache()
+            self.streamer = FewStepStreamer(
+                self.model, width=self.size[0], height=self.size[1],
+                max_frames=1024, device=self.device, dtype=self.cfg.param_dtype,
+                window_frames=window,
+                # the world's K/V is pinned; the window bounds only the events after it
+                cache_frames=nw + window + block + 2,
+                block_frames=block, num_steps=steps, shift=shift, sampler=sampler,
+            )
+            self.streamer.latent_norm = latent_norm
+            self._set_text(emb)
+
+            self.decoder.reset()
+            self._frames_emitted = self._blocks = 0
+            self._pending_prompt = None
+            self._gen = torch.Generator(device=self.device).manual_seed(int(seed))
+
+            t0 = time.time()
+            self.streamer.set_world(world_lat.to(self.device, torch.float32))
+            pixels = self.decoder.decode(world_lat.to(self.device, torch.float32))
+            self._timings = {"start_s": round(time.time() - t0, 3)}
+            self._frames_emitted = pixels.shape[0]
+
+            self._stop.clear()
+            self._worker = threading.Thread(target=self._run, args=(pixels,),
+                                            daemon=True, name="livewan-worker")
+            self._worker.start()
+            self.status.state, self.status.detail = "streaming", ""
+        return self.status
+
+    def steer(self, idx=None, text=None, crossfade=0):
+        if self.streamer is None or self.status.state != "streaming":
+            raise RuntimeError("no stream is running")
+        emb, label, source = self.resolve_prompt(idx, text)
+        self._pending_prompt = (emb, int(crossfade))
+        self.status.prompt, self.status.prompt_source = label, source
+        return self.status
+
+    def stop(self):
+        with self.lock:
+            self._stop_worker()
+            self.status = Status(state="idle", detail="stopped")
+
+    def _stop_worker(self):
+        self._stop.set()
+        if self._worker and self._worker.is_alive():
+            self._worker.join(timeout=15)
+        self._worker = None
+        self._drain()
+
+    def _drain(self):
+        try:
+            while True:
+                self.frames.get_nowait()
+        except queue.Empty:
+            pass
+
+    def _push(self, pixels):
+        arr = pixels.numpy()
+        for i in range(arr.shape[0]):
+            data = _encode_jpeg(arr[i], self.jpeg_quality)
+            while not self._stop.is_set():
+                try:
+                    self.frames.put(data, timeout=0.25)
+                    break
+                except queue.Full:
+                    continue
+            if self._stop.is_set():
+                return
+
+    def _run(self, world_pixels):
+        blend = None
+        try:
+            self._push(world_pixels)
+            while not self._stop.is_set():
+                if self._pending_prompt is not None:
+                    emb, xf = self._pending_prompt
+                    self._pending_prompt = None
+                    if xf and self._cur_emb is not None:
+                        blend = [self._cur_emb.clone(), emb.to(self.device), 0, xf]
+                    else:
+                        blend = None
+                        self._set_text(emb)
+                st = self.streamer
+                if st.n_world + st.n_frames + st.block_frames > st.max_frames:
+                    raise StreamExhausted(
+                        f"stream reached the {st.max_frames}-latent-frame RoPE "
+                        "ceiling; start a new stream")
+
+                t0 = time.time()
+                z = st.generate_block(generator=self._gen)
+                t_gen = time.time()
+                pixels = self.decoder.decode(z[0].float())
+                self._blocks += 1
+                self._frames_emitted += pixels.shape[0]
+                self._timings = {
+                    "gen_s": round(t_gen - t0, 4),
+                    "decode_s": round(time.time() - t_gen, 4),
+                    "total_s": round(time.time() - t0, 4),
+                }
+                self.status.stats = self.stats()
+
+                if blend:
+                    src, dst, i, n = blend
+                    i += 1
+                    if i >= n:
+                        blend = None
+                        self._set_text(dst)
+                    else:
+                        blend[2] = i
+                        self._set_text(src * (1 - i / n) + dst * (i / n))
+                self._push(pixels)
+        except StreamExhausted as e:
+            self.status.state, self.status.detail = "idle", str(e)
+        except Exception as e:  # surface, never die silently
+            self.status.state, self.status.error = "error", f"{type(e).__name__}: {e}"
+
+
+
+    def stats(self):
+        st = self.streamer
+        return {
+            "frames": self._frames_emitted,
+            "blocks": self._blocks,
+            "latent_frames": (st.n_world + st.n_frames) if st else 0,
+            "latent_frames_max": st.max_frames if st else 0,
+            "seconds": self._frames_emitted / FPS,
+            "kv_mb": (st.cache.memory_bytes() / 1e6) if st else 0.0,
+            **self._timings,
+        }
+
+    def info(self):
+        return {
+            "step": self.step,
+            "weights": self.weights.name,
+            "worlds": [
+                {"idx": i, "prompt": w["prompt"], "generated": w["generated"],
+                 "seconds": 81 / FPS, "gen_seconds": w["seconds"]}
+                for i, w in sorted(self.worlds.items())
+            ],
+            "prompts": self.bank.catalogue(),
+            "fps": FPS,
+            "encoder_loaded": self.encoder.loaded if self.encoder else False,
+            "worldgen": self.worldgen_available,
+            "buffer_frames": self.maxframes,
+            "defaults": DEFAULTS,
+        }

wanstreamer/serve/paths.py

@@ -0,0 +1,22 @@
+"""Where the weights live.
+
+`setup.sh` downloads into the repository directory, so the defaults resolve there
+too and a fresh clone runs without arguments. Each can be overridden by an
+environment variable, or by the matching command-line flag.
+"""
+
+import os
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def _p(env, *default):
+    return str(Path(os.environ.get(env) or ROOT.joinpath(*default)))
+
+
+ASSETS = _p("LIVEWAN_ASSETS", "assets")               # the LiveWan HF repo
+BASE_DIR = _p("LIVEWAN_BASE_DIR", "wan21_13b")        # VAE, umt5, base transformer
+WAN_REPO = _p("LIVEWAN_WAN_REPO", "wan21_repo")       # the Wan2.1 reference code
+WORLDS_DIR = _p("LIVEWAN_WORLDS_DIR", "generated_worlds")
+WEIGHTS = str(Path(ASSETS) / "checkpoints/t14b_b64/latest.pt")

wanstreamer/serve/server.py

@@ -0,0 +1,290 @@
+"""LiveWan browser demo server.
+
+    livewan-serve --port 17070   # or: python -m wanstreamer.serve.server
+
+Serves a single steerable stream: pick a world (or generate a new one from text),
+then steer it with the 96-prompt bank or free text while it runs.
+"""
+
+import argparse
+import asyncio
+import queue
+import threading
+import time
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
+from fastapi.responses import HTMLResponse, Response
+from pydantic import BaseModel
+
+from . import paths
+from .auth import TokenAuth
+from .conditioning import theme_of
+from .engine import WORLDS, Engine
+import cv2
+
+INDEX = Path(__file__).resolve().parent / "web" / "index.html"
+
+# Populated by main(); the endpoints read it rather than module-level globals built at
+# import time, so importing this module has no side effects.
+_S = {"engine": None, "boot": {"messages": [], "done": False, "error": None},
+      "thumbs": {}, "clients": set()}
+
+
+def E() -> Engine:
+    return _S["engine"]
+
+
+def build_parser():
+    p = argparse.ArgumentParser(prog="livewan-serve", description=__doc__,
+                                formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--assets", default=paths.ASSETS, help="the LiveWan weights/data dir")
+    p.add_argument("--wan-repo", default=paths.WAN_REPO, help="Wan2.1 reference checkout")
+    p.add_argument("--base-dir", default=paths.BASE_DIR, help="VAE / umt5 / base model")
+    p.add_argument("--weights", default=None, help="defaults to checkpoints/t14b_b64/latest.pt")
+    p.add_argument("--worlds-dir", default=paths.WORLDS_DIR, help="where generated worlds go")
+    p.add_argument("--host", default="127.0.0.1")
+    p.add_argument("--port", type=int, default=17070)
+    p.add_argument("--jpeg-quality", type=int, default=85)
+    p.add_argument("--buffer-seconds", type=float, default=1.5)
+    p.add_argument("--no-compile", action="store_true",
+                   help="skip the one-off VAE decoder compile (~50 ms/block slower)")
+    p.add_argument("--no-worldgen", action="store_true",
+                   help="disable text->world generation (saves 2.6 GB VRAM)")
+    p.add_argument("--token", default="1234",
+                   help="shared token for ?token=/Bearer/cookie auth; empty string disables")
+    return p
+
+
+
+class StartReq(BaseModel):
+    world: int
+    prompt_idx: int | None = None
+    prompt_text: str | None = None
+    seed: int = 0
+    steps: int = 2
+    window: int = 6
+    latent_norm: float = 1.0
+
+
+class SteerReq(BaseModel):
+    prompt_idx: int | None = None
+    prompt_text: str | None = None
+    crossfade: int = 0
+
+
+class WorldReq(BaseModel):
+    prompt_text: str | None = None
+    prompt_idx: int | None = None
+    steps: int = 30
+    seed: int = 0
+    guide: float = 5.0
+
+
+
+def _make_thumb(wid):
+    """One representative frame per world, cached in memory for the picker."""
+    
+
+    entry = E().worlds.get(int(wid))
+    if not entry:
+        return
+    d = torch.load(entry["path"], map_location="cpu", weights_only=False)
+    frame = d["pixels"][40].numpy()[:, :, ::-1]
+    _S["thumbs"][int(wid)] = cv2.imencode(
+        ".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 82])[1].tobytes()
+
+
+def _boot_load():
+    try:
+        E().load(progress=lambda m: _S["boot"]["messages"].append(m))
+        for w in list(E().worlds):
+            _make_thumb(w)
+        _S["boot"]["done"] = True
+    except Exception as e:
+        _S["boot"]["error"] = f"{type(e).__name__}: {e}"
+
+
+@asynccontextmanager
+async def _lifespan(_app):
+    task = asyncio.create_task(_pump())
+    yield
+    task.cancel()
+
+
+app = FastAPI(title="LiveWan", lifespan=_lifespan)
+
+
+
+@app.get("/", response_class=HTMLResponse)
+def index():
+    return INDEX.read_text()
+
+
+@app.get("/api/boot")
+def boot():
+    b = _S["boot"]
+    return {"done": b["done"], "error": b["error"], "messages": b["messages"][-6:],
+            "detail": E().status.detail if E() else ""}
+
+
+@app.get("/api/info")
+def info():
+    if not _S["boot"]["done"]:
+        raise HTTPException(503, "still loading")
+    d = E().info()
+    d["world_themes"] = {w: theme_of(w) for w in WORLDS}
+    return d
+
+
+@app.get("/api/status")
+def status():
+    e = E()
+    s = e.status
+    return {
+        "state": s.state, "detail": s.detail, "world": s.world,
+        "prompt": s.prompt, "prompt_source": s.prompt_source,
+        "stats": s.stats, "error": s.error, "progress": s.progress,
+        "buffered": e.frames.qsize(),
+        "encoder_loaded": e.encoder.loaded if e.encoder else False,
+    }
+
+
+@app.get("/api/thumb/{world}")
+def thumb(world: int):
+    if world not in _S["thumbs"]:
+        raise HTTPException(404, "no thumbnail")
+    return Response(_S["thumbs"][world], media_type="image/jpeg")
+
+
+@app.post("/api/start")
+def start(r: StartReq):
+    if not _S["boot"]["done"]:
+        raise HTTPException(503, "still loading")
+    try:
+        E().start(r.world, r.prompt_idx, r.prompt_text, seed=r.seed, steps=r.steps,
+                  window=r.window, latent_norm=r.latent_norm)
+    except Exception as e:
+        raise HTTPException(400, f"{type(e).__name__}: {e}")
+    return status()
+
+
+@app.post("/api/steer")
+def steer(r: SteerReq):
+    try:
+        E().steer(r.prompt_idx, r.prompt_text, crossfade=r.crossfade)
+    except Exception as e:
+        raise HTTPException(400, f"{type(e).__name__}: {e}")
+    return status()
+
+
+@app.post("/api/world/new")
+def world_new(r: WorldReq):
+    """Generate a brand-new opening world from text. Slow: ~52 s at 30 steps."""
+    if not _S["boot"]["done"]:
+        raise HTTPException(503, "still loading")
+    try:
+        wid = E().generate_world(r.prompt_text, r.prompt_idx, steps=r.steps,
+                                 seed=r.seed, guide=r.guide)
+        _make_thumb(wid)
+    except Exception as e:
+        raise HTTPException(400, f"{type(e).__name__}: {e}")
+    return {"world": wid, **E().info()}
+
+
+@app.delete("/api/world/{wid}")
+def world_delete(wid: int):
+    try:
+        E().delete_world(wid)
+        _S["thumbs"].pop(wid, None)
+    except Exception as e:
+        raise HTTPException(400, f"{type(e).__name__}: {e}")
+    return E().info()
+
+
+@app.post("/api/stop")
+def stop():
+    E().stop()
+    return status()
+
+
+
+def _next_frame(timeout=0.1):
+    try:
+        return E().frames.get(True, timeout)
+    except queue.Empty:
+        return None
+
+
+async def _pump():
+    """Single reader of the frame queue, broadcasting to every viewer.
+
+    One consumer only. If each socket drained the queue itself, two viewers would
+    each get half the frames.
+    """
+    loop = asyncio.get_running_loop()
+    clients = _S["clients"]
+    last = 0.0
+    while True:
+        if not clients or E() is None:
+            await asyncio.sleep(0.1)
+            continue
+        frame = await loop.run_in_executor(None, _next_frame, 0.1)
+        dead = []
+        if frame is not None:
+            for c in list(clients):
+                try:
+                    await c.send_bytes(frame)
+                except Exception:
+                    dead.append(c)
+        if time.time() - last > 0.5:
+            last = time.time()
+            payload = {"t": "status", **status()}
+            for c in list(clients):
+                try:
+                    await c.send_json(payload)
+                except Exception:
+                    dead.append(c)
+        for c in dead:
+            clients.discard(c)
+
+
+@app.websocket("/ws")
+async def ws(sock: WebSocket):
+    await sock.accept()
+    _S["clients"].add(sock)
+    try:
+        await sock.send_json({"t": "status", **status()})
+        while True:
+            await sock.receive_text()  # client keepalive; ignored
+    except (WebSocketDisconnect, RuntimeError):
+        pass
+    finally:
+        _S["clients"].discard(sock)
+
+
+
+def main(argv=None):
+    args = build_parser().parse_args(argv)
+    _S["engine"] = Engine(
+        args.assets, args.wan_repo, args.base_dir, args.weights,
+        jpeg_quality=args.jpeg_quality, buffer_seconds=args.buffer_seconds,
+        worlds_dir=args.worlds_dir, allow_worldgen=not args.no_worldgen,
+        compile_vae=not args.no_compile,
+    )
+    threading.Thread(target=_boot_load, daemon=True, name="livewan-boot").start()
+
+    served = TokenAuth(app, args.token) if args.token else app
+    if args.token:
+        print(f"token auth on; open  http://{args.host}:{args.port}/?token={args.token}")
+    else:
+        print(f"no auth; open  http://{args.host}:{args.port}/")
+    uvicorn.run(served, host=args.host, port=args.port, log_level="warning",
+                ws_max_size=16 * 2**20)
+
+
+if __name__ == "__main__":
+    main()

wanstreamer/serve/streamdecode.py

@@ -0,0 +1,58 @@
+"""Streaming wrapper around the Wan2.1 VAE decoder.
+
+The offline path (`scripts/demo.py`) decodes the whole latent sequence in one call at
+the end of a run. A live stream cannot: it has to emit pixels every block. Stock
+`WanVAE_.decode` clears its causal-conv feature cache on entry and exit, so calling it
+once per block would restart the temporal convolutions and seam every 3 latent frames.
+The decoder is already causal and already walks the sequence one latent frame at a
+time ( the only thing between it and a continuous stream is that `clear_cache()`).
+This keeps the cache alive across calls instead.
+
+Frame arithmetic: the first latent frame of a stream decodes to 1 pixel frame, every
+later one to 4 (temporal stride 4). So a 21-frame world -> 81 pixel frames, and each
+subsequent 3-frame block -> 12 pixel frames = 750 ms at 16 fps.
+"""
+
+import torch
+
+
+class StreamingVAEDecoder:
+    def __init__(self, vae_path, wan_repo=None, device="cuda", dtype=torch.float16):
+        from wan.modules.vae import WanVAE
+
+        self.wrapper = WanVAE(vae_pth=str(vae_path), dtype=dtype, device=device)
+        self.model = self.wrapper.model
+        self.device, self.dtype = device, dtype
+        self.scale = self.wrapper.scale
+        self.reset()
+
+    def reset(self):
+        self.model.clear_cache()
+
+    @torch.no_grad()
+    def decode(self, z):
+        """z: [C, F, H, W] latents -> uint8 pixels [T, H*8, W*8, 3] (RGB).
+
+        Continues the previous call's temporal context; call `reset()` to cut.
+        """
+        z = z.to(self.device, torch.float32).unsqueeze(0).clamp(-4, 4)
+        mean, inv_std = self.scale
+        z = z / inv_std.view(1, -1, 1, 1, 1).float() + mean.view(1, -1, 1, 1, 1).float()
+
+        with torch.amp.autocast("cuda", dtype=self.dtype):
+            x = self.model.conv2(z)
+            outs = []
+            for i in range(x.shape[2]):
+                self.model._conv_idx = [0]
+                outs.append(self.model.decoder(
+                    x[:, :, i:i + 1], feat_cache=self.model._feat_map,
+                    feat_idx=self.model._conv_idx))
+            out = torch.cat(outs, dim=2)
+
+        out = out.float().clamp_(-1, 1).squeeze(0)  # [3, T, H, W]
+        out = ((out.permute(1, 2, 3, 0) + 1) * 127.5).clamp(0, 255).to(torch.uint8)
+        return out.cpu()
+
+    def memory_bytes(self):
+        return sum(t.numel() * t.element_size()
+                   for t in (self.model._feat_map or []) if torch.is_tensor(t))

wanstreamer/serve/web/index.html

@@ -0,0 +1,822 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width,initial-scale=1">
+<title>LiveWan — steerable streaming video</title>
+<style>
+  :root{
+    --bg:#0c0d10; --panel:#14161b; --panel2:#1b1e25; --line:#282c36;
+    --ink:#e8eaf0; --dim:#9aa1b1; --faint:#666d7e;
+    --accent:#5b8cff; --accent2:#37d39b; --warn:#f2b544; --bad:#ef5f6b;
+    --radius:12px;
+  }
+  *{box-sizing:border-box}
+  body{margin:0;background:var(--bg);color:var(--ink);
+    font:14px/1.5 ui-sans-serif,system-ui,-apple-system,"Segoe UI",Inter,sans-serif}
+  button,input,select{font:inherit;color:inherit}
+  .wrap{max-width:1320px;margin:0 auto;padding:20px}
+
+  header{display:flex;align-items:baseline;gap:14px;flex-wrap:wrap;margin-bottom:18px}
+  h1{font-size:20px;margin:0;letter-spacing:-.01em}
+  h1 span{color:var(--accent)}
+  .sub{color:var(--dim);font-size:13px}
+  .pill{margin-left:auto;display:flex;align-items:center;gap:8px;background:var(--panel);
+    border:1px solid var(--line);border-radius:999px;padding:5px 12px;font-size:12px;color:var(--dim)}
+  .dot{width:8px;height:8px;border-radius:50%;background:var(--faint)}
+  .dot.live{background:var(--accent2);box-shadow:0 0 0 3px rgba(55,211,155,.16)}
+  .dot.load{background:var(--warn);animation:pulse 1s infinite}
+  .dot.err{background:var(--bad)}
+  @keyframes pulse{50%{opacity:.35}}
+
+  .grid{display:grid;grid-template-columns:minmax(0,1fr) 380px;gap:18px}
+  @media(max-width:1040px){.grid{grid-template-columns:1fr}}
+
+  .stage{background:#000;border:1px solid var(--line);border-radius:var(--radius);
+    position:relative;overflow:hidden;aspect-ratio:640/368}
+  .stage canvas{width:100%;height:100%;display:block}
+  .overlay{position:absolute;inset:0;display:flex;flex-direction:column;align-items:center;
+    justify-content:center;gap:0;background:rgba(8,9,12,.96);backdrop-filter:blur(6px);
+    text-align:center;padding:24px;z-index:10}
+  .overlay.hidden{display:none}
+  .overlay:before{content:"";position:absolute;inset:0;
+    background:radial-gradient(circle at 50% 42%,rgba(91,140,255,.14),transparent 62%)}
+  .spinner{width:46px;height:46px;border:3px solid rgba(255,255,255,.09);
+    border-top-color:var(--accent);border-radius:50%;animation:spin .75s linear infinite;
+    margin-bottom:20px;position:relative}
+  @keyframes spin{to{transform:rotate(360deg)}}
+  .ov-title{font-size:20px;font-weight:650;letter-spacing:-.01em;position:relative}
+  .ov-detail{color:var(--dim);font-size:13.5px;max-width:460px;margin-top:8px;position:relative}
+  .ov-log{color:var(--faint);font-size:11.5px;font-family:ui-monospace,SFMono-Regular,Menlo,monospace;
+    margin-top:14px;position:relative}
+  /* staged progress: shows which step of a multi-stage load is running */
+  .ov-steps{display:flex;gap:0;margin-top:20px;position:relative;align-items:center}
+  .ov-step{display:flex;align-items:center;gap:7px;font-size:11.5px;color:var(--faint);
+    padding:0 4px}
+  .ov-step i{width:7px;height:7px;border-radius:50%;background:var(--line);flex:none;
+    transition:background .25s,box-shadow .25s}
+  .ov-step.done{color:var(--dim)}
+  .ov-step.done i{background:var(--accent2)}
+  .ov-step.now{color:var(--ink);font-weight:600}
+  .ov-step.now i{background:var(--accent);box-shadow:0 0 0 4px rgba(91,140,255,.2);
+    animation:pulse 1.1s infinite}
+  .ov-sep{width:22px;height:1px;background:var(--line)}
+  .ov-bar{width:min(330px,72%);height:3px;border-radius:2px;background:rgba(255,255,255,.08);
+    margin-top:20px;overflow:hidden;position:relative}
+  .ov-bar i{display:block;height:100%;border-radius:2px;background:var(--accent);
+    width:30%;animation:slide 1.5s ease-in-out infinite}
+  @keyframes slide{0%{transform:translateX(-110%)}100%{transform:translateX(370%)}}
+  .ov-bar.hide{display:none}
+  .ov-hint{margin-top:16px;font-size:11.5px;color:var(--faint);max-width:420px;
+    position:relative;line-height:1.6}
+
+  .bar{display:flex;gap:14px;flex-wrap:wrap;margin-top:10px;padding:10px 13px;
+    background:var(--panel);border:1px solid var(--line);border-radius:var(--radius);font-size:12px}
+  .bar div{display:flex;gap:6px}
+  .bar b{font-weight:600;font-variant-numeric:tabular-nums}
+  .bar span{color:var(--faint)}
+
+  .now{margin-top:10px;padding:12px 14px;background:var(--panel);border:1px solid var(--line);
+    border-radius:var(--radius)}
+  .now .lbl{font-size:11px;text-transform:uppercase;letter-spacing:.07em;color:var(--faint)}
+  .now .txt{margin-top:3px;font-size:13.5px}
+  .tag{display:inline-block;font-size:10.5px;padding:1px 7px;border-radius:999px;
+    border:1px solid var(--line);color:var(--dim);margin-left:7px;vertical-align:1px}
+
+  .card{background:var(--panel);border:1px solid var(--line);border-radius:var(--radius);
+    padding:14px;margin-bottom:14px}
+  .card h2{margin:0 0 3px;font-size:13px;letter-spacing:.02em}
+  .card p.hint{margin:0 0 11px;color:var(--dim);font-size:12px}
+
+  .worlds{display:grid;grid-template-columns:1fr 1fr;gap:9px}
+  .world{position:relative;border:1.5px solid var(--line);border-radius:9px;overflow:hidden;
+    cursor:pointer;background:#000;padding:0;text-align:left;transition:border-color .15s}
+  .world:hover{border-color:var(--accent)}
+  .world.on{border-color:var(--accent2)}
+  .world img{width:100%;aspect-ratio:640/368;object-fit:cover;display:block;opacity:.85}
+  .world .cap{padding:6px 8px;font-size:11px;color:var(--dim);line-height:1.35}
+  .world .cap span{display:-webkit-box;-webkit-line-clamp:2;-webkit-box-orient:vertical;
+    overflow:hidden;height:2.7em}
+  .wbadge{position:absolute;top:5px;left:5px;font-size:9.5px;padding:2px 6px;border-radius:999px;
+    background:rgba(55,211,155,.9);color:#06281c;font-weight:650;letter-spacing:.02em}
+  .wdel{position:absolute;top:4px;right:4px;width:18px;height:18px;border-radius:50%;
+    background:rgba(0,0,0,.62);color:#fff;font-size:13px;line-height:17px;text-align:center;
+    opacity:0;transition:opacity .15s}
+  .world:hover .wdel{opacity:1}
+  .wdel:hover{background:var(--bad)}
+
+  .tabs{display:flex;gap:4px;margin-bottom:11px;background:var(--panel2);padding:3px;border-radius:9px}
+  .tabs button{flex:1;background:none;border:0;padding:6px;border-radius:7px;color:var(--dim);
+    cursor:pointer;font-size:12.5px}
+  .tabs button.on{background:var(--panel);color:var(--ink);box-shadow:0 1px 3px rgba(0,0,0,.4)}
+
+  input[type=text],input[type=search],input[type=number]{width:100%;background:var(--panel2);
+    border:1px solid var(--line);border-radius:8px;padding:8px 10px;outline:none}
+  input:focus{border-color:var(--accent)}
+  .row{display:flex;gap:8px;align-items:center}
+  .row+.row{margin-top:8px}
+
+  .btn{background:var(--accent);border:0;color:#fff;padding:8px 14px;border-radius:8px;
+    cursor:pointer;font-weight:600;font-size:13px;white-space:nowrap}
+  .btn:disabled{opacity:.4;cursor:not-allowed}
+  .btn.ghost{background:var(--panel2);border:1px solid var(--line);color:var(--ink);font-weight:500}
+
+  .chips{max-height:262px;overflow-y:auto;margin-top:9px;padding-right:3px}
+  .chips::-webkit-scrollbar{width:7px}
+  .chips::-webkit-scrollbar-thumb{background:var(--line);border-radius:4px}
+  .theme{font-size:10.5px;text-transform:uppercase;letter-spacing:.07em;color:var(--faint);
+    margin:9px 0 5px;position:sticky;top:0;background:var(--panel);padding:2px 0}
+  .chip{display:block;width:100%;text-align:left;background:var(--panel2);border:1px solid transparent;
+    border-radius:7px;padding:6px 9px;margin-bottom:4px;cursor:pointer;font-size:12.2px;color:var(--dim);
+    line-height:1.4}
+  .chip:hover{border-color:var(--accent);color:var(--ink)}
+  .chip.on{border-color:var(--accent2);color:var(--ink)}
+  .chip em{font-style:normal;color:var(--faint);margin-right:6px;font-variant-numeric:tabular-nums}
+
+  details{border-top:1px solid var(--line);margin-top:11px;padding-top:10px}
+  summary{cursor:pointer;color:var(--dim);font-size:12px;list-style:none}
+  summary::-webkit-details-marker{display:none}
+  summary:before{content:"▸ ";color:var(--faint)}
+  details[open] summary:before{content:"▾ "}
+  .adv{margin-top:10px;display:grid;grid-template-columns:1fr 1fr;gap:9px}
+  .adv label{font-size:11.5px;color:var(--dim);display:block;margin-bottom:3px}
+
+  .pprow{display:grid;grid-template-columns:74px 1fr 42px;align-items:center;gap:9px;
+    margin-bottom:7px}
+  .pprow label{font-size:11.5px;color:var(--dim)}
+  .pprow b{font-size:11.5px;text-align:right;font-variant-numeric:tabular-nums;color:var(--dim)}
+  .pprow input[type=range]{width:100%;accent-color:var(--accent)}
+
+  .ftfrom{margin-top:11px}
+  .ftlabel{font-size:11px;text-transform:uppercase;letter-spacing:.07em;color:var(--faint)}
+  .ftchips{display:flex;gap:5px;margin-top:6px;flex-wrap:wrap}
+  .ftchips button{flex:1 1 0;min-width:0;background:var(--panel2);border:1px solid var(--line);
+    border-radius:7px;padding:5px 7px;cursor:pointer;font-size:11px;color:var(--dim);
+    white-space:nowrap;overflow:hidden;text-overflow:ellipsis;transition:border-color .15s}
+  .ftchips button:hover{border-color:var(--accent);color:var(--ink)}
+  .ftchips button.on{border-color:var(--accent2);color:var(--ink);background:rgba(55,211,155,.08)}
+
+  .tip{display:inline-flex;align-items:center;justify-content:center;width:14px;height:14px;
+    border-radius:50%;border:1px solid var(--line);color:var(--faint);font-size:9.5px;
+    cursor:help;margin-left:5px;vertical-align:1px;position:relative}
+  .tip:hover .tiptext{opacity:1;visibility:visible}
+  .tiptext{position:absolute;bottom:130%;left:50%;transform:translateX(-50%);width:250px;
+    background:#20242d;border:1px solid var(--line);border-radius:8px;padding:8px 10px;
+    font-size:11.5px;color:var(--ink);line-height:1.5;opacity:0;visibility:hidden;
+    transition:opacity .13s;z-index:20;text-align:left;box-shadow:0 6px 20px rgba(0,0,0,.5);
+    font-weight:400}
+  .note{margin-top:9px;padding:8px 10px;border-radius:8px;font-size:11.5px;line-height:1.5;
+    background:rgba(242,181,68,.09);border:1px solid rgba(242,181,68,.28);color:#f0cd8b}
+  .note.hidden{display:none}
+  .err{background:rgba(239,95,107,.1);border-color:rgba(239,95,107,.3);color:#f5a3aa}
+  footer{margin-top:20px;color:var(--faint);font-size:11.5px;line-height:1.7}
+  footer code{background:var(--panel);padding:1px 5px;border-radius:4px}
+</style>
+</head>
+<body>
+<!-- Post-processing runs in the browser compositor, not on the GPU that is generating
+     frames: an unsharp convolution here, brightness/contrast/saturation as CSS filter
+     functions on the canvas. Per-frame JS cost is zero, so the stream stays at 1x. -->
+<svg width="0" height="0" style="position:absolute" aria-hidden="true"><defs>
+  <filter id="pp-sharpen" x="-5%" y="-5%" width="110%" height="110%"
+          color-interpolation-filters="sRGB">
+    <feConvolveMatrix id="ppKernel" order="3" preserveAlpha="true" divisor="1"
+                      kernelMatrix="0 0 0  0 1 0  0 0 0"/>
+  </filter>
+</defs></svg>
+
+<div class="wrap">
+  <header>
+    <h1>Live<span>Wan</span></h1>
+    <div class="sub">streaming text-to-video · Wan2.1-1.3B student · <span id="ckpt">step 3000</span></div>
+    <div class="pill"><i class="dot" id="dot"></i><span id="pilltext">connecting</span></div>
+  </header>
+
+  <div class="grid">
+    <div>
+      <div class="stage">
+        <canvas id="cv" width="640" height="368"></canvas>
+        <div class="overlay" id="ov">
+          <div class="spinner" id="ovSpin"></div>
+          <div class="ov-title" id="ovTitle">Loading the model</div>
+          <div class="ov-detail" id="ovDetail">This takes about 90 seconds on first start.</div>
+          <div class="ov-steps" id="ovSteps"></div>
+          <div class="ov-bar" id="ovBar"><i></i></div>
+          <div class="ov-log" id="ovLog"></div>
+          <div class="ov-hint" id="ovHint"></div>
+        </div>
+      </div>
+
+      <div class="bar">
+        <div><span>block</span><b id="mBlock">—</b></div>
+        <div><span>speed</span><b id="mSpeed">—</b></div>
+        <div><span>streamed</span><b id="mLen">—</b></div>
+        <div><span>latent frames</span><b id="mLat">—</b></div>
+        <div><span>K/V cache</span><b id="mKv">—</b></div>
+        <div><span>buffer</span><b id="mBuf">—</b></div>
+      </div>
+
+      <div class="now">
+        <div class="lbl">Now conditioning on</div>
+        <div class="txt" id="nowPrompt">—<span class="tag" id="nowTag"></span></div>
+      </div>
+
+      <div class="note hidden" id="note"></div>
+
+      <div class="card" style="margin-top:10px">
+        <h2>What you are looking at</h2>
+        <p class="hint" style="margin:0">
+          A 1.3 B student, distilled from a Wan2.1-14B teacher, generating video
+          <em>continuously</em> instead of as a fixed clip. Every 750 ms it produces
+          another block of 12 frames conditioned on a rolling K/V cache of what it
+          already made — so the picture never restarts.
+        </p>
+        <div style="display:grid;grid-template-columns:1fr 1fr;gap:11px;margin-top:12px">
+          <div style="background:var(--panel2);border-radius:9px;padding:10px 11px">
+            <div style="font-size:11px;color:var(--accent2);font-weight:600;margin-bottom:3px">STEER · continues</div>
+            <div style="font-size:11.5px;color:var(--dim);line-height:1.5">
+              Swaps the text, keeps the cache. The scene morphs in place. Best between
+              related scenes.
+            </div>
+          </div>
+          <div style="background:var(--panel2);border-radius:9px;padding:10px 11px">
+            <div style="font-size:11px;color:var(--accent);font-weight:600;margin-bottom:3px">SCENE · cuts</div>
+            <div style="font-size:11.5px;color:var(--dim);line-height:1.5">
+              Clears the cache and reopens from a cached world. Use it when you want a
+              genuinely different picture.
+            </div>
+          </div>
+        </div>
+      </div>
+    </div>
+
+    <div>
+      <div class="card">
+        <h2>1 · Scene
+          <span class="tip">?<span class="tiptext">A <b>scene</b> is one of four cached worlds — 21 latent
+            frames the base model generated once, so a stream can open instantly instead of
+            waiting ~60 s. Picking one <b>cuts</b>: the stream restarts and the K/V cache is cleared.</span></span>
+        </h2>
+        <p class="hint">Pick a starting world. This cuts to a new stream.</p>
+        <div class="worlds" id="worlds"></div>
+
+        <details id="genBox" style="display:none">
+          <summary>New scene from text</summary>
+          <p class="hint" style="margin:9px 0">
+            Generates a brand-new world with the Wan2.1 base model — the same way the four
+            above were made. This is the one slow step in the project: about a minute, not
+            real time. It is then streamable forever like any other scene.
+          </p>
+          <input type="text" id="genText" placeholder="A lighthouse on a cliff in a storm…">
+          <div class="row" style="margin-top:8px">
+            <label style="font-size:11.5px;color:var(--dim)">Steps
+              <span class="tip">?<span class="tiptext">Base-model denoising steps. 30 takes
+                ~52 s and looks good; 50 takes ~87 s for a little more detail.</span></span>
+            </label>
+            <select id="genSteps" style="background:var(--panel2);border:1px solid var(--line);
+                    border-radius:8px;padding:6px 8px">
+              <option value="20">20 · ~35 s</option>
+              <option value="30" selected>30 · ~52 s</option>
+              <option value="50">50 · ~87 s</option>
+            </select>
+            <input type="number" id="genSeed" value="0" title="seed" style="width:74px">
+            <button class="btn" id="btnGen" style="flex:1">Generate</button>
+          </div>
+        </details>
+      </div>
+
+      <div class="card">
+        <h2>2 · Steer
+          <span class="tip">?<span class="tiptext">Steering swaps the text conditioning but <b>keeps</b> the
+            K/V cache, so the scene carries on and morphs instead of cutting. Nearby scenes blend
+            beautifully; a violent jump (day → night) tends to smear.</span></span>
+        </h2>
+        <p class="hint">Change the text mid-stream. The scene continues.</p>
+
+        <div class="tabs">
+          <button class="on" data-tab="bank">Prompt bank</button>
+          <button data-tab="text">Free text</button>
+        </div>
+
+        <div id="tabBank">
+          <input type="search" id="search" placeholder="Search 96 prompts…" autocomplete="off">
+          <div class="chips" id="chips"></div>
+        </div>
+
+        <div id="tabText" style="display:none">
+          <input type="text" id="freeText" placeholder="A lighthouse in a storm, waves breaking…">
+
+          <div class="ftfrom">
+            <span class="ftlabel">Opens from
+              <span class="tip">?<span class="tiptext">Every stream has to open from one of the four
+                cached worlds — they supply the first 21 latent frames. Your text conditions what
+                happens <b>next</b>, it cannot invent the opening picture, so pick the world that
+                starts closest to what you want.</span></span>
+            </span>
+            <div class="ftchips" id="ftWorlds"></div>
+          </div>
+
+          <div class="row" style="margin-top:10px">
+            <button class="btn" id="btnFree">Steer the live stream</button>
+            <button class="btn ghost" id="btnFreeStart">Open as new scene</button>
+          </div>
+          <div class="note" style="display:block">
+            The opening frames always come from the world above — text steers from there rather than
+            creating the picture, so a prompt far from that world will not teleport.
+            Free text is encoded with umt5-xxl on this machine (11.4 GB, loaded once on first use);
+            the 96 bank prompts are the conditioning the model was trained and measured under.
+          </div>
+        </div>
+
+        <div class="row" style="margin-top:11px">
+          <label style="font-size:11.5px;color:var(--dim);flex:1">
+            Crossfade
+            <span class="tip">?<span class="tiptext">Blend the old and new conditioning over N blocks
+              instead of switching at once. Softens big jumps; 0 is the instant swap the model card
+              describes.</span></span>
+            <input type="range" id="xfade" min="0" max="12" value="0" style="width:100%">
+          </label>
+          <b id="xfadeVal" style="font-size:12px;width:56px;text-align:right">instant</b>
+        </div>
+
+        <details>
+          <summary>Advanced</summary>
+          <div class="adv">
+            <div><label>Denoising steps</label><input type="number" id="steps" value="2" min="1" max="8"></div>
+            <div><label>Seed</label><input type="number" id="seed" value="0"></div>
+            <div><label>K/V window (blocks)</label><input type="number" id="window" value="6" min="1" max="12"></div>
+            <div><label>Latent norm</label><input type="number" id="lnorm" value="1.0" step="0.1" min="0" max="1"></div>
+          </div>
+          <p class="hint" style="margin-top:9px">Defaults match the published command. Steps above 2 cost
+            proportionally more time without fixing hard steers; latent norm 0 lets long streams blow out.</p>
+          <button class="btn ghost" id="btnStop" style="margin-top:4px">Stop stream</button>
+        </details>
+      </div>
+
+      <div class="card">
+        <h2>3 · Look
+          <span class="tip">?<span class="tiptext">Display-only post-processing, applied in
+            your browser as frames arrive. It costs the generator nothing, so the stream
+            stays at 1× — but it also does not change what the model produced. Useful
+            against the softness long streams drift into.</span></span>
+        </h2>
+        <p class="hint">Post-processing on the picture. Costs no GPU time.</p>
+
+        <div class="tabs" id="ppPresets">
+          <button class="on" data-pp="off">Off</button>
+          <button data-pp="subtle">Subtle</button>
+          <button data-pp="crisp">Crisp</button>
+          <button data-pp="punch">Punch</button>
+        </div>
+
+        <div id="ppSliders">
+          <div class="pprow"><label>Sharpen</label>
+            <input type="range" id="ppSharp" min="0" max="150" value="0"><b id="ppSharpV">0</b></div>
+          <div class="pprow"><label>Contrast</label>
+            <input type="range" id="ppCon" min="70" max="150" value="100"><b id="ppConV">1.00</b></div>
+          <div class="pprow"><label>Saturation</label>
+            <input type="range" id="ppSat" min="0" max="200" value="100"><b id="ppSatV">1.00</b></div>
+          <div class="pprow"><label>Brightness</label>
+            <input type="range" id="ppBri" min="70" max="130" value="100"><b id="ppBriV">1.00</b></div>
+        </div>
+
+        <div class="row" style="margin-top:10px">
+          <button class="btn ghost" id="ppCompare" style="flex:1">Hold to compare</button>
+          <button class="btn ghost" id="ppReset">Reset</button>
+        </div>
+      </div>
+    </div>
+  </div>
+
+  <footer>
+    Block = 3 latent frames = 12 pixel frames = 750 ms of 640×368 at 16 fps.
+    Streams are capped at 1024 latent frames (~4.3 min) by <code>WanModel.freqs</code>.
+    Quality holds for roughly a minute; after that sharpness drifts.
+  </footer>
+</div>
+
+<script>
+const $ = s => document.querySelector(s);
+const cv = $('#cv'), ctx = cv.getContext('2d');
+let INFO = null, state = 'boot', curWorld = null, curIdx = null, tab = 'bank';
+
+/* ---------- frame buffer: decouples network jitter from playback ---------- */
+const buf = [];
+const TARGET = 1000/16;
+let playing = false;
+
+function pushFrame(blob){
+  createImageBitmap(blob).then(bm => {
+    buf.push(bm);
+    while (buf.length > 48) buf.shift().close();   // hard cap: never lag by >3 s
+    if (!playing) { playing = true; requestAnimationFrame(tick); }
+  }).catch(()=>{});
+}
+let last = 0;
+function tick(now){
+  if (!last) last = now;
+  let step = TARGET;
+  if (buf.length > 32) step = TARGET * 0.75;      // gently catch up when behind
+  if (now - last >= step) {
+    if (buf.length) {
+      const bm = buf.shift();
+      ctx.drawImage(bm, 0, 0, cv.width, cv.height);
+      bm.close();
+      last = now;
+    } else { last = now; }
+  }
+  requestAnimationFrame(tick);
+}
+
+/* ------------------------------- transport ------------------------------- */
+function connect(){
+  const proto = location.protocol === 'https:' ? 'wss' : 'ws';
+  const ws = new WebSocket(`${proto}://${location.host}/ws`);
+  ws.binaryType = 'blob';
+  ws.onmessage = ev => {
+    if (typeof ev.data === 'string') applyStatus(JSON.parse(ev.data));
+    else pushFrame(ev.data);
+  };
+  ws.onclose = () => setTimeout(connect, 1200);
+  setInterval(() => { if (ws.readyState === 1) ws.send('.'); }, 15000);
+}
+
+/* --------------------------------- boot ---------------------------------- */
+async function boot(){
+  let r;
+  try { r = await (await fetch('/api/boot')).json(); }
+  catch(e){ overlay('Cannot reach the server', String(e), {isErr:true}); return; }
+
+  if (r.error){ overlay('Failed to load', r.error, {isErr:true}); return; }
+  if (!r.done){
+    const msg = (r.messages||[]).slice(-1)[0] || '';
+    overlay('Starting up', 'Loading 2.6 GB of weights and compiling the VAE decoder. About 90 seconds on first start — it only happens once.',
+            {stages:'boot', active:msg, log:msg,
+             hint:'The four scene thumbnails appear as soon as this finishes.'});
+    setTimeout(boot, 600);
+    return;
+  }
+  INFO = await (await fetch('/api/info')).json();
+  $('#ckpt').textContent = `step ${INFO.step}`;
+  buildWorlds(); buildChips(); buildFtWorlds();
+  overlay('Pick a scene', 'Choose one of the four cached worlds on the right to open a stream.',
+          {isErr:true});
+  $('#ovSpin').style.display = 'none';
+  $('#dot').className = 'dot'; $('#pilltext').textContent = 'ready';
+  connect();
+}
+
+/* The two multi-stage waits worth showing progress for. Each entry is matched
+   against the server's own progress message, so the dots track real work. */
+const STAGES = {
+  boot: [
+    ['prompt bank',  /prompt bank/i],
+    ['student',      /student/i],
+    ['VAE',          /VAE/i],
+    ['compile',      /compil/i],
+  ],
+  start: [
+    ['world',        /opening the world/i],
+    ['K\/V cache',   /priming/i],
+    ['first frames', /^$/],
+  ],
+};
+
+function overlay(title, detail, opts={}){
+  const {isErr=false, stages=null, active='', log='', hint='', pct=null} = opts;
+  const ov = $('#ov');
+  ov.classList.remove('hidden');
+  $('#ovTitle').textContent = title;
+  $('#ovDetail').textContent = detail || '';
+  $('#ovSpin').style.display = isErr ? 'none' : '';
+  const bar = $('#ovBar'), fill = bar.firstElementChild;
+  bar.classList.toggle('hide', isErr || (!stages && pct === null));
+  if (pct === null){ fill.style.animation = ''; fill.style.width = '30%'; }
+  else { fill.style.animation = 'none'; fill.style.width = Math.round(pct*100)+'%'; }
+  $('#ovLog').textContent = log;
+  $('#ovHint').textContent = hint;
+
+  const el = $('#ovSteps');
+  el.innerHTML = '';
+  if (!stages) return;
+  const list = STAGES[stages] || [];
+  let idx = list.findIndex(([, re]) => re.test(active));
+  if (idx < 0) idx = 0;
+  list.forEach(([label], i) => {
+    if (i) { const s = document.createElement('div'); s.className='ov-sep'; el.appendChild(s); }
+    const d = document.createElement('div');
+    d.className = 'ov-step ' + (i < idx ? 'done' : i === idx ? 'now' : '');
+    d.innerHTML = `<i></i>${label}`;
+    el.appendChild(d);
+  });
+}
+const hideOverlay = () => $('#ov').classList.add('hidden');
+
+/* --------------------------------- build --------------------------------- */
+function buildWorlds(){
+  const el = $('#worlds'); el.innerHTML = '';
+  INFO.worlds.forEach(w => {
+    const b = document.createElement('button');
+    b.className = 'world' + (w.idx === curWorld ? ' on' : '');
+    b.dataset.world = w.idx;
+    const badge = w.generated
+      ? `<span class="wbadge">generated${w.gen_seconds ? ' · '+Math.round(w.gen_seconds)+'s' : ''}</span>` : '';
+    const del = w.generated ? `<span class="wdel" title="delete this world">×</span>` : '';
+    b.innerHTML = `<img src="/api/thumb/${w.idx}" alt="">${badge}${del}
+                   <div class="cap"><span>${w.prompt}</span></div>`;
+    b.onclick = ev => {
+      if (ev.target.classList.contains('wdel')) { ev.stopPropagation(); delWorld(w.idx); return; }
+      startWorld(w.idx);
+    };
+    el.appendChild(b);
+  });
+  $('#genBox').style.display = INFO.worldgen ? '' : 'none';
+}
+
+async function delWorld(w){
+  if (!confirm('Delete this generated world?')) return;
+  const r = await fetch('/api/world/'+w, {method:'DELETE'});
+  if (r.ok){ INFO = {...INFO, ...(await r.json())}; if (curWorld===w) curWorld=null;
+             buildWorlds(); buildFtWorlds(); }
+}
+
+/* Which world the free-text tab opens from. Kept explicit and visible, because a
+   typed prompt cannot produce the opening frames -- a cached world always does. */
+let ftWorld = null;
+const WORLD_SHORT = {0:'Woman', 44:'Owl', 60:'Waterfall', 82:'Court'};
+
+/* Short label for the picker: a hand-picked word for the four shipped worlds,
+   otherwise the first couple of meaningful words of the prompt it was made from. */
+const STOP = new Set(['a','an','the','of','in','on','at','with','and','to']);
+function worldLabel(w){
+  if (WORLD_SHORT[w.idx]) return WORLD_SHORT[w.idx];
+  const words = (w.prompt||'').split(/\s+/).filter(x => x && !STOP.has(x.toLowerCase()));
+  return words.slice(0,2).join(' ').replace(/[,.]$/,'') || `World ${w.idx}`;
+}
+
+function setFtWorld(w){
+  ftWorld = w;
+  document.querySelectorAll('#ftWorlds button').forEach(b =>
+    b.classList.toggle('on', +b.dataset.world === w));
+}
+
+function buildFtWorlds(){
+  const el = $('#ftWorlds'); el.innerHTML = '';
+  INFO.worlds.forEach(w => {
+    const b = document.createElement('button');
+    b.dataset.world = w.idx;
+    b.textContent = worldLabel(w);
+    b.title = w.prompt;
+    b.onclick = () => setFtWorld(w.idx);
+    el.appendChild(b);
+  });
+  setFtWorld(ftWorld ?? curWorld ?? INFO.worlds[0].idx);
+}
+
+function buildChips(filter=''){
+  const el = $('#chips'); el.innerHTML = '';
+  const f = filter.trim().toLowerCase();
+  let theme = null;
+  INFO.prompts.filter(p => !f || p.text.toLowerCase().includes(f)).forEach(p => {
+    if (p.theme !== theme){
+      theme = p.theme;
+      const h = document.createElement('div');
+      h.className = 'theme'; h.textContent = theme;
+      el.appendChild(h);
+    }
+    const b = document.createElement('button');
+    b.className = 'chip' + (p.idx === curIdx ? ' on' : '');
+    b.innerHTML = `<em>${p.idx}</em>${p.text}`;
+    b.onclick = () => steer({prompt_idx: p.idx}, p);
+    el.appendChild(b);
+  });
+  if (!el.children.length) el.innerHTML = '<div class="theme">no match</div>';
+}
+
+/* -------------------------------- actions -------------------------------- */
+function advanced(){
+  return {
+    seed: +$('#seed').value || 0,
+    steps: +$('#steps').value || 2,
+    window: +$('#window').value || 6,
+    latent_norm: parseFloat($('#lnorm').value),
+  };
+}
+
+async function post(url, body){
+  const r = await fetch(url, {method:'POST', headers:{'Content-Type':'application/json'},
+                             body: JSON.stringify(body)});
+  if (!r.ok){
+    const t = await r.text();
+    note(t.replace(/^\{"detail":"|"\}$/g,''), true);
+    throw new Error(t);
+  }
+  return r.json();
+}
+
+async function startWorld(w, text){
+  // Clicking a scene card opens it on that world's own bank prompt. If free text is
+  // sitting unused in the other tab, say so rather than silently ignoring it.
+  const pending = $('#freeText').value.trim();
+  // Only the four shipped worlds have ids that double as bank indices. A generated
+  // world's id (1000+) is not a prompt, so send no prompt at all and let the server
+  // use the embedding stored with the world.
+  const entry = (INFO?.worlds || []).find(x => x.idx === w);
+  const isGen = !!(entry && entry.generated);
+  curWorld = w; curIdx = (text || isGen) ? null : w;
+  setFtWorld(w);
+  if (!text && pending)
+    note('Opened this scene on its own bank prompt. To use the text you typed, ' +
+         'go to Free text and press "Open as new scene" (or "Steer the live stream").');
+  else note('');
+  document.querySelectorAll('.world').forEach(e =>
+    e.classList.toggle('on', +e.dataset.world === w));
+  const needsT5 = !!text && !INFO.encoder_loaded;
+  overlay('Opening the scene',
+          needsT5 ? 'Loading umt5-xxl (11.4 GB) to encode your prompt, then priming the K/V cache. First free-text prompt only.'
+                  : 'Priming the K/V cache from 21 cached latent frames, then generating the first block.',
+          {stages:'start', active:'opening the world',
+           hint:'Takes about 5 seconds' + (needsT5 ? ', plus ~15 s for the text encoder.' : '.')});
+  buf.splice(0).forEach(b=>b.close());
+  const body = {world: w, ...advanced()};
+  if (text) body.prompt_text = text;
+  else if (!isGen) body.prompt_idx = w;
+  await post('/api/start', body);
+  buildChips($('#search').value);
+}
+
+async function steer(what, p){
+  if (state !== 'streaming'){
+    note('Start a scene first — steering needs a running stream.', true); return;
+  }
+  if (p){ curIdx = p.idx; buildChips($('#search').value); }
+  note('');
+  await post('/api/steer', {...what, crossfade: +$('#xfade').value});
+  const wTheme = INFO.world_themes[curWorld];
+  if (p && wTheme && p.theme !== wTheme)
+    note(`Steering from ${wTheme.toLowerCase()} to ${p.theme.toLowerCase()} is a big jump — expect the picture to smear. A crossfade softens it; cutting to a new scene stays clean.`);
+}
+
+/* --------------------------------- status -------------------------------- */
+function applyStatus(s){
+  state = s.state;
+  const d = $('#dot');
+  d.className = 'dot ' + (s.state==='streaming' ? 'live'
+                        : s.state==='loading' ? 'load'
+                        : s.state==='error' ? 'err' : '');
+  $('#pilltext').textContent = s.state==='streaming' ? 'live'
+                            : s.state==='error' ? 'error' : (s.detail || s.state);
+  // hold the overlay until frames are actually arriving, not merely promised
+  if (s.state === 'streaming' && buf.length) hideOverlay();
+  else if (s.state === 'streaming')
+    overlay('Opening the scene', 'Generating the first block…',
+            {stages:'start', active:''});
+  if (s.state === 'loading')
+    overlay('Opening the scene', s.detail || 'Working…',
+            {stages:'start', active:s.detail || ''});
+  if (s.state === 'generating')
+    overlay('Generating a new world',
+            s.detail || 'Running the Wan2.1 base model…',
+            {pct: s.progress ?? 0,
+             log: s.progress != null ? Math.round(s.progress*100)+'%' : '',
+             hint:'This is the one slow step — the base model, not the student. '
+                + 'Once it exists you can stream from it forever at 1x.'});
+  if (s.state === 'error'){ overlay('Stream failed', s.error, {isErr:true}); note(s.error, true); }
+  if (s.state === 'idle' && s.detail && s.detail !== 'ready' && s.detail !== 'stopped')
+    overlay('Stream ended', s.detail, {isErr:true});
+
+  if (s.prompt){
+    $('#nowPrompt').childNodes[0].nodeValue = s.prompt;
+    $('#nowTag').textContent = s.prompt_source === 'text' ? 'free text' : 'bank';
+  }
+  const st = s.stats || {};
+  if (st.total_s){
+    $('#mBlock').textContent = Math.round(st.total_s*1000)+' ms';
+    $('#mSpeed').textContent = (0.75/st.total_s).toFixed(2)+'×';
+  }
+  if (st.seconds != null) $('#mLen').textContent = st.seconds.toFixed(1)+' s';
+  if (st.latent_frames != null)
+    $('#mLat').textContent = st.latent_frames+' / '+st.latent_frames_max;
+  if (st.kv_mb) $('#mKv').textContent = Math.round(st.kv_mb)+' MB';
+  $('#mBuf').textContent = buf.length+' fr';
+}
+
+function note(msg, isErr){
+  const n = $('#note');
+  n.className = 'note' + (msg ? '' : ' hidden') + (isErr ? ' err' : '');
+  n.textContent = msg;
+}
+
+/* ---------------------------------- wire --------------------------------- */
+$('#search').oninput = e => buildChips(e.target.value);
+$('#xfade').oninput = e => {
+  const v = +e.target.value;
+  $('#xfadeVal').textContent = v ? v+' blocks' : 'instant';
+};
+document.querySelectorAll('.tabs button').forEach(b => b.onclick = () => {
+  tab = b.dataset.tab;
+  document.querySelectorAll('.tabs button').forEach(x => x.classList.toggle('on', x===b));
+  $('#tabBank').style.display = tab==='bank' ? '' : 'none';
+  $('#tabText').style.display = tab==='text' ? '' : 'none';
+});
+$('#btnFree').onclick = async () => {
+  const t = $('#freeText').value.trim();
+  if (!t) return note('Type a prompt first.', true);
+  if (!INFO.encoder_loaded) note('Loading umt5-xxl (11.4 GB) — first free-text prompt only…');
+  curIdx = null;
+  await steer({prompt_text: t});
+  INFO.encoder_loaded = true;
+};
+$('#btnFreeStart').onclick = async () => {
+  const t = $('#freeText').value.trim();
+  if (!t) return note('Type a prompt first.', true);
+  if (!INFO.encoder_loaded) note('Loading umt5-xxl (11.4 GB) — first free-text prompt only…');
+  await startWorld(ftWorld ?? curWorld ?? 60, t);
+  INFO.encoder_loaded = true;
+};
+$('#btnGen').onclick = async () => {
+  const t = $('#genText').value.trim();
+  if (!t) return note('Describe the scene you want first.', true);
+  const steps = +$('#genSteps').value, seed = +$('#genSeed').value || 0;
+  note('');
+  overlay('Generating a new world', 'Encoding the prompt…', {pct:0});
+  try {
+    const r = await post('/api/world/new', {prompt_text:t, steps, seed});
+    INFO = {...INFO, ...r};
+    buildWorlds(); buildFtWorlds();
+    await startWorld(r.world);            // open the thing we just made
+  } catch(e){
+    overlay('Generation failed', String(e), {isErr:true});
+  }
+};
+
+$('#btnStop').onclick = async () => {
+  await post('/api/stop', {});
+  overlay('Stopped', 'Pick a scene to start again.', {isErr:true});
+};
+$('#freeText').addEventListener('keydown', e => { if (e.key==='Enter') $('#btnFree').click(); });
+
+/* ----------------------------- post-processing ---------------------------- */
+const PP_PRESETS = {
+  off:    {sharp:0,   con:100, sat:100, bri:100},
+  subtle: {sharp:35,  con:105, sat:106, bri:100},
+  crisp:  {sharp:80,  con:110, sat:104, bri:100},
+  punch:  {sharp:110, con:120, sat:130, bri:102},
+};
+let ppBypass = false;
+
+function ppApply(){
+  const a = ppBypass ? 0 : +$('#ppSharp').value / 100;
+  const con = ppBypass ? 1 : +$('#ppCon').value / 100;
+  const sat = ppBypass ? 1 : +$('#ppSat').value / 100;
+  const bri = ppBypass ? 1 : +$('#ppBri').value / 100;
+
+  // unsharp: centre 1+4a, four-neighbour -a. Sums to 1, so overall exposure holds.
+  $('#ppKernel').setAttribute('kernelMatrix',
+    `0 ${-a} 0  ${-a} ${1 + 4*a} ${-a}  0 ${-a} 0`);
+
+  const parts = [];
+  if (a > 0) parts.push('url(#pp-sharpen)');
+  if (bri !== 1) parts.push(`brightness(${bri})`);
+  if (con !== 1) parts.push(`contrast(${con})`);
+  if (sat !== 1) parts.push(`saturate(${sat})`);
+  cv.style.filter = parts.join(' ') || 'none';
+
+  $('#ppSharpV').textContent = Math.round(a * 100);
+  $('#ppConV').textContent = con.toFixed(2);
+  $('#ppSatV').textContent = sat.toFixed(2);
+  $('#ppBriV').textContent = bri.toFixed(2);
+}
+
+function ppSet(name){
+  const p = PP_PRESETS[name]; if (!p) return;
+  $('#ppSharp').value = p.sharp; $('#ppCon').value = p.con;
+  $('#ppSat').value = p.sat;     $('#ppBri').value = p.bri;
+  ppApply();
+}
+
+document.querySelectorAll('#ppPresets button').forEach(b => b.onclick = () => {
+  document.querySelectorAll('#ppPresets button').forEach(x => x.classList.toggle('on', x===b));
+  ppSet(b.dataset.pp);
+});
+['ppSharp','ppCon','ppSat','ppBri'].forEach(id => $('#'+id).oninput = () => {
+  // hand-editing a slider means no preset is exactly active any more
+  document.querySelectorAll('#ppPresets button').forEach(x => x.classList.remove('on'));
+  ppApply();
+});
+const ppHold = on => { ppBypass = on; ppApply(); };
+$('#ppCompare').addEventListener('mousedown', () => ppHold(true));
+$('#ppCompare').addEventListener('touchstart', e => { e.preventDefault(); ppHold(true); });
+['mouseup','mouseleave','touchend','touchcancel'].forEach(ev =>
+  $('#ppCompare').addEventListener(ev, () => ppHold(false)));
+$('#ppReset').onclick = () => {
+  document.querySelectorAll('#ppPresets button').forEach(x =>
+    x.classList.toggle('on', x.dataset.pp === 'off'));
+  ppSet('off');
+};
+ppApply();
+
+boot();
+</script>
+</body>
+</html>

wanstreamer/serve/worldgen.py

@@ -0,0 +1,101 @@
+"""Generate new opening worlds from text with the Wan2.1-1.3B base model.
+
+A "world" is just the 21 latent frames a stream is primed from. The four shipped
+ones were produced this way once and cached (`base_seconds` in each file is how long
+it took), so nothing stops any prompt from having one.
+
+This is the base model, not the student: full bidirectional attention over the
+whole sequence, CFG, and a real 30-50 step schedule. It is slow by design (~52 s at
+30 steps on an A100-40GB) and is the only part of the project that is not real time.
+The student then extends the result indefinitely at ~1x.
+"""
+
+import time
+from pathlib import Path
+
+import torch
+
+LATENT_FRAMES = 21  # -> 81 pixel frames, the shipped world length
+
+
+def _sampling_sigmas(steps, shift):
+    """Wan's own shifted inference schedule -- correct for the *undistilled* base
+    model, unlike the uniform few-step spacing the student is rolled out with."""
+    import numpy as np
+
+    s = np.linspace(1, 0, steps + 1)[:steps]
+    return list(shift * s / (1 + (shift - 1) * s)) + [0.0]
+
+
+class WorldGenerator:
+    """Lazy holder for the base model. 5.7 GB on disk, ~2.6 GB resident in bf16."""
+
+    def __init__(self, weights, cfg, device="cuda"):
+        self.weights = Path(weights)
+        self.cfg = cfg
+        self.device = device
+        self._model = None
+
+    @property
+    def loaded(self):
+        return self._model is not None
+
+    @property
+    def available(self):
+        return self.weights.exists()
+
+    def load(self):
+        if self._model is not None:
+            return
+        if not self.available:
+            raise FileNotFoundError(
+                f"{self.weights} not found — world generation needs the Wan2.1 base "
+                "transformer. Re-run setup.sh without SKIP_BASE=1.")
+        from safetensors.torch import load_file
+        from wan.modules.model import WanModel
+
+        c = self.cfg
+        m = WanModel(dim=c.dim, ffn_dim=c.ffn_dim, freq_dim=c.freq_dim,
+                     num_heads=c.num_heads, num_layers=c.num_layers,
+                     window_size=c.window_size, qk_norm=True,
+                     cross_attn_norm=True, eps=1e-6)
+        m.load_state_dict(load_file(str(self.weights)), strict=True)
+        self._model = m.to(self.device, c.param_dtype).eval().requires_grad_(False)
+
+    def unload(self):
+        self._model = None
+        torch.cuda.empty_cache()
+
+    @torch.no_grad()
+    def generate(self, pos, neg, size=(640, 368), steps=30, guide=5.0, shift=5.0,
+                 seed=0, progress=None):
+        """pos/neg: [1, 512, 4096] umt5 embeddings -> latents [16, 21, h, w].
+
+        `progress(i, steps)` fires after every step so the UI can show a real
+        determinate bar rather than a spinner for a minute.
+        """
+        self.load()
+        m = self._model
+        dtype = self.cfg.param_dtype
+        h, w = size[1] // 8, size[0] // 8
+        seq_len = LATENT_FRAMES * (h // 2) * (w // 2)
+
+        t0 = time.time()
+        g = torch.Generator(device=self.device).manual_seed(int(seed))
+        x = torch.randn((16, LATENT_FRAMES, h, w), generator=g,
+                        device=self.device, dtype=torch.float32)
+        ctx_p = [pos.to(self.device, dtype).squeeze(0)]
+        ctx_n = [neg.to(self.device, dtype).squeeze(0)]
+
+        sig = _sampling_sigmas(steps, shift)
+        for i in range(steps):
+            t = torch.full((1,), sig[i] * 1000.0, device=self.device)
+            with torch.amp.autocast("cuda", dtype=dtype):
+                vc = m([x.to(dtype)], t=t, context=ctx_p, seq_len=seq_len)[0].float()
+                vu = m([x.to(dtype)], t=t, context=ctx_n, seq_len=seq_len)[0].float()
+            # classifier free guidance: the student never needs this, the base does
+            x = x + (sig[i + 1] - sig[i]) * (vu + guide * (vc - vu))
+            if progress:
+                progress(i + 1, steps)
+
+        return x, time.time() - t0

wanstreamer/stream.py

@@ -0,0 +1,298 @@
+"""Few-step block-causal streaming: the deployment loop, and the loop the
+self-forcing trainer rolls out through.
+
+One class serves both, because that identity *is* the method. Self-forcing only
+means anything if the context the student is trained on is the context it will
+actually receive -- its own previous output, encoded into the same K/V cache, at
+the same noise levels, with the same number of steps.
+
+Per streaming unit (a block of `block_frames` latent frames):
+
+    z = noise
+    for t, t_next in schedule:            # e.g. 1.0 -> .75 -> .5 -> .25 -> 0
+        x0 = z - t * G(z, t | clean K/V of world + committed events)
+        z  = (1 - t_next) * x0 + t_next * eps            (fresh eps; last step
+                                                          leaves x0 as-is)
+    G(x0, t=0)  -> write this block's clean K/V, commit
+
+The re-run at t=0 is what keeps the cache clean: every committed key/value in
+the prefix was produced from a *clean* latent, which is the condition Stage 1
+trained under. It costs one extra forward per unit and is worth it.
+
+The world's K/V is pinned; eviction under a bounded window drops the oldest
+events only, so memory and per-unit cost are flat in stream length while scene
+and subject identity persist (v0.3's "world + event stream").
+"""
+import time
+
+import torch
+
+from . import blockcausal as bc
+from .core import make_rope_table, make_cache, latent_geometry
+from .graphrunner import GraphedForward, SteadyStateGraphs
+
+
+def few_step_schedule(num_steps, t_max=1.0, shift=1.0):
+    """Flow fractions, highest first.
+
+    shift == 1 gives the uniform spacing DMD-distilled few-step video models use
+    (CausVid / Self-Forcing: 1.0, .75, .5, .25). shift > 1 reproduces Wan's own
+    inference schedule, which is the right one for evaluating a many-step,
+    not-yet-distilled model.
+    """
+    ts = [t_max * (num_steps - i) / num_steps for i in range(num_steps)]
+    if shift == 1.0:
+        return ts
+    return [shift * t / (1 + (shift - 1) * t) for t in ts]
+
+
+class FewStepStreamer:
+    def __init__(self, model, width=640, height=368, max_frames=256,
+                 device='cuda', dtype=torch.bfloat16, window_frames=None,
+                 time_scale=1000.0, cache_frames=None, block_frames=3,
+                 num_steps=4, rope=None, shift=1.0, sampler='renoise',
+                 cuda_graphs=False, noisy_context=False):
+        self.model = model
+        self.device = device
+        self.dtype = dtype
+        self.time_scale = time_scale
+        self.block_frames = block_frames
+        self.schedule = few_step_schedule(num_steps, shift=shift)
+        self.sampler = sampler
+        self.h_lat, self.w_lat, self.hp, self.wp = latent_geometry(width, height)
+        self.S = self.hp * self.wp
+        self.max_frames = max_frames
+        self.window_frames = window_frames
+        self.rope = rope if rope is not None else make_rope_table(
+            model, self.hp, self.wp, max_frames, device)
+        self.cache = make_cache(model, self.S, cache_frames or max_frames,
+                                device, dtype)
+        self.kv = bc.BufferKV(self.cache)
+        self.n_world = 0
+        self.n_frames = 0
+        # Training-only: pushes the event window this far from the pinned world
+        # in RoPE index space, simulating a long stream. At deployment the gap
+        # arises naturally from evicted events, so this stays 0.
+        self.rope_gap = 0
+        # per-channel moment matching towards the world (0 = off); see _renorm
+        self.latent_norm = 0.0
+        self.ref_stats = None
+        self.graphs = SteadyStateGraphs() if cuda_graphs else None
+        # "Noisy context": commit the K/V the LAST denoising step already wrote
+        # instead of re-running the finished block at t=0. That drops the cost
+        # of a unit from N+1 forwards to N -- a third of the latency at N=2,
+        # which is the single largest inference-cost lever left in this loop.
+        # It is NOT a free inference-time swap: every committed key then comes
+        # from a latent at the schedule's lowest noise level rather than a clean
+        # one, and Stage 1 and DMD both trained against a clean prefix. Turning
+        # it on without retraining under the same convention changes the
+        # conditioning the student was distilled for. Exposed here so the
+        # latency claim can be measured before that training is paid for.
+        self.noisy_context = noisy_context
+
+    # ------------------------------------------------------------------ context
+    def set_text(self, ctx_emb, ctx_lens=None):
+        self.ctx = ctx_emb
+        # Resolve "is a cross-attention mask needed?" ONCE, here, and store None
+        # if it is not. Deciding it inside the forward costs a `.item()` on a
+        # device tensor -- a host sync, which is illegal during CUDA graph
+        # capture. Contexts are padded to the full text length (upstream's own
+        # WanModel.forward passes context_lens=None for exactly this reason), so
+        # None is not an approximation, it is the same computation.
+        full = ctx_emb.shape[1]
+        if ctx_lens is not None and int(torch.as_tensor(ctx_lens).min()) < full:
+            self.ctx_lens = ctx_lens
+        else:
+            self.ctx_lens = None
+
+    def _eager(self, z, tbl, e, e0):
+        return bc.block_forward(self.model, z, 0.0, 0, self.rope, self.ctx,
+                                self.ctx_lens, kv=self.kv, dtype=self.dtype,
+                                time_scale=self.time_scale, tbl=tbl, emb=(e, e0))
+
+    def _fwd(self, z, t, t_start):
+        tbl = self.rope.span(t_start, z.shape[2])
+        e, e0 = bc.time_embed(self.model, t, self.device, self.time_scale)
+        if self.graphs is not None and not torch.is_grad_enabled():
+            # Key on the cache length: that is the only thing that changes the
+            # shapes, and a bounded window makes it constant once saturated.
+            g = self.graphs.get(
+                (self.cache.length, z.shape[2]),
+                lambda: GraphedForward(self._eager, tuple(z.shape),
+                                       tuple(tbl.shape), tuple(e.shape),
+                                       tuple(e0.shape), self.device))
+            if g is not None:
+                return g(z, tbl, e, e0)
+        return self._eager(z, tbl, e, e0)
+
+    def _commit(self, n):
+        self.cache.commit(n * self.S)
+        self.n_frames += n
+        if self.window_frames is not None:
+            self.trim_to_window(self.window_frames)
+
+    def trim_to_window(self, window_frames):
+        """Evict the oldest events so only `window_frames` of them remain.
+
+        Deployment calls this from `_commit` after every unit, which is what
+        makes per-unit cost flat in stream length. The self-forcing trainer
+        cannot: it records denoising steps during a rollout and redoes them
+        under gradient afterwards against `buffer[:prefix]`, and an eviction in
+        between silently changes what those tokens are (see
+        StreamingKVCache.evictions). So the trainer leaves `window_frames` None,
+        rolls one iteration's blocks with the buffer growing monotonically, and
+        calls this itself at the iteration boundary -- the one point where no
+        recorded step is still owed its prefix.
+
+        The world stays pinned either way; only events are dropped.
+        """
+        protect = self.n_world * self.S
+        budget = (self.n_world + window_frames) * self.S
+        excess = self.cache.num_tokens - budget
+        if excess > 0:
+            self.cache.evict_front(excess, protect=protect)
+
+    # -------------------------------------------------------------------- world
+    @torch.no_grad()
+    def set_world(self, world_latents):
+        """Prime the cache from clean world latents, attended bidirectionally."""
+        self.cache.reset()
+        self.n_frames = 0
+        z = world_latents.to(self.device).unsqueeze(0) \
+            if world_latents.dim() == 4 else world_latents.to(self.device)
+        F = z.shape[2]
+        self._fwd(z, 0.0, 0)
+        self._commit(F)
+        self.n_world = F
+        # Per-channel moments of the world: the in-distribution reference for
+        # `latent_norm`. Computed over (F, H, W) per channel.
+        w = z[0].float()
+        self.ref_stats = (w.mean(dim=(1, 2, 3), keepdim=True).unsqueeze(0),
+                          w.std(dim=(1, 2, 3), keepdim=True).unsqueeze(0))
+        return F
+
+    def _renorm(self, z, strength):
+        """Pull a finished block's PER-CHANNEL moments back towards the world's.
+
+        Long autoregressive rollouts drift, and here the drift is measurably a
+        per-channel one: the decoded stream develops colour casts and saturated
+        blotches long before it loses structure. A distilled student inherits
+        this from the CFG-guided real score it was matched to -- CFG moves mass
+        towards higher-contrast, more saturated samples, and in a feedback loop
+        (output becomes context becomes output) that bias compounds.
+
+        Matching the first two moments per channel against the world is the
+        cheapest correction that targets exactly that, and it is applied BEFORE
+        the block's K/V is recomputed and committed, so the cached history stays
+        in distribution and the correction cannot itself accumulate. It cannot
+        fix a wrong *direction*, only a wrong per-channel scale and offset.
+        """
+        if strength <= 0 or getattr(self, 'ref_stats', None) is None:
+            return z
+        mu, sd = self.ref_stats
+        zm = z.mean(dim=(2, 3, 4), keepdim=True)
+        zs = z.std(dim=(2, 3, 4), keepdim=True)
+        z_n = (z - zm) / (zs + 1e-5) * sd.to(z.dtype) + mu.to(z.dtype)
+        return (1.0 - strength) * z + strength * z_n
+
+    # ------------------------------------------------------------------- events
+    def generate_block(self, generator=None, record=None, n_frames=None):
+        """Emit one streaming unit. Returns clean latents [1, C, b, H, W]."""
+        b = n_frames or self.block_frames
+        t0 = bc.event_rope_index(self.n_frames, self.n_world, self.rope_gap)
+        z = torch.randn(1, 16, b, self.h_lat, self.w_lat, device=self.device,
+                        dtype=torch.float32, generator=generator)
+        sched = self.schedule
+        for i, t in enumerate(sched):
+            if record is not None:
+                record.append({'block_start': t0, 'step': i, 't': t,
+                               'z': z.detach()})
+            v = self._fwd(z, t, t0).float()
+            x0 = z - t * v
+            tn = sched[i + 1] if i + 1 < len(sched) else 0.0
+            if tn == 0.0:
+                z = x0
+            elif self.sampler == 'renoise':
+                # DMD-distilled few-step sampler: re-noise the x0 prediction to
+                # the next level with FRESH noise. This is the sampler the
+                # student is distilled under, so it is also the one it must be
+                # rolled out with.
+                eps = torch.randn(z.shape, device=z.device, dtype=z.dtype,
+                                  generator=generator)
+                z = (1 - tn) * x0 + tn * eps
+            else:
+                z = z + (tn - t) * v          # deterministic Euler, for the
+                #                               undistilled many-step model
+        z = self._renorm(z, self.latent_norm)
+        if not self.noisy_context:
+            self._fwd(z, 0.0, t0)      # clean K/V for the committed block
+        # else: the scratch region still holds the K/V written by the final
+        # denoising forward, so committing promotes that instead and the extra
+        # pass is skipped entirely. (`_renorm` is then not reflected in the
+        # cache, which is another reason the two options are not interchangeable
+        # without retraining.)
+        self._commit(b)
+        return z
+
+    @torch.no_grad()
+    def stream(self, num_blocks, seed=0, log_every=0, n_frames=None):
+        g = torch.Generator(device=self.device).manual_seed(seed)
+        lats, times = [], []
+        for u in range(num_blocks):
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            z = self.generate_block(generator=g, n_frames=n_frames)
+            torch.cuda.synchronize()
+            dt = time.perf_counter() - t0
+            lats.append(z[0])
+            times.append(dt)
+            if log_every and (u + 1) % log_every == 0:
+                b = n_frames or self.block_frames
+                print(f'  unit {u+1}/{num_blocks}: {dt*1000:7.1f} ms '
+                      f'({dt/b*1000:6.1f} ms/latent frame, '
+                      f'cache {self.cache.num_tokens} tok)', flush=True)
+        return torch.cat(lats, dim=1), times
+
+    # --------------------------------------------------- self-forcing rollout
+    @torch.no_grad()
+    def rollout_record(self, num_blocks, generator=None):
+        """Roll out `num_blocks` units and keep what the DMD step needs.
+
+        Returns (x0_blocks, records). `records[i]` holds, for every denoising
+        step of block i, the exact input latent and timestep, plus the cache
+        length at that point -- enough to recompute any single step under
+        gradient without re-running the rollout.
+        """
+        out, recs = [], []
+        for _ in range(num_blocks):
+            rec = []
+            prefix = self.cache.num_tokens
+            ev = self.cache.evictions
+            z = self.generate_block(generator=generator, record=rec)
+            for r in rec:
+                r['prefix'] = prefix
+                r['evictions'] = ev
+            out.append(z)
+            recs.append(rec)
+        return out, recs
+
+    def prefix_intact(self, rec_step):
+        """Is this recorded step's prefix still literally `buffer[:prefix]`?
+
+        False once anything has been evicted since it was recorded -- the tokens
+        at those indices are now different tokens. The trainer must not take a
+        gradient through a step for which this is False; the forward would
+        succeed and quietly train against the wrong context.
+        """
+        return rec_step.get('evictions', 0) == self.cache.evictions
+
+    def recompute_step(self, rec_step, prefix_kv):
+        """Differentiably redo one recorded denoising step against a detached
+        prefix. Returns that step's x0 prediction."""
+        z, t = rec_step['z'], rec_step['t']
+        v = bc.block_forward(self.model, z, t, rec_step['block_start'],
+                             self.rope, self.ctx, self.ctx_lens, kv=prefix_kv,
+                             dtype=self.dtype, time_scale=self.time_scale,
+                             prefix_upto=rec_step['prefix'],
+                             grad_checkpoint=True)
+        return z - t * v.float()

wanstreamer/text.py

@@ -0,0 +1,45 @@
+"""Prompt encoding with a content-addressed cache.
+
+The T5 encoder (umt5-xxl) runs on CPU here to keep the GPU free, which costs
+~40 s per prompt -- worth caching. The original cache key in run_streaming.py was
+`abs(hash(prompt)) % 10**8`; Python salts `hash()` for str per process
+(PYTHONHASHSEED), so the key changed on every invocation and the cache never hit.
+That is why diag/ accumulated 19 ctx_*.pt files for ~3 distinct prompts. sha1 of
+the text is stable across processes.
+"""
+import gc
+import hashlib
+import os
+
+import torch
+
+
+def prompt_key(prompt):
+    return hashlib.sha1(prompt.encode('utf-8')).hexdigest()[:16]
+
+
+def encode_prompt(prompt, neg_prompt, ckpt_dir, cfg, cache_dir, text_len=512):
+    """Return (pos, neg) embeddings [1, text_len, dim], cached on disk."""
+    os.makedirs(cache_dir, exist_ok=True)
+    path = os.path.join(cache_dir, f'ctx_{prompt_key(prompt)}.pt')
+    if os.path.exists(path):
+        return torch.load(path, map_location='cpu')
+
+    from wan.modules.t5 import T5EncoderModel
+    print(f'Loading T5 (CPU) to encode prompt -> {os.path.basename(path)} ...')
+    t5 = T5EncoderModel(text_len=text_len, dtype=cfg.t5_dtype,
+                        device=torch.device('cpu'),
+                        checkpoint_path=f'{ckpt_dir}/{cfg.t5_checkpoint}',
+                        tokenizer_path=f'{ckpt_dir}/google/umt5-xxl')
+
+    def enc(p):
+        x = t5([p], torch.device('cpu'))[0].float()
+        if x.shape[0] < text_len:
+            x = torch.cat([x, torch.zeros(text_len - x.shape[0], x.shape[1])], 0)
+        return x[:text_len].unsqueeze(0)
+
+    blob = {'pos': enc(prompt), 'neg': enc(neg_prompt), 'prompt': prompt}
+    del t5
+    gc.collect()
+    torch.save(blob, path)
+    return blob