Blame for models/predictor_v3.py

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