File size: 3,931 Bytes
59630ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | 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) # (B, N, C)
x = self.unpatchify(
rearrange(x, "b (t p) c -> (b t) p c", p=self.num_patches)
) # (B * T, H, W, C)
x = rearrange(
x, "(b t) h w c -> b t c h w", b=input_batch_size
) # (B, T, C, H, W)
return x
|