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