HY / models /predictor_v2.py

Upload batch 46: 223 files (7.45 GiB)

ba798d3 verified about 1 month ago

24.3 kB

"""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()
        }

	"""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()
	}