| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from einops import rearrange |
| import numpy as np |
| from typing import Tuple |
|
|
| from .unet_causal_3d_blocks import get_down_block3d, CausalConv3d |
|
|
| class ControlNetCausalConditioningEmbedding(nn.Module): |
| def __init__(self, conditioning_embedding_channels: int, conditioning_channels: int = 3, block_out_channels: Tuple[int, ...] = (16, 32, 96, 256)): |
| super().__init__() |
| self.conv_in = CausalConv3d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1) |
| self.blocks = nn.ModuleList([]) |
|
|
| for i in range(len(block_out_channels) - 1): |
| channel_in = block_out_channels[i] |
| channel_out = block_out_channels[i + 1] |
| self.blocks.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1)) |
| self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2)) |
|
|
| self.conv_out = nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1) |
|
|
| def forward(self, conditioning): |
| embedding = self.conv_in(conditioning) |
| embedding = F.silu(embedding) |
|
|
| for block in self.blocks: |
| embedding = block(embedding) |
| embedding = F.silu(embedding) |
|
|
| embedding = self.conv_out(embedding) |
|
|
| return embedding |
|
|
| class MiniHunyuanEncoder(nn.Module): |
| ''' |
| a direct copy of hunyuan encoder |
| ''' |
| def __init__( |
| self, |
| in_channels = 3, |
| out_channels = 3, |
| down_block_types = ['DownEncoderBlockCausal3D', 'DownEncoderBlockCausal3D', 'DownEncoderBlockCausal3D', 'DownEncoderBlockCausal3D'], |
| block_out_channels = [128, 256, 512, 512], |
| layers_per_block = 2, |
| norm_num_groups = 32, |
| act_fn: str = "silu", |
| time_compression_ratio: int = 4, |
| spatial_compression_ratio: int = 8, |
| ): |
| super().__init__() |
| self.layers_per_block = layers_per_block |
| self.conv_in = CausalConv3d( |
| in_channels, block_out_channels[0], kernel_size=3, stride=1) |
| self.mid_block = None |
| self.down_blocks = nn.ModuleList([]) |
|
|
| |
| output_channel = block_out_channels[0] |
| for i, down_block_type in enumerate(down_block_types): |
| input_channel = output_channel |
| output_channel = block_out_channels[i] |
| is_final_block = i == len(block_out_channels) - 1 |
| num_spatial_downsample_layers = int( |
| np.log2(spatial_compression_ratio)) |
| num_time_downsample_layers = int(np.log2(time_compression_ratio)) |
|
|
| if time_compression_ratio == 4: |
| add_spatial_downsample = bool( |
| i < num_spatial_downsample_layers) |
| add_time_downsample = bool(i >= ( |
| len(block_out_channels) - 1 - num_time_downsample_layers) and not is_final_block) |
| elif time_compression_ratio == 8: |
| add_spatial_downsample = bool( |
| i < num_spatial_downsample_layers) |
| add_time_downsample = bool(i < num_time_downsample_layers) |
| else: |
| raise ValueError( |
| f"Unsupported time_compression_ratio: {time_compression_ratio}") |
|
|
| downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) |
| downsample_stride_T = (2, ) if add_time_downsample else (1, ) |
| downsample_stride = tuple( |
| downsample_stride_T + downsample_stride_HW) |
| down_block = get_down_block3d( |
| down_block_type, |
| num_layers=self.layers_per_block, |
| in_channels=input_channel, |
| out_channels=output_channel, |
| add_downsample=bool( |
| add_spatial_downsample or add_time_downsample), |
| downsample_stride=downsample_stride, |
| resnet_eps=1e-6, |
| downsample_padding=0, |
| resnet_act_fn=act_fn, |
| resnet_groups=norm_num_groups, |
| ) |
| self.down_blocks.append(down_block) |
|
|
| self.conv_out = CausalConv3d(block_out_channels[-1], out_channels, kernel_size=3) |
|
|
| def forward(self, sample): |
| assert len(sample.shape) == 5, "The input tensor should have 5 dimensions" |
| sample = self.conv_in(sample) |
| |
| for down_block in self.down_blocks: |
| sample = down_block(sample) |
| sample = self.conv_out(sample) |
| return sample |
|
|
|
|
| class ControlNetConditioningEmbedding(nn.Module): |
| """ |
| Quoting from https://arxiv.org/abs/2302.05543: "Stable Diffusion uses a pre-processing method similar to VQ-GAN |
| [11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized |
| training. This requires ControlNets to convert image-based conditions to 64 × 64 feature space to match the |
| convolution size. We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides |
| (activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full |
| model) to encode image-space conditions ... into feature maps ..." |
| """ |
|
|
| def __init__( |
| self, |
| conditioning_embedding_channels: int, |
| conditioning_channels: int = 3, |
| block_out_channels: Tuple[int, ...] = (16, 32, 96, 256), |
| ): |
| super().__init__() |
|
|
| self.conv_in = nn.Conv2d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1) |
|
|
| self.blocks = nn.ModuleList([]) |
|
|
| for i in range(len(block_out_channels) - 1): |
| channel_in = block_out_channels[i] |
| channel_out = block_out_channels[i + 1] |
| self.blocks.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1)) |
| self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2)) |
|
|
| self.conv_out = nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1) |
|
|
| def forward(self, conditioning): |
| embedding = self.conv_in(conditioning) |
| embedding = F.silu(embedding) |
|
|
| for block in self.blocks: |
| embedding = block(embedding) |
| embedding = F.silu(embedding) |
|
|
| embedding = self.conv_out(embedding) |
|
|
| return embedding |
|
|
|
|
| class InflatedGroupNorm(nn.GroupNorm): |
| def forward(self, x): |
| video_length = x.shape[2] |
|
|
| x = rearrange(x, "b c f h w -> (b f) c h w") |
| x = super().forward(x) |
| x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length) |
|
|
| return x |
|
|
| class InflatedConv3d(nn.Conv2d): |
| def forward(self, x): |
| video_length = x.shape[2] |
|
|
| x = rearrange(x, "b c f h w -> (b f) c h w") |
| x = super().forward(x) |
| x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length) |
|
|
| return x |
|
|
|
|
| class ResnetBlockInflated(nn.Module): |
| def __init__(self, *, in_channels, out_channels=None, dropout=0.0, groups=32, groups_out=None, pre_norm=True, eps=1e-6, non_linearity="swish", output_scale_factor=1.0): |
| super().__init__() |
| self.pre_norm = pre_norm |
| self.pre_norm = True |
| self.in_channels = in_channels |
| out_channels = in_channels if out_channels is None else out_channels |
| self.out_channels = out_channels |
| self.output_scale_factor = output_scale_factor |
|
|
| if groups_out is None: |
| groups_out = groups |
|
|
| self.norm1 = InflatedGroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) |
| self.conv1 = InflatedConv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) |
| self.norm2 = InflatedGroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True) |
| self.dropout = torch.nn.Dropout(dropout) |
| self.conv2 = InflatedConv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) |
|
|
| if non_linearity == "swish": |
| self.nonlinearity = lambda x: F.silu(x) |
| elif non_linearity == "silu": |
| self.nonlinearity = nn.SiLU() |
|
|
| def forward(self, input_tensor, temb): |
| if temb is not None: |
| print("Warning: temb is None in ResnetBlockInflated") |
| hidden_states = input_tensor |
|
|
| hidden_states = self.norm1(hidden_states) |
| hidden_states = self.nonlinearity(hidden_states) |
|
|
| hidden_states = self.conv1(hidden_states) |
|
|
| if temb is not None: |
| hidden_states = hidden_states + temb |
|
|
| hidden_states = self.norm2(hidden_states) |
| hidden_states = self.nonlinearity(hidden_states) |
| hidden_states = self.dropout(hidden_states) |
| hidden_states = self.conv2(hidden_states) |
|
|
| output_tensor = (input_tensor + hidden_states) / self.output_scale_factor |
|
|
| return output_tensor |
|
|
| class DownEncoderBlockInflated(nn.Module): |
| def __init__(self, *, num_layers: int, in_channels: int, out_channels: int, add_downsample: bool, downsample_stride: tuple = (1, 2, 2), |
| resnet_eps: float = 1e-6, resnet_act_fn: str = "silu", resnet_groups: int = 32): |
| super().__init__() |
|
|
| self.resnets = nn.ModuleList([ResnetBlockInflated( |
| in_channels=in_channels if i == 0 else out_channels, |
| out_channels=out_channels, |
| eps=resnet_eps, |
| non_linearity=resnet_act_fn, |
| groups=resnet_groups, |
| ) for i in range(num_layers)]) |
|
|
| self.downsamplers = nn.ModuleList() |
| if add_downsample: |
| self.downsamplers.append( |
| InflatedConv3d( |
| out_channels, |
| out_channels, |
| kernel_size=3, |
| stride=2, |
| padding=1, |
| ) |
| ) |
| self.down_stride = downsample_stride |
| else: |
| self.down_stride = (1, 1, 1) |
|
|
| def forward(self, x, temb=None): |
| for resnet in self.resnets: |
| x = resnet(x, temb) |
|
|
| for down in self.downsamplers: |
| x = down(x) |
| return x |
|
|
|
|
| class SFT(nn.Module): |
| def __init__( |
| self, in_channels, out_channels, intermediate_channels=128, groups=32, eps=1e-6): |
| super().__init__() |
| self.out_channels = out_channels |
| self.norm = InflatedGroupNorm(groups, out_channels, eps, affine=True) |
| self.mlp_shared = nn.Sequential(InflatedConv3d(in_channels, intermediate_channels, kernel_size=3, stride=1, padding=1), nn.SiLU()) |
| self.mlp_gamma = InflatedConv3d(intermediate_channels, out_channels, kernel_size=3, stride=1, padding=1) |
| self.mlp_beta = InflatedConv3d(intermediate_channels, out_channels, kernel_size=3, stride=1, padding=1) |
|
|
| def forward(self, hidden_state, condition): |
| """ |
| hidden_state : (B, Cout, T, H, W) |
| condition : (B, Cin, 1, H, W) |
| """ |
| hidden_state = self.norm(hidden_state) |
|
|
| actv = self.mlp_shared(condition) |
| gamma = self.mlp_gamma(actv) |
| beta = self.mlp_beta(actv) |
|
|
| return torch.addcmul(beta, hidden_state, 1 + gamma) |
|
|
|
|
| class MiniEncoder2D(nn.Module): |
|
|
| def __init__( |
| self, |
| in_channels: int = 3, |
| out_channels: int = 3, |
| down_block_types: list = ( |
| "DownEncoderBlockInflated", |
| "DownEncoderBlockInflated", |
| "DownEncoderBlockInflated", |
| "DownEncoderBlockInflated", |
| ), |
| block_out_channels: list = (128, 256, 512, 512), |
| layers_per_block: int = 2, |
| norm_num_groups: int = 32, |
| act_fn: str = "silu", |
| spatial_compression_ratio: int = 8, |
| ): |
| super().__init__() |
|
|
| |
| |
| |
| self.conv_in = InflatedConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) |
|
|
| self.down_blocks = nn.ModuleList() |
| output_channel = block_out_channels[0] |
| num_spatial_down_layers = int(np.log2(spatial_compression_ratio)) |
|
|
| for i, block_type in enumerate(down_block_types): |
| input_channel = output_channel |
| output_channel = block_out_channels[i] |
| |
|
|
| add_spatial_downsample = bool(i < num_spatial_down_layers) |
|
|
| downsample_stride = (1, 2, 2) if add_spatial_downsample else (1, 1, 1) |
|
|
| down_block = DownEncoderBlockInflated( |
| num_layers=layers_per_block, |
| in_channels=input_channel, |
| out_channels=output_channel, |
| add_downsample=add_spatial_downsample, |
| downsample_stride=downsample_stride, |
| resnet_eps=1e-6, |
| resnet_act_fn=act_fn, |
| resnet_groups=norm_num_groups, |
| ) |
| self.down_blocks.append(down_block) |
|
|
| self.conv_out = InflatedConv3d(output_channel, out_channels, kernel_size=3, stride=1, padding=1) |
|
|
| def forward(self, x): |
| |
| x = self.conv_in(x) |
|
|
| for block in self.down_blocks: |
| x = block(x) |
|
|
| return self.conv_out(x) |
|
|
|
|
| class Driven_Ref_PoseEncoder(nn.Module): |
| def __init__( |
| self, in_channels = 3, out_channels = 3, |
| down_block_types = ['DownEncoderBlockCausal3D', 'DownEncoderBlockCausal3D', 'DownEncoderBlockCausal3D', 'DownEncoderBlockCausal3D'], |
| block_out_channels = [128, 256, 512, 512], layers_per_block = 2, norm_num_groups = 32, |
| act_fn: str = "silu", time_compression_ratio: int = 4, spatial_compression_ratio: int = 8, |
| ): |
| super().__init__() |
| self.layers_per_block = layers_per_block |
| self.conv_in = CausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1) |
| self.mid_block = None |
| self.down_blocks = nn.ModuleList([]) |
|
|
| |
| output_channel = block_out_channels[0] |
| for i, down_block_type in enumerate(down_block_types): |
| input_channel = output_channel |
| output_channel = block_out_channels[i] |
| is_final_block = i == len(block_out_channels) - 1 |
| num_spatial_downsample_layers = int( |
| np.log2(spatial_compression_ratio)) |
| num_time_downsample_layers = int(np.log2(time_compression_ratio)) |
|
|
| if time_compression_ratio == 4: |
| add_spatial_downsample = bool( |
| i < num_spatial_downsample_layers) |
| add_time_downsample = bool(i >= ( |
| len(block_out_channels) - 1 - num_time_downsample_layers) and not is_final_block) |
| elif time_compression_ratio == 8: |
| add_spatial_downsample = bool( |
| i < num_spatial_downsample_layers) |
| add_time_downsample = bool(i < num_time_downsample_layers) |
| else: |
| raise ValueError( |
| f"Unsupported time_compression_ratio: {time_compression_ratio}") |
|
|
| downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) |
| downsample_stride_T = (2, ) if add_time_downsample else (1, ) |
| downsample_stride = tuple( |
| downsample_stride_T + downsample_stride_HW) |
| down_block = get_down_block3d( |
| down_block_type, |
| num_layers=self.layers_per_block, |
| in_channels=input_channel, |
| out_channels=output_channel, |
| add_downsample=bool( |
| add_spatial_downsample or add_time_downsample), |
| downsample_stride=downsample_stride, |
| resnet_eps=1e-6, |
| downsample_padding=0, |
| resnet_act_fn=act_fn, |
| resnet_groups=norm_num_groups, |
| attention_head_dim=output_channel, |
| ) |
| self.down_blocks.append(down_block) |
|
|
| self.conv_out = CausalConv3d(block_out_channels[-1], out_channels, kernel_size=3) |
|
|
| self.ref_pose_encoder = MiniEncoder2D( |
| in_channels = in_channels, |
| out_channels = out_channels, |
| block_out_channels = block_out_channels, |
| norm_num_groups = norm_num_groups, |
| layers_per_block = layers_per_block, |
| spatial_compression_ratio = spatial_compression_ratio, |
| ) |
| self.sft_layers = nn.ModuleList() |
| for i, ch in enumerate(block_out_channels): |
| if i == 0: |
| self.sft_layers.append(None) |
| else: |
| self.sft_layers.append( |
| SFT( |
| in_channels=ch, |
| out_channels=ch, |
| intermediate_channels=max(8, ch // 2), |
| groups=norm_num_groups, |
| ) |
| ) |
|
|
| def forward(self, driven_pose, ref_pose): |
| |
| |
| ref_pose_cond, ref_feats = self.ref_pose_encoder(ref_pose) |
|
|
| x = self.conv_in(driven_pose) |
| for i, down_block in enumerate(self.down_blocks): |
| x = down_block(x) |
|
|
| if self.sft_layers[i] is not None: |
| cond_feat = ref_feats[i] |
| x = self.sft_layers[i](x, cond_feat) |
|
|
| driven_pose_cond = self.conv_out(x) |
| return driven_pose_cond, ref_pose_cond |
|
|
