338 lines
12 KiB
Python
338 lines
12 KiB
Python
import os, math, torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from torchdiffeq import odeint
|
|
from ..basemodel import BaseModel
|
|
|
|
from timm.layers import use_fused_attn
|
|
from timm.models.vision_transformer import Mlp
|
|
|
|
|
|
def enc_dec_mask(T, S, frame_width = 1, expansion = 2):
|
|
mask = torch.ones(T, S)
|
|
for i in range(T):
|
|
mask[i, max(0, (i - expansion) * frame_width):(i + expansion + 1) * frame_width] = 0
|
|
return mask == 1
|
|
|
|
|
|
def get_sinusoid_encoding_table(n_position, d_hid, padding_idx=None):
|
|
"""
|
|
Sinusoidal position encoding table.
|
|
Args:
|
|
n_position (int): the length of the input sequence
|
|
d_hid (int): the dimension of the hidden state
|
|
"""
|
|
def cal_angle(position, hid_idx):
|
|
return position / (10000 ** (2 * (hid_idx // 2) / d_hid))
|
|
|
|
def get_posi_angle_vec(position):
|
|
return [cal_angle(position, hid_j) for hid_j in range(d_hid)]
|
|
|
|
sinusoid_table = torch.Tensor([get_posi_angle_vec(pos_i) for pos_i in range(n_position)])
|
|
sinusoid_table[:, 0::2] = torch.sin(sinusoid_table[:, 0::2]) # dim 2i
|
|
sinusoid_table[:, 1::2] = torch.cos(sinusoid_table[:, 1::2]) # dim 2i+1
|
|
if padding_idx is not None: sinusoid_table[padding_idx] = 0.
|
|
return sinusoid_table
|
|
|
|
|
|
class Attention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_heads: int = 8,
|
|
qkv_bias: bool = False,
|
|
qk_norm: bool = False,
|
|
attn_drop: float = 0.,
|
|
proj_drop: float = 0.,
|
|
norm_layer: nn.Module = nn.LayerNorm,
|
|
) -> None:
|
|
|
|
super().__init__()
|
|
assert dim % num_heads == 0, 'dim should be divisible by num_heads'
|
|
self.num_heads = num_heads
|
|
self.head_dim = dim // num_heads
|
|
self.scale = self.head_dim ** -0.5
|
|
self.fused_attn = use_fused_attn()
|
|
|
|
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
|
self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
|
|
self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
|
|
self.attn_drop = nn.Dropout(attn_drop)
|
|
self.proj = nn.Linear(dim, dim)
|
|
self.proj_drop = nn.Dropout(proj_drop)
|
|
|
|
def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
|
|
B, N, C = x.shape
|
|
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
|
|
q, k, v = qkv.unbind(0)
|
|
q, k = self.q_norm(q), self.k_norm(k)
|
|
|
|
if self.fused_attn:
|
|
x = F.scaled_dot_product_attention(
|
|
q, k, v,
|
|
attn_mask = ~mask,
|
|
dropout_p=self.attn_drop.p if self.training else 0.,
|
|
)
|
|
else:
|
|
q = q * self.scale
|
|
attn = q @ k.transpose(-2, -1)
|
|
attn = attn.softmax(dim=-1)
|
|
attn = self.attn_drop(attn)
|
|
x = attn @ v
|
|
|
|
x = x.transpose(1, 2).reshape(B, N, C)
|
|
x = self.proj(x)
|
|
x = self.proj_drop(x)
|
|
return x
|
|
|
|
class TimestepEmbedder(nn.Module):
|
|
"""
|
|
Embeds scalar timesteps into vector representations.
|
|
"""
|
|
def __init__(self, hidden_size, frequency_embedding_size = 256):
|
|
super().__init__()
|
|
self.mlp = nn.Sequential(
|
|
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_size, hidden_size, bias=True),
|
|
)
|
|
self.frequency_embedding_size = frequency_embedding_size
|
|
|
|
@staticmethod
|
|
def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
|
|
"""
|
|
Create sinusoidal timestep embeddings.
|
|
:param t: a 1-D Tensor of N indices, one per batch element.
|
|
These may be fractional.
|
|
:param dim: the dimension of the output.
|
|
:param max_period: controls the minimum frequency of the embeddings.
|
|
:return: an (N, D) Tensor of positional embeddings.
|
|
"""
|
|
# https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
|
|
half = dim // 2
|
|
freqs = torch.exp(
|
|
-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
|
|
).to(device=t.device)
|
|
args = t[:, None].float() * freqs[None]
|
|
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
|
if dim % 2:
|
|
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
|
return embedding
|
|
|
|
def forward(self, t: torch.Tensor) -> torch.Tensor:
|
|
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
|
|
t_emb = self.mlp(t_freq)
|
|
return t_emb
|
|
|
|
class SequenceEmbed(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim_w,
|
|
dim_h,
|
|
norm_layer=None,
|
|
bias=True,
|
|
):
|
|
super().__init__()
|
|
|
|
self.proj = nn.Linear(dim_w, dim_h, bias=bias)
|
|
self.norm = norm_layer(dim_h) if norm_layer else nn.Identity()
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.norm(self.proj(x))
|
|
|
|
|
|
class FMTBlock(nn.Module):
|
|
"""
|
|
A FMT block inspried by DiT Block
|
|
"""
|
|
def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs) -> None:
|
|
super().__init__()
|
|
self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
|
self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)
|
|
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
|
mlp_hidden_dim = int(hidden_size * mlp_ratio)
|
|
approx_gelu = lambda: nn.GELU(approximate="tanh")
|
|
self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
|
|
self.adaLN_modulation = nn.Sequential(
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_size, 6 * hidden_size, bias=True)
|
|
)
|
|
|
|
def framewise_modulate(self, x, shift, scale) -> torch.Tensor:
|
|
return x * (1 + scale) + shift
|
|
|
|
def forward(self, x, c, mask=None) -> torch.Tensor:
|
|
assert mask is not None
|
|
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=-1)
|
|
x = x + gate_msa * self.attn(self.framewise_modulate(self.norm1(x), shift_msa, scale_msa), mask = mask)
|
|
x = x + gate_mlp * self.mlp(self.framewise_modulate(self.norm2(x), shift_mlp, scale_mlp))
|
|
return x
|
|
|
|
class Decoder(nn.Module):
|
|
"""
|
|
The final decoder of FlowMatchingTransformer.
|
|
"""
|
|
def __init__(self, hidden_size, dim_w):
|
|
super().__init__()
|
|
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
|
self.adaLN_modulation = nn.Sequential(
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_size, 2 * hidden_size, bias=True)
|
|
)
|
|
self.linear = nn.Linear(hidden_size, dim_w, bias=True)
|
|
|
|
def framewise_modulate(self, x, shift, scale) -> torch.Tensor:
|
|
return x * (1 + scale) + shift
|
|
|
|
def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
|
|
shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
|
|
x = self.framewise_modulate(self.norm_final(x), shift, scale)
|
|
return self.linear(x)
|
|
|
|
|
|
class FlowMatchingTransformer(BaseModel):
|
|
"""
|
|
Flow Matching Transformer (FMT)
|
|
"""
|
|
def __init__(self, opt) -> None:
|
|
super().__init__()
|
|
self.opt = opt
|
|
|
|
self.num_frames_for_clip = int(self.opt.wav2vec_sec * self.opt.fps)
|
|
self.num_prev_frames = int(opt.num_prev_frames)
|
|
self.num_total_frames = self.num_prev_frames + self.num_frames_for_clip
|
|
|
|
self.hidden_size = opt.dim_h
|
|
self.mlp_ratio = opt.mlp_ratio
|
|
self.fmt_depth = opt.fmt_depth
|
|
self.num_heads = opt.num_heads
|
|
|
|
self.x_embedder = SequenceEmbed(opt.dim_w, self.hidden_size)
|
|
|
|
# video time position encoding
|
|
self.pos_embed = nn.Parameter(torch.zeros(1, self.num_total_frames, self.hidden_size), requires_grad=False)
|
|
|
|
# flow trajectory time encoding
|
|
self.t_embedder = TimestepEmbedder(self.hidden_size)
|
|
self.c_embedder = nn.Linear(opt.dim_w + opt.dim_a + opt.dim_e, self.hidden_size)
|
|
|
|
# define FMT blocks
|
|
self.blocks = nn.ModuleList([FMTBlock(self.hidden_size, self.num_heads, mlp_ratio=self.mlp_ratio) for _ in range(self.fmt_depth)])
|
|
self.decoder = Decoder(self.hidden_size, self.opt.dim_w)
|
|
self.initialize_weights()
|
|
|
|
# define alignment mask
|
|
alignment_mask = enc_dec_mask(self.num_total_frames, self.num_total_frames, 1, expansion=opt.attention_window).to(opt.rank)
|
|
self.register_buffer('alignment_mask', alignment_mask)
|
|
|
|
|
|
def initialize_weights(self) -> None:
|
|
def _basic_init(module):
|
|
if isinstance(module, nn.Linear):
|
|
torch.nn.init.xavier_uniform_(module.weight)
|
|
if module.bias is not None:
|
|
nn.init.constant_(module.bias, 0)
|
|
|
|
self.apply(_basic_init)
|
|
|
|
pos_embed = get_sinusoid_encoding_table(self.num_total_frames, self.hidden_size)
|
|
self.pos_embed.data.copy_(pos_embed.unsqueeze(0))
|
|
|
|
w = self.x_embedder.proj.weight.data
|
|
nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
|
|
nn.init.constant_(self.x_embedder.proj.bias, 0)
|
|
|
|
# Initialize timestep embedding MLP:
|
|
nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
|
|
nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
|
|
|
|
# Zero-out adaLN modulation layers in FMT blocks:
|
|
for block in self.blocks:
|
|
nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
|
|
nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
|
|
|
|
# Zero-out output layers:
|
|
nn.init.constant_(self.decoder.adaLN_modulation[-1].weight, 0)
|
|
nn.init.constant_(self.decoder.adaLN_modulation[-1].bias, 0)
|
|
nn.init.constant_(self.decoder.linear.weight, 0)
|
|
nn.init.constant_(self.decoder.linear.bias, 0)
|
|
|
|
def sequence_embedder(self, sequence, dropout_prob, train=False) -> torch.Tensor:
|
|
if train:
|
|
batch_id_for_drop = torch.where(torch.rand(sequence.shape[0], device=sequence.device) < dropout_prob)
|
|
sequence[batch_id_for_drop] = 0
|
|
return sequence
|
|
|
|
|
|
def forward(self, t, x, wa, wr, we, prev_x = None, prev_wa = None, train = True, **kwargs) -> torch.Tensor:
|
|
"""
|
|
Forward pass of ConditionalFlowMatchingTransformer.
|
|
|
|
t: (B,) tensor of diffusion timesteps [0, 1]
|
|
x: (B, L, 512) : tensor of sequence of motion latent
|
|
|
|
wa: (B, L, 512) / tensor sequence of wa latent
|
|
wp: (B, L, 6) / tensor sequence of wp latent
|
|
wr: (B, 512) / tensor of reference motion latent (i.e., r -> s)
|
|
we: (B, 1, 7) / tensor of emotion latent
|
|
|
|
prev_x: (B, L', 512) / previous x for auto-regressive generation
|
|
prev_wa: (B, L', 512) / previous audio for auto-regressive generation
|
|
"""
|
|
|
|
# time encoding
|
|
t = self.t_embedder(t).unsqueeze(1) # (N, D)
|
|
|
|
# condition encoding
|
|
wa = self.sequence_embedder(wa, dropout_prob = self.opt.audio_dropout_prob, train=train)
|
|
wr = self.sequence_embedder(wr.unsqueeze(1), dropout_prob = self.opt.ref_dropout_prob, train=train)
|
|
we = self.sequence_embedder(we, dropout_prob = self.opt.emotion_dropout_prob, train=train)
|
|
|
|
# previous condition encoding
|
|
if prev_x is not None:
|
|
prev_x = self.sequence_embedder(prev_x, dropout_prob=0.5, train=train)
|
|
prev_wa = self.sequence_embedder(prev_wa, dropout_prob=0.5, train=train)
|
|
|
|
x = torch.cat([prev_x, x], dim=1)
|
|
wa = torch.cat([prev_wa, wa], dim=1)
|
|
|
|
x = self.x_embedder(x)
|
|
x = x + self.pos_embed # (N, L + L', D), where T = opt.wav2vec_sec * opt.fps, D = dim_w
|
|
|
|
wr = wr.repeat(1, wa.shape[1], 1)
|
|
we = we.repeat(1, wa.shape[1], 1)
|
|
|
|
c = torch.cat([wr, wa, we], dim=-1)
|
|
c = self.c_embedder(c)
|
|
c = t + c
|
|
|
|
# forwarding FMT Blocks
|
|
for block in self.blocks:
|
|
x = block(x, c, self.alignment_mask) # (N, T, D)
|
|
return self.decoder(x, c)
|
|
|
|
@torch.no_grad()
|
|
def forward_with_cfv(self, t, x, wa, wr, we, prev_x, prev_wa, a_cfg_scale=1.0, r_cfg_scale=1.0, e_cfg_scale=1.0, **kwargs) -> torch.Tensor:
|
|
if a_cfg_scale != 1.0 or r_cfg_scale != 1.0 or e_cfg_scale != 1.0:
|
|
null_wa = torch.zeros_like(wa)
|
|
null_we = torch.zeros_like(we)
|
|
null_wr = torch.zeros_like(wr)
|
|
|
|
audio_cat = torch.cat([null_wa, wa, wa], dim=0) # concat along batch
|
|
ref_cat = torch.cat([wr, wr, wr], dim=0) # concat along batch
|
|
emotion_cat = torch.cat([null_we, we, null_we], dim=0) # concat along batch
|
|
x = torch.cat([x, x, x], dim=0) # concat along batch
|
|
|
|
prev_x_cat = torch.cat([prev_x, prev_x, prev_x], dim=0)
|
|
prev_wa_cat = torch.cat([prev_wa, prev_wa, prev_wa], dim=0)
|
|
|
|
model_output = self.forward(t, x, audio_cat, ref_cat, emotion_cat, prev_x_cat, prev_wa_cat, train=False)
|
|
uncond, all_cond, audio_uncond_emotion = torch.chunk(model_output, chunks=3, dim=0)
|
|
|
|
# Classifier-free vector field (cfv) incremental manner
|
|
return uncond + a_cfg_scale * (audio_uncond_emotion - uncond) + e_cfg_scale * (all_cond - audio_uncond_emotion)
|
|
else:
|
|
return self.forward(t, x, wa, wr, we, prev_x, prev_wa, train = False)
|
|
|