honey126/VoxAI
0
1"""2ein notation:3b - batch4n - sequence5nt - text sequence6nw - raw wave length7d - dimension8"""9 10from __future__ import annotations11from typing import Optional12import math13 14import torch15from torch import nn16import torch.nn.functional as F17import torchaudio18 19from x_transformers.x_transformers import apply_rotary_pos_emb20 21 22# raw wav to mel spec23 24 25class MelSpec(nn.Module):26 def __init__(27 self,28 filter_length=1024,29 hop_length=256,30 win_length=1024,31 n_mel_channels=100,32 target_sample_rate=24_000,33 normalize=False,34 power=1,35 norm=None,36 center=True,37 ):38 super().__init__()39 self.n_mel_channels = n_mel_channels40 41 self.mel_stft = torchaudio.transforms.MelSpectrogram(42 sample_rate=target_sample_rate,43 n_fft=filter_length,44 win_length=win_length,45 hop_length=hop_length,46 n_mels=n_mel_channels,47 power=power,48 center=center,49 normalized=normalize,50 norm=norm,51 )52 53 self.register_buffer("dummy", torch.tensor(0), persistent=False)54 55 def forward(self, inp):56 if len(inp.shape) == 3:57 inp = inp.squeeze(1) # 'b 1 nw -> b nw'58 59 assert len(inp.shape) == 260 61 if self.dummy.device != inp.device:62 self.to(inp.device)63 64 mel = self.mel_stft(inp)65 mel = mel.clamp(min=1e-5).log()66 return mel67 68 69# sinusoidal position embedding70 71 72class SinusPositionEmbedding(nn.Module):73 def __init__(self, dim):74 super().__init__()75 self.dim = dim76 77 def forward(self, x, scale=1000):78 device = x.device79 half_dim = self.dim // 280 emb = math.log(10000) / (half_dim - 1)81 emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)82 emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)83 emb = torch.cat((emb.sin(), emb.cos()), dim=-1)84 return emb85 86 87# convolutional position embedding88 89 90class ConvPositionEmbedding(nn.Module):91 def __init__(self, dim, kernel_size=31, groups=16):92 super().__init__()93 assert kernel_size % 2 != 094 self.conv1d = nn.Sequential(95 nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),96 nn.Mish(),97 nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),98 nn.Mish(),99 )100 101 def forward(self, x: float["b n d"], mask: bool["b n"] | None = None): # noqa: F722102 if mask is not None:103 mask = mask[..., None]104 x = x.masked_fill(~mask, 0.0)105 106 x = x.permute(0, 2, 1)107 x = self.conv1d(x)108 out = x.permute(0, 2, 1)109 110 if mask is not None:111 out = out.masked_fill(~mask, 0.0)112 113 return out114 115 116# rotary positional embedding related117 118 119def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):120 # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning121 # has some connection to NTK literature122 # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/123 # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py124 theta *= theta_rescale_factor ** (dim / (dim - 2))125 freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))126 t = torch.arange(end, device=freqs.device) # type: ignore127 freqs = torch.outer(t, freqs).float() # type: ignore128 freqs_cos = torch.cos(freqs) # real part129 freqs_sin = torch.sin(freqs) # imaginary part130 return torch.cat([freqs_cos, freqs_sin], dim=-1)131 132 133def get_pos_embed_indices(start, length, max_pos, scale=1.0):134 # length = length if isinstance(length, int) else length.max()135 scale = scale * torch.ones_like(start, dtype=torch.float32) # in case scale is a scalar136 pos = (137 start.unsqueeze(1)138 + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()139 )140 # avoid extra long error.141 pos = torch.where(pos < max_pos, pos, max_pos - 1)142 return pos143 144 145# Global Response Normalization layer (Instance Normalization ?)146 147 148class GRN(nn.Module):149 def __init__(self, dim):150 super().__init__()151 self.gamma = nn.Parameter(torch.zeros(1, 1, dim))152 self.beta = nn.Parameter(torch.zeros(1, 1, dim))153 154 def forward(self, x):155 Gx = torch.norm(x, p=2, dim=1, keepdim=True)156 Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)157 return self.gamma * (x * Nx) + self.beta + x158 159 160# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py161# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108162 163 164class ConvNeXtV2Block(nn.Module):165 def __init__(166 self,167 dim: int,168 intermediate_dim: int,169 dilation: int = 1,170 ):171 super().__init__()172 padding = (dilation * (7 - 1)) // 2173 self.dwconv = nn.Conv1d(174 dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation175 ) # depthwise conv176 self.norm = nn.LayerNorm(dim, eps=1e-6)177 self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers178 self.act = nn.GELU()179 self.grn = GRN(intermediate_dim)180 self.pwconv2 = nn.Linear(intermediate_dim, dim)181 182 def forward(self, x: torch.Tensor) -> torch.Tensor:183 residual = x184 x = x.transpose(1, 2) # b n d -> b d n185 x = self.dwconv(x)186 x = x.transpose(1, 2) # b d n -> b n d187 x = self.norm(x)188 x = self.pwconv1(x)189 x = self.act(x)190 x = self.grn(x)191 x = self.pwconv2(x)192 return residual + x193 194 195# AdaLayerNormZero196# return with modulated x for attn input, and params for later mlp modulation197 198 199class AdaLayerNormZero(nn.Module):200 def __init__(self, dim):201 super().__init__()202 203 self.silu = nn.SiLU()204 self.linear = nn.Linear(dim, dim * 6)205 206 self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)207 208 def forward(self, x, emb=None):209 emb = self.linear(self.silu(emb))210 shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)211 212 x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]213 return x, gate_msa, shift_mlp, scale_mlp, gate_mlp214 215 216# AdaLayerNormZero for final layer217# return only with modulated x for attn input, cuz no more mlp modulation218 219 220class AdaLayerNormZero_Final(nn.Module):221 def __init__(self, dim):222 super().__init__()223 224 self.silu = nn.SiLU()225 self.linear = nn.Linear(dim, dim * 2)226 227 self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)228 229 def forward(self, x, emb):230 emb = self.linear(self.silu(emb))231 scale, shift = torch.chunk(emb, 2, dim=1)232 233 x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]234 return x235 236 237# FeedForward238 239 240class FeedForward(nn.Module):241 def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"):242 super().__init__()243 inner_dim = int(dim * mult)244 dim_out = dim_out if dim_out is not None else dim245 246 activation = nn.GELU(approximate=approximate)247 project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)248 self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))249 250 def forward(self, x):251 return self.ff(x)252 253 254# Attention with possible joint part255# modified from diffusers/src/diffusers/models/attention_processor.py256 257 258class Attention(nn.Module):259 def __init__(260 self,261 processor: JointAttnProcessor | AttnProcessor,262 dim: int,263 heads: int = 8,264 dim_head: int = 64,265 dropout: float = 0.0,266 context_dim: Optional[int] = None, # if not None -> joint attention267 context_pre_only=None,268 ):269 super().__init__()270 271 if not hasattr(F, "scaled_dot_product_attention"):272 raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")273 274 self.processor = processor275 276 self.dim = dim277 self.heads = heads278 self.inner_dim = dim_head * heads279 self.dropout = dropout280 281 self.context_dim = context_dim282 self.context_pre_only = context_pre_only283 284 self.to_q = nn.Linear(dim, self.inner_dim)285 self.to_k = nn.Linear(dim, self.inner_dim)286 self.to_v = nn.Linear(dim, self.inner_dim)287 288 if self.context_dim is not None:289 self.to_k_c = nn.Linear(context_dim, self.inner_dim)290 self.to_v_c = nn.Linear(context_dim, self.inner_dim)291 if self.context_pre_only is not None:292 self.to_q_c = nn.Linear(context_dim, self.inner_dim)293 294 self.to_out = nn.ModuleList([])295 self.to_out.append(nn.Linear(self.inner_dim, dim))296 self.to_out.append(nn.Dropout(dropout))297 298 if self.context_pre_only is not None and not self.context_pre_only:299 self.to_out_c = nn.Linear(self.inner_dim, dim)300 301 def forward(302 self,303 x: float["b n d"], # noised input x # noqa: F722304 c: float["b n d"] = None, # context c # noqa: F722305 mask: bool["b n"] | None = None, # noqa: F722306 rope=None, # rotary position embedding for x307 c_rope=None, # rotary position embedding for c308 ) -> torch.Tensor:309 if c is not None:310 return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope)311 else:312 return self.processor(self, x, mask=mask, rope=rope)313 314 315# Attention processor316 317 318class AttnProcessor:319 def __init__(self):320 pass321 322 def __call__(323 self,324 attn: Attention,325 x: float["b n d"], # noised input x # noqa: F722326 mask: bool["b n"] | None = None, # noqa: F722327 rope=None, # rotary position embedding328 ) -> torch.FloatTensor:329 batch_size = x.shape[0]330 331 # `sample` projections.332 query = attn.to_q(x)333 key = attn.to_k(x)334 value = attn.to_v(x)335 336 # apply rotary position embedding337 if rope is not None:338 freqs, xpos_scale = rope339 q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)340 341 query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)342 key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)343 344 # attention345 inner_dim = key.shape[-1]346 head_dim = inner_dim // attn.heads347 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)348 key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)349 value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)350 351 # mask. e.g. inference got a batch with different target durations, mask out the padding352 if mask is not None:353 attn_mask = mask354 attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n'355 attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])356 else:357 attn_mask = None358 359 x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)360 x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)361 x = x.to(query.dtype)362 363 # linear proj364 x = attn.to_out[0](x)365 # dropout366 x = attn.to_out[1](x)367 368 if mask is not None:369 mask = mask.unsqueeze(-1)370 x = x.masked_fill(~mask, 0.0)371 372 return x373 374 375# Joint Attention processor for MM-DiT376# modified from diffusers/src/diffusers/models/attention_processor.py377 378 379class JointAttnProcessor:380 def __init__(self):381 pass382 383 def __call__(384 self,385 attn: Attention,386 x: float["b n d"], # noised input x # noqa: F722387 c: float["b nt d"] = None, # context c, here text # noqa: F722388 mask: bool["b n"] | None = None, # noqa: F722389 rope=None, # rotary position embedding for x390 c_rope=None, # rotary position embedding for c391 ) -> torch.FloatTensor:392 residual = x393 394 batch_size = c.shape[0]395 396 # `sample` projections.397 query = attn.to_q(x)398 key = attn.to_k(x)399 value = attn.to_v(x)400 401 # `context` projections.402 c_query = attn.to_q_c(c)403 c_key = attn.to_k_c(c)404 c_value = attn.to_v_c(c)405 406 # apply rope for context and noised input independently407 if rope is not None:408 freqs, xpos_scale = rope409 q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)410 query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)411 key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)412 if c_rope is not None:413 freqs, xpos_scale = c_rope414 q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)415 c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale)416 c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale)417 418 # attention419 query = torch.cat([query, c_query], dim=1)420 key = torch.cat([key, c_key], dim=1)421 value = torch.cat([value, c_value], dim=1)422 423 inner_dim = key.shape[-1]424 head_dim = inner_dim // attn.heads425 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)426 key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)427 value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)428 429 # mask. e.g. inference got a batch with different target durations, mask out the padding430 if mask is not None:431 attn_mask = F.pad(mask, (0, c.shape[1]), value=True) # no mask for c (text)432 attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n'433 attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])434 else:435 attn_mask = None436 437 x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)438 x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)439 x = x.to(query.dtype)440 441 # Split the attention outputs.442 x, c = (443 x[:, : residual.shape[1]],444 x[:, residual.shape[1] :],445 )446 447 # linear proj448 x = attn.to_out[0](x)449 # dropout450 x = attn.to_out[1](x)451 if not attn.context_pre_only:452 c = attn.to_out_c(c)453 454 if mask is not None:455 mask = mask.unsqueeze(-1)456 x = x.masked_fill(~mask, 0.0)457 # c = c.masked_fill(~mask, 0.) # no mask for c (text)458 459 return x, c460 461 462# DiT Block463 464 465class DiTBlock(nn.Module):466 def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1):467 super().__init__()468 469 self.attn_norm = AdaLayerNormZero(dim)470 self.attn = Attention(471 processor=AttnProcessor(),472 dim=dim,473 heads=heads,474 dim_head=dim_head,475 dropout=dropout,476 )477 478 self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)479 self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")480 481 def forward(self, x, t, mask=None, rope=None): # x: noised input, t: time embedding482 # pre-norm & modulation for attention input483 norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)484 485 # attention486 attn_output = self.attn(x=norm, mask=mask, rope=rope)487 488 # process attention output for input x489 x = x + gate_msa.unsqueeze(1) * attn_output490 491 norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]492 ff_output = self.ff(norm)493 x = x + gate_mlp.unsqueeze(1) * ff_output494 495 return x496 497 498# MMDiT Block https://arxiv.org/abs/2403.03206499 500 501class MMDiTBlock(nn.Module):502 r"""503 modified from diffusers/src/diffusers/models/attention.py504 505 notes.506 _c: context related. text, cond, etc. (left part in sd3 fig2.b)507 _x: noised input related. (right part)508 context_pre_only: last layer only do prenorm + modulation cuz no more ffn509 """510 511 def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, context_pre_only=False):512 super().__init__()513 514 self.context_pre_only = context_pre_only515 516 self.attn_norm_c = AdaLayerNormZero_Final(dim) if context_pre_only else AdaLayerNormZero(dim)517 self.attn_norm_x = AdaLayerNormZero(dim)518 self.attn = Attention(519 processor=JointAttnProcessor(),520 dim=dim,521 heads=heads,522 dim_head=dim_head,523 dropout=dropout,524 context_dim=dim,525 context_pre_only=context_pre_only,526 )527 528 if not context_pre_only:529 self.ff_norm_c = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)530 self.ff_c = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")531 else:532 self.ff_norm_c = None533 self.ff_c = None534 self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)535 self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")536 537 def forward(self, x, c, t, mask=None, rope=None, c_rope=None): # x: noised input, c: context, t: time embedding538 # pre-norm & modulation for attention input539 if self.context_pre_only:540 norm_c = self.attn_norm_c(c, t)541 else:542 norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t)543 norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t)544 545 # attention546 x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope)547 548 # process attention output for context c549 if self.context_pre_only:550 c = None551 else: # if not last layer552 c = c + c_gate_msa.unsqueeze(1) * c_attn_output553 554 norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]555 c_ff_output = self.ff_c(norm_c)556 c = c + c_gate_mlp.unsqueeze(1) * c_ff_output557 558 # process attention output for input x559 x = x + x_gate_msa.unsqueeze(1) * x_attn_output560 561 norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None]562 x_ff_output = self.ff_x(norm_x)563 x = x + x_gate_mlp.unsqueeze(1) * x_ff_output564 565 return c, x566 567 568# time step conditioning embedding569 570 571class TimestepEmbedding(nn.Module):572 def __init__(self, dim, freq_embed_dim=256):573 super().__init__()574 self.time_embed = SinusPositionEmbedding(freq_embed_dim)575 self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))576 577 def forward(self, timestep: float["b"]): # noqa: F821578 time_hidden = self.time_embed(timestep)579 time_hidden = time_hidden.to(timestep.dtype)580 time = self.time_mlp(time_hidden) # b d581 return time582