| import torch.nn as nn
|
| import torch.nn.functional as F
|
| import torch
|
| import math
|
| from einops import rearrange
|
|
|
| from ..wanvideo.modules.model import WanRMSNorm, attention
|
| from ..multitalk.multitalk import RotaryPositionalEmbedding1D, normalize_and_scale
|
|
|
| class FeedForwardSwiGLU(nn.Module):
|
| def __init__(
|
| self,
|
| dim: int,
|
| hidden_dim: int,
|
| multiple_of: int = 256,
|
| ):
|
| super().__init__()
|
| hidden_dim = int(2 * hidden_dim / 3)
|
| hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
|
|
|
| self.dim = dim
|
| self.hidden_dim = hidden_dim
|
| self.w1 = nn.Linear(dim, hidden_dim, bias=False)
|
| self.w2 = nn.Linear(hidden_dim, dim, bias=False)
|
| self.w3 = nn.Linear(dim, hidden_dim, bias=False)
|
|
|
| def forward(self, x):
|
| return self.w2(F.silu(self.w1(x)) * self.w3(x))
|
|
|
| class TimestepEmbedder(nn.Module):
|
| """
|
| Embeds scalar timesteps into vector representations.
|
| """
|
|
|
| def __init__(self, t_embed_dim, frequency_embedding_size=256):
|
| super().__init__()
|
| self.t_embed_dim = t_embed_dim
|
| self.frequency_embedding_size = frequency_embedding_size
|
| self.mlp = nn.Sequential(
|
| nn.Linear(frequency_embedding_size, t_embed_dim, bias=True),
|
| nn.SiLU(),
|
| nn.Linear(t_embed_dim, t_embed_dim, bias=True),
|
| )
|
|
|
| @staticmethod
|
| def timestep_embedding(t, dim, max_period=10000):
|
| """
|
| 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.
|
| """
|
| half = dim // 2
|
| freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half)
|
| freqs = freqs.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, dtype):
|
| t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
|
| if t_freq.dtype != dtype:
|
| t_freq = t_freq.to(dtype)
|
| t_emb = self.mlp(t_freq)
|
| return t_emb
|
|
|
|
|
| class SingleStreamAttention(nn.Module):
|
| def __init__(
|
| self,
|
| dim: int,
|
| encoder_hidden_states_dim: int,
|
| num_heads: int,
|
| qkv_bias: bool,
|
| qk_norm: bool,
|
| attn_drop: float = 0.0,
|
| proj_drop: float = 0.0,
|
| eps: float = 1e-6,
|
| class_range: int = 24,
|
| class_interval: int = 4,
|
| attention_mode: str = "sdpa",
|
| ) -> None:
|
| super().__init__()
|
| assert dim % num_heads == 0, "dim should be divisible by num_heads"
|
| self.dim = dim
|
| self.encoder_hidden_states_dim = encoder_hidden_states_dim
|
| self.num_heads = num_heads
|
| self.head_dim = dim // num_heads
|
| self.scale = self.head_dim**-0.5
|
|
|
| self.q_linear = nn.Linear(dim, dim, bias=qkv_bias)
|
| self.q_norm = WanRMSNorm(self.head_dim, eps=eps) 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)
|
|
|
| self.kv_linear = nn.Linear(encoder_hidden_states_dim, dim * 2, bias=qkv_bias)
|
| self.k_norm = WanRMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()
|
|
|
| self.attention_mode = attention_mode
|
|
|
|
|
| self.class_interval = class_interval
|
| self.class_range = class_range
|
| self.rope_h1 = (0, self.class_interval)
|
| self.rope_h2 = (self.class_range - self.class_interval, self.class_range)
|
| self.rope_bak = int(self.class_range // 2)
|
| self.rope_1d = RotaryPositionalEmbedding1D(self.head_dim)
|
|
|
| def _process_cross_attn(self, x, cond, frames_num=None, x_ref_attn_map=None):
|
|
|
| N_t = frames_num
|
| out_dtype = x.dtype
|
| x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)
|
|
|
|
|
| B, N, C = x.shape
|
| q = self.q_linear(x)
|
| q_shape = (B, N, self.num_heads, self.head_dim)
|
| q = q.view(q_shape).permute((0, 2, 1, 3))
|
| q = self.q_norm(q.to(self.q_norm.weight.dtype)).to(q.dtype)
|
|
|
|
|
| if x_ref_attn_map is not None:
|
| max_values = x_ref_attn_map.max(1).values[:, None, None]
|
| min_values = x_ref_attn_map.min(1).values[:, None, None]
|
| max_min_values = torch.cat([max_values, min_values], dim=2)
|
| human1_max_value, human1_min_value = max_min_values[0, :, 0].max(), max_min_values[0, :, 1].min()
|
| human2_max_value, human2_min_value = max_min_values[1, :, 0].max(), max_min_values[1, :, 1].min()
|
|
|
| human1 = normalize_and_scale(x_ref_attn_map[0], (human1_min_value, human1_max_value), (self.rope_h1[0], self.rope_h1[1]))
|
| human2 = normalize_and_scale(x_ref_attn_map[1], (human2_min_value, human2_max_value), (self.rope_h2[0], self.rope_h2[1]))
|
| back = torch.full((x_ref_attn_map.size(1),), self.rope_bak, dtype=human1.dtype).to(human1.device)
|
| max_indices = x_ref_attn_map.argmax(dim=0)
|
| normalized_map = torch.stack([human1, human2, back], dim=1)
|
| normalized_pos = normalized_map[range(x_ref_attn_map.size(1)), max_indices]
|
|
|
| q = rearrange(q, "(B N_t) H S C -> B H (N_t S) C", N_t=N_t)
|
| q = self.rope_1d(q, normalized_pos)
|
| q = rearrange(q, "B H (N_t S) C -> (B N_t) H S C", N_t=N_t)
|
|
|
|
|
| _, N_a, _ = cond.shape
|
| encoder_kv = self.kv_linear(cond)
|
| encoder_kv_shape = (B, N_a, 2, self.num_heads, self.head_dim)
|
| encoder_kv = encoder_kv.view(encoder_kv_shape).permute((2, 0, 3, 1, 4))
|
|
|
| encoder_k, encoder_v = encoder_kv.unbind(0)
|
| encoder_k = self.k_norm(encoder_k.to(self.k_norm.weight.dtype)).to(encoder_k.dtype)
|
|
|
|
|
|
|
| if x_ref_attn_map is not None:
|
| per_frame = torch.zeros(N_a, dtype=encoder_k.dtype).to(encoder_k.device)
|
| per_frame[:per_frame.size(0)//2] = (self.rope_h1[0] + self.rope_h1[1]) / 2
|
| per_frame[per_frame.size(0)//2:] = (self.rope_h2[0] + self.rope_h2[1]) / 2
|
| encoder_pos = torch.concat([per_frame]*N_t, dim=0)
|
| encoder_k = rearrange(encoder_k, "(B N_t) H S C -> B H (N_t S) C", N_t=N_t)
|
| encoder_k = self.rope_1d(encoder_k, encoder_pos)
|
| encoder_k = rearrange(encoder_k, "B H (N_t S) C -> (B N_t) H S C", N_t=N_t)
|
|
|
|
|
|
|
|
|
| q = rearrange(q, "B H M K -> B M H K")
|
| encoder_k = rearrange(encoder_k, "B H M K -> B M H K")
|
| encoder_v = rearrange(encoder_v, "B H M K -> B M H K")
|
| x = attention(q, encoder_k, encoder_v, attention_mode=self.attention_mode)
|
| x = rearrange(x, "B M H K -> B H M K")
|
|
|
|
|
| x_output_shape = (B, N, C)
|
| x = x.transpose(1, 2)
|
| x = x.reshape(x_output_shape)
|
| x = self.proj(x)
|
| x = self.proj_drop(x)
|
|
|
|
|
| x = rearrange(x, "(B N_t) S C -> B (N_t S) C", N_t=N_t)
|
|
|
| return x.type(out_dtype)
|
|
|
| def forward(self, x, cond, num_latent_frames=None, num_cond_latents=None, x_ref_attn_map=None, human_num=None):
|
|
|
| B, N, C = x.shape
|
| if (num_cond_latents is None or num_cond_latents == 0):
|
|
|
| output = self._process_cross_attn(x, cond, num_latent_frames, x_ref_attn_map)
|
| return None, output
|
| elif num_cond_latents is not None and num_cond_latents > 0:
|
|
|
| num_cond_latents_thw = num_cond_latents * (N // num_latent_frames)
|
| x_noise = x[:, num_cond_latents_thw:]
|
| cond = rearrange(cond, "(B N_t) M C -> B N_t M C", B=B)
|
| cond = cond[:, num_cond_latents:]
|
| cond = rearrange(cond, "B N_t M C -> (B N_t) M C")
|
| frames_num = num_latent_frames - num_cond_latents
|
| if human_num is not None and human_num == 2:
|
|
|
| output_noise = self._process_cross_attn(x_noise, cond, frames_num, x_ref_attn_map)
|
| else:
|
|
|
| output_noise = self._process_cross_attn(x_noise, cond, frames_num)
|
| output_cond = torch.zeros((B, num_cond_latents_thw, C), dtype=output_noise.dtype, device=output_noise.device)
|
| return output_cond, output_noise
|
| else:
|
| raise NotImplementedError
|
|
|
