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