# coding: utf-8 """ @File : models6/hm_control.py @Author : Songkey @Email : songkey@pku.edu.cn @Date : 8/14/2024 @Desc : """ from typing import Tuple import torch import torch.nn as nn import torch.nn.functional as F from diffusers.models.modeling_utils import ModelMixin from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.embeddings import TimestepEmbedding, Timesteps from einops import rearrange from .hm_blocks import SKCrossAttention, STKAttentionV5, SmallUnet, SmallUnetV5 class HMControlNetBase(ModelMixin, ConfigMixin): @register_to_config def __init__( self, in_channels: int = 3, cross_attention_dim: int = 2048, ): super().__init__() self.emo_embedding = Timesteps(1024, True, 0) self.emo_pre_proj = TimestepEmbedding(548, 16) self.emo_proj = TimestepEmbedding(1024, cross_attention_dim) self.exp_embedding = Timesteps(1024, True, 0) self.exp_pre_proj = TimestepEmbedding(51, 16) self.face_pre_proj = TimestepEmbedding(3, 16) self.exp_proj = TimestepEmbedding(1024, cross_attention_dim) self.face_proj = TimestepEmbedding(1024, cross_attention_dim) self.ctrl_attn = SmallUnet(in_channels=in_channels, out_channels=320, mid_channels=[320, 640, 1280], cross_attention_dim=cross_attention_dim, temporal_attn=False) def forward(self, condition, drive_coeff=None, face_parts=None, emo_embedding=None): bs, video_length = condition.size(0), condition.size(2) condition = rearrange(condition, "b c f h w -> (b f) c h w") condition = F.interpolate(condition, size=(condition.size(2)//8, condition.size(3)//8), mode='bilinear', align_corners=False) condition = rearrange(condition, "(b f) c h w -> b c f h w", f=video_length) if drive_coeff is None: emo_coeff = rearrange(emo_embedding * 20., "b f c -> (b f c)") emo_embedding = self.emo_embedding(emo_coeff) emo_embedding = rearrange(emo_embedding, "(b f c) d -> (b f) d c", b=bs, f=video_length) emo_embedding = self.emo_pre_proj(emo_embedding.to(dtype=condition.dtype)) emo_embedding = rearrange(emo_embedding, "(b f) d c -> (b f) c d", f=video_length) drive_embedding = self.emo_proj(emo_embedding) else: drive_coeff = rearrange(drive_coeff * 500., "b f c -> (b f c)") exp_embedding = self.exp_embedding(drive_coeff) exp_embedding = rearrange(exp_embedding, "(b f c) d -> (b f) d c", b=bs, f=video_length) exp_embedding = self.exp_pre_proj(exp_embedding.to(dtype=condition.dtype)) exp_embedding = rearrange(exp_embedding, "(b f) d c -> (b f) c d", f=video_length) face_parts = rearrange(face_parts, "b f c d -> (b f) d c") face_embedding = self.face_pre_proj(face_parts) face_embedding = rearrange(face_embedding, "(b f) d c -> (b f) c d", f=video_length) drive_embedding = self.exp_proj(exp_embedding) + self.face_proj(face_embedding) drive_embedding = rearrange(drive_embedding, "(b f) c d -> b f c d", f=video_length) return self.ctrl_attn(condition, drive_embedding) class HM5ControlNetBase(ModelMixin, ConfigMixin): @register_to_config def __init__( self, in_channels: int = 3, cross_attention_dim: int = 2048, ): super().__init__() self.emo_embedding = Timesteps(1024, True, 0) self.emo_pre_proj = TimestepEmbedding(548, 16) self.emo_proj = TimestepEmbedding(1024, cross_attention_dim) self.exp_embedding = Timesteps(1024, True, 0) self.exp_pre_proj = TimestepEmbedding(51, 16) self.face_pre_proj = TimestepEmbedding(3, 16) self.exp_proj = TimestepEmbedding(1024, cross_attention_dim) self.face_proj = TimestepEmbedding(1024, cross_attention_dim) self.ctrl_attn = SmallUnetV5(in_channels=in_channels, out_channels=320, mid_channels=[320, 640, 1280], cross_attention_dim=cross_attention_dim, temporal_attn=False) def forward(self, condition, drive_coeff=None, face_parts=None, emo_embedding=None): bs, video_length = condition.size(0), condition.size(2) condition = rearrange(condition, "b c f h w -> (b f) c h w") condition = F.interpolate(condition, size=(condition.size(2)//8, condition.size(3)//8), mode='bilinear', align_corners=False) condition = rearrange(condition, "(b f) c h w -> b c f h w", f=video_length) if drive_coeff is None: emo_coeff = rearrange(emo_embedding * 20., "b f c -> (b f c)") emo_embedding = self.emo_embedding(emo_coeff) emo_embedding = rearrange(emo_embedding, "(b f c) d -> (b f) d c", b=bs, f=video_length) emo_embedding = self.emo_pre_proj(emo_embedding.to(dtype=condition.dtype)) emo_embedding = rearrange(emo_embedding, "(b f) d c -> (b f) c d", f=video_length) drive_embedding = self.emo_proj(emo_embedding) else: drive_coeff = rearrange(drive_coeff * 500., "b f c -> (b f c)") exp_embedding = self.exp_embedding(drive_coeff) exp_embedding = rearrange(exp_embedding, "(b f c) d -> (b f) d c", b=bs, f=video_length) exp_embedding = self.exp_pre_proj(exp_embedding.to(dtype=condition.dtype)) exp_embedding = rearrange(exp_embedding, "(b f) d c -> (b f) c d", f=video_length) face_parts = rearrange(face_parts, "b f c d -> (b f) d c") face_embedding = self.face_pre_proj(face_parts) face_embedding = rearrange(face_embedding, "(b f) d c -> (b f) c d", f=video_length) drive_embedding = self.exp_proj(exp_embedding) + self.face_proj(face_embedding) drive_embedding = rearrange(drive_embedding, "(b f) c d -> b f c d", f=video_length) return self.ctrl_attn(condition, drive_embedding) class HM4SD15ControlProj(ModelMixin, ConfigMixin): @register_to_config def __init__( self, ): super().__init__() self.map640_to_640 = nn.Conv2d(640, 640, kernel_size=3, padding=1, bias=False) self.map640_to_1280 = nn.Conv2d(640, 1280, kernel_size=3, padding=1, bias=False) self.map320_to_320 = nn.Conv2d(320, 320, kernel_size=3, padding=1, bias=False) self.map320_to_640 = nn.Conv2d(320, 640, kernel_size=3, padding=1, bias=False) self.map320_to_1280 = nn.Conv2d(320, 1280, kernel_size=3, padding=1, bias=False) self.map1280_to_1280 = nn.Conv2d(1280, 1280, kernel_size=3, padding=1, bias=False) self.map1280_to_1280_2 = nn.Conv2d(1280, 1280, kernel_size=3, padding=1, bias=False) self.map1280_to_1280_3 = nn.Conv2d(1280, 1280, kernel_size=3, padding=1, bias=False) def forward(self, control_dict): video_length = control_dict['feat_0'].size(2) ret_dict = dict( up3_0=rearrange( self.map1280_to_1280( rearrange(control_dict['feat_0'], "b c f h w -> (b f) c h w"), ), "(b f) c h w -> b c f h w", f=video_length), up3_1=rearrange( self.map640_to_1280( rearrange(control_dict['feat_1'], "b c f h w -> (b f) c h w"), ), "(b f) c h w -> b c f h w", f=video_length), up3_2=rearrange( self.map320_to_1280( rearrange(control_dict['feat_2'], "b c f h w -> (b f) c h w"), ), "(b f) c h w -> b c f h w", f=video_length), up3_3=rearrange( self.map320_to_640( rearrange(control_dict['feat_3'], "b c f h w -> (b f) c h w"), ), "(b f) c h w -> b c f h w", f=video_length), down3_0=rearrange( self.map320_to_320( rearrange(control_dict['feat_2'], "b c f h w -> (b f) c h w"), ), "(b f) c h w -> b c f h w", f=video_length), down3_1=rearrange( self.map640_to_640( rearrange(control_dict['feat_1'], "b c f h w -> (b f) c h w"), ), "(b f) c h w -> b c f h w", f=video_length), down3_2=rearrange( self.map1280_to_1280_2( rearrange(control_dict['feat_0'], "b c f h w -> (b f) c h w"), ), "(b f) c h w -> b c f h w", f=video_length), down3_3=rearrange( self.map1280_to_1280_3( rearrange(control_dict['feat_0'], "b c f h w -> (b f) c h w"), ), "(b f) c h w -> b c f h w", f=video_length), ) return ret_dict class HMProjector(nn.Module): def __init__( self, in_channels: int = 3, min_dim: int = 2048, ): super().__init__() self.proj = nn.Conv2d(in_channels, min_dim, kernel_size=3, padding=1, bias=False) self.attn = STKAttentionV5(in_channels=min_dim, mid_channels=min_dim // 4, temporal_attn=False, time_embed_dim=None) def forward(self, x): f = x.size(2) x = rearrange(x, "b c f h w -> (b f) c h w") x = self.proj(x) x = rearrange(x, "(b f) c h w -> b f h w c", f=f) x = self.attn(x) return rearrange(x, "b f h w c -> b c f h w") class HM5SD15ControlProj(ModelMixin, ConfigMixin): @register_to_config def __init__( self, ): super().__init__() self.map640_to_640 = HMProjector(640, 640) # nn.Conv2d(640, 640, kernel_size=3, padding=1, bias=False) self.map640_to_1280 = HMProjector(640, 1280) # nn.Conv2d(640, 1280, kernel_size=3, padding=1, bias=False) self.map320_to_320 = HMProjector(320, 320) # nn.Conv2d(320, 320, kernel_size=3, padding=1, bias=False) self.map320_to_640 = HMProjector(320, 640) # nn.Conv2d(320, 640, kernel_size=3, padding=1, bias=False) self.map320_to_1280 = HMProjector(320, 1280) # nn.Conv2d(320, 1280, kernel_size=3, padding=1, bias=False) self.map1280_to_1280 = HMProjector(1280, 1280) # nn.Conv2d(1280, 1280, kernel_size=3, padding=1, bias=False) self.map1280_to_1280_2 = HMProjector(1280, 1280) # nn.Conv2d(1280, 1280, kernel_size=3, padding=1, bias=False) self.map1280_to_1280_3 = HMProjector(1280, 1280) # nn.Conv2d(1280, 1280, kernel_size=3, padding=1, bias=False) def forward(self, control_dict): ret_dict = dict( up3_0=self.map1280_to_1280(control_dict['feat_0']), up3_1=self.map640_to_1280(control_dict['feat_1']), up3_2=self.map320_to_1280(control_dict['feat_2']), up3_3=self.map320_to_640(control_dict['feat_3']), down3_0=self.map320_to_320(control_dict['feat_2']), down3_1=self.map640_to_640(control_dict['feat_1']), down3_2=self.map1280_to_1280_2(control_dict['feat_0']), down3_3=self.map1280_to_1280_3(control_dict['feat_0']), ) return ret_dict class HMControlNet(ModelMixin, ConfigMixin): @register_to_config def __init__( self, embedding_channels: int = 1280, input_channels: int = 3, scale_factor: int = 8, cross_attention_dim: int = 320, block_out_channels: Tuple[int] = (128, 320, 640, 1280), ): super().__init__() self.scale_factor = scale_factor self.embedding_channels = embedding_channels self.cross_attention_dim = cross_attention_dim self.conv_in = nn.Conv2d(input_channels, block_out_channels[0], kernel_size=3, padding=1, bias=False) self.exp_embedding = Timesteps(cross_attention_dim, True, 0) self.exp_proj = TimestepEmbedding(cross_attention_dim, cross_attention_dim) self.face_proj = TimestepEmbedding(1024, cross_attention_dim) self.blocks_down = nn.ModuleList([]) for i in range(1, len(block_out_channels)): channel_in = block_out_channels[i-1] channel_out = block_out_channels[i] self.blocks_down.append( SKCrossAttention( channel_in=channel_in, channel_out=channel_out, cross_attention_dim=cross_attention_dim, num_positional_embeddings=64, num_positional_embeddings_hidden=64 ) ) def forward(self, condition, drive_coeff, face_parts): bs, _, video_length, h, w = condition.shape condition = rearrange(condition, "b c f h w -> (b f) c h w") conditioning = F.interpolate(condition, size=(h // self.scale_factor, w // self.scale_factor), mode='bilinear', align_corners=False) embedding = self.conv_in(conditioning) drive_coeff = rearrange(drive_coeff * 500., "b f c -> (b f c)") drive_embedding = self.exp_embedding(drive_coeff).to(dtype=embedding.dtype) drive_embedding = rearrange(drive_embedding, "(b f c) d -> b f c d", b=bs, f=video_length, d=self.cross_attention_dim) face_parts = rearrange(face_parts, "b f c d -> (b f) c d") face_embedding = self.face_proj(face_parts) face_embedding = rearrange(face_embedding, "(b f) c d -> b f c d", f=video_length) drive_embedding = torch.cat([face_embedding, drive_embedding], dim=2) drive_embedding = rearrange(drive_embedding, "b f c d -> (b f) c d") drive_embedding = self.exp_proj(drive_embedding) ret_dict = {} for idx, block in enumerate(self.blocks_down): embedding = block(embedding, drive_embedding) ret_dict[f'down_{idx}'] = rearrange(embedding, "(b f) c h w -> b c f h w", f=video_length) return ret_dict class HMControlNet2(ModelMixin, ConfigMixin): @register_to_config def __init__( self, embedding_channels: int = 1280, input_channels: int = 3, scale_factor: int = 8, cross_attention_dim: int = 320, block_out_channels: Tuple[int] = (128, 320, 640, 1280), ): super().__init__() self.scale_factor = scale_factor self.embedding_channels = embedding_channels self.cross_attention_dim = cross_attention_dim self.conv_in = nn.Conv2d(input_channels, block_out_channels[0], kernel_size=3, padding=1, bias=False) self.exp_embedding = Timesteps(cross_attention_dim, True, 0) self.emo_proj = TimestepEmbedding(cross_attention_dim, cross_attention_dim) self.blocks_down = nn.ModuleList([]) for i in range(1, len(block_out_channels)): channel_in = block_out_channels[i-1] channel_out = block_out_channels[i] self.blocks_down.append( SKCrossAttention( channel_in=channel_in, channel_out=channel_out, cross_attention_dim=cross_attention_dim, num_positional_embeddings=64, num_positional_embeddings_hidden=1024 ) ) def forward(self, condition, emo_embedding): bs, _, video_length, h, w = condition.shape condition = rearrange(condition, "b c f h w -> (b f) c h w") condition = F.interpolate(condition, size=(h // self.scale_factor, w // self.scale_factor), mode='bilinear', align_corners=False) embedding = self.conv_in(condition) drive_coeff = rearrange(emo_embedding * 20., "b f c -> (b f c)") drive_embedding = self.exp_embedding(drive_coeff).to(dtype=embedding.dtype) drive_embedding = rearrange(drive_embedding, "(b f c) d -> b f c d", b=bs, f=video_length, d=self.cross_attention_dim) drive_embedding = rearrange(drive_embedding, "b f c d -> (b f) c d") drive_embedding = self.emo_proj(drive_embedding) ret_dict = {} for idx, block in enumerate(self.blocks_down): embedding = block(embedding, drive_embedding) ret_dict[f'down2_{idx}'] = rearrange(embedding, "(b f) c h w -> b c f h w", f=video_length) return ret_dict class HMV2ControlNet(ModelMixin, ConfigMixin): def __init__( self, embedding_channels: int = 1280, input_channels: int = 3, scale_factor: int = 4, cross_attention_dim: int = 320, block_out_channels: Tuple[int] = (128, 640, 1280, 1280, 1280), ): super().__init__() self.scale_factor = scale_factor self.embedding_channels = embedding_channels self.cross_attention_dim = cross_attention_dim self.conv_in = nn.Conv2d(input_channels, block_out_channels[0], kernel_size=3, padding=1, bias=False) self.exp_embedding = Timesteps(cross_attention_dim, True, 0) self.exp_proj = TimestepEmbedding(cross_attention_dim, cross_attention_dim) self.face_proj = TimestepEmbedding(1024, cross_attention_dim) self.blocks_down = nn.ModuleList([]) for i in range(1, len(block_out_channels)): channel_in = block_out_channels[i - 1] channel_out = block_out_channels[i] self.blocks_down.append( SKCrossAttention( channel_in=channel_in, channel_out=channel_out, cross_attention_dim=cross_attention_dim, num_positional_embeddings=64, num_positional_embeddings_hidden=1024 ) ) def forward(self, condition, drive_coeff, face_parts): bs, _, video_length, h, w = condition.shape conditioning = rearrange(condition, "b c f h w -> (b f) c h w") conditioning = F.interpolate(conditioning, size=(h // self.scale_factor, w // self.scale_factor), mode='bilinear', align_corners=False) embedding = self.conv_in(conditioning) drive_coeff = rearrange(drive_coeff * 500., "b f c -> (b f c)") drive_embedding = self.exp_embedding(drive_coeff).to(dtype=embedding.dtype) drive_embedding = rearrange(drive_embedding, "(b f c) d -> b f c d", b=bs, f=video_length, d=self.cross_attention_dim) face_parts = rearrange(face_parts, "b f c d -> (b f) c d") face_embedding = self.face_proj(face_parts) face_embedding = rearrange(face_embedding, "(b f) c d -> b f c d", f=video_length) drive_embedding = torch.cat([face_embedding, drive_embedding], dim=2) drive_embedding = rearrange(drive_embedding, "b f c d -> (b f) c d") drive_embedding = self.exp_proj(drive_embedding) ret_dict = {} for idx, block in enumerate(self.blocks_down): embedding = block(embedding, drive_embedding) ret_dict[f'up_v2_{len(self.blocks_down) - idx - 1}'] = rearrange(embedding, "(b f) c h w -> b c f h w", f=video_length) return ret_dict class HMV3ControlNet(ModelMixin, ConfigMixin): def __init__( self, embedding_channels: int = 1280, input_channels: int = 3, scale_factor: int = 4, cross_attention_dim: int = 320, block_out_channels: Tuple[int] = (128, 640, 1280, 1280, 1280), ): super().__init__() self.scale_factor = scale_factor self.embedding_channels = embedding_channels self.cross_attention_dim = cross_attention_dim self.conv_in = nn.Conv2d(input_channels, block_out_channels[0], kernel_size=3, padding=1, bias=False) self.conv_up2_down0 = nn.Conv2d(1280, 320, kernel_size=1, padding=0, bias=False) self.conv_up1_down1 = nn.Conv2d(1280, 640, kernel_size=1, padding=0, bias=False) self.emo_embedding = Timesteps(cross_attention_dim, True, 0) self.exp_embedding = Timesteps(cross_attention_dim, True, 0) self.exp_proj = TimestepEmbedding(cross_attention_dim, cross_attention_dim) self.face_proj = TimestepEmbedding(1024, cross_attention_dim) self.emo_proj = TimestepEmbedding(cross_attention_dim, cross_attention_dim) self.blocks_down = nn.ModuleList([]) for i in range(1, len(block_out_channels)): channel_in = block_out_channels[i-1] channel_out = block_out_channels[i] self.blocks_down.append( SKCrossAttention( channel_in=channel_in, channel_out=channel_out, cross_attention_dim=cross_attention_dim, num_positional_embeddings=64, num_positional_embeddings_hidden=1024 ) ) def forward(self, condition, drive_coeff=None, face_parts=None, emo_embedding=None): bs, _, video_length, h, w = condition.shape conditioning = rearrange(condition, "b c f h w -> (b f) c h w") conditioning = F.interpolate(conditioning, size=(h // self.scale_factor, w // self.scale_factor), mode='bilinear', align_corners=False) embedding = self.conv_in(conditioning) if drive_coeff is None or face_parts is None: emo_coeff = rearrange(emo_embedding * 20., "b f c -> (b f c)") emo_embedding = self.emo_embedding(emo_coeff).to(dtype=embedding.dtype) emo_embedding = rearrange(emo_embedding, "(b f c) d -> b f c d", b=bs, f=video_length, d=self.cross_attention_dim) emo_embedding = rearrange(emo_embedding, "b f c d -> (b f) c d") drive_embedding = self.emo_proj(emo_embedding) else: drive_coeff = rearrange(drive_coeff * 500., "b f c -> (b f c)") drive_embedding = self.exp_embedding(drive_coeff).to(dtype=embedding.dtype) drive_embedding = rearrange(drive_embedding, "(b f c) d -> b f c d", b=bs, f=video_length, d=self.cross_attention_dim) face_parts = rearrange(face_parts, "b f c d -> (b f) c d") face_embedding = self.face_proj(face_parts) face_embedding = rearrange(face_embedding, "(b f) c d -> b f c d", f=video_length) drive_embedding = torch.cat([face_embedding, drive_embedding], dim=2) drive_embedding = rearrange(drive_embedding, "b f c d -> (b f) c d") drive_embedding = self.exp_proj(drive_embedding) ret_dict = {} for idx, block in enumerate(self.blocks_down): embedding = block(embedding, drive_embedding) up_idx = len(self.blocks_down)-idx-1 ret_dict[f'up3_{up_idx}'] = rearrange(embedding, "(b f) c h w -> b c f h w", f=video_length) if up_idx == 2: ret_dict[f'down3_0'] = rearrange(self.conv_up2_down0(embedding), "(b f) c h w -> b c f h w", f=video_length) if up_idx == 1: ret_dict[f'down3_1'] = rearrange(self.conv_up1_down1(embedding), "(b f) c h w -> b c f h w", f=video_length) if up_idx == 0: ret_dict[f'down3_2'] = ret_dict[f'up3_0'] ret_dict[f'down3_3'] = ret_dict[f'up3_0'] return ret_dict class HMV2ControlNet2(ModelMixin, ConfigMixin): def __init__( self, embedding_channels: int = 1280, input_channels: int = 3, scale_factor: int = 4, cross_attention_dim: int = 320, block_out_channels: Tuple[int] = (128, 640, 1280, 1280, 1280), ): super().__init__() self.scale_factor = scale_factor self.embedding_channels = embedding_channels self.cross_attention_dim = cross_attention_dim self.conv_in = nn.Conv2d(input_channels, block_out_channels[0], kernel_size=3, padding=1, bias=False) self.exp_embedding = Timesteps(cross_attention_dim, True, 0) self.emo_proj = TimestepEmbedding(cross_attention_dim, cross_attention_dim) # self.face_proj = TimestepEmbedding(1024, cross_attention_dim) self.blocks_down = nn.ModuleList([]) for i in range(1, len(block_out_channels)): channel_in = block_out_channels[i-1] channel_out = block_out_channels[i] self.blocks_down.append( SKCrossAttention( channel_in=channel_in, channel_out=channel_out, cross_attention_dim=cross_attention_dim, num_positional_embeddings=64, num_positional_embeddings_hidden=1024 ) ) def forward(self, condition, emo_embedding): bs, _, video_length, h, w = condition.shape conditioning = rearrange(condition, "b c f h w -> (b f) c h w") conditioning = F.interpolate(conditioning, size=(h // self.scale_factor, w // self.scale_factor), mode='bilinear', align_corners=False) embedding = self.conv_in(conditioning) drive_coeff = rearrange(emo_embedding * 20., "b f c -> (b f c)") drive_embedding = self.exp_embedding(drive_coeff).to(dtype=embedding.dtype) drive_embedding = rearrange(drive_embedding, "(b f c) d -> b f c d", b=bs, f=video_length, d=self.cross_attention_dim) drive_embedding = rearrange(drive_embedding, "b f c d -> (b f) c d") drive_embedding = self.emo_proj(drive_embedding) ret_dict = {} for idx, block in enumerate(self.blocks_down): embedding = block(embedding, drive_embedding) ret_dict[f'up2_v2_{len(self.blocks_down)-idx-1}'] = rearrange(embedding, "(b f) c h w -> b c f h w", f=video_length) return ret_dict