Cccccz/HY
0
1"""DisCa-style dense image/text feature Predictor for HY-WorldPlay AR."""2 3from __future__ import annotations4 5from dataclasses import asdict, dataclass6from math import prod7from pathlib import Path8from typing import Any9 10import torch11from einops import repeat12from safetensors import safe_open13from torch import nn14from torch.utils.checkpoint import checkpoint15 16from hyvideo.models.transformers.modules.activation_layers import get_activation_layer17from hyvideo.models.transformers.modules.embed_layers import PatchEmbed18from hyvideo.models.transformers.modules.mlp_layers import FinalLayer19from hyvideo.models.transformers.modules.posemb_layers import get_nd_rotary_pos_embed20from hyvideo.models.transformers.worldplay_1_5_transformer import MMDoubleStreamBlock21 22 23@dataclass(frozen=True)24class PredictorConfig:25 hidden_size: int = 204826 heads_num: int = 1627 mlp_width_ratio: float = 4.028 mlp_act_type: str = "gelu_tanh"29 qkv_bias: bool = True30 qk_norm: bool = True31 qk_norm_type: str = "rms"32 attn_mode: str = "flash"33 patch_size: tuple[int, int, int] = (1, 1, 1)34 in_channels: int = 3235 out_channels: int = 3236 concat_condition: bool = True37 rope_dim_list: tuple[int, int, int] = (16, 56, 56)38 rope_theta: float = 256.039 source_block_ids: tuple[int, int] = (1, 52)40 latent_height: int = 3041 latent_width: int = 5242 latent_frames: int = 443 44 45class FeatureFusion(nn.Module):46 def __init__(self, hidden_size: int) -> None:47 super().__init__()48 self.current_norm = nn.LayerNorm(hidden_size, eps=1e-6)49 self.cached_norm = nn.LayerNorm(hidden_size, eps=1e-6)50 self.mlp = nn.Sequential(51 nn.Linear(2 * hidden_size, hidden_size),52 nn.SiLU(),53 nn.Linear(hidden_size, hidden_size),54 )55 56 def forward(self, current: torch.Tensor, cached: torch.Tensor) -> torch.Tensor:57 if current.shape != cached.shape:58 raise ValueError(f"Fusion shape mismatch: {current.shape} != {cached.shape}")59 return self.mlp(60 torch.cat([self.current_norm(current), self.cached_norm(cached)], dim=-1)61 )62 63 64class HYWorldPlayPredictor(nn.Module):65 """Two full double-stream blocks over fused current/cached image and txt features."""66 67 def __init__(self, config: PredictorConfig | None = None) -> None:68 super().__init__()69 self.predictor_config = config or PredictorConfig()70 cfg = self.predictor_config71 self.img_in = PatchEmbed(72 list(cfg.patch_size),73 cfg.in_channels,74 cfg.hidden_size,75 is_reshape_temporal_channels=False,76 concat_condition=cfg.concat_condition,77 )78 self.img_fusion = FeatureFusion(cfg.hidden_size)79 self.txt_fusion = FeatureFusion(cfg.hidden_size)80 self.double_blocks = nn.ModuleList(81 [82 MMDoubleStreamBlock(83 cfg.hidden_size,84 cfg.heads_num,85 mlp_width_ratio=cfg.mlp_width_ratio,86 mlp_act_type=cfg.mlp_act_type,87 attn_mode=cfg.attn_mode,88 qk_norm=cfg.qk_norm,89 qk_norm_type=cfg.qk_norm_type,90 qkv_bias=cfg.qkv_bias,91 )92 for _ in cfg.source_block_ids93 ]94 )95 # HY-WorldPlay action checkpoints add the ProPE output projection after96 # constructing the base Hunyuan blocks (see transformer.add_action_parameters).97 for block in self.double_blocks:98 block.img_attn_prope_proj = nn.Linear(99 cfg.hidden_size, cfg.hidden_size, bias=cfg.qkv_bias100 )101 self.residual_out = nn.Linear(cfg.hidden_size, cfg.hidden_size)102 self.final_layer = FinalLayer(103 cfg.hidden_size,104 list(cfg.patch_size),105 cfg.out_channels,106 get_activation_layer("silu"),107 )108 nn.init.zeros_(self.residual_out.weight)109 nn.init.zeros_(self.residual_out.bias)110 self.gradient_checkpointing = False111 self.attn_param: dict[str, Any] = {112 "thw": [cfg.latent_frames, cfg.latent_height, cfg.latent_width],113 "win_type": "fixed",114 "win_ratio": 0,115 }116 self.img_in.requires_grad_(False)117 self.final_layer.requires_grad_(False)118 119 @property120 def config_dict(self) -> dict[str, Any]:121 return asdict(self.predictor_config)122 123 def enable_gradient_checkpointing(self, enabled: bool = True) -> None:124 self.gradient_checkpointing = enabled125 126 def train(self, mode: bool = True):127 super().train(mode)128 self.img_in.eval()129 self.final_layer.eval()130 return self131 132 @staticmethod133 def _load_prefixed_module(134 module: nn.Module,135 handle,136 prefix: str,137 ) -> None:138 target = module.state_dict()139 loaded = {}140 for name in target:141 key = prefix + name142 if key not in handle.keys():143 raise KeyError(f"Missing Teacher checkpoint key: {key}")144 loaded[name] = handle.get_tensor(key)145 result = module.load_state_dict(loaded, strict=True)146 if result.missing_keys or result.unexpected_keys:147 raise RuntimeError(f"Unexpected load result for {prefix}: {result}")148 149 def load_teacher_initialization(self, checkpoint_path: str | Path) -> None:150 path = str(Path(checkpoint_path).resolve())151 with safe_open(path, framework="pt", device="cpu") as handle:152 self._load_prefixed_module(self.img_in, handle, "img_in.")153 self._load_prefixed_module(self.final_layer, handle, "final_layer.")154 for predictor_block, teacher_id in zip(155 self.double_blocks, self.predictor_config.source_block_ids156 ):157 self._load_prefixed_module(158 predictor_block, handle, f"double_blocks.{teacher_id}."159 )160 self.img_in.requires_grad_(False).eval()161 self.final_layer.requires_grad_(False).eval()162 nn.init.zeros_(self.residual_out.weight)163 nn.init.zeros_(self.residual_out.bias)164 165 def _vision_rope(166 self,167 rope_temporal_size: int,168 start_rope_start_idx: int,169 *,170 device: torch.device,171 dtype: torch.dtype,172 ) -> tuple[torch.Tensor, torch.Tensor]:173 cfg = self.predictor_config174 cos, sin = get_nd_rotary_pos_embed(175 list(cfg.rope_dim_list),176 (rope_temporal_size, cfg.latent_height, cfg.latent_width),177 theta=cfg.rope_theta,178 use_real=True,179 theta_rescale_factor=1,180 )181 tokens_per_frame = cfg.latent_height * cfg.latent_width182 start = start_rope_start_idx * tokens_per_frame183 end = (start_rope_start_idx + cfg.latent_frames) * tokens_per_frame184 cos = cos[start:end].to(device=device, dtype=dtype)185 sin = sin[start:end].to(device=device, dtype=dtype)186 expected = cfg.latent_frames * tokens_per_frame187 if cos.shape[0] != expected:188 raise ValueError(f"RoPE tokens {cos.shape[0]} != {expected}")189 return cos, sin190 191 def _run_block(192 self,193 block: MMDoubleStreamBlock,194 source_block_id: int,195 img: torch.Tensor,196 txt: torch.Tensor,197 vec: torch.Tensor,198 vec_txt: torch.Tensor,199 freqs_cis: tuple[torch.Tensor, torch.Tensor],200 viewmats: torch.Tensor,201 Ks: torch.Tensor,202 ) -> tuple[torch.Tensor, torch.Tensor]:203 def block_forward(img_arg: torch.Tensor, txt_arg: torch.Tensor):204 return block(205 bi_inference=True,206 ar_txt_inference=False,207 ar_vision_inference=False,208 img=img_arg,209 txt=txt_arg,210 vec_txt=vec_txt,211 vec=vec,212 freqs_cis=freqs_cis,213 text_mask=None,214 attn_param=self.attn_param,215 is_flash=False,216 block_idx=source_block_id,217 viewmats=viewmats,218 Ks=Ks,219 )220 221 if self.gradient_checkpointing and self.training:222 return checkpoint(block_forward, img, txt, use_reentrant=False)223 return block_forward(img, txt)224 225 def unpatchify(self, x: torch.Tensor) -> torch.Tensor:226 cfg = self.predictor_config227 batch = x.shape[0]228 expected = cfg.latent_frames * cfg.latent_height * cfg.latent_width229 if x.shape[1] != expected:230 raise ValueError(f"Output tokens {x.shape[1]} != {expected}")231 x = x.reshape(232 batch,233 cfg.latent_frames,234 cfg.latent_height,235 cfg.latent_width,236 cfg.out_channels,237 *cfg.patch_size,238 )239 x = torch.einsum("nthwcopq->nctohpwq", x)240 return x.reshape(241 batch,242 cfg.out_channels,243 cfg.latent_frames * cfg.patch_size[0],244 cfg.latent_height * cfg.patch_size[1],245 cfg.latent_width * cfg.patch_size[2],246 )247 248 def hidden_to_velocity(249 self,250 hidden: torch.Tensor,251 frame_condition: torch.Tensor,252 ) -> torch.Tensor:253 cfg = self.predictor_config254 spatial_tokens = cfg.latent_height * cfg.latent_width255 condition_tokens = frame_condition.repeat_interleave(spatial_tokens, dim=1)256 vec = condition_tokens.reshape(-1, cfg.hidden_size)257 return self.unpatchify(self.final_layer(hidden, vec))258 259 def forward(260 self,261 *,262 target_model_input: torch.Tensor,263 anchor_hidden: torch.Tensor,264 current_txt: torch.Tensor,265 cached_txt: torch.Tensor,266 target_frame_condition: torch.Tensor,267 vec_txt: torch.Tensor,268 target_viewmats: torch.Tensor,269 target_Ks: torch.Tensor,270 rope_temporal_size: torch.Tensor | int,271 start_rope_start_idx: torch.Tensor | int,272 ) -> dict[str, torch.Tensor]:273 cfg = self.predictor_config274 batch = target_model_input.shape[0]275 if batch != 1:276 raise ValueError("Predictor v1 currently requires micro-batch 1")277 if target_model_input.shape[1:] != (278 65,279 cfg.latent_frames,280 cfg.latent_height,281 cfg.latent_width,282 ):283 raise ValueError(f"Unexpected target_model_input: {target_model_input.shape}")284 285 with torch.no_grad():286 current_img = self.img_in(target_model_input)287 if current_img.shape != anchor_hidden.shape:288 raise ValueError(f"Current/anchor mismatch: {current_img.shape} != {anchor_hidden.shape}")289 290 if current_txt.shape[0] != batch:291 current_txt = current_txt.expand(batch, -1, -1)292 if cached_txt.shape[0] != batch:293 cached_txt = cached_txt.expand(batch, -1, -1)294 if vec_txt.shape[0] != batch:295 vec_txt = vec_txt.expand(batch, -1)296 297 img = self.img_fusion(current_img, anchor_hidden)298 txt = self.txt_fusion(current_txt, cached_txt)299 spatial_tokens = cfg.latent_height * cfg.latent_width300 if target_frame_condition.shape != (batch, cfg.latent_frames, cfg.hidden_size):301 raise ValueError(f"Unexpected frame condition: {target_frame_condition.shape}")302 condition_tokens = target_frame_condition.repeat_interleave(spatial_tokens, dim=1)303 vec = condition_tokens.reshape(-1, cfg.hidden_size)304 viewmats = repeat(305 target_viewmats,306 "B T M N -> B (T H W) M N",307 H=cfg.latent_height,308 W=cfg.latent_width,309 )310 Ks = repeat(311 target_Ks,312 "B T M N -> B (T H W) M N",313 H=cfg.latent_height,314 W=cfg.latent_width,315 )316 rope_size = int(rope_temporal_size.reshape(-1)[0].item()) if torch.is_tensor(rope_temporal_size) else int(rope_temporal_size)317 rope_start = int(start_rope_start_idx.reshape(-1)[0].item()) if torch.is_tensor(start_rope_start_idx) else int(start_rope_start_idx)318 freqs_cis = self._vision_rope(319 rope_size,320 rope_start,321 device=img.device,322 dtype=img.dtype,323 )324 self.attn_param["thw"] = [cfg.latent_frames, cfg.latent_height, cfg.latent_width]325 for block, source_id in zip(self.double_blocks, cfg.source_block_ids):326 img, txt = self._run_block(327 block,328 source_id,329 img,330 txt,331 vec,332 vec_txt,333 freqs_cis,334 viewmats,335 Ks,336 )337 338 delta_hidden = self.residual_out(img)339 pred_hidden = anchor_hidden + delta_hidden340 # Frozen parameters still allow the velocity loss to backpropagate to pred_hidden.341 pred_tokens = self.final_layer(pred_hidden, vec)342 pred_velocity = self.unpatchify(pred_tokens)343 return {344 "pred_hidden": pred_hidden,345 "pred_velocity": pred_velocity,346 "delta_hidden": delta_hidden,347 "pred_txt": txt,348 }349 350 def trainable_parameter_count(self) -> int:351 return sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad)352 353 def trainable_parameter_breakdown(self) -> dict[str, int]:354 groups = {355 "img_fusion": self.img_fusion,356 "txt_fusion": self.txt_fusion,357 "double_blocks": self.double_blocks,358 "residual_out": self.residual_out,359 }360 return {361 name: sum(parameter.numel() for parameter in module.parameters() if parameter.requires_grad)362 for name, module in groups.items()363 }364