271 lines
12 KiB
Python
271 lines
12 KiB
Python
# 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
|
|
|
|
|
|
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 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 |