"""Predictor-v3 with configurable previous-chunk feature fusion.""" from __future__ import annotations import math import torch from torch import nn from .predictor_v2 import HYWorldPlayPredictorV2, PredictorV2Config from .predictor_atc import ATCFeatureFusion class SourceConditionedTokenGate(nn.Module): """Predict one scalar gate per token without materializing a 3C concat.""" def __init__( self, hidden_size: int, gate_hidden_dim: int, initial_probability: float, ) -> None: super().__init__() if gate_hidden_dim <= 0: raise ValueError("gate_hidden_dim must be positive") if not 0.0 < initial_probability < 1.0: raise ValueError("gate_initial_probability must be in (0, 1)") self.current_in = nn.Linear(hidden_size, gate_hidden_dim) self.source_in = nn.Linear(hidden_size, gate_hidden_dim) self.condition_in = nn.Linear(hidden_size, gate_hidden_dim) self.activation = nn.SiLU() self.out = nn.Linear(gate_hidden_dim, 1) nn.init.normal_(self.out.weight, mean=0.0, std=1e-3) nn.init.constant_( self.out.bias, math.log(initial_probability / (1.0 - initial_probability)), ) def forward( self, current: torch.Tensor, source: torch.Tensor, condition: torch.Tensor, ) -> torch.Tensor: if current.shape != source.shape or current.shape != condition.shape: raise ValueError( "Token gate inputs must have identical [B, N, C] shapes; " f"current={current.shape}, source={source.shape}, " f"condition={condition.shape}" ) hidden = ( self.current_in(current) + self.source_in(source) + self.condition_in(condition) ) return torch.sigmoid(self.out(self.activation(hidden))) class TripleFeatureFusion(nn.Module): """Fuse current/anchor tokens with configurable previous-chunk states.""" VALID_MODES = { "none", "same_step", "previous_step", "both", "both_delta_gated" } VALID_FUSION_TYPES = { "concat", "projection_add", "token_gate_prev", "concat_residual_gate_prev", } def __init__( self, hidden_size: int, previous_condition_mode: str = "same_step", predictor_fusion_type: str = "concat", gate_hidden_dim: int = 256, gate_initial_probability: float = 0.8, ) -> None: super().__init__() mode = str(previous_condition_mode).lower() if mode not in self.VALID_MODES: raise ValueError( f"Unknown previous_condition_mode {mode!r}; " f"expected one of {sorted(self.VALID_MODES)}" ) fusion_type = str(predictor_fusion_type).lower() if fusion_type not in self.VALID_FUSION_TYPES: raise ValueError( f"Unknown predictor_fusion_type {fusion_type!r}; " f"expected one of {sorted(self.VALID_FUSION_TYPES)}" ) if fusion_type != "concat" and mode != "same_step": raise ValueError( f"predictor_fusion_type={fusion_type!r} currently requires " "previous_condition_mode='same_step'" ) self.previous_condition_mode = mode self.predictor_fusion_type = fusion_type self.latest_gate_stats: dict[str, torch.Tensor] = {} self.current_norm = nn.LayerNorm(hidden_size, eps=1e-6) self.anchor_norm = nn.LayerNorm(hidden_size, eps=1e-6) if mode != "none": self.previous_norm = nn.LayerNorm(hidden_size, eps=1e-6) if mode in {"both", "both_delta_gated"}: self.previous_step_norm = nn.LayerNorm(hidden_size, eps=1e-6) input_count = { "none": 2, "same_step": 3, "previous_step": 3, "both": 4, # Keep the proven same-step fusion as the base. The extra # previous-step condition enters only through a zero-initialized # gated residual below, so initialization can be exactly # equivalent to a trained same-step checkpoint. "both_delta_gated": 3, }[mode] if fusion_type in {"concat", "concat_residual_gate_prev"}: self.mlp = nn.Sequential( nn.Linear(input_count * hidden_size, 2 * hidden_size), nn.SiLU(), nn.Linear(2 * hidden_size, hidden_size), ) if fusion_type == "concat_residual_gate_prev": self.gate_p = SourceConditionedTokenGate( hidden_size, gate_hidden_dim, gate_initial_probability, ) self.residual_mlp = nn.Sequential( nn.Linear(hidden_size, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size), ) nn.init.zeros_(self.residual_mlp[-1].weight) nn.init.zeros_(self.residual_mlp[-1].bias) else: self.proj_e = nn.Linear(hidden_size, hidden_size) self.proj_a = nn.Linear(hidden_size, hidden_size) self.proj_p = nn.Linear(hidden_size, hidden_size) for projection in (self.proj_e, self.proj_a, self.proj_p): nn.init.eye_(projection.weight) nn.init.zeros_(projection.bias) self.proj_norm_e = nn.LayerNorm( hidden_size, eps=1e-6, elementwise_affine=False ) self.proj_norm_a = nn.LayerNorm( hidden_size, eps=1e-6, elementwise_affine=False ) self.proj_norm_p = nn.LayerNorm( hidden_size, eps=1e-6, elementwise_affine=False ) self.fusion_norm = nn.LayerNorm(hidden_size, eps=1e-6) if fusion_type == "token_gate_prev": self.gate_p = SourceConditionedTokenGate( hidden_size, gate_hidden_dim, gate_initial_probability, ) if fusion_type == "concat" and mode == "both_delta_gated": self.delta_mlp = nn.Sequential( nn.Linear(hidden_size, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size), ) self.delta_gate = nn.Parameter(torch.zeros(hidden_size)) nn.init.zeros_(self.delta_mlp[-1].weight) nn.init.zeros_(self.delta_mlp[-1].bias) def forward( self, current: torch.Tensor, anchor: torch.Tensor, previous_same_step: torch.Tensor | None = None, previous_previous_step: torch.Tensor | None = None, condition_tokens: torch.Tensor | None = None, ) -> torch.Tensor: if current.shape != anchor.shape: raise ValueError( "Fusion current/anchor shape mismatch: " f"current={current.shape}, anchor={anchor.shape}" ) features = [self.current_norm(current), self.anchor_norm(anchor)] mode = self.previous_condition_mode if mode in {"same_step", "both", "both_delta_gated"}: if previous_same_step is None or previous_same_step.shape != current.shape: raise ValueError( "same-step previous hidden must match current shape; got " f"{None if previous_same_step is None else previous_same_step.shape} " f"versus {current.shape}" ) features.append(self.previous_norm(previous_same_step)) if mode in {"previous_step", "both", "both_delta_gated"}: if ( previous_previous_step is None or previous_previous_step.shape != current.shape ): raise ValueError( "previous-step previous hidden must match current shape; got " f"{None if previous_previous_step is None else previous_previous_step.shape} " f"versus {current.shape}" ) normalizer = ( self.previous_step_norm if mode in {"both", "both_delta_gated"} else self.previous_norm ) normalized_previous_step = normalizer(previous_previous_step) if mode in {"previous_step", "both"}: features.append(normalized_previous_step) if self.predictor_fusion_type in {"projection_add", "token_gate_prev"}: self.latest_gate_stats = {} z_e = self.proj_norm_e(self.proj_e(features[0])) z_a = self.proj_norm_a(self.proj_a(features[1])) z_p = self.proj_norm_p(self.proj_p(features[2])) if self.predictor_fusion_type == "projection_add": return self.fusion_norm(z_e + z_a + z_p) if condition_tokens is None: raise ValueError("token_gate_prev requires condition_tokens") gate_p = self.gate_p(z_e, z_p, condition_tokens) self._update_gate_stats(gate_p) return self.fusion_norm(z_e + z_a + gate_p * z_p) self.latest_gate_stats = {} base = self.mlp(torch.cat(features, dim=-1)) if self.predictor_fusion_type == "concat_residual_gate_prev": if condition_tokens is None: raise ValueError( "concat_residual_gate_prev requires condition_tokens" ) previous_anchor_delta = features[2] - features[1] gate_p = self.gate_p( features[0], previous_anchor_delta, condition_tokens ) self._update_gate_stats(gate_p) return base + gate_p * self.residual_mlp(previous_anchor_delta) if mode != "both_delta_gated": return base delta = normalized_previous_step - self.previous_norm(previous_same_step) gate = torch.sigmoid(self.delta_gate).view(1, 1, -1) return base + gate * self.delta_mlp(delta) def _update_gate_stats(self, gate_p: torch.Tensor) -> None: flat_gate = gate_p.detach().float().reshape(-1) quantiles = torch.quantile( flat_gate, torch.tensor( (0.1, 0.5, 0.9), device=flat_gate.device, dtype=flat_gate.dtype ), ) self.latest_gate_stats = { "gate_p_mean": flat_gate.mean(), "gate_p_std": flat_gate.std(unbiased=False), "gate_p_p10": quantiles[0], "gate_p_p50": quantiles[1], "gate_p_p90": quantiles[2], "gate_p_near_zero": (flat_gate < 0.05).float().mean(), "gate_p_near_one": (flat_gate > 0.95).float().mean(), } class HYWorldPlayPredictorV3(HYWorldPlayPredictorV2): """Fresh v3 Predictor with configurable previous-chunk conditions.""" requires_previous_chunk_hidden = True def __init__(self, config: PredictorV2Config | None = None) -> None: super().__init__(config) self.previous_condition_mode = self.predictor_config.previous_condition_mode arch = self.predictor_config.predictor_arch.lower() if arch not in {"baseline", "atc"}: raise ValueError("predictor_arch must be 'baseline' or 'atc'") self.predictor_arch = arch if arch == "atc": if self.previous_condition_mode != "same_step": raise ValueError("ATC requires previous_condition_mode='same_step'") self.requires_previous_chunk_hidden = True cfg = self.predictor_config transport_heads = cfg.atc_transport_heads or cfg.heads_num if cfg.atc_transport_mode == "global": head_dim = cfg.hidden_size // transport_heads if cfg.hidden_size % transport_heads or head_dim != sum(cfg.rope_dim_list): raise ValueError( "Global ATC transport reuses the Teacher RoPE, so " f"hidden_size/heads ({cfg.hidden_size}/{transport_heads}) must equal " f"sum(rope_dim_list)={sum(cfg.rope_dim_list)}" ) self.img_fusion = ATCFeatureFusion( hidden_size=self.predictor_config.hidden_size, mlp_hidden_size=self.predictor_config.atc_mlp_hidden_dim, gate_hidden_size=self.predictor_config.atc_gate_hidden_dim, gate_initial_probability=( self.predictor_config.atc_gate_initial_probability ), transport_heads=transport_heads, transport_window=self.predictor_config.atc_transport_window, transport_query_chunk_size=( self.predictor_config.atc_transport_query_chunk_size ), transport_initial_scale=( self.predictor_config.atc_transport_initial_scale ), frames=self.predictor_config.latent_frames, height=self.predictor_config.latent_height, width=self.predictor_config.latent_width, variant=self.predictor_config.atc_variant, transport_mode=self.predictor_config.atc_transport_mode, previous_scope=self.predictor_config.atc_previous_scope, attn_mode=self.predictor_config.attn_mode, ) else: self.requires_previous_chunk_hidden = self.previous_condition_mode != "none" self.img_fusion = TripleFeatureFusion( self.predictor_config.hidden_size, previous_condition_mode=self.previous_condition_mode, predictor_fusion_type=self.predictor_config.predictor_fusion_type, gate_hidden_dim=self.predictor_config.gate_hidden_dim, gate_initial_probability=self.predictor_config.gate_initial_probability, ) def fuse_features( self, current_img: torch.Tensor, anchor_hidden: torch.Tensor, previous_chunk_hidden: torch.Tensor | None = None, previous_step_chunk_hidden: torch.Tensor | None = None, condition_tokens: torch.Tensor | None = None, ) -> torch.Tensor: return self.img_fusion( current_img, anchor_hidden, previous_same_step=previous_chunk_hidden, previous_previous_step=previous_step_chunk_hidden, condition_tokens=condition_tokens, )