miniworld/denoiser.py

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)