TingtingLiao / streaming-code-world-model

Streaming Code World Model 
https://github.com/user-attachments/assets/37da6610-90eb-4174-9f16-fa5e0b628b1c
 4-step streaming generation; proxy inset bottom-right. 
---
This project implements a streaming version of Code World Model (CWM).
Each 0.71 s chunk is generated causally from the start image, the proxy so far and the previous chunks, with no access
to the future, and a 4-step distilled model renders a chunk in 5.1 s on one H200.
## ✨ Highlights
- **Causal streaming** — one chunk per VAE clip, a sliding 6-chunk attention span, KV cache moved by RoPE rotation.
- **Diffusion forcing with a noisy history** — slows long-horizon drift.
- **4-step DMD self forcing** — 9× faster than 35-step sampling, and more accurate.
- **Fast** — FlashAttention-3 in training and sampling, optional FP8; 5.1 s per chunk on one H200 (not yet real time).
## 📊 Results
Four held-out CARLA clips, 24 chunks (17 s) each; pixel MSE to ground truth (lower is better), one seed.
Inference runs on **one NVIDIA H200 (141 GB)** at 1344×768; peak GPU memory is **123 GB** in bf16 and **108 GB** with
`--fp8` (nvidia-smi peak, GiB); 80 GB GPUs are not supported yet.
| model | steps | mean | 0–4 s | 13–17 s | s / chunk |
|:---|:---:|---:|---:|---:|---:|
| streaming diffusion forcing | 35 | 2231 | 1335 | 3332 | 47.9 |
| ↳ + DMD, released CWM as teacher | 4 | 1741 | **1049** | **2501** | 5.2 |
| **↳ + DMD, streaming model as teacher** | **4** | **1660** | 1070 | 2555 | **5.1** |
## 🔧 Installation
 [code] 
- **`cwm` env** — PyTorch ≥ 2.9 (CUDA 12.8), FlashAttention-3, `transformers`, and from the
 CWM repository its patched musubi-tuner, `cwm_h3_inference` and the
 released models.
- **`carla` env** (data only) — CARLA 0.9.15.
- open-code-world-model and Qwen3-VL-32B-Instruct (captions).
 [code] 
## 🚀 Sampling
Checkpoints (LoRAs for the released MiniMax-H3) are on
Hugging Face: `streaming_df.safetensors` (35 steps) and
`streaming_dmd_4step.safetensors` (4 steps).
 [code] 
`--fp8` runs the large linears in FP8. By default each chunk enters the history from its last denoising step (one
forward fewer, and the noisy history keeps long streams from drifting); `--no-history-from-last-step` uses a clean
forward instead.
## 🏋️ Training
 [code] 
Both trainings resume from the latest checkpoint in their output directory. The Slurm scripts assume 8-GPU nodes.
 How it works 
- **Chunks.** One chunk = one 17-frame VAE clip (5 latents). Chunk *k* is generated inside its span (itself and the
 5 chunks before it), with positions counted from the span's first chunk, so they never grow with the stream.
- **Causal conditioning.** Video chunk *c* sees the caption, the start image, proxy chunks ≤ *c* and earlier video
 chunks of its span. The caption describes the start image only.
- **Diffusion forcing.** Independent per-chunk noise levels, so the model learns not to trust its history blindly.
- **DMD self forcing.** The student streams 8–14 chunks on its own output; a frozen teacher scores the last spa