models/predictor_v2.py

"""AR-context Predictor-v2 for HY-WorldPlay skipped denoising steps."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Mapping

import torch
from einops import repeat
from safetensors import safe_open
from torch import nn
from torch.utils.checkpoint import checkpoint

from hyvideo.models.transformers.modules.activation_layers import get_activation_layer
from hyvideo.models.transformers.modules.embed_layers import PatchEmbed
from hyvideo.models.transformers.modules.mlp_layers import FinalLayer
from hyvideo.models.transformers.modules.posemb_layers import get_nd_rotary_pos_embed
from hyvideo.models.transformers.worldplay_1_5_transformer import MMDoubleStreamBlock

from .predictor import FeatureFusion


@dataclass(frozen=True)
class PredictorV2Config:
    hidden_size: int = 2048
    heads_num: int = 16
    mlp_width_ratio: float = 4.0
    mlp_act_type: str = "gelu_tanh"
    qkv_bias: bool = True
    qk_norm: bool = True
    qk_norm_type: str = "rms"
    attn_mode: str = "flash"
    patch_size: tuple[int, int, int] = (1, 1, 1)
    in_channels: int = 32
    out_channels: int = 32
    concat_condition: bool = True
    rope_dim_list: tuple[int, int, int] = (16, 56, 56)
    rope_theta: float = 256.0
    source_block_ids: tuple[int, int] = (1, 52)
    teacher_block_count: int = 54
    latent_height: int = 30
    latent_width: int = 52
    latent_frames: int = 4
    previous_condition_mode: str = "same_step"
    predictor_fusion_type: str = "concat"
    gate_hidden_dim: int = 256
    gate_initial_probability: float = 0.8
    predictor_arch: str = "baseline"
    atc_variant: str = "full"
    atc_mlp_hidden_dim: int = 4096
    atc_gate_hidden_dim: int = 512
    atc_gate_initial_probability: float = 0.3
    # Transport heads default to the Teacher's head count so head_dim equals
    # sum(rope_dim_list) and the Teacher RoPE can be reused (global mode).
    atc_transport_heads: int | None = None
    atc_transport_window: int = 3
    atc_transport_query_chunk_size: int = 128
    atc_transport_initial_scale: float = 0.1
    # "global": Causal-Forcing RoPE'd global cross-chunk attention (default);
    # "local": the original 3x3 local window transport (legacy checkpoints).
    atc_transport_mode: str = "global"
    # Keys of the global transport: the whole previous chunk or its last frame.
    atc_previous_scope: str = "chunk"


class HYWorldPlayPredictorV2(nn.Module):
    """Two Teacher-initialized AR vision blocks reading exact selected-layer memory."""

    uses_context_kv = True

    def __init__(self, config: PredictorV2Config | None = None) -> None:
        super().__init__()
        self.predictor_config = config or PredictorV2Config()
        cfg = self.predictor_config
        if len(cfg.source_block_ids) != 2 or len(set(cfg.source_block_ids)) != 2:
            raise ValueError("Predictor v2 requires exactly two distinct source blocks")
        if any(block_id < 0 or block_id >= cfg.teacher_block_count for block_id in cfg.source_block_ids):
            raise ValueError(f"Invalid source blocks: {cfg.source_block_ids}")

        self.img_in = PatchEmbed(
            list(cfg.patch_size),
            cfg.in_channels,
            cfg.hidden_size,
            is_reshape_temporal_channels=False,
            concat_condition=cfg.concat_condition,
        )
        self.img_fusion = FeatureFusion(cfg.hidden_size)
        self.double_blocks = nn.ModuleList(
            [
                MMDoubleStreamBlock(
                    cfg.hidden_size,
                    cfg.heads_num,
                    mlp_width_ratio=cfg.mlp_width_ratio,
                    mlp_act_type=cfg.mlp_act_type,
                    attn_mode=cfg.attn_mode,
                    qk_norm=cfg.qk_norm,
                    qk_norm_type=cfg.qk_norm_type,
                    qkv_bias=cfg.qkv_bias,
                )
                for _ in cfg.source_block_ids
            ]
        )
        for block in self.double_blocks:
            block.img_attn_prope_proj = nn.Linear(
                cfg.hidden_size, cfg.hidden_size, bias=cfg.qkv_bias
            )
            # AR vision inference consumes cached text K/V and never executes the
            # text stream. Keep those Teacher weights frozen and out of checkpoints.
            for name, parameter in block.named_parameters():
                if name.startswith("txt_"):
                    parameter.requires_grad_(False)

        self.residual_out = nn.Linear(cfg.hidden_size, cfg.hidden_size)
        self.final_layer = FinalLayer(
            cfg.hidden_size,
            list(cfg.patch_size),
            cfg.out_channels,
            get_activation_layer("silu"),
        )
        nn.init.zeros_(self.residual_out.weight)
        nn.init.zeros_(self.residual_out.bias)
        self.gradient_checkpointing = False
        self.attn_param: dict[str, Any] = {
            "thw": [cfg.latent_frames, cfg.latent_height, cfg.latent_width],
            "win_type": "fixed",
            "win_ratio": 0,
            "enable_sparse": False,
            "enable_decode_sparse": False,
            "enable_prefill_sparse": False,
            "append_vision_cache": False,
        }
        self.img_in.requires_grad_(False)
        self.final_layer.requires_grad_(False)

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

    def enable_gradient_checkpointing(self, enabled: bool = True) -> None:
        self.gradient_checkpointing = enabled

    def train(self, mode: bool = True):
        super().train(mode)
        self.img_in.eval()
        self.final_layer.eval()
        return self

    @staticmethod
    def _load_prefixed_module(module: nn.Module, handle, prefix: str) -> None:
        target = module.state_dict()
        loaded = {}
        keys = set(handle.keys())
        for name in target:
            key = prefix + name
            if key not in keys:
                raise KeyError(f"Missing Teacher checkpoint key: {key}")
            loaded[name] = handle.get_tensor(key)
        result = module.load_state_dict(loaded, strict=True)
        if result.missing_keys or result.unexpected_keys:
            raise RuntimeError(f"Unexpected Teacher load result for {prefix}: {result}")

    def load_teacher_initialization(self, checkpoint_path: str | Path) -> None:
        path = str(Path(checkpoint_path).resolve())
        with safe_open(path, framework="pt", device="cpu") as handle:
            self._load_prefixed_module(self.img_in, handle, "img_in.")
            self._load_prefixed_module(self.final_layer, handle, "final_layer.")
            for predictor_block, teacher_id in zip(
                self.double_blocks, self.predictor_config.source_block_ids
            ):
                self._load_prefixed_module(
                    predictor_block, handle, f"double_blocks.{teacher_id}."
                )
        self.img_in.requires_grad_(False).eval()
        self.final_layer.requires_grad_(False).eval()
        for block in self.double_blocks:
            for name, parameter in block.named_parameters():
                if name.startswith("txt_"):
                    parameter.requires_grad_(False)
        nn.init.zeros_(self.residual_out.weight)
        nn.init.zeros_(self.residual_out.bias)

    def _vision_rope_single(
        self,
        rope_temporal_size: int,
        start_rope_start_idx: int,
        *,
        device: torch.device,
        dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        cfg = self.predictor_config
        cos, sin = get_nd_rotary_pos_embed(
            list(cfg.rope_dim_list),
            (rope_temporal_size, cfg.latent_height, cfg.latent_width),
            theta=cfg.rope_theta,
            use_real=True,
            theta_rescale_factor=1,
        )
        tokens_per_frame = cfg.latent_height * cfg.latent_width
        start = start_rope_start_idx * tokens_per_frame
        end = (start_rope_start_idx + cfg.latent_frames) * tokens_per_frame
        cos = cos[start:end].to(device=device, dtype=dtype)
        sin = sin[start:end].to(device=device, dtype=dtype)
        expected = cfg.latent_frames * tokens_per_frame
        if cos.shape[0] != expected:
            raise ValueError(f"RoPE tokens {cos.shape[0]} != {expected}")
        return cos, sin

    def _vision_rope_span(
        self,
        rope_temporal_size: int,
        start_frame: int,
        frames: int,
        *,
        device: torch.device,
        dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """RoPE tables for ``frames`` latent frames starting at ``start_frame``."""
        cfg = self.predictor_config
        if start_frame < 0 or start_frame + frames > rope_temporal_size:
            raise ValueError(
                f"RoPE span [{start_frame}, {start_frame + frames}) outside "
                f"temporal size {rope_temporal_size}"
            )
        cos, sin = get_nd_rotary_pos_embed(
            list(cfg.rope_dim_list),
            (rope_temporal_size, cfg.latent_height, cfg.latent_width),
            theta=cfg.rope_theta,
            use_real=True,
            theta_rescale_factor=1,
        )
        tokens_per_frame = cfg.latent_height * cfg.latent_width
        start = start_frame * tokens_per_frame
        end = (start_frame + frames) * tokens_per_frame
        return cos[start:end].to(device=device, dtype=dtype), sin[start:end].to(device=device, dtype=dtype)

    def _atc_key_rope_batch(
        self,
        rope_temporal_size: torch.Tensor | int,
        start_rope_start_idx: torch.Tensor | int,
        *,
        batch: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Batched RoPE tables for the ATC transport keys (previous chunk or its last frame)."""
        cfg = self.predictor_config
        frames = cfg.latent_frames if cfg.atc_previous_scope == "chunk" else 1
        def values(value, name):
            result = ([int(v) for v in value.detach().reshape(-1).cpu().tolist()]
                      if torch.is_tensor(value) else [int(value)])
            if len(result) == 1:
                result *= batch
            if len(result) != batch:
                raise ValueError(f"{name} contains {len(result)} values for batch {batch}")
            return result
        sizes = values(rope_temporal_size, "rope_temporal_size")
        starts = values(start_rope_start_idx, "start_rope_start_idx")
        pairs = []
        for size, start in zip(sizes, starts):
            key_start = start - frames
            if key_start < 0:
                raise ValueError(
                    "ATC transport needs a real previous chunk; "
                    f"start_rope_start_idx={start} < key frames {frames}"
                )
            pairs.append(self._vision_rope_span(size, key_start, frames, device=device, dtype=dtype))
        return (torch.stack([p[0] for p in pairs]), torch.stack([p[1] for p in pairs]))

    def _vision_rope_batch(
        self,
        rope_temporal_size: torch.Tensor | int,
        start_rope_start_idx: torch.Tensor | int,
        *,
        batch: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        def values(value: torch.Tensor | int, name: str) -> list[int]:
            if torch.is_tensor(value):
                result = [int(item) for item in value.detach().reshape(-1).cpu().tolist()]
            else:
                result = [int(value)]
            if len(result) == 1:
                result *= batch
            if len(result) != batch:
                raise ValueError(f"{name} contains {len(result)} values for batch {batch}")
            return result

        sizes = values(rope_temporal_size, "rope_temporal_size")
        starts = values(start_rope_start_idx, "start_rope_start_idx")
        pairs = [
            self._vision_rope_single(size, start, device=device, dtype=dtype)
            for size, start in zip(sizes, starts)
        ]
        return (
            torch.stack([pair[0] for pair in pairs], dim=0),
            torch.stack([pair[1] for pair in pairs], dim=0),
        )

    def _make_cache(
        self,
        context_kv: Mapping[int, Mapping[str, torch.Tensor]],
    ) -> list[dict[str, torch.Tensor] | None]:
        cfg = self.predictor_config
        result: list[dict[str, torch.Tensor] | None] = [None] * cfg.teacher_block_count
        for block_id in cfg.source_block_ids:
            if block_id not in context_kv:
                raise ValueError(f"Missing Predictor context block {block_id}")
            cache = context_kv[block_id]
            missing = {"k_txt", "v_txt", "k_vision", "v_vision"}.difference(cache)
            if missing:
                raise ValueError(f"Block {block_id} missing context tensors: {sorted(missing)}")
            if cache["k_txt"].shape != cache["v_txt"].shape:
                raise ValueError(f"Block {block_id} text K/V mismatch")
            k_vision = cache["k_vision"]
            v_vision = cache["v_vision"]
            if (k_vision is None) != (v_vision is None):
                raise ValueError(f"Block {block_id} has incomplete vision KV")
            if k_vision is None:
                k_txt = cache["k_txt"]
                empty_shape = (
                    2 * int(k_txt.shape[0]),
                    int(k_txt.shape[1]),
                    0,
                    int(k_txt.shape[3]),
                )
                k_vision = k_txt.new_empty(empty_shape)
                v_vision = cache["v_txt"].new_empty(empty_shape)
            if k_vision.shape != v_vision.shape:
                raise ValueError(f"Block {block_id} vision K/V mismatch")
            batch = int(cache["k_txt"].shape[0])
            text_valid_mask = cache.get("text_valid_mask")
            if text_valid_mask is None:
                text_valid_mask = torch.ones(
                    batch,
                    cache["k_txt"].shape[2],
                    dtype=torch.bool,
                    device=cache["k_txt"].device,
                )
            context_valid_mask = cache.get("context_valid_mask")
            if context_valid_mask is None:
                context_valid_mask = torch.ones(
                    batch,
                    k_vision.shape[2],
                    dtype=torch.bool,
                    device=k_vision.device,
                )
            if tuple(text_valid_mask.shape) != (batch, cache["k_txt"].shape[2]):
                raise ValueError(f"Block {block_id} text mask/KV mismatch")
            if tuple(context_valid_mask.shape) != (batch, k_vision.shape[2]):
                raise ValueError(f"Block {block_id} context mask/KV mismatch")
            if k_vision.shape[0] != 2 * batch:
                raise ValueError(
                    f"Block {block_id} vision KV batch {k_vision.shape[0]} != {2 * batch}"
                )
            result[block_id] = {
                "k_txt": cache["k_txt"],
                "v_txt": cache["v_txt"],
                "k_vision": k_vision,
                "v_vision": v_vision,
                "text_valid_mask": text_valid_mask.bool(),
                "context_valid_mask": context_valid_mask.bool(),
            }
        return result

    def _run_block(
        self,
        block: MMDoubleStreamBlock,
        source_block_id: int,
        img: torch.Tensor,
        vec: torch.Tensor,
        freqs_cis: tuple[torch.Tensor, torch.Tensor],
        viewmats: torch.Tensor,
        Ks: torch.Tensor,
        kv_cache: list[dict[str, torch.Tensor] | None],
    ) -> torch.Tensor:
        def block_forward(img_arg: torch.Tensor) -> torch.Tensor:
            output, _ = block(
                bi_inference=False,
                ar_txt_inference=False,
                ar_vision_inference=True,
                img=img_arg,
                vec=vec,
                freqs_cis=freqs_cis,
                attn_param=self.attn_param,
                block_idx=source_block_id,
                viewmats=viewmats,
                Ks=Ks,
                kv_cache=kv_cache,
                cache_vision=False,
            )
            return output

        if self.gradient_checkpointing and self.training:
            return checkpoint(block_forward, img, use_reentrant=False)
        return block_forward(img)

    def unpatchify(self, x: torch.Tensor) -> torch.Tensor:
        cfg = self.predictor_config
        batch = x.shape[0]
        expected = cfg.latent_frames * cfg.latent_height * cfg.latent_width
        if x.shape[1] != expected:
            raise ValueError(f"Output tokens {x.shape[1]} != {expected}")
        x = x.reshape(
            batch,
            cfg.latent_frames,
            cfg.latent_height,
            cfg.latent_width,
            cfg.out_channels,
            *cfg.patch_size,
        )
        x = torch.einsum("nthwcopq->nctohpwq", x)
        return x.reshape(
            batch,
            cfg.out_channels,
            cfg.latent_frames * cfg.patch_size[0],
            cfg.latent_height * cfg.patch_size[1],
            cfg.latent_width * cfg.patch_size[2],
        )

    def hidden_to_velocity(
        self,
        hidden: torch.Tensor,
        frame_condition: torch.Tensor,
    ) -> torch.Tensor:
        cfg = self.predictor_config
        spatial_tokens = cfg.latent_height * cfg.latent_width
        condition_tokens = frame_condition.repeat_interleave(spatial_tokens, dim=1)
        vec = condition_tokens.reshape(-1, cfg.hidden_size)
        return self.unpatchify(self.final_layer(hidden, vec))

    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:
        if previous_chunk_hidden is not None or previous_step_chunk_hidden is not None:
            raise ValueError("Predictor v2 does not consume previous-chunk hidden states")
        return self.img_fusion(current_img, anchor_hidden)

    def forward(
        self,
        *,
        target_model_input: torch.Tensor,
        anchor_hidden: torch.Tensor,
        target_frame_condition: torch.Tensor,
        target_viewmats: torch.Tensor,
        target_Ks: torch.Tensor,
        rope_temporal_size: torch.Tensor | int,
        start_rope_start_idx: torch.Tensor | int,
        context_kv: Mapping[int, Mapping[str, torch.Tensor]],
        previous_chunk_hidden: torch.Tensor | None = None,
        previous_step_chunk_hidden: torch.Tensor | None = None,
        anchor_distance: torch.Tensor | None = None,
        return_debug: bool = False,
    ) -> dict[str, torch.Tensor]:
        cfg = self.predictor_config
        batch = target_model_input.shape[0]
        expected_input = (
            65,
            cfg.latent_frames,
            cfg.latent_height,
            cfg.latent_width,
        )
        if tuple(target_model_input.shape[1:]) != expected_input:
            raise ValueError(f"Unexpected target_model_input: {target_model_input.shape}")

        # ``img_in`` is frozen, but its input must remain differentiable for
        # recursive rollout training: later-step losses backpropagate through
        # the scheduler sample into earlier Predictor calls.  With ordinary
        # offline inputs ``requires_grad`` is false, so this creates no extra
        # graph there.
        current_img = self.img_in(target_model_input)
        if current_img.shape != anchor_hidden.shape:
            raise ValueError(f"Current/anchor mismatch: {current_img.shape} != {anchor_hidden.shape}")
        if tuple(target_frame_condition.shape) != (
            batch,
            cfg.latent_frames,
            cfg.hidden_size,
        ):
            raise ValueError(f"Unexpected frame condition: {target_frame_condition.shape}")

        spatial_tokens = cfg.latent_height * cfg.latent_width
        condition_tokens = target_frame_condition.repeat_interleave(
            spatial_tokens, dim=1
        )
        freqs_cis = self._vision_rope_batch(
            rope_temporal_size,
            start_rope_start_idx,
            batch=batch,
            device=current_img.device,
            dtype=current_img.dtype,
        )
        atc_prepared = None
        if getattr(self.img_fusion, "is_atc", False):
            key_freqs = None
            if getattr(self.img_fusion, "transport_mode", "local") == "global":
                key_freqs = self._atc_key_rope_batch(
                    rope_temporal_size,
                    start_rope_start_idx,
                    batch=batch,
                    device=current_img.device,
                    dtype=current_img.dtype,
                )
            atc_prepared = self.img_fusion.prepare_features(
                current_img,
                anchor_hidden,
                previous_chunk_hidden,
                condition_tokens,
                anchor_distance,
                query_freqs=freqs_cis,
                key_freqs=key_freqs,
            )
            img = atc_prepared["backbone_input"]
        else:
            img = self.fuse_features(
                current_img,
                anchor_hidden,
                previous_chunk_hidden=previous_chunk_hidden,
                previous_step_chunk_hidden=previous_step_chunk_hidden,
                condition_tokens=condition_tokens,
            )
        vec = condition_tokens.reshape(-1, cfg.hidden_size)
        viewmats = repeat(
            target_viewmats,
            "B T M N -> B (T H W) M N",
            H=cfg.latent_height,
            W=cfg.latent_width,
        )
        Ks = repeat(
            target_Ks,
            "B T M N -> B (T H W) M N",
            H=cfg.latent_height,
            W=cfg.latent_width,
        )
        kv_cache = self._make_cache(context_kv)
        for source_id in cfg.source_block_ids:
            if kv_cache[source_id]["k_txt"].shape[0] != batch:
                raise ValueError(f"Block {source_id} context batch does not match input batch")
        self.attn_param["thw"] = [cfg.latent_frames, cfg.latent_height, cfg.latent_width]
        for block, source_id in zip(self.double_blocks, cfg.source_block_ids):
            img = self._run_block(
                block,
                source_id,
                img,
                vec,
                freqs_cis,
                viewmats,
                Ks,
                kv_cache,
            )

        delta_hidden = self.residual_out(img)
        atc_result = None
        if atc_prepared is None:
            pred_hidden = anchor_hidden + delta_hidden
        else:
            atc_result = self.img_fusion.finalize(
                anchor_hidden, delta_hidden, atc_prepared
            )
            pred_hidden = atc_result["pred_hidden"]
        pred_tokens = self.final_layer(pred_hidden, vec)
        pred_velocity = self.unpatchify(pred_tokens)
        result = {
            "pred_hidden": pred_hidden,
            "pred_velocity": pred_velocity,
            "delta_hidden": delta_hidden,
        }
        gate_stats = getattr(self.img_fusion, "latest_gate_stats", None)
        if gate_stats:
            result["gate_stats"] = gate_stats
        if atc_result is not None:
            result["atc_stats"] = atc_result["atc_stats"]
            if return_debug:
                result.update(
                    delta_anchor=delta_hidden,
                    delta_correction=atc_result["delta_correction"],
                    gate=atc_result["gate"],
                    transported_prev=atc_result["transported_prev"],
                )
        return result

    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]:
        groups = {
            "img_fusion": self.img_fusion,
            "ar_vision_blocks": self.double_blocks,
            "residual_out": self.residual_out,
        }
        return {
            name: sum(parameter.numel() for parameter in module.parameters() if parameter.requires_grad)
            for name, module in groups.items()
        }