MiniWorld camera-controlled world model demo

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/deck.png filter=lfs diff=lfs merge=lfs -text
+examples/garden.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,56 @@
 ---
-title: Miniworld Simulator
-emoji: 🐨
-colorFrom: purple
-colorTo: purple
+title: MiniWorld Simulator
+emoji: 🌍
+colorFrom: indigo
+colorTo: red
 sdk: gradio
 sdk_version: 6.22.0
-python_version: '3.12'
 app_file: app.py
+python_version: "3.12"
+short_description: Drive a video world model with a camera path from one frame
+startup_duration_timeout: 45m
+license: apache-2.0
+models:
+  - zhaoyian01/MiniWorld
+  - Wan-AI/Wan2.2-TI2V-5B
+tags:
+  - world-model
+  - video-generation
+  - camera-control
+  - image-to-video
+  - streaming-generation
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MiniWorld — camera-controlled world model
+
+Interactive demo of [**MiniWorld**](https://huggingface.co/papers/2608.01127)
+(`zhaoyian01/MiniWorld`, 1B RealEstate10K checkpoint): give it a single frame
+and a procedural camera path, and it autoregressively rolls the world forward.
+
+This Space reproduces the authors' reference inference route
+([`zhao-yian/MiniWorld`](https://github.com/zhao-yian/MiniWorld)):
+
+```bash
+python -m miniworld.sample --dataset re10k \
+  --init_image <frame> --custom_camera_trajectory orbit_right \
+  --wm_model 1B --total_len 64 --trajectory_magnitude 3.0
+```
+
+* 240×320 frames, Wan2.2 VAE (16× spatial / 4× temporal), 48 latent channels.
+* AR-diffusion streaming sampler: `df_chunk_size=2`, `df_ardiff_step=5`,
+  8 in-flight chunks, 24 cached chunks, 1 attention-sink frame, CFG 2.0,
+  100 sampling steps (the streaming schedule caps the effective per-chunk
+  steps at 40) — i.e. exactly `scripts/sample_re10k.sh`.
+* Camera conditioning is the repo's procedural trajectory builder turned into
+  ray-encoding features (`freq=15`, unnormalized translations), so no
+  ground-truth poses or reference video are needed.
+* Streaming causal VAE decode-on-commit, matching the reference pipeline.
+
+The action-conditioned DROID checkpoint is intentionally not exposed: its
+conditioning requires per-dataset `q01/q99` action normalization statistics that
+only ship with the LeRobot DROID dataset, so there is no faithful dataset-free
+input for it.
+
+Example frames are the first frames of the authors' own RealEstate10K rollout
+grid (`assets/demo_re10k.mp4` in the model repo).

app.py

@@ -0,0 +1,498 @@
+"""MiniWorld — camera-controlled video world model simulator (ZeroGPU).
+
+Mirrors the authors' reference inference path
+
+    python -m miniworld.sample --dataset re10k --custom_camera_trajectory ...
+
+one-to-one: a single init image is Wan2.2-VAE-encoded into the clean seed
+latent, a procedural camera path is turned into ray-encoding conditioning, and
+the AR-diffusion denoiser rolls the world forward chunk-by-chunk with a
+position-bounded streaming KV cache and streaming VAE decode.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import tempfile
+import time
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # noqa: E402  (must precede any torch / CUDA work)
+
+import gradio as gr  # noqa: E402
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+from einops import rearrange  # noqa: E402
+from huggingface_hub import hf_hub_download  # noqa: E402
+from PIL import Image  # noqa: E402
+
+# MiniWorld checkpoints are plain `torch.save` dicts that carry a `meta` blob of
+# plain-python objects next to the tensors, so they need the full unpickler.
+_ORIG_TORCH_LOAD = torch.load
+
+
+def _torch_load(*args, **kwargs):
+    kwargs.setdefault("weights_only", False)
+    return _ORIG_TORCH_LOAD(*args, **kwargs)
+
+
+torch.load = _torch_load
+
+from miniworld.conditioning.actions import (  # noqa: E402
+    ConditioningConfig,
+    build_cond_seq_for_batch,
+)
+from miniworld.conditioning.trajectories import build_custom_trajectory  # noqa: E402
+from miniworld.denoiser import DenoiserConfig, build_denoiser_from_mode  # noqa: E402
+from miniworld.vae.codec import StreamingVAEDecoder, vae_encode  # noqa: E402
+from miniworld.vae.wan22_vae import Wan2_2_VAE  # noqa: E402
+
+# --------------------------------------------------------------------------- #
+#                     Constants (match scripts/sample_re10k.sh)               #
+# --------------------------------------------------------------------------- #
+MINIWORLD_REPO = "zhaoyian01/MiniWorld"
+MINIWORLD_CKPT = "MiniWorld_1b_re10k.pt"
+VAE_REPO = "Wan-AI/Wan2.2-TI2V-5B"
+VAE_FILE = "Wan2.2_VAE.pth"
+
+RESIZE_H, RESIZE_W = 240, 320
+SPATIAL_DOWNSAMPLE = 16
+LATENT_CHANNELS = 48
+POSE_ENC_FREQ = 15
+DF_CHUNK_SIZE = 2
+DF_ARDIFF_STEP = 5
+STREAM_INFLIGHT_CHUNKS = 8
+STREAM_MAX_CACHE_CHUNKS = 24
+STREAM_SINK_SIZE = 1
+SAMPLE_HISTORY_LEN = 1
+SAVE_FPS = 8
+MAX_SEED = np.iinfo(np.int32).max
+
+H_LAT, W_LAT = RESIZE_H // SPATIAL_DOWNSAMPLE, RESIZE_W // SPATIAL_DOWNSAMPLE
+
+TRAJECTORIES = [
+    "orbit_right",
+    "orbit_left",
+    "pan_right",
+    "pan_left",
+    "forward",
+    "backward",
+    "tilt_up",
+    "tilt_down",
+    "spiral",
+    "zoom_in",
+    "zoom_out",
+    "static",
+]
+
+# --------------------------------------------------------------------------- #
+#                              Model construction                             #
+# --------------------------------------------------------------------------- #
+print("Fetching Wan2.2 VAE ...", flush=True)
+vae_path = hf_hub_download(VAE_REPO, VAE_FILE)
+print("Fetching MiniWorld-1B (RealEstate10K) ...", flush=True)
+ckpt_path = hf_hub_download(MINIWORLD_REPO, MINIWORLD_CKPT)
+
+_ckpt = torch.load(ckpt_path, map_location="cpu")
+_weights = _ckpt.get("ema_model", _ckpt.get("model"))
+_meta = _ckpt.get("meta", {}) or {}
+if _weights is None:
+    raise RuntimeError("MiniWorld checkpoint has neither `ema_model` nor `model`")
+
+
+def _resolve_latent_frames() -> int:
+    for key in ("latent_frames", "trained_num_frames"):
+        val = int(_meta.get(key, 0) or 0)
+        if val > 0:
+            return val
+    freqs = _weights.get("net.feat_rope.freqs_cos")
+    tokens_per_frame = H_LAT * W_LAT
+    if freqs is not None and freqs.shape[0] % tokens_per_frame == 0:
+        return int(freqs.shape[0] // tokens_per_frame)
+    raise RuntimeError("Cannot determine the checkpoint's latent frame count")
+
+
+LATENT_FRAMES = _resolve_latent_frames()
+WM_MODEL = str(_meta.get("wm_model") or "1B")
+TRAINED_NUM_FRAMES = int(_meta.get("trained_num_frames", 0) or 0) or LATENT_FRAMES
+MAX_TOTAL_LEN = TRAINED_NUM_FRAMES
+print(
+    f"[Checkpoint] wm_model={WM_MODEL} latent_frames={LATENT_FRAMES} "
+    f"trained_num_frames={TRAINED_NUM_FRAMES}",
+    flush=True,
+)
+
+denoiser = build_denoiser_from_mode(
+    DenoiserConfig(
+        wm_model=WM_MODEL,
+        latent_size=(H_LAT, W_LAT),
+        latent_channels=LATENT_CHANNELS,
+        latent_frames=LATENT_FRAMES,
+        wm_mlp_ratio=4.0,
+        wm_use_qknorm=True,
+        wm_use_checkpoint=False,
+        cond_dim=4 * 6 * 2 * POSE_ENC_FREQ,
+        cond_per_token=True,
+        adaln_mode="adaln_lora",
+        cond_dropout_prob=0.1,
+        timestep_baseshift=2.667,
+        timestep_shift=-1.0,
+        num_sampling_steps=100,
+        cfg_scale=2.0,
+        cfg_interval_min=0.2,
+        cfg_interval_max=1.0,
+        df_chunk_size=DF_CHUNK_SIZE,
+        df_ardiff_step=DF_ARDIFF_STEP,
+    )
+).eval()
+denoiser.load_state_dict(_weights, strict=True)
+denoiser.trained_num_frames = TRAINED_NUM_FRAMES
+print("[Checkpoint] loaded: all keys matched", flush=True)
+del _ckpt, _weights
+
+denoiser = denoiser.to("cuda")
+
+vae = Wan2_2_VAE(vae_pth=vae_path, device="cuda")
+vae.model.requires_grad_(False)
+vae.model.eval()
+
+_COND_CFG = ConditioningConfig(
+    use_pose_cond=True, use_action_cond=False, pose_enc_freq=POSE_ENC_FREQ
+)
+
+
+# --------------------------------------------------------------------------- #
+#                                  Helpers                                    #
+# --------------------------------------------------------------------------- #
+def _prepare_init_frame(image) -> torch.Tensor:
+    """PIL / ndarray -> ``(H, W, C)`` float32 in [-1, 1] (== `load_init_image`)."""
+    if image is None:
+        raise gr.Error("Please provide an initial frame.")
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    arr = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
+    img = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
+    if tuple(img.shape[-2:]) != (RESIZE_H, RESIZE_W):
+        img = torch.nn.functional.interpolate(
+            img, size=(RESIZE_H, RESIZE_W), mode="bilinear", align_corners=False
+        )
+    return img.squeeze(0).permute(1, 2, 0).contiguous() * 2.0 - 1.0
+
+
+def _write_mp4(frames: np.ndarray, fps: int) -> str:
+    import imageio.v2 as imageio
+
+    path = os.path.join(tempfile.mkdtemp(), "miniworld.mp4")
+    writer = imageio.get_writer(
+        path,
+        fps=fps,
+        codec="libx264",
+        quality=8,
+        macro_block_size=1,
+        ffmpeg_params=["-pix_fmt", "yuv420p"],
+    )
+    try:
+        for frame in frames:
+            writer.append_data(frame)
+    finally:
+        writer.close()
+    return path
+
+
+def _rollout_tflops(total_len: int, steps: int) -> float:
+    """Replay the streaming schedule to cost a rollout in DiT TFLOPs.
+
+    The AR-diffusion schedule is not linear in ``total_len`` (short rollouts
+    that fit inside the in-flight window run the *full* sampler), so the ZeroGPU
+    reservation is derived from the same bookkeeping the sampler does.
+    """
+    chunk = DF_CHUNK_SIZE
+    ar = DF_ARDIFF_STEP
+    inflight = STREAM_INFLIGHT_CHUNKS
+    max_cache_frames = STREAM_MAX_CACHE_CHUNKS * chunk
+    total_chunks = (total_len + chunk - 1) // chunk
+    eff = steps if total_chunks <= inflight else min(steps, inflight * ar)
+
+    prev = [0] * total_chunks
+    masks = []
+    n_rows = 0
+    while any(p != eff for p in prev):
+        row = [0] * total_chunks
+        for i in range(total_chunks):
+            row[i] = prev[i] + 1 if (i == 0 or prev[i - 1] == eff) else row[i - 1] - ar
+            row[i] = max(0, min(eff, row[i]))
+        masks.append([row[i] != prev[i] for i in range(total_chunks)])
+        prev = row
+        n_rows += 1
+        if n_rows > 4000:  # safety valve
+            break
+
+    terminal = min(inflight, total_chunks)
+    committed = 0
+    cache_frames = 0
+    tflops = 0.0
+    # per-forward TFLOPs for a 1B DiT: 0.6 per query frame (linear layers) plus
+    # 0.01548 per (query frame x key frame) (attention), at 300 tokens/frame.
+    for step in range(n_rows):
+        if terminal < total_chunks and masks[step][terminal]:
+            terminal += 1
+        win_sc = max(0, terminal - inflight)
+        while committed < win_sc:
+            frames = min((committed + 1) * chunk, total_len) - committed * chunk
+            tflops += 2 * frames * (0.6 + 0.01548 * (cache_frames + frames))
+            cache_frames = min(cache_frames + frames, max_cache_frames)
+            committed += 1
+        if terminal <= win_sc:
+            continue
+        q_frames = min(terminal * chunk, total_len) - win_sc * chunk
+        tflops += 2 * q_frames * (0.6 + 0.01548 * (cache_frames + q_frames))
+    return tflops
+
+
+# Calibrated against measured wall-clock on ZeroGPU (see README).
+_TFLOPS_PER_SEC = 90.0
+_VAE_SEC_PER_LATENT_FRAME = 0.10
+_FIXED_OVERHEAD_SEC = 15.0
+
+
+def _duration(*args, **kwargs) -> int:
+    total_len, steps = 32, 100
+    if len(args) >= 4:
+        total_len = int(args[3])
+    if len(args) >= 8:
+        steps = int(args[7])
+    total_len = int(kwargs.get("total_len", total_len))
+    steps = int(kwargs.get("num_sampling_steps", steps))
+    total_len = max(4, min(total_len, 64))
+    est = (
+        _FIXED_OVERHEAD_SEC
+        + _rollout_tflops(total_len, steps) / _TFLOPS_PER_SEC
+        + _VAE_SEC_PER_LATENT_FRAME * total_len
+    )
+    return int(min(400, math.ceil(est * 1.15)))
+
+
+# --------------------------------------------------------------------------- #
+#                                 Inference                                   #
+# --------------------------------------------------------------------------- #
+@spaces.GPU(duration=_duration)
+@torch.no_grad()
+def simulate(
+    image,
+    trajectory: str = "orbit_right",
+    magnitude: float = 3.0,
+    total_len: int = 32,
+    seed: int = 0,
+    randomize_seed: bool = True,
+    cfg_scale: float = 2.0,
+    num_sampling_steps: int = 100,
+    focal_norm: float = 0.5,
+    progress=gr.Progress(track_tqdm=True),
+):
+    init_frame = _prepare_init_frame(image)
+    total_len = max(4, min(int(total_len), MAX_TOTAL_LEN))
+
+    if randomize_seed:
+        seed = int(np.random.randint(0, MAX_SEED))
+    seed = int(seed) % (MAX_SEED + 1)
+
+    device = torch.device("cuda")
+    denoiser.cfg_scale = float(cfg_scale)
+    denoiser.steps = int(num_sampling_steps)
+
+    videos = init_frame.unsqueeze(0).unsqueeze(0).to(device)  # (1, 1, H, W, C)
+    poses = (
+        build_custom_trajectory(
+            trajectory,
+            num_frames=4 * (total_len - 1) + 1,
+            focal_norm=float(focal_norm),
+            magnitude=float(magnitude),
+        )
+        .unsqueeze(0)
+        .to(device)
+    )
+
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    noise = torch.randn(
+        1, LATENT_CHANNELS, total_len, H_LAT, W_LAT,
+        generator=generator, dtype=torch.float32,
+    ).to(device)
+
+    start = time.perf_counter()
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+        latents = vae_encode(
+            vae, rearrange(videos, "b t h w c -> b c t h w").contiguous()
+        )
+        _, c_latent, _, h_lat, w_lat = latents.shape
+        full_latents = latents.new_zeros(1, c_latent, total_len, h_lat, w_lat)
+        full_latents[:, :, :1] = latents[:, :, :1]
+
+        cond_seq = build_cond_seq_for_batch(
+            cfg=_COND_CFG,
+            poses=poses,
+            actions=None,
+            t_latent=total_len,
+            h_lat=h_lat,
+            w_lat=w_lat,
+        )
+
+        _, pred_rgb = denoiser.generate_eval_latents_streaming(
+            full_latents,
+            cond_seq,
+            total_len=total_len,
+            history_len=SAMPLE_HISTORY_LEN,
+            max_cache_chunks=STREAM_MAX_CACHE_CHUNKS,
+            inflight_chunks=STREAM_INFLIGHT_CHUNKS,
+            sink_frames=STREAM_SINK_SIZE,
+            stream_decoder=StreamingVAEDecoder(vae),
+            noise=noise.to(full_latents.dtype),
+        )
+    elapsed = time.perf_counter() - start
+
+    video = ((pred_rgb[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1.0) * 127.5).to(
+        torch.uint8
+    )
+    frames = video.cpu().numpy()
+    path = _write_mp4(frames, SAVE_FPS)
+    n = int(frames.shape[0])
+    print(f"[Timing] total_len={total_len} steps={num_sampling_steps}: "
+          f"{elapsed:.2f}s (reserved {_duration(None, trajectory, magnitude, total_len, seed, False, cfg_scale, num_sampling_steps)}s)",
+          flush=True)
+    return (
+        path,
+        seed,
+        f"**{n} frames** @ {SAVE_FPS} fps ({n / SAVE_FPS:.1f}s) · "
+        f"{total_len} latent frames · `{trajectory}` @ {magnitude:g} · "
+        f"seed `{seed}` · {elapsed:.1f}s of GPU time",
+    )
+
+
+# --------------------------------------------------------------------------- #
+#                                     UI                                      #
+# --------------------------------------------------------------------------- #
+CSS = "#col-container { max-width: 1060px; margin: 0 auto; }"
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            f"""
+            # 🌍 MiniWorld · camera-controlled world model
+
+            Hand MiniWorld-1B **one frame and a camera path** and it rolls the
+            world forward autoregressively — no text prompt, no reference video,
+            no ground-truth poses. A position-bounded streaming KV cache plus
+            causal Wan2.2 VAE decoding keep the horizon open, so a rollout can
+            run to {4 * (MAX_TOTAL_LEN - 1) + 1} frames from a
+            {TRAINED_NUM_FRAMES}-latent-frame checkpoint.
+
+            Model: [`zhaoyian01/MiniWorld`](https://huggingface.co/zhaoyian01/MiniWorld)
+            (RealEstate10K, {WM_MODEL}) · Paper:
+            [MiniWorld: Democratizing the Training of Video World Models from Scratch](https://huggingface.co/papers/2608.01127)
+            · Code: [zhao-yian/MiniWorld](https://github.com/zhao-yian/MiniWorld)
+            """
+        )
+
+        with gr.Row():
+            with gr.Column():
+                image = gr.Image(
+                    label="Initial frame",
+                    type="pil",
+                    height=270,
+                    sources=["upload", "clipboard"],
+                )
+                trajectory = gr.Dropdown(
+                    label="Camera trajectory",
+                    choices=TRAJECTORIES,
+                    value="orbit_right",
+                )
+                magnitude = gr.Slider(
+                    label="Motion strength",
+                    minimum=0.5,
+                    maximum=8.0,
+                    step=0.5,
+                    value=3.0,
+                    info="3.0 gives clear, stable motion over a 32-frame rollout. "
+                    "Scale up for longer rollouts; lower it if late frames smear.",
+                )
+                total_len = gr.Slider(
+                    label="Rollout length (latent frames)",
+                    minimum=20,
+                    maximum=MAX_TOTAL_LEN,
+                    step=4,
+                    value=32,
+                    info=f"Each latent frame decodes to 4 RGB frames at {SAVE_FPS} fps; "
+                    f"{MAX_TOTAL_LEN} → {4 * (MAX_TOTAL_LEN - 1) + 1} frames.",
+                )
+                run_button = gr.Button("Simulate", variant="primary")
+
+            with gr.Column():
+                result = gr.Video(
+                    label="Rollout", autoplay=True, loop=True, height=270
+                )
+                info = gr.Markdown()
+
+        with gr.Accordion("Advanced settings", open=False):
+            with gr.Row():
+                seed = gr.Slider(
+                    label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            with gr.Row():
+                cfg_scale = gr.Slider(
+                    label="Guidance scale (CFG)",
+                    minimum=1.0,
+                    maximum=5.0,
+                    step=0.1,
+                    value=2.0,
+                )
+                num_sampling_steps = gr.Slider(
+                    label="Sampling steps",
+                    minimum=20,
+                    maximum=100,
+                    step=10,
+                    value=100,
+                    info="Effective steps per chunk are capped by the streaming "
+                    "schedule at in-flight chunks × AR step = 40.",
+                )
+            focal_norm = gr.Slider(
+                label="Normalized focal length",
+                minimum=0.3,
+                maximum=1.2,
+                step=0.05,
+                value=0.5,
+                info="0.5 matches typical RealEstate10K intrinsics; smaller = wider FOV.",
+            )
+
+        gr.Examples(
+            examples=[
+                ["examples/kitchen.png", "orbit_right", 3.0, 32],
+                ["examples/deck.png", "forward", 3.0, 32],
+                ["examples/garden.png", "pan_left", 3.0, 32],
+            ],
+            inputs=[image, trajectory, magnitude, total_len],
+            outputs=[result, seed, info],
+            fn=simulate,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+    gr.on(
+        triggers=[run_button.click],
+        fn=simulate,
+        inputs=[
+            image,
+            trajectory,
+            magnitude,
+            total_len,
+            seed,
+            randomize_seed,
+            cfg_scale,
+            num_sampling_steps,
+            focal_norm,
+        ],
+        outputs=[result, seed, info],
+    )
+
+demo.queue().launch(mcp_server=True)

miniworld/__init__.py

@@ -0,0 +1,6 @@
+"""Public MiniWorld package."""
+
+from miniworld.miniworld import MiniWorldModel, MiniWorldModels
+
+__all__ = ["MiniWorldModel", "MiniWorldModels"]
+

miniworld/conditioning/__init__.py

@@ -0,0 +1,6 @@
+"""Conditioning utilities for MiniWorld."""
+
+from miniworld.conditioning.actions import ConditioningConfig, build_cond_seq_from_actions, build_cond_seq_for_batch
+
+__all__ = ["ConditioningConfig", "build_cond_seq_from_actions", "build_cond_seq_for_batch"]
+

miniworld/conditioning/actions.py

@@ -0,0 +1,58 @@
+"""Conditioning helpers for robot-action and camera-pose inputs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from miniworld.conditioning.poses import compute_ray_encoding, downsample_poses_to_latent
+
+
+@dataclass
+class ConditioningConfig:
+    """Configuration for converting raw batch conditions into model conditions."""
+
+    use_pose_cond: bool = False
+    use_action_cond: bool = False
+    pose_enc_freq: int = 15
+
+
+def build_cond_seq_from_actions(actions: torch.Tensor) -> torch.Tensor:
+    """Map raw per-frame actions to WAN-VAE latent-frame conditioning.
+
+    WAN-VAE compresses time as 4x+1. Latent frame 0 is the seed frame and uses a
+    zero action slot; each following latent frame receives the four raw actions
+    that drive its decoded frame group.
+    """
+    batch, num_actions, action_dim = actions.shape
+    if num_actions % 4 != 0:
+        raise ValueError(f"Expected action length 4n for real actions, got {num_actions}")
+    num_generated_latents = num_actions // 4
+    cond = actions.new_zeros(batch, num_generated_latents + 1, 4 * action_dim)
+    if num_generated_latents > 0:
+        cond[:, 1:, :] = actions.reshape(batch, num_generated_latents, 4 * action_dim)
+    return cond
+
+
+def build_cond_seq_for_batch(
+    *,
+    cfg: ConditioningConfig,
+    poses: Optional[torch.Tensor],
+    actions: Optional[torch.Tensor],
+    t_latent: int,
+    h_lat: int,
+    w_lat: int,
+) -> torch.Tensor:
+    """Build the MiniWorld conditioning tensor for one batch."""
+    if cfg.use_pose_cond:
+        if poses is None:
+            raise ValueError("Pose conditioning requested but batch has no poses")
+        latent_poses = downsample_poses_to_latent(poses, t_latent)
+        return compute_ray_encoding(latent_poses, h_lat, w_lat, freq=cfg.pose_enc_freq, normalize_trans=False)
+    if cfg.use_action_cond:
+        if actions is None:
+            raise ValueError("Action conditioning requested but batch has no actions")
+        return build_cond_seq_from_actions(actions)
+    raise ValueError("Either pose or action conditioning must be enabled")

miniworld/conditioning/poses.py

@@ -0,0 +1,218 @@
+"""Camera pose -> ray-encoding utilities for pose-conditioned world model.
+
+Adapted from `GeometryForcing/utils/geometry_utils.py` with the following
+simplifications:
+  * Only the bits needed for ``ray_encoding`` conditioning are kept (the
+    variant the user picked as the best-performing one).
+  * ``rays`` accepts independent ``(h_res, w_res)`` so non-square latents
+    (e.g. 15x20) are supported without distorting the intrinsics.
+
+All functions follow this convention:
+  * Raw camera pose layout: ``(B, T, 16)`` = ``[K(4), R(9 + T(3))]`` where the
+    first 4 columns are normalised intrinsics ``(fx, fy, px, py)`` (pixel-coords
+    divided by image size) and the last 12 columns are a flattened ``3x4``
+    world-to-camera extrinsics matrix in row-major.
+  * Ray encoding output: ``(B, T, 180, H_lat, W_lat)`` (6 ray dims * 2 trig fns
+    * 15 NeRF frequencies = 180). This matches what `DiT3DPose` consumes when
+    ``conditioning_type=ray_encoding``.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Tuple
+
+import torch
+from einops import einsum, rearrange, repeat
+
+
+def _split_pose16(raw_poses: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """``(B, T, 16)`` -> ``R (B, T, 3, 3)``, ``T (B, T, 3)``, ``K (B, T, 4)``."""
+    assert raw_poses.shape[-1] == 16, f"expected 16-dim pose, got {raw_poses.shape[-1]}"
+    K, RT = raw_poses.split([4, 12], dim=-1)
+    RT = rearrange(RT, "b t (i j) -> b t i j", i=3, j=4)
+    R = RT[..., :3, :3]
+    T = RT[..., :3, 3]
+    return R, T, K
+
+
+def _normalize_by_first(R: torch.Tensor, T: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Re-express all poses so the first frame is the world origin."""
+    R_ref = R[:, 0]  # (B, 3, 3)
+    T_ref = T[:, 0]  # (B, 3)
+    R_inv = rearrange(R_ref, "b i j -> b j i")
+    R_new = einsum(R, R_inv, "b t i j1, b j1 j2 -> b t i j2")
+    T_new = T - einsum(R_new, T_ref, "b t i j, b j -> b t i")
+    return R_new, T_new
+
+
+def _compute_rays(
+    R: torch.Tensor,
+    T: torch.Tensor,
+    K: torch.Tensor,
+    h_res: int,
+    w_res: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Per-pixel ray origin / direction in world coords.
+
+    Args:
+        R: ``(B, T, 3, 3)`` world->cam rotation.
+        T: ``(B, T, 3)`` world->cam translation.
+        K: ``(B, T, 4)`` normalised intrinsics ``(fx, fy, px, py)``.
+        h_res, w_res: target ray grid resolution (independent so non-square
+            latents are handled correctly).
+
+    Returns:
+        origin: ``(B, T, H, W, 3)``
+        direction: ``(B, T, H, W, 3)`` (unnormalised; norm encodes depth scale)
+    """
+    device, dtype = K.device, K.dtype
+
+    coord_w, coord_h = torch.meshgrid(
+        torch.linspace(0, w_res - 1, w_res, device=device, dtype=dtype),
+        torch.linspace(0, h_res - 1, h_res, device=device, dtype=dtype),
+        indexing="xy",
+    )  # (H, W) each
+    coord_w = rearrange(coord_w, "h w -> 1 1 h w") + 0.5
+    coord_h = rearrange(coord_h, "h w -> 1 1 h w") + 0.5
+
+    # Normalised K -> pixel-space K (separate W / H scaling for non-square grids).
+    fx = (K[..., 0] * w_res).view(*K.shape[:-1], 1, 1)  # (B, T, 1, 1)
+    fy = (K[..., 1] * h_res).view(*K.shape[:-1], 1, 1)
+    px = (K[..., 2] * w_res).view(*K.shape[:-1], 1, 1)
+    py = (K[..., 3] * h_res).view(*K.shape[:-1], 1, 1)
+
+    x = (coord_w - px) / fx
+    y = (coord_h - py) / fy
+    z = torch.ones_like(x)
+    direction = torch.stack([x, y, z], dim=-1)  # (B, T, H, W, 3)
+
+    R_inv = rearrange(R, "b t i j -> b t j i")
+    direction = einsum(R_inv, direction, "b t i j, b t h w j -> b t h w i")
+
+    origin = -einsum(R_inv, T, "b t i j, b t j -> b t i")
+    origin = repeat(origin, "b t i -> b t h w i", h=h_res, w=w_res).clone()
+    return origin, direction
+
+
+def _normalize_translation_scale(T: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    """Per-clip translation-scale normalisation (CameraCtrl / lingbot style).
+
+    Monocular SfM poses (e.g. RealEstate10K) have an arbitrary, per-clip metric
+    scale, so raw camera translations vary wildly in magnitude across clips.
+    Since the ray origin feeds NeRF frequency encoding ``sin(2^k pi x)`` -- which
+    is very sensitive to the absolute magnitude of ``x`` -- this inconsistency
+    hurts learning.  We rescale each clip so its largest camera displacement is
+    ~1, making camera motion scale-invariant across clips.
+
+    Args:
+        T: ``(B, S, 3)`` camera translations, already expressed relative to the
+            first frame (so the first frame sits at the origin).
+        eps: guard so static / near-static clips (max norm ~ 0) are left
+            unchanged ("only normalize when moving").
+
+    Returns:
+        ``(B, S, 3)`` translations divided by the per-clip max translation norm.
+    """
+    max_norm = torch.norm(T, dim=-1).amax(dim=1, keepdim=True)  # (B, 1)
+    scale = torch.where(max_norm > eps, max_norm, torch.ones_like(max_norm))
+    return T / scale.unsqueeze(-1)
+
+
+def _nerf_pos_encoding(x: torch.Tensor, freq: int) -> torch.Tensor:
+    """NeRF-style sin/cos positional encoding along the last dim."""
+    scale = (
+        2 ** torch.linspace(0, freq - 1, freq, device=x.device, dtype=x.dtype)
+        * math.pi
+    )
+    encoding = rearrange(x[..., None] * scale, "b t h w i s -> b t h w (i s)")
+    return torch.sin(torch.cat([encoding, encoding + 0.5 * math.pi], dim=-1))
+
+
+@torch.no_grad()
+@torch.autocast(device_type="cuda", enabled=False)  # always fp32 for geometry
+def compute_ray_encoding(
+    raw_poses: torch.Tensor,
+    h_lat: int,
+    w_lat: int,
+    freq: int = 15,
+    normalize_trans: bool = False,
+) -> torch.Tensor:
+    """End-to-end raw poses -> ray-encoding feature volume.
+
+    Args:
+        raw_poses: either ``(B, T, 16)`` or ``(B, T, K, 16)``. MiniWorld's RE10K
+            pipeline uses ``K=4`` poses inside each WAN-VAE latent chunk.
+        h_lat, w_lat: latent spatial size (= model input H, W after VAE).
+        freq: NeRF frequency count. ``freq=15`` gives ``6 * 2 * 15 = 180``
+            channels per pose.
+        normalize_trans: if True, rescale each clip's camera
+            translations so the largest displacement is ~1 (see
+            ``_normalize_translation_scale``). Disabled by default to match
+            DFoT's RealEstate10K preprocessing; static clips are untouched.
+
+    Returns:
+        ``(B, T, K * 6 * 2 * freq, H_lat, W_lat)`` float32. The ``K`` poses are
+        ray-encoded independently then concatenated along the channel axis
+        (so the spatial ``y_embedder`` sees ``K * 180`` channels). For the
+        common ``(B, T, 16)`` input ``K=1`` and the output channel count is
+        ``180``.
+    """
+    assert raw_poses.dim() in (3, 4), (
+        f"raw_poses must be (B, T, 16) or (B, T, K, 16); got {raw_poses.shape}"
+    )
+    raw_poses = raw_poses.float()
+    if raw_poses.dim() == 3:
+        b, t_lat, _ = raw_poses.shape
+        k_per_lat = 1
+        flat = raw_poses  # (B, T, 16)
+    else:
+        b, t_lat, k_per_lat, _ = raw_poses.shape
+        # Flatten K into the time axis so we can reuse the single-pose pipeline
+        # (one shared normalisation anchor = first pose in the sequence).
+        flat = raw_poses.reshape(b, t_lat * k_per_lat, 16)
+
+    R, T, K = _split_pose16(flat)
+    R, T = _normalize_by_first(R, T)
+    if normalize_trans:
+        T = _normalize_translation_scale(T)
+    origin, direction = _compute_rays(R, T, K, h_res=h_lat, w_res=w_lat)
+    enc = torch.cat(
+        [
+            _nerf_pos_encoding(origin, freq),
+            _nerf_pos_encoding(direction, freq),
+        ],
+        dim=-1,
+    )  # (B, T*K, H, W, 6 * 2 * freq)
+
+    if k_per_lat == 1:
+        return rearrange(enc, "b t h w c -> b t c h w").contiguous()
+    return rearrange(
+        enc, "b (t k) h w c -> b t (k c) h w", t=t_lat, k=k_per_lat,
+    ).contiguous()
+
+
+def downsample_poses_to_latent(
+    raw_poses: torch.Tensor,
+    t_latent: int,
+) -> torch.Tensor:
+    """Map per-raw-frame poses to four poses per WAN-style latent frame.
+
+    The causal WAN VAE encodes ``T_raw = 4*(T_lat-1)+1`` raw frames into
+    ``T_lat`` latents with the temporal grouping:
+        * latent 0      -> raw [0]
+        * latent j (>0) -> raw [4j-3, 4j-2, 4j-1, 4j]
+
+    Latent 0 has only raw[0], so it is duplicated four times to keep the output
+    shape consistent with action conditioning: ``(B, T_lat, 4, 16)``.
+    """
+    idx_per_latent = [[0, 0, 0, 0]]
+    for j in range(1, t_latent):
+        idx_per_latent.append([4 * j - 3, 4 * j - 2, 4 * j - 1, 4 * j])
+    idx_flat = [i for chunk in idx_per_latent for i in chunk]
+    assert raw_poses.shape[1] > max(idx_flat), (
+        f"raw_poses has only {raw_poses.shape[1]} frames; need at least "
+        f"{max(idx_flat) + 1} to build {t_latent} latent poses."
+    )
+    b = raw_poses.shape[0]
+    return raw_poses[:, idx_flat].view(b, t_latent, 4, 16).contiguous()

miniworld/conditioning/trajectories.py

@@ -0,0 +1,299 @@
+"""Procedural camera-trajectory utilities for pose-conditioned WM inference.
+
+Builds ``(T, 16)`` pose tensors compatible with
+``model.pose_utils.compute_ray_encoding`` so we can drive the world model
+with **any** camera path (no GT video / poses required).
+
+Pose layout (per frame, matches ``model.pose_utils._split_pose16``):
+    ``[fx, fy, px, py, R(9, row-major), T(3)]`` = ``[K(4), RT(12)]``
+where K is normalised (intrinsics divided by image W / H) and
+(R, T) is world->camera (OpenCV convention: x=right, y=down, z=forward).
+
+All trajectories here **start at the identity pose** at frame 0
+(R=I, T=0). ``compute_ray_encoding`` will re-anchor the first frame as
+the world origin anyway, so only relative camera motion w.r.t. frame 0
+ever reaches the model.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+from pathlib import Path
+from typing import Callable, Tuple
+
+import numpy as np
+import torch
+
+
+###############################################################################
+#                       Low-level rotation / look-at                          #
+###############################################################################
+
+
+def _yaw(a: float) -> np.ndarray:
+    """Rotation around world-down axis (=+y). ``a > 0`` -> camera pans right
+    (the world's +x moves toward camera-forward)."""
+    c, s = math.cos(a), math.sin(a)
+    # camera basis in world: right=(c,0,-s), down=(0,1,0), forward=(s,0,c)
+    # R_w2c rows = [right; down; forward]
+    return np.array(
+        [
+            [c, 0.0, -s],
+            [0.0, 1.0, 0.0],
+            [s, 0.0, c],
+        ],
+        dtype=np.float64,
+    )
+
+
+def _pitch(a: float) -> np.ndarray:
+    """Rotation around world-right axis (=+x). ``a > 0`` -> camera tilts up."""
+    c, s = math.cos(a), math.sin(a)
+    # camera basis: right=(1,0,0), down=(0,c,s), forward=(0,-s,c)
+    return np.array(
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, c, s],
+            [0.0, -s, c],
+        ],
+        dtype=np.float64,
+    )
+
+
+def _look_at(
+    eye: np.ndarray,
+    target: np.ndarray,
+    world_up: np.ndarray = np.array([0.0, -1.0, 0.0]),
+) -> Tuple[np.ndarray, np.ndarray]:
+    """world->camera (R, T) for a camera at ``eye`` looking at ``target``.
+
+    OpenCV convention: camera frame is (right, down, forward) = (+x, +y, +z).
+    ``world_up`` points along the world's "visual up" direction; in OpenCV
+    image-y is down, so the canonical world-up is ``(0, -1, 0)``.
+    """
+    fwd = target - eye
+    n = float(np.linalg.norm(fwd))
+    if n < 1e-8:
+        # Degenerate: fall back to identity orientation.
+        R = np.eye(3, dtype=np.float64)
+        T = -R @ eye
+        return R, T
+    fwd = fwd / n
+
+    down_world = -world_up
+    right = np.cross(down_world, fwd)
+    rn = float(np.linalg.norm(right))
+    if rn < 1e-6:
+        # forward parallel to up -> pick any perpendicular right
+        right = np.array([1.0, 0.0, 0.0])
+        if abs(float(fwd @ right)) > 0.99:
+            right = np.array([0.0, 0.0, 1.0])
+    else:
+        right = right / rn
+    down = np.cross(fwd, right)
+
+    R = np.stack([right, down, fwd], axis=0).astype(np.float64)  # world->cam rows
+    T = -R @ eye
+    return R, T
+
+
+###############################################################################
+#                       Trajectory primitives                                 #
+###############################################################################
+
+
+def _build_RT(traj_fn: Callable[[float], Tuple[np.ndarray, np.ndarray]], num_frames: int) -> np.ndarray:
+    """Sample ``traj_fn(s)`` at ``num_frames`` evenly-spaced ``s`` in [0, 1]
+    and return a ``(num_frames, 12)`` row-major flattened RT.
+
+    ``traj_fn(0.0)`` is expected to return the identity pose (R=I, T=0) so
+    frame 0 anchors the world origin cleanly.
+    """
+    out = np.zeros((num_frames, 12), dtype=np.float32)
+    for i in range(num_frames):
+        s = i / max(num_frames - 1, 1)
+        R, T = traj_fn(s)
+        out[i, :9] = R.reshape(-1)
+        out[i, 9:] = T.reshape(-1)
+    return out
+
+
+SUPPORTED_TRAJECTORIES = (
+    "static",
+    "forward",
+    "backward",
+    "pan_left",
+    "pan_right",
+    "tilt_up",
+    "tilt_down",
+    "orbit_right",
+    "orbit_left",
+    "spiral",
+    "zoom_in",
+    "zoom_out",
+)
+
+
+def build_custom_trajectory(
+    traj_type: str,
+    num_frames: int,
+    focal_norm: float = 0.7,
+    magnitude: float = 1.0,
+) -> torch.Tensor:
+    """Build a ``(num_frames, 16)`` pose sequence for a named procedural path.
+
+    Args:
+        traj_type: one of :data:`SUPPORTED_TRAJECTORIES`.
+        num_frames: number of *raw* frames the pose sequence must cover
+            (i.e. ``eval_t_dataset = 4*(total_len-1)+2`` -- e.g. 126 for
+            ``total_len=32``). ``compute_ray_encoding`` indexes into this.
+        focal_norm: normalised focal length (``fx = fy = focal_norm``). RE10K
+            videos typically sit near 0.5-1.0; smaller = wider FOV.
+        magnitude: global scaling. At ``magnitude=1.0`` the defaults are:
+              * translate ~0.5 units.
+              * rotate up to 30 deg (pan / tilt).
+              * orbit / spiral: 60 deg arc on a radius-``magnitude`` circle.
+              * zoom: focal scales linearly to 1.5x (in) / 0.67x (out).
+            For **rawscale** RE10K checkpoints (``normalize_trans=False``),
+            ``magnitude=1.0`` looks nearly static.
+
+    Returns:
+        ``(num_frames, 16)`` float32 tensor on cpu.
+    """
+    if traj_type not in SUPPORTED_TRAJECTORIES:
+        raise ValueError(
+            f"Unknown trajectory '{traj_type}'. "
+            f"Supported: {SUPPORTED_TRAJECTORIES}"
+        )
+
+    I = np.eye(3, dtype=np.float64)
+    Z = np.zeros(3, dtype=np.float64)
+    PI = math.pi
+
+    # ----- translation / rotation only paths (K is constant) -----
+    def f_static(s):
+        return I, Z
+
+    def f_forward(s):
+        # camera center moves to (0, 0, +d) in world; T = -R @ c = (0,0,-d)
+        d = 0.5 * magnitude * s
+        return I, np.array([0.0, 0.0, -d])
+
+    def f_backward(s):
+        d = 0.5 * magnitude * s
+        return I, np.array([0.0, 0.0, d])
+
+    def f_pan_right(s):
+        return _yaw(+(PI / 6) * magnitude * s), Z
+
+    def f_pan_left(s):
+        return _yaw(-(PI / 6) * magnitude * s), Z
+
+    def f_tilt_up(s):
+        return _pitch(+(PI / 6) * magnitude * s), Z
+
+    def f_tilt_down(s):
+        return _pitch(-(PI / 6) * magnitude * s), Z
+
+    # Orbit / spiral are anchored so that frame 0 is exactly (R=I, T=0):
+    # the camera starts at the world origin looking at a target one unit
+    # away along +z (= (0, 0, r)), and pivots around that target while
+    # keeping it in view.
+
+    def _orbit(sign: float):
+        def fn(s):
+            # Radius scales with magnitude so rawscale mag>>1 also translates
+            # farther (angle alone on r=1 caps |T| at ~2).
+            a = sign * (PI / 3) * s
+            r = 1.0 * magnitude
+            target = np.array([0.0, 0.0, r])
+            eye = np.array([r * math.sin(a), 0.0, r * (1.0 - math.cos(a))])
+            return _look_at(eye, target)
+        return fn
+
+    def f_spiral(s):
+        a = (PI / 3) * s
+        r = 1.0 * magnitude
+        target = np.array([0.0, 0.0, r])
+        eye = np.array(
+            [r * math.sin(a), -0.2 * magnitude * s, r * (1.0 - math.cos(a))]
+        )
+        return _look_at(eye, target)
+
+    rt_dispatch = {
+        "static": f_static,
+        "forward": f_forward,
+        "backward": f_backward,
+        "pan_left": f_pan_left,
+        "pan_right": f_pan_right,
+        "tilt_up": f_tilt_up,
+        "tilt_down": f_tilt_down,
+        "orbit_left": _orbit(-1.0),
+        "orbit_right": _orbit(+1.0),
+        "spiral": f_spiral,
+        # zoom paths keep RT = identity, vary K instead
+        "zoom_in": f_static,
+        "zoom_out": f_static,
+    }
+    RT = _build_RT(rt_dispatch[traj_type], num_frames)  # (T, 12)
+
+    # ----- intrinsics K (T, 4) -----
+    K = np.zeros((num_frames, 4), dtype=np.float32)
+    for i in range(num_frames):
+        s = i / max(num_frames - 1, 1)
+        if traj_type == "zoom_in":
+            scale = 1.0 + 0.5 * magnitude * s  # up to 1.5x at magnitude=1
+        elif traj_type == "zoom_out":
+            scale = 1.0 / (1.0 + 0.5 * magnitude * s)  # down to ~0.67x
+        else:
+            scale = 1.0
+        K[i, 0] = focal_norm * scale  # fx
+        K[i, 1] = focal_norm * scale  # fy
+        K[i, 2] = 0.5  # px at image center
+        K[i, 3] = 0.5  # py at image center
+
+    pose16 = np.concatenate([K, RT], axis=1)  # (T, 16)
+    return torch.from_numpy(pose16).to(torch.float32)
+
+
+###############################################################################
+#                       Init-image loading                                    #
+###############################################################################
+
+
+def load_init_image(path: str, resize_h: int, resize_w: int) -> torch.Tensor:
+    """Load a single image (or first frame of a video) and return it as a
+    ``(H, W, C)`` float32 tensor in ``[-1, 1]`` -- the same format that
+    ``SimpleVideoDataset`` produces for a single frame.
+
+    Supported inputs:
+      * PIL-readable still images (.jpg / .png / .webp / ...).
+      * Video files (.mp4 / .mov / ...). First frame is taken.
+    """
+    p = Path(path)
+    assert p.exists(), f"--init_image not found: {path}"
+    suffix = p.suffix.lower()
+
+    if suffix in (".mp4", ".mov", ".avi", ".mkv", ".webm"):
+        import torchvision.io
+        frames, _, _ = torchvision.io.read_video(
+            os.fspath(p), pts_unit="sec", output_format="TCHW",
+        )
+        if frames.shape[0] == 0:
+            raise RuntimeError(f"--init_image video decoded 0 frames: {path}")
+        img = frames[0:1].float() / 255.0  # (1, C, H, W)
+    else:
+        from PIL import Image
+        with Image.open(p) as im:
+            im = im.convert("RGB")
+            arr = np.asarray(im, dtype=np.float32) / 255.0  # (H, W, C)
+        img = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)  # (1, C, H, W)
+
+    if tuple(img.shape[-2:]) != (resize_h, resize_w):
+        img = torch.nn.functional.interpolate(
+            img, size=(resize_h, resize_w), mode="bilinear", align_corners=False,
+        )
+    img = img.squeeze(0).permute(1, 2, 0).contiguous()  # (H, W, C)
+    img = img * 2.0 - 1.0
+    return img

miniworld/denoiser.py

@@ -0,0 +1,1073 @@
+from __future__ import annotations
+
+import time
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from miniworld.vae.codec import print0 as _print0
+
+
+class IncrementalTimesteps:
+    """AR-Diffusion style combinatorial timestep sampler.
+    """
+
+    def __init__(self, F: int, T: int):
+        self.F = F
+        self.T = T
+
+        mat = torch.zeros((T, F), dtype=torch.float64)
+        for t in range(T):
+            mat[t, F - 1] = 1
+        for f in range(F - 2, -1, -1):
+            mat[T - 1, f] = 1
+            for t in range(T - 2, -1, -1):
+                mat[t, f] = mat[t + 1, f] + mat[t, f + 1]
+        self.mat_s = mat.numpy()
+
+        mat = torch.zeros((T, F), dtype=torch.float64)
+        for t in range(T):
+            mat[t, 0] = 1
+        for f in range(1, F):
+            mat[0, f] = 1
+            for t in range(1, T):
+                mat[t, f] = mat[t - 1, f] + mat[t, f - 1]
+        self.mat_e = mat.numpy()
+
+    def sample_stepseq_from_mid(self):
+        timesteps = torch.zeros(self.F, dtype=torch.long)
+        cur_f = np.random.randint(self.F)
+        timesteps[cur_f] = np.random.randint(self.T)
+
+        for f in range(cur_f - 1, -1, -1):
+            candidate_weights = self.mat_e[: int(timesteps[f + 1]) + 1, f]
+            prob_sequence = candidate_weights / candidate_weights.sum()
+            cur_step = np.random.choice(range(0, int(timesteps[f + 1]) + 1), p=prob_sequence)
+            timesteps[f] = int(cur_step)
+
+        for f in range(cur_f + 1, self.F):
+            candidate_weights = self.mat_s[int(timesteps[f - 1]):, f]
+            prob_sequence = candidate_weights / candidate_weights.sum()
+            cur_step = np.random.choice(range(int(timesteps[f - 1]), self.T), p=prob_sequence)
+            timesteps[f] = int(cur_step)
+        return timesteps
+
+
+class DenoiserConfig:
+    def __init__(self, **kwargs):
+        self.wm_model: str = "1B"
+        self.latent_size: int = 16
+        self.latent_channels: int = 48
+        self.latent_frames: int = 9
+
+        self.wm_mlp_ratio: float = 4.0
+        self.wm_use_qknorm: bool = True
+        self.wm_use_checkpoint: bool = True
+        self.cond_dim: int = 0
+        # When True, y is treated as per-token spatial conditioning
+        # ``(B, T, cond_dim, H_lat, W_lat)`` (e.g. ray-encoding for camera
+        # pose). When False (default), y is the per-frame ``(B, T, cond_dim)``
+        # latent-action condition.
+        self.cond_per_token: bool = False
+
+        # Structured action/pose dropout for classifier-free guidance training.
+        self.adaln_mode: str = "adaln_lora"
+        self.cond_dropout_prob: float = 0.0
+        # Route the true first latent frame (seed / initial observation, no
+        # preceding action) through the learned null_action (action mode only).
+        self.action_null_first: bool = True
+
+        # Long-video finetune / streaming inference metadata.
+        # ``trained_num_frames`` defaults to ``latent_frames`` and is saved in the
+        # ckpt meta so streaming inference can assert the active window
+        # (cache + in-flight) never exceeds it.
+        self.trained_num_frames: int = -1  # -1 => fallback to latent_frames at runtime
+
+        # Training timesteps: t = sigmoid(P_mean + P_std * z), z ~ N(0, 1).
+        # P_std <= 0 falls back to uniform.
+        self.P_mean: float = 0.0
+        self.P_std: float = 1.0
+        self.timestep_shift: float = -1.0  # -1 = auto from per-chunk token count; >0 = manual override
+        self.timestep_baseshift: float = 2.667  # shift at _REF_TOKENS; see Denoiser.__init__
+
+        # sample
+        self.num_sampling_steps: int = 50
+        self.cfg_scale: float = 1.0
+        self.cfg_interval_min: float = 0.1
+        self.cfg_interval_max: float = 1.0
+
+        self.df_chunk_size: int = 2
+        self.df_train_time_bins: int = 50
+        self.df_ardiff_step: int = 1
+
+        for k, v in kwargs.items():
+            if hasattr(self, k):
+                setattr(self, k, v)
+    
+
+
+class Denoiser(nn.Module):
+    """World model Denoiser.
+
+    Args:
+
+    Return:
+        diffusion loss
+    """
+
+    def __init__(self, cfg: DenoiserConfig) -> None:
+        super().__init__()
+
+        self.cfg = cfg
+        from miniworld.miniworld import MiniWorldModels
+
+        if cfg.wm_model not in MiniWorldModels:
+            raise ValueError(
+                f"Unknown MiniWorld model {cfg.wm_model!r}. "
+                f"Choose one of {sorted(MiniWorldModels)}."
+            )
+        self.net = MiniWorldModels[cfg.wm_model](
+            input_size=cfg.latent_size,
+            in_channels=cfg.latent_channels,
+            num_frames=cfg.latent_frames,
+            mlp_ratio=cfg.wm_mlp_ratio,
+            use_qknorm=cfg.wm_use_qknorm,
+            use_rope=True,
+            use_abs_pos=False,
+            use_checkpoint=cfg.wm_use_checkpoint,
+            cond_dim=cfg.cond_dim,
+            cond_per_token=cfg.cond_per_token,
+            adaln_mode=cfg.adaln_mode,
+            cond_dropout_prob=cfg.cond_dropout_prob,
+            action_null_first=cfg.action_null_first,
+        )
+
+        self.trained_num_frames = (
+            cfg.trained_num_frames if cfg.trained_num_frames > 0 else cfg.latent_frames
+        )
+
+        # SD3-style timestep shift, scaled by the tokens denoised jointly at one
+        # noise level (a single chunk). Under diffusion forcing every chunk has
+        # its own t, so a longer window must not move the training t distribution.
+        latent_size = cfg.latent_size
+        if isinstance(latent_size, (tuple, list)):
+            h_lat, w_lat = latent_size
+        else:
+            h_lat = w_lat = int(latent_size)
+        n_tokens = int(cfg.df_chunk_size) * h_lat * w_lat
+        if cfg.timestep_shift > 0:
+            self.timestep_shift = cfg.timestep_shift
+        else:
+            _REF_TOKENS = 600  # 2 * 15 * 20: chunk_size=2 at 240x320 @16x downsample
+            # Nothing derives the 2.667 default; it is the knob for how hard
+            # training leans towards high-noise timesteps.
+            self.timestep_shift = cfg.timestep_baseshift * (n_tokens / _REF_TOKENS) ** 0.5
+        _print0(f"[Denoiser] latent=({cfg.latent_frames}, {h_lat}, {w_lat}), "
+                f"chunk_size={cfg.df_chunk_size}, chunk_tokens={n_tokens}, "
+                f"timestep_shift={self.timestep_shift:.4f}")
+        # Scheme B: when the net is MiniWorld and its internal structured
+        # dropout is enabled, CFG uses the model's *learned null* token for the
+        # unconditional branch (train + infer), instead of zeroing cond_seq.
+        # This keeps the train-time null and infer-time uncond identical.
+        self.use_model_null_cfg = cfg.cond_dropout_prob > 0.0
+        # Filled by generate_* so callers can report pipeline throughput.
+        self.last_eval_meta: Dict[str, object] = {}
+
+        self.steps = cfg.num_sampling_steps
+        self.cfg_scale = cfg.cfg_scale
+        self.cfg_interval_min = cfg.cfg_interval_min
+        self.cfg_interval_max = cfg.cfg_interval_max
+        self.df_chunk_size = int(cfg.df_chunk_size)
+        self.df_train_time_bins = max(2, int(cfg.df_train_time_bins))
+        self.df_ardiff_step = int(cfg.df_ardiff_step)
+        if self.df_ardiff_step <= 0:
+            raise ValueError("df_ardiff_step must be > 0 for MiniWorld AR-diffusion")
+        self.condition_noise_max_t = 0.05
+        self.P_mean = float(cfg.P_mean)
+        self.P_std = float(cfg.P_std)
+        self._df_train_step_samplers: Dict[int, IncrementalTimesteps] = {}
+
+    def _set_last_eval_meta(
+        self,
+        *,
+        path: str,
+        total_chunks: int,
+        n_ctx_chunks: int,
+        num_outer_steps: int,
+        effective_steps: Optional[int] = None,
+    ) -> None:
+        self.last_eval_meta = {
+            "path": path,
+            "total_chunks": int(total_chunks),
+            "n_ctx_chunks": int(n_ctx_chunks),
+            "gen_chunks": int(max(0, total_chunks - n_ctx_chunks)),
+            "num_outer_steps": int(num_outer_steps),
+            "ar_step": int(self.df_ardiff_step),
+            "chunk_size": int(self.df_chunk_size),
+            "effective_steps": (
+                int(effective_steps) if effective_steps is not None else int(self.steps)
+            ),
+        }
+
+    def _make_uncond(self, cond_seq: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Return ``(cond_for_uncond, cond_drop)`` for the CFG unconditional pass.
+
+        When structured dropout is trained, keep the real conditioning tensor
+        and force the model's learned null token via ``cond_drop=all-True``.
+        """
+        if self.use_model_null_cfg:
+            b = cond_seq.shape[0]
+            return cond_seq, torch.ones(b, dtype=torch.bool, device=cond_seq.device)
+        return torch.zeros_like(cond_seq), None
+
+    def drop_cond(self, cond_seq: torch.Tensor) -> torch.Tensor:
+        return cond_seq
+
+    def _build_chunk_slices(self, t: int) -> List[slice]:
+        if t <= 0:
+            raise ValueError(f"t must be positive, got {t}")
+        chunk_size = self.df_chunk_size
+        assert chunk_size > 0
+
+        chunk_slices: List[slice] = []
+        start = 0
+        while start < t:
+            end = min(t, start + chunk_size)
+            chunk_slices.append(slice(start, end))
+            start = end
+        return chunk_slices
+
+    def _get_df_train_step_sampler(self, num_chunks: int) -> IncrementalTimesteps:
+        sampler = self._df_train_step_samplers.get(num_chunks)
+        if sampler is None:
+            sampler = IncrementalTimesteps(num_chunks, self.df_train_time_bins)
+            self._df_train_step_samplers[num_chunks] = sampler
+        return sampler
+
+    def _sample_df_chunk_timesteps(self, num_chunks: int, device: torch.device) -> torch.Tensor:
+        if num_chunks <= 0:
+            return torch.zeros(0, device=device, dtype=torch.long)
+        sampler = self._get_df_train_step_sampler(num_chunks)
+        sampled = sampler.sample_stepseq_from_mid()
+        return sampled.to(device=device, dtype=torch.long)
+
+    def _broadcast_chunk_values_to_frames(
+        self,
+        chunk_values: torch.Tensor,
+        chunk_slices: List[slice],
+        t: int,
+    ) -> torch.Tensor:
+        b = chunk_values.shape[0]
+        frame_values = torch.zeros(b, t, device=chunk_values.device, dtype=chunk_values.dtype)
+        for chunk_idx, chunk_slice in enumerate(chunk_slices):
+            frame_values[:, chunk_slice] = chunk_values[:, chunk_idx].unsqueeze(1)
+        return frame_values
+
+    def _build_async_step_index_matrix(
+        self,
+        total_chunks: int,
+        num_steps: int,
+        device: torch.device,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # for ar diffusion inference
+        if total_chunks <= 0:
+            step_index = torch.full((1, total_chunks), num_steps, device=device, dtype=torch.long)
+            update_mask = torch.zeros((1, total_chunks), device=device, dtype=torch.bool)
+            return step_index, update_mask
+
+        ar_step = int(self.df_ardiff_step)
+        pre_row = torch.zeros(total_chunks, dtype=torch.long)
+        rows: List[torch.Tensor] = []
+        masks: List[torch.Tensor] = []
+
+        while not torch.all(pre_row == num_steps):
+            new_row = torch.zeros_like(pre_row)
+            for idx in range(total_chunks):
+                if idx == 0 or pre_row[idx - 1] == num_steps:
+                    new_row[idx] = pre_row[idx] + 1
+                else:
+                    new_row[idx] = new_row[idx - 1] - ar_step
+            new_row = new_row.clamp(0, num_steps)
+            masks.append(new_row != pre_row)
+            rows.append(new_row.clone())
+            pre_row = new_row
+
+        step_index = torch.stack(rows, dim=0).to(device=device)
+        update_mask = torch.stack(masks, dim=0).to(device=device)
+        return step_index, update_mask
+
+    def _build_chunk_sampling_schedule(
+        self,
+        total_chunks: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_context_chunks: int = 1,
+        effective_steps: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # for ar diffusion inference
+        num_steps = effective_steps if effective_steps is not None else int(self.steps)
+        ts = torch.linspace(1.0, 0.0, num_steps + 1, device=device, dtype=dtype)
+        ts = self.shift_timestep(ts, self.timestep_shift)
+        step_index, update_mask = self._build_async_step_index_matrix(
+            total_chunks=total_chunks,
+            num_steps=num_steps,
+            device=device,
+        )
+
+        current_lookup = torch.cat([ts[:1], ts[:-1]], dim=0)
+        next_lookup = ts
+        t_chunk = current_lookup[step_index]
+        t_next_chunk = next_lookup[step_index]
+        for ci in range(min(n_context_chunks, total_chunks)):
+            t_chunk[:, ci] = 0
+            t_next_chunk[:, ci] = 0
+            update_mask[:, ci] = False
+        return t_chunk, t_next_chunk, update_mask
+
+    def _compute_fifo_valid_intervals(
+        self,
+        update_mask: torch.Tensor,
+        total_chunks: int,
+        max_chunks_in_window: int,
+    ) -> List[Tuple[int, int]]:
+        """Compute per-step FIFO window bounds (chunk-level).
+
+        Mirrors AR-Diffusion ``fifoddim.py``'s ``valid_interval`` logic.
+        The window starts covering chunks ``[0, max_chunks_in_window)`` and
+        slides right by one chunk each time a new chunk at the window
+        boundary becomes active (``update_mask`` turns True).
+
+        Returns a list of ``(start_chunk, end_chunk)`` tuples, one per
+        outer iteration.
+        """
+        terminal = min(max_chunks_in_window, total_chunks)
+        intervals: List[Tuple[int, int]] = []
+        for i in range(update_mask.shape[0]):
+            if terminal < total_chunks and bool(update_mask[i, terminal]):
+                terminal += 1
+            start = max(0, terminal - max_chunks_in_window)
+            intervals.append((start, terminal))
+        return intervals
+
+
+    def _build_diffusion_forcing_timesteps(
+        self,
+        b: int,
+        t: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        """Build per-frame timesteps for diffusion forcing training.
+
+        Clean-context length is sampled per example:
+          Mode A (p=0.5): only the first frame is clean
+          Mode B (p=0.5): the entire first chunk is clean
+
+        Returns:
+            t_frame:      (B, T)
+            chunk_slices: list of slices
+            chunk_t:      (B, num_chunks)
+            clean_mask:   (B, T)  1 on clean context frames, else 0
+        """
+        chunk_slices = self._build_chunk_slices(t)
+        num_chunks = len(chunk_slices)
+        chunk_t = torch.zeros(b, num_chunks, device=device, dtype=dtype)
+        scale = float(max(self.df_train_time_bins - 1, 1))
+
+        for sample_idx in range(b):
+            if num_chunks <= 0:
+                continue
+            seq1 = self._sample_df_chunk_timesteps(num_chunks, device=device)
+            chunk_t[sample_idx, :] = seq1.to(dtype=dtype) / scale
+
+        chunk_t = self.logit_normal_warp(chunk_t)
+        chunk_t = self.shift_timestep(chunk_t, self.timestep_shift)
+        t_frame = self._broadcast_chunk_values_to_frames(chunk_t, chunk_slices, t)
+
+        clean_mask = torch.zeros(b, t, device=device, dtype=dtype)
+        cond_noise = self.sample_condition_t((b,), device=device, dtype=dtype)
+
+        for sample_idx in range(b):
+            if num_chunks <= 0:
+                continue
+            if torch.rand(1).item() < 0.5:
+                # Mode A: only first frame is clean
+                t_frame[sample_idx, 0] = cond_noise[sample_idx]
+                clean_mask[sample_idx, 0] = 1.0
+            else:
+                # Mode B: entire first chunk is clean
+                first_sl = chunk_slices[0]
+                t_frame[sample_idx, first_sl] = cond_noise[sample_idx]
+                clean_mask[sample_idx, first_sl] = 1.0
+
+        return t_frame, chunk_slices, chunk_t, clean_mask
+
+    def _get_df_action_guidance_scale(self, chunk_t: torch.Tensor) -> torch.Tensor:
+        # Apply cfg_scale when chunk_t is inside
+        # (cfg_interval_min, cfg_interval_max]; else 1.0.  Upper bound
+        # is inclusive so that the first denoising step (chunk_t == 1.0) still
+        # receives CFG, matching diffusion-forcing guidance semantics.
+        low = self.cfg_interval_min
+        high = self.cfg_interval_max
+        interval_mask = (chunk_t <= high) & ((low == 0.0) | (chunk_t > low))
+        action_scale = torch.where(
+            interval_mask,
+            torch.full_like(chunk_t, self.cfg_scale),
+            torch.ones_like(chunk_t),
+        )
+        return action_scale
+
+    def logit_normal_warp(self, u: torch.Tensor) -> torch.Tensor:
+        """Give the training timesteps a logit-normal density.
+
+        ``u`` is the uniform bin grid from ``IncrementalTimesteps``. The
+        logit-normal inverse CDF is monotone, so it reshapes the density without
+        disturbing the non-decreasing noise ordering across chunks.
+        """
+        if self.P_std <= 0.0:
+            return u
+        z = torch.special.ndtri(u.to(torch.float64))
+        return torch.sigmoid(self.P_mean + self.P_std * z).to(dtype=u.dtype)
+
+    @staticmethod
+    def shift_timestep(t: torch.Tensor, shift: float) -> torch.Tensor:
+        """SD3-style timestep shift: t' = shift*t / (1 + (shift-1)*t).
+        Maps [0,1]->[0,1]; shift>1 biases towards higher t (more noise)."""
+        if shift == 1.0:
+            return t
+        return shift * t / (1.0 + (shift - 1.0) * t)
+
+    def sample_condition_t(self, shape: Tuple[int, ...], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        if self.condition_noise_max_t <= 0.0:
+            return torch.zeros(shape, device=device, dtype=dtype)
+        return torch.rand(shape, device=device, dtype=dtype) * self.condition_noise_max_t
+
+    def forward_diffusion_forcing(
+        self,
+        latents: torch.Tensor,
+        cond_seq: torch.Tensor,
+        history_len: int = 1,
+        return_pred: bool = False,
+    ) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # ``history_len`` is accepted for API compatibility with train/sample CLI,
+        # but training clean-context length is sampled via Mode A/B (see
+        # ``_build_diffusion_forcing_timesteps``). Inference uses ``history_len``
+        # in ``generate_eval_latents_streaming``.
+        assert int(history_len) > 0, f"history_len must be > 0, got {history_len}"
+        assert latents.dim() == 5, f"latents must be (B, C, T, H, W), got {latents.shape}"
+        b, _, t, _, _ = latents.shape
+        assert cond_seq.shape[0] == b and cond_seq.shape[1] == t, (
+            f"cond_seq shape {cond_seq.shape} must match (B, T, D) with B={b}, T={t}"
+        )
+
+        cond_seq = self.drop_cond(cond_seq)
+        device = latents.device
+
+        t_frame, _, _, clean_mask = self._build_diffusion_forcing_timesteps(
+            b=b,
+            t=t,
+            device=device,
+            dtype=latents.dtype,
+        )
+
+        noise = torch.randn_like(latents)
+        v_target = latents - noise
+
+        t_view = t_frame.view(b, 1, t, 1, 1)
+        z = (1.0 - t_view) * latents + t_view * noise
+        v_pred = self.net(
+            z,
+            t_frame,
+            cond_seq,
+            temporal_causal=True,
+            chunk_size=self.df_chunk_size,
+        )
+
+        # --- v_loss (per-frame, excluding clean context) ---
+        diff = (v_target - v_pred) ** 2
+        diff = diff.mean(dim=(1, 3, 4))  # (B, T)
+        loss_mask = 1.0 - clean_mask  # 0 on clean frames, 1 on noisy frames
+        v_loss = (diff * loss_mask).sum(dim=1) / loss_mask.sum(dim=1).clamp_min(1.0)
+        v_loss = v_loss.mean()
+        if not return_pred:
+            return v_loss
+
+        x_pred = z + (1.0 - t_view) * v_pred
+        clean_mask_5d = clean_mask.view(b, 1, t, 1, 1)
+        x_pred = x_pred * (1.0 - clean_mask_5d) + latents * clean_mask_5d
+        return v_loss, x_pred.detach(), t_frame.max(dim=1).values.detach()
+
+
+class DiffusionForcingDenoiser(Denoiser):
+    def forward(
+        self,
+        latents: torch.Tensor,
+        cond_seq: torch.Tensor,
+        history_len: int = 1,
+        return_pred: bool = False,
+    ) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Run one diffusion-forcing training step.
+
+        Returns the scalar loss, or ``(loss, x_pred, t_noise)`` when
+        ``return_pred`` is set: the detached ``(B, C, T, H, W)`` one-step clean
+        latent and the ``(B,)`` peak noise level, for logging videos.
+        """
+        return super().forward_diffusion_forcing(
+            latents=latents,
+            cond_seq=cond_seq,
+            history_len=history_len,
+            return_pred=return_pred,
+        )
+
+    # ------------------------------------------------------------------
+    #  Streaming AR-diffusion inference with KV cache
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _append_kv_cache(
+        cache: List[Optional[Tuple[torch.Tensor, torch.Tensor]]],
+        new_kv: List[Optional[Tuple[torch.Tensor, torch.Tensor]]],
+    ) -> List[Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        depth = len(cache)
+        out: List[Optional[Tuple[torch.Tensor, torch.Tensor]]] = [None] * depth
+        for i in range(depth):
+            k_new, v_new = new_kv[i]
+            if cache[i] is None:
+                out[i] = (k_new, v_new)
+            else:
+                k_old, v_old = cache[i]
+                out[i] = (
+                    torch.cat([k_old, k_new], dim=-2),
+                    torch.cat([v_old, v_new], dim=-2),
+                )
+        return out
+
+    @staticmethod
+    def _evict_and_shift_cache(
+        cache: List[Optional[Tuple[torch.Tensor, torch.Tensor]]],
+        drop_frames: int,
+        tokens_per_frame: int,
+        rope_module,
+        sink_frames: int = 0,
+    ) -> List[Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        """Evict ``drop_frames`` frames from the cache and renumber positions.
+
+        With ``sink_frames == 0`` (default): drop the leading ``drop_frames``
+        frames and re-rotate the remaining K so positions restart at 0 (pure
+        FIFO sliding window).
+
+        With ``sink_frames > 0`` (StreamingLLM-style attention sink): the first
+        ``sink_frames`` frames are pinned at positions ``[0, sink_frames)`` and
+        never dropped or re-rotated; only the *post-sink* window frames are
+        evicted (oldest first) and shifted down by ``drop_frames`` so they sit
+        contiguously right behind the sink (positions ``[sink_frames, ...)``).
+        Net layout stays contiguous ``[0, cache_frames)`` and inside the trained
+        RoPE range, while the true origin frame(s) stay resident as an anchor.
+        """
+        if drop_frames <= 0:
+            return cache
+        sink_frames = max(0, sink_frames)
+        sink_tokens = sink_frames * tokens_per_frame
+        drop_tokens = drop_frames * tokens_per_frame
+        depth = len(cache)
+        out: List[Optional[Tuple[torch.Tensor, torch.Tensor]]] = [None] * depth
+        for i in range(depth):
+            if cache[i] is None:
+                continue
+            k_old, v_old = cache[i]
+            # sink slice: kept verbatim (no eviction, no RoPE shift).
+            k_sink = k_old[..., :sink_tokens, :]
+            v_sink = v_old[..., :sink_tokens, :]
+            # window slice: drop the oldest ``drop_frames`` right after the sink,
+            # then renumber survivors down by ``drop_frames``.
+            k_rest = k_old[..., sink_tokens + drop_tokens:, :]
+            v_rest = v_old[..., sink_tokens + drop_tokens:, :]
+            if k_rest.numel() > 0:
+                k_rest = rope_module.rope_shift_time(-drop_frames, k_rest)
+            if sink_tokens > 0:
+                k_new = torch.cat([k_sink, k_rest], dim=-2)
+                v_new = torch.cat([v_sink, v_rest], dim=-2)
+            else:
+                k_new, v_new = k_rest, v_rest
+            out[i] = (k_new, v_new)
+        return out
+
+    @torch.no_grad()
+    def generate_eval_latents_streaming(
+        self,
+        latents: torch.Tensor,
+        cond_seq: torch.Tensor,
+        total_len: int,
+        history_len: int = 1,
+        max_cache_chunks: int = 16,
+        inflight_chunks: int = 4,
+        sink_frames: int = 0,
+        stream_decoder=None,
+        collect_stream_timing: bool = False,
+        noise: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        """Streaming AR-diffusion with KV cache (position-bounded, renumbered from 0).
+
+        The active attention window at any time is exactly
+        ``(max_cache_chunks + inflight_chunks) * df_chunk_size`` frames, which must
+        fit inside ``self.trained_num_frames`` to avoid RoPE extrapolation.
+
+        Committed chunks live in a per-block KV cache, logically numbered at
+        temporal positions ``[0, cache_frames)``.  In-flight chunks sit at
+        ``[cache_frames, cache_frames + inflight_frames)``.  When a chunk
+        completes denoising (its ``t`` hits 0) it is committed: we run a t=0
+        forward to obtain its K/V, append them to the cache, and if the cache
+        overflows we drop the leading frames and :meth:`rope_shift_time` the
+        remaining K to renumber positions back to 0.
+
+        When ``stream_decoder`` is provided (Wan2.2 ``StreamingVAEDecoder``), each
+        committed latent chunk is VAE-decoded immediately (decode-on-commit).
+        Concatenating those RGB chunks is bit-exact with batch ``vae_decode`` of
+        the same latents, so clip metrics are unchanged.
+
+        Args:
+            latents: ``(B, C, T_full, H, W)`` -- first ``history_len`` frames are
+                used as clean visual context.
+            cond_seq: ``(B, T_full, D)`` full action sequence.
+            total_len: number of latent frames to produce.
+            history_len: clean context length in frames; must be ``> 0`` and at
+                most ``max_cache_chunks * df_chunk_size`` for full-chunk
+                prefill. When ``history_len % df_chunk_size != 0`` (e.g.
+                image-to-video with ``history_len=1`` and ``df_chunk_size=2``),
+                the leading ``history_len // df_chunk_size`` chunks are
+                pre-filled into the KV cache, and the remaining
+                ``history_len % df_chunk_size`` frames are pinned to ``t=0``
+                inside the first in-flight chunk.
+            max_cache_chunks: max number of committed chunks retained in the
+                cache at any one time.
+            inflight_chunks: number of chunks simultaneously being denoised.
+            sink_frames: StreamingLLM-style attention-sink size in frames. 0
+                (default) = pure sliding window (no resident anchor). >0 pins the
+                first ``sink_frames`` committed frames (the true origin / clean
+                context) at cache positions ``[0, sink_frames)`` permanently;
+                they are never evicted or re-rotated, so a long rollout always
+                retains them as an anchor. Must be <= the cache capacity.
+            stream_decoder: optional streaming VAE decoder with
+                ``begin() / step(latents) / end()``. When set, returns
+                ``(latents, rgb_video)`` with RGB in ``[-1, 1]``; otherwise
+                returns latents only.
+            noise: optional ``(B, C, >=total_len, H, W)`` initial noise; pass a
+                fixed tensor to make repeated rollouts comparable. Sampled from
+                the global RNG when omitted.
+        """
+        del kwargs
+        device = latents.device
+        dtype = latents.dtype
+        net = self.net
+        chunk_size = self.df_chunk_size
+        inflight_frames = inflight_chunks * chunk_size
+        max_cache_frames = max_cache_chunks * chunk_size
+        active_frames = max_cache_frames + inflight_frames
+        sink_frames = max(0, int(sink_frames))
+        assert sink_frames <= max_cache_frames, (
+            f"[StreamingGen] sink_frames={sink_frames} exceeds cache capacity "
+            f"max_cache_frames={max_cache_frames}. Increase stream_max_cache_chunks."
+        )
+
+        trained_num_frames = int(getattr(self, "trained_num_frames", 0))
+        if trained_num_frames <= 0:
+            trained_num_frames = self.cfg.latent_frames
+        assert active_frames <= trained_num_frames, (
+            f"[StreamingGen] active window (cache={max_cache_frames} + "
+            f"inflight={inflight_frames} = {active_frames}) exceeds "
+            f"trained_num_frames={trained_num_frames}. Reduce "
+            f"--stream_max_cache_chunks or --stream_inflight_chunks."
+        )
+        assert not net.use_abs_pos, (
+            "[StreamingGen] requires use_abs_pos=False (relative RoPE only). "
+            "MiniWorld checkpoints should be trained with RoPE-only positioning."
+        )
+        assert self.df_ardiff_step > 0, (
+            "[StreamingGen] requires df_ardiff_step > 0 (AR-diffusion schedule)."
+        )
+
+        total_len = min(total_len, latents.shape[2], cond_seq.shape[1])
+
+        ctx_len = int(history_len)
+        assert ctx_len > 0, f"[StreamingGen] history_len must be > 0, got {ctx_len}"
+        # Sub-chunk ctx (e.g. ctx_len=1, chunk_size=2 for i2v) is supported via
+        # per-frame t=0 pinning inside the first in-flight chunk; we only
+        # pre-fill the *whole* leading chunks into the KV cache. The leftover
+        # ``ctx_len - n_full_ctx_frames`` frames stay in the in-flight window
+        # cache capacity.
+        n_full_ctx_chunks = ctx_len // chunk_size
+        n_full_ctx_frames = n_full_ctx_chunks * chunk_size
+        assert n_full_ctx_frames <= max_cache_frames, (
+            f"[StreamingGen] full-chunk history ({n_full_ctx_frames} frames "
+            f"= {n_full_ctx_chunks} chunks) exceeds cache capacity "
+            f"({max_cache_frames} frames = {max_cache_chunks} chunks). "
+            f"Increase --stream_max_cache_chunks."
+        )
+
+        b, c_ch, _, h, w = latents.shape
+        p_t, p_h, p_w = net.x_embedder.patch_size
+        _, h_total, w_total = net.x_embedder.input_size
+        grid_h = h_total // p_h
+        grid_w = w_total // p_w
+        tokens_per_frame = grid_h * grid_w
+        rope_module = net.feat_rope
+
+        use_cfg = float(self.cfg_scale) > 1.0
+        depth = net.depth
+        cache_cond: List[Optional[Tuple[torch.Tensor, torch.Tensor]]] = [None] * depth
+        cache_uncond: List[Optional[Tuple[torch.Tensor, torch.Tensor]]] = [None] * depth if use_cfg else []
+        cache_frames = 0
+
+        # --- Output buffer ---
+        if noise is None:
+            z_global = torch.randn(b, c_ch, total_len, h, w, device=device, dtype=dtype)
+        else:
+            assert noise.shape[:2] == (b, c_ch) and noise.shape[3:] == (h, w), (
+                f"[StreamingGen] noise shape {tuple(noise.shape)} does not match "
+                f"latents {tuple(latents.shape)}"
+            )
+            assert noise.shape[2] >= total_len, (
+                f"[StreamingGen] noise covers {noise.shape[2]} frames, need {total_len}"
+            )
+            # Cloned because the rollout denoises this buffer in place.
+            z_global = noise[:, :, :total_len].to(device=device, dtype=dtype).clone()
+        if ctx_len > 0:
+            z_global[:, :, :ctx_len] = latents[:, :, :ctx_len]
+
+        timing_enabled = bool(collect_stream_timing)
+        timing_start = None
+        dit_chunk_events: List[Dict[str, object]] = []
+        vae_chunk_events: List[Dict[str, object]] = []
+
+        def _timing_now() -> float:
+            if torch.cuda.is_available():
+                torch.cuda.synchronize(device)
+            return time.perf_counter()
+
+        if timing_enabled:
+            timing_start = _timing_now()
+
+        # --- Optional streaming VAE decode (decode-on-commit) ---
+        rgb_parts: List[torch.Tensor] = []
+        decoded_frames = 0
+
+        def _stream_decode_upto(end_frame: int, *, chunk_idx: Optional[int] = None, step_idx: Optional[int] = None) -> None:
+            nonlocal decoded_frames
+            if stream_decoder is None or end_frame <= decoded_frames:
+                return
+            start_frame = decoded_frames
+            t0 = _timing_now() if timing_enabled else None
+            rgb_parts.append(stream_decoder.step(z_global[:, :, start_frame:end_frame]))
+            t1 = _timing_now() if timing_enabled else None
+            if timing_enabled and timing_start is not None and t0 is not None and t1 is not None:
+                vae_chunk_events.append(
+                    {
+                        "chunk_idx": int(chunk_idx) if chunk_idx is not None else None,
+                        "step_idx": int(step_idx) if step_idx is not None else None,
+                        "start_frame": int(start_frame),
+                        "end_frame": int(end_frame),
+                        "generated": bool(chunk_idx is not None and chunk_idx >= n_full_ctx_chunks),
+                        "start_sec": float(t0 - timing_start),
+                        "end_sec": float(t1 - timing_start),
+                        "duration_sec": float(t1 - t0),
+                    }
+                )
+            decoded_frames = end_frame
+
+        if stream_decoder is not None:
+            stream_decoder.begin()
+
+        try:
+            # --- Pre-fill cache with clean history context ---
+            # Only fully-aligned ctx chunks go into the cache. Sub-chunk leftover
+            # (n_partial_ctx_frames) is pinned via per-frame t=0 inside the first
+            # in-flight chunk, see the in-flight forward block below.
+            if n_full_ctx_frames > 0:
+                ctx_frames = z_global[:, :, :n_full_ctx_frames]
+                ctx_cond = cond_seq[:, :n_full_ctx_frames]
+                ctx_t = torch.zeros(b, n_full_ctx_frames, device=device, dtype=dtype)
+
+                _, kv_cond_ctx = net.forward_with_cache(
+                    ctx_frames, ctx_t, ctx_cond,
+                    past_kv_list=None, current_position_offset=0,
+                    return_kv=True, chunk_size=chunk_size,
+                )
+                cache_cond = list(kv_cond_ctx)
+                if use_cfg:
+                    ctx_uncond, ctx_drop_uncond = self._make_uncond(ctx_cond)
+                    _, kv_uncond_ctx = net.forward_with_cache(
+                        ctx_frames, ctx_t, ctx_uncond,
+                        past_kv_list=None, current_position_offset=0,
+                        return_kv=True, chunk_size=chunk_size, cond_drop=ctx_drop_uncond,
+                    )
+                    cache_uncond = list(kv_uncond_ctx)
+                cache_frames = n_full_ctx_frames
+                # Decode clean context immediately (same order as batch decode).
+                _stream_decode_upto(n_full_ctx_frames)
+
+            # --- Global AR schedule ---
+            # Keep the final partial chunk. Training uses the same chunk layout
+            # (e.g. T=9, chunk_size=2 -> four 2-frame chunks plus one 1-frame
+            # chunk), so dropping it at inference changes the requested video
+            # length and the learned schedule.
+            total_chunks = (total_len + chunk_size - 1) // chunk_size
+            # Residence cap: a chunk can be updated for ~inflight*ar outer steps
+            # before the FIFO window must slide past it. When the *entire*
+            # sequence fits in the inflight window, nothing is force-evicted
+            # mid-denoise, so use the full sampler length (e.g. T=64, 100 steps).
+            residence_cap = inflight_chunks * max(self.df_ardiff_step, 1)
+            if total_chunks <= inflight_chunks:
+                effective_steps = int(self.steps)
+            else:
+                effective_steps = min(int(self.steps), residence_cap)
+            t_chunk_sched, t_next_chunk_sched, chunk_update_mask = (
+                self._build_chunk_sampling_schedule(
+                    total_chunks=total_chunks,
+                    device=device, dtype=dtype,
+                    n_context_chunks=n_full_ctx_chunks,
+                    effective_steps=effective_steps,
+                )
+            )
+            valid_intervals = self._compute_fifo_valid_intervals(
+                chunk_update_mask, total_chunks, max_chunks_in_window=inflight_chunks,
+            )
+
+            num_outer_steps = t_chunk_sched.shape[0]
+            self._set_last_eval_meta(
+                path="streaming",
+                total_chunks=total_chunks,
+                n_ctx_chunks=n_full_ctx_chunks,
+                num_outer_steps=num_outer_steps,
+                effective_steps=effective_steps,
+            )
+            self.last_eval_meta["cfg_enabled"] = bool(use_cfg)
+            self.last_eval_meta["stream_timing_enabled"] = bool(timing_enabled)
+            _print0(f"[StreamingGen] total_len={total_len}, total_chunks={total_chunks}, "
+                    f"ctx_chunks={n_full_ctx_chunks}, chunk_size={chunk_size}, "
+                    f"inflight_chunks={inflight_chunks}, max_cache_chunks={max_cache_chunks}, "
+                    f"trained_num_frames={trained_num_frames}, active_frames={active_frames}, "
+                    f"sink_frames={sink_frames}, "
+                    f"effective_steps={effective_steps}, outer_steps={num_outer_steps}, "
+                    f"ar_step={self.df_ardiff_step}, "
+                    f"stream_decode={stream_decoder is not None}, cfg_enabled={use_cfg}")
+
+            last_win_sc = n_full_ctx_chunks
+            committed_chunks = set(range(n_full_ctx_chunks))
+            # VAE decode is tied to schedule completion (t_next==0), not KV
+            # eviction. With max_cache=0 + full inflight, the window may never
+            # slide, but chunks still finish and should decode immediately.
+            next_decode_ci = n_full_ctx_chunks
+
+            def _decode_finished_chunks(step_idx: int) -> None:
+                nonlocal next_decode_ci
+                if stream_decoder is None:
+                    return
+                while next_decode_ci < total_chunks:
+                    if float(t_next_chunk_sched[step_idx, next_decode_ci]) > 0.0:
+                        break
+                    end_f = min((next_decode_ci + 1) * chunk_size, total_len)
+                    if timing_enabled and timing_start is not None:
+                        t_dit = _timing_now()
+                        dit_chunk_events.append(
+                            {
+                                "chunk_idx": int(next_decode_ci),
+                                "step_idx": int(step_idx),
+                                "end_frame": int(end_f),
+                                "generated": bool(next_decode_ci >= n_full_ctx_chunks),
+                                "complete_sec": float(t_dit - timing_start),
+                            }
+                        )
+                    _stream_decode_upto(end_f, chunk_idx=next_decode_ci, step_idx=step_idx)
+                    _print0(
+                        f"[StreamingGen]   decoded chunk {next_decode_ci}/{total_chunks} | "
+                        f"latent_frames={decoded_frames}/{total_len} | "
+                        f"step={step_idx}/{num_outer_steps}"
+                    )
+                    next_decode_ci += 1
+
+            for step in range(num_outer_steps):
+                win_sc, win_ec = valid_intervals[step]
+                # Enforce: committed chunks stay inside cache coverage.
+                # win_sc should equal committed chunks count.  If win_sc < committed
+                # (shouldn't happen), clamp.
+                win_sc = max(win_sc, n_full_ctx_chunks)
+
+                # --- Commit newly-finished chunks into the cache ---
+                while last_win_sc < win_sc:
+                    ci = last_win_sc
+                    gsl = slice(ci * chunk_size, min((ci + 1) * chunk_size, total_len))
+                    commit_frames = z_global[:, :, gsl]
+                    commit_cond = cond_seq[:, gsl]
+                    # t=0: chunk has finished denoising, treat as clean ctx going forward.
+                    t_commit = torch.zeros(
+                        b, commit_frames.shape[2], device=device, dtype=dtype,
+                    )
+
+                    _, kv_cond_new = net.forward_with_cache(
+                        commit_frames, t_commit, commit_cond,
+                        past_kv_list=cache_cond,
+                        current_position_offset=cache_frames,
+                        return_kv=True, chunk_size=chunk_size,
+                    )
+                    cache_cond = self._append_kv_cache(cache_cond, kv_cond_new)
+                    if use_cfg:
+                        commit_uncond, commit_drop_uncond = self._make_uncond(commit_cond)
+                        _, kv_uncond_new = net.forward_with_cache(
+                            commit_frames, t_commit, commit_uncond,
+                            past_kv_list=cache_uncond,
+                            current_position_offset=cache_frames,
+                            return_kv=True, chunk_size=chunk_size, cond_drop=commit_drop_uncond,
+                        )
+                        cache_uncond = self._append_kv_cache(cache_uncond, kv_uncond_new)
+                    cache_frames += commit_frames.shape[2]
+
+                    if cache_frames > max_cache_frames:
+                        # Never evict into the resident sink region.
+                        drop = min(cache_frames - max_cache_frames,
+                                   cache_frames - sink_frames)
+                        if drop > 0:
+                            cache_cond = self._evict_and_shift_cache(
+                                cache_cond, drop, tokens_per_frame, rope_module,
+                                sink_frames=sink_frames,
+                            )
+                            if use_cfg:
+                                cache_uncond = self._evict_and_shift_cache(
+                                    cache_uncond, drop, tokens_per_frame, rope_module,
+                                    sink_frames=sink_frames,
+                                )
+                            cache_frames -= drop
+
+                    committed_chunks.add(ci)
+                    _print0(f"[StreamingGen]   committed chunk {ci}/{total_chunks} | "
+                            f"cache_frames={cache_frames} | step={step}/{num_outer_steps}")
+                    last_win_sc += 1
+
+                if win_ec <= win_sc:
+                    _decode_finished_chunks(step)
+                    continue
+
+                # --- In-flight forward ---
+                win_sf = win_sc * chunk_size
+                win_ef = min(win_ec * chunk_size, total_len)
+                inflight_z = z_global[:, :, win_sf:win_ef].clone()
+                inflight_cond = cond_seq[:, win_sf:win_ef]
+
+                n_inflight = win_ec - win_sc
+                inflight_chunk_slices = self._build_chunk_slices(win_ef - win_sf)
+
+                t_chunks = t_chunk_sched[step, win_sc:win_ec]
+                t_next_chunks = t_next_chunk_sched[step, win_sc:win_ec]
+                t_frame = self._broadcast_chunk_values_to_frames(
+                    t_chunks.unsqueeze(0).expand(b, -1),
+                    inflight_chunk_slices, win_ef - win_sf,
+                )
+                t_next_frame = self._broadcast_chunk_values_to_frames(
+                    t_next_chunks.unsqueeze(0).expand(b, -1),
+                    inflight_chunk_slices, win_ef - win_sf,
+                )
+
+                # Pin sub-chunk context frames inside this window to t=0 so dt=0
+                # and they aren't perturbed by the velocity update.
+                ctx_in_inflight = min(max(0, ctx_len - win_sf), win_ef - win_sf)
+                if ctx_in_inflight > 0:
+                    t_frame[:, :ctx_in_inflight] = 0.0
+                    t_next_frame[:, :ctx_in_inflight] = 0.0
+                    inflight_z[:, :, :ctx_in_inflight] = latents[
+                        :, :, win_sf:win_sf + ctx_in_inflight
+                    ]
+
+                v_cond_pred, _ = net.forward_with_cache(
+                    inflight_z, t_frame, inflight_cond,
+                    past_kv_list=cache_cond,
+                    current_position_offset=cache_frames,
+                    return_kv=False, chunk_size=chunk_size,
+                )
+                if use_cfg:
+                    inflight_uncond, inflight_drop_uncond = self._make_uncond(inflight_cond)
+                    v_uncond_pred, _ = net.forward_with_cache(
+                        inflight_z, t_frame, inflight_uncond,
+                        past_kv_list=cache_uncond,
+                        current_position_offset=cache_frames,
+                        return_kv=False, chunk_size=chunk_size, cond_drop=inflight_drop_uncond,
+                    )
+                else:
+                    v_uncond_pred = None
+
+                update_mask_row = chunk_update_mask[step, win_sc:win_ec]
+                for lci in range(n_inflight):
+                    if not bool(update_mask_row[lci]):
+                        continue
+                    sl = inflight_chunk_slices[lci]
+                    gsl = slice(
+                        (win_sc + lci) * chunk_size,
+                        min((win_sc + lci + 1) * chunk_size, total_len),
+                    )
+                    gl = gsl.stop - gsl.start
+                    if use_cfg:
+                        chunk_t_val = t_frame[:, sl].mean(dim=1)
+                        action_scale = self._get_df_action_guidance_scale(chunk_t_val)
+                        action_scale = action_scale.view(-1, 1, 1, 1, 1)
+                        assert v_uncond_pred is not None
+                        v_chunk = (
+                            v_uncond_pred[:, :, sl]
+                            + action_scale * (v_cond_pred[:, :, sl] - v_uncond_pred[:, :, sl])
+                        )
+                    else:
+                        v_chunk = v_cond_pred[:, :, sl]
+                    dt = (t_next_frame[:, sl] - t_frame[:, sl]).view(b, 1, -1, 1, 1)[:, :, :gl]
+                    z_global[:, :, gsl] = (
+                        z_global[:, :, gsl] - dt * v_chunk[:, :, :gl]
+                    )
+
+                # Re-pin clean ctx frames (numerical safety; dt should already be
+                # 0 for them, but FP error can drift otherwise).
+                if ctx_len > 0:
+                    z_global[:, :, :ctx_len] = latents[:, :, :ctx_len]
+
+                # Decode as soon as each leading chunk's schedule hits t=0.
+                _decode_finished_chunks(step)
+
+            # Flush any remaining (e.g. final partial) frames for VAE.
+            if decoded_frames < total_len:
+                flush_ci = (total_len - 1) // chunk_size
+                if timing_enabled and timing_start is not None:
+                    t_dit = _timing_now()
+                    dit_chunk_events.append(
+                        {
+                            "chunk_idx": int(flush_ci),
+                            "step_idx": int(num_outer_steps),
+                            "end_frame": int(total_len),
+                            "generated": bool(flush_ci >= n_full_ctx_chunks),
+                            "complete_sec": float(t_dit - timing_start),
+                        }
+                    )
+                _stream_decode_upto(total_len, chunk_idx=flush_ci, step_idx=num_outer_steps)
+            if timing_enabled and timing_start is not None:
+                self.last_eval_meta["stream_timing"] = {
+                    "enabled": True,
+                    "cfg_enabled": bool(use_cfg),
+                    "start_sec": 0.0,
+                    "dit_chunk_events": dit_chunk_events,
+                    "vae_chunk_events": vae_chunk_events,
+                }
+            _print0(
+                f"[StreamingGen] done. committed={len(committed_chunks)}/{total_chunks} "
+                f"chunks, decoded_latent_frames={decoded_frames}/{total_len}, "
+                f"effective_steps={effective_steps}."
+            )
+
+            if stream_decoder is not None:
+                assert rgb_parts, (
+                    "[StreamingGen] stream_decoder was set but no RGB chunks were produced"
+                )
+                return z_global, torch.cat(rgb_parts, dim=2)
+            return z_global
+        finally:
+            if stream_decoder is not None:
+                stream_decoder.end()
+
+
+def build_denoiser_from_mode(cfg: DenoiserConfig) -> Denoiser:
+    """Build the public MiniWorld AR-diffusion denoiser."""
+    return DiffusionForcingDenoiser(cfg)

miniworld/miniworld.py

@@ -0,0 +1,1044 @@
+"""MiniWorld: action / pose-conditioned streaming video DiT.
+
+The model combines RoPE-only video tokens, action/pose conditioning streams,
+AdaLN-LoRA modulation, and structured condition dropout for streaming world
+modeling.
+
+  * **RoPE-only** positioning (absolute ``pos_embed`` removed entirely) so
+    train / streaming inference share the same position scheme.
+  * **AdaLN-LoRA modulation** (``adaln_mode="adaln_lora"``, default): a single
+    model-level modulation MLP produces a shared ``(B, T, 6D)`` term reused by
+    every block, plus a cheap per-block low-rank refinement ``D -> r -> 6D``.
+    This is the parameter-efficient middle ground between FLUX.2's
+    fully-shared modulation and the classic per-block full ``D -> 6D`` AdaLN.
+    Two more modes are provided for ablation: ``"fully_shared"`` (FLUX.2 style)
+    and ``"per_block"`` (classic DiT).
+  * **Separated conditioning streams** instead of the old
+    ``c_token = t_emb + y_emb``:
+      - timestep  -> its own encoder, drives the base modulation;
+      - action    -> DreamDojo-style encoder, *added* into the timestep /
+        AdaLN stream at per-latent-frame granularity (global-per-frame signal);
+      - pose      -> ray-encoding, injected as a *separate* per-token spatial
+        modulation stream (lingbot-style), kept out of the timestep stream.
+  * **Structured condition dropout** with a learned null embedding, for
+    classifier-free guidance training.
+
+The forward signature is designed for ``miniworld.denoiser.Denoiser``.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch import Tensor
+from torch.utils.checkpoint import checkpoint
+
+from flash_attn import flash_attn_func
+
+# FlexAttention is optional; disabled by default
+create_block_mask = None
+flex_attention = None
+
+
+# --------------------------------------------------------------------------- #
+#                           Attention masks (streaming)                        #
+# --------------------------------------------------------------------------- #
+def _build_temporal_chunkwise_attn_mask(
+    seq_len: int,
+    tokens_per_frame: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    chunk_size: int,
+) -> torch.Tensor:
+    """Block-causal (chunk-wise) additive mask over the temporal axis."""
+    token_idx = torch.arange(seq_len, device=device)
+    frame_idx = token_idx // tokens_per_frame
+    chunk_idx = frame_idx // chunk_size
+    mask = chunk_idx.unsqueeze(1) >= chunk_idx.unsqueeze(0)
+    float_mask = torch.zeros((1, 1, seq_len, seq_len), device=device, dtype=dtype)
+    float_mask.masked_fill_(~mask.unsqueeze(0).unsqueeze(0), float("-inf"))
+    return float_mask
+
+
+def _build_cached_block_causal_mask(
+    n_past: int,
+    n_cur: int,
+    tokens_per_frame: int,
+    chunk_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Additive mask for streaming forward with a KV cache.
+
+    Query length is ``n_cur`` (in-flight tokens); key/value length is
+    ``n_past + n_cur``.  Past cache is always visible; current tokens use
+    block-causal attention along the temporal chunk axis.
+    """
+    total_kv = n_past + n_cur
+    float_mask = torch.zeros((1, 1, n_cur, total_kv), device=device, dtype=dtype)
+    if n_cur == 0:
+        return float_mask
+    token_idx = torch.arange(n_cur, device=device)
+    chunk_idx = (token_idx // tokens_per_frame) // chunk_size
+    cur_mask = chunk_idx.unsqueeze(1) >= chunk_idx.unsqueeze(0)  # (n_cur, n_cur)
+    float_mask[0, 0, :, n_past:] = torch.where(
+        cur_mask,
+        torch.zeros((), device=device, dtype=dtype),
+        torch.full((), float("-inf"), device=device, dtype=dtype),
+    )
+    return float_mask
+
+
+# --------------------------------------------------------------------------- #
+#                                 Basic layers                                 #
+# --------------------------------------------------------------------------- #
+def modulate(x: Tensor, shift: Optional[Tensor], scale: Tensor) -> Tensor:
+    """AdaLN modulation. ``shift`` / ``scale`` may be ``(B, D)`` or ``(B, N, D)``."""
+    if scale.dim() == 2:
+        scale = scale.unsqueeze(1)
+        if shift is not None:
+            shift = shift.unsqueeze(1)
+    if shift is None:
+        return x * (1 + scale)
+    return x * (1 + scale) + shift
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x: Tensor) -> Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._norm(x.float()).type_as(x) * self.weight
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1, x2 = self.w12(x).chunk(2, dim=-1)
+        return self.w3(F.silu(x1) * x2)
+
+
+class PatchEmbed3D(nn.Module):
+    """(B, C, T, H, W) -> (B, N, D) via a 3D conv patchifier."""
+
+    def __init__(
+        self,
+        input_size: int | Tuple[int, int, int],
+        patch_size: int | Tuple[int, int, int],
+        in_chans: int,
+        embed_dim: int,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        if isinstance(input_size, int):
+            input_size = (input_size, input_size, input_size)
+        if isinstance(patch_size, int):
+            patch_size = (patch_size, patch_size, patch_size)
+        elif len(patch_size) == 2:
+            patch_size = (1, patch_size[0], patch_size[1])
+
+        self.input_size = input_size
+        self.patch_size = patch_size
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
+        self.num_patches = (
+            (input_size[0] // patch_size[0])
+            * (input_size[1] // patch_size[1])
+            * (input_size[2] // patch_size[2])
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.proj(x)  # (B, D, T', H', W')
+        x = x.flatten(2).transpose(1, 2)  # (B, N, D)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = RMSNorm(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim) if qk_norm else nn.Identity()
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(
+        self,
+        x: Tensor,
+        rope: Optional[Callable] = None,
+        attn_mask: Optional[Tensor] = None,
+        past_kv: Optional[Tuple[Tensor, Tensor]] = None,
+        return_kv: bool = False,
+    ):
+        B, N, C = x.shape
+        in_dtype = x.dtype
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        if rope is not None:
+            q = rope(q)
+            k = rope(k)
+
+        k_current, v_current = k, v
+        if past_kv is not None:
+            k_past, v_past = past_kv
+            k = torch.cat([k_past.to(dtype=k.dtype, device=k.device), k], dim=-2)
+            v = torch.cat([v_past.to(dtype=v.dtype, device=v.device), v], dim=-2)
+
+        if attn_mask is None and past_kv is None:
+            # flash-attn fast path expects (B, N, num_heads, head_dim)
+            q = q.transpose(1, 2).to(torch.bfloat16)
+            k = k.transpose(1, 2).to(torch.bfloat16)
+            v = v.transpose(1, 2).to(torch.bfloat16)
+            x = flash_attn_func(q, k, v, causal=False).transpose(1, 2)
+        else:
+            x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+
+        x = x.transpose(1, 2).reshape(B, N, C).to(in_dtype)
+        x = self.proj_drop(self.proj(x))
+        if return_kv:
+            return x, (k_current, v_current)
+        return x
+
+
+# --------------------------------------------------------------------------- #
+#                             3D rotary embedding                              #
+# --------------------------------------------------------------------------- #
+def broadcat(tensors, dim: int = -1):
+    num_tensors = len(tensors)
+    shape_lens = {len(t.shape) for t in tensors}
+    assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
+    shape_len = list(shape_lens)[0]
+    dim = (dim + shape_len) if dim < 0 else dim
+    dims = list(zip(*map(lambda t: list(t.shape), tensors)))
+    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+    assert all(len(set(t[1])) <= 2 for t in expandable_dims), "invalid broadcast dims"
+    max_dims = [(t[0], max(t[1])) for t in expandable_dims]
+    expanded_dims = [(t[0], (t[1],) * num_tensors) for t in max_dims]
+    expanded_dims.insert(dim, (dim, dims[dim]))
+    expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
+    tensors = [t[0].expand(*t[1]) for t in zip(tensors, expandable_shapes)]
+    return torch.cat(tensors, dim=dim)
+
+
+def rotate_half(x: Tensor) -> Tensor:
+    x = rearrange(x, "... (d r) -> ... d r", r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, "... d r -> ... (d r)")
+
+
+class VisionRotaryEmbeddingFast3D(nn.Module):
+    """Axial 3D RoPE (time / height / width), borrowed from EVA/lightning_wm."""
+
+    def __init__(self, dim: int, num_frames: int, frame_height: int, frame_width: int, theta: int = 10000) -> None:
+        super().__init__()
+        dim_h = (dim // 3) // 2 * 2
+        dim_w = (dim // 3) // 2 * 2
+        dim_t = dim - dim_h - dim_w
+        if dim_t % 2 != 0:
+            dim_t -= 1
+
+        freqs_t = 1.0 / (theta ** (torch.arange(0, dim_t, 2)[: (dim_t // 2)].float() / dim_t))
+        freqs_h = 1.0 / (theta ** (torch.arange(0, dim_h, 2)[: (dim_h // 2)].float() / dim_h))
+        freqs_w = 1.0 / (theta ** (torch.arange(0, dim_w, 2)[: (dim_w // 2)].float() / dim_w))
+
+        self.register_buffer("base_freqs_t", freqs_t)
+        self.register_buffer("base_freqs_h", freqs_h)
+        self.register_buffer("base_freqs_w", freqs_w)
+        self.frame_height = frame_height
+        self.frame_width = frame_width
+        self.dim = dim
+        self.dim_t = dim_t
+        self.dim_h = dim_h
+        self.dim_w = dim_w
+        self._num_frames = num_frames
+
+        freqs_cos, freqs_sin = self._build_freqs(
+            num_frames, freqs_t, freqs_h, freqs_w, frame_height, frame_width, dim
+        )
+        self.register_buffer("freqs_cos", freqs_cos)
+        self.register_buffer("freqs_sin", freqs_sin)
+
+    @staticmethod
+    def _build_freqs(num_frames, freqs_t, freqs_h, freqs_w, fh, fw, dim, start_frame: int = 0):
+        device = freqs_t.device
+        t_time = torch.arange(start_frame, start_frame + num_frames, device=device, dtype=torch.float32)
+        t_height = torch.arange(fh, device=device, dtype=torch.float32)
+        t_width = torch.arange(fw, device=device, dtype=torch.float32)
+        ft = repeat(torch.einsum("n,d->nd", t_time, freqs_t), "... n -> ... (n r)", r=2)
+        fht = repeat(torch.einsum("n,d->nd", t_height, freqs_h), "... n -> ... (n r)", r=2)
+        fwt = repeat(torch.einsum("n,d->nd", t_width, freqs_w), "... n -> ... (n r)", r=2)
+        freqs = broadcat(
+            (ft.view(num_frames, 1, 1, -1), fht.view(1, fh, 1, -1), fwt.view(1, 1, fw, -1)), dim=-1
+        )
+        return freqs.cos().view(-1, dim), freqs.sin().view(-1, dim)
+
+    def _get_freqs(self, num_frames: int, device: torch.device, start_frame: int = 0):
+        if start_frame == 0 and num_frames == self._num_frames:
+            return self.freqs_cos, self.freqs_sin
+        return self._build_freqs(
+            num_frames,
+            self.base_freqs_t.to(device),
+            self.base_freqs_h.to(device),
+            self.base_freqs_w.to(device),
+            self.frame_height,
+            self.frame_width,
+            self.dim,
+            start_frame=start_frame,
+        )
+
+    def forward(self, t: Tensor, num_frames_override: Optional[int] = None, start_frame: int = 0) -> Tensor:
+        num_frames = num_frames_override if num_frames_override is not None else self._num_frames
+        cos, sin = self._get_freqs(num_frames, t.device, start_frame=start_frame)
+        return t * cos + rotate_half(t) * sin
+
+    def rope_shift_time(self, delta: int, cached: Tensor) -> Tensor:
+        """Re-rotate cached K/Q on the temporal axis by ``delta`` frames.
+
+        Given a tensor originally RoPE-rotated at temporal positions
+        ``[p, ..., p+N-1]``, returns it rotated as if positions were
+        ``[p+delta, ..., p+delta+N-1]``. Use ``delta=-k`` after evicting ``k``
+        leading frames from a streaming cache to renumber positions back to 0.
+        Only the temporal slice (first ``dim_t`` dims) is rotated; spatial dims
+        receive identity rotation.
+        """
+        if delta == 0 or cached.numel() == 0:
+            return cached
+        device = cached.device
+        out_dtype = cached.dtype
+
+        base_freqs_t = self.base_freqs_t.to(device=device, dtype=torch.float32)
+        angle_t = float(delta) * base_freqs_t
+        angle_t_rep = repeat(angle_t, "n -> (n r)", r=2)
+        cos_t = angle_t_rep.cos()
+        sin_t = angle_t_rep.sin()
+
+        rest = self.dim - self.dim_t
+        cos_rest = torch.ones(rest, device=device, dtype=torch.float32)
+        sin_rest = torch.zeros(rest, device=device, dtype=torch.float32)
+
+        cos_full = torch.cat([cos_t, cos_rest], dim=-1).to(out_dtype)
+        sin_full = torch.cat([sin_t, sin_rest], dim=-1).to(out_dtype)
+
+        return cached * cos_full + rotate_half(cached) * sin_full
+
+
+# --------------------------------------------------------------------------- #
+#                              Timestep embedding                              #
+# --------------------------------------------------------------------------- #
+class TimestepEmbedder(nn.Module):
+    """Scalar timestep -> D-dim vector (sinusoidal + MLP)."""
+
+    def __init__(self, hidden_size: int, freq_dim: int = 256) -> None:
+        super().__init__()
+        self.freq_dim = freq_dim
+        self.mlp = nn.Sequential(
+            nn.Linear(freq_dim, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+
+    @staticmethod
+    def timestep_embedding(t: Tensor, dim: int, max_period: int = 10000) -> Tensor:
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(half, device=t.device, dtype=torch.float32) / half
+        )
+        args = t[:, None].float() * freqs[None]
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2 == 1:
+            emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
+        return emb
+
+    def forward(self, t: Tensor) -> Tensor:
+        if t.dim() == 1:
+            return self.mlp(self.timestep_embedding(t, self.freq_dim))
+        if t.dim() == 2:
+            b, s = t.shape
+            emb = self.mlp(self.timestep_embedding(t.reshape(-1), self.freq_dim))
+            return emb.view(b, s, -1)
+        raise ValueError(f"Unsupported timestep shape: {t.shape}")
+
+
+# --------------------------------------------------------------------------- #
+#                          Conditioning encoders                              #
+# --------------------------------------------------------------------------- #
+class ActionEncoder(nn.Module):
+    """DreamDojo-style action encoder for a global per-frame action signal.
+
+    Input action condition is ``(B, T, cond_dim)`` where each latent frame ``t``
+    already packs its chunk of raw actions (the training pipeline builds
+    ``cond_dim = num_action_per_latent * action_dim``).  Two MLP heads produce:
+
+      * ``emb_B_T_D``  -- added into the timestep embedding stream;
+      * ``mod_B_T_MD`` -- added into the (shared) AdaLN modulation stream,
+        where ``M = n_mod_chunks`` (6 here: attn shift/scale/gate + mlp
+        shift/scale/gate).
+
+    This mirrors Cosmos' ``action_embedder_B_D`` / ``action_embedder_B_3D``
+    but targets a 6-chunk modulation layout.
+    """
+
+    def __init__(self, cond_dim: int, hidden_size: int, n_mod_chunks: int = 6, hidden_mult: int = 4) -> None:
+        super().__init__()
+        hidden = hidden_size * hidden_mult
+        act = lambda: nn.GELU(approximate="tanh")
+        self.to_emb = nn.Sequential(
+            nn.Linear(cond_dim, hidden), act(), nn.Linear(hidden, hidden_size)
+        )
+        self.to_mod = nn.Sequential(
+            nn.Linear(cond_dim, hidden), act(), nn.Linear(hidden, n_mod_chunks * hidden_size)
+        )
+
+    def forward(self, action_B_T_C: Tensor) -> Tuple[Tensor, Tensor]:
+        return self.to_emb(action_B_T_C), self.to_mod(action_B_T_C)
+
+
+class PoseEncoder(nn.Module):
+    """Ray-encoding -> per-token spatial AdaLN modulation (lingbot-style).
+
+    The pose condition is a per-pixel ray-encoding volume
+    ``(B, T, cond_dim, H_lat, W_lat)`` (see ``pose_utils.compute_ray_encoding``,
+    e.g. cond_dim = 180 for origin+direction with 15 NeRF frequencies).
+
+    Unlike ``ActionEncoder`` (a per-frame signal folded into the timestep
+    stream), pose is spatially varying, so it drives its *own* per-token
+    modulation stream ``(B, N, MD)`` that is added on top of the timestep /
+    action modulation inside every block.  A residual MLP over the patchified
+    ray features mirrors lingbot's ``cam_injector`` before producing scale/shift.
+    """
+
+    def __init__(self, cond_dim: int, hidden_size: int, patch_size: int, n_mod_chunks: int = 6) -> None:
+        super().__init__()
+        self.patchify = nn.Conv2d(cond_dim, hidden_size, kernel_size=patch_size, stride=patch_size, bias=True)
+        self.res_mlp = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size)
+        )
+        self.to_mod = nn.Linear(hidden_size, n_mod_chunks * hidden_size, bias=True)
+
+    def forward(self, pose_B_T_C_H_W: Tensor, b: int, grid_t: int) -> Tensor:
+        """Return per-token modulation ``(B, N, MD)`` with N = grid_t * h * w."""
+        y = rearrange(pose_B_T_C_H_W, "b t c h w -> (b t) c h w")
+        y = self.patchify(y)  # ((B*T), D, h', w')
+        y = rearrange(y, "(b t) d h w -> b (t h w) d", b=b, t=grid_t)
+        y = y + self.res_mlp(y)  # residual (lingbot cam_injector style)
+        return self.to_mod(y)  # (B, N, MD)
+
+
+# --------------------------------------------------------------------------- #
+#                          Modulation (shared / lora)                         #
+# --------------------------------------------------------------------------- #
+_MODES = ("adaln_lora", "fully_shared", "per_block")
+
+
+class BlockModulation(nn.Module):
+    """Per-block modulation producer, respecting the model-wide ``adaln_mode``.
+
+    * ``adaln_lora``   : ``shared_mod + lora(emb)`` where ``lora = SiLU -> D->r
+      -> r->MD`` (zero-init, so a block starts exactly at ``shared_mod``).
+    * ``fully_shared`` : ``shared_mod`` (no per-block params; FLUX.2 style).
+    * ``per_block``    : ``full(emb)`` with ``full = SiLU -> D->MD`` (classic
+      per-block AdaLN, zero-init).
+    """
+
+    def __init__(self, hidden_size: int, adaln_mode: str, n_mod_chunks: int = 6, lora_dim: int = 256) -> None:
+        super().__init__()
+        assert adaln_mode in _MODES, f"adaln_mode must be one of {_MODES}"
+        self.adaln_mode = adaln_mode
+        out = n_mod_chunks * hidden_size
+        if adaln_mode == "adaln_lora":
+            self.lora = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(hidden_size, lora_dim, bias=False),
+                nn.Linear(lora_dim, out, bias=False),
+            )
+        elif adaln_mode == "per_block":
+            self.full = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, out, bias=True))
+        # fully_shared: no parameters
+
+    def forward(self, emb: Tensor, shared_mod: Optional[Tensor], pose_mod: Optional[Tensor]) -> Tensor:
+        if self.adaln_mode == "adaln_lora":
+            mod = shared_mod + self.lora(emb)
+        elif self.adaln_mode == "fully_shared":
+            mod = shared_mod
+        else:  # per_block
+            mod = self.full(emb)
+        if pose_mod is not None:
+            mod = mod + pose_mod
+        return mod
+
+
+# --------------------------------------------------------------------------- #
+#                                   Blocks                                     #
+# --------------------------------------------------------------------------- #
+class MiniWorldBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        adaln_mode: str,
+        mlp_ratio: float = 4.0,
+        use_qknorm: bool = False,
+        lora_dim: int = 256,
+    ) -> None:
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size)
+        self.norm2 = RMSNorm(hidden_size)
+        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=use_qknorm)
+        mlp_hidden = int(hidden_size * mlp_ratio)
+        self.mlp = SwiGLUFFN(hidden_size, int(2 / 3 * mlp_hidden))
+        self.modulation = BlockModulation(hidden_size, adaln_mode, n_mod_chunks=6, lora_dim=lora_dim)
+
+    def forward(
+        self,
+        x: Tensor,
+        emb: Tensor,
+        shared_mod: Optional[Tensor],
+        pose_mod: Optional[Tensor],
+        feat_rope: Optional[Callable] = None,
+        attn_mask: Optional[Tensor] = None,
+        past_kv: Optional[Tuple[Tensor, Tensor]] = None,
+        return_kv: bool = False,
+    ):
+        mod = self.modulation(emb, shared_mod, pose_mod)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=-1)
+
+        attn_result = self.attn(
+            modulate(self.norm1(x), shift_msa, scale_msa),
+            rope=feat_rope,
+            attn_mask=attn_mask,
+            past_kv=past_kv,
+            return_kv=return_kv,
+        )
+        if return_kv:
+            attn_out, new_kv = attn_result
+        else:
+            attn_out, new_kv = attn_result, None
+        x = x + gate_msa * attn_out
+        x = x + gate_mlp * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        if return_kv:
+            return x, new_kv
+        return x
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int) -> None:
+        super().__init__()
+        self.norm = RMSNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x: Tensor, emb: Tensor) -> Tensor:
+        cond = self.adaLN_modulation(emb)
+        if cond.dim() == 2:
+            shift, scale = cond.chunk(2, dim=1)
+        else:
+            shift, scale = cond.chunk(2, dim=-1)
+        return self.linear(modulate(self.norm(x), shift, scale))
+
+
+# --------------------------------------------------------------------------- #
+#                                Main model                                    #
+# --------------------------------------------------------------------------- #
+class MiniWorldModel(nn.Module):
+    """Action / pose-conditioned streaming video DiT (RoPE-only)."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_size: int,
+        cond_dim: int,
+        depth: int,
+        num_heads: int,
+        patch_size: int,
+        input_size: int | Tuple[int, int],
+        num_frames: int = 9,
+        mlp_ratio: float = 4.0,
+        use_qknorm: bool = True,
+        use_checkpoint: bool = False,
+        cond_per_token: bool = False,
+        adaln_mode: str = "adaln_lora",
+        adaln_lora_dim: int = 256,
+        cond_dropout_prob: float = 0.0,
+        action_null_first: bool = True,
+        # Kept for checkpoint compatibility; MiniWorld always uses RoPE.
+        use_rope: bool = True,
+        use_abs_pos: bool = False,
+    ) -> None:
+        super().__init__()
+        assert adaln_mode in _MODES, f"adaln_mode must be one of {_MODES}"
+        assert use_rope, "MiniWorldModel is RoPE-only; use_rope must be True."
+        assert not use_abs_pos, "MiniWorldModel is RoPE-only; use_abs_pos must be False."
+
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+        self.cond_per_token = cond_per_token
+        self.adaln_mode = adaln_mode
+        self.cond_dropout_prob = cond_dropout_prob
+        # Route the true first latent frame (the seed / initial observation,
+        # which has no preceding action) through the learned ``null_action``.
+        self.action_null_first = action_null_first
+        # RoPE-only: kept as attributes for downstream code / streaming asserts.
+        self.use_rope = True
+        self.use_abs_pos = False
+
+        input_size = (
+            (num_frames, input_size, input_size)
+            if isinstance(input_size, int)
+            else (num_frames,) + tuple(input_size)
+        )
+        self.x_embedder = PatchEmbed3D(
+            input_size=input_size,
+            patch_size=(1, patch_size, patch_size) if isinstance(patch_size, int) else patch_size,
+            in_chans=in_channels,
+            embed_dim=hidden_size,
+            bias=True,
+        )
+
+        # ---- conditioning streams -------------------------------------- #
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        # RMSNorm on the (timestep [+ action]) embedding before it drives the
+        # AdaLN heads / final layer. Mirrors Cosmos/DreamDojo ``t_embedding_norm``:
+        # keeps the affine embedding well-conditioned once the action encoder
+        # sums a second, independently-scaled signal into the timestep stream.
+        self.emb_norm = RMSNorm(hidden_size)
+        # shared modulation head (base AdaLN term reused across all blocks)
+        self.shared_mod = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
+
+        if cond_per_token:
+            # pose (ray-encoding) -> separate per-token spatial modulation.
+            self.pose_encoder = PoseEncoder(cond_dim, hidden_size, patch_size, n_mod_chunks=6)
+            self.action_encoder = None
+        else:
+            # action -> DreamDojo-style, folded into timestep / AdaLN stream.
+            self.action_encoder = ActionEncoder(cond_dim, hidden_size, n_mod_chunks=6)
+            self.pose_encoder = None
+            # learned null action for classifier-free guidance dropout.
+            self.null_action = nn.Parameter(torch.zeros(1, 1, cond_dim))
+
+        head_dim = hidden_size // num_heads
+        self.feat_rope = VisionRotaryEmbeddingFast3D(
+            dim=head_dim,
+            num_frames=num_frames,
+            frame_height=input_size[1] // patch_size,
+            frame_width=input_size[2] // patch_size,
+        )
+
+        self.blocks = nn.ModuleList(
+            [
+                MiniWorldBlock(
+                    hidden_size=hidden_size,
+                    num_heads=num_heads,
+                    adaln_mode=adaln_mode,
+                    mlp_ratio=mlp_ratio,
+                    use_qknorm=use_qknorm,
+                    lora_dim=adaln_lora_dim,
+                )
+                for _ in range(depth)
+            ]
+        )
+        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
+        self.initialize_weights()
+
+    # ------------------------------------------------------------------ #
+    def initialize_weights(self) -> None:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # AdaLN-zero: shared modulation starts at 0 -> identity blocks.
+        nn.init.constant_(self.shared_mod[-1].weight, 0)
+        nn.init.constant_(self.shared_mod[-1].bias, 0)
+
+        # zero-init per-block modulation refinement so blocks start at shared_mod.
+        for block in self.blocks:
+            if self.adaln_mode == "adaln_lora":
+                nn.init.constant_(block.modulation.lora[-1].weight, 0)
+            elif self.adaln_mode == "per_block":
+                nn.init.constant_(block.modulation.full[-1].weight, 0)
+                nn.init.constant_(block.modulation.full[-1].bias, 0)
+
+        # action stream: zero-init the modulation head so action ramps in.
+        if self.action_encoder is not None:
+            nn.init.constant_(self.action_encoder.to_mod[-1].weight, 0)
+            nn.init.constant_(self.action_encoder.to_mod[-1].bias, 0)
+
+        # pose stream: zero-init so pose modulation ramps in from identity.
+        if self.pose_encoder is not None:
+            nn.init.constant_(self.pose_encoder.to_mod.weight, 0)
+            nn.init.constant_(self.pose_encoder.to_mod.bias, 0)
+
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    # ------------------------------------------------------------------ #
+    def unpatchify(self, x: Tensor) -> Tensor:
+        b, n, _ = x.shape
+        c = self.out_channels
+        p_t, p_h, p_w = self.x_embedder.patch_size
+        t_in, h_in, w_in = self.x_embedder.input_size
+        grid_t, grid_h, grid_w = t_in // p_t, h_in // p_h, w_in // p_w
+        assert n == grid_t * grid_h * grid_w, f"seq len {n} != grid {grid_t}x{grid_h}x{grid_w}"
+        x = x.reshape(b, grid_t, grid_h, grid_w, p_t, p_h, p_w, c)
+        x = torch.einsum("nthwpqrc->nctphqwr", x)
+        return x.reshape(b, c, grid_t * p_t, grid_h * p_h, grid_w * p_w)
+
+    # ------------------------------------------------------------------ #
+    def _resolve_drop_mask(self, b: int, device, cond_drop: Optional[Tensor]) -> Optional[Tensor]:
+        """Resolve a per-sample CFG drop mask ``(B,)`` bool, or None.
+
+        Explicit ``cond_drop`` wins; otherwise sample from ``cond_dropout_prob``
+        while training.  Shared by the action and pose streams so a dropped
+        sample is fully unconditional.
+        """
+        if cond_drop is not None:
+            return cond_drop.to(device=device, dtype=torch.bool).view(b)
+        if self.training and self.cond_dropout_prob > 0:
+            return torch.rand(b, device=device) < self.cond_dropout_prob
+        return None
+
+    def _build_conditioning(
+        self,
+        t: Tensor,
+        y: Optional[Tensor],
+        b: int,
+        grid_t: int,
+        n: int,
+        frame_ids: Tensor,
+        device,
+        dtype,
+        cond_drop: Optional[Tensor],
+        frame_offset: int = 0,
+    ) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
+        """Return ``(emb_tok, shared_mod_tok, pose_mod_tok)`` all in per-token layout.
+
+        * ``emb_tok``       : ``(B, N, D)`` timestep(+action) embedding.
+        * ``shared_mod_tok``: ``(B, N, 6D)`` shared AdaLN modulation, or None
+          (``per_block`` mode does not use it).
+        * ``pose_mod_tok``  : ``(B, N, 6D)`` pose modulation, or None.
+
+        ``frame_offset`` is the absolute temporal index of this window's first
+        latent frame (0 for whole-clip / training / the first streaming window;
+        >0 for later streaming windows). It gates the ``action_null_first``
+        behaviour so only the true global frame 0 is treated as action-free.
+        """
+        # ---- timestep -> per-token embedding --------------------------- #
+        per_token_t = False
+        if t.dim() == 1:
+            t = t.view(b, 1).expand(b, grid_t)
+        elif t.dim() == 2:
+            if t.size(1) == n:
+                per_token_t = True
+            elif t.size(1) != grid_t:
+                raise ValueError(f"t shape {t.shape} != frames {grid_t} or tokens {n}")
+        else:
+            raise ValueError(f"Unsupported timestep shape: {t.shape}")
+        t_emb = self.t_embedder(t)  # (B, grid_t, D) or (B, N, D)
+
+        drop_mask = self._resolve_drop_mask(b, device, cond_drop)  # (B,) bool or None
+
+        # ---- action -> fold into timestep / modulation stream ---------- #
+        if self.action_encoder is not None:
+            null = self.null_action.to(device=device, dtype=dtype)
+            if y is None:
+                y = null.expand(b, grid_t, -1)
+            else:
+                if y.size(1) == 1:
+                    y = y.expand(b, grid_t, -1)
+                assert y.size(1) == grid_t, f"action T {y.size(1)} != latent frames {grid_t}"
+                if drop_mask is not None:
+                    m = drop_mask.view(b, 1, 1).to(y.dtype)
+                    y = y * (1 - m) + null * m
+            # The true first latent frame is the seed / initial observation and
+            # has no preceding action -> use the learned null there. Only when
+            # this window actually starts at global frame 0 (frame_offset == 0),
+            # so later streaming windows keep their real per-frame actions.
+            if self.action_null_first and frame_offset == 0 and grid_t > 0:
+                y = y.clone()
+                y[:, 0:1, :] = null
+            a_emb, a_mod = self.action_encoder(y)  # (B,grid_t,D), (B,grid_t,6D)
+            t_emb = t_emb + a_emb
+            action_mod = a_mod
+        else:
+            action_mod = None
+
+        # ---- normalize the combined timestep(+action) embedding -------- #
+        # Applied unconditionally (part of the timestep pipeline; also helps the
+        # pose / no-action paths). The 6D modulation deltas stay un-normed.
+        t_emb = self.emb_norm(t_emb)
+
+        # ---- broadcast per-frame -> per-token -------------------------- #
+        emb_tok = t_emb if per_token_t else t_emb[:, frame_ids, :]  # (B, N, D)
+
+        shared_mod_tok: Optional[Tensor] = None
+        if self.adaln_mode in ("adaln_lora", "fully_shared"):
+            shared = self.shared_mod(emb_tok)  # (B, N, 6D)
+            if action_mod is not None:
+                shared = shared + action_mod[:, frame_ids, :]
+            shared_mod_tok = shared
+        elif action_mod is not None:
+            # per_block mode has no shared term; route action modulation
+            # through the pose channel (both are additive per-token deltas).
+            action_mod = action_mod[:, frame_ids, :]
+
+        # ---- pose -> per-token spatial modulation ---------------------- #
+        pose_mod_tok: Optional[Tensor] = None
+        if self.pose_encoder is not None and y is not None:
+            pose_mod_tok = self.pose_encoder(y, b, grid_t)  # (B, N, 6D)
+            if drop_mask is not None:
+                # dropped samples become unconditional -> zero pose modulation.
+                pose_mod_tok = pose_mod_tok * (~drop_mask).view(b, 1, 1).to(pose_mod_tok.dtype)
+
+        # in per_block mode, fold action delta into the pose channel
+        if self.adaln_mode == "per_block" and action_mod is not None:
+            pose_mod_tok = action_mod if pose_mod_tok is None else pose_mod_tok + action_mod
+
+        return emb_tok, shared_mod_tok, pose_mod_tok
+
+    # ------------------------------------------------------------------ #
+    def forward(
+        self,
+        x: Tensor,
+        t: Optional[Tensor] = None,
+        y: Optional[Tensor] = None,
+        use_fp16: bool = False,
+        temporal_causal: bool = False,
+        chunk_size: Optional[int] = None,
+        cond_drop: Optional[Tensor] = None,
+        frame_offset: int = 0,
+    ):
+        """Forward pass.
+
+        Args:
+            x: ``(B, C, T, H, W)`` latent video.
+            t: ``(B,)`` / ``(B, T')`` per-frame or ``(B, N)`` per-token timesteps.
+            y: condition. If ``cond_per_token`` -> ``(B, T', cond_dim, H, W)``
+               ray-encoding; else ``(B, T', cond_dim)`` action.
+            cond_drop: optional per-sample bool ``(B,)`` forcing the null / uncond
+               condition (for CFG). When None and training, sampled from
+               ``cond_dropout_prob``.
+
+        Returns:
+            ``v_pred`` of shape ``(B, C, T, H, W)``.
+        """
+        x = self.x_embedder(x)
+
+        p_t, p_h, p_w = self.x_embedder.patch_size
+        t_in, h_in, w_in = self.x_embedder.input_size
+        grid_t, grid_h, grid_w = t_in // p_t, h_in // p_h, w_in // p_w
+        tokens_per_frame = grid_h * grid_w
+        b, n, _ = x.shape
+        assert n == grid_t * tokens_per_frame, f"token len {n} != grid {grid_t}x{grid_h}x{grid_w}"
+        chunk_size = 1 if chunk_size is None else chunk_size
+
+        attn_mask = None
+        if temporal_causal:
+            attn_mask = _build_temporal_chunkwise_attn_mask(
+                seq_len=n,
+                tokens_per_frame=tokens_per_frame,
+                device=x.device,
+                dtype=x.dtype,
+                chunk_size=chunk_size,
+            )
+
+        frame_ids = torch.arange(n, device=x.device, dtype=torch.long) // tokens_per_frame
+        emb_tok, shared_mod_tok, pose_mod_tok = self._build_conditioning(
+            t, y, b, grid_t, n, frame_ids, x.device, x.dtype, cond_drop,
+            frame_offset=frame_offset,
+        )
+
+        for block in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint(
+                    block, x, emb_tok, shared_mod_tok, pose_mod_tok, self.feat_rope, attn_mask,
+                    use_reentrant=True,
+                )
+            else:
+                x = block(x, emb_tok, shared_mod_tok, pose_mod_tok, self.feat_rope, attn_mask)
+
+        # final layer uses the per-frame(-broadcast) timestep embedding
+        x = self.final_layer(x, emb_tok)
+        x = self.unpatchify(x)
+        return x
+
+    # ------------------------------------------------------------------ #
+    @torch.no_grad()
+    def forward_with_cache(
+        self,
+        x: Tensor,
+        t: Tensor,
+        y: Optional[Tensor] = None,
+        past_kv_list: Optional[List[Optional[Tuple[Tensor, Tensor]]]] = None,
+        current_position_offset: int = 0,
+        return_kv: bool = False,
+        chunk_size: int = 1,
+        cond_drop: Optional[Tensor] = None,
+    ):
+        """Streaming forward with optional KV cache injection (RoPE-only)."""
+        b, c_in, t_cur, h_in, w_in = x.shape
+        x = self.x_embedder(x)
+
+        p_t, p_h, p_w = self.x_embedder.patch_size
+        _, h_total, w_total = self.x_embedder.input_size
+        grid_h, grid_w = h_total // p_h, w_total // p_w
+        assert h_in == h_total and w_in == w_total, (
+            f"forward_with_cache expects {h_total}x{w_total}, got {h_in}x{w_in}"
+        )
+        grid_t_cur = t_cur // p_t
+        tokens_per_frame = grid_h * grid_w
+        n_cur = grid_t_cur * tokens_per_frame
+        assert x.shape[1] == n_cur, f"patch embed produced {x.shape[1]} tokens, expected {n_cur}"
+
+        if past_kv_list is None:
+            past_kv_list = [None] * self.depth
+        assert len(past_kv_list) == self.depth
+
+        n_past = 0
+        first_past = next((kv for kv in past_kv_list if kv is not None), None)
+        if first_past is not None:
+            n_past = int(first_past[0].shape[-2])
+            assert n_past == current_position_offset * tokens_per_frame, (
+                f"past token len {n_past} != offset*tokens_per_frame "
+                f"{current_position_offset * tokens_per_frame}"
+            )
+
+        attn_mask = _build_cached_block_causal_mask(
+            n_past=n_past,
+            n_cur=n_cur,
+            tokens_per_frame=tokens_per_frame,
+            chunk_size=max(1, int(chunk_size)),
+            device=x.device,
+            dtype=x.dtype,
+        )
+
+        def feat_rope_current(tt: Tensor) -> Tensor:
+            return self.feat_rope(tt, num_frames_override=grid_t_cur, start_frame=int(current_position_offset))
+
+        frame_ids = torch.arange(n_cur, device=x.device, dtype=torch.long) // tokens_per_frame
+        emb_tok, shared_mod_tok, pose_mod_tok = self._build_conditioning(
+            t, y, b, grid_t_cur, n_cur, frame_ids, x.device, x.dtype, cond_drop,
+            frame_offset=int(current_position_offset),
+        )
+
+        new_kv_list: List[Optional[Tuple[Tensor, Tensor]]] = [None] * self.depth
+        for idx, block in enumerate(self.blocks):
+            out = block(
+                x, emb_tok, shared_mod_tok, pose_mod_tok, feat_rope_current, attn_mask,
+                past_kv_list[idx], return_kv,
+            )
+            if return_kv:
+                x, new_kv_list[idx] = out
+            else:
+                x = out
+
+        x = self.final_layer(x, emb_tok)
+        c_out = self.out_channels
+        x_vid = x.reshape(b, grid_t_cur, grid_h, grid_w, p_t, p_h, p_w, c_out)
+        x_vid = torch.einsum("nthwpqrc->nctphqwr", x_vid)
+        v_pred = x_vid.reshape(b, c_out, grid_t_cur * p_t, grid_h * p_h, grid_w * p_w)
+
+        if return_kv:
+            return v_pred, new_kv_list
+        return v_pred, None
+
+
+# --------------------------------------------------------------------------- #
+#                               Factory configs                               #
+# --------------------------------------------------------------------------- #
+# Scaling-law ladder. All use patch_size=1; approx param counts are for the
+# action mode (cond_dim~128); pose mode differs by only a few tens of M.
+#
+#   public name  hidden  depth  heads  head_dim   ~params
+#   B              768     12     12       64       ~0.12B
+#   L             1024     24     16       64       ~0.39B
+#   0.5B          1152     28     16       72       ~0.55B
+#   1B            1536     28     12      128       ~0.96B
+#   3B            2560     32     20      128       ~2.9B
+def MiniWorld_B(**kwargs):
+    """~0.12B. hidden=768, depth=12, heads=12 (head_dim=64)."""
+    return MiniWorldModel(depth=12, hidden_size=768, num_heads=12, patch_size=1, **kwargs)
+
+
+def MiniWorld_L(**kwargs):
+    """~0.39B. hidden=1024, depth=24, heads=16 (head_dim=64)."""
+    return MiniWorldModel(depth=24, hidden_size=1024, num_heads=16, patch_size=1, **kwargs)
+
+
+def MiniWorld_0_5B(**kwargs):
+    """~0.55B. hidden=1152, depth=28, heads=16 (head_dim=72)."""
+    return MiniWorldModel(depth=28, hidden_size=1152, num_heads=16, patch_size=1, **kwargs)
+
+
+def MiniWorld_1B(**kwargs):
+    """~0.96B. hidden=1536, depth=28, heads=12 (head_dim=128)."""
+    return MiniWorldModel(depth=28, hidden_size=1536, num_heads=12, patch_size=1, **kwargs)
+
+
+def MiniWorld_3B(**kwargs):
+    """~2.9B. hidden=2560, depth=32, heads=20 (head_dim=128)."""
+    return MiniWorldModel(depth=32, hidden_size=2560, num_heads=20, patch_size=1, **kwargs)
+
+
+MiniWorldModels = {
+    "B": MiniWorld_B,
+    "L": MiniWorld_L,
+    "0.5B": MiniWorld_0_5B,
+    "1B": MiniWorld_1B,
+    "3B": MiniWorld_3B,
+}

miniworld/vae/__init__.py

@@ -0,0 +1,6 @@
+"""VAE helpers for MiniWorld."""
+
+from miniworld.vae.codec import StreamingVAEDecoder, load_wan22_vae, vae_decode, vae_encode
+
+__all__ = ["StreamingVAEDecoder", "load_wan22_vae", "vae_decode", "vae_encode"]
+

miniworld/vae/codec.py

@@ -0,0 +1,98 @@
+"""WAN2.2 VAE loading, encode/decode, and streaming decode helpers."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Mapping
+
+import torch
+from torch.distributed import get_rank
+
+from miniworld.vae.wan22_vae import Wan2_2_VAE
+
+
+def is_main_process() -> bool:
+    """Return whether the current process should print user-facing logs."""
+    return int(os.environ.get("RANK", "0")) == 0
+
+
+def print0(message: str) -> None:
+    """Print only from rank 0."""
+    if is_main_process():
+        print(message, flush=True)
+
+
+def get_rank_id() -> int:
+    """Return distributed rank, or 0 when torch.distributed is inactive."""
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        return int(get_rank())
+    return 0
+
+
+def _get_arg(args: Mapping[str, Any] | object, name: str, default: Any = None) -> Any:
+    if isinstance(args, Mapping):
+        return args.get(name, default)
+    return getattr(args, name, default)
+
+
+def load_wan22_vae(args: Mapping[str, Any] | object) -> Wan2_2_VAE:
+    """Load a frozen WAN2.2 VAE on the current CUDA rank."""
+    checkpoint = _get_arg(args, "vae_checkpoint")
+    if checkpoint is None:
+        raise ValueError("vae_checkpoint is required")
+    checkpoint_path = Path(checkpoint)
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"WAN2.2 VAE checkpoint not found: {checkpoint_path}")
+
+    device = torch.device(f"cuda:{get_rank_id()}" if torch.cuda.is_available() else "cpu")
+    vae = Wan2_2_VAE(vae_pth=os.fspath(checkpoint_path), device=device)
+    vae.model.requires_grad_(False)
+    vae.model.eval()
+    print0(f"WAN2.2 VAE parameters: {sum(p.numel() for p in vae.model.parameters()):,}")
+    return vae
+
+
+def vae_encode(vae: Wan2_2_VAE, video: torch.Tensor) -> torch.Tensor:
+    """Encode RGB video in ``[-1, 1]`` to WAN2.2 latents."""
+    total_frames = video.shape[2]
+    target_frames = ((total_frames - 1) // 4) * 4 + 1
+    return vae.encode(video[:, :, :target_frames])
+
+
+@torch.no_grad()
+def vae_decode(vae: Wan2_2_VAE, latents: torch.Tensor) -> torch.Tensor:
+    """Decode WAN2.2 latents to RGB video in ``[-1, 1]``."""
+    return vae.decode(latents)
+
+
+class StreamingVAEDecoder:
+    """Causal streaming WAN2.2 decode session."""
+
+    def __init__(self, vae: Wan2_2_VAE) -> None:
+        self.vae = vae
+        self._active = False
+
+    def begin(self) -> None:
+        self.vae.decode_stream_begin()
+        self._active = True
+
+    def step(self, latents_chunk: torch.Tensor) -> torch.Tensor:
+        """Decode one latent chunk to RGB frames."""
+        if not self._active:
+            raise RuntimeError("StreamingVAEDecoder.begin() must be called before step()")
+        return self.vae.decode_stream_step(latents_chunk)
+
+    def end(self) -> None:
+        if self._active:
+            self.vae.decode_stream_end()
+            self._active = False
+
+    def decode_all(self, latents: torch.Tensor) -> torch.Tensor:
+        """Stream-decode a complete latent tensor."""
+        self.begin()
+        try:
+            return self.step(latents)
+        finally:
+            self.end()
+

miniworld/vae/wan22_vae.py

@@ -0,0 +1,1093 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+__all__ = [
+    "Wan2_2_VAE",
+]
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (
+            self.padding[2],
+            self.padding[2],
+            self.padding[1],
+            self.padding[1],
+            2 * self.padding[0],
+            0,
+        )
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x):
+        return (F.normalize(x, dim=(1 if self.channel_first else -1)) *
+                self.scale * self.gamma + self.bias)
+
+
+class Upsample(nn.Upsample):
+
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+    def __init__(self, dim, mode):
+        assert mode in (
+            "none",
+            "upsample2d",
+            "upsample3d",
+            "downsample2d",
+            "downsample3d",
+        )
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+                # nn.Conv2d(dim, dim//2, 3, padding=1)
+            )
+            self.time_conv = CausalConv3d(
+                dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(
+                dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == "upsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = "Rep"
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and
+                            feat_cache[idx] != "Rep"):
+                        # cache last frame of last two chunk
+                        cache_x = torch.cat(
+                            [
+                                feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                    cache_x.device),
+                                cache_x,
+                            ],
+                            dim=2,
+                        )
+                    if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and
+                            feat_cache[idx] == "Rep"):
+                        cache_x = torch.cat(
+                            [
+                                torch.zeros_like(cache_x).to(cache_x.device),
+                                cache_x
+                            ],
+                            dim=2,
+                        )
+                    if feat_cache[idx] == "Rep":
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
+                                    3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.resample(x)
+        x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+
+        if self.mode == "downsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    x = self.time_conv(
+                        torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        one_matrix = torch.eye(c1, c2)
+        init_matrix = one_matrix
+        nn.init.zeros_(conv_weight)
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix  # * 0.5
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # layers
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1),
+        )
+        self.shortcut = (
+            CausalConv3d(in_dim, out_dim, 1)
+            if in_dim != out_dim else nn.Identity())
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+        # zero out the last layer params
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.norm(x)
+        # compute query, key, value
+        q, k, v = (
+            self.to_qkv(x).reshape(b * t, 1, c * 3,
+                                   -1).permute(0, 1, 3,
+                                               2).contiguous().chunk(3, dim=-1))
+
+        # apply attention
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+        )
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+        # output
+        x = self.proj(x)
+        x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
+        return x + identity
+
+
+def patchify(x, patch_size):
+    if patch_size == 1:
+        return x
+    if x.dim() == 4:
+        x = rearrange(
+            x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(
+            x,
+            "b c f (h q) (w r) -> b (c r q) f h w",
+            q=patch_size,
+            r=patch_size,
+        )
+    else:
+        raise ValueError(f"Invalid input shape: {x.shape}")
+
+    return x
+
+
+def unpatchify(x, patch_size):
+    if patch_size == 1:
+        return x
+
+    if x.dim() == 4:
+        x = rearrange(
+            x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(
+            x,
+            "b (c r q) f h w -> b c f (h q) (w r)",
+            q=patch_size,
+            r=patch_size,
+        )
+    return x
+
+
+class AvgDown3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        factor_t,
+        factor_s=1,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+
+        assert in_channels * self.factor % out_channels == 0
+        self.group_size = in_channels * self.factor // out_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
+        pad = (0, 0, 0, 0, pad_t, 0)
+        x = F.pad(x, pad)
+        B, C, T, H, W = x.shape
+        x = x.view(
+            B,
+            C,
+            T // self.factor_t,
+            self.factor_t,
+            H // self.factor_s,
+            self.factor_s,
+            W // self.factor_s,
+            self.factor_s,
+        )
+        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
+        x = x.view(
+            B,
+            C * self.factor,
+            T // self.factor_t,
+            H // self.factor_s,
+            W // self.factor_s,
+        )
+        x = x.view(
+            B,
+            self.out_channels,
+            self.group_size,
+            T // self.factor_t,
+            H // self.factor_s,
+            W // self.factor_s,
+        )
+        x = x.mean(dim=2)
+        return x
+
+
+class DupUp3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        factor_t,
+        factor_s=1,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+
+        assert out_channels * self.factor % in_channels == 0
+        self.repeats = out_channels * self.factor // in_channels
+
+    def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
+        x = x.repeat_interleave(self.repeats, dim=1)
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            self.factor_t,
+            self.factor_s,
+            self.factor_s,
+            x.size(2),
+            x.size(3),
+            x.size(4),
+        )
+        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            x.size(2) * self.factor_t,
+            x.size(4) * self.factor_s,
+            x.size(6) * self.factor_s,
+        )
+        if first_chunk:
+            x = x[:, :, self.factor_t - 1:, :, :]
+        return x
+
+
+class Down_ResidualBlock(nn.Module):
+
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 dropout,
+                 mult,
+                 temperal_downsample=False,
+                 down_flag=False):
+        super().__init__()
+
+        # Shortcut path with downsample
+        self.avg_shortcut = AvgDown3D(
+            in_dim,
+            out_dim,
+            factor_t=2 if temperal_downsample else 1,
+            factor_s=2 if down_flag else 1,
+        )
+
+        # Main path with residual blocks and downsample
+        downsamples = []
+        for _ in range(mult):
+            downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+
+        # Add the final downsample block
+        if down_flag:
+            mode = "downsample3d" if temperal_downsample else "downsample2d"
+            downsamples.append(Resample(out_dim, mode=mode))
+
+        self.downsamples = nn.Sequential(*downsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        x_copy = x.clone()
+        for module in self.downsamples:
+            x = module(x, feat_cache, feat_idx)
+
+        return x + self.avg_shortcut(x_copy)
+
+
+class Up_ResidualBlock(nn.Module):
+
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 dropout,
+                 mult,
+                 temperal_upsample=False,
+                 up_flag=False):
+        super().__init__()
+        # Shortcut path with upsample
+        if up_flag:
+            self.avg_shortcut = DupUp3D(
+                in_dim,
+                out_dim,
+                factor_t=2 if temperal_upsample else 1,
+                factor_s=2 if up_flag else 1,
+            )
+        else:
+            self.avg_shortcut = None
+
+        # Main path with residual blocks and upsample
+        upsamples = []
+        for _ in range(mult):
+            upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+
+        # Add the final upsample block
+        if up_flag:
+            mode = "upsample3d" if temperal_upsample else "upsample2d"
+            upsamples.append(Resample(out_dim, mode=mode))
+
+        self.upsamples = nn.Sequential(*upsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+        x_main = x.clone()
+        for module in self.upsamples:
+            x_main = module(x_main, feat_cache, feat_idx)
+        if self.avg_shortcut is not None:
+            x_shortcut = self.avg_shortcut(x, first_chunk)
+            return x_main + x_shortcut
+        else:
+            return x_main
+
+
+class Encoder3d(nn.Module):
+
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        # init block
+        self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_down_flag = (
+                temperal_downsample[i]
+                if i < len(temperal_downsample) else False)
+            downsamples.append(
+                Down_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks,
+                    temperal_downsample=t_down_flag,
+                    down_flag=i != len(dim_mult) - 1,
+                ))
+            scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout),
+            AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout),
+        )
+
+        # # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, z_dim, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device),
+                        cache_x,
+                    ],
+                    dim=2,
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## downsamples
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+
+        return x
+
+
+class Decoder3d(nn.Module):
+
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_upsample=[False, True, True],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2**(len(dim_mult) - 2)
+        # init block
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout),
+            AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout),
+        )
+
+        # upsample blocks
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_up_flag = temperal_upsample[i] if i < len(
+                temperal_upsample) else False
+            upsamples.append(
+                Up_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks + 1,
+                    temperal_upsample=t_up_flag,
+                    up_flag=i != len(dim_mult) - 1,
+                ))
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, 12, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device),
+                        cache_x,
+                    ],
+                    dim=2,
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## upsamples
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx, first_chunk)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class WanVAE_(nn.Module):
+
+    def __init__(
+        self,
+        dim=160,
+        dec_dim=256,
+        z_dim=16,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        # modules
+        self.encoder = Encoder3d(
+            dim,
+            z_dim * 2,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_downsample,
+            dropout,
+        )
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(
+            dec_dim,
+            z_dim,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_upsample,
+            dropout,
+        )
+
+    def forward(self, x, scale=[0, 1]):
+        mu = self.encode(x, scale)
+        x_recon = self.decode(mu, scale)
+        return x_recon, mu
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        x = patchify(x, patch_size=2)
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(
+                    x[:, :, :1, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        self.clear_cache()
+        return mu
+
+    def decode(self, z, scale):
+        """Batch decode. Numerically identical to streaming begin/step/end."""
+        self.decode_stream_begin(scale)
+        out = self.decode_stream_step(z)
+        self.decode_stream_end()
+        return out
+
+    def decode_stream_begin(self, scale):
+        """Start a causal streaming decode session (keeps feat_cache)."""
+        self.clear_cache()
+        self._stream_scale = scale
+        self._stream_frame_idx = 0
+
+    def decode_stream_step(self, z):
+        """Decode one or more latent frames with the live feat_cache.
+
+        Args:
+            z: ``(B, z_dim, T, H, W)`` latent chunk (T >= 1), same scale space
+                as ``encode`` / batch ``decode`` inputs.
+
+        Returns:
+            RGB video ``(B, 3, T_rgb, H', W')`` in ``[-1, 1]`` for this chunk.
+        """
+        assert hasattr(self, "_stream_scale"), (
+            "decode_stream_begin() must be called before decode_stream_step()"
+        )
+        scale = self._stream_scale
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+
+        outs = []
+        for i in range(z.shape[2]):
+            # CausalConv3d(k=1): per-frame conv2 matches full-sequence conv2.
+            x_i = self.conv2(z[:, :, i:i + 1, :, :])
+            self._conv_idx = [0]
+            out_i = self.decoder(
+                x_i,
+                feat_cache=self._feat_map,
+                feat_idx=self._conv_idx,
+                first_chunk=(self._stream_frame_idx == 0),
+            )
+            outs.append(out_i)
+            self._stream_frame_idx += 1
+        out = torch.cat(outs, dim=2)
+        return unpatchify(out, patch_size=2)
+
+    def decode_stream_end(self):
+        """Finish a streaming decode session and drop feat_cache."""
+        self.clear_cache()
+        if hasattr(self, "_stream_scale"):
+            del self._stream_scale
+        if hasattr(self, "_stream_frame_idx"):
+            del self._stream_frame_idx
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        # cache encode
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", **kwargs):
+    # params
+    cfg = dict(
+        dim=dim,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, True],
+        dropout=0.0,
+    )
+    cfg.update(**kwargs)
+
+    # init model
+    with torch.device("meta"):
+        model = WanVAE_(**cfg)
+
+    # load checkpoint
+    logging.info(f"loading {pretrained_path}")
+    model.load_state_dict(
+        torch.load(pretrained_path, map_location=device), assign=True)
+
+    return model
+
+
+class Wan2_2_VAE:
+
+    def __init__(
+        self,
+        z_dim=48,
+        c_dim=160,
+        vae_pth=None,
+        dim_mult=[1, 2, 4, 4],
+        temperal_downsample=[False, True, True],
+        dtype=torch.float,
+        device="cuda",
+    ):
+
+        self.dtype = dtype
+        self.device = device
+
+        mean = torch.tensor(
+            [
+                -0.2289,
+                -0.0052,
+                -0.1323,
+                -0.2339,
+                -0.2799,
+                0.0174,
+                0.1838,
+                0.1557,
+                -0.1382,
+                0.0542,
+                0.2813,
+                0.0891,
+                0.1570,
+                -0.0098,
+                0.0375,
+                -0.1825,
+                -0.2246,
+                -0.1207,
+                -0.0698,
+                0.5109,
+                0.2665,
+                -0.2108,
+                -0.2158,
+                0.2502,
+                -0.2055,
+                -0.0322,
+                0.1109,
+                0.1567,
+                -0.0729,
+                0.0899,
+                -0.2799,
+                -0.1230,
+                -0.0313,
+                -0.1649,
+                0.0117,
+                0.0723,
+                -0.2839,
+                -0.2083,
+                -0.0520,
+                0.3748,
+                0.0152,
+                0.1957,
+                0.1433,
+                -0.2944,
+                0.3573,
+                -0.0548,
+                -0.1681,
+                -0.0667,
+            ],
+            dtype=dtype,
+            device=device,
+        )
+        std = torch.tensor(
+            [
+                0.4765,
+                1.0364,
+                0.4514,
+                1.1677,
+                0.5313,
+                0.4990,
+                0.4818,
+                0.5013,
+                0.8158,
+                1.0344,
+                0.5894,
+                1.0901,
+                0.6885,
+                0.6165,
+                0.8454,
+                0.4978,
+                0.5759,
+                0.3523,
+                0.7135,
+                0.6804,
+                0.5833,
+                1.4146,
+                0.8986,
+                0.5659,
+                0.7069,
+                0.5338,
+                0.4889,
+                0.4917,
+                0.4069,
+                0.4999,
+                0.6866,
+                0.4093,
+                0.5709,
+                0.6065,
+                0.6415,
+                0.4944,
+                0.5726,
+                1.2042,
+                0.5458,
+                1.6887,
+                0.3971,
+                1.0600,
+                0.3943,
+                0.5537,
+                0.5444,
+                0.4089,
+                0.7468,
+                0.7744,
+            ],
+            dtype=dtype,
+            device=device,
+        )
+        self.scale = [mean, 1.0 / std]
+
+        # init model
+        self.model = (
+            _video_vae(
+                pretrained_path=vae_pth,
+                z_dim=z_dim,
+                dim=c_dim,
+                dim_mult=dim_mult,
+                temperal_downsample=temperal_downsample,
+            ).eval().requires_grad_(False).to(device))
+
+    def encode(self, videos):
+        try:
+            with torch.amp.autocast(device_type='cuda', dtype=self.dtype):
+                if not isinstance(videos, list):
+                    return self.model.encode(videos, self.scale).float()
+                else:
+                    return [
+                        self.model.encode(u.unsqueeze(0),
+                                        self.scale).float().squeeze(0)
+                        for u in videos
+                    ]
+        except TypeError as e:
+            logging.info(e)
+            return None
+
+    def decode(self, zs):
+        try:
+            with torch.amp.autocast(device_type='cuda', dtype=self.dtype):
+                if not isinstance(zs, list):
+                    return self.model.decode(zs, self.scale).float().clamp_(-1, 1)
+                else:
+                    return [
+                        self.model.decode(u.unsqueeze(0),
+                                        self.scale).float().clamp_(-1,
+                                                                    1).squeeze(0)
+                        for u in zs
+                    ]
+        except TypeError as e:
+            logging.info(e)
+            return None
+
+    def decode_stream_begin(self):
+        """Begin causal streaming decode (feat_cache retained across steps)."""
+        self.model.decode_stream_begin(self.scale)
+
+    def decode_stream_step(self, z):
+        """Decode a latent chunk ``(B,C,T,H,W)``; returns RGB in ``[-1, 1]``."""
+        with torch.amp.autocast(device_type='cuda', dtype=self.dtype):
+            return self.model.decode_stream_step(z).float().clamp_(-1, 1)
+
+    def decode_stream_end(self):
+        """End streaming decode and clear feat_cache."""
+        self.model.decode_stream_end()

requirements.txt

@@ -0,0 +1,6 @@
+numpy
+einops
+Pillow
+imageio
+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