| 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
|
|
|
