model/predictor_v4.py

"""Wan Predictor-v4 for skipped Self-Forcing denoising steps.

The Predictor is initialized from an already-loaded ``CausalWanModel``.  It
keeps the Teacher patch/time/head modules frozen, trains two copied causal Wan
blocks, and predicts a residual over the same-chunk anchor hidden state.

Two history paths are supported:

* online F-P-P-F inference can pass the generator's existing KV caches;
* offline training can rebuild selected-layer history KV from clean-pass
  self-attention prefeatures with :meth:`build_history_kv_cache`.

The ordinary ``state_dict`` API is intentionally unchanged.  Use
``trainable_state_dict``/``checkpoint_dict`` for compact Predictor checkpoints
that omit frozen Teacher weights.
"""

from __future__ import annotations

import copy
import math
from collections import OrderedDict
from collections.abc import Mapping, Sequence
from dataclasses import asdict, dataclass
from typing import Any

import torch
from torch import nn

from wan.modules.causal_model import (
    CausalWanAttentionBlock,
    CausalWanModel,
    causal_rope_apply,
)
from wan.modules.model import sinusoidal_embedding_1d


@dataclass(frozen=True)
class WanPredictorV4Config:
    """Serializable architecture metadata derived from the loaded Teacher."""

    format_version: int
    model_type: str
    patch_size: tuple[int, int, int]
    in_dim: int
    dim: int
    ffn_dim: int
    freq_dim: int
    out_dim: int
    num_heads: int
    num_layers: int
    local_attn_size: int
    sink_size: int
    qk_norm: bool
    cross_attn_norm: bool
    eps: float
    source_block_ids: tuple[int, int]
    spatial_grid: tuple[int, int]

    @property
    def tokens_per_frame(self) -> int:
        return math.prod(self.spatial_grid)

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


class TripleFeatureFusion(nn.Module):
    """Fuse target-latent, same-chunk anchor, and previous-chunk features."""

    def __init__(self, dim: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.current_norm = nn.LayerNorm(dim, eps=eps)
        self.anchor_norm = nn.LayerNorm(dim, eps=eps)
        self.previous_norm = nn.LayerNorm(dim, eps=eps)
        self.mlp = nn.Sequential(
            nn.Linear(3 * dim, 2 * dim),
            nn.SiLU(),
            nn.Linear(2 * dim, dim),
        )

    def forward(
        self,
        current: torch.Tensor,
        anchor: torch.Tensor,
        previous: torch.Tensor,
    ) -> torch.Tensor:
        if current.shape != anchor.shape or current.shape != previous.shape:
            raise ValueError(
                "TripleFeatureFusion requires identical [B, L, D] shapes, got "
                f"current={tuple(current.shape)}, anchor={tuple(anchor.shape)}, "
                f"previous={tuple(previous.shape)}"
            )
        return self.mlp(
            torch.cat(
                (
                    self.current_norm(current),
                    self.anchor_norm(anchor),
                    self.previous_norm(previous),
                ),
                dim=-1,
            )
        )


class _FrozenHistoryProjector(nn.Module):
    """Frozen copy of one Teacher self-attention K/V projection path."""

    def __init__(self, teacher_block: CausalWanAttentionBlock) -> None:
        super().__init__()
        self.k = copy.deepcopy(teacher_block.self_attn.k)
        self.v = copy.deepcopy(teacher_block.self_attn.v)
        self.norm_k = copy.deepcopy(teacher_block.self_attn.norm_k)
        self.requires_grad_(False)

    def forward(
        self,
        self_attn_input: torch.Tensor,
        *,
        num_heads: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch, tokens, dim = self_attn_input.shape
        if dim % num_heads:
            raise ValueError(f"Hidden dim {dim} is not divisible by {num_heads} heads")
        head_dim = dim // num_heads
        key = self.norm_k(self.k(self_attn_input)).view(
            batch, tokens, num_heads, head_dim
        )
        value = self.v(self_attn_input).view(
            batch, tokens, num_heads, head_dim
        )
        return key, value


class SelfForcingPredictorV4(nn.Module):
    """Two-block Wan Predictor used for the middle denoising steps of F-P-P-F."""

    requires_previous_chunk_hidden = True
    uses_history_kv = True
    uses_clean_prefeature = True
    checkpoint_format_version = 1

    def __init__(
        self,
        teacher_model: CausalWanModel,
        *,
        source_block_ids: tuple[int, int] = (1, 28),
        spatial_grid: tuple[int, int] = (30, 52),
    ) -> None:
        super().__init__()
        teacher_model = self._unwrap_teacher(teacher_model)
        if teacher_model.model_type != "t2v":
            raise NotImplementedError("SelfForcingPredictorV4 currently supports Wan T2V")
        if len(source_block_ids) != 2 or len(set(source_block_ids)) != 2:
            raise ValueError("Predictor-v4 requires exactly two distinct source blocks")
        if any(index < 0 or index >= len(teacher_model.blocks) for index in source_block_ids):
            raise ValueError(
                f"Invalid source blocks {source_block_ids} for "
                f"{len(teacher_model.blocks)} Teacher blocks"
            )
        if len(spatial_grid) != 2 or any(int(size) <= 0 for size in spatial_grid):
            raise ValueError(f"Invalid Predictor token spatial grid: {spatial_grid}")

        first_self_attn = teacher_model.blocks[0].self_attn
        self.predictor_config = WanPredictorV4Config(
            format_version=self.checkpoint_format_version,
            model_type=str(teacher_model.model_type),
            patch_size=tuple(int(item) for item in teacher_model.patch_size),
            in_dim=int(teacher_model.in_dim),
            dim=int(teacher_model.dim),
            ffn_dim=int(teacher_model.ffn_dim),
            freq_dim=int(teacher_model.freq_dim),
            out_dim=int(teacher_model.out_dim),
            num_heads=int(teacher_model.num_heads),
            num_layers=len(teacher_model.blocks),
            local_attn_size=int(teacher_model.local_attn_size),
            sink_size=int(first_self_attn.sink_size),
            qk_norm=bool(teacher_model.qk_norm),
            cross_attn_norm=bool(teacher_model.cross_attn_norm),
            eps=float(teacher_model.eps),
            source_block_ids=tuple(int(item) for item in source_block_ids),
            spatial_grid=tuple(int(item) for item in spatial_grid),
        )
        cfg = self.predictor_config

        # Frozen modules are copied rather than referenced so that calling
        # Predictor.train()/to() cannot alter the loaded generator.
        self.patch_embedding = copy.deepcopy(teacher_model.patch_embedding)
        self.time_embedding = copy.deepcopy(teacher_model.time_embedding)
        self.time_projection = copy.deepcopy(teacher_model.time_projection)
        self.head = copy.deepcopy(teacher_model.head)
        self.predictor_blocks = nn.ModuleList(
            [copy.deepcopy(teacher_model.blocks[index]) for index in source_block_ids]
        )
        self.history_projectors = nn.ModuleDict(
            {
                str(index): _FrozenHistoryProjector(teacher_model.blocks[index])
                for index in source_block_ids
            }
        )

        self.feature_fusion = TripleFeatureFusion(cfg.dim, cfg.eps)
        self.residual_out = nn.Linear(cfg.dim, cfg.dim)

        self._freeze_teacher_modules()
        self.predictor_blocks.requires_grad_(True)
        # Text K/V come from the Teacher cross-attention cache.  Predictor only
        # executes the query/output side, so keep unused cache-building weights
        # frozen and out of the optimizer/checkpoint.
        for block in self.predictor_blocks:
            block.cross_attn.k.requires_grad_(False)
            block.cross_attn.v.requires_grad_(False)
            block.cross_attn.norm_k.requires_grad_(False)

        reference = teacher_model.patch_embedding.weight
        self.feature_fusion.to(device=reference.device, dtype=reference.dtype)
        self.residual_out.to(device=reference.device, dtype=reference.dtype)
        nn.init.zeros_(self.residual_out.weight)
        nn.init.zeros_(self.residual_out.bias)

        # Match CausalWanModel: RoPE frequencies are runtime state rather than a
        # persistent buffer, so compact checkpoints contain no derived table.
        self._freqs = teacher_model.freqs.detach().clone()

    @staticmethod
    def _unwrap_teacher(model: Any) -> CausalWanModel:
        current = model
        visited: set[int] = set()
        while id(current) not in visited:
            visited.add(id(current))
            if isinstance(current, CausalWanModel):
                return current
            wrapped = getattr(current, "module", None)
            if wrapped is not None:
                current = wrapped
                continue
            nested = getattr(current, "model", None)
            if nested is not None:
                current = nested
                continue
            break
        raise TypeError(
            "teacher_model must be CausalWanModel (normally generator.model), "
            f"got {type(model)!r}"
        )

    @classmethod
    def from_teacher(
        cls,
        teacher_model: CausalWanModel,
        *,
        source_block_ids: tuple[int, int] = (1, 28),
        spatial_grid: tuple[int, int] = (30, 52),
    ) -> "SelfForcingPredictorV4":
        """Initialize all copied/frozen/trainable weights from a loaded Teacher."""

        return cls(
            teacher_model,
            source_block_ids=source_block_ids,
            spatial_grid=spatial_grid,
        )

    @property
    def config_dict(self) -> dict[str, Any]:
        return self.predictor_config.to_dict()

    @property
    def source_block_ids(self) -> tuple[int, int]:
        return self.predictor_config.source_block_ids

    def _freeze_teacher_modules(self) -> None:
        for module in (
            self.patch_embedding,
            self.time_embedding,
            self.time_projection,
            self.head,
            self.history_projectors,
        ):
            module.requires_grad_(False)
            module.eval()

    @torch.no_grad()
    def sync_frozen_from_teacher(
        self,
        teacher_model: CausalWanModel,
    ) -> None:
        """Refresh only the frozen Teacher-derived Predictor parameters.

        Joint DMD changes the Full Generator after Predictor construction.  A
        compact Predictor checkpoint is reconstructed from that updated Full
        model at inference time, so the frozen training-time copies must track
        it as well.  Predictor-owned trainable blocks/fusion are never
        overwritten here.
        """

        teacher_model = self._unwrap_teacher(teacher_model)
        for destination, source in (
            (self.patch_embedding, teacher_model.patch_embedding),
            (self.time_embedding, teacher_model.time_embedding),
            (self.time_projection, teacher_model.time_projection),
            (self.head, teacher_model.head),
        ):
            destination.load_state_dict(source.state_dict(), strict=True)

        for position, source_id in enumerate(self.source_block_ids):
            teacher_block = teacher_model.blocks[source_id]
            while hasattr(teacher_block, "module"):
                teacher_block = teacher_block.module
            history = self.history_projectors[str(source_id)]
            history.k.load_state_dict(
                teacher_block.self_attn.k.state_dict(), strict=True
            )
            history.v.load_state_dict(
                teacher_block.self_attn.v.state_dict(), strict=True
            )
            history.norm_k.load_state_dict(
                teacher_block.self_attn.norm_k.state_dict(), strict=True
            )

            predictor_cross = self.predictor_blocks[position].cross_attn
            teacher_cross = teacher_block.cross_attn
            predictor_cross.k.load_state_dict(
                teacher_cross.k.state_dict(), strict=True
            )
            predictor_cross.v.load_state_dict(
                teacher_cross.v.state_dict(), strict=True
            )
            predictor_cross.norm_k.load_state_dict(
                teacher_cross.norm_k.state_dict(), strict=True
            )

        self._freeze_teacher_modules()
        for block in self.predictor_blocks:
            block.cross_attn.k.requires_grad_(False)
            block.cross_attn.v.requires_grad_(False)
            block.cross_attn.norm_k.requires_grad_(False)

    def train(self, mode: bool = True) -> "SelfForcingPredictorV4":
        super().train(mode)
        # Frozen layers have no stochastic operations today, but pinning their
        # mode makes the intended boundary robust to future Wan changes.
        for module in (
            self.patch_embedding,
            self.time_embedding,
            self.time_projection,
            self.head,
            self.history_projectors,
        ):
            module.eval()
        return self

    def _runtime_freqs(self, device: torch.device) -> torch.Tensor:
        if self._freqs.device != device:
            self._freqs = self._freqs.to(device)
        return self._freqs

    @staticmethod
    def _scalar_int(value: int | torch.Tensor, name: str) -> int:
        if torch.is_tensor(value):
            if value.numel() != 1:
                raise ValueError(f"{name} must be scalar, got shape {tuple(value.shape)}")
            value = value.detach().item()
        result = int(value)
        if result < 0:
            raise ValueError(f"{name} must be non-negative, got {result}")
        return result

    @staticmethod
    def _start_values(
        value: int | torch.Tensor,
        *,
        batch: int,
        name: str,
    ) -> list[int]:
        if torch.is_tensor(value):
            values = [int(item) for item in value.detach().reshape(-1).cpu().tolist()]
        else:
            values = [int(value)]
        if len(values) == 1:
            values *= batch
        if len(values) != batch:
            raise ValueError(f"{name} has {len(values)} values for batch {batch}")
        if any(item < 0 for item in values):
            raise ValueError(f"{name} must contain non-negative frame indices")
        return values

    def _rope_history_key(
        self,
        key: torch.Tensor,
        *,
        start_frames: int | torch.Tensor,
    ) -> torch.Tensor:
        cfg = self.predictor_config
        batch, tokens = key.shape[:2]
        if tokens % cfg.tokens_per_frame:
            raise ValueError(
                f"History tokens {tokens} are not divisible by "
                f"{cfg.tokens_per_frame} tokens/frame"
            )
        frames = tokens // cfg.tokens_per_frame
        starts = self._start_values(start_frames, batch=batch, name="start_frames")
        freqs = self._runtime_freqs(key.device)
        grid = torch.tensor(
            [[frames, *cfg.spatial_grid]],
            dtype=torch.long,
            device=key.device,
        )
        if len(set(starts)) == 1:
            return causal_rope_apply(
                key,
                grid.expand(batch, -1),
                freqs,
                start_frame=starts[0],
            )
        return torch.cat(
            [
                causal_rope_apply(
                    key[index : index + 1],
                    grid,
                    freqs,
                    start_frame=start,
                )
                for index, start in enumerate(starts)
            ],
            dim=0,
        )

    def _project_history_part(
        self,
        block_id: int,
        prefeature: torch.Tensor,
        *,
        start_frames: int | torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        cfg = self.predictor_config
        if prefeature.ndim != 3 or prefeature.shape[-1] != cfg.dim:
            raise ValueError(
                f"Block {block_id} prefeature must be [B, S, {cfg.dim}], got "
                f"{tuple(prefeature.shape)}"
            )
        projector = self.history_projectors[str(block_id)]
        projector_device = projector.k.weight.device
        if prefeature.device != projector_device:
            raise ValueError(
                f"Block {block_id} prefeature is on {prefeature.device}, "
                f"projector is on {projector_device}"
            )
        prefeature = prefeature.to(dtype=projector.k.weight.dtype)
        # Clean-history tensors are fixed offline Teacher data.  Avoid retaining
        # a graph through several GiB of reconstructed cache.
        with torch.no_grad():
            key, value = projector(prefeature, num_heads=cfg.num_heads)
            key = self._rope_history_key(key, start_frames=start_frames)
        return key, value

    def build_history_kv_cache(
        self,
        clean_prefeature_by_block: Mapping[
            int | str, torch.Tensor | Sequence[torch.Tensor]
        ],
        *,
        current_start: int | torch.Tensor,
        current_tokens: int | torch.Tensor,
        start_frames: int | torch.Tensor | Sequence[int | torch.Tensor] = 0,
        cache_capacity: int | None = None,
    ) -> dict[int, dict[str, torch.Tensor]]:
        """Rebuild selected-layer clean-history caches for Predictor training.

        ``clean_prefeature_by_block`` may contain one already-concatenated
        ``[B, S, D]`` tensor per block, or a sequence of chunk tensors.  For a
        sequence, ``start_frames`` can be the matching sequence ``0, 3, ...``.
        ``current_start`` is the global token offset used by Wan inference and
        ``current_tokens`` reserves the writable current-chunk cache region.
        """

        current_start_int = self._scalar_int(current_start, "current_start")
        current_tokens_int = self._scalar_int(current_tokens, "current_tokens")
        if current_tokens_int == 0:
            raise ValueError("current_tokens must be positive")

        cfg = self.predictor_config
        result: dict[int, dict[str, torch.Tensor]] = {}
        expected_batch: int | None = None
        expected_history_tokens: int | None = None
        for block_id in self.source_block_ids:
            value = clean_prefeature_by_block.get(block_id)
            if value is None:
                value = clean_prefeature_by_block.get(str(block_id))
            if value is None:
                raise ValueError(f"Missing clean prefeature for block {block_id}")

            if torch.is_tensor(value):
                if isinstance(start_frames, Sequence) and not torch.is_tensor(start_frames):
                    if len(start_frames) != 1:
                        raise ValueError(
                            "Already-concatenated prefeatures require one start_frames value"
                        )
                    part_start = start_frames[0]
                else:
                    part_start = start_frames
                keys, values = self._project_history_part(
                    block_id, value, start_frames=part_start
                )
            else:
                parts = list(value)
                if not parts:
                    raise ValueError(f"Block {block_id} has no history prefeatures")
                if isinstance(start_frames, Sequence) and not torch.is_tensor(start_frames):
                    starts = list(start_frames)
                    if len(starts) != len(parts):
                        raise ValueError(
                            f"start_frames has {len(starts)} entries for "
                            f"{len(parts)} history chunks"
                        )
                else:
                    starts = []
                    next_start: int | torch.Tensor = start_frames
                    for part in parts:
                        starts.append(next_start)
                        if torch.is_tensor(next_start) and next_start.numel() > 1:
                            next_start = next_start + (
                                part.shape[1] // self.predictor_config.tokens_per_frame
                            )
                        else:
                            next_start = self._scalar_int(next_start, "start_frames") + (
                                part.shape[1] // self.predictor_config.tokens_per_frame
                            )
                projected = [
                    self._project_history_part(
                        block_id, part, start_frames=part_start
                    )
                    for part, part_start in zip(parts, starts)
                ]
                keys = torch.cat([item[0] for item in projected], dim=1)
                values = torch.cat([item[1] for item in projected], dim=1)

            batch, history_tokens = keys.shape[:2]
            if expected_batch is None:
                expected_batch = batch
                expected_history_tokens = history_tokens
            elif batch != expected_batch or history_tokens != expected_history_tokens:
                raise ValueError(
                    "Selected blocks must have the same history shape, got "
                    f"block {block_id}: batch={batch}, tokens={history_tokens}; "
                    f"expected batch={expected_batch}, tokens={expected_history_tokens}"
                )
            if history_tokens > current_start_int:
                raise ValueError(
                    f"History has {history_tokens} tokens but current_start is "
                    f"{current_start_int}"
                )

            required_capacity = history_tokens + current_tokens_int
            capacity = required_capacity if cache_capacity is None else int(cache_capacity)
            if capacity < required_capacity:
                raise ValueError(
                    f"cache_capacity {capacity} is smaller than required "
                    f"{required_capacity}"
                )
            cache_k = keys.new_zeros(
                batch, capacity, cfg.num_heads, cfg.dim // cfg.num_heads
            )
            cache_v = values.new_zeros(
                batch, capacity, cfg.num_heads, cfg.dim // cfg.num_heads
            )
            cache_k[:, :history_tokens].copy_(keys)
            cache_v[:, :history_tokens].copy_(values)
            result[block_id] = {
                "k": cache_k,
                "v": cache_v,
                "global_end_index": torch.tensor(
                    [current_start_int], dtype=torch.long, device=keys.device
                ),
                "local_end_index": torch.tensor(
                    [history_tokens], dtype=torch.long, device=keys.device
                ),
            }
        return result

    def _select_cache(
        self,
        caches: Mapping[Any, Any] | Sequence[Any],
        *,
        source_id: int,
        source_position: int,
        name: str,
    ) -> Mapping[str, Any]:
        if isinstance(caches, Mapping):
            selected = caches.get(source_id)
            if selected is None:
                selected = caches.get(str(source_id))
        else:
            if len(caches) == self.predictor_config.num_layers:
                selected = caches[source_id]
            elif len(caches) == len(self.source_block_ids):
                selected = caches[source_position]
            else:
                selected = None
        if selected is None:
            raise ValueError(f"{name} is missing source block {source_id}")
        if not isinstance(selected, Mapping):
            raise TypeError(f"{name}[{source_id}] must be a mapping")
        return selected

    def _selected_crossattn_cache(
        self,
        caches: Mapping[Any, Any] | Sequence[Any],
        *,
        source_id: int,
        source_position: int,
    ) -> dict[str, Any]:
        selected = self._select_cache(
            caches,
            source_id=source_id,
            source_position=source_position,
            name="crossattn_cache",
        )
        missing = {"k", "v"}.difference(selected)
        if missing:
            raise ValueError(
                f"crossattn_cache block {source_id} is missing {sorted(missing)}"
            )
        if selected["k"].shape != selected["v"].shape:
            raise ValueError(f"crossattn_cache block {source_id} K/V shape mismatch")
        # Offline files contain K/V but need not serialize the runtime flag.
        # A shallow wrapper avoids mutating the generator-owned dictionary.
        return {**selected, "is_init": True}

    def _time_condition(
        self,
        target_timestep: torch.Tensor,
        reference: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        cfg = self.predictor_config
        time_embedding = self.time_embedding(
            sinusoidal_embedding_1d(
                cfg.freq_dim, target_timestep.flatten()
            ).type_as(reference)
        )
        block_condition = self.time_projection(time_embedding).unflatten(
            1, (6, cfg.dim)
        ).unflatten(0, target_timestep.shape)
        head_condition = time_embedding.unflatten(
            0, target_timestep.shape
        ).unsqueeze(2)
        return block_condition, head_condition

    def _unpatchify(
        self,
        tokens: torch.Tensor,
        grid_sizes: torch.Tensor,
    ) -> torch.Tensor:
        cfg = self.predictor_config
        outputs = []
        for sample, grid in zip(tokens, grid_sizes.tolist()):
            sample = sample[: math.prod(grid)].view(
                *grid, *cfg.patch_size, cfg.out_dim
            )
            # [f, h, w, p, q, r, c] -> [f, p, c, h, q, w, r] so the
            # following reshape merges each grid axis with its patch axis.
            sample = torch.einsum("fhwpqrc->fpchqwr", sample)
            outputs.append(
                sample.reshape(
                    grid[0] * cfg.patch_size[0],
                    cfg.out_dim,
                    grid[1] * cfg.patch_size[1],
                    grid[2] * cfg.patch_size[2],
                )
            )
        return torch.stack(outputs)

    def forward(
        self,
        *,
        target_latent: torch.Tensor,
        target_timestep: torch.Tensor,
        anchor_hidden: torch.Tensor,
        previous_chunk_hidden: torch.Tensor,
        kv_cache: Mapping[Any, Any] | Sequence[Any],
        crossattn_cache: Mapping[Any, Any] | Sequence[Any],
        current_start: int | torch.Tensor,
    ) -> dict[str, torch.Tensor]:
        """Predict one skipped denoising step.

        Args:
            target_latent: Noisy target chunk in ``[B, F, C, H, W]`` layout.
            target_timestep: Per-frame timestep tensor ``[B, F]``.
            anchor_hidden: Same-chunk preceding-step final hidden ``[B, L, D]``.
            previous_chunk_hidden: Previous-chunk same-step hidden ``[B, L, D]``.
            kv_cache: Full 30-layer list or selected-layer mapping/list.
            crossattn_cache: Full list or selected cached text K/V.
            current_start: Current chunk's global token offset.
        """

        cfg = self.predictor_config
        if target_latent.ndim != 5:
            raise ValueError(
                f"target_latent must be [B, F, C, H, W], got {tuple(target_latent.shape)}"
            )
        batch, frames, channels, height, width = target_latent.shape
        if channels != cfg.in_dim:
            raise ValueError(f"target_latent channels {channels} != {cfg.in_dim}")
        expected_timestep = (batch, frames // cfg.patch_size[0])
        if tuple(target_timestep.shape) != expected_timestep:
            raise ValueError(
                f"target_timestep shape {tuple(target_timestep.shape)} != "
                f"{expected_timestep}"
            )
        if target_timestep.device != target_latent.device:
            raise ValueError("target_timestep and target_latent must share a device")
        if target_latent.device != self.patch_embedding.weight.device:
            raise ValueError(
                f"target_latent is on {target_latent.device}, Predictor is on "
                f"{self.patch_embedding.weight.device}"
            )

        latent_cf = target_latent.permute(0, 2, 1, 3, 4).to(
            dtype=self.patch_embedding.weight.dtype
        )
        # Do not wrap frozen patch/time/head modules in no_grad: recursive
        # rollout losses must still backpropagate to an earlier target latent.
        current = self.patch_embedding(latent_cf)
        grid_sizes = torch.tensor(
            [current.shape[2:]] * batch,
            dtype=torch.long,
            device=current.device,
        )
        current = current.flatten(2).transpose(1, 2)
        expected_hidden = (batch, current.shape[1], cfg.dim)
        if tuple(anchor_hidden.shape) != expected_hidden:
            raise ValueError(
                f"anchor_hidden shape {tuple(anchor_hidden.shape)} != {expected_hidden}"
            )
        if tuple(previous_chunk_hidden.shape) != expected_hidden:
            raise ValueError(
                "previous_chunk_hidden shape "
                f"{tuple(previous_chunk_hidden.shape)} != {expected_hidden}"
            )
        current = current.to(dtype=anchor_hidden.dtype)
        hidden = self.feature_fusion(
            current, anchor_hidden, previous_chunk_hidden
        )
        block_condition, head_condition = self._time_condition(
            target_timestep, current
        )
        seq_lens = torch.full(
            (batch,), current.shape[1], dtype=torch.long, device=current.device
        )
        current_start_int = self._scalar_int(current_start, "current_start")
        freqs = self._runtime_freqs(current.device)

        for position, (source_id, block) in enumerate(
            zip(self.source_block_ids, self.predictor_blocks)
        ):
            selected_kv = self._select_cache(
                kv_cache,
                source_id=source_id,
                source_position=position,
                name="kv_cache",
            )
            selected_cross = self._selected_crossattn_cache(
                crossattn_cache,
                source_id=source_id,
                source_position=position,
            )
            hidden = block(
                hidden,
                e=block_condition,
                seq_lens=seq_lens,
                grid_sizes=grid_sizes,
                freqs=freqs,
                context=None,
                context_lens=None,
                block_mask=None,
                kv_cache=selected_kv,
                crossattn_cache=selected_cross,
                current_start=current_start_int,
                cache_start=current_start_int,
            )

        delta_hidden = self.residual_out(hidden)
        pred_hidden = anchor_hidden + delta_hidden
        pred_tokens = self.head(pred_hidden, head_condition)
        pred_flow = self._unpatchify(pred_tokens, grid_sizes)
        return {
            "pred_hidden": pred_hidden,
            "pred_flow": pred_flow,
            "delta_hidden": delta_hidden,
        }

    def trainable_parameter_count(self) -> int:
        return sum(
            parameter.numel()
            for parameter in self.parameters()
            if parameter.requires_grad
        )

    def trainable_parameter_breakdown(self) -> dict[str, int]:
        modules = {
            "feature_fusion": self.feature_fusion,
            "predictor_blocks": self.predictor_blocks,
            "residual_out": self.residual_out,
        }
        return {
            name: sum(
                parameter.numel()
                for parameter in module.parameters()
                if parameter.requires_grad
            )
            for name, module in modules.items()
        }

    def trainable_state_dict(
        self,
        *,
        keep_vars: bool = False,
    ) -> OrderedDict[str, torch.Tensor]:
        """Return only optimizer-owned parameters, suitable for safetensors."""

        trainable = {
            name for name, parameter in self.named_parameters()
            if parameter.requires_grad
        }
        state = super().state_dict(keep_vars=keep_vars)
        return OrderedDict(
            (name, value) for name, value in state.items() if name in trainable
        )

    def load_trainable_state_dict(
        self,
        state_dict: Mapping[str, torch.Tensor],
        *,
        strict: bool = True,
    ) -> None:
        """Load a compact state into a fresh Predictor initialized from Teacher."""

        expected = set(self.trainable_state_dict())
        received = set(state_dict)
        if strict:
            missing = sorted(expected.difference(received))
            unexpected = sorted(received.difference(expected))
            if missing or unexpected:
                raise RuntimeError(
                    "Predictor trainable checkpoint mismatch: "
                    f"missing={missing}, unexpected={unexpected}"
                )
        filtered = {
            name: tensor for name, tensor in state_dict.items() if name in expected
        }
        self.load_state_dict(filtered, strict=False)

    def checkpoint_dict(self) -> dict[str, Any]:
        """Build a compact torch-save payload with architecture metadata."""

        return {
            "format": "self_forcing_wan_predictor_v4",
            "format_version": self.checkpoint_format_version,
            "config": self.config_dict,
            "trainable_state_dict": self.trainable_state_dict(),
        }

    def load_checkpoint_dict(
        self,
        checkpoint: Mapping[str, Any],
        *,
        strict: bool = True,
    ) -> None:
        if checkpoint.get("format") != "self_forcing_wan_predictor_v4":
            raise ValueError(f"Unsupported Predictor checkpoint: {checkpoint.get('format')}")
        saved_config = dict(checkpoint.get("config", {}))
        if strict and saved_config != self.config_dict:
            raise ValueError(
                "Predictor checkpoint config does not match the Teacher/config "
                f"used for reconstruction: saved={saved_config}, current={self.config_dict}"
            )
        state = checkpoint.get("trainable_state_dict")
        if not isinstance(state, Mapping):
            raise ValueError("Predictor checkpoint has no trainable_state_dict")
        self.load_trainable_state_dict(state, strict=strict)


__all__ = [
    "SelfForcingPredictorV4",
    "TripleFeatureFusion",
    "WanPredictorV4Config",
]