ForgeWM few-step action-conditioned world model demo

.gitattributes

@@ -36,3 +36,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 demo_images/cave.png filter=lfs diff=lfs merge=lfs -text
 demo_images/forest.png filter=lfs diff=lfs merge=lfs -text
 demo_images/plains.png filter=lfs diff=lfs merge=lfs -text
+examples/cave.png filter=lfs diff=lfs merge=lfs -text
+examples/forest.png filter=lfs diff=lfs merge=lfs -text
+examples/plains.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,58 +1,60 @@
 ---
-title: ForgeWM World Model
+title: ForgeWM
 emoji: 🎮
 colorFrom: yellow
 colorTo: blue
 sdk: gradio
 sdk_version: 6.25.0
 app_file: app.py
+short_description: Few-step action-conditioned Minecraft world model
 python_version: "3.12"
-short_description: Drive a Minecraft world model with keyboard + mouse
 startup_duration_timeout: 1h
 pinned: false
 license: apache-2.0
+models:
+  - ForgeWM/ForgeWM
+  - Skywork/Matrix-Game-2.0
 ---
 
-# 🎮 ForgeWM — a Minecraft world model you can drive
+# ForgeWM — Progressive Causal Training for Few-Step Action-Conditioned Video World Models
 
-Interactive demo of **[ForgeWM](https://huggingface.co/ForgeWM/ForgeWM)**
-(*ForgeWM: Progressive Causal Training for Few-Step Action-Conditioned Video
-World Models*).
+Interactive demo for [`ForgeWM/ForgeWM`](https://huggingface.co/ForgeWM/ForgeWM)
+([paper](https://huggingface.co/papers/2608.14022) ·
+[code](https://github.com/asdfo123/ForgeWM) ·
+[project page](https://asdfo123.github.io/ForgeWM/)).
 
-Give the model **one reference frame** plus an **action track** (keyboard /
-mouse actions, one per second) and it rolls the world forward autoregressively,
-one block-causal chunk at a time.
+Feed the model **one Minecraft frame** plus a short **action script**
+(`forward`, `turn_right`, `look_up`, …) and it rolls the world forward
+autoregressively: a block-causal diffusion transformer denoises 3 latent frames
+(≈1 second of 12 fps video) per step, conditioned on the keyboard/mouse actions
+for that window and a sliding KV cache of the past.
 
-## What's running
+Three released students are selectable:
 
-- **Backbone**: Matrix-Game 2 / Wan2.1-1.3B block-causal DiT (30 layers, dim 1536)
-  with the MG2 action module (keyboard cross-attention + mouse channel condition).
-- **Students**: `ForgeWM-4` (`stage3/model.pt`, 4 denoising steps) and
-  `ForgeWM-1` (`1step/model.pt`, 1 step + the paper's First-Frame Enhancement
-  on chunk 0), distilled with progressive causal training
-  (bidirectional SFT → teacher-forced causal AR → consistency distillation →
-  on-policy DMD).
-- **Rollout regime**: 352×640, 12 fps, chunks of 3 latent frames (12 pixel
-  frames), sliding local attention window of 6 latent frames, `sink_size=0`,
-  `timestep_shift=5.0`, warped denoising schedule — matching
-  `inference.py` / `pipeline/causal_inference.py` upstream.
+| Student | Denoising steps | First-Frame Enhancement |
+|---|---|---|
+| ForgeWM-4 | 4 | – |
+| ForgeWM-2 | 2 | 4-step schedule on block 0 |
+| ForgeWM-1 | 1 | 4-step schedule on block 0 |
 
-## Difference from the reference CLI
+## Fidelity notes
 
-The upstream `inference.py` holds **one** action for the whole clip. Because
-`CausalInferencePipeline.cond_current()` slices the action tensors per raw
-frame, this Space feeds a **time-varying action track** instead: one action per
-one-second causal chunk, so you can compose e.g.
-`forward x3, turn_right x2, forward`. Action semantics (key flags, ±0.10
-camera deltas) are taken verbatim from the reference `make_action()` palette.
+The app mirrors the repo's own `inference.py` and
+`pipeline/causal_inference.py`: 352×640, `num_frame_per_block=3`,
+`local_attn_size=6`, `sink_size=0`, `warp_denoising_step=true`, the published
+`denoising_step_list` / `denoising_step_list_first_chunk` schedules, MG2-style
+conditioning (CLIP ViT-H visual context + 4-channel mask concatenated with the
+first-frame latent), and the exact Minecraft action palette (`CAM_VALUE=0.10`).
 
-## Credits
+Attention runs on PyTorch SDPA rather than FlashAttention-2. In the KV-cached
+inference path the repo calls `attention(q, k, v)` with no causal/window
+arguments, so the two are numerically equivalent — causality is enforced by the
+cache layout, not the kernel.
 
-- ForgeWM — https://github.com/asdfo123/ForgeWM (Apache-2.0). The three
-  reference frames in `demo_images/` are the ForgeWM repo's own demo images.
-- Matrix-Game 2.0 — https://huggingface.co/Skywork/Matrix-Game-2.0
-- Wan2.1 — https://github.com/Wan-Video/Wan2.1
-- GameFactory — the Minecraft action-conditioned training data source.
+## Credits
 
-Model code under `wan/`, `utils/`, `pipeline/`, `configs/` is vendored from the
-ForgeWM repository; see `LICENSE` and `NOTICE`.
+- Model and reference frames (`examples/*.png`): the
+  [ForgeWM repository](https://github.com/asdfo123/ForgeWM), Apache-2.0.
+- Base weights: [Skywork/Matrix-Game-2.0](https://huggingface.co/Skywork/Matrix-Game-2.0),
+  distributed under Skywork's own terms.
+- Vendored `wan/`, `pipeline/`, `utils/` are Apache-2.0; see `NOTICE`.

app.py

@@ -1,549 +1,483 @@
-"""ForgeWM — a playable Minecraft world model you drive with keyboard + mouse.
+"""ForgeWM — few-step, action-conditioned Minecraft world model.
 
-Faithful port of the reference implementation in
-https://github.com/asdfo123/ForgeWM (`inference.py` + `pipeline/causal_inference.py`)
-onto ZeroGPU, extended so that the action condition can change over the rollout
-instead of being a single action held for the whole clip.
+Gradio demo for `ForgeWM/ForgeWM` (paper: "ForgeWM: Progressive Causal Training
+for Few-Step Action-Conditioned Video World Models").
+
+Faithful to the repo's own `inference.py` / `pipeline/causal_inference.py`:
+same 352x640 resolution, 3-latent-frame causal blocks, sliding window of 6
+latent frames, the released `warp_denoising_step` schedules, and First-Frame
+Enhancement for the 1-/2-step students.
 """
 
 import os
 
-# NOTE: expandable_segments is deliberately NOT enabled. On this Space it makes
-# the very first attention allocation inside the ZeroGPU (MIG-virtualised)
-# worker fail with `NVML_SUCCESS == r INTERNAL ASSERT FAILED`; the plain caching
-# allocator handles the ~12 GB working set fine.
-os.environ.pop("PYTORCH_CUDA_ALLOC_CONF", None)
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 
-import spaces  # noqa: E402  (must precede torch / CUDA-touching imports)
+import spaces  # noqa: E402  (must precede torch)
 
-import random  # noqa: E402
-import re  # noqa: E402
+import gc  # noqa: E402
+import shutil  # noqa: E402
 import tempfile  # noqa: E402
-import threading  # noqa: E402
 import time  # noqa: E402
+from typing import List, Tuple  # noqa: E402
 
 import gradio as gr  # noqa: E402
 import imageio.v2 as imageio  # noqa: E402
 import numpy as np  # noqa: E402
 import torch  # noqa: E402
-from huggingface_hub import hf_hub_download, snapshot_download  # noqa: E402
+import torch.nn.functional as F  # noqa: E402
+from huggingface_hub import hf_hub_download  # noqa: E402
 from omegaconf import OmegaConf  # noqa: E402
 from PIL import Image  # noqa: E402
-from torchvision.transforms import (  # noqa: E402
-    Compose,
-    InterpolationMode,
-    Normalize,
-    Resize,
-    ToTensor,
-)
 
-# ─── Weights ──────────────────────────────────────────────────────────────────
+from pipeline import CausalInferencePipeline
+from utils.wan_wrapper import WanVAEWrapper
 
-MG2_REPO = "Skywork/Matrix-Game-2.0"
-FORGEWM_REPO = "ForgeWM/ForgeWM"
+# ────────────────────────────── constants ──────────────────────────────
+HEIGHT, WIDTH = 352, 640
+FPS = 12
+VAE_TCR = 4                 # VAE temporal compression ratio
+FRAMES_PER_BLOCK = 3        # latent frames per causal block
+CAM_VALUE = 0.10            # camera delta magnitude, from inference.py
+
+MINECRAFT_ACTIONS = [
+    "forward", "back", "left", "right",
+    "turn_right", "turn_left", "look_up", "look_down",
+    "forward_turn_right", "random", "no_action",
+]
 
-BASE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "mg2_base")
-os.makedirs(BASE_DIR, exist_ok=True)
+BASE_REPO = "Skywork/Matrix-Game-2.0"
+FORGEWM_REPO = "ForgeWM/ForgeWM"
 
+VARIANTS = {
+    "ForgeWM-4 · 4 steps": ("configs/stage3_dmd.yaml", "stage3/model.pt"),
+    "ForgeWM-2 · 2 steps": ("configs/stage3_dmd_2step.yaml", "2step/model.pt"),
+    "ForgeWM-1 · 1 step": ("configs/stage3_dmd_1step.yaml", "1step/model.pt"),
+}
+DEFAULT_VARIANT = "ForgeWM-4 · 4 steps"
+
+# Measured ballpark on the ZeroGPU Blackwell card; used only to size the
+# @spaces.GPU reservation.
+SECONDS_PER_CHUNK = {
+    "ForgeWM-4 · 4 steps": 1.5,
+    "ForgeWM-2 · 2 steps": 1.0,
+    "ForgeWM-1 · 1 step": 0.8,
+}
+FIXED_OVERHEAD = 22.0       # CLIP + VAE encode/decode + mp4 mux
 
-def _link(src: str, name: str) -> str:
-    dst = os.path.join(BASE_DIR, name)
-    if os.path.islink(dst):
-        os.remove(dst)
-    if not os.path.exists(dst):
-        os.symlink(src, dst)
-    return dst
-
-
-print("[boot] fetching Matrix-Game-2.0 base weights …", flush=True)
-_link(hf_hub_download(MG2_REPO, "base_model/base_config.json"), "base_config.json")
-_link(
-    hf_hub_download(MG2_REPO, "base_model/diffusion_pytorch_model.safetensors"),
-    "diffusion_pytorch_model.safetensors",
-)
-VAE_PATH = hf_hub_download(MG2_REPO, "Wan2.1_VAE.pth")
-CLIP_PATH = hf_hub_download(
-    MG2_REPO, "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"
-)
-TOKENIZER_DIR = snapshot_download(MG2_REPO, allow_patterns=["xlm-roberta-large/*"])
-TOKENIZER_DIR = os.path.join(TOKENIZER_DIR, "xlm-roberta-large")
-
-print("[boot] fetching ForgeWM student checkpoints …", flush=True)
-CKPT_4STEP = hf_hub_download(FORGEWM_REPO, "stage3/model.pt")
-CKPT_1STEP = hf_hub_download(FORGEWM_REPO, "1step/model.pt")
-
-# ─── Model ────────────────────────────────────────────────────────────────────
-
-from pipeline import CausalInferencePipeline  # noqa: E402
-from utils.wan_wrapper import WanVAEWrapper  # noqa: E402
-
-DEVICE = "cuda"
-DTYPE = torch.bfloat16
-HEIGHT, WIDTH = 352, 640
-FPS = 12
-NUM_FRAME_PER_BLOCK = 3          # latent frames per causal chunk
-RAW_PER_BLOCK = 12               # pixel frames per causal chunk (VAE tcr = 4)
-CAM_VALUE = 0.10                 # reference camera delta magnitude
-MAX_CHUNKS = 12
+CKPT_DIR = os.path.join(os.getcwd(), "ckpts", "MG2-base")
 
-torch.set_grad_enabled(False)
 
+# ─────────────────────────── weight preparation ───────────────────────────
+def _link(src: str, dst: str) -> None:
+    if os.path.lexists(dst):
+        os.remove(dst)
+    os.symlink(src, dst)
 
-def _load_config(name: str):
-    cfg = OmegaConf.merge(
-        OmegaConf.load("configs/default.yaml"), OmegaConf.load(f"configs/{name}")
-    )
-    cfg.model_kwargs.model_name = BASE_DIR
-    return cfg
 
+def _purge(path: str) -> None:
+    """Drop a hub blob from local disk once its tensors are in memory."""
+    try:
+        real = os.path.realpath(path)
+        if os.path.isfile(real):
+            os.remove(real)
+        if os.path.islink(path):
+            os.remove(path)
+    except OSError as exc:  # pragma: no cover
+        print(f"[disk] could not purge {path}: {exc}")
+
+
+def _disk() -> str:
+    total, used, free = shutil.disk_usage("/")
+    return f"disk {used / 2**30:.1f}G used / {free / 2**30:.1f}G free"
+
+
+os.makedirs(os.path.join(CKPT_DIR, "xlm-roberta-large"), exist_ok=True)
+
+print(f"[setup] fetching Matrix-Game-2.0 base weights … ({_disk()})", flush=True)
+_dit = hf_hub_download(BASE_REPO, "base_model/diffusion_pytorch_model.safetensors")
+_link(_dit, os.path.join(CKPT_DIR, "diffusion_pytorch_model.safetensors"))
+_link(hf_hub_download(BASE_REPO, "base_model/base_config.json"),
+      os.path.join(CKPT_DIR, "base_config.json"))
+_vae_path = hf_hub_download(BASE_REPO, "Wan2.1_VAE.pth")
+_clip_path = hf_hub_download(
+    BASE_REPO, "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth")
+for _tok in ("sentencepiece.bpe.model", "special_tokens_map.json",
+             "tokenizer.json", "tokenizer_config.json"):
+    _link(hf_hub_download(BASE_REPO, f"xlm-roberta-large/{_tok}"),
+          os.path.join(CKPT_DIR, "xlm-roberta-large", _tok))
+print(f"[setup] base weights ready ({_disk()})", flush=True)
+
+# ─────────────────────────────── models ───────────────────────────────
+torch.set_grad_enabled(False)
+DTYPE = torch.bfloat16
 
-print("[boot] building VAE + CLIP …", flush=True)
+print("[setup] building VAE + CLIP …", flush=True)
 VAE = WanVAEWrapper(
-    vae_path=VAE_PATH, clip_checkpoint_path=CLIP_PATH, clip_tokenizer_path=TOKENIZER_DIR
-)
-VAE = VAE.to(device=DEVICE, dtype=DTYPE).eval()
-VAE.clip.model = VAE.clip.model.to(DEVICE)
-
-
-def _build_pipeline(config_name: str, checkpoint_path: str):
-    cfg = _load_config(config_name)
-    pipe = CausalInferencePipeline(cfg, device=DEVICE, vae=VAE)
-    state = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
-    gen_sd = state.get("generator", state.get("generator_ema", state)) if isinstance(
-        state, dict
-    ) else state
-    fixed = {
-        k.replace("._fsdp_wrapped_module.", ".").replace(
-            "._checkpoint_wrapped_module.", "."
-        ): v
-        for k, v in gen_sd.items()
-    }
+    vae_path=_vae_path,
+    clip_checkpoint_path=_clip_path,
+    clip_tokenizer_path=os.path.join(CKPT_DIR, "xlm-roberta-large"),
+).eval()
+VAE = VAE.to(device="cuda", dtype=DTYPE)
+# `WanVAEWrapper.clip` is a plain object, not a submodule, so `.to()` above does
+# not reach it.  Move it explicitly (in fp32, as the repo does) so ZeroGPU packs
+# it with everything else instead of paying a 4.8 GB host copy per request.
+VAE.clip.model = VAE.clip.model.to("cuda")
+gc.collect()
+print(f"[setup] VAE + CLIP on GPU ({_disk()})", flush=True)
+
+
+def _build_variant(name: str) -> CausalInferencePipeline:
+    cfg_path, ckpt_file = VARIANTS[name]
+    config = OmegaConf.merge(OmegaConf.load("configs/default.yaml"),
+                             OmegaConf.load(cfg_path))
+    config.model_kwargs.model_name = CKPT_DIR
+
+    pipe = CausalInferencePipeline(config, device=torch.device("cuda"), vae=VAE)
+
+    ckpt = hf_hub_download(FORGEWM_REPO, ckpt_file)
+    try:
+        state = torch.load(ckpt, map_location="cpu", weights_only=True)
+    except Exception:
+        state = torch.load(ckpt, map_location="cpu", weights_only=False)
+    gen_sd = state.get("generator", state.get("generator_ema", state)) \
+        if isinstance(state, dict) else state
+    fixed = {k.replace("._fsdp_wrapped_module.", ".")
+              .replace("._checkpoint_wrapped_module.", "."): v
+             for k, v in gen_sd.items()}
     missing, unexpected = pipe.generator.load_state_dict(fixed, strict=False)
-    print(
-        f"[boot] {config_name}: missing={len(missing)} unexpected={len(unexpected)}",
-        flush=True,
-    )
+    print(f"[{name}] loaded: missing={len(missing)} unexpected={len(unexpected)}",
+          flush=True)
     del state, gen_sd, fixed
-    pipe.generator = pipe.generator.to(device=DEVICE, dtype=DTYPE).eval()
-    return pipe
-
-
-print("[boot] building ForgeWM-4 (stage3) …", flush=True)
-PIPE_4 = _build_pipeline("stage3_dmd.yaml", CKPT_4STEP)
-print("[boot] building ForgeWM-1 (1step) …", flush=True)
-PIPE_1 = _build_pipeline("stage3_dmd_1step.yaml", CKPT_1STEP)
-
-MODELS = {
-    "ForgeWM-4 · 4-step (best quality)": PIPE_4,
-    "ForgeWM-1 · 1-step (fastest)": PIPE_1,
-}
-DEFAULT_MODEL = "ForgeWM-4 · 4-step (best quality)"
-
-_GPU_LOCK = threading.Lock()
-
-# ─── Action palette (Minecraft / MG2 schema: 2-D mouse, 6 key flags) ──────────
-
-ACTIONS = [
-    "forward",
-    "back",
-    "left",
-    "right",
-    "turn_left",
-    "turn_right",
-    "look_up",
-    "look_down",
-    "forward_turn_right",
-    "random",
-    "no_action",
-]
-ALIASES = {
-    "w": "forward",
-    "s": "back",
-    "a": "left",
-    "d": "right",
-    "idle": "no_action",
-    "none": "no_action",
-    "stay": "no_action",
-    "turnleft": "turn_left",
-    "turnright": "turn_right",
-    "lookup": "look_up",
-    "lookdown": "look_down",
-}
-
-BUTTONS = [
-    ("⬆️ Forward (W)", "forward"),
-    ("⬇️ Back (S)", "back"),
-    ("⬅️ Strafe left (A)", "left"),
-    ("➡️ Strafe right (D)", "right"),
-    ("↺ Turn left", "turn_left"),
-    ("↻ Turn right", "turn_right"),
-    ("🔼 Look up", "look_up"),
-    ("🔽 Look down", "look_down"),
-    ("⬆️↻ Forward + turn right", "forward_turn_right"),
-    ("⏸️ Stand still", "no_action"),
-    ("🎲 Random", "random"),
-]
+    gc.collect()
+    _purge(ckpt)
 
+    pipe.generator = pipe.generator.to(device="cuda", dtype=DTYPE).eval()
+    pipe.vae = VAE
+    gc.collect()
+    print(f"[{name}] ready ({_disk()})", flush=True)
+    return pipe
 
-_REPEAT_ONLY = re.compile(r"^[x*×]?(\d+)$")
-_REPEAT_SUFFIX = re.compile(r"^(.*?)[\s_\-]*[x*×](\d+)$")
 
+PIPELINES = {}
+for _name in VARIANTS:
+    try:
+        PIPELINES[_name] = _build_variant(_name)
+    except Exception as exc:  # keep the Space usable if one student fails
+        print(f"[setup] FAILED to load {_name}: {exc!r}", flush=True)
 
-def parse_script(script: str) -> list:
-    """Parse an action script into one action per one-second chunk.
+if not PIPELINES:
+    raise RuntimeError("No ForgeWM student could be loaded.")
+if DEFAULT_VARIANT not in PIPELINES:
+    DEFAULT_VARIANT = next(iter(PIPELINES))
 
-    Accepts comma / whitespace separated action names, with an optional
-    ``xN`` repeat, attached or detached: ``forward x3, turn_right x2, forward``.
-    """
-    tokens = [t for t in re.split(r"[,\s]+", (script or "").strip().lower()) if t]
-    actions = []
-    for tok in tokens:
-        m = _REPEAT_ONLY.match(tok)
-        if m:
-            if not actions:
-                raise gr.Error(f"'{tok}' has no action to repeat.")
-            actions.extend([actions[-1]] * max(0, int(m.group(1)) - 1))
-            continue
-        repeat = 1
-        m = _REPEAT_SUFFIX.match(tok)
-        if m and m.group(1):
-            tok, repeat = m.group(1), int(m.group(2))
-        tok = ALIASES.get(tok, tok)
-        if tok not in ACTIONS:
-            raise gr.Error(
-                f"Unknown action '{tok}'. Available actions: {', '.join(ACTIONS)}."
-            )
-        actions.extend([tok] * max(1, repeat))
-    if not actions:
-        raise gr.Error(
-            "The action track is empty — click the action buttons to build a "
-            "sequence, or type e.g. 'forward x3, turn_right x2'."
-        )
-    return actions[:MAX_CHUNKS]
+# The base DiT safetensors were only needed to instantiate the architecture;
+# every parameter has since been overwritten by the ForgeWM checkpoints.
+_purge(os.path.join(CKPT_DIR, "diffusion_pytorch_model.safetensors"))
+print(f"[setup] {len(PIPELINES)} variant(s) ready ({_disk()})", flush=True)
 
 
-def _chunk_of_raw_frame(r: int) -> int:
-    """Map a raw (pixel) frame index onto its causal chunk index.
+# ──────────────────────────── action scripting ────────────────────────────
+def make_action(action_type: str, num_raw_frames: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Minecraft action palette (mouse_dim_in=2), verbatim from inference.py."""
+    mouse = torch.zeros(1, num_raw_frames, 2)
+    keyboard = torch.zeros(1, num_raw_frames, 6)
 
-    Chunk k covers latent frames 3k..3k+2, which the pipeline feeds with raw
-    action frames up to ``1 + 4 * (3k + 2)``. Chunk 0 therefore owns raw frames
-    0..8 and every later chunk owns 12 raw frames.
-    """
-    return 0 if r < 9 else (r + 3) // RAW_PER_BLOCK
+    if action_type == "forward":
+        keyboard[:, :, 0] = 1.0
+    elif action_type == "back":
+        keyboard[:, :, 1] = 1.0
+    elif action_type == "left":
+        keyboard[:, :, 2] = 1.0
+    elif action_type == "right":
+        keyboard[:, :, 3] = 1.0
+    elif action_type == "turn_right":
+        mouse[:, :, 1] = CAM_VALUE
+    elif action_type == "turn_left":
+        mouse[:, :, 1] = -CAM_VALUE
+    elif action_type == "look_up":
+        mouse[:, :, 0] = CAM_VALUE
+    elif action_type == "look_down":
+        mouse[:, :, 0] = -CAM_VALUE
+    elif action_type == "forward_turn_right":
+        keyboard[:, :, 0] = 1.0
+        mouse[:, :, 1] = CAM_VALUE
+    elif action_type == "random":
+        torch.manual_seed(42)
+        mouse = (torch.rand(1, num_raw_frames, 2) - 0.5) * (2 * CAM_VALUE)
+        keyboard[:, :, :4] = (torch.rand(1, num_raw_frames, 4) > 0.5).float()
+    elif action_type == "no_action":
+        pass
+    else:
+        raise gr.Error(f"Unknown action '{action_type}'.")
+    return mouse, keyboard
 
 
-def build_action_tensors(actions, num_raw_frames, cam_value, seed):
-    """Build [1, T, 2] mouse and [1, T, 6] keyboard tensors from the script.
+def _chunk_start(chunk_index: int) -> int:
+    """First raw (pixel) frame owned by causal block `chunk_index`.
 
-    Semantics are exactly the reference `make_action()` palette, applied
-    per-chunk instead of once for the whole clip.
+    Latent frame 0 maps to raw frame 0; latent frame f>=1 maps to raw frames
+    4f-3 … 4f.  A block spans 3 latent frames, so block c covers raw frames
+    [12c-3, 12c+9) — and [0, 9) for c == 0.
     """
-    mouse = torch.zeros(1, num_raw_frames, 2)
-    keyboard = torch.zeros(1, num_raw_frames, 6)
-    rng = torch.Generator().manual_seed(int(seed))
-    for r in range(num_raw_frames):
-        a = actions[min(_chunk_of_raw_frame(r), len(actions) - 1)]
-        if a == "forward":
-            keyboard[0, r, 0] = 1.0
-        elif a == "back":
-            keyboard[0, r, 1] = 1.0
-        elif a == "left":
-            keyboard[0, r, 2] = 1.0
-        elif a == "right":
-            keyboard[0, r, 3] = 1.0
-        elif a == "turn_right":
-            mouse[0, r, 1] = cam_value
-        elif a == "turn_left":
-            mouse[0, r, 1] = -cam_value
-        elif a == "look_up":
-            mouse[0, r, 0] = cam_value
-        elif a == "look_down":
-            mouse[0, r, 0] = -cam_value
-        elif a == "forward_turn_right":
-            keyboard[0, r, 0] = 1.0
-            mouse[0, r, 1] = cam_value
-        elif a == "random":
-            mouse[0, r] = (torch.rand(2, generator=rng) - 0.5) * (2 * cam_value)
-            keyboard[0, r, :4] = (torch.rand(4, generator=rng) > 0.5).float()
-        # "no_action" leaves both channels at zero
+    return 0 if chunk_index == 0 else 12 * chunk_index - 3
+
+
+def build_action_track(segments: List[Tuple[str, int]]) -> Tuple[torch.Tensor, torch.Tensor]:
+    total_chunks = sum(n for _, n in segments)
+    num_raw = _chunk_start(total_chunks)
+    mouse = torch.zeros(1, num_raw, 2)
+    keyboard = torch.zeros(1, num_raw, 6)
+    chunk = 0
+    for action, count in segments:
+        for _ in range(count):
+            lo, hi = _chunk_start(chunk), _chunk_start(chunk + 1)
+            m, k = make_action(action, hi - lo)
+            mouse[:, lo:hi] = m
+            keyboard[:, lo:hi] = k
+            chunk += 1
     return mouse, keyboard
 
 
-_TRANSFORM = Compose(
-    [
-        Resize((HEIGHT, WIDTH), interpolation=InterpolationMode.BILINEAR),
-        ToTensor(),
-        Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-    ]
-)
-
+# ───────────────────────────── preprocessing ─────────────────────────────
+def load_reference_frame(image_path: str) -> torch.Tensor:
+    """Aspect-preserving resize + centre crop to 352x640, scaled to [-1, 1]."""
+    image = Image.open(image_path).convert("RGB")
+    arr = torch.from_numpy(np.asarray(image)).permute(2, 0, 1)[None].float() / 255.0
+    _, _, h, w = arr.shape
+    if h / w > HEIGHT / WIDTH:
+        new_w, new_h = WIDTH, max(HEIGHT, int(round(h * WIDTH / w)))
+    else:
+        new_h, new_w = HEIGHT, max(WIDTH, int(round(w * HEIGHT / h)))
+    arr = F.interpolate(arr, size=(new_h, new_w), mode="bilinear", align_corners=False)
+    top, left = (new_h - HEIGHT) // 2, (new_w - WIDTH) // 2
+    arr = arr[:, :, top:top + HEIGHT, left:left + WIDTH]
+    arr = (arr - 0.5) / 0.5
+    return arr.unsqueeze(0).to(device="cuda", dtype=DTYPE)  # [1, 1, 3, H, W]
+
+
+def build_conditional_dict(pipe, pixel, num_frames, mouse_cond, keyboard_cond):
+    """MG2-style conditioning: CLIP context + first-frame latent + mask."""
+    num_pixel_frames = (num_frames - 1) * VAE_TCR + 1
+    visual_context = pipe.vae.encode_visual_context_from_pixels(pixel).to(DTYPE)
 
-def build_conditional_dict(pixel, num_frames, mouse_cond, keyboard_cond):
-    """Reference `build_conditional_dict()` — MG2 three-pathway I2V condition."""
-    num_pixel_frames = (num_frames - 1) * 4 + 1
-    visual_context = VAE.encode_visual_context_from_pixels(pixel).to(DTYPE)
     first_frame = pixel[:, 0:1]
-    pad_pix = torch.zeros(
-        1, num_pixel_frames - 1, 3, pixel.shape[3], pixel.shape[4],
-        device=pixel.device, dtype=pixel.dtype,
-    )
-    padded = torch.cat([first_frame, pad_pix], dim=1).permute(0, 2, 1, 3, 4)
-    img_cond = VAE.encode_to_latent(padded).to(DTYPE)
+    pad = torch.zeros(1, num_pixel_frames - 1, 3, pixel.shape[3], pixel.shape[4],
+                      device=pixel.device, dtype=DTYPE)
+    padded = torch.cat([first_frame, pad], dim=1).permute(0, 2, 1, 3, 4)
+    img_cond = pipe.vae.encode_to_latent(padded).to(DTYPE)
 
     _, _, _, h_lat, w_lat = img_cond.shape
-    mask = torch.zeros(1, num_frames, 4, h_lat, w_lat, device=DEVICE, dtype=DTYPE)
+    mask = torch.zeros(1, num_frames, 4, h_lat, w_lat,
+                       device=pixel.device, dtype=DTYPE)
     mask[:, 0:1] = 1
-    cond_concat = torch.cat([mask, img_cond], dim=2)
     return {
         "visual_context": visual_context,
-        "cond_concat": cond_concat,
-        "mouse_condition": mouse_cond.to(device=DEVICE, dtype=DTYPE),
-        "keyboard_condition": keyboard_cond.to(device=DEVICE, dtype=DTYPE),
+        "cond_concat": torch.cat([mask, img_cond], dim=2),
+        "mouse_condition": mouse_cond.to(device="cuda", dtype=DTYPE),
+        "keyboard_condition": keyboard_cond.to(device="cuda", dtype=DTYPE),
     }
 
 
-def _estimate_duration(reference_image, action_script="forward", *args, **kwargs):
-    try:
-        chunks = len(parse_script(action_script or "forward"))
-    except Exception:
-        chunks = 7
-    return int(min(180, 40 + 6 * chunks))
+# ──────────────────────────────── core ────────────────────────────────
+@torch.no_grad()
+def _rollout(image_path, variant, segments, seed):
+    if image_path is None:
+        raise gr.Error("Please provide a reference frame.")
+    if variant not in PIPELINES:
+        variant = DEFAULT_VARIANT
+    segments = [(a, int(n)) for a, n in segments if int(n) > 0]
+    total_chunks = sum(n for _, n in segments)
+    if total_chunks < 1:
+        raise gr.Error("Give at least one action segment a non-zero length.")
+    if total_chunks > 12:
+        raise gr.Error("Keep the total rollout at 12 chunks or fewer.")
+
+    pipe = PIPELINES[variant]
+    num_frames = total_chunks * FRAMES_PER_BLOCK
+
+    pixel = load_reference_frame(image_path)
+    mouse, keyboard = build_action_track(segments)
+    cond = build_conditional_dict(pipe, pixel, num_frames, mouse, keyboard)
+
+    torch.manual_seed(int(seed))
+    noise = torch.randn([1, num_frames, 16, HEIGHT // 8, WIDTH // 8],
+                        device="cuda", dtype=DTYPE)
+
+    start = time.time()
+    video = pipe.inference(noise=noise, conditional_dict=cond, return_latents=False)
+    torch.cuda.synchronize()
+    elapsed = time.time() - start
+    pipe.vae.model.clear_cache()
+
+    frames = (video[0].permute(0, 2, 3, 1).float().cpu().numpy() * 255)
+    frames = frames.clip(0, 255).astype(np.uint8)
+    out_path = os.path.join(tempfile.mkdtemp(), "forgewm.mp4")
+    writer = imageio.get_writer(out_path, fps=FPS, codec="libx264",
+                                quality=8, macro_block_size=None)
+    for frame in frames:
+        writer.append_data(frame)
+    writer.close()
+
+    script = " → ".join(f"`{a}` ×{n}" for a, n in segments)
+    info = (
+        f"**{variant}** · {total_chunks} causal blocks · {num_frames} latent "
+        f"frames → {len(frames)} pixel frames ({len(frames) / FPS:.1f}s @ {FPS} fps)\n\n"
+        f"Action script: {script}\n\n"
+        f"Rollout + VAE decode: **{elapsed:.2f}s** "
+        f"({1000 * elapsed / total_chunks:.0f} ms per block, decode included) · seed `{int(seed)}`"
+    )
+    return out_path, info
+
+
+def _estimate(variant, chunks) -> int:
+    per = SECONDS_PER_CHUNK.get(variant, 1.5)
+    return int(FIXED_OVERHEAD + per * max(1, min(int(chunks), 12)) + 8)
+
+
+def _duration_main(*args):
+    variant = args[1] if len(args) > 1 else DEFAULT_VARIANT
+    chunks = sum(int(args[i]) for i in (3, 5, 7) if len(args) > i)
+    return _estimate(variant, chunks)
 
 
-@spaces.GPU(duration=_estimate_duration)
+def _duration_example(*args):
+    chunks = sum(int(args[i]) for i in (2, 4, 6) if len(args) > i)
+    return _estimate(DEFAULT_VARIANT, chunks)
+
+
+@spaces.GPU(duration=_duration_main)
 def generate(
-    reference_image: str,
-    action_script: str = "forward",
-    model_choice: str = DEFAULT_MODEL,
-    camera_speed: float = CAM_VALUE,
+    image: str,
+    variant: str,
+    action_1: str,
+    chunks_1: int,
+    action_2: str,
+    chunks_2: int,
+    action_3: str,
+    chunks_3: int,
     seed: int = 0,
-    randomize_seed: bool = True,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Roll out the ForgeWM world model from a reference frame and an action script.
+    """Roll out an action-conditioned Minecraft video from one reference frame.
 
     Args:
-        reference_image: Path to the opening frame the world is generated from.
-        action_script: Comma-separated actions, one per second of video, with an
-            optional `xN` repeat (e.g. "forward x3, turn_right x2"). Valid
-            actions: forward, back, left, right, turn_left, turn_right, look_up,
-            look_down, forward_turn_right, random, no_action.
-        model_choice: Which few-step student to run (4-step or 1-step).
-        camera_speed: Magnitude of the mouse/camera delta per frame.
-        seed: RNG seed for the initial noise.
-        randomize_seed: Draw a fresh random seed instead of using `seed`.
+        image: Path to the reference frame that anchors the world.
+        variant: Which few-step ForgeWM student to sample with.
+        action_1: Action held during the first segment of the rollout.
+        chunks_1: Length of the first segment, in 3-latent-frame blocks (~1s each).
+        action_2: Action held during the second segment.
+        chunks_2: Length of the second segment, in blocks.
+        action_3: Action held during the third segment.
+        chunks_3: Length of the third segment, in blocks.
+        seed: Random seed for the initial noise.
 
     Returns:
-        A tuple of (path to the generated mp4, a short markdown run summary).
+        An mp4 of the generated rollout and a markdown summary of the run.
     """
-    if reference_image is None:
-        raise gr.Error("Please provide a reference frame to start the world from.")
-
-    actions = parse_script(action_script)
-    num_chunks = len(actions)
-    num_frames = num_chunks * NUM_FRAME_PER_BLOCK
-    num_raw_frames = (num_frames - 1) * 4 + 1
-
-    if randomize_seed:
-        seed = random.randint(0, 2**31 - 1)
-    seed = int(seed)
-
-    pipeline = MODELS.get(model_choice, PIPE_4)
-
-    image = Image.open(reference_image).convert("RGB")
-    pixel = _TRANSFORM(image).unsqueeze(0).unsqueeze(0).to(device=DEVICE, dtype=DTYPE)
-
-    mouse_cond, keyboard_cond = build_action_tensors(
-        actions, num_raw_frames, float(camera_speed), seed
+    return _rollout(
+        image, variant,
+        [(action_1, chunks_1), (action_2, chunks_2), (action_3, chunks_3)],
+        seed,
     )
 
-    t0 = time.perf_counter()
-    # `torch.set_grad_enabled` is thread-local, so the module-scope call does not
-    # reach the ZeroGPU worker thread — without this guard the 30-layer causal
-    # rollout keeps every activation alive and OOMs.
-    with _GPU_LOCK, torch.no_grad():
-        conditional_dict = build_conditional_dict(
-            pixel, num_frames, mouse_cond, keyboard_cond
-        )
-        torch.manual_seed(seed)
-        noise = torch.randn(
-            [1, num_frames, 16, HEIGHT // 8, WIDTH // 8], device=DEVICE, dtype=DTYPE
-        )
-        video = pipeline.inference(
-            noise=noise, conditional_dict=conditional_dict, return_latents=False
-        )
-        video_np = (
-            video[0].permute(0, 2, 3, 1).cpu().float().numpy() * 255
-        ).clip(0, 255).astype(np.uint8)
-        if hasattr(pipeline.vae, "model"):
-            pipeline.vae.model.clear_cache()
-        try:
-            free_b, total_b = torch.cuda.mem_get_info()
-            print(
-                f"[gen] peak={torch.cuda.max_memory_allocated() / 2**30:.1f}GiB "
-                f"free={free_b / 2**30:.1f}/{total_b / 2**30:.1f}GiB",
-                flush=True,
-            )
-        except Exception as exc:  # pragma: no cover - diagnostics only
-            print(f"[gen] memory probe failed: {exc}", flush=True)
-    elapsed = time.perf_counter() - t0
-
-    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-    writer = imageio.get_writer(
-        out_path, fps=FPS, codec="libx264", quality=8, macro_block_size=None
-    )
-    for frame in video_np:
-        writer.append_data(frame)
-    writer.close()
 
-    info = (
-        f"**{model_choice}** · {len(video_np)} frames @ {FPS} fps "
-        f"({len(video_np) / FPS:.1f}s, {WIDTH}×{HEIGHT}) · "
-        f"{num_chunks} causal chunk(s) · seed `{seed}` · "
-        f"generated in **{elapsed:.1f}s**  \n"
-        f"Actions: `{' → '.join(actions)}`"
+@spaces.GPU(duration=_duration_example)
+def generate_example(
+    image: str,
+    action_1: str,
+    chunks_1: int,
+    action_2: str,
+    chunks_2: int,
+    action_3: str,
+    chunks_3: int,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Run a bundled example with the default student and seed 0."""
+    return _rollout(
+        image, DEFAULT_VARIANT,
+        [(action_1, chunks_1), (action_2, chunks_2), (action_3, chunks_3)],
+        0,
     )
-    return out_path, info
-
 
-# ─── UI ───────────────────────────────────────────────────────────────────────
 
+# ───────────────────────────────── UI ─────────────────────────────────
 CSS = """
-#col-container { max-width: 1180px; margin: 0 auto; }
-.dark .gradio-container { color: var(--body-text-color); }
-#pad button { min-width: 0 !important; }
-"""
-
-INTRO = """# 🎮 ForgeWM — a Minecraft world model you can drive
-
-**[ForgeWM](https://huggingface.co/ForgeWM/ForgeWM)** is a few-step, action-conditioned video world model:
-give it one frame plus a keyboard/mouse action track and it rolls the world forward, one causal chunk at a time.
-Distilled with progressive causal training (bidirectional SFT → teacher-forced causal AR → consistency
-distillation → on-policy DMD) from a Matrix-Game 2 / Wan2.1-1.3B backbone.
-
-Build an action track below — **each action lasts one second** of generated video.
-
-📄 [Paper](https://huggingface.co/papers/2608.14022) · 💻 [Code](https://github.com/asdfo123/ForgeWM) · 🌐 [Project page](https://asdfo123.github.io/ForgeWM/)
+.gradio-container { max-width: 1200px !important; }
 """
 
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="ForgeWM") as demo:
+    gr.Markdown(
+        "# 🎮 ForgeWM — few-step action-conditioned world model\n"
+        "Give it **one Minecraft frame** and a short **action script**; the "
+        "block-causal diffusion transformer rolls the world forward at "
+        "**1, 2 or 4 denoising steps** per 3-frame block.\n\n"
+        "[Paper](https://huggingface.co/papers/2608.14022) · "
+        "[Model](https://huggingface.co/ForgeWM/ForgeWM) · "
+        "[Code](https://github.com/asdfo123/ForgeWM) · "
+        "[Project page](https://asdfo123.github.io/ForgeWM/)"
+    )
 
-def append_action(script: str, action: str) -> str:
-    current = [t for t in (script or "").replace(",", " ").split() if t]
-    if len(current) >= MAX_CHUNKS:
-        gr.Warning(f"The action track is capped at {MAX_CHUNKS} seconds.")
-        return script
-    current.append(action)
-    return ", ".join(current)
-
-
-def undo_action(script: str) -> str:
-    current = [t for t in (script or "").replace(",", " ").split() if t]
-    return ", ".join(current[:-1])
-
-
-with gr.Blocks(title="ForgeWM world model") as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(INTRO)
-
-        with gr.Row():
-            with gr.Column(scale=1):
-                reference_image = gr.Image(
-                    label="Reference frame (the world starts here)",
-                    type="filepath",
-                    height=300,
-                )
-                action_script = gr.Textbox(
-                    label="Action track — one action per second",
-                    value="forward, forward, turn_right, turn_right, forward",
-                    lines=2,
-                    info="Click the controls below, or type directly "
-                    "(e.g. 'forward x3, turn_right x2').",
-                )
-                with gr.Row(elem_id="pad"):
-                    undo_btn = gr.Button("⤺ Undo", size="sm", scale=1)
-                    clear_btn = gr.Button("✕ Clear", size="sm", scale=1)
-                with gr.Row(elem_id="pad"):
-                    pad_buttons = [
-                        (gr.Button(label, size="sm"), action)
-                        for label, action in BUTTONS[:4]
-                    ]
-                with gr.Row(elem_id="pad"):
-                    pad_buttons += [
-                        (gr.Button(label, size="sm"), action)
-                        for label, action in BUTTONS[4:8]
-                    ]
-                with gr.Row(elem_id="pad"):
-                    pad_buttons += [
-                        (gr.Button(label, size="sm"), action)
-                        for label, action in BUTTONS[8:]
-                    ]
-                run_btn = gr.Button("▶ Roll out the world", variant="primary")
-
-            with gr.Column(scale=1):
-                video_out = gr.Video(
-                    label="Generated rollout", autoplay=True, loop=True, height=380
-                )
-                info_out = gr.Markdown()
-
-        with gr.Accordion("Advanced settings", open=False):
-            model_choice = gr.Radio(
-                choices=list(MODELS.keys()),
-                value=DEFAULT_MODEL,
-                label="Few-step student",
-                info="ForgeWM-1 runs a single denoising step per chunk (with the "
-                "paper's First-Frame Enhancement on chunk 0); ForgeWM-4 runs four.",
-            )
-            camera_speed = gr.Slider(
-                0.02, 0.30, value=CAM_VALUE, step=0.01,
-                label="Camera speed (mouse delta per frame)",
+    with gr.Row():
+        with gr.Column(scale=1):
+            image = gr.Image(label="Reference frame", type="filepath",
+                             height=280, sources=["upload", "clipboard"])
+            variant = gr.Radio(
+                choices=list(PIPELINES.keys()),
+                value=DEFAULT_VARIANT,
+                label="Student",
+                info="Fewer steps = faster. 1-/2-step use First-Frame Enhancement.",
             )
+            gr.Markdown("### Action script  \nEach block is 12 frames ≈ 1 second.")
             with gr.Row():
-                seed = gr.Number(label="Seed", value=0, precision=0)
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-        gr.Examples(
-            examples=[
-                ["demo_images/forest.png", "forward, forward, turn_right, turn_right, forward"],
-                ["demo_images/plains.png", "forward x4, look_up, forward x2"],
-                ["demo_images/cave.png", "forward x3, turn_left, turn_left, forward x2"],
-                ["demo_images/forest.png", "random x7"],
-            ],
-            inputs=[reference_image, action_script],
-            outputs=[video_out, info_out],
-            fn=generate,
-            cache_examples=True,
-            cache_mode="lazy",
-            label="Examples (reference frames from the ForgeWM repo)",
-        )
-
-        gr.Markdown(
-            "Rollouts run at 352×640, 12 fps, in causal chunks of 3 latent frames "
-            "(12 pixel frames) with a 6-frame sliding attention window — exactly the "
-            "regime the Stage-3 students were trained in."
-        )
-
-    for btn, action in pad_buttons:
-        btn.click(
-            lambda s, a=action: append_action(s, a),
-            inputs=action_script,
-            outputs=action_script,
-            api_name=False,
-        )
-    undo_btn.click(undo_action, inputs=action_script, outputs=action_script,
-                   api_name=False)
-    clear_btn.click(lambda: "", outputs=action_script, api_name=False)
-
-    run_btn.click(
-        generate,
-        inputs=[reference_image, action_script, model_choice, camera_speed, seed,
-                randomize_seed],
-        outputs=[video_out, info_out],
-        api_name="generate",
-        concurrency_limit=1,
+                action_1 = gr.Dropdown(MINECRAFT_ACTIONS, value="forward",
+                                       label="Segment 1", scale=2)
+                chunks_1 = gr.Slider(0, 8, value=3, step=1, label="blocks", scale=1)
+            with gr.Row():
+                action_2 = gr.Dropdown(MINECRAFT_ACTIONS, value="turn_right",
+                                       label="Segment 2", scale=2)
+                chunks_2 = gr.Slider(0, 8, value=2, step=1, label="blocks", scale=1)
+            with gr.Row():
+                action_3 = gr.Dropdown(MINECRAFT_ACTIONS, value="forward",
+                                       label="Segment 3", scale=2)
+                chunks_3 = gr.Slider(0, 8, value=2, step=1, label="blocks", scale=1)
+            with gr.Accordion("Advanced", open=False):
+                seed = gr.Slider(0, 2**31 - 1, value=0, step=1, label="Seed")
+            run = gr.Button("Roll out the world", variant="primary")
+
+        with gr.Column(scale=1):
+            video = gr.Video(label="Generated rollout", autoplay=True, loop=True)
+            info = gr.Markdown()
+
+    gr.Examples(
+        examples=[
+            ["examples/forest.png", "forward", 3, "turn_right", 2, "forward", 2],
+            ["examples/plains.png", "forward", 3, "look_up", 1, "forward_turn_right", 3],
+            ["examples/cave.png", "forward", 2, "turn_left", 2, "forward", 3],
+        ],
+        inputs=[image, action_1, chunks_1, action_2, chunks_2, action_3, chunks_3],
+        outputs=[video, info],
+        fn=generate_example,
+        cache_examples=True,
+        cache_mode="lazy",
+        label="Examples (reference frames from the ForgeWM repo)",
     )
 
-if __name__ == "__main__":
-    demo.launch(
-        theme=gr.themes.Citrus(), css=CSS, mcp_server=True, show_error=True
+    gr.Markdown(
+        "Rollouts are 352×640 at 12 fps. Attention uses a 6-latent-frame sliding "
+        "window with a block-causal KV cache, exactly as the released students "
+        "were trained. Base weights: "
+        "[Skywork/Matrix-Game-2.0](https://huggingface.co/Skywork/Matrix-Game-2.0)."
     )
+
+    run.click(
+        fn=generate,
+        inputs=[image, variant, action_1, chunks_1, action_2, chunks_2,
+                action_3, chunks_3, seed],
+        outputs=[video, info],
+        concurrency_limit=1,
+        api_name="generate",
+    )
+
+demo.queue(max_size=12).launch(mcp_server=True, show_error=True)

pipeline/__init__.py

@@ -1,3 +1,13 @@
+from .causal_diffusion_inference import CausalDiffusionInferencePipeline
 from .causal_inference import CausalInferencePipeline
+from .self_forcing_training import SelfForcingTrainingPipeline
+from .teacher_forcing_training import TeacherForcingTrainingPipeline
+from .bidirectional_training import BidirectionalTrainingPipeline
 
-__all__ = ["CausalInferencePipeline"]
+__all__ = [
+    "CausalDiffusionInferencePipeline",
+    "CausalInferencePipeline",
+    "SelfForcingTrainingPipeline",
+    "TeacherForcingTrainingPipeline",
+    "BidirectionalTrainingPipeline",
+]

pipeline/bidirectional_training.py

@@ -0,0 +1,160 @@
+from utils.wan_wrapper import WanDiffusionWrapper
+from utils.scheduler import SchedulerInterface
+from typing import List, Optional
+import torch
+import torch.distributed as dist
+
+
+class BidirectionalTrainingPipeline:
+    def __init__(self,
+                 denoising_step_list: List[int],
+                 scheduler: SchedulerInterface,
+                 generator: WanDiffusionWrapper,
+                 num_frame_per_block=3,
+                 independent_first_frame: bool = False,
+                 same_step_across_blocks: bool = False,
+                 last_step_only: bool = False,
+                 num_max_frames: int = 21,
+                 context_noise: int = 0,
+                 spatial_self: bool = True,
+                 **kwargs):
+        super().__init__()
+        self.scheduler = scheduler
+        self.generator = generator
+        self.denoising_step_list = denoising_step_list
+        if self.denoising_step_list[-1] == 0:
+            self.denoising_step_list = self.denoising_step_list[:-1]  # remove the zero timestep for inference
+
+        # Wan specific hyperparameters
+        self.num_transformer_blocks = 30
+        self.frame_seq_length = 1560
+        self.num_frame_per_block = num_frame_per_block
+        self.context_noise = context_noise
+        self.i2v = False
+
+        self.kv_cache1 = None
+        self.kv_cache2 = None
+        self.independent_first_frame = independent_first_frame
+        self.same_step_across_blocks = same_step_across_blocks
+        self.last_step_only = last_step_only
+        self.kv_cache_size = num_max_frames * self.frame_seq_length
+        
+        self.spatial_self = spatial_self
+
+    def generate_and_sync_list(self, num_blocks, num_denoising_steps, device):
+        rank = dist.get_rank() if dist.is_initialized() else 0
+
+        if rank == 0:
+            # Generate random indices
+            indices = torch.randint(
+                low=0,
+                high=num_denoising_steps,
+                size=(num_blocks,),
+                device=device
+            )
+            # In our training, self.last_step_only is False
+            if self.last_step_only:
+                indices = torch.ones_like(indices) * (num_denoising_steps - 1)
+        else:
+            indices = torch.empty(num_blocks, dtype=torch.long, device=device)
+
+        dist.broadcast(indices, src=0)  # Broadcast the random indices to all ranks
+        return indices.tolist()
+
+    def inference_with_trajectory(
+            self,
+            noise: torch.Tensor,
+            clean_image_or_video: torch.Tensor = None, # same shape as noise
+            initial_latent: Optional[torch.Tensor] = None,
+            return_sim_step: bool = False,
+            **conditional_dict
+    ) -> torch.Tensor:
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        if not self.independent_first_frame or (self.independent_first_frame and initial_latent is not None):
+            # If the first frame is independent and the first frame is provided, then the number of frames in the
+            # noise should still be a multiple of num_frame_per_block
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        else:
+            # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames + num_input_frames  # add the initial latent frames
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+        
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks 
+        num_denoising_steps = len(self.denoising_step_list)
+        exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device)
+        start_gradient_frame_index = num_output_frames - 21 # always 0 as long as we train 21 latent frames
+        if start_gradient_frame_index != 0:
+            raise NotImplementedError("start_gradient_frame_index is always 0 as long as we train 21 latent frames")
+        
+        
+        noisy_input = noise
+        for index, current_timestep in enumerate(self.denoising_step_list):
+            # self.same_step_across_blocks is True
+            if self.same_step_across_blocks:
+                exit_flag = (index == exit_flags[0])
+            else:
+                raise NotImplementedError('Here t is a scalar denoting that all chunks are at the same t, but in the future we may set t a tensor denoting different chunks')  # Only backprop at the randomly selected timestep (consistent across all ranks)
+            timestep = torch.ones(
+                [batch_size, self.num_frame_per_block*num_blocks],
+                device=noise.device,
+                dtype=torch.int64) * current_timestep
+
+            if not exit_flag:
+                with torch.no_grad():
+                    _,denoised_pred = self.generator(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep
+                    )
+                    next_timestep = self.denoising_step_list[index + 1]
+                    noisy_input = self.scheduler.add_noise(
+                        denoised_pred.flatten(0, 1),
+                        torch.randn_like(denoised_pred.flatten(0, 1)),
+                        next_timestep * torch.ones(
+                            [batch_size * self.num_frame_per_block*num_blocks], device=noise.device, dtype=torch.long)
+                    ).unflatten(0, denoised_pred.shape[:2])
+                    print('denoise')
+            else:
+                _,output = self.generator(
+                            noisy_image_or_video=noisy_input,
+                            conditional_dict=conditional_dict,
+                            timestep=timestep
+                        )
+                print('final denoise')
+                break
+        # ======================= SF -> TF modification ends  ============================
+        
+        # Step 3.5: Return the denoised timestep
+        if not self.same_step_across_blocks: # Useless, never met
+            denoised_timestep_from, denoised_timestep_to = None, None
+        # T -> \tau_1 -> \tau_2 ->...-> \tau —— enable grad ——> 0
+        # denoised_timestep_from = \tau
+        # denoised_timestep_to = next timestep smaller than \tau
+        # These are just engineering tricks
+        # to align DMD timestep sampling with the actual denoising range used by the generator
+        elif exit_flags[0] == len(self.denoising_step_list) - 1:
+            # corner case when \tau is the smallest non-zero timestep
+            denoised_timestep_to = 0
+            denoised_timestep_from = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+        else:
+            denoised_timestep_to = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0] + 1].cuda()).abs(), dim=0).item()
+            denoised_timestep_from = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+        
+        if return_sim_step: # False
+            return output, denoised_timestep_from, denoised_timestep_to, exit_flags[0] + 1
+
+        return output, denoised_timestep_from, denoised_timestep_to
+
+    

pipeline/causal_diffusion_inference.py

@@ -0,0 +1,637 @@
+from tqdm import tqdm
+from typing import List, Optional
+import torch
+
+from wan.utils.fm_solvers import FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps
+from wan.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+def cond_current(conditional_dict, current_start_frame, num_frame_per_block, vae_time_compression_ratio=4):
+    raw_end = 1 + vae_time_compression_ratio * (current_start_frame + num_frame_per_block - 1)
+    current = {
+        "visual_context": conditional_dict["visual_context"],
+        "cond_concat": conditional_dict["cond_concat"][:, current_start_frame:current_start_frame + num_frame_per_block],
+    }
+    mouse_condition = conditional_dict.get("mouse_condition", conditional_dict.get("mouse_cond"))
+    keyboard_condition = conditional_dict.get("keyboard_condition", conditional_dict.get("keyboard_cond"))
+    if mouse_condition is not None:
+        current["mouse_condition"] = mouse_condition[:, :raw_end]
+    if keyboard_condition is not None:
+        current["keyboard_condition"] = keyboard_condition[:, :raw_end]
+    return current
+
+
+class CausalDiffusionInferencePipeline(torch.nn.Module):
+    def __init__(
+            self,
+            args,
+            device,
+            generator=None,
+            text_encoder=None,
+            vae=None,
+            need_vae = True
+    ):
+        super().__init__()
+        # Step 1: Initialize all models
+        model_kwargs = dict(getattr(args, "model_kwargs", {}))
+        action_config = getattr(args, "action_config", None)
+        if action_config is not None:
+            action_config = dict(action_config)
+            action_config.pop("local_attn_size", None)
+            model_kwargs["action_config"] = action_config
+        # Match the training-time sliding-window regime; see
+        # CausalInferencePipeline for the full rationale.
+        if "local_attn_size" not in model_kwargs and hasattr(args, "local_attn_size"):
+            model_kwargs["local_attn_size"] = args.local_attn_size
+        if "sink_size" not in model_kwargs and hasattr(args, "sink_size"):
+            model_kwargs["sink_size"] = args.sink_size
+        self.generator = WanDiffusionWrapper(
+            **model_kwargs, is_causal=True) if generator is None else generator
+        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
+        if need_vae:
+            self.vae = WanVAEWrapper() if vae is None else vae
+
+        # Step 2: Initialize scheduler
+        self.num_train_timesteps = args.num_train_timestep
+        self.sampling_steps = 50
+        self.sample_solver = 'unipc'
+        self.shift = args.timestep_shift
+
+        self.num_transformer_blocks = 30
+        self.frame_seq_length = None
+
+        self.kv_cache_pos = None
+        self.kv_cache_neg = None
+        self.kv_cache_mouse_pos = None
+        self.kv_cache_keyboard_pos = None
+        self.crossattn_cache_pos = None
+        self.crossattn_cache_neg = None
+        self.args = args
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.independent_first_frame = args.independent_first_frame
+        self.local_attn_size = self.generator.model.local_attn_size
+
+        print(f"KV inference with {self.num_frame_per_block} frames per block")
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+    def inference(
+        self,
+        noise: torch.Tensor,
+        conditional_dict: dict,
+        initial_latent: Optional[torch.Tensor] = None,
+        return_latents: bool = False,
+        start_frame_index: Optional[int] = 0,
+        return_video=True
+    ) -> torch.Tensor:
+        """
+        Perform inference on the given noise and Stage1 conditional inputs.
+        Inputs:
+            noise (torch.Tensor): The input noise tensor of shape
+                (batch_size, num_output_frames, num_channels, height, width).
+            conditional_dict (dict): MG2-style conditioning with visual_context, cond_concat,
+                and optional mouse/keyboard action sequences.
+            initial_latent (torch.Tensor): The initial latent tensor of shape
+                (batch_size, num_input_frames, num_channels, height, width).
+                If num_input_frames is 1, perform image to video.
+                If num_input_frames is greater than 1, perform video extension.
+            return_latents (bool): Whether to return the latents.
+            start_frame_index (int): In long video generation, where does the current window start?
+        Outputs:
+            video (torch.Tensor): The generated video tensor of shape
+                (batch_size, num_frames, num_channels, height, width). It is normalized to be in the range [0, 1].
+        """
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        if not self.independent_first_frame or (self.independent_first_frame and initial_latent is not None):
+            # If the first frame is independent and the first frame is provided, then the number of frames in the
+            # noise should still be a multiple of num_frame_per_block
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        elif self.independent_first_frame and initial_latent is None:
+            # Using a [1, 4, 4, 4, 4, 4] model to generate a video without image conditioning
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames + num_input_frames  # add the initial latent frames
+        self.frame_seq_length = (height // self.generator.model.patch_size[1]) * (width // self.generator.model.patch_size[2])
+        vae_time_compression_ratio = getattr(self.args.action_config, 'vae_time_compression_ratio', 4) if getattr(self.args, 'action_config', None) else 4
+
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+
+        # Step 1: Initialize KV cache to all zeros
+        if self.kv_cache_pos is None:
+            self._initialize_kv_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device,
+                num_output_frames=num_output_frames,
+            )
+            self._initialize_kv_cache_mouse_and_keyboard(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device,
+                num_output_frames=num_output_frames,
+            )
+            self._initialize_crossattn_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device
+            )
+        else:
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache_pos[block_index]["is_init"] = False
+            for block_index in range(len(self.kv_cache_pos)):
+                self.kv_cache_pos[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache_pos[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache_mouse_pos[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache_mouse_pos[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache_keyboard_pos[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache_keyboard_pos[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+
+        # Step 2: Cache context feature
+        current_start_frame = start_frame_index
+        cache_start_frame = 0
+        if initial_latent is not None:
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            if self.independent_first_frame:
+                # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
+                assert (num_input_frames - 1) % self.num_frame_per_block == 0
+                num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block
+                output[:, :1] = initial_latent[:, :1]
+                self.generator(
+                    noisy_image_or_video=initial_latent[:, :1],
+                    conditional_dict=cond_current(conditional_dict, current_start_frame, 1, vae_time_compression_ratio),
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache_pos,
+                    kv_cache_mouse=self.kv_cache_mouse_pos,
+                    kv_cache_keyboard=self.kv_cache_keyboard_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length
+                )
+                current_start_frame += 1
+                cache_start_frame += 1
+            else:
+                # Assume num_input_frames is self.num_frame_per_block * num_input_blocks
+                assert num_input_frames % self.num_frame_per_block == 0
+                num_input_blocks = num_input_frames // self.num_frame_per_block
+
+            for block_index in range(num_input_blocks):
+                current_ref_latents = \
+                    initial_latent[:, cache_start_frame:cache_start_frame + self.num_frame_per_block]
+                output[:, cache_start_frame:cache_start_frame + self.num_frame_per_block] = current_ref_latents
+                self.generator(
+                    noisy_image_or_video=current_ref_latents,
+                    conditional_dict=cond_current(conditional_dict, current_start_frame, self.num_frame_per_block, vae_time_compression_ratio),
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache_pos,
+                    kv_cache_mouse=self.kv_cache_mouse_pos,
+                    kv_cache_keyboard=self.kv_cache_keyboard_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length
+                )
+                current_start_frame += self.num_frame_per_block
+                cache_start_frame += self.num_frame_per_block
+
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+        for current_num_frames in all_num_frames:
+            noisy_input = noise[
+                :, cache_start_frame - num_input_frames:cache_start_frame + current_num_frames - num_input_frames]
+            latents = noisy_input
+
+            # Step 3.1: Spatial denoising loop
+            sample_scheduler = self._initialize_sample_scheduler(noise)
+            for _, t in enumerate(tqdm(sample_scheduler.timesteps)):
+                latent_model_input = latents
+                timestep = t * torch.ones(
+                    [batch_size, current_num_frames], device=noise.device, dtype=torch.float32
+                )
+
+                flow_pred, _ = self.generator(
+                    noisy_image_or_video=latent_model_input,
+                    conditional_dict=cond_current(conditional_dict, current_start_frame, current_num_frames, vae_time_compression_ratio),
+                    timestep=timestep,
+                    kv_cache=self.kv_cache_pos,
+                    kv_cache_mouse=self.kv_cache_mouse_pos,
+                    kv_cache_keyboard=self.kv_cache_keyboard_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length
+                )
+
+                temp_x0 = sample_scheduler.step(
+                    flow_pred,
+                    t,
+                    latents,
+                    return_dict=False)[0]
+                latents = temp_x0
+
+            # Step 3.2: record the model's output
+            output[:, cache_start_frame:cache_start_frame + current_num_frames] = latents
+
+            # Step 3.3: rerun with timestep zero to update KV cache using clean context
+            self.generator(
+                noisy_image_or_video=latents,
+                conditional_dict=cond_current(conditional_dict, current_start_frame, current_num_frames, vae_time_compression_ratio),
+                timestep=timestep * 0,
+                kv_cache=self.kv_cache_pos,
+                kv_cache_mouse=self.kv_cache_mouse_pos,
+                kv_cache_keyboard=self.kv_cache_keyboard_pos,
+                crossattn_cache=self.crossattn_cache_pos,
+                current_start=current_start_frame * self.frame_seq_length,
+                cache_start=cache_start_frame * self.frame_seq_length
+            )
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+            cache_start_frame += current_num_frames
+
+        # Step 4: Decode the output
+        if return_video:
+            video = self.vae.decode_to_pixel(output)
+            video = (video * 0.5 + 0.5).clamp(0, 1)
+
+            if return_latents:
+                return video, output
+            else:
+                return video
+        else:
+            return output
+
+    
+    def inference_for_cd(
+        self,
+        noise: torch.Tensor,
+        conditional_dict: dict,
+        record_step_indices: List[int],
+        initial_latent: Optional[torch.Tensor] = None,
+        start_frame_index: int = 0
+    ) -> torch.Tensor:
+        """Run causal denoising and record selected latent states for consistency distillation."""
+        self.sampling_steps = 48
+        batch_size, num_frames, _, height, width = noise.shape
+
+        if (not self.independent_first_frame) or (self.independent_first_frame and initial_latent is not None):
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        else:
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        self.frame_seq_length = (height // self.generator.model.patch_size[1]) * (width // self.generator.model.patch_size[2])
+        vae_time_compression_ratio = getattr(self.args.action_config, 'vae_time_compression_ratio', 4) if getattr(self.args, 'action_config', None) else 4
+
+        if self.kv_cache_pos is None:
+            self._initialize_kv_cache(batch_size=batch_size, dtype=noise.dtype, device=noise.device,
+                                      num_output_frames=num_frames + num_input_frames)
+            self._initialize_kv_cache_mouse_and_keyboard(batch_size=batch_size, dtype=noise.dtype, device=noise.device,
+                                                          num_output_frames=num_frames + num_input_frames)
+            self._initialize_crossattn_cache(batch_size=batch_size, dtype=noise.dtype, device=noise.device)
+        else:
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache_pos[block_index]["is_init"] = False
+            for block_index in range(len(self.kv_cache_pos)):
+                self.kv_cache_pos[block_index]["global_end_index"] = torch.tensor([0], dtype=torch.long, device=noise.device)
+                self.kv_cache_pos[block_index]["local_end_index"] = torch.tensor([0], dtype=torch.long, device=noise.device)
+                self.kv_cache_mouse_pos[block_index]["global_end_index"] = torch.tensor([0], dtype=torch.long, device=noise.device)
+                self.kv_cache_mouse_pos[block_index]["local_end_index"] = torch.tensor([0], dtype=torch.long, device=noise.device)
+                self.kv_cache_keyboard_pos[block_index]["global_end_index"] = torch.tensor([0], dtype=torch.long, device=noise.device)
+                self.kv_cache_keyboard_pos[block_index]["local_end_index"] = torch.tensor([0], dtype=torch.long, device=noise.device)
+
+        sample_scheduler_probe = self._initialize_sample_scheduler(noise)
+        total_steps = len(sample_scheduler_probe.timesteps)
+        record_step_indices = sorted(set(int(i) for i in record_step_indices))
+        if not record_step_indices:
+            raise ValueError("record_step_indices must be non-empty")
+        if record_step_indices[0] < 0 or record_step_indices[-1] >= total_steps:
+            raise ValueError(f"record_step_indices out of range: valid=[0,{total_steps - 1}], got={record_step_indices}")
+        record_set = set(record_step_indices)
+
+        current_start_frame = start_frame_index
+        cache_start_frame = 0
+        if initial_latent is not None:
+            timestep0 = torch.zeros([batch_size, 1], device=noise.device, dtype=torch.int64)
+            if self.independent_first_frame:
+                assert (num_input_frames - 1) % self.num_frame_per_block == 0
+                num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block
+                self.generator(
+                    noisy_image_or_video=initial_latent[:, :1],
+                    conditional_dict=cond_current(conditional_dict, current_start_frame, 1, vae_time_compression_ratio),
+                    timestep=timestep0,
+                    kv_cache=self.kv_cache_pos,
+                    kv_cache_mouse=self.kv_cache_mouse_pos,
+                    kv_cache_keyboard=self.kv_cache_keyboard_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length,
+                )
+                current_start_frame += 1
+                cache_start_frame += 1
+            else:
+                assert num_input_frames % self.num_frame_per_block == 0
+                num_input_blocks = num_input_frames // self.num_frame_per_block
+
+            for _ in range(num_input_blocks):
+                current_ref_latents = initial_latent[:, cache_start_frame:cache_start_frame + self.num_frame_per_block]
+                self.generator(
+                    noisy_image_or_video=current_ref_latents,
+                    conditional_dict=cond_current(conditional_dict, current_start_frame, self.num_frame_per_block, vae_time_compression_ratio),
+                    timestep=timestep0,
+                    kv_cache=self.kv_cache_pos,
+                    kv_cache_mouse=self.kv_cache_mouse_pos,
+                    kv_cache_keyboard=self.kv_cache_keyboard_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length,
+                )
+                current_start_frame += self.num_frame_per_block
+                cache_start_frame += self.num_frame_per_block
+
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+
+        full_chunk_record = []
+        for current_num_frames in all_num_frames:
+            latents = noise[:, cache_start_frame - num_input_frames:cache_start_frame + current_num_frames - num_input_frames]
+            chunk_records = []
+            sample_scheduler = self._initialize_sample_scheduler(noise)
+            current_cond = cond_current(conditional_dict, current_start_frame, current_num_frames, vae_time_compression_ratio)
+
+            for progress_id, t in enumerate(tqdm(sample_scheduler.timesteps)):
+                if progress_id in record_set:
+                    print(f"{progress_id}: {t} saved")
+                    chunk_records.append(latents.detach().clone())
+
+                timestep = t * torch.ones([batch_size, current_num_frames], device=noise.device, dtype=torch.float32)
+                flow_pred, _ = self.generator(
+                    noisy_image_or_video=latents,
+                    conditional_dict=current_cond,
+                    timestep=timestep,
+                    kv_cache=self.kv_cache_pos,
+                    kv_cache_mouse=self.kv_cache_mouse_pos,
+                    kv_cache_keyboard=self.kv_cache_keyboard_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length,
+                )
+                latents = sample_scheduler.step(flow_pred, t, latents, return_dict=False)[0]
+
+            chunk_records.append(latents.detach().clone())
+            full_chunk_record.append(torch.stack(chunk_records, dim=1))
+
+            timestep0 = torch.zeros([batch_size, current_num_frames], device=noise.device, dtype=torch.float32)
+            self.generator(
+                noisy_image_or_video=latents,
+                conditional_dict=current_cond,
+                timestep=timestep0,
+                kv_cache=self.kv_cache_pos,
+                kv_cache_mouse=self.kv_cache_mouse_pos,
+                kv_cache_keyboard=self.kv_cache_keyboard_pos,
+                crossattn_cache=self.crossattn_cache_pos,
+                current_start=current_start_frame * self.frame_seq_length,
+                cache_start=cache_start_frame * self.frame_seq_length,
+            )
+
+            current_start_frame += current_num_frames
+            cache_start_frame += current_num_frames
+
+        return torch.cat(full_chunk_record, dim=2)
+    
+    
+    def inference_for_genuine_cd(
+        self,
+        noisy_input: torch.Tensor,
+        conditional_dict: dict,
+        initial_latent: Optional[torch.Tensor] = None,
+        timestep_idx=0,
+        sampling_steps=48,
+        chunksize=3
+    ) -> torch.Tensor:
+        batch_size, num_frames, _, height, width = noisy_input.shape
+        assert num_frames == chunksize
+
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        self.frame_seq_length = (height // self.generator.model.patch_size[1]) * (width // self.generator.model.patch_size[2])
+        vae_time_compression_ratio = getattr(self.args.action_config, 'vae_time_compression_ratio', 4) if getattr(self.args, 'action_config', None) else 4
+
+        if self.kv_cache_pos is None:
+            self._initialize_kv_cache(
+                batch_size=batch_size,
+                dtype=noisy_input.dtype,
+                device=noisy_input.device,
+                num_output_frames=num_frames + num_input_frames,
+            )
+            self._initialize_kv_cache_mouse_and_keyboard(
+                batch_size=batch_size,
+                dtype=noisy_input.dtype,
+                device=noisy_input.device,
+                num_output_frames=num_frames + num_input_frames,
+            )
+            self._initialize_crossattn_cache(
+                batch_size=batch_size,
+                dtype=noisy_input.dtype,
+                device=noisy_input.device
+            )
+        else:
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache_pos[block_index]["is_init"] = False
+            for block_index in range(len(self.kv_cache_pos)):
+                self.kv_cache_pos[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noisy_input.device)
+                self.kv_cache_pos[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noisy_input.device)
+                self.kv_cache_mouse_pos[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noisy_input.device)
+                self.kv_cache_mouse_pos[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noisy_input.device)
+                self.kv_cache_keyboard_pos[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noisy_input.device)
+                self.kv_cache_keyboard_pos[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noisy_input.device)
+
+        current_start_frame = 0
+        cache_start_frame = 0
+        timestep0 = torch.zeros([batch_size, 1], device=noisy_input.device, dtype=torch.int64)
+
+        if initial_latent is not None:
+            if self.independent_first_frame:
+                assert (num_input_frames - 1) % chunksize == 0
+                num_input_blocks = (num_input_frames - 1) // chunksize
+                self.generator(
+                    noisy_image_or_video=initial_latent[:, :1],
+                    conditional_dict=cond_current(conditional_dict, current_start_frame, 1, vae_time_compression_ratio),
+                    timestep=timestep0,
+                    kv_cache=self.kv_cache_pos,
+                    kv_cache_mouse=self.kv_cache_mouse_pos,
+                    kv_cache_keyboard=self.kv_cache_keyboard_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length,
+                )
+                current_start_frame += 1
+                cache_start_frame += 1
+            else:
+                assert num_input_frames % chunksize == 0
+                num_input_blocks = num_input_frames // chunksize
+
+            for _ in range(num_input_blocks):
+                current_ref_latents = initial_latent[:, cache_start_frame:cache_start_frame + chunksize]
+                self.generator(
+                    noisy_image_or_video=current_ref_latents,
+                    conditional_dict=cond_current(conditional_dict, current_start_frame, chunksize, vae_time_compression_ratio),
+                    timestep=timestep0,
+                    kv_cache=self.kv_cache_pos,
+                    kv_cache_mouse=self.kv_cache_mouse_pos,
+                    kv_cache_keyboard=self.kv_cache_keyboard_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length,
+                )
+                current_start_frame += chunksize
+                cache_start_frame += chunksize
+
+        sample_scheduler = self._initialize_sample_scheduler(noisy_input, sampling_steps=sampling_steps)
+        t = sample_scheduler.timesteps[timestep_idx]
+        timestep = t * torch.ones(
+            [batch_size, chunksize], device=noisy_input.device, dtype=torch.float32
+        )
+        flow_pred, _ = self.generator(
+            noisy_image_or_video=noisy_input,
+            conditional_dict=cond_current(conditional_dict, current_start_frame, chunksize, vae_time_compression_ratio),
+            timestep=timestep,
+            kv_cache=self.kv_cache_pos,
+            kv_cache_mouse=self.kv_cache_mouse_pos,
+            kv_cache_keyboard=self.kv_cache_keyboard_pos,
+            crossattn_cache=self.crossattn_cache_pos,
+            current_start=current_start_frame * self.frame_seq_length,
+            cache_start=cache_start_frame * self.frame_seq_length,
+        )
+
+        latents = sample_scheduler.step(
+            flow_pred,
+            t,
+            noisy_input,
+            return_dict=False)[0]
+
+        return latents
+
+    
+
+    def _initialize_kv_cache(self, batch_size, dtype, device, num_output_frames=21):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+
+        For local_attn_size != -1, sliding-window cache sized to the window.
+        For local_attn_size == -1, cache must fit the entire rollout — Stage-3
+        uses 22 frames; the previous hardcoded 15 was undersized and caused
+        an out-of-bounds slice when block 5+ writes past frame 15.
+        """
+        kv_cache_pos = []
+        kv_cache_neg = []
+        if self.local_attn_size != -1:
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        else:
+            # full causal: cover the whole rollout
+            kv_cache_size = num_output_frames * self.frame_seq_length
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache_pos.append({
+                "k": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+            })
+            kv_cache_neg.append({
+                "k": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+            })
+
+        self.kv_cache_pos = kv_cache_pos  # always store the clean cache
+        self.kv_cache_neg = kv_cache_neg  # always store the clean cache
+
+    def _initialize_kv_cache_mouse_and_keyboard(self, batch_size, dtype, device, num_output_frames=21):
+        kv_cache_mouse_pos = []
+        kv_cache_keyboard_pos = []
+        # Same sizing logic as main kv_cache: -1 means cover entire rollout.
+        kv_cache_size = self.local_attn_size if self.local_attn_size != -1 else num_output_frames
+        for _ in range(self.num_transformer_blocks):
+            kv_cache_keyboard_pos.append({
+                "k": torch.zeros([batch_size, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+            })
+            kv_cache_mouse_pos.append({
+                "k": torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+            })
+        self.kv_cache_mouse_pos = kv_cache_mouse_pos
+        self.kv_cache_keyboard_pos = kv_cache_keyboard_pos
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        crossattn_cache_pos = []
+        crossattn_cache_neg = []
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache_pos.append({
+                "k": torch.zeros([batch_size, 257, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, 257, 12, 128], dtype=dtype, device=device),
+                "is_init": False
+            })
+            crossattn_cache_neg.append({
+                "k": torch.zeros([batch_size, 257, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, 257, 12, 128], dtype=dtype, device=device),
+                "is_init": False
+            })
+
+        self.crossattn_cache_pos = crossattn_cache_pos
+        self.crossattn_cache_neg = crossattn_cache_neg
+
+    def _initialize_sample_scheduler(self, noise, sampling_steps=-1):
+        if sampling_steps == -1:
+            sampling_steps = self.sampling_steps
+        if self.sample_solver == 'unipc':
+            sample_scheduler = FlowUniPCMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1,
+                use_dynamic_shifting=False)
+            sample_scheduler.set_timesteps(
+                sampling_steps, device=noise.device, shift=self.shift)
+            self.timesteps = sample_scheduler.timesteps
+        elif self.sample_solver == 'dpm++':
+            sample_scheduler = FlowDPMSolverMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1,
+                use_dynamic_shifting=False)
+            sampling_sigmas = get_sampling_sigmas(sampling_steps, self.shift)
+            self.timesteps, _ = retrieve_timesteps(
+                sample_scheduler,
+                device=noise.device,
+                sigmas=sampling_sigmas)
+        else:
+            raise NotImplementedError("Unsupported solver.")
+        return sample_scheduler

pipeline/self_forcing_training.py

@@ -0,0 +1,495 @@
+from utils.wan_wrapper import WanDiffusionWrapper
+from utils.scheduler import SchedulerInterface
+from typing import List, Optional
+import torch
+import torch.distributed as dist
+
+
+def _slice_cond_for_block(conditional_dict: dict, current_start_frame: int,
+                          num_frames_in_block: int,
+                          vae_time_compression_ratio: int = 4) -> dict:
+    """Return a shallow copy of conditional_dict with cond_concat /
+    mouse_condition / keyboard_condition sliced for the current block.
+
+    MG2's CausalWanModel concatenates cond_concat to x along the channel
+    dim, so cond_concat must have the SAME F dim as the noisy video
+    being fed in (= num_frames_in_block during autoregressive rollout).
+
+    visual_context stays as-is (it's the first-frame CLIP embedding,
+    not per-frame).
+
+    mouse/keyboard are sampled at pixel-frame rate; slice up to the
+    pixel frame corresponding to the end of the current block.
+    """
+    if "cond_concat" not in conditional_dict:
+        return conditional_dict
+    current = dict(conditional_dict)
+    current["cond_concat"] = conditional_dict["cond_concat"][
+        :, current_start_frame:current_start_frame + num_frames_in_block]
+    # Slice mouse/keyboard up to the corresponding pixel frame boundary.
+    raw_end = 1 + vae_time_compression_ratio * (
+        current_start_frame + num_frames_in_block - 1)
+    for k in ("mouse_condition", "keyboard_condition", "mouse_cond", "keyboard_cond"):
+        if k in conditional_dict and conditional_dict[k] is not None:
+            current[k] = conditional_dict[k][:, :raw_end]
+    return current
+
+
+class SelfForcingTrainingPipeline:
+    def __init__(self,
+                 denoising_step_list: List[int],
+                 scheduler: SchedulerInterface,
+                 generator: WanDiffusionWrapper,
+                 num_frame_per_block=3,
+                 independent_first_frame: bool = False,
+                 same_step_across_blocks: bool = False,
+                 last_step_only: bool = False,
+                 num_max_frames: int = 21,
+                 context_noise: int = 0,
+                 frame_seq_length: Optional[int] = None,
+                 vae_time_compression_ratio: int = 4,
+                 use_action: bool = False,
+                 denoising_step_list_first_chunk: Optional[List[int]] = None,
+                 **kwargs):
+        super().__init__()
+        self.scheduler = scheduler
+        self.generator = generator
+        self.denoising_step_list = denoising_step_list
+        if self.denoising_step_list[-1] == 0:
+            self.denoising_step_list = self.denoising_step_list[:-1]  # remove the zero timestep for inference
+
+        # ─── First-Frame Enhancement (FFE) ───────────────────────────────
+        # The 1-/2-step students spend the full 4-step schedule on chunk 0
+        # and the cheap main schedule on every later chunk.  Chunk 0 anchors
+        # the whole rollout (it is the only block conditioned purely on the
+        # reference frame), so paying for it once buys a large quality gain
+        # at negligible amortized cost.  None => one schedule throughout,
+        # which is what the 4-step student uses.
+        #
+        # NOTE (train/inference asymmetry, kept deliberately): the main
+        # `denoising_step_list` reaching this pipeline has already been
+        # remapped onto the shifted flow-matching grid when
+        # `warp_denoising_step: true`, whereas the first-chunk list is
+        # consumed here as the raw integers from the config.  At inference
+        # (see pipeline/causal_inference.py) BOTH lists are warped.  The
+        # released 1-/2-step checkpoints were trained under exactly this
+        # asymmetry, so we reproduce it rather than silently "fix" it: block 0
+        # is masked out of the DMD gradient
+        # (`gradient_mask[:, :num_frame_per_block] = False`) and the DMD
+        # timestep window below is derived from the interior schedule, so the
+        # first-chunk timesteps only shape the rollout that later blocks
+        # condition on.  Changing this would invalidate the published numbers.
+        self.denoising_step_list_first_chunk = denoising_step_list_first_chunk
+        if self.denoising_step_list_first_chunk is not None:
+            if self.denoising_step_list_first_chunk[-1] == 0:
+                self.denoising_step_list_first_chunk = self.denoising_step_list_first_chunk[:-1]
+
+        # Wan specific hyperparameters.
+        # num_transformer_blocks: MG2 base has num_layers=30.
+        self.num_transformer_blocks = 30
+        # frame_seq_length is (latent_H / patch_H) * (latent_W / patch_W)
+        # with patch_size=(1,2,2):
+        #   - 480p (60x104 latent) -> 30 * 52 = 1560
+        #   - 360p (44x80  latent) -> 22 * 40 =  880  ← MG2 native
+        # Prefer to pass this explicitly from config; otherwise fall back
+        # to the legacy 1560 default for back-compat with T2V Wan-1.3B configs.
+        self.frame_seq_length = frame_seq_length if frame_seq_length is not None else 1560
+        self.num_frame_per_block = num_frame_per_block
+        self.context_noise = context_noise
+        self.i2v = False
+        self.vae_time_compression_ratio = vae_time_compression_ratio
+
+        self.kv_cache1 = None
+        self.kv_cache2 = None
+        self.kv_cache_mouse = None
+        self.kv_cache_keyboard = None
+        self.use_action = use_action
+        # When the generator is MG2 (action-conditioned), the cross-attention
+        # context is CLIP visual_context of length 257, not legacy T5 text
+        # embeddings of length 512.
+        self._use_mg2_ctx = use_action
+        self.independent_first_frame = independent_first_frame
+        self.same_step_across_blocks = same_step_across_blocks
+        self.last_step_only = last_step_only
+        # IMPORTANT: train-time KV cache is sized to the FULL rollout
+        # (num_max_frames frames), regardless of local_attn_size.  This
+        # mirrors SF original (Self-Forcing/pipeline/self_forcing_training.py:39).
+        #
+        # Why: the sliding window during training is enforced by the
+        # _prepare_blockwise_causal_attn_mask FUNCTION (causal_model.py:535-540),
+        # NOT by physically truncating the KV cache.  Training cache stays
+        # full-size so the eviction-shift code path in causal_model.py L228
+        # never triggers — that path mutates kv_cache["k"]/["v"] in-place,
+        # which breaks gradient_checkpointing's recompute (the recomputed
+        # forward sees a different cache state than the original forward,
+        # producing a CheckpointError "saved metadata != recomputed metadata").
+        #
+        # Inference (pipeline/causal_inference.py:339) sizes the cache to
+        # `local_attn_size * frame_seq_length` and DOES trigger eviction —
+        # because at inference the cache is single-pass and has no
+        # backward to recompute.  Training and inference behaviors differ;
+        # the mask function is what aligns them semantically.
+        self.num_max_frames = num_max_frames
+        self.kv_cache_size = num_max_frames * self.frame_seq_length
+
+    def generate_and_sync_list(self, num_blocks, num_denoising_steps, device):
+        rank = dist.get_rank() if dist.is_initialized() else 0
+
+        if rank == 0:
+            # Generate random indices
+            indices = torch.randint(
+                low=0,
+                high=num_denoising_steps,
+                size=(num_blocks,),
+                device=device
+            )
+            # In our training, self.last_step_only is False
+            if self.last_step_only:
+                indices = torch.ones_like(indices) * (num_denoising_steps - 1)
+        else:
+            indices = torch.empty(num_blocks, dtype=torch.long, device=device)
+
+        dist.broadcast(indices, src=0)  # Broadcast the random indices to all ranks
+        return indices.tolist()
+
+    def inference_with_trajectory(
+            self,
+            noise: torch.Tensor,
+            clean_image_or_video: torch.Tensor = None, # same shape as noise
+            initial_latent: Optional[torch.Tensor] = None,
+            return_sim_step: bool = False,
+            **conditional_dict
+    ) -> torch.Tensor:
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        if not self.independent_first_frame or (self.independent_first_frame and initial_latent is not None):
+            # If the first frame is independent and the first frame is provided, then the number of frames in the
+            # noise should still be a multiple of num_frame_per_block
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        else:
+            # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames + num_input_frames  # add the initial latent frames
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+
+        # Step 1: Initialize KV cache to all zeros
+        self._initialize_kv_cache(
+            batch_size=batch_size, dtype=noise.dtype, device=noise.device
+        )
+        self._initialize_crossattn_cache(
+            batch_size=batch_size, dtype=noise.dtype, device=noise.device
+        )
+        if self.use_action:
+            self._initialize_kv_cache_mouse_and_keyboard(
+                batch_size=batch_size, dtype=noise.dtype, device=noise.device
+            )
+
+
+        # Step 2: Cache context feature
+        current_start_frame = 0
+        if initial_latent is not None: # Never met
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            # Cast initial_latent to the same dtype as noise so the ref-frame
+            # DiT forward doesn't hit a float/bfloat16 mismatch in
+            # patch_embedding (FSDP keeps model params bf16 under
+            # mixed_precision).
+            initial_latent = initial_latent.to(dtype=noise.dtype)
+            # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
+            output[:, :1] = initial_latent
+            # The ref-frame init is a 1-frame forward.  If the model's
+            # num_frame_per_block > 1, temporarily flip it to 1 so that the
+            # ActionModule's internal assertions (which assume one forward
+            # call processes exactly num_frame_per_block latent frames) hold.
+            # NOTE: the generator is FSDP-wrapped, so writing to
+            # `self.generator.model.num_frame_per_block` actually mutates
+            # the outer FSDP shell.  To reach the inner CausalWanModel
+            # (which `kwargs["num_frame_per_block"] = self.num_frame_per_block`
+            # reads at forward time), we walk through `_fsdp_wrapped_module`
+            # if present.
+            def _inner_model(m):
+                # Unwrap FSDP / checkpoint wrappers to get the real CausalWanModel
+                while hasattr(m, "_fsdp_wrapped_module"):
+                    m = m._fsdp_wrapped_module
+                if hasattr(m, "_checkpoint_wrapped_module"):
+                    m = m._checkpoint_wrapped_module
+                return m
+            _inner = _inner_model(self.generator.model)
+            _orig_nfpb = getattr(_inner, "num_frame_per_block", 1)
+            if _orig_nfpb != 1:
+                _inner.num_frame_per_block = 1
+            try:
+                with torch.no_grad():
+                    self.generator(
+                        noisy_image_or_video=initial_latent,
+                        conditional_dict=_slice_cond_for_block(
+                            conditional_dict, current_start_frame, 1,
+                            self.vae_time_compression_ratio),
+                        timestep=timestep * 0,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        kv_cache_mouse=self.kv_cache_mouse,
+                        kv_cache_keyboard=self.kv_cache_keyboard,
+                        current_start=current_start_frame * self.frame_seq_length
+                    )
+            finally:
+                if _orig_nfpb != 1:
+                    _inner.num_frame_per_block = _orig_nfpb
+            current_start_frame += 1
+
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        # In out training, self.independent_first_frame is False
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+        num_denoising_steps = len(self.denoising_step_list)
+        has_first = self.denoising_step_list_first_chunk is not None
+        # FFE gives chunk 0 a schedule of a different length, so its exit
+        # flag has to be drawn from its own range.  Every other chunk shares
+        # the interior range as before.
+        exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device)
+        if has_first:
+            exit_flag_first = self.generate_and_sync_list(
+                1, len(self.denoising_step_list_first_chunk), device=noise.device)[0]
+        else:
+            exit_flag_first = None
+        start_gradient_frame_index = num_output_frames - 21
+
+        # for block_index in range(num_blocks):
+        for block_index, current_num_frames in enumerate(all_num_frames):
+            # Slice the per-frame conditioning (cond_concat / mouse / kbd)
+            # to cover only the current block's latent frames.  visual_context
+            # is first-frame CLIP, stays as-is.
+            block_cond = _slice_cond_for_block(
+                conditional_dict, current_start_frame, current_num_frames,
+                self.vae_time_compression_ratio,
+            )
+
+            if True:
+                noisy_input = noise[
+                    :, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames]
+
+                # FFE: chunk 0 runs the (longer) first-chunk schedule; all
+                # later chunks run the main schedule.
+                if has_first and block_index == 0:
+                    current_denoising_list = self.denoising_step_list_first_chunk
+                    current_exit_flag_index = exit_flag_first
+                else:
+                    current_denoising_list = self.denoising_step_list
+                    current_exit_flag_index = (
+                        exit_flags[0] if self.same_step_across_blocks else exit_flags[block_index])
+
+                # Step 3.1: Spatial denoising loop
+                for index, current_timestep in enumerate(current_denoising_list):
+                    exit_flag = (index == current_exit_flag_index)
+                    timestep = torch.ones(
+                        [batch_size, current_num_frames],
+                        device=noise.device,
+                        dtype=torch.int64) * current_timestep
+
+                    if not exit_flag:
+                        with torch.no_grad():
+                            _, denoised_pred = self.generator(
+                                noisy_image_or_video=noisy_input,
+                                conditional_dict=block_cond,
+                                timestep=timestep,
+                                kv_cache=self.kv_cache1,
+                                crossattn_cache=self.crossattn_cache,
+                                kv_cache_mouse=self.kv_cache_mouse,
+                                kv_cache_keyboard=self.kv_cache_keyboard,
+                                current_start=current_start_frame * self.frame_seq_length
+                            )
+                            next_timestep = current_denoising_list[index + 1]
+                            noisy_input = self.scheduler.add_noise(
+                                denoised_pred.flatten(0, 1),
+                                torch.randn_like(denoised_pred.flatten(0, 1)),
+                                next_timestep * torch.ones(
+                                    [batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+                            ).unflatten(0, denoised_pred.shape[:2])
+                    else:
+                        if current_start_frame < start_gradient_frame_index:
+                            with torch.no_grad():
+                                _, denoised_pred = self.generator(
+                                    noisy_image_or_video=noisy_input,
+                                    conditional_dict=block_cond,
+                                    timestep=timestep,
+                                    kv_cache=self.kv_cache1,
+                                    crossattn_cache=self.crossattn_cache,
+                                    kv_cache_mouse=self.kv_cache_mouse,
+                                    kv_cache_keyboard=self.kv_cache_keyboard,
+                                    current_start=current_start_frame * self.frame_seq_length
+                                )
+                        else: # enable grad
+                            _, denoised_pred = self.generator(
+                                noisy_image_or_video=noisy_input,
+                                conditional_dict=block_cond,
+                                timestep=timestep,
+                                kv_cache=self.kv_cache1,
+                                crossattn_cache=self.crossattn_cache,
+                                kv_cache_mouse=self.kv_cache_mouse,
+                                kv_cache_keyboard=self.kv_cache_keyboard,
+                                current_start=current_start_frame * self.frame_seq_length
+                            )
+                        break
+
+            # Step 3.2: record the model's output
+            output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred
+
+            # Step 3.3: rerun with timestep zero to update the cache
+            #
+            # ─── Cache-refresh fix ───────────────────────────────────────────
+            # Previously we skipped this rerun for grad-enabled blocks (i.e.
+            # all chunks 1..6 when num_training_frames=22 since
+            # start_gradient_frame_index = num_output_frames - 21 = 1).
+            # That created a train/inference distribution shift on the KV
+            # cache: at inference each chunk goes through the full 4-step
+            # denoising chain so its cache is near-clean (≈ t=0); at training
+            # we used to leave cache at the exit_flags step (often noisy).
+            #
+            # CF orig / minWM / SF orig all run this refresh unconditionally.
+            # Restoring the original behavior closes the gap and empirically
+            # fixes HUD shrinkage / OOD drift on long-video rollouts.
+            #
+            # The previous "skip" branch is kept commented below for
+            # reference; do not re-enable without re-analyzing.
+            #
+            # was_grad_block = (current_start_frame >= start_gradient_frame_index)
+            # if was_grad_block:
+            #     current_start_frame += current_num_frames
+            #     continue
+            # ─────────────────────────────────────────────────────────────────
+
+
+            context_timestep = torch.ones_like(timestep) * self.context_noise
+            # add context noise
+            denoised_pred = self.scheduler.add_noise(
+                denoised_pred.flatten(0, 1),
+                torch.randn_like(denoised_pred.flatten(0, 1)),
+                context_timestep * torch.ones(
+                    [batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+            ).unflatten(0, denoised_pred.shape[:2])
+            with torch.no_grad():
+                self.generator(
+                    noisy_image_or_video=denoised_pred,
+                    conditional_dict=block_cond,
+                    timestep=context_timestep,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    kv_cache_mouse=self.kv_cache_mouse,
+                    kv_cache_keyboard=self.kv_cache_keyboard,
+                    current_start=current_start_frame * self.frame_seq_length
+                )
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+
+        # Step 3.5: Return the denoised timestep
+        if not self.same_step_across_blocks: # Useless, never met
+            denoised_timestep_from, denoised_timestep_to = None, None
+        # T -> \tau_1 -> \tau_2 ->...-> \tau —— enable grad ——> 0
+        # denoised_timestep_from = \tau
+        # denoised_timestep_to = next timestep smaller than \tau
+        # These are just engineering tricks
+        # to align DMD timestep sampling with the actual denoising range used by the generator
+        # Under FFE this window is deliberately derived from the INTERIOR
+        # schedule (`exit_flags[0]`), never from the first-chunk one: chunk 0
+        # is excluded from the DMD gradient, so the loss should be aligned
+        # with the schedule the supervised blocks actually ran.
+        elif exit_flags[0] == len(self.denoising_step_list) - 1:
+            # corner case when \tau is the smallest non-zero timestep
+            denoised_timestep_to = 0
+            denoised_timestep_from = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+        else:
+            denoised_timestep_to = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0] + 1].cuda()).abs(), dim=0).item()
+            denoised_timestep_from = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+
+        if return_sim_step: # False
+            return output, denoised_timestep_from, denoised_timestep_to, exit_flags[0] + 1
+
+        return output, denoised_timestep_from, denoised_timestep_to
+
+    def _initialize_kv_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        MG2 base: num_heads=12, head_dim=128.
+        """
+        kv_cache1 = []
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache1.append({
+                "k": torch.zeros([batch_size, self.kv_cache_size, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, self.kv_cache_size, 12, 128], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+            })
+
+        self.kv_cache1 = kv_cache1  # always store the clean cache
+
+    def _initialize_kv_cache_mouse_and_keyboard(self, batch_size, dtype, device):
+        """Initialize per-block KV caches for the MG2 ActionModule's
+        mouse / keyboard attention.  Only needed when the generator is
+        an action-conditioned MG2 model.
+
+        Shapes follow pipeline/causal_diffusion_inference.py:
+          - keyboard: [B, cache_size, 16, 64]          (heads_num=16, head_dim=64)
+          - mouse:    [B * frame_seq, cache_size, 16, 64]  (per spatial token)
+        """
+        kv_cache_mouse = []
+        kv_cache_keyboard = []
+        # Mouse/keyboard kv_cache is sized per-LATENT-FRAME (not per-token).
+        #
+        # IMPORTANT: at TRAIN time we always size to num_max_frames (the full
+        # rollout), regardless of local_attn_size.  Same reason as the main
+        # KV cache (see __init__ comment): if we sized to local_attn_size
+        # here, the eviction code path in ActionModule would mutate
+        # k/v in-place between forward and gradient_checkpointing's
+        # recompute, producing
+        #   CheckpointError: Recomputed values have different metadata
+        #   saved [880, 6, 16, 64]  vs  recomputed [880, 3, 16, 64]
+        # The sliding window is enforced semantically by the attention
+        # mask, NOT by physically truncating the cache.  Inference uses
+        # local_attn_size sizing (single-pass, no backward) — see
+        # pipeline/causal_inference.py for the inference path.
+        kv_cache_size = self.num_max_frames
+        for _ in range(self.num_transformer_blocks):
+            kv_cache_keyboard.append({
+                "k": torch.zeros([batch_size, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+            })
+            kv_cache_mouse.append({
+                "k": torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, 16, 64], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+            })
+        self.kv_cache_mouse = kv_cache_mouse
+        self.kv_cache_keyboard = kv_cache_keyboard
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+
+        For MG2 (I2V, CLIP visual_context) the context length is 257
+        (16x16 ViT tokens + cls token).  For legacy T2V (T5 text) it
+        was 512.  We pick based on whether action_config is present —
+        a proxy for "this is MG2".
+        """
+        ctx_len = 257 if getattr(self, "_use_mg2_ctx", True) else 512
+        crossattn_cache = []
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append({
+                "k": torch.zeros([batch_size, ctx_len, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, ctx_len, 12, 128], dtype=dtype, device=device),
+                "is_init": False
+            })
+        self.crossattn_cache = crossattn_cache

pipeline/teacher_forcing_training.py

@@ -0,0 +1,210 @@
+from utils.wan_wrapper import WanDiffusionWrapper
+from utils.scheduler import SchedulerInterface
+from typing import List, Optional
+import torch
+import torch.distributed as dist
+
+
+class TeacherForcingTrainingPipeline:
+    def __init__(self,
+                 denoising_step_list: List[int],
+                 scheduler: SchedulerInterface,
+                 generator: WanDiffusionWrapper,
+                 num_frame_per_block=3,
+                 independent_first_frame: bool = False,
+                 same_step_across_blocks: bool = False,
+                 last_step_only: bool = False,
+                 num_max_frames: int = 21,
+                 context_noise: int = 0,
+                 spatial_self: bool = True,
+                 **kwargs):
+        super().__init__()
+        self.scheduler = scheduler
+        self.generator = generator
+        self.denoising_step_list = denoising_step_list
+        if self.denoising_step_list[-1] == 0:
+            self.denoising_step_list = self.denoising_step_list[:-1]  # remove the zero timestep for inference
+
+        # Wan specific hyperparameters
+        self.num_transformer_blocks = 30
+        self.frame_seq_length = 1560
+        self.num_frame_per_block = num_frame_per_block
+        self.context_noise = context_noise
+        self.i2v = False
+
+        self.kv_cache1 = None
+        self.kv_cache2 = None
+        self.independent_first_frame = independent_first_frame
+        self.same_step_across_blocks = same_step_across_blocks
+        self.last_step_only = last_step_only
+        self.kv_cache_size = num_max_frames * self.frame_seq_length
+        
+        self.spatial_self = spatial_self
+
+    def generate_and_sync_list(self, num_blocks, num_denoising_steps, device):
+        rank = dist.get_rank() if dist.is_initialized() else 0
+
+        if rank == 0:
+            # Generate random indices
+            indices = torch.randint(
+                low=0,
+                high=num_denoising_steps,
+                size=(num_blocks,),
+                device=device
+            )
+            # In our training, self.last_step_only is False
+            if self.last_step_only:
+                indices = torch.ones_like(indices) * (num_denoising_steps - 1)
+        else:
+            indices = torch.empty(num_blocks, dtype=torch.long, device=device)
+
+        dist.broadcast(indices, src=0)  # Broadcast the random indices to all ranks
+        return indices.tolist()
+
+    def inference_with_trajectory(
+            self,
+            noise: torch.Tensor,
+            clean_image_or_video: torch.Tensor, # same shape as noise
+            initial_latent: Optional[torch.Tensor] = None,
+            return_sim_step: bool = False,
+            **conditional_dict
+    ) -> torch.Tensor:
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        if not self.independent_first_frame or (self.independent_first_frame and initial_latent is not None):
+            # If the first frame is independent and the first frame is provided, then the number of frames in the
+            # noise should still be a multiple of num_frame_per_block
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        else:
+            # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames + num_input_frames  # add the initial latent frames
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+        
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks 
+        num_denoising_steps = len(self.denoising_step_list)
+        exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device)
+        start_gradient_frame_index = num_output_frames - 21 # always 0 as long as we train 21 latent frames
+        if start_gradient_frame_index != 0:
+            raise NotImplementedError("start_gradient_frame_index is always 0 as long as we train 21 latent frames")
+        
+        if self.spatial_self:
+            noisy_input = noise
+            # Step 3.1: Spatial denoising loop
+            # Such a loop corresponds to the truncated denoising algorithm:
+            #    T -> \tau_1 -> \tau_2 ->...-> \tau —— enable grad ——> 0
+            # For many-step model, we certainly cannot use this method, but for 4-step DMD, 
+            # we can inherit it for a fair comaprison. Note that as long as the conditions 
+            # are clean GT rather than self-generated frames, we can perform TF. So this 
+            # method does not conflict with TF in the frame- dimension.
+            for index, current_timestep in enumerate(self.denoising_step_list):
+                # self.same_step_across_blocks is True
+                if self.same_step_across_blocks:
+                    exit_flag = (index == exit_flags[0])
+                else:
+                    raise NotImplementedError('Here t is a scalar denoting that all chunks are at the same t, but in the future we may set t a tensor denoting different chunks')  # Only backprop at the randomly selected timestep (consistent across all ranks)
+                timestep = torch.ones(
+                    [batch_size, self.num_frame_per_block*num_blocks],
+                    device=noise.device,
+                    dtype=torch.int64) * current_timestep
+
+                if not exit_flag:
+                    with torch.no_grad():
+                        _,denoised_pred = self.generator(
+                            noisy_image_or_video=noisy_input,
+                            conditional_dict=conditional_dict,
+                            timestep=timestep,
+                            clean_x = clean_image_or_video
+                        )
+                        next_timestep = self.denoising_step_list[index + 1]
+                        noisy_input = self.scheduler.add_noise(
+                            denoised_pred.flatten(0, 1),
+                            torch.randn_like(denoised_pred.flatten(0, 1)),
+                            next_timestep * torch.ones(
+                                [batch_size * self.num_frame_per_block*num_blocks], device=noise.device, dtype=torch.long)
+                        ).unflatten(0, denoised_pred.shape[:2])
+                else:
+                    # for getting real output
+                    # with torch.set_grad_enabled(current_start_frame >= start_gradient_frame_index):
+                    # enable grad
+                    
+                    _,output = self.generator(
+                                noisy_image_or_video=noisy_input,
+                                conditional_dict=conditional_dict,
+                                timestep=timestep,
+                                clean_x = clean_image_or_video
+                            )
+                    break
+            # ======================= SF -> TF modification ends  ============================
+            
+            # Step 3.5: Return the denoised timestep
+            if not self.same_step_across_blocks: # Useless, never met
+                denoised_timestep_from, denoised_timestep_to = None, None
+            # T -> \tau_1 -> \tau_2 ->...-> \tau —— enable grad ——> 0
+            # denoised_timestep_from = \tau
+            # denoised_timestep_to = next timestep smaller than \tau
+            # These are just engineering tricks
+            # to align DMD timestep sampling with the actual denoising range used by the generator
+            elif exit_flags[0] == len(self.denoising_step_list) - 1:
+                # corner case when \tau is the smallest non-zero timestep
+                denoised_timestep_to = 0
+                denoised_timestep_from = 1000 - torch.argmin(
+                    (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+            else:
+                denoised_timestep_to = 1000 - torch.argmin(
+                    (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0] + 1].cuda()).abs(), dim=0).item()
+                denoised_timestep_from = 1000 - torch.argmin(
+                    (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+        else:
+            print('add noise from gt')
+            current_timestep = self.denoising_step_list[exit_flags[0]]
+            timestep = torch.ones(
+                [batch_size, self.num_frame_per_block*num_blocks],
+                device=noise.device,
+                dtype=torch.int64) * current_timestep
+            
+            noisy_input = self.scheduler.add_noise(
+                        clean_image_or_video,
+                        torch.randn_like(clean_image_or_video),
+                        timestep,
+                    )
+            assert clean_image_or_video.shape == noisy_input.shape
+            
+            
+            _,output = self.generator(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        clean_x = clean_image_or_video
+                    )
+            
+            # T -> \tau_1 -> \tau_2 ->...-> \tau —— enable grad ——> 0
+            # denoised_timestep_from = \tau
+            # denoised_timestep_to = next timestep smaller than \tau
+            # These are just engineering tricks
+            # to align DMD timestep sampling with the actual denoising range used by the generator
+            if exit_flags[0] == len(self.denoising_step_list) - 1:
+                # corner case when \tau is the smallest non-zero timestep
+                denoised_timestep_to = 0
+                denoised_timestep_from = 1000 - torch.argmin(
+                    (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+            else:
+                denoised_timestep_to = 1000 - torch.argmin(
+                    (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0] + 1].cuda()).abs(), dim=0).item()
+                denoised_timestep_from = 1000 - torch.argmin(
+                    (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+
+            
+        if return_sim_step: # False
+            return output, denoised_timestep_from, denoised_timestep_to, exit_flags[0] + 1
+
+        return output, denoised_timestep_from, denoised_timestep_to
+
+    

requirements.txt

@@ -1,16 +1,15 @@
 torchvision
-einops
 omegaconf
+einops
+numpy
+pillow
+tqdm
+imageio
+imageio-ffmpeg
+safetensors
 diffusers
 transformers
 accelerate
-safetensors
-sentencepiece
-numpy
-pillow
 ftfy
 regex
-tqdm
-imageio[ffmpeg]
-imageio-ffmpeg
-https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt211-cu130-cp312/flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl
+sentencepiece