"""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", ]