| from typing import Callable, Optional, Tuple |
|
|
| import numpy as np |
| import torch |
| from einops import rearrange |
| from torch import nn |
|
|
|
|
| def patchify(x, patch_embedding, check_patchify_match=None, check_patchify_match_prefix="Patchify"): |
| output_dtype = x.dtype |
| input_dtype = output_dtype |
| weight = getattr(patch_embedding, "weight", None) |
| bias = getattr(patch_embedding, "bias", None) |
| if weight is not None and getattr(weight.dtype, "is_floating_point", False): |
| input_dtype = weight.dtype |
| if bias is not None and getattr(bias.dtype, "is_floating_point", False): |
| input_dtype = bias.dtype |
| if x.dtype != input_dtype: |
| x = x.to(input_dtype) |
| x = patch_embedding(x) |
| if x.dtype != output_dtype: |
| x = x.to(output_dtype) |
| b, c, f, h, w = x.shape |
| x = rearrange(x, "b c f h w -> b (f h w) c").contiguous() |
| if check_patchify_match is not None and (f, h, w) != check_patchify_match: |
| raise AssertionError(f"{check_patchify_match_prefix}: x={(f, h, w)} and patchify={check_patchify_match} don't match.") |
| return x, (f, h, w) |
|
|
|
|
| class PatchEmbedding(nn.Module): |
| def __init__(self, init_mode: str = "zero_init", in_channels: Optional[int] = None, wan_patch_embedding: nn.Conv3d = None): |
| super().__init__() |
| if init_mode not in ("zero_init", "wan_patch_embed", "wan_patch_embed_frozen"): |
| raise ValueError(f"Unsupported Vista4D patch embedding init mode: {init_mode}") |
| if wan_patch_embedding is None: |
| raise ValueError("wan_patch_embedding is required") |
|
|
| out_channels, base_in_channels, p1, p2, p3 = wan_patch_embedding.weight.shape |
| if in_channels is None: |
| in_channels = base_in_channels |
| elif in_channels != base_in_channels: |
| init_mode = "zero_init" |
|
|
| if init_mode == "wan_patch_embed_frozen": |
| self.patch_embedding = None |
| else: |
| self.patch_embedding = nn.Conv3d(in_channels, out_channels, kernel_size=(p1, p2, p3), stride=(p1, p2, p3), bias=True) |
| if init_mode == "zero_init": |
| nn.init.zeros_(self.patch_embedding.weight) |
| nn.init.zeros_(self.patch_embedding.bias) |
| else: |
| self.patch_embedding.weight = nn.Parameter(wan_patch_embedding.weight.clone().detach()) |
| self.patch_embedding.bias = nn.Parameter(wan_patch_embedding.bias.clone().detach()) |
|
|
| self.init_mode = init_mode |
| self.out_channels = out_channels |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| wan_patch_embedding: Optional[Callable] = None, |
| check_patchify_match: Optional[Tuple[int, int, int]] = None, |
| check_patchify_match_prefix: str = "Patchify", |
| ): |
| patch_embedding = self.patch_embedding |
| if patch_embedding is None: |
| if wan_patch_embedding is None: |
| raise ValueError("wan_patch_embedding cannot be None with init_mode='wan_patch_embed_frozen'") |
| patch_embedding = wan_patch_embedding |
| return patchify(x, patch_embedding, check_patchify_match=check_patchify_match, check_patchify_match_prefix=check_patchify_match_prefix) |
|
|
|
|
| class RGBMaskPatchEmbedding(nn.Module): |
| def __init__( |
| self, |
| rgb_init_mode: Optional[str] = "wan_patch_embed", |
| mask_init_mode: Optional[str] = None, |
| wan_patch_embedding: nn.Conv3d = None, |
| rgb_in_channels: Optional[int] = None, |
| mask_in_channels: Optional[int] = None, |
| ): |
| super().__init__() |
| if rgb_init_mode is not None: |
| self.rgb_patchify = PatchEmbedding(init_mode=rgb_init_mode, wan_patch_embedding=wan_patch_embedding, in_channels=rgb_in_channels) |
| if mask_init_mode is not None: |
| self.mask_patchify = PatchEmbedding(init_mode=mask_init_mode, wan_patch_embedding=wan_patch_embedding, in_channels=mask_in_channels) |
| if self.mask_patchify.init_mode != "zero_init": |
| out_channels = wan_patch_embedding.weight.shape[0] |
| self.projector = nn.Linear(self.mask_patchify.out_channels, out_channels, bias=True) |
|
|
| def forward( |
| self, |
| rgb_latents: Optional[torch.Tensor] = None, |
| mask_latents: Optional[torch.Tensor] = None, |
| wan_patch_embedding: Optional[Callable] = None, |
| check_patchify_match: Optional[Tuple[int, int, int]] = None, |
| check_patchify_match_prefix: str = "Patch embedding", |
| ): |
| def is_batch_none(value): |
| return value is None or (isinstance(value, (list, tuple, np.ndarray, torch.Tensor)) and any(item is None for item in value)) |
|
|
| output_latents = 0.0 |
| patchify_shape = None |
|
|
| if hasattr(self, "rgb_patchify") and not is_batch_none(rgb_latents): |
| rgb_latents, patchify_shape = self.rgb_patchify( |
| rgb_latents, |
| wan_patch_embedding=wan_patch_embedding, |
| check_patchify_match=check_patchify_match, |
| check_patchify_match_prefix=f"{check_patchify_match_prefix}, RGB", |
| ) |
| output_latents = output_latents + rgb_latents |
|
|
| if hasattr(self, "mask_patchify") and not is_batch_none(mask_latents): |
| mask_latents, mask_shape = self.mask_patchify( |
| mask_latents, |
| wan_patch_embedding=wan_patch_embedding, |
| check_patchify_match=check_patchify_match, |
| check_patchify_match_prefix=f"{check_patchify_match_prefix}, mask", |
| ) |
| if patchify_shape is None: |
| patchify_shape = mask_shape |
| if hasattr(self, "projector"): |
| mask_latents = self.projector(mask_latents) |
| output_latents = output_latents + mask_latents |
|
|
| return output_latents, patchify_shape |
|
|
|
|
| class LatentEncoder(nn.Module): |
| def __init__( |
| self, |
| source_init_mode: str = "wan_patch_embed", |
| point_cloud_init_mode: str = "wan_patch_embed", |
| mask_init_mode: str = "zero_init", |
| use_source_masks: bool = True, |
| use_point_cloud_masks: bool = True, |
| wan_patch_embedding: nn.Conv3d = None, |
| rgb_in_channels: Optional[int] = None, |
| mask_in_channels: int = 2 * 4 * 8 * 8, |
| ): |
| super().__init__() |
| self.output_patch_embedding = RGBMaskPatchEmbedding( |
| rgb_init_mode="wan_patch_embed_frozen", |
| mask_init_mode=None, |
| wan_patch_embedding=wan_patch_embedding, |
| rgb_in_channels=rgb_in_channels, |
| mask_in_channels=None, |
| ) |
| self.source_patch_embedding = RGBMaskPatchEmbedding( |
| rgb_init_mode=source_init_mode, |
| mask_init_mode=mask_init_mode if use_source_masks else None, |
| wan_patch_embedding=wan_patch_embedding, |
| rgb_in_channels=rgb_in_channels, |
| mask_in_channels=mask_in_channels, |
| ) |
| self.point_cloud_patch_embedding = RGBMaskPatchEmbedding( |
| rgb_init_mode=point_cloud_init_mode, |
| mask_init_mode=mask_init_mode if use_point_cloud_masks else None, |
| wan_patch_embedding=wan_patch_embedding, |
| rgb_in_channels=rgb_in_channels, |
| mask_in_channels=mask_in_channels, |
| ) |
|
|
| def forward( |
| self, |
| wan_patch_embedding_fn: Callable, |
| x: torch.Tensor, |
| source_video_latents: Optional[torch.Tensor] = None, |
| source_mask_latents: Optional[torch.Tensor] = None, |
| point_cloud_video_latents: Optional[torch.Tensor] = None, |
| point_cloud_mask_latents: Optional[torch.Tensor] = None, |
| ): |
| x, patchify_shape = self.output_patch_embedding(rgb_latents=x, mask_latents=None, wan_patch_embedding=wan_patch_embedding_fn) |
| source_latents, _ = self.source_patch_embedding( |
| rgb_latents=source_video_latents, |
| mask_latents=source_mask_latents, |
| wan_patch_embedding=wan_patch_embedding_fn, |
| check_patchify_match=patchify_shape, |
| check_patchify_match_prefix="Source patch embedding", |
| ) |
| point_cloud_latents, _ = self.point_cloud_patch_embedding( |
| rgb_latents=point_cloud_video_latents, |
| mask_latents=point_cloud_mask_latents, |
| wan_patch_embedding=wan_patch_embedding_fn, |
| check_patchify_match=patchify_shape, |
| check_patchify_match_prefix="Point cloud patch embedding", |
| ) |
| return x, source_latents, point_cloud_latents, patchify_shape |
|
|
