| from typing import Optional |
| import torch |
| from torch import nn |
| from omegaconf import DictConfig |
| from einops import rearrange, repeat |
| from ..modules.embeddings import ( |
| RandomDropoutPatchEmbed, |
| RandomEmbeddingDropout, |
| ) |
| from .dit3d import DiT3D |
|
|
|
|
| class DiT3DPose(DiT3D): |
|
|
| def __init__( |
| self, |
| cfg: DictConfig, |
| x_shape: torch.Size, |
| max_tokens: int, |
| external_cond_dim: int, |
| use_causal_mask=True, |
| ): |
| self.conditioning_modeling = cfg.conditioning.modeling |
| self.conditioning_type = cfg.conditioning.type |
| self.conditioning_dropout = cfg.external_cond_dropout |
| super().__init__( |
| cfg, |
| x_shape, |
| max_tokens, |
| cfg.conditioning.dim, |
| use_causal_mask, |
| ) |
|
|
| @property |
| def in_channels(self) -> int: |
| return ( |
| self.x_shape[0] + self.external_cond_dim |
| if self.conditioning_modeling == "concat" |
| else self.x_shape[0] |
| ) |
|
|
| @property |
| def external_cond_emb_dim(self) -> int: |
| return self.cfg.hidden_size |
|
|
| def _build_external_cond_embedding(self) -> Optional[nn.Module]: |
| if self.conditioning_type == "global": |
| return super()._build_external_cond_embedding() |
| match self.conditioning_modeling: |
| case "concat": |
| return RandomEmbeddingDropout( |
| p=self.conditioning_dropout, |
| ) |
| case "film": |
| return RandomDropoutPatchEmbed( |
| dropout_prob=self.conditioning_dropout, |
| img_size=self.x_shape[1], |
| patch_size=self.cfg.patch_size, |
| in_chans=self.external_cond_dim, |
| embed_dim=self.external_cond_emb_dim, |
| bias=True, |
| ) |
| case _: |
| raise ValueError( |
| f"Unknown external condition modeling: {self.conditioning_modeling}" |
| ) |
|
|
| def initialize_weights(self) -> None: |
| super().initialize_weights() |
| if self.conditioning_type != "global" and self.conditioning_modeling == "film": |
| self._patch_embedder_init(self.external_cond_embedding.patch_embedder) |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| noise_levels: torch.Tensor, |
| external_cond: Optional[torch.Tensor] = None, |
| external_cond_mask: Optional[torch.Tensor] = None, |
| ) -> torch.Tensor: |
| assert ( |
| external_cond is not None |
| ), "External condition (camera pose) is required for DiT3DPose model." |
| input_batch_size = x.shape[0] |
| external_cond_emb = self.external_cond_embedding( |
| external_cond, external_cond_mask |
| ) |
| if self.conditioning_modeling == "concat": |
| x = torch.cat( |
| [x, external_cond_emb], |
| dim=2, |
| ) |
| x = rearrange(x, "b t c h w -> (b t) c h w") |
| x = self.patch_embedder(x) |
| x = rearrange(x, "(b t) p c -> b (t p) c", b=input_batch_size) |
|
|
| emb = self.noise_level_pos_embedding(noise_levels) |
| emb = repeat(emb, "b t c -> b (t p) c", p=self.num_patches) |
|
|
| if self.conditioning_modeling == "film": |
| if self.conditioning_type == "global": |
| external_cond_emb = repeat( |
| external_cond_emb, "b t c -> b (t p) c", p=self.num_patches |
| ) |
| else: |
| external_cond_emb = rearrange(external_cond_emb, "b t p c -> b (t p) c") |
| emb = emb + external_cond_emb |
|
|
| x = self.dit_base(x, emb) |
| x = self.unpatchify( |
| rearrange(x, "b (t p) c -> (b t) p c", p=self.num_patches) |
| ) |
| x = rearrange( |
| x, "(b t) h w c -> b t c h w", b=input_batch_size |
| ) |
| return x |
|
|
