| """Frozen Teacher projectors that rebuild regular-RoPE and ProPE Context K/V.""" |
|
|
| from __future__ import annotations |
|
|
| import copy |
|
|
| import torch |
| from einops import rearrange |
| from torch import nn |
|
|
| from hyvideo.models.transformers.modules.posemb_layers import ( |
| apply_rotary_emb, |
| get_nd_rotary_pos_embed, |
| ) |
| from hyvideo.prope.camera_rope import _prepare_apply_fns_all_dim |
|
|
|
|
| class FrozenTeacherContextProjector(nn.Module): |
| """Rebuild the exact cache layout ``[regular batches, ProPE batches]``.""" |
|
|
| def __init__(self, teacher_transformer: nn.Module, block_ids: tuple[int, ...]) -> None: |
| super().__init__() |
| self.block_ids = tuple(int(value) for value in block_ids) |
| self.heads = int(teacher_transformer.heads_num) |
| self.head_dim = int(teacher_transformer.hidden_size // self.heads) |
| self.rope_dim_list = tuple(teacher_transformer.rope_dim_list) |
| self.rope_theta = float(teacher_transformer.rope_theta) |
| self.patch_size = tuple(teacher_transformer.patch_size) |
| self.latent_height = int(getattr(teacher_transformer, "latent_height", 30)) |
| self.latent_width = int(getattr(teacher_transformer, "latent_width", 52)) |
| self.k_proj = nn.ModuleDict() |
| self.v_proj = nn.ModuleDict() |
| self.k_norm = nn.ModuleDict() |
| for block_id in self.block_ids: |
| block = teacher_transformer.double_blocks[block_id] |
| key = str(block_id) |
| self.k_proj[key] = copy.deepcopy(block.img_attn_k) |
| self.v_proj[key] = copy.deepcopy(block.img_attn_v) |
| self.k_norm[key] = copy.deepcopy(block.img_attn_k_norm) |
| self.eval().requires_grad_(False) |
|
|
| def _rope(self, frames: int, height: int, width: int, device: torch.device): |
| cos, sin = get_nd_rotary_pos_embed( |
| list(self.rope_dim_list), |
| (frames // self.patch_size[0], height, width), |
| theta=self.rope_theta, |
| use_real=True, |
| theta_rescale_factor=1, |
| ) |
| return cos.to(device), sin.to(device) |
|
|
| def forward_block( |
| self, |
| block_id: int, |
| img_modulated: torch.Tensor, |
| context_viewmats: torch.Tensor, |
| context_Ks: torch.Tensor, |
| context_valid_mask: torch.Tensor, |
| *, |
| latent_height: int | None = None, |
| latent_width: int | None = None, |
| ) -> dict[str, torch.Tensor]: |
| key = str(int(block_id)) |
| if key not in self.k_proj: |
| raise ValueError(f"Projector has no block {block_id}") |
| latent_height = self.latent_height if latent_height is None else int(latent_height) |
| latent_width = self.latent_width if latent_width is None else int(latent_width) |
| batch, tokens, _ = img_modulated.shape |
| tokens_per_frame = latent_height * latent_width |
| if tokens % tokens_per_frame: |
| raise ValueError("Context features are not frame-aligned") |
| frames = tokens // tokens_per_frame |
| if tuple(context_valid_mask.shape) != (batch, tokens): |
| raise ValueError("Context valid mask shape mismatch") |
| if tuple(context_viewmats.shape) != (batch, frames, 4, 4): |
| raise ValueError("Context viewmats shape mismatch") |
| if tuple(context_Ks.shape) != (batch, frames, 3, 3): |
| raise ValueError("Context Ks shape mismatch") |
| if tokens == 0: |
| empty = img_modulated.new_empty((2 * batch, self.heads, 0, self.head_dim)) |
| return {"k_vision": empty, "v_vision": empty.clone()} |
|
|
| k = rearrange( |
| self.k_proj[key](img_modulated), |
| "b s (h d) -> b s h d", |
| h=self.heads, |
| ) |
| v = rearrange( |
| self.v_proj[key](img_modulated), |
| "b s (h d) -> b s h d", |
| h=self.heads, |
| ) |
| k = self.k_norm[key](k).to(v) |
| cos, sin = self._rope(frames, latent_height, latent_width, img_modulated.device) |
| _, regular_k = apply_rotary_emb(k, k, (cos, sin), head_first=False) |
| regular_k = regular_k.permute(0, 2, 1, 3) |
| regular_v = v.permute(0, 2, 1, 3) |
|
|
| token_viewmats = context_viewmats.repeat_interleave(tokens_per_frame, dim=1) |
| token_Ks = context_Ks.repeat_interleave(tokens_per_frame, dim=1) |
| _, apply_kv, _ = _prepare_apply_fns_all_dim( |
| head_dim=self.head_dim, |
| viewmats=token_viewmats, |
| Ks=token_Ks, |
| patches_x=None, |
| patches_y=None, |
| image_width=None, |
| image_height=None, |
| ) |
| prope_k = apply_kv(k.permute(0, 2, 1, 3)) |
| prope_v = apply_kv(v.permute(0, 2, 1, 3)) |
| k_cache = torch.cat([regular_k, prope_k], dim=0) |
| v_cache = torch.cat([regular_v, prope_v], dim=0) |
| branch_mask = torch.cat([context_valid_mask, context_valid_mask], dim=0) |
| branch_mask = branch_mask[:, None, :, None] |
| k_cache = k_cache.masked_fill(~branch_mask, 0) |
| v_cache = v_cache.masked_fill(~branch_mask, 0) |
| return {"k_vision": k_cache, "v_vision": v_cache} |
|
|
| def forward(self, context_by_block: dict[int, dict[str, torch.Tensor]]): |
| return { |
| int(block_id): self.forward_block( |
| int(block_id), |
| values["img_modulated"], |
| values["context_viewmats"], |
| values["context_Ks"], |
| values["context_valid_mask"], |
| ) |
| for block_id, values in context_by_block.items() |
| } |
|
|
