app.py

"""Streaming Code World Model -- causal, chunk-by-chunk world-model video on ZeroGPU.

Reference implementation: https://github.com/TingtingLiao/streaming-code-world-model (rtcwm), sampling exactly as
`python -m rtcwm.eval.open_world generate --steps 4 --rule x0 --shift 5 --kv-reuse --fp8` with the
`streaming_dmd_4step` LoRA on the released MiniMax-H3 Ref2VA model. Hardware adaptations (RTX PRO 6000, sm_120):
  * FlashAttention-3 / flex_attention -> torch SDPA over the same key prefixes (rtcwm.model.attention.dense_df);
  * the blocks' qkv / out / fc1 / fc2 in FP8 exactly as rtcwm.model.fp8 (the authors' --fp8 option);
  * each block's AdaLN projection only ever sees the time embedding of a handful of fixed times (condition rows
    1.0 / 0.999 and the 4 sampling steps), so its outputs are precomputed once at startup instead of keeping the
    26 GB of AdaLN weights on the GPU.
"""
import os
import sys

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
APP_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, APP_DIR)

import spaces  # noqa: E402  (before torch)

import json  # noqa: E402
import math  # noqa: E402
import random  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
from pathlib import Path  # noqa: E402

import av  # noqa: E402
import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402
from PIL import Image  # noqa: E402
from safetensors import safe_open  # noqa: E402
from safetensors.torch import load_file  # noqa: E402

from musubi_tuner.minimax_h3.model import load_h3_transformer  # noqa: E402
from musubi_tuner.minimax_h3.video_vae import load_video_vae  # noqa: E402
from rtcwm.model import sampler as rt_sampler  # noqa: E402
from rtcwm.model.layout import CHUNK_LATENTS, CLIP_FRAMES, TARGET_LATENTS  # noqa: E402
from rtcwm.model.packing import audio_frames  # noqa: E402

H3_REPO, H3_REV = "Comfy-Org/MiniMax-H3", "0543966fbdce5ba05709a8f2031c94bdba629b4a"
LORA_REPO, LORA_FILE = "Luffuly/streaming-code-world-model", "streaming_dmd_4step.safetensors"
STEPS, SHIFT, RULE = 4, 5.0, "x0"        # the released 4-step DMD sampling (model card)
MAX_CHUNKS = 24
FPS = 24
EX_DIR = Path(APP_DIR) / "examples"

# --------------------------------------------------------------------------------------------- examples
SCENES = json.loads((EX_DIR / "scenes.json").read_text())         # name -> {id, chunks, caption}
SILENCE = load_file(str(EX_DIR / "silence_audio.safetensors"))["audio"]
PALETTE = np.array([[0, 0, 0], [135, 206, 235], [30, 90, 200], [150, 120, 80], [90, 90, 90], [60, 160, 60],
                    [200, 120, 80], [220, 200, 60], [230, 40, 40], [200, 80, 200], [40, 220, 220], [250, 250, 250]],
                   dtype=np.uint8)                                   # rtcwm.eval.open_world.PALETTE


def _scene_tensors(name: str) -> dict:
    with safe_open(str(EX_DIR / f"ex{SCENES[name]['id']:02d}.safetensors"), "pt") as f:
        return {k: f.get_tensor(k) for k in f.keys()}


SCENE_DATA = {name: _scene_tensors(name) for name in SCENES}


def load_window(name: str, j: int) -> dict:
    """rtcwm.train.dataset.load_window(..., start=0, with_target=False) on the precomputed caches."""
    d = SCENE_DATA[name]
    duv = d["duv"][:, CHUNK_LATENTS * j:CHUNK_LATENTS * j + TARGET_LATENTS]
    assert duv.shape[1] == TARGET_LATENTS, (name, j, duv.shape)
    ref = torch.nn.functional.pad(duv.unsqueeze(0), (0, 1, 0, 0), value=0.0)
    audio = SILENCE[..., :audio_frames(TARGET_LATENTS)].unsqueeze(0).clone()
    return dict(audio=audio, anchor=d["anchor"].unsqueeze(0), ref=ref, hidden=d["hidden"].unsqueeze(0),
                tags=d["tags"].unsqueeze(0), stream_start=j == 0)


# --------------------------------------------------------------------------------------------- load
from streaming_runtime import prepare_transformer  # noqa: E402

dit = hf_hub_download(H3_REPO, "diffusion_models/minimax_h3_ref2va_bf16.safetensors", revision=H3_REV, local_dir="/tmp/h3")
lora_path = hf_hub_download(LORA_REPO, LORA_FILE)
transformer = load_h3_transformer(dit, device="cpu", dtype=torch.bfloat16, attn_mode="torch")
transformer = prepare_transformer(transformer, load_file(lora_path)).to("cuda")
vae_path = hf_hub_download(H3_REPO, "vae/minimax_h3_video_vae_fp16.safetensors", revision=H3_REV, local_dir="/tmp/h3")
vae = load_video_vae(vae_path, device="cpu", dtype=torch.float16).eval().requires_grad_(False).to("cuda")
print("models ready", flush=True)


# --------------------------------------------------------------------------------------------- helpers
def _write_mp4(frames: np.ndarray, crf: int = 18) -> str:
    path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    with av.open(path, "w") as c:
        s = c.add_stream("libx264", rate=FPS)
        s.width, s.height, s.pix_fmt = frames.shape[2], frames.shape[1], "yuv420p"
        s.options = {"crf": str(crf)}
        for f in frames:
            for p in s.encode(av.VideoFrame.from_ndarray(np.ascontiguousarray(f), format="rgb24")):
                c.mux(p)
        for p in s.encode():
            c.mux(p)
    return path


def _to_uint8(px: torch.Tensor) -> np.ndarray:
    """[3,F,H,W] in [-1,1] -> [F,H,W,3] uint8 (rtcwm.eval.video.decode)."""
    return ((px.float().clamp(-1, 1) + 1) * 127.5).round().byte().permute(1, 2, 3, 0).cpu().numpy()


_PROXY_CACHE = {}


def proxy_video(name: str, chunks: int) -> str:
    """The Semantic-ID half of the proxy the stream follows, coloured as in rtcwm.eval.open_world.proxy_frames."""
    key = (name, int(chunks))
    if key not in _PROXY_CACHE:
        sem = np.load(EX_DIR / f"ex{SCENES[name]['id']:02d}_semantic.npz")["semantic_id"][: int(chunks) * CLIP_FRAMES]
        frames = PALETTE[np.clip(sem, 0, 11)]
        frames = np.stack([np.asarray(Image.fromarray(f).resize((672, 384), Image.NEAREST)) for f in frames])
        _PROXY_CACHE[key] = _write_mp4(frames, crf=23)
    return _PROXY_CACHE[key]


def _duration(scene, chunks=8, seed=0, *args, **kwargs):
    """Measured on ZeroGPU xlarge: ~4-6.5 s per chunk while the 6-chunk span fills, 7.6 s once it is full,
    ~1.7 s per chunk of VAE decode + mp4, ~8 s fixed; +15%."""
    n = max(1, min(MAX_CHUNKS, int(chunks)))
    stream = sum(4.0 + 0.5 * k if k < 6 else 7.7 for k in range(n))
    return int(math.ceil(1.15 * (8 + stream + 1.7 * n)))


# --------------------------------------------------------------------------------------------- inference
@spaces.GPU(duration=_duration, size="xlarge")
def generate(scene: str, chunks: int = 8, seed: int = 0):
    """Stream a world-model video chunk by chunk from a released Code World Model scene.

    Each 17-frame chunk (0.71 s at 24 fps) is generated causally from the start image, its caption, the
    Depth + Semantic-ID proxy up to that chunk and the previous chunks (KV cache, 6-chunk attention span),
    with the 4-step DMD streaming model.

    Args:
        scene: one of the released example scenes.
        chunks: number of 0.71 s chunks to stream (1-24).
        seed: random seed.

    Returns:
        (latest-chunk preview frame, final mp4 path) -- yielded progressively.
    """
    chunks = max(1, min(MAX_CHUNKS, int(chunks)))
    dev = torch.device("cuda")
    gen = torch.Generator().manual_seed(int(seed))
    t_start = time.time()
    torch.cuda.reset_peak_memory_stats()
    events = []

    def on_chunk(k, secs):
        events.append((k, secs))
        print(f"{scene}: chunk {k + 1}/{chunks} {secs:.1f}s", flush=True)

    # stream the chunks (rtcwm.model.sampler.generate_stream), previewing each finished chunk
    import queue
    import threading
    q = queue.Queue()
    result = {}

    def preview_chunk(k, secs):
        on_chunk(k, secs)
        q.put(k)

    def worker():
        try:
            with torch.inference_mode():
                result["lat"] = rt_sampler.generate_stream(
                    transformer, lambda j: load_window(scene, j), chunks, STEPS, gen, dev, RULE, SHIFT,
                    history_noise=0.0, kv_reuse=True, on_chunk=preview_chunk)
        except Exception as e:  # noqa: BLE001
            result["err"] = e
        q.put(None)

    th = threading.Thread(target=worker, daemon=True)
    th.start()
    while True:
        k = q.get()
        if k is None:
            break
        yield gr.update(label=f"Streaming… chunk {k + 1}/{chunks} ({(k + 1) * CLIP_FRAMES / FPS:.1f} s generated)"), None
    th.join()
    if "err" in result:
        raise result["err"]
    lat = result["lat"]
    gen_s = time.time() - t_start
    with torch.inference_mode():
        px = vae.decode(lat.to(dev, torch.float16))[0]
    frames = _to_uint8(px)
    path = _write_mp4(frames)
    peak = torch.cuda.max_memory_allocated() / 2**30
    print(f"{scene}: {chunks} chunks, {len(frames)} frames, stream {gen_s:.1f}s, total {time.time() - t_start:.1f}s, "
          f"peak {peak:.1f} GiB", flush=True)
    yield gr.update(label=f"Done: {chunks} chunks ({len(frames) / FPS:.1f} s) in {time.time() - t_start:.0f} s",
                    value=Image.fromarray(frames[-1]).resize((672, 384))), path


# --------------------------------------------------------------------------------------------- UI
def show_scene(scene: str, chunks: int = 8):
    """Preview a scene's start image, caption and proxy."""
    s = SCENES[scene]
    return str(EX_DIR / f"ex{s['id']:02d}_anchor.png"), s["caption"], proxy_video(scene, chunks)


CSS = """
#col-container { max-width: 1280px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""
NAMES = list(SCENES)

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Streaming Code World Model\n"
            "Causal, chunk-by-chunk video generation from a live control proxy: each 0.71 s chunk is generated from "
            "the start image, its caption, the Depth + Semantic-ID proxy **so far** and the previous chunks — never "
            "the future. 4-step DMD streaming LoRA on MiniMax-H3 Ref2VA + "
            "[Code World Model](https://github.com/buaacyw/code-world-model).  \n"
            "[GitHub](https://github.com/TingtingLiao/streaming-code-world-model) · "
            "[Checkpoints](https://huggingface.co/Luffuly/streaming-code-world-model)"
        )
        with gr.Row():
            with gr.Column(scale=1):
                scene = gr.Dropdown(NAMES, value=NAMES[0], label="Scene (released CWM example)")
                chunks = gr.Slider(1, MAX_CHUNKS, value=8, step=1, label="Chunks to stream (0.71 s each)")
                run = gr.Button("Stream", variant="primary")
                anchor = gr.Image(label="Start image (Picture 1)", interactive=False, height=240)
                proxy = gr.Video(label="Control proxy (Semantic-ID view)", interactive=False, autoplay=True, loop=True,
                                 height=240)
                with gr.Accordion("Advanced", open=False):
                    seed = gr.Number(value=0, precision=0, label="Seed")
                    randomize = gr.Checkbox(value=False, label="Randomize seed")
                    caption = gr.Textbox(label="Static caption (pre-encoded by Qwen3-VL-32B)", lines=6,
                                         interactive=False)
            with gr.Column(scale=2):
                status = gr.Image(label="Stream status", interactive=False, height=200)
                video = gr.Video(label="Generated stream (1344×768, 24 fps)", autoplay=True)

        def _pick_seed(s, r):
            return random.randint(0, 2**31 - 1) if r else int(s)

        scene.change(show_scene, [scene, chunks], [anchor, caption, proxy], api_name=False)
        chunks.release(show_scene, [scene, chunks], [anchor, caption, proxy], api_name=False)
        run.click(_pick_seed, [seed, randomize], seed, api_name=False).then(
            generate, [scene, chunks, seed], [status, video], api_name="generate")
        demo.load(show_scene, [scene, chunks], [anchor, caption, proxy], api_name=False)

        gr.Examples(examples=[[n, 8] for n in NAMES], inputs=[scene, chunks], outputs=[status, video], fn=generate,
                    cache_examples=True, cache_mode="lazy", label="Scenes")

demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)