Blame for space/vendor/models/diffusion_forcing_position_wan.py

"""Position-conditioned diffusion-forcing models with an isolated engine.



The original models.diffusion_forcing_wan module is intentionally not modified.

This module carries its own position-capable engine so direct per-frame root

motion and WAN-packed root conditioning can share identical offline/stream

logic without changing the base DiffForcingWanModel behavior.  Root motion means

only XZ difference plus angular difference; pelvis/root pose remains in the

diffused pose state.

"""

import math
from fractions import Fraction
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn

from .tools.t5 import T5EncoderModel
from .tools.wan_model import WanModel
from .normalization import load_statistics
from .quadratic_loss import load_quadratic_loss_matrix, quadratic_error_sum

EPSILON = 0.05
# Most-negative rope frame index supported by rope_params_with_negative(1024, ...).
# Clamping at this value keeps trim_stream's shifted ids inside the cached range.
ROPE_MIN_INDEX = -1024


class TriangularTimeScheduler:
    def __init__(self, config):
        self.steps = config["steps"]
        self.chunk_size = config["chunk_size"]
        self.random_epsilon = config.get("random_epsilon", 0.00)  # schedule jittering
        self.noise_type = config.get("noise_type", "linear")
        self.sigma_type = config.get("sigma_type", "zero")  # "zero", "memoryless"

        if self.noise_type == "exponential" or self.noise_type == "exponential_rev":
            self.exp_max = config.get("exp_max", 5.0)
        elif self.noise_type == "diffusion":
            self.T = config.get("T", 1000)
            self.beta_start = config.get("beta_start", 0.0001)
            self.beta_end = config.get("beta_end", 0.02)

        if self.sigma_type == "memoryless":
            self.sigma_scale = config.get("sigma_scale", 1.0)
        self.content_len = config.get("content_len", None)
        # True (default): t_max covers the tail ramp-down in place (pairs
        # with extra_len=0). False: original L/c; cover the tail with
        # extra_len padding (typically chunk_size-1) at generate time.
        self.t_max_tail = bool(config.get("t_max_tail", True))
        # For simplicity we require steps to be divisible by chunk_size, so that time windows align well.

    def _t_max_exact(self, valid_len):
        """Exact rational t upper bound (single source of truth).

        t_max_tail: 1 + (L-1)/c so the last frame's band is covered in

        place; otherwise the original L/c (tail covered via extra_len)."""
        if self.t_max_tail:
            return Fraction(valid_len + self.chunk_size - 1, self.chunk_size)
        return Fraction(valid_len, self.chunk_size)

    def get_total_steps(self, seq_len):
        # Exact arithmetic: deriving via float double-rounds and can lose the
        # final step (e.g. steps=30, c=30, L=94).
        total = self.steps * self._t_max_exact(seq_len)
        assert total.denominator == 1, (
            f"inference grid does not close: steps={self.steps} * t_max"
            f"({seq_len}+{self.chunk_size}-1)/{self.chunk_size} is not an "
            f"integer; use steps divisible by chunk_size"
        )
        return int(total)

    def get_t_max(self, valid_len):
        """Float view of _t_max_exact (for continuous t sampling)."""
        return float(self._t_max_exact(valid_len))

    def get_time_steps(self, device, valid_len, current_step):
        """current_step (int) → one t scalar shared by the whole batch,

        returned as a per-sample List. Inference-only (generate / stream);

        training's stratified sampling bypasses this and calls the

        lower-level helpers directly."""
        t = current_step * (1 / self.steps)
        return [torch.tensor(t, device=device) for _ in range(len(valid_len))]

    def get_time_schedules(self, device, valid_len, time_steps, training=False):
        time_schedules = []
        time_schedules_derivative = []
        for i in range(len(valid_len)):
            t = time_steps[i].item()
            current_time_schedules = torch.clamp(
                -torch.arange(valid_len[i], device=device) / self.chunk_size + t,
                min=0.0,
                max=1.0,
            )
            current_time_schedules_derivative = torch.ones_like(
                current_time_schedules
            ) * (1 / self.steps)
            if training:
                current_time_schedules = torch.clamp(
                    current_time_schedules
                    + torch.randn_like(current_time_schedules) * self.random_epsilon,
                    min=0.0,
                    max=1.0,
                )
            time_schedules.append(current_time_schedules)
            time_schedules_derivative.append(current_time_schedules_derivative)
        return time_schedules, time_schedules_derivative

    def get_windows(self, valid_len, time_steps, training=False):
        # Grid t (inference) lands exactly on band boundaries; +eps makes the
        # half-open [0, 1) level convention robust to float rounding there.
        # Continuous t (training) never hits boundaries; eps would only shift
        # the band by eps/chunk_size, so it is skipped.
        eps = 0.0 if training else 0.5 * (1 / (self.steps * self.chunk_size))
        input_start, input_end, output_start, output_end = [], [], [], []
        for i in range(len(time_steps)):
            t = time_steps[i].item()
            start_index = max(
                0,
                math.floor((t - 1) * self.chunk_size + eps) + 1,
            )
            end_index = min(
                valid_len[i],
                math.floor(t * self.chunk_size + eps) + 1,
            )
            # An exact boundary hit at the sequence end gives start == end;
            # the last frame is still in the band, so fall back to [end-1, end).
            start_index = min(start_index, max(end_index - 1, 0))

            if self.content_len is not None:
                input_start.append(max(0, end_index - self.content_len))
            else:
                input_start.append(0)
            input_end.append(end_index)
            output_start.append(start_index)
            output_end.append(end_index)
        return input_start, input_end, output_start, output_end

    def get_noise_levels(self, device, valid_len, time_schedules):
        alpha = []
        dalpha = []
        dlog_alpha = []
        beta = []
        dbeta = []
        dlog_beta = []
        sigma = []
        for i in range(len(valid_len)):
            t = time_schedules[i]
            if self.noise_type == "linear":
                alpha_i = t
                dalpha_i = torch.ones_like(alpha_i)
                dlog_alpha_i = dalpha_i / torch.clamp(alpha_i, min=EPSILON)
                beta_i = 1 - t
                dbeta_i = -torch.ones_like(beta_i)
                dlog_beta_i = dbeta_i / torch.clamp(beta_i, min=EPSILON)
            elif self.noise_type == "exponential":
                # "eps" prediction
                k = self.exp_max
                alpha_i = torch.exp(-k * (1 - t))
                dalpha_i = k * alpha_i
                dlog_alpha_i = k * torch.ones_like(alpha_i)
                beta_i = 1 - alpha_i
                dbeta_i = -dalpha_i
                dlog_beta_i = dbeta_i / torch.clamp(beta_i, min=EPSILON)
            elif self.noise_type == "exponential_rev":
                # "x0" prediction
                k = self.exp_max
                beta_i = torch.exp(-k * t)
                dbeta_i = -k * beta_i
                dlog_beta_i = -k * torch.ones_like(beta_i)
                alpha_i = 1 - beta_i
                dalpha_i = -dbeta_i
                dlog_alpha_i = dalpha_i / torch.clamp(alpha_i, min=EPSILON)
            elif self.noise_type == "diffusion":
                t_rev = 1.0 - t
                beta_rate = (
                    self.beta_start + t_rev * (self.beta_end - self.beta_start)
                ) * self.T
                Gamma = (
                    self.beta_start * t_rev
                    + 0.5 * (self.beta_end - self.beta_start) * t_rev * t_rev
                ) * self.T
                alpha_i = torch.exp(-0.5 * Gamma)
                dalpha_i = 0.5 * beta_rate * alpha_i
                dlog_alpha_i = 0.5 * beta_rate
                beta_i = torch.sqrt(torch.clamp(1 - torch.exp(-Gamma), min=0.0))
                dbeta_i = (
                    -0.5
                    * torch.exp(-Gamma)
                    * beta_rate
                    / torch.clamp(beta_i, min=EPSILON)
                )
                dlog_beta_i = dbeta_i / torch.clamp(beta_i, min=EPSILON)
            else:
                raise ValueError(f"Unknown noise type: {self.noise_type}")
            alpha.append(torch.clamp(alpha_i, min=0.0, max=1.0))
            dalpha.append(dalpha_i)
            dlog_alpha.append(dlog_alpha_i)
            beta.append(torch.clamp(beta_i, min=0.0, max=1.0))
            dbeta.append(dbeta_i)
            dlog_beta.append(dlog_beta_i)
            if self.sigma_type == "zero":
                sigma_i = torch.zeros_like(t)
            elif self.sigma_type == "memoryless":
                if (
                    self.noise_type == "linear"
                    or self.noise_type == "exponential"
                    or self.noise_type == "exponential_rev"
                ):
                    sigma_i = self.sigma_scale * torch.sqrt(
                        torch.clamp(2 * dlog_alpha_i * beta_i, min=0.0)
                    )
                elif self.noise_type == "diffusion":
                    sigma_i = self.sigma_scale * torch.sqrt(
                        torch.clamp(2 * dlog_alpha_i, min=0.0)
                    )
                else:
                    sigma_i = self.sigma_scale * torch.sqrt(
                        torch.clamp(
                            2 * beta_i * (dlog_alpha_i * beta_i - dbeta_i), min=0.0
                        )
                    )
            sigma.append(sigma_i)
        return alpha, dalpha, beta, dbeta, sigma, dlog_alpha, dlog_beta

    def slice_input(self, x, input_start, input_end):
        """Inference-side: slice each sample's current state to the input window."""
        return [x[i][:, input_start[i] : input_end[i], ...] for i in range(len(x))]

    def prepare(self, x, device, valid_len, current_step):
        """Inference-side helper: current_step (int) → all schedule/noise/window

        coefficients sliced to their input/output windows, plus xt sliced to

        input window. Training-side stratified sampling does not go through

        this helper; it calls the lower-level scheduler methods directly.



        Returns dict with keys:

            time_schedules, time_schedules_derivative,

            alpha, dalpha, beta, dbeta, sigma, dlog_alpha, dlog_beta,

            input_start, input_end, output_start, output_end, xt

        """
        time_steps = self.get_time_steps(device, valid_len, current_step)
        time_schedules, time_schedules_derivative = self.get_time_schedules(
            device, valid_len, time_steps,
        )
        alpha, dalpha, beta, dbeta, sigma, dlog_alpha, dlog_beta = \
            self.get_noise_levels(device, valid_len, time_schedules)
        input_start, input_end, output_start, output_end = \
            self.get_windows(valid_len, time_steps)
        xt = self.slice_input(x, input_start, input_end)

        # Slice all coefficients to their respective windows
        batch_size = len(valid_len)
        time_schedules = [time_schedules[i][input_start[i]:input_end[i]] for i in range(batch_size)]
        time_schedules_derivative = [time_schedules_derivative[i][output_start[i]:output_end[i]] for i in range(batch_size)]
        alpha = [alpha[i][output_start[i]:output_end[i]] for i in range(batch_size)]
        dalpha = [dalpha[i][output_start[i]:output_end[i]] for i in range(batch_size)]
        beta = [beta[i][output_start[i]:output_end[i]] for i in range(batch_size)]
        dbeta = [dbeta[i][output_start[i]:output_end[i]] for i in range(batch_size)]
        sigma = [sigma[i][output_start[i]:output_end[i]] for i in range(batch_size)]
        dlog_alpha = [dlog_alpha[i][output_start[i]:output_end[i]] for i in range(batch_size)]
        dlog_beta = [dlog_beta[i][output_start[i]:output_end[i]] for i in range(batch_size)]

        return {
            "time_schedules": time_schedules,
            "time_schedules_derivative": time_schedules_derivative,
            "input_start": input_start,
            "input_end": input_end,
            "output_start": output_start,
            "output_end": output_end,
            "alpha": alpha,
            "dalpha": dalpha,
            "beta": beta,
            "dbeta": dbeta,
            "sigma": sigma,
            "dlog_alpha": dlog_alpha,
            "dlog_beta": dlog_beta,
            "xt": xt,
        }

    # --- Streaming support ---

    def get_committable(self, total_frames):
        """Given total accumulated conditions, return how many frames can be committed.

        Currently, we suppose steps % chunk_size == 0 for simplicity."""
        committable_length = max(0, total_frames - self.chunk_size + 1)
        committable_steps = total_frames * (self.steps // self.chunk_size)
        return committable_length, committable_steps

    def get_step_rollback(self, seq_len):
        """Get the step count to subtract when wrapping the buffer by seq_len.

        Corresponds to how many steps were consumed by seq_len frames."""
        steps = seq_len * (self.steps // self.chunk_size)
        return steps


class T5TextCrossModule(nn.Module):
    """Cross-attention module for T5 text conditioning."""

    def __init__(

        self,

        len=512,

        dim=4096,

        t5_size="xxl",

        checkpoint_path=None,

        tokenizer_path=None,

        drop_out=0.1,

        cross_rope=False,

        input_keys={

            "text": "text",

            "text_end": "text_end",

        },

        cache_encoded=True,

    ):
        assert checkpoint_path is not None and tokenizer_path is not None, (
            "T5 checkpoint and tokenizer paths must be provided."
        )
        super().__init__()
        self.len = len
        self.dim = dim
        self.cross_attn_norm = True
        self.cross_rope = cross_rope
        self.drop_out = drop_out
        self.input_keys = input_keys

        self.text_encoder = T5EncoderModel(
            text_len=len,
            dtype=torch.bfloat16,
            device=torch.device("cpu"),
            checkpoint_path=checkpoint_path,
            tokenizer_path=tokenizer_path,
            shard_fn=None,
            t5_size=t5_size,
        )
        # When False, text_cache is cleared after each encode/get_context call
        # so it doesn't accumulate across training steps (~190KB per caption,
        # OOM on 1M+ unique captions).
        self.cache_encoded = bool(cache_encoded)
        self.text_cache = {}

    def encode(self, text_list, device):
        """Encode text list with cache. Returns List[Tensor]."""
        # Deduplicate uncached texts
        texts_to_encode = []
        for text in text_list:
            if text not in self.text_cache and text not in texts_to_encode:
                texts_to_encode.append(text)

        # Batch encode deduplicated texts
        if texts_to_encode:
            self.text_encoder.model.to(device)
            encoded = self.text_encoder(texts_to_encode, device)
            for text, feature in zip(texts_to_encode, encoded):
                self.text_cache[text] = feature.cpu()

        # Collect from cache
        out = [self.text_cache[text].to(device) for text in text_list]
        if not self.cache_encoded:
            self.text_cache.clear()
        return out

    def get_context(self, x, valid_len, device, param_dtype, training=False):
        """Mask-based unified context.



        Both uniform (text=str) and multi-segment (text=List[str] + text_end)

        inputs are funneled into a single representation: for each sample,

        concatenate all (unique) segments along the K dimension and produce a

        (T_i, S_i) cross-attn mask telling which tokens each query frame may see.



        Returns:

            context: List[Tensor (S_i, D_kv)]   (S_i = num_seg_i * L_text)

            metadata: dict with

              'full_text': List[str] (logging only)

              'cross_attn_mask': List[BoolTensor (T_i, S_i)]

              'text_k_rope_ids' (cross_rope=True): List[Tuple]; f_ids

                  has shape (S_i,) and each segment's L_text tokens share the

                  segment-start-frame value. Null samples carry f_ids=0

                  throughout (single segment, L_text tokens).

              'text_q_rope_ids' (cross_rope=True): List[Tuple|None]

                  per-sample Q rope override; non-None only for null samples

                  (zeros over T_i) so cross-attn phase = 0 (matches K=0).

        """
        text_key = self.input_keys.get("text", "text")
        text_end_key = self.input_keys.get("text_end", "text_end")
        metadata = {}

        text_input = x[text_key] if text_key in x else ["" for _ in range(len(valid_len))]
        text_end_input = x.get(text_end_key, None)

        # Normalize each sample to (text_list, text_end_list, is_null).
        # text_end_list always starts with 0 and ends with valid_len[i].
        norm = []
        for i in range(len(valid_len)):
            is_null = training and np.random.rand() <= self.drop_out
            if is_null:
                norm.append(([""], [0, valid_len[i]], True))
                continue
            if isinstance(text_input[i], list):
                tl = text_input[i]
                assert text_end_input is not None, (
                    "text_end required when text is List[str] (multi-segment)"
                )
                te = [0] + [min(int(t), valid_len[i]) for t in text_end_input[i]]
            else:
                tl = [text_input[i]]
                te = [0, valid_len[i]]
            # Force last segment to extend to valid_len so batch-padding
            # frames (generate passes the batch-max seq_len as every sample's
            # generated_len) get text K coverage instead of an all-False mask.
            # Mirrors PositionCrossModule. No-op for the single-segment
            # branch where te[-1] is already valid_len[i].
            te[-1] = valid_len[i]
            norm.append((tl, te, False))

        # Encode all unique texts (cache-aware dedup)
        texts_to_encode = []
        for tl, _, _ in norm:
            for t in tl:
                if t not in self.text_cache and t not in texts_to_encode:
                    texts_to_encode.append(t)
        if texts_to_encode:
            self.text_encoder.model.to(device)
            encoded = self.text_encoder(texts_to_encode, device)
            for t, f in zip(texts_to_encode, encoded):
                self.text_cache[t] = f.cpu()

        all_context = []
        all_mask = []
        all_text_k_rope_ids = [] if self.cross_rope else None
        text_q_rope_override = [] if self.cross_rope else None
        full_text = []

        for i, (tl, te, is_null) in enumerate(norm):
            seg_tensors = [self.text_cache[t].to(device).to(param_dtype) for t in tl]
            # T5 returns variable-length tokens per caption (trimmed to its
            # actual mask), so segments may have different K-dim lengths.
            # Use cumulative offsets, NOT k * L_text (would mis-stride and
            # leak/lose K tokens whenever segments have different lengths).
            seg_K_lens = [t.size(0) for t in seg_tensors]
            seg_K_offsets = [0]
            for L in seg_K_lens:
                seg_K_offsets.append(seg_K_offsets[-1] + L)
            seg_lens = [t - b for t, b in zip(te[1:], te[:-1])]

            context_i = torch.cat(seg_tensors, dim=0)  # (sum(seg_K_lens), D_kv)
            S_i = context_i.size(0)
            all_context.append(context_i)

            full_text.append(
                " ////////// ".join([f"{u} //dur:{d}" for u, d in zip(tl, seg_lens)])
            )

            # Frame f belongs to segment k where te[k] <= f < te[k+1];
            # mark mask[f, seg_K_offsets[k]:seg_K_offsets[k+1]] = True.
            T_i = valid_len[i]
            mask = torch.zeros(T_i, S_i, dtype=torch.bool, device=device)
            for k, (s_f, e_f) in enumerate(zip(te[:-1], te[1:])):
                s_f = max(0, s_f)
                e_f = min(T_i, e_f)
                if s_f >= e_f:
                    continue
                mask[s_f:e_f, seg_K_offsets[k]:seg_K_offsets[k+1]] = True
            all_mask.append(mask)

            if self.cross_rope:
                # K rope: per-token segment-start-frame; null samples → 0.
                f_ids = torch.empty(S_i, dtype=torch.long, device=device)
                if is_null:
                    f_ids.zero_()
                else:
                    for k, s_f in enumerate(te[:-1]):
                        f_ids[seg_K_offsets[k]:seg_K_offsets[k+1]] = s_f
                all_text_k_rope_ids.append((f_ids, None, None))

                # Q rope override: real → None (use self-attn default arange);
                # null → zeros so phase = Q(0) - K(0) = 0 (rope-off equivalent).
                if is_null:
                    text_q_rope_override.append(
                        (torch.zeros(T_i, dtype=torch.long, device=device), None, None)
                    )
                else:
                    text_q_rope_override.append(None)

        metadata["full_text"] = full_text
        metadata["cross_attn_mask"] = all_mask
        if self.cross_rope:
            metadata["text_k_rope_ids"] = all_text_k_rope_ids
            metadata["text_q_rope_ids"] = text_q_rope_override
        if not self.cache_encoded:
            self.text_cache.clear()
        return all_context, metadata

    def get_null_context(self, real_context, device, param_dtype, valid_len):
        """Build null context with the mask-path convention.



        Each sample gets a single null segment of L_text tokens (encoded("")),

        an all-True (T_i, L_text) mask, K rope = 0, Q rope override = 0.

        This matches the original uniform-null math (phase = 0 = rope-off

        equivalent) and is the only configuration consistent with the mask

        path's unified shape.



        Args:

            real_context: List[Tensor], used to read batch size.

            valid_len: List[int], per-sample query frame count (T_i).



        Returns:

            null_context: List[Tensor (L_text, D_kv)]

            metadata: dict with cross_attn_mask + (when self.cross_rope=True)

                text_k_rope_ids and text_q_rope_ids.

        """
        batch_size = len(real_context)
        encoded_null = self.encode([""], device)[0].to(param_dtype)
        L_text = encoded_null.size(0)
        metadata = {}

        null_context = [encoded_null for _ in range(batch_size)]
        metadata["cross_attn_mask"] = [
            torch.ones(valid_len[i], L_text, dtype=torch.bool, device=device)
            for i in range(batch_size)
        ]

        if self.cross_rope:
            metadata["text_k_rope_ids"] = [
                (torch.zeros(L_text, dtype=torch.long, device=device), None, None)
                for _ in range(batch_size)
            ]
            metadata["text_q_rope_ids"] = [
                (torch.zeros(valid_len[i], dtype=torch.long, device=device), None, None)
                for i in range(batch_size)
            ]
        return null_context, metadata

    # --- Streaming state management (mask path) ---

    def init_stream(self, batch_size):
        # Per-sample list of segments; each entry is (text, start_frame_in_buffer).
        # End frame of segment k is start of segment k+1 (or stream_frames[i]).
        self.stream_segments = [[] for _ in range(batch_size)]
        # Total frames pushed via update_stream (one per call).
        self.stream_frames = [0 for _ in range(batch_size)]

    def update_stream(self, x, device, param_dtype):
        """Add one frame of conditioning. If text matches the latest segment's

        text, extend (implicit via stream_frames++); otherwise start a new

        segment whose start_frame = current frame index.



        Pre-encodes any uncached texts so get_stream_context is cache-only.

        """
        text_key = self.input_keys.get("text", "text")
        text_input = x[text_key]

        # Cache uncached texts (incl. "" so null path can pull from cache too).
        to_encode = []
        for t in list(text_input) + [""]:
            if t not in self.text_cache and t not in to_encode:
                to_encode.append(t)
        if to_encode:
            self.text_encoder.model.to(device)
            encoded = self.text_encoder(to_encode, device)
            for t, f in zip(to_encode, encoded):
                self.text_cache[t] = f.cpu()

        for i in range(len(self.stream_segments)):
            cur_text = text_input[i] if i < len(text_input) else ""
            cur_frame = self.stream_frames[i]
            if (
                len(self.stream_segments[i]) > 0
                and self.stream_segments[i][-1][0] == cur_text
            ):
                # Same as latest segment → frame counter handles extension.
                pass
            else:
                self.stream_segments[i].append((cur_text, cur_frame))
            self.stream_frames[i] += 1

    def get_stream_context(self, start_index, end_index, device, param_dtype):
        """Return mask-form context for buffer window [start_index, end_index).



        Returns:

            contexts: List[Tensor (S_i, D_kv)]

            metadata: {

                'cross_attn_mask': List[BoolTensor (window_len, S_i)]

                'text_k_rope_ids' (cross_rope=True): List[Tuple]

                    f_ids = segment_start_frame (buffer-absolute), repeated L_text

                    times within each segment.

                'text_q_rope_ids' (cross_rope=True): List[None] (real path).

            }

        """
        window_len = end_index - start_index
        contexts = []
        masks = []
        rope_ids = [] if self.cross_rope else None
        text_q_rope_overrides = [] if self.cross_rope else None

        for i in range(len(self.stream_segments)):
            segs = self.stream_segments[i]
            if len(segs) == 0:
                raise RuntimeError(
                    f"sample {i}: stream not initialized (call update_stream first)"
                )
            seg_starts = [s for _, s in segs]
            seg_ends = seg_starts[1:] + [self.stream_frames[i]]

            # Active = segments overlapping [start_index, end_index).
            active = []
            for k, (text, s_start) in enumerate(segs):
                s_end = seg_ends[k]
                if s_end <= start_index or s_start >= end_index:
                    continue
                active.append((text, s_start, s_end))
            if len(active) == 0:
                # Window entirely outside any segment — fall back to null seg
                # covering the whole window so attention still has valid K.
                active = [("", start_index, end_index)]

            seg_tensors = [
                self.text_cache[text].to(device).to(param_dtype)
                for (text, _, _) in active
            ]
            # T5 returns variable-length per-caption tokens; segments may
            # have different K-dim lengths. Use cumulative offsets, not
            # k * L_text (would mis-stride and leak/lose K tokens).
            seg_K_lens = [t.size(0) for t in seg_tensors]
            seg_K_offsets = [0]
            for L in seg_K_lens:
                seg_K_offsets.append(seg_K_offsets[-1] + L)
            context_i = torch.cat(seg_tensors, dim=0)
            S_i = context_i.size(0)
            contexts.append(context_i)

            mask = torch.zeros(window_len, S_i, dtype=torch.bool, device=device)
            for k, (_, s_start, s_end) in enumerate(active):
                local_start = max(0, s_start - start_index)
                local_end = min(window_len, s_end - start_index)
                if local_start < local_end:
                    mask[local_start:local_end, seg_K_offsets[k]:seg_K_offsets[k+1]] = True
            masks.append(mask)

            if self.cross_rope:
                f_ids = torch.empty(S_i, dtype=torch.long, device=device)
                for k, (_, s_start, _) in enumerate(active):
                    f_ids[seg_K_offsets[k]:seg_K_offsets[k+1]] = s_start
                rope_ids.append((f_ids, None, None))
                text_q_rope_overrides.append(None)

        metadata = {"cross_attn_mask": masks}
        if self.cross_rope:
            metadata["text_k_rope_ids"] = rope_ids
            metadata["text_q_rope_ids"] = text_q_rope_overrides
        return contexts, metadata

    def trim_stream(self, trim_len):
        """Shift segment starts back by trim_len; drop segments fully out of buffer.



        Negative starts are kept (clamped to ROPE_MIN_INDEX) so K rope can still

        encode pre-buffer positions via rope_params_with_negative.

        """
        for i in range(len(self.stream_segments)):
            segs = self.stream_segments[i]
            if len(segs) == 0:
                self.stream_frames[i] -= trim_len
                continue
            seg_starts = [s for _, s in segs]
            seg_ends = seg_starts[1:] + [self.stream_frames[i]]
            new_segs = []
            for k, (text, s_start) in enumerate(segs):
                s_end = seg_ends[k]
                new_start = s_start - trim_len
                new_end = s_end - trim_len
                if new_end <= 0:
                    continue  # fully before buffer
                new_segs.append((text, max(new_start, ROPE_MIN_INDEX)))
            self.stream_segments[i] = new_segs
            self.stream_frames[i] -= trim_len


class DiffForcingWanModel(nn.Module):
    def __init__(

        self,

        input_dim=256,

        mean_path=None,

        std_path=None,

        loss_W=None,

        loss_w_coefficient=1.0,

        hidden_dim=1024,

        ffn_dim=2048,

        freq_dim=256,

        num_heads=8,

        num_layers=8,

        time_embedding_scale=1.0,

        attn_type="full",     # "causal" | "full" | "partial"

        rope_channel_split=[1, 0, 0],

        spatial_shape=(1, 1),

        prediction_type="vel",  # "vel", "x0", "eps"

        text_config={

            "len": 512,

            "dim": 4096,

            "cross_rope": False,

        },

        schedule_config={

            "noise_type": "linear",

            "chunk_size": 5,

            "steps": 10,

            "random_epsilon": 0.00,

        },

        cfg_config={

            "text_scale": 5.0,

            "null_scale": -4.0,

        },

        input_keys={

            "feature": "feature",

            "feature_length": "feature_length",

            "text": "text",

            "text_end": "text_end",

        },

        model_input_dim=None,

    ):
        super().__init__()
        self.input_keys = input_keys

        self.mean_path = mean_path
        self.std_path = std_path
        self.input_dim = input_dim
        loss_matrix = load_quadratic_loss_matrix(
            loss_W, loss_w_coefficient, self.input_dim
        )
        if loss_matrix is None:
            self.loss_matrix = None
        else:
            self.register_buffer("loss_matrix", loss_matrix)
        # The diffusion state/output dimensionality remains ``input_dim``.
        # Conditioned variants may expose extra clean channels only to the
        # denoiser (packed root + latent), without diffusing or
        # scoring those channels.
        self.model_input_dim = (
            self.input_dim if model_input_dim is None else int(model_input_dim)
        )
        self.spatial_shape = tuple(spatial_shape)
        self.spatial_size = math.prod(self.spatial_shape)
        self.hidden_dim = hidden_dim
        self.ffn_dim = ffn_dim
        self.freq_dim = freq_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.time_embedding_scale = time_embedding_scale
        assert attn_type in ("causal", "full", "partial"), (
            f"attn_type must be one of 'causal'/'full'/'partial', got {attn_type!r}"
        )
        self.attn_type = attn_type
        self.rope_channel_split = rope_channel_split
        self.prediction_type = prediction_type
        self.cfg_config = cfg_config
        self.schedule_config = schedule_config
        # The multi-window mask is only sound for attn_type='partial': it
        # forces clean x clean to be causal and W_k x clean to be a strict
        # prefix, which contradicts full's "no restriction" and causal's
        # "global sequence causal". Force n_windows=1 under these attn_types
        # so the packed sequence degenerates to a single window and the
        # multi-window mask is bypassed in forward.
        if self.attn_type in ("full", "causal"):
            cfg_n = int(self.schedule_config.get("train_n_windows", 10))
            if cfg_n != 1:
                import warnings
                warnings.warn(
                    f"attn_type={self.attn_type!r} is incompatible with "
                    f"multi-window training (multi-window mask differs from "
                    f"full/causal inference mask). "
                    f"Forcing train_n_windows from {cfg_n} to 1.",
                    stacklevel=2,
                )
                self.schedule_config["train_n_windows"] = 1
        self.time_scheduler = TriangularTimeScheduler(schedule_config)
        # Cross-attention module (text)
        self.text_module = T5TextCrossModule(**text_config)

        mean, std = load_statistics(self.mean_path, self.std_path, input_dim)
        self.register_buffer("mean", mean)
        self.register_buffer("std", std)

        self.model = WanModel(
            patch_size=(1, 1, 1),
            text_len=self.text_module.len,
            text_dim=self.text_module.dim,
            cross_attn_norm=self.text_module.cross_attn_norm,
            cross_rope=self.text_module.cross_rope,
            in_dim=self.model_input_dim,
            dim=self.hidden_dim,
            ffn_dim=self.ffn_dim,
            freq_dim=self.freq_dim,
            out_dim=self.input_dim,
            num_heads=self.num_heads,
            num_layers=self.num_layers,
            window_size=(-1, -1),
            qk_norm=True,
            eps=1e-6,
            causal=(self.attn_type == "causal"),
            rope_channel_split=self.rope_channel_split,
        )
        self.param_dtype = torch.float32

    # ------------------------------------------------------------------
    # Conditioning extension points
    # ------------------------------------------------------------------
    # The base model deliberately keeps these hooks as no-ops.  Position
    # conditioning lives in models/diffusion_forcing_position_wan.py while
    # this class remains bit-for-bit compatible for existing configs.

    def _sample_training_noise(self, x0, inputs, batch_index):
        return torch.randn_like(x0)

    def _get_diffusion_state(self, inputs):
        return inputs["feature"]

    def _prepare_training_condition(

        self, inputs, normalized_feature, valid_len, device

    ):
        return None

    def _condition_training_input(

        self, noisy, clean, condition, batch_index, start, end

    ):
        return noisy

    def _prepare_generation_condition(self, inputs, valid_len, seq_len, device):
        return None

    def _condition_generation_input(self, noisy, condition, starts, ends):
        return noisy

    def _apply_generation_state_condition(self, generated, condition):
        return generated

    def _finalize_generation_output(self, generated, condition, valid_len):
        return generated

    def _generation_output_metadata(self, condition, valid_len):
        return {}

    def _init_stream_condition(

        self, batch_size, buf_len, device, history_position=None,

        history_latent=None,

    ):
        return None

    def _update_stream_condition(self, inputs, device):
        return None

    def _commit_stream_condition(self, index):
        return None

    def _rollback_stream_condition(self, trim_len):
        return None

    def _condition_stream_input(self, noisy, starts, ends):
        return noisy

    def _apply_stream_state_condition(self):
        return None

    def _finalize_stream_output(self, generated, start, end):
        return generated

    def _stream_output_metadata(self, start, end):
        return {}

    def _extract_inputs(self, x):
        """Extract inputs from x using input_keys mapping."""
        inputs = {}
        for internal_key, external_key in self.input_keys.items():
            if external_key in x:
                inputs[internal_key] = x[external_key]
        return inputs

    def preprocess(self, x):
        """Convert last-channel format to channel-first, padding to 4D (C, T, H, W).

        (T, C) -> (C, T, 1, 1)

        (T, H, C) -> (C, T, H, 1)

        (T, H, W, C) -> (C, T, H, W)

        """
        for i in range(len(x)):
            ndim = x[i].ndim
            if ndim == 2:  # (T, C)
                x[i] = x[i].permute(1, 0)[:, :, None, None]
            elif ndim == 3:  # (T, H, C)
                x[i] = x[i].permute(2, 0, 1)[:, :, :, None]
            elif ndim == 4:  # (T, H, W, C)
                x[i] = x[i].permute(3, 0, 1, 2)
        return x

    def postprocess(self, x):
        """Reverse of preprocess: channel-first 4D back to last-channel, stripping padding dims.

        (C, T, 1, 1) -> (T, C)

        (C, T, H, 1) -> (T, H, C)

        (C, T, H, W) -> (T, H, W, C)

        """
        for i in range(len(x)):
            shape = x[i].shape  # (C, T, H, W)
            if shape[2] == 1 and shape[3] == 1:  # (C, T, 1, 1) -> (T, C)
                x[i] = x[i][:, :, 0, 0].permute(1, 0)
            elif shape[3] == 1:  # (C, T, H, 1) -> (T, H, C)
                x[i] = x[i][:, :, :, 0].permute(1, 2, 0)
            else:  # (C, T, H, W) -> (T, H, W, C)
                x[i] = x[i].permute(1, 2, 3, 0)
        return x

    @staticmethod
    def _shift_rope_ids(rope_ids, starts):
        """Subtract starts[i] from f_ids of each tuple, making K rope window-relative

        so cross-attn phase matches Q's default-arange convention even when input

        window doesn't start at frame 0 (content_len truncation).

        Handles per-frame (List[List[Tuple]]), uniform (List[Tuple]), or per-sample

        optional (List[Tuple|None]). None per-sample entries pass through.

        Tuples with f_ids=None (null/CFG placeholders) are also passed through.

        """
        if rope_ids is None:
            return None

        def shift_tuple(tup, s):
            f, h, w = tup
            return ((f - s) if f is not None else None, h, w)

        result = []
        for sample_ids, s in zip(rope_ids, starts):
            if sample_ids is None:
                result.append(None)
            elif isinstance(sample_ids, list):
                result.append([shift_tuple(t, s) for t in sample_ids])
            else:
                result.append(shift_tuple(sample_ids, s))
        return result

    def _patch_seq_len(self, frame_seq_len):
        """Frame-count seq_len → patch-count seq_len."""
        return frame_seq_len * self.spatial_size

    def _expand_schedule_to_patch(self, time_schedules_input):
        """Per-frame schedule → per-patch by repeating each frame value spatial_size times.

        No-op when spatial_size == 1.

        """
        if self.spatial_size == 1:
            return time_schedules_input
        return [s.repeat_interleave(self.spatial_size) for s in time_schedules_input]

    def _expand_context_to_patch(self, context):
        """Per-frame nested context → per-patch by repeating each frame's context spatial_size times.

        Uniform per-sample context (List[Tensor]) is returned unchanged.

        No-op when spatial_size == 1.

        """
        if self.spatial_size == 1:
            return context
        if len(context) > 0 and isinstance(context[0], (list, tuple)):
            return [
                [u for u in sample_ctx for _ in range(self.spatial_size)]
                for sample_ctx in context
            ]
        return context

    def _expand_mask_to_patch(self, masks):
        """Cross-attn mask Q dim (frames) → patch dim by repeating each row

        spatial_size times. K dim untouched. No-op when spatial_size == 1

        or masks is None.

        """
        if masks is None or self.spatial_size == 1:
            return masks
        return [m.repeat_interleave(self.spatial_size, dim=0) for m in masks]

    def _build_partial_mask(self, input_start, input_end, output_start, output_end, device):
        """Per-sample partial-causal attention mask (used when attn_type='partial').



        Layout per sample:

          [input_start, output_start) → history: causal sliding window of size history_len

          [output_start, output_end)  → target : bidirectional + sees all history

        When spatial > 1, repeat-interleave to patch level (F·H·W, F·H·W).

        """
        spatial = self.spatial_size
        masks = []
        for is_, ie_, os_, _oe in zip(input_start, input_end, output_start, output_end):
            L_frame = ie_ - is_
            history_len = os_ - is_
            sliding_window = history_len
            m = torch.zeros(L_frame, L_frame, dtype=torch.bool, device=device)
            for j in range(history_len):
                m[j, max(0, j - sliding_window + 1):j + 1] = True
            m[history_len:, :history_len] = True   # target -> history (fully visible)
            m[history_len:, history_len:] = True   # target x target (bidirectional)
            if spatial > 1:
                m = m.repeat_interleave(spatial, dim=0).repeat_interleave(spatial, dim=1)
            masks.append(m)
        return masks

    def _build_multi_window_attn_mask(self, clean_len, windows, clean_frame_start, device):
        """Self-attn mask for multi-window packed training.



        Packed sequence layout per sample: [clean (clean_len tokens)] + concat over k

        of W_k tokens, where W_k = frames [windows[k][0], windows[k][1]).

        The clean section corresponds to frames

        [clean_frame_start, clean_frame_start + clean_len).

        - clean × clean: causal (clean[i] sees clean[0..i]).

        - clean × W: blocked (clean doesn't depend on noised tokens).

        - W_k × clean[0:s_k - clean_frame_start]: visible — W_k sees every

          clean token in the packed sequence whose absolute frame index is

          less than s_k (strict causal prefix in frame coords).

        - W_k × W_k: bidirectional (denoising chunk full attention).

        - W_k × W_j, j != k: blocked (independent denoising instances).

        When spatial > 1, repeat-interleave to patch level

        (L_frame·spatial, L_frame·spatial).

        """
        sizes = [e - s for s, e in windows]
        total_len = clean_len + sum(sizes)
        mask = torch.zeros(total_len, total_len, dtype=torch.bool, device=device)
        mask[:clean_len, :clean_len] = torch.tril(
            torch.ones(clean_len, clean_len, dtype=torch.bool, device=device)
        )
        cursor = clean_len
        for s, e in windows:
            L = e - s
            s_packed = s - clean_frame_start    # packed col for frame s
            if s_packed > 0:
                mask[cursor:cursor + L, :s_packed] = True
            mask[cursor:cursor + L, cursor:cursor + L] = True
            cursor += L
        if self.spatial_size > 1:
            mask = mask.repeat_interleave(
                self.spatial_size, dim=0
            ).repeat_interleave(self.spatial_size, dim=1)
        return mask

    @staticmethod
    def _gather_for_packed(per_frame_tensor, clean_frame_start, clean_len, windows):
        """Gather rows from a per-frame tensor (T_i, ...) into the packed

        layout [clean section | W_0 | W_1 | ...]. The clean section uses

        rows [clean_frame_start, clean_frame_start + clean_len); each W_k

        uses rows [s_k, e_k). Returns a tensor of shape (packed_len, ...)

        with packed_len = clean_len + sum(L_k).

        """
        parts = []
        if clean_len > 0:
            parts.append(per_frame_tensor[clean_frame_start:clean_frame_start + clean_len])
        for s, e in windows:
            parts.append(per_frame_tensor[s:e])
        return torch.cat(parts, dim=0)

    def _gather_text_q_rope_packed(self, q_rope_ids_full,

                                   clean_frame_starts, clean_lens, all_windows):
        """Gather per-sample text_q_rope_ids tuples (per-sample List[Tuple|None])

        to the packed layout. Per-sample None entries pass through; non-None

        tuples have their f_ids (length T_i) gathered to packed_len following

        the packed layout. h_ids / w_ids are forwarded unchanged.

        """
        if q_rope_ids_full is None:
            return None
        result = []
        for i, item in enumerate(q_rope_ids_full):
            if item is None:
                result.append(None)
                continue
            f_ids, h_ids, w_ids = item
            if f_ids is None:
                result.append(item)
                continue
            gathered = self._gather_for_packed(
                f_ids, clean_frame_starts[i], clean_lens[i], all_windows[i],
            )
            result.append((gathered, h_ids, w_ids))
        return result

    def _build_multi_window_rope_ids(self, clean_len, windows, clean_frame_start, device):
        """Self-attn rope f_ids: per-token position **relative to clean_frame_start**.



        Clean: arange(0, clean_len)。

        W_k: arange(s_k - clean_frame_start, e_k - clean_frame_start)。

        Using positions relative to clean_frame_start matches the inference

        path's default arange(0, e) convention. The Q-K rope phase is still

        the frame difference (the constant offset cancels out after

        subtraction), so the semantics are unchanged.

        """
        parts = [torch.arange(clean_len, device=device, dtype=torch.long)]
        for s, e in windows:
            parts.append(torch.arange(
                s - clean_frame_start, e - clean_frame_start,
                device=device, dtype=torch.long,
            ))
        return (torch.cat(parts), None, None)

    def forward(self, x):
        """Multi-window packed training.



        Per sample (valid_len = T_i), divide t-range [0, (T_i+c-1)/c) into

        `train_n_windows` equal buckets (default 10, configurable via

        schedule_config); each bucket yields one uniform t_k. Windows are

        derived from t_k via `time_scheduler.get_windows`, so the choice of

        n_windows stays consistent with the inference path's formula and

        training cost is independent of sequence length.

        The packed self-attn sequence is

            [clean feature | W_0 | W_1 | ... | W_{N-1}]

        wired together by `_build_multi_window_attn_mask` and per-token rope

        ids from `_build_multi_window_rope_ids`. Cross-attn is uniform (full

        True mask, no rope override, no cross-rope). Loss is per-token MSE

        averaged across all window tokens in the batch.



        When spatial > 1, the metadata (time / attn_mask / cross_attn_mask /

        seq_len) is expanded to patch granularity; rope_ids' f_ids stays at

        frame level since rope_apply broadcasts across H/W internally.

        """
        x = self._extract_inputs(x)
        feature_original = self._get_diffusion_state(x)
        feature_length = x["feature_length"]
        feature_original = (feature_original - self.mean) / self.std
        batch_size = feature_original.shape[0]
        pad_seq_len = feature_original.shape[1]
        device = feature_original.device
        n_windows = int(self.schedule_config.get("train_n_windows", 10))

        valid_len = []
        feature = []
        for i in range(batch_size):
            length = min(feature_length[i].item(), pad_seq_len)
            valid_len.append(length)
            feature.append(feature_original[i, :length, ...])
        feature = self.preprocess(feature)  # list of (C, T_i, 1, 1)
        training_condition = self._prepare_training_condition(
            x, feature_original, valid_len, device,
        )

        # Text context + metadata. cross_attn_mask handles multi-segment Q->K
        # visibility; text_k_rope_ids / text_q_rope_ids carry per-segment rope
        # info when cross_rope=True (None otherwise).
        context, ctx_meta = self.text_module.get_context(
            x, valid_len, device, self.param_dtype, training=True,
        )
        cross_attn_mask_full = ctx_meta["cross_attn_mask"]
        text_k_rope_ids_full = ctx_meta.get("text_k_rope_ids")
        text_q_rope_ids_full = ctx_meta.get("text_q_rope_ids")

        all_xt = []
        all_t = []
        all_rope = []
        all_attn_mask = []
        all_cross_attn_mask = []
        all_targets = []  # per-sample list of (target_tensor, seq_start, seq_end)
        # Per-sample packed layout (used after the loop to gather Q-side text rope
        # and to shift K-side rope into the packed coordinate frame).
        all_clean_frame_starts = []
        all_clean_lens = []
        all_packed_windows = []

        for i in range(batch_size):
            T_i = valid_len[i]
            x0_i = feature[i]
            eps_i = self._sample_training_noise(x0_i, x, i)

            # Split t-range [0, t_max) into n_windows equal buckets; one
            # uniform sample per bucket. Bucket size scales with t_max so the
            # window count stays constant across sequence lengths.
            t_max = self.time_scheduler.get_t_max(T_i)
            bucket = t_max / n_windows
            t_values = [
                torch.tensor(
                    np.random.uniform(k * bucket, (k + 1) * bucket),
                    device=device,
                )
                for k in range(n_windows)
            ]
            input_start_all, _, output_start, output_end = self.time_scheduler.get_windows(
                [T_i] * n_windows, t_values, training=True,
            )
            windows = list(zip(output_start, output_end))

            # t_values is non-decreasing => window starts are
            # monotonically non-decreasing => clean section bounds:
            #   - clean_frame_start = min(input_start): the earliest window's
            #     history start
            #   - clean_frame_end   = max(output_start) = windows[-1][0]
            # Frames beyond clean_frame_end are not read by any W_k, and
            # frames before clean_frame_start are never visible to any W_k
            # (mask gates by frame index), so trimming the buffer to this
            # span is sufficient.
            clean_frame_start = min(input_start_all)
            clean_frame_end = windows[-1][0]
            clean_len = clean_frame_end - clean_frame_start

            # N schedules over T_i frames, one per t_k. Reuses existing
            # scheduler so noise_type / random_epsilon stay consistent with
            # generate.
            schedules, _ = self.time_scheduler.get_time_schedules(
                device, [T_i] * n_windows, t_values, training=True,
            )
            alphas, dalphas, betas, dbetas, _sg, _dla, _dlb = (
                self.time_scheduler.get_noise_levels(
                    device, [T_i] * n_windows, schedules,
                )
            )
            # Clean section: same α·x0 + β·noise noising path as every window,
            # so the latent agrees with the (possibly random_epsilon-jittered)
            # schedule that's also handed to the model as time embedding.
            clean_schedule, _ = self.time_scheduler.get_time_schedules(
                device, [T_i], [torch.tensor(t_max, device=device)],
                training=True,
            )
            clean_alphas, _, clean_betas, _, _, _, _ = (
                self.time_scheduler.get_noise_levels(
                    device, [T_i], clean_schedule,
                )
            )

            # Assemble packed sequence [clean | W_0 | ... | W_{N-1}].
            ca = clean_alphas[0][clean_frame_start:clean_frame_end]
            cb = clean_betas[0][clean_frame_start:clean_frame_end]
            x0_c = x0_i[:, clean_frame_start:clean_frame_end, ...]
            eps_c = eps_i[:, clean_frame_start:clean_frame_end, ...]
            xt_c = (
                x0_c * ca[None, :, None, None]
                + eps_c * cb[None, :, None, None]
            )
            xt_c = self._condition_training_input(
                xt_c,
                x0_c,
                training_condition,
                i,
                clean_frame_start,
                clean_frame_end,
            )
            parts_xt = [xt_c]
            parts_time = [clean_schedule[0][clean_frame_start:clean_frame_end]]
            targets = []
            cursor = clean_len
            for k in range(n_windows):
                s, e = windows[k]
                L = e - s
                a = alphas[k][s:e]
                b = betas[k][s:e]
                x0_w = x0_i[:, s:e, ...]
                eps_w = eps_i[:, s:e, ...]
                xt_w = (
                    x0_w * a[None, :, None, None]
                    + eps_w * b[None, :, None, None]
                )
                xt_w = self._condition_training_input(
                    xt_w, x0_w, training_condition, i, s, e,
                )
                parts_xt.append(xt_w)
                parts_time.append(schedules[k][s:e])

                if self.prediction_type == "vel":
                    da = dalphas[k][s:e]
                    db = dbetas[k][s:e]
                    tgt = (
                        x0_w * da[None, :, None, None]
                        + eps_w * db[None, :, None, None]
                    )
                elif self.prediction_type == "x0":
                    tgt = x0_w
                elif self.prediction_type == "eps":
                    tgt = eps_w
                else:
                    raise ValueError(
                        f"Unknown prediction_type: {self.prediction_type!r}"
                    )
                targets.append((tgt, cursor, cursor + L))
                cursor += L

            xt_i = torch.cat(parts_xt, dim=1)
            assert xt_i.size(0) == self.model_input_dim, (
                f"conditioned denoiser input has {xt_i.size(0)} channels; "
                f"expected model_input_dim={self.model_input_dim}"
            )
            time_i = torch.cat(parts_time) * self.time_embedding_scale

            all_xt.append(xt_i)
            all_t.append(time_i)
            all_rope.append(
                self._build_multi_window_rope_ids(
                    clean_len, windows, clean_frame_start, device,
                )
            )
            all_attn_mask.append(
                self._build_multi_window_attn_mask(
                    clean_len, windows, clean_frame_start, device,
                )
            )
            # Cross-attn mask: gather per-frame rows from cross_attn_mask_full
            # to follow the packed [clean | W_k] layout. Frame-level here;
            # spatial expansion happens after the loop.
            all_cross_attn_mask.append(
                self._gather_for_packed(
                    cross_attn_mask_full[i], clean_frame_start, clean_len, windows,
                )
            )
            all_targets.append(targets)
            all_clean_frame_starts.append(clean_frame_start)
            all_clean_lens.append(clean_len)
            all_packed_windows.append(windows)

        # Expand frame-level time / cross_attn_mask to patch-level
        # (no-ops when spatial_size == 1).
        all_t = self._expand_schedule_to_patch(all_t)
        all_cross_attn_mask = self._expand_mask_to_patch(all_cross_attn_mask)
        # seq_len = max(pad_seq_len, max_packed). In n=1 mode
        # packed_length <= T_i <= pad_seq_len so pad_seq_len dominates and
        # the model sees the same padded geometry as the inference path.
        # With n>1 (partial only) packed can exceed pad_seq_len; the max
        # keeps the upper bound correct.
        seq_len_for_model = self._patch_seq_len(
            max(pad_seq_len, max(xt.size(1) for xt in all_xt))
        )
        text_seq_len = max(int(ctx.size(0)) for ctx in context)

        # Cross-attn rope (only active when cross_rope=True): K rope encodes
        # per-token segment-start-frame in absolute coords, Q rope override
        # exists for null samples (zeros over T_i). Both must be expressed in
        # the same packed coordinate frame as self-attn rope_ids (which is
        # relative to clean_frame_start), so:
        #   - K rope: shift by clean_frame_start (subtract) per sample
        #   - Q rope: gather per-frame f_ids to packed layout, then shift
        text_k_rope_ids = self._shift_rope_ids(
            text_k_rope_ids_full, all_clean_frame_starts,
        )
        text_q_rope_ids = self._gather_text_q_rope_packed(
            text_q_rope_ids_full, all_clean_frame_starts, all_clean_lens,
            all_packed_windows,
        )
        text_q_rope_ids = self._shift_rope_ids(
            text_q_rope_ids, all_clean_frame_starts,
        )

        # The multi-window mask is only meaningful for attn_type='partial'.
        # Under full/causal (where __init__ already forced n_windows=1), the
        # packed sequence is a single window, so we let WanModel apply
        # default full attention or its internal causal mask.
        final_attn_mask = all_attn_mask if self.attn_type == "partial" else None
        predicted_result = self.model(
            all_xt,
            all_t,
            context,
            seq_len_for_model,
            attn_mask=final_attn_mask,
            text_k_rope_ids=text_k_rope_ids,
            text_q_rope_ids=text_q_rope_ids,
            cross_attn_mask=all_cross_attn_mask,
            text_seq_len=text_seq_len,
            y=None,
            rope_ids=all_rope,
        )

        # Per-sample mean -> batch mean. Within each sample sum SE across
        # all W_k tokens and divide by the token count to get per-sample
        # MSE, then average across the batch.
        loss = 0.0
        for i in range(batch_size):
            pred_i = predicted_result[i]
            sample_se_sum = pred_i.new_zeros(())
            sample_n = 0
            for tgt, s, e in all_targets[i]:
                error = pred_i[:, s:e, ...] - tgt
                if self.loss_matrix is None:
                    sample_se_sum = sample_se_sum + (error ** 2).sum()
                else:
                    sample_se_sum = sample_se_sum + quadratic_error_sum(
                        error, self.loss_matrix
                    )
                sample_n += error.numel()
            loss = loss + sample_se_sum / max(sample_n, 1)
        loss = loss / batch_size
        return {"total": loss, "mse": loss}

    def generate(self, x):
        """

        Generation - Diffusion Forcing inference

        Uses triangular noise schedule, progressively generating from left to right



        Generation process:

        1. Start from t=0, gradually increase t

        2. Each t corresponds to a noise schedule: clean on left, noisy on right, gradient in middle

        3. After each denoising step, t increases slightly and continues

        """
        x = self._extract_inputs(x)
        feature_length = x["feature_length"]  # (B,)
        batch_size = len(feature_length)
        # extra_len pads the rollout with disposable frames so real frames
        # finish under full-width windows; 0 relies on the tail ramp-down.
        extra_len = self.schedule_config.get("extra_len", 0)
        seq_len = max(feature_length).item() + extra_len
        if (not getattr(self.time_scheduler, "t_max_tail", True)
                and extra_len < self.time_scheduler.chunk_size - 1):
            import warnings
            warnings.warn(
                f"t_max_tail=False with extra_len={extra_len} < chunk_size-1="
                f"{self.time_scheduler.chunk_size - 1}: the last frames of "
                f"every generated sequence will not finish denoising",
                stacklevel=2,
            )
        device = next(self.parameters()).device
        valid_len = []
        for i in range(batch_size):
            length = min(feature_length[i].item(), seq_len)
            valid_len.append(length)
        generated_len = [seq_len for _ in range(batch_size)]
        generation_condition = self._prepare_generation_condition(
            x, valid_len, seq_len, device,
        )

        # Initialize entire sequence as pure noise
        generated = torch.randn(
            batch_size, seq_len, *self.spatial_shape, self.input_dim, device=device
        )
        generated = [generated[i] for i in range(batch_size)]
        generated = self.preprocess(generated)
        self._apply_generation_state_condition(generated, generation_condition)

        # Precompute real text context (null built per-step from window-sliced real).
        text_context, metadata = self.text_module.get_context(
            x,
            generated_len,
            device,
            self.param_dtype,
            training=False,
        )
        full_text = metadata["full_text"]
        cross_attn_mask_full = metadata["cross_attn_mask"]
        text_k_rope_ids_full = metadata.get("text_k_rope_ids")

        total_steps = self.time_scheduler.get_total_steps(seq_len)
        # Progressively advance from t=0 to t=max_t
        for step in range(total_steps):
            s = self.time_scheduler.prepare(
                generated, device, generated_len, current_step=step,
            )
            time_schedules = s["time_schedules"]
            time_schedules_derivative = s["time_schedules_derivative"]
            alpha = s["alpha"]
            dalpha = s["dalpha"]
            beta = s["beta"]
            dbeta = s["dbeta"]
            sigma = s["sigma"]
            dlog_alpha = s["dlog_alpha"]
            dlog_beta = s["dlog_beta"]
            input_start_index = s["input_start"]
            input_end_index = s["input_end"]
            output_start_index = s["output_start"]
            output_end_index = s["output_end"]
            xt = s["xt"]
            xt = self._condition_generation_input(
                xt, generation_condition, input_start_index, input_end_index,
            )

            # time_schedules already sliced to input window by prepare()
            time_schedules_input = [
                time_schedules[i] * self.time_embedding_scale
                for i in range(batch_size)
            ]

            # Slice cross_attn_mask Q dim to input window (K untouched).
            window_cross_attn_mask = [
                cross_attn_mask_full[i][input_start_index[i]:input_end_index[i], :]
                for i in range(batch_size)
            ]
            # Shift K rope to window-relative coords (lockstep with forward path).
            window_text_k_rope = self._shift_rope_ids(text_k_rope_ids_full, input_start_index)

            # CFG: text_scale * pred_text + null_scale * pred_null
            ts_patch = self._expand_schedule_to_patch(time_schedules_input)
            patch_seq_len = self._patch_seq_len(seq_len)
            attn_mask = (
                self._build_partial_mask(input_start_index, input_end_index,
                                         output_start_index, output_end_index, device)
                if self.attn_type == "partial" else None
            )
            # Build null context matching window-sliced real shape.
            window_lens = [
                input_end_index[i] - input_start_index[i] for i in range(batch_size)
            ]
            window_null_context, window_null_meta = self.text_module.get_null_context(
                text_context, device, self.param_dtype, valid_len=window_lens,
            )
            window_null_mask = window_null_meta["cross_attn_mask"]
            window_null_text_k_rope = window_null_meta.get("text_k_rope_ids")
            window_null_text_q_rope = window_null_meta.get("text_q_rope_ids")

            # Doubled batch: [text..., null...]. One model.forward, split outputs.
            combined_ctx = self._expand_context_to_patch(text_context + window_null_context)
            combined_cross_attn_mask = self._expand_mask_to_patch(
                window_cross_attn_mask + window_null_mask
            )
            if window_text_k_rope is not None and window_null_text_k_rope is not None:
                combined_text_k_rope = window_text_k_rope + window_null_text_k_rope
            else:
                combined_text_k_rope = None
            # Real branch has no Q-override (cross falls back to self_attn rope);
            # null branch emits one when cross_rope=True. Concat [None]*B for real.
            if window_null_text_q_rope is not None:
                combined_text_q_rope = [None] * batch_size + list(window_null_text_q_rope)
            else:
                combined_text_q_rope = None
            combined_mask = attn_mask + attn_mask if attn_mask is not None else None

            combined_text_seq_len = max(
                int(c.size(0)) for c in (text_context + window_null_context)
            )
            pred_all = self.model(
                xt + xt,
                ts_patch + ts_patch,
                combined_ctx,
                patch_seq_len,
                attn_mask=combined_mask,
                text_k_rope_ids=combined_text_k_rope,
                text_q_rope_ids=combined_text_q_rope,
                cross_attn_mask=combined_cross_attn_mask,
                text_seq_len=combined_text_seq_len,
                y=None,
            )
            pred_text = pred_all[:batch_size]
            pred_null = pred_all[batch_size:]
            predicted_result = [
                self.cfg_config["text_scale"] * pt + self.cfg_config["null_scale"] * pn
                for pt, pn in zip(pred_text, pred_null)
            ]

            # All noise coefficients already sliced to output window by prepare()
            for i in range(batch_size):
                os, oe = output_start_index[i], output_end_index[i]
                pred_os = os - input_start_index[i]
                pred_oe = oe - input_start_index[i]
                predicted_result_i = predicted_result[i][:, pred_os:pred_oe, ...]
                generated_i = generated[i][:, os:oe, ...]
                dt = time_schedules_derivative[i][None, :, None, None]
                alpha_i = alpha[i][None, :, None, None]
                dalpha_i = dalpha[i][None, :, None, None]
                beta_i = beta[i][None, :, None, None]
                dbeta_i = dbeta[i][None, :, None, None]
                sigma_i = sigma[i][None, :, None, None]
                dlog_alpha_i = dlog_alpha[i][None, :, None, None]
                dlog_beta_i = dlog_beta[i][None, :, None, None]
                if self.prediction_type == "vel":
                    vel = predicted_result_i
                elif self.prediction_type == "x0":
                    vel = (
                        predicted_result_i * (-dlog_beta_i * alpha_i + dalpha_i)
                        + generated_i * dlog_beta_i
                    )
                elif self.prediction_type == "eps":
                    vel = (
                        predicted_result_i * (-dlog_alpha_i * beta_i + dbeta_i)
                        + generated_i * dlog_alpha_i
                    )
                st = (vel - generated_i * dlog_alpha_i) / (
                    (beta_i * dlog_alpha_i - dbeta_i) * beta_i
                )
                generated[i][:, os:oe, ...] += (
                    vel * dt
                    + st * 0.5 * sigma_i**2 * dt
                    + sigma_i * torch.sqrt(dt) * torch.randn_like(generated_i)
                )
            self._apply_generation_state_condition(generated, generation_condition)

        generated = self.postprocess(generated)  # list of (T, C)
        y_hat_out = []
        for i in range(batch_size):
            single_generated = generated[i][: valid_len[i], :] * self.std + self.mean
            y_hat_out.append(single_generated)
        y_hat_out = self._finalize_generation_output(
            y_hat_out, generation_condition, valid_len,
        )
        out = {}
        out["generated"] = y_hat_out
        out["text"] = full_text
        out.update(self._generation_output_metadata(generation_condition, valid_len))

        return out

    def init_generated(

        self,

        seq_len,

        batch_size=1,

        schedule_config={},

        history_latent=None,

        history_position=None,

    ):
        """Initialize streaming generation state.



        Args:

            seq_len: Model window size (how many frames WanModel processes per step).

            schedule_config: Optional schedule config overrides.

            history_latent: optional list of (h, ...) un-normalized latent tensors per sample

                            (typically vae.encode(motion) output). If provided:

                              - normalize via self.mean/self.std

                              - apply triangular noise matching the schedule at t=1

                                (ts[i] = clamp(1 - i/chunk_size, 0, 1) for i in [0, h))

                              - inject into buffer[:, :h, ...]

                              - set current_step = steps so wavefront starts just past history

                            If None: cold start (buffer all noise, current_step=0).



        Buffer is 2*seq_len. Model window is always buffer[0:seq_len].

        When conditions overflow seq_len, shift buffer by seq_len and restart.

        """
        self.schedule_config.update(schedule_config)
        content_len = self.schedule_config.get("content_len", None)
        if content_len is None:
            self.schedule_config["content_len"] = seq_len
        else:
            self.schedule_config["content_len"] = min(seq_len, content_len)
        self.time_scheduler = TriangularTimeScheduler(self.schedule_config)

        self.batch_size = batch_size
        self.seq_len = seq_len
        self.buf_len = seq_len * 2
        self.current_step = 0
        self.current_commit = 0
        self.condition_frames = 0

        device = next(self.parameters()).device
        # Initialize entire buffer as pure noise
        generated = torch.randn(
            batch_size, self.buf_len, *self.spatial_shape, self.input_dim, device=device
        )
        generated = [generated[i] for i in range(batch_size)]
        self.generated = self.preprocess(generated)

        # Initialize streaming state for cross module
        self.text_module.init_stream(self.batch_size)
        self._init_stream_condition(
            batch_size,
            self.buf_len,
            device,
            history_position=history_position,
            history_latent=history_latent,
        )

        # Optional: inject GT history with triangular re-noising matching schedule at t=1
        if history_latent is not None:
            h = history_latent[0].shape[0]
            chunk_size = self.time_scheduler.chunk_size
            # Triangular ts: ts[0]=1 (clean), ts[chunk_size-1]=1/chunk_size (most noisy),
            # ts[i>=chunk_size]=0 (full noise — beyond effective transition zone)
            ts_history = torch.clamp(
                1 - torch.arange(h, device=device, dtype=torch.float32) / chunk_size,
                min=0.0, max=1.0,
            )
            alpha_h, _, beta_h, _, _, _, _ = self.time_scheduler.get_noise_levels(
                device, [h], [ts_history],
            )
            alpha_b = alpha_h[0][None, :, None, None]
            beta_b = beta_h[0][None, :, None, None]
            for i in range(batch_size):
                z = (history_latent[i].to(device) - self.mean) / self.std
                z = self.preprocess([z])[0]
                noise = torch.randn_like(z)
                z_noised = z * alpha_b + noise * beta_b
                self.generated[i][:, :h, ...] = z_noised
            # Conditioned subclasses may need to restore channels that must
            # remain clean after the base model's triangular history re-noising.
            self._apply_stream_state_condition()
            # t starts from 1 → wavefront just past the history transition zone
            self.current_step = self.time_scheduler.steps

    def _rollback(self):
        """Shift buffer by seq_len when conditions overflow the window."""
        for i in range(self.batch_size):
            self.generated[i][:, : self.seq_len, ...] = self.generated[i][
                :, self.seq_len :, ...
            ].clone()
            self.generated[i][:, self.seq_len :, ...] = torch.randn_like(
                self.generated[i][:, self.seq_len :, ...]
            )
        self.current_step -= self.time_scheduler.get_step_rollback(self.seq_len)
        self.condition_frames -= self.seq_len
        self.current_commit -= self.seq_len
        self.text_module.trim_stream(self.seq_len)
        self._rollback_stream_condition(self.seq_len)

    @torch.no_grad()
    def stream_generate_step(self, x):
        """

        Streaming generation step. Each call provides 1 frame of conditions.

        The scheduler determines committable frames from accumulated conditions.



        Returns:

            dict with "generated": list of one (N, C) tensor, or [] if nothing to commit.

        """
        x = self._extract_inputs(x)
        device = next(self.parameters()).device
        self.generated = [g.to(device) for g in self.generated]

        # 1. Update conditions (1 frame per call)
        self._update_stream_condition(x, device)
        self.text_module.update_stream(x, device, self.param_dtype)
        self.condition_frames += 1

        # 2. Rollback if conditions overflow the window
        if self.condition_frames > self.buf_len:
            self._rollback()
        self._commit_stream_condition(self.condition_frames - 1)

        # 3. Determine how many frames can be committed
        committable_length, committable_steps = self.time_scheduler.get_committable(
            self.condition_frames
        )
        while self.current_step < committable_steps:
            s = self.time_scheduler.prepare(
                self.generated, device, [self.buf_len] * self.batch_size,
                current_step=self.current_step,
            )
            time_schedules = s["time_schedules"]
            time_schedules_derivative = s["time_schedules_derivative"]
            alpha = s["alpha"]
            dalpha = s["dalpha"]
            beta = s["beta"]
            dbeta = s["dbeta"]
            sigma = s["sigma"]
            dlog_alpha = s["dlog_alpha"]
            dlog_beta = s["dlog_beta"]
            is_ = s["input_start"]
            ie_ = s["input_end"]
            os_ = s["output_start"]
            oe_ = s["output_end"]
            xt = s["xt"]
            xt = self._condition_stream_input(xt, is_, ie_)

            # time_schedules already sliced to input window by prepare()
            time_schedules_input = [
                time_schedules[0] * self.time_embedding_scale
            ] * self.batch_size

            # CFG: batch text + null in one forward pass
            text_context, text_meta = self.text_module.get_stream_context(
                is_[0], ie_[0], device, self.param_dtype,
            )
            text_mask = text_meta["cross_attn_mask"]
            text_k_rope = text_meta.get("text_k_rope_ids")
            # Q-rope override (per-sample List[Tuple|None]): real-text stream
            # entries are None and pass through _shift unchanged.
            text_q_rope = text_meta.get("text_q_rope_ids")

            window_len = ie_[0] - is_[0]
            null_context, null_meta = self.text_module.get_null_context(
                text_context, device, self.param_dtype,
                valid_len=[window_len] * self.batch_size,
            )
            null_mask = null_meta["cross_attn_mask"]
            null_text_k_rope = null_meta.get("text_k_rope_ids")
            null_text_q_rope = null_meta.get("text_q_rope_ids")

            ts_patch = self._expand_schedule_to_patch(time_schedules_input)
            ctx_patch = self._expand_context_to_patch(text_context + null_context)
            cross_mask_patch = self._expand_mask_to_patch(text_mask + null_mask)
            # K-rope ids: real is buffer-absolute → shift to window-relative; null
            # comes from get_null_context already window-relative (zeros).
            text_k_rope = self._shift_rope_ids(
                text_k_rope, [is_[0]] * self.batch_size,
            )
            text_q_rope = self._shift_rope_ids(
                text_q_rope, [is_[0]] * self.batch_size,
            )
            if text_k_rope is not None and null_text_k_rope is not None:
                text_k_rope_patch = text_k_rope + null_text_k_rope
            else:
                text_k_rope_patch = None
            # Q-override: real samples carry None (cross uses self_attn rope),
            # null samples carry zeros tuple (phase=0 vs K=0).
            if text_q_rope is not None or null_text_q_rope is not None:
                text_q_rope_patch = (
                    (text_q_rope if text_q_rope is not None
                     else [None] * self.batch_size)
                    + (null_text_q_rope if null_text_q_rope is not None
                       else [None] * self.batch_size)
                )
            else:
                text_q_rope_patch = None

            if self.attn_type == "partial":
                masks = self._build_partial_mask(is_, ie_, os_, oe_, device)
                attn_mask = masks + masks   # CFG doubled batch (text + null), same mask
            else:
                attn_mask = None
            stream_text_seq_len = max(
                int(c.size(0)) for c in (text_context + null_context)
            )
            pred_all = self.model(
                xt + xt,
                ts_patch + ts_patch,
                ctx_patch,
                self._patch_seq_len(self.seq_len),
                attn_mask=attn_mask,
                text_k_rope_ids=text_k_rope_patch,
                text_q_rope_ids=text_q_rope_patch,
                cross_attn_mask=cross_mask_patch,
                text_seq_len=stream_text_seq_len,
                y=None,
            )
            pred_text = pred_all[: self.batch_size]
            pred_null = pred_all[self.batch_size :]
            predicted_result = [
                self.cfg_config["text_scale"] * pt + self.cfg_config["null_scale"] * pn
                for pt, pn in zip(pred_text, pred_null)
            ]

            # All noise coefficients already sliced to output window by prepare()
            os_idx, oe_idx = os_[0], oe_[0]
            pred_os_idx = os_idx - is_[0]
            pred_oe_idx = oe_idx - is_[0]
            dt = time_schedules_derivative[0][None, :, None, None]
            alpha_i = alpha[0][None, :, None, None]
            dalpha_i = dalpha[0][None, :, None, None]
            beta_i = beta[0][None, :, None, None]
            dbeta_i = dbeta[0][None, :, None, None]
            sigma_i = sigma[0][None, :, None, None]
            dlog_alpha_i = dlog_alpha[0][None, :, None, None]
            dlog_beta_i = dlog_beta[0][None, :, None, None]
            for i in range(self.batch_size):
                predicted_result_i = predicted_result[i][:, pred_os_idx:pred_oe_idx, ...]
                generated_i = self.generated[i][:, os_idx:oe_idx, ...]
                if self.prediction_type == "vel":
                    vel = predicted_result_i
                elif self.prediction_type == "x0":
                    vel = (
                        predicted_result_i * (-dlog_beta_i * alpha_i + dalpha_i)
                        + generated_i * dlog_beta_i
                    )
                elif self.prediction_type == "eps":
                    vel = (
                        predicted_result_i * (-dlog_alpha_i * beta_i + dbeta_i)
                        + generated_i * dlog_alpha_i
                    )
                st = (vel - generated_i * dlog_alpha_i) / (
                    (beta_i * dlog_alpha_i - dbeta_i) * beta_i
                )
                self.generated[i][:, os_idx:oe_idx, ...] += (
                    vel * dt
                    + st * 0.5 * sigma_i**2 * dt
                    + sigma_i * torch.sqrt(dt) * torch.randn_like(generated_i)
                )
            self._apply_stream_state_condition()
            self.current_step += 1

        # 5. Extract newly committed frames
        if self.current_commit < committable_length:
            commit_start = self.current_commit
            output = [
                self.generated[i][:, commit_start:committable_length, ...]
                for i in range(self.batch_size)
            ]
            output = self.postprocess(output)
            output = [o * self.std + self.mean for o in output]
            output = self._finalize_stream_output(
                output, commit_start, committable_length,
            )
            self.current_commit = committable_length
            result = {"generated": output}
            result.update(
                self._stream_output_metadata(commit_start, committable_length)
            )
            return result
        else:
            empty = [
                torch.zeros(self.input_dim, 0, *self.spatial_shape, device=device)
                for _ in range(self.batch_size)
            ]
            empty = self.postprocess(empty)
            empty = self._finalize_stream_output(
                empty, self.current_commit, self.current_commit,
            )
            result = {"generated": empty}
            result.update(
                self._stream_output_metadata(
                    self.current_commit, self.current_commit,
                )
            )
            return result

# Position-specific root packing and model contracts.
from .vae_wan_position import (
    POSITION_ROOT_DIMS,
    WAN_ROOT_GROUP_SIZE,
    pack_root_condition,
    resolve_position_representation,
)

# Channels that are identically zero in the ground truth have std ~0.  Dividing
# by that std turns float noise into an O(1) training target, so below this
# floor the channel is centred but not scaled.  The floor sits in the measured
# gap observed in the supported motion statistics.
STD_NORMALIZATION_FLOOR = 1e-3


def guard_std(std, floor=STD_NORMALIZATION_FLOOR):
    """Replace near-zero standard deviations with 1 (i.e. do not divide)."""
    std = torch.as_tensor(std).float()
    return torch.where(std < floor, torch.ones_like(std), std)


def _to_batched_tensor(value, device, dtype=torch.float32):
    if isinstance(value, (list, tuple)):
        tensors = [torch.as_tensor(v, device=device, dtype=dtype) for v in value]
        if not tensors:
            raise ValueError("position condition cannot be empty")
        if tensors[0].ndim == 0:
            if any(tensor.ndim != 0 for tensor in tensors):
                raise ValueError("position list mixes scalar and non-scalar values")
            return torch.stack(tensors, dim=0)
        if tensors[0].ndim == 1:
            return torch.stack(tensors, dim=0)
        lengths = {int(t.size(0)) for t in tensors}
        if len(lengths) != 1:
            raise ValueError(
                "variable-length position lists must be padded/collated before use"
            )
        return torch.stack(tensors, dim=0)
    return torch.as_tensor(value, device=device, dtype=dtype)


class _PositionStreamBufferMixin:
    """Fixed-size, rollback-safe raw condition buffer for stream generation."""

    stream_condition_dim = 0

    def _assert_stream_schedule(self):
        """Subclass hook: validated again after init_generated rebuilds it."""
        return None

    def init_generated(self, seq_len, *args, **kwargs):
        # The base rebuilds time_scheduler from a caller-supplied
        # schedule_config, so schedule invariants must be re-checked.
        result = super().init_generated(seq_len, *args, **kwargs)
        self._assert_stream_schedule()
        return result

    def _init_stream_condition(

        self,

        batch_size,

        buf_len,

        device,

        history_position=None,

        history_latent=None,

    ):
        self._stream_position = torch.zeros(
            batch_size,
            buf_len,
            self.stream_condition_dim,
            device=device,
            dtype=torch.float32,
        )
        self._stream_position_valid = torch.zeros(
            batch_size, buf_len, device=device, dtype=torch.bool
        )
        self._pending_stream_position = None

        if history_position is not None and history_latent is None:
            raise ValueError("history_position requires matching history_latent")
        history_length = None
        if history_latent is not None:
            if len(history_latent) != batch_size:
                raise ValueError(
                    f"history latent batch={len(history_latent)}, expected {batch_size}"
                )
            history_lengths = [int(torch.as_tensor(v).size(0)) for v in history_latent]
            if len(set(history_lengths)) != 1:
                raise ValueError(
                    f"all history latent lengths must match, got {history_lengths}"
                )
            history_length = history_lengths[0]
            if history_length > buf_len:
                raise ValueError(
                    f"history latent length={history_length} exceeds buffer={buf_len}"
                )

        if history_position is None and history_latent is not None:
            history_position = self._position_from_history(history_latent)
        if history_position is not None:
            history = self._coerce_stream_history(history_position, device)
            if history.size(0) != batch_size:
                raise ValueError(
                    f"history position batch={history.size(0)}, expected {batch_size}"
                )
            if history.size(1) > buf_len:
                raise ValueError(
                    f"history position length={history.size(1)} exceeds buffer={buf_len}"
                )
            if history_length is not None and history.size(1) != history_length:
                raise ValueError(
                    f"history position length={history.size(1)} does not match "
                    f"history latent length={history_length}"
                )
            self._stream_position[:, : history.size(1)] = history
            self._stream_position_valid[:, : history.size(1)] = True

    def _position_from_history(self, history_latent):
        return None

    def _coerce_stream_history(self, position, device):
        value = _to_batched_tensor(position, device)
        if value.ndim == 2:
            value = value.unsqueeze(0)
        if value.ndim != 3 or value.size(-1) != self.stream_condition_dim:
            raise ValueError(
                f"history position must be (B,T,{self.stream_condition_dim}), "
                f"got {tuple(value.shape)}"
            )
        if not torch.isfinite(value).all():
            raise ValueError("history position contains NaN or Inf")
        return value.float()

    def _update_stream_condition(self, inputs, device):
        # Validation happens before text/counters are mutated in the base class.
        value = self._parse_stream_position(inputs, device)
        if value.shape != (self.batch_size, self.stream_condition_dim):
            raise ValueError(
                f"stream position must resolve to "
                f"({self.batch_size},{self.stream_condition_dim}), "
                f"got {tuple(value.shape)}"
            )
        if not torch.isfinite(value).all():
            raise ValueError("stream position contains NaN or Inf")
        self._pending_stream_position = value.float()

    def _commit_stream_condition(self, index):
        if self._pending_stream_position is None:
            raise RuntimeError("stream position validation/commit transaction is missing")
        if not 0 <= index < self.buf_len:
            raise IndexError(
                f"stream position index {index} outside buffer [0,{self.buf_len})"
            )
        self._stream_position[:, index] = self._pending_stream_position
        self._stream_position_valid[:, index] = True
        self._pending_stream_position = None

    def _rollback_stream_condition(self, trim_len):
        keep = self.buf_len - trim_len
        self._stream_position[:, :keep] = self._stream_position[
            :, trim_len:
        ].clone()
        self._stream_position_valid[:, :keep] = self._stream_position_valid[
            :, trim_len:
        ].clone()
        self._stream_position[:, keep:] = 0
        self._stream_position_valid[:, keep:] = False

    def _stream_position_slice(self, start, end):
        valid = self._stream_position_valid[:, start:end]
        if valid.numel() and not bool(valid.all()):
            missing = (~valid).nonzero(as_tuple=False)[0].tolist()
            raise RuntimeError(
                f"stream root condition missing inside active window "
                f"[{start},{end}); first missing batch/local-index={missing}"
            )
        return self._stream_position[:, start:end]


class _RootConditionedDiffusion(
    _PositionStreamBufferMixin, DiffForcingWanModel
):
    """Diffusion whose root channels are pure conditioning.



    The root is never noised, never predicted and never scored: the diffused

    state holds only the non-root channels, the root is concatenated in front

    of the denoiser input at every training window and every denoise step, and

    the emitted output is ``[supplied root, network state]``.  Root output is

    therefore bit-exact with the supplied root.



    Subclasses define how a batch's raw root maps onto one condition vector per

    state unit (``condition_dim``): the direct model uses one root vector per

    frame, while the latent model packs four root vectors per latent token.

    """

    condition_dim = 0

    def _split_statistics(self, mean_path, std_path, expected_width):
        """Split one statistics pair into (condition, state) halves.



        Both halves come from the same file, so the condition is normalized

        with per-dimension statistics of the exact values it carries.

        """
        if (mean_path is None) != (std_path is None):
            raise ValueError("mean_path and std_path must be provided together")
        mean, std = load_statistics(mean_path, std_path, expected_width)
        mean, std = mean.reshape(-1), std.reshape(-1)
        if mean.numel() != expected_width or std.numel() != expected_width:
            raise ValueError(
                f"statistics must have {expected_width} values, got "
                f"{mean.numel()}/{std.numel()}"
            )
        std = guard_std(std)
        return (
            mean[: self.condition_dim].clone(),
            std[: self.condition_dim].clone(),
            mean[self.condition_dim :].clone(),
            std[self.condition_dim :].clone(),
        )

    def _register_statistics(self, mean_path, std_path, expected_width):
        cond_mean, cond_std, state_mean, state_std = self._split_statistics(
            mean_path, std_path, expected_width
        )
        self.mean = state_mean
        self.std = state_std
        self.register_buffer("position_mean", cond_mean)
        self.register_buffer("position_std", cond_std)

    def _assert_stream_schedule(self):
        """Reject schedules that commit frames without ever denoising them.



        ``get_committable`` yields zero committable steps when

        ``steps < chunk_size``, so the loop body never runs while frames are

        still emitted.

        """
        steps = int(self.time_scheduler.steps)
        chunk_size = int(self.time_scheduler.chunk_size)
        if steps < chunk_size or steps % chunk_size != 0:
            raise ValueError(
                "streaming requires steps to be a positive multiple of "
                f"chunk_size, got steps={steps}, chunk_size={chunk_size}"
            )

    # ------------------------------------------------------------------
    # Condition extraction (subclass hooks)
    # ------------------------------------------------------------------
    def _split_feature(self, feature):
        """Return (condition_part, state_part) of a full-width feature."""
        raise NotImplementedError

    def _pack_position_batch(self, position, valid_len, seq_len, device):
        """Return (B, seq_len, condition_dim) raw conditions."""
        raise NotImplementedError

    def _extract_inputs(self, x):
        inputs = super()._extract_inputs(x)
        # Position-only and streaming callers commonly use the literal key
        # without adding it to input_keys.  Keep custom mappings authoritative,
        # then fall back to the public ``position`` key.
        if "position" not in inputs:
            position_key = self.input_keys.get("position", "position")
            if position_key in x:
                inputs["position"] = x[position_key]
            elif "position" in x:
                inputs["position"] = x["position"]
        feature = inputs.get("feature")
        if feature is not None:
            width = feature.size(-1)
            if width == self.condition_dim + self.input_dim:
                condition, state = self._split_feature(feature)
                inputs["feature"] = state
                inputs.setdefault("position", condition)
            elif width != self.input_dim:
                raise ValueError(
                    f"feature must have {self.condition_dim + self.input_dim} "
                    f"(root+state) or {self.input_dim} (state) channels, got {width}"
                )
        return inputs

    def _get_raw_position(self, inputs):
        position = inputs.get("position")
        if position is None:
            raise KeyError(
                "root conditioning is required: supply input_keys.position or a "
                "full-width feature whose leading channels carry the root"
            )
        return position

    # ------------------------------------------------------------------
    # Training
    # ------------------------------------------------------------------
    def _prepare_training_condition(

        self, inputs, normalized_feature, valid_len, device

    ):
        packed = self._pack_position_batch(
            self._get_raw_position(inputs),
            valid_len,
            normalized_feature.size(1),
            device,
        )
        normalized = (packed - self.position_mean) / self.position_std
        return [
            normalized[i].transpose(0, 1)[:, :, None, None]
            for i in range(normalized.size(0))
        ]

    def _condition_training_input(

        self, noisy, clean, condition, batch_index, start, end

    ):
        root = condition[batch_index][:, start:end]
        return torch.cat([root.to(noisy), noisy], dim=0)

    # ------------------------------------------------------------------
    # Offline generation
    # ------------------------------------------------------------------
    def _prepare_generation_condition(self, inputs, valid_len, seq_len, device):
        packed = self._pack_position_batch(
            self._get_raw_position(inputs), valid_len, seq_len, device
        )
        normalized = (packed - self.position_mean) / self.position_std
        return {
            "raw": packed,
            "normalized": [
                normalized[i].transpose(0, 1)[:, :, None, None]
                for i in range(normalized.size(0))
            ],
        }

    def _condition_generation_input(self, noisy, condition, starts, ends):
        return [
            torch.cat(
                [
                    condition["normalized"][i][:, starts[i] : ends[i]].to(noisy[i]),
                    noisy[i],
                ],
                dim=0,
            )
            for i in range(len(noisy))
        ]

    def _finalize_generation_output(self, generated, condition, valid_len):
        return [
            torch.cat(
                [condition["raw"][i, : valid_len[i]].to(generated[i]), generated[i]],
                dim=-1,
            )
            for i in range(len(generated))
        ]

    def _generation_output_metadata(self, condition, valid_len):
        return {
            "position": [
                condition["raw"][i, : valid_len[i]].clone()
                for i in range(len(valid_len))
            ]
        }

    # ------------------------------------------------------------------
    # Streaming
    # ------------------------------------------------------------------
    def _condition_stream_input(self, noisy, starts, ends):
        result = []
        for i, item in enumerate(noisy):
            raw = self._stream_position_slice(starts[i], ends[i])[i]
            norm = (raw - self.position_mean) / self.position_std
            root = norm.transpose(0, 1)[:, :, None, None].to(item)
            result.append(torch.cat([root, item], dim=0))
        return result

    def _finalize_stream_output(self, generated, start, end):
        raw = self._stream_position_slice(start, end)
        return [
            torch.cat([raw[i].to(generated[i]), generated[i]], dim=-1)
            for i in range(len(generated))
        ]

    def _stream_output_metadata(self, start, end):
        # Clone: the slice is a view of the rollback-mutated ring buffer.
        raw = self._stream_position_slice(start, end)
        return {"position": [raw[i].clone() for i in range(self.batch_size)]}


class DiffForcingPositionWanModel(_RootConditionedDiffusion):
    """Direct full-feature diffusion conditioned on clean root motion.



    ``input_dim`` keeps its config meaning (the complete feature width); the

    diffused state contains every non-root channel.  The denoiser reads the

    clean root plus noisy pose state and predicts only the pose state, so the

    loss never covers root motion and emitted root channels are exactly the

    supplied values.

    """

    def __init__(

        self,

        *args,

        input_dim=138,

        representation=None,

        root_dim=None,

        mean_path=None,

        std_path=None,

        root_mean_path=None,

        root_std_path=None,

        **kwargs,

    ):
        self.feature_dim = int(input_dim)
        self.representation, self.root_dim = resolve_position_representation(
            representation,
            feature_dim=self.feature_dim,
            root_dim=root_dim,
        )
        self.condition_dim = self.root_dim
        if not 0 < self.root_dim < self.feature_dim:
            raise ValueError(
                f"root_dim must be in (0,{self.feature_dim}), got {self.root_dim}"
            )
        kwargs.pop("model_input_dim", None)
        super().__init__(
            *args,
            input_dim=self.feature_dim - self.root_dim,
            model_input_dim=self.feature_dim,
            mean_path=None,
            std_path=None,
            **kwargs,
        )
        self._assert_stream_schedule()
        self.stream_condition_dim = self.root_dim
        self._register_statistics(mean_path, std_path, self.feature_dim)
        if root_mean_path is not None or root_std_path is not None:
            # Optional per-dimension root statistics can override the leading
            # values from the full-feature statistics.
            if root_mean_path is None or root_std_path is None:
                raise ValueError(
                    "root_mean_path and root_std_path must be provided together"
                )
            root_mean, root_std = load_statistics(
                root_mean_path, root_std_path, self.root_dim
            )
            root_mean, root_std = root_mean.reshape(-1), root_std.reshape(-1)
            if root_mean.numel() < self.root_dim or root_std.numel() < self.root_dim:
                raise ValueError(
                    f"root statistics must have at least {self.root_dim} values"
                )
            self.position_mean = root_mean[: self.root_dim].clone()
            self.position_std = guard_std(root_std[: self.root_dim]).clone()

    def _split_feature(self, feature):
        return feature[..., : self.root_dim], feature[..., self.root_dim :]

    def _pack_position_batch(self, position, valid_len, seq_len, device):
        position = _to_batched_tensor(position, device)
        if position.ndim == 2:
            position = position.unsqueeze(0)
        if position.ndim != 3 or position.size(-1) < self.root_dim:
            raise ValueError(
                f"direct root must be (B,T,>={self.root_dim}), got "
                f"{tuple(position.shape)}"
            )
        if position.size(0) != len(valid_len):
            raise ValueError(
                f"root batch={position.size(0)}, expected {len(valid_len)}"
            )
        raw = position[..., : self.root_dim].float()
        padded = []
        for i, length in enumerate(valid_len):
            if raw.size(1) < length:
                raise ValueError(
                    f"sample {i} has {raw.size(1)} root frames but needs {length}"
                )
            item = raw[i, :length]
            if seq_len > length:
                if length == 0:
                    raise ValueError("cannot pad an empty valid root sequence")
                item = torch.cat([item, item[-1:].expand(seq_len - length, -1)], dim=0)
            padded.append(item)
        packed = torch.stack(padded)
        if not torch.isfinite(packed).all():
            raise ValueError("root condition contains NaN or Inf")
        return packed

    def _position_from_history(self, history_latent):
        tensors = [torch.as_tensor(v) for v in history_latent]
        if not tensors:
            return None
        width = int(tensors[0].size(-1))
        if width == self.feature_dim:
            return torch.stack([v[..., : self.root_dim] for v in tensors])
        return None

    def init_generated(

        self,

        seq_len,

        batch_size=1,

        schedule_config={},

        history_latent=None,

        history_position=None,

    ):
        if history_latent is not None:
            # A full-width history carries its own root, which must be pulled
            # out here: the diffused state has no root channels for the base
            # class to recover it from.
            split_state = []
            inferred_position = []
            for item in history_latent:
                item = torch.as_tensor(item)
                width = int(item.size(-1))
                if width == self.feature_dim:
                    inferred_position.append(item[..., : self.root_dim])
                    split_state.append(item[..., self.root_dim :])
                elif width == self.input_dim:
                    split_state.append(item)
                else:
                    raise ValueError(
                        f"history must be {self.feature_dim}D (root+state) or "
                        f"{self.input_dim}D (state), got {width}D"
                    )
            if inferred_position:
                if len(inferred_position) != len(split_state):
                    raise ValueError("cannot mix full-width and state-only history")
                if history_position is not None:
                    raise ValueError(
                        "history_position must be omitted for a full-width history"
                    )
                history_position = inferred_position
            elif history_position is None:
                raise ValueError("a state-only history requires history_position")
            history_latent = split_state
        return super().init_generated(
            seq_len,
            batch_size=batch_size,
            schedule_config=schedule_config,
            history_latent=history_latent,
            history_position=history_position,
        )

    def _parse_stream_position(self, inputs, device):
        position = _to_batched_tensor(self._get_raw_position(inputs), device)
        if position.ndim == 3 and position.size(1) == 1:
            position = position[:, 0]
        if position.ndim == 1:
            position = position.unsqueeze(0)
        if position.ndim != 2 or position.size(-1) < self.root_dim:
            raise ValueError(
                f"each direct stream step needs (B,{self.root_dim}) root, "
                f"got {tuple(position.shape)}"
            )
        return position[..., : self.root_dim]


class LatentDiffForcingPositionWanModel(_RootConditionedDiffusion):
    """Latent diffusion conditioned on the WAN-causal packed root.



    ``input_dim`` keeps its config meaning (the VAE latent width); the diffused

    state is exactly that latent.  The denoiser reads ``position_dim + latent``

    channels and predicts only the latent, so the emitted token is

    ``[supplied packed root, network latent]`` and its root is bit-exact.



    The packed root width is inferred, in order, from ``representation``, an

    explicit ``root_dim``/``position_dim``, or the token statistics width.  The

    last path keeps existing configs working when they do not forward the

    top-level representation into model parameters.

    """

    def __init__(

        self,

        *args,

        input_dim=16,

        representation=None,

        position_dim=None,

        root_dim=None,

        root_group_size=WAN_ROOT_GROUP_SIZE,

        mean_path=None,

        std_path=None,

        **kwargs,

    ):
        self.root_group_size = int(root_group_size)
        if self.root_group_size != WAN_ROOT_GROUP_SIZE:
            raise ValueError(
                "the WAN position ABI groups exactly "
                f"{WAN_ROOT_GROUP_SIZE} root frames per latent, got "
                f"root_group_size={self.root_group_size}"
            )
        self.latent_dim = int(input_dim)
        if self.latent_dim <= 0:
            raise ValueError(f"input_dim must be positive, got {self.latent_dim}")

        resolved_root_dim = None if root_dim is None else int(root_dim)
        self.representation = None
        if representation is not None:
            self.representation, resolved_root_dim = (
                resolve_position_representation(
                    representation,
                    root_dim=resolved_root_dim,
                )
            )

        resolved_position_dim = (
            None if position_dim is None else int(position_dim)
        )
        if resolved_position_dim is not None:
            if resolved_position_dim <= 0:
                raise ValueError(
                    f"position_dim must be positive, got {resolved_position_dim}"
                )
            if resolved_position_dim % self.root_group_size != 0:
                raise ValueError(
                    "position_dim must be divisible by root_group_size "
                    f"({self.root_group_size}), got {resolved_position_dim}"
                )
            position_root_dim = resolved_position_dim // self.root_group_size
            if (
                resolved_root_dim is not None
                and resolved_root_dim != position_root_dim
            ):
                raise ValueError(
                    "position_dim and root_dim disagree: "
                    f"{resolved_position_dim} != {resolved_root_dim} * "
                    f"{self.root_group_size}"
                )
            resolved_root_dim = position_root_dim
        elif resolved_root_dim is not None:
            resolved_position_dim = resolved_root_dim * self.root_group_size
        elif (
            mean_path is not None and std_path is not None
            and Path(mean_path).is_file() and Path(std_path).is_file()
        ):
            mean_width = int(np.asarray(np.load(mean_path)).size)
            std_width = int(np.asarray(np.load(std_path)).size)
            if mean_width != std_width:
                raise ValueError(
                    "token statistics widths must match, got "
                    f"{mean_width}/{std_width}"
                )
            resolved_position_dim = mean_width - self.latent_dim
            if (
                resolved_position_dim <= 0
                or resolved_position_dim % self.root_group_size != 0
            ):
                raise ValueError(
                    f"cannot infer packed root width from {mean_width}D token "
                    f"statistics, latent_dim={self.latent_dim}, and "
                    f"root_group_size={self.root_group_size}"
                )
            resolved_root_dim = resolved_position_dim // self.root_group_size
            if resolved_root_dim not in set(POSITION_ROOT_DIMS.values()):
                raise ValueError(
                    f"cannot infer a supported root width from {mean_width}D "
                    f"token statistics: inferred root_dim={resolved_root_dim}, "
                    f"supported root widths are "
                    f"{sorted(set(POSITION_ROOT_DIMS.values()))}; pass "
                    "root_dim/position_dim explicitly for a custom schema"
                )
        else:
            # Historical root3/root12 default for configs without token stats.
            _, resolved_root_dim = resolve_position_representation()
            resolved_position_dim = resolved_root_dim * self.root_group_size

        # If a representation supplied the root width, also validate any width
        # inferred from position_dim or token statistics.
        self.representation, resolved_root_dim = resolve_position_representation(
            self.representation,
            root_dim=resolved_root_dim,
        )
        self.root_dim = resolved_root_dim
        self.position_dim = resolved_position_dim
        self.condition_dim = self.position_dim
        if self.position_dim != self.root_dim * self.root_group_size:
            raise ValueError(
                "position_dim must equal root_dim * root_group_size "
                f"({self.root_dim}*{self.root_group_size})"
            )
        kwargs.pop("model_input_dim", None)
        super().__init__(
            *args,
            input_dim=self.latent_dim,
            model_input_dim=self.position_dim + self.latent_dim,
            mean_path=None,
            std_path=None,
            **kwargs,
        )
        self._assert_stream_schedule()
        self.stream_condition_dim = self.position_dim
        self._register_statistics(
            mean_path, std_path, self.position_dim + self.latent_dim
        )

    def _split_feature(self, feature):
        return feature[..., : self.position_dim], feature[..., self.position_dim :]

    def _pack_position_batch(self, position, valid_len, seq_len, device):
        position = _to_batched_tensor(position, device)
        if position.ndim == 2:
            position = position.unsqueeze(0)
        if position.ndim != 3:
            raise ValueError(
                f"latent root must be (B,L,{self.position_dim}) or (B,T,>="
                f"{self.root_dim}), got {tuple(position.shape)}"
            )
        if position.size(0) != len(valid_len):
            raise ValueError(
                f"root batch={position.size(0)}, expected {len(valid_len)}"
            )
        packed_items = []
        for i, length in enumerate(valid_len):
            item = position[i : i + 1]
            if item.size(-1) == self.position_dim:
                if item.size(1) < length:
                    raise ValueError(
                        f"sample {i} has {item.size(1)} packed roots but needs {length}"
                    )
                packed = item[:, :length]
            else:
                required_frames = (
                    0 if length == 0 else 1 + self.root_group_size * (length - 1)
                )
                if item.size(-1) < self.root_dim or item.size(1) < required_frames:
                    raise ValueError(
                        f"sample {i} needs {required_frames} raw root frames for "
                        f"{length} latents, got {item.size(1)}"
                    )
                # Trim to this sample's own frames: in a collated batch the
                # tensor is padded to the batch maximum, and pack_root_condition
                # clamps its taps to the frame count it is given, so an untrimmed
                # item would let a short sample's tail taps land on pad zeros.
                item = item[:, :required_frames]
                packed = pack_root_condition(
                    item[..., : self.root_dim],
                    latent_length=length,
                    root_dim=self.root_dim,
                    group_size=self.root_group_size,
                    first_chunk=True,
                )
            packed = packed[0]
            if seq_len > length:
                if length == 0:
                    raise ValueError("cannot pad an empty valid root sequence")
                packed = torch.cat(
                    [packed, packed[-1:].expand(seq_len - length, -1)], dim=0
                )
            packed_items.append(packed)
        packed = torch.stack(packed_items).float()
        if not torch.isfinite(packed).all():
            raise ValueError("latent root condition contains NaN or Inf")
        return packed

    def _position_from_history(self, history_latent):
        tensors = [torch.as_tensor(v) for v in history_latent]
        if not tensors:
            return None
        width = int(tensors[0].size(-1))
        if width == self.position_dim + self.latent_dim:
            return torch.stack([v[..., : self.position_dim] for v in tensors])
        return None

    def init_generated(

        self,

        seq_len,

        batch_size=1,

        schedule_config={},

        history_latent=None,

        history_position=None,

    ):
        if history_latent is not None:
            token_dim = self.position_dim + self.latent_dim
            split_state = []
            inferred_position = []
            for item in history_latent:
                item = torch.as_tensor(item)
                width = int(item.size(-1))
                if width == token_dim:
                    inferred_position.append(item[..., : self.position_dim])
                    split_state.append(item[..., self.position_dim :])
                elif width == self.latent_dim:
                    split_state.append(item)
                else:
                    raise ValueError(
                        f"latent history must be {token_dim}D (root+latent) or "
                        f"{self.latent_dim}D (latent), got {width}D"
                    )
            if inferred_position:
                if len(inferred_position) != len(split_state):
                    raise ValueError("cannot mix token-width and latent-only history")
                if history_position is not None:
                    raise ValueError(
                        "history_position must be omitted for a token-width history"
                    )
                history_position = inferred_position
            elif history_position is None:
                raise ValueError("a latent-only history requires history_position")
            history_latent = split_state
        return super().init_generated(
            seq_len,
            batch_size=batch_size,
            schedule_config=schedule_config,
            history_latent=history_latent,
            history_position=history_position,
        )

    def _parse_stream_position(self, inputs, device):
        position = _to_batched_tensor(self._get_raw_position(inputs), device)
        if position.ndim == 3 and position.shape[-2:] == (
            self.root_group_size,
            self.root_dim,
        ):
            position = position.flatten(-2)
        elif position.ndim == 3 and position.size(1) == 1:
            position = position[:, 0]
        if position.ndim == 1:
            position = position.unsqueeze(0)
        if position.ndim != 2 or position.size(-1) != self.position_dim:
            raise ValueError(
                f"each latent stream step needs (B,{self.position_dim}) or "
                f"(B,{self.root_group_size},{self.root_dim}) root, got "
                f"{tuple(position.shape)}"
            )
        return position


# Concise aliases for configs and downstream imports.
DiffForcingWanPositionModel = DiffForcingPositionWanModel
LatentDiffForcingWanPositionModel = LatentDiffForcingPositionWanModel