dskill/DiffRhythm
2
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 torch17import torch.nn.functional as F18import torchaudio19 20from x_transformers.x_transformers import apply_rotary_pos_emb21 22 23 24class FiLMLayer(nn.Module):25 """26 Feature-wise Linear Modulation (FiLM) layer27 Reference: https://arxiv.org/abs/1709.0787128 """29 def __init__(self, in_channels, cond_channels):30 31 super(FiLMLayer, self).__init__()32 self.in_channels = in_channels33 self.film = nn.Conv1d(cond_channels, in_channels * 2, 1)34 35 def forward(self, x, c):36 gamma, beta = torch.chunk(self.film(c.unsqueeze(2)), chunks=2, dim=1)37 gamma = gamma.transpose(1, 2)38 beta = beta.transpose(1, 2)39 # print(gamma.shape, beta.shape)40 return gamma * x + beta41 42# raw wav to mel spec43 44 45class MelSpec(nn.Module):46 def __init__(47 self,48 filter_length=1024,49 hop_length=256,50 win_length=1024,51 n_mel_channels=100,52 target_sample_rate=24_000,53 normalize=False,54 power=1,55 norm=None,56 center=True,57 ):58 super().__init__()59 self.n_mel_channels = n_mel_channels60 61 self.mel_stft = torchaudio.transforms.MelSpectrogram(62 sample_rate=target_sample_rate,63 n_fft=filter_length,64 win_length=win_length,65 hop_length=hop_length,66 n_mels=n_mel_channels,67 power=power,68 center=center,69 normalized=normalize,70 norm=norm,71 )72 73 self.register_buffer("dummy", torch.tensor(0), persistent=False)74 75 def forward(self, inp):76 if len(inp.shape) == 3:77 inp = inp.squeeze(1) # 'b 1 nw -> b nw'78 79 assert len(inp.shape) == 280 81 if self.dummy.device != inp.device:82 self.to(inp.device)83 84 mel = self.mel_stft(inp)85 mel = mel.clamp(min=1e-5).log()86 return mel87 88 89# sinusoidal position embedding90 91 92class SinusPositionEmbedding(nn.Module):93 def __init__(self, dim):94 super().__init__()95 self.dim = dim96 97 def forward(self, x, scale=1000):98 device = x.device99 half_dim = self.dim // 2100 emb = math.log(10000) / (half_dim - 1)101 emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)102 emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)103 emb = torch.cat((emb.sin(), emb.cos()), dim=-1)104 return emb105 106 107# convolutional position embedding108 109 110class ConvPositionEmbedding(nn.Module):111 def __init__(self, dim, kernel_size=31, groups=16):112 super().__init__()113 assert kernel_size % 2 != 0114 self.conv1d = nn.Sequential(115 nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),116 nn.Mish(),117 nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),118 nn.Mish(),119 )120 121 def forward(self, x: float["b n d"], mask: bool["b n"] | None = None): # noqa: F722122 if mask is not None:123 mask = mask[..., None]124 x = x.masked_fill(~mask, 0.0)125 126 x = x.permute(0, 2, 1)127 x = self.conv1d(x)128 out = x.permute(0, 2, 1)129 130 if mask is not None:131 out = out.masked_fill(~mask, 0.0)132 133 return out134 135 136# rotary positional embedding related137 138 139def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):140 # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning141 # has some connection to NTK literature142 # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/143 # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py144 theta *= theta_rescale_factor ** (dim / (dim - 2))145 freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))146 t = torch.arange(end, device=freqs.device) # type: ignore147 freqs = torch.outer(t, freqs).float() # type: ignore148 freqs_cos = torch.cos(freqs) # real part149 freqs_sin = torch.sin(freqs) # imaginary part150 return torch.cat([freqs_cos, freqs_sin], dim=-1)151 152 153def get_pos_embed_indices(start, length, max_pos, scale=1.0):154 # length = length if isinstance(length, int) else length.max()155 scale = scale * torch.ones_like(start, dtype=torch.float32) # in case scale is a scalar156 pos = (157 start.unsqueeze(1)158 + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()159 )160 # avoid extra long error.161 pos = torch.where(pos < max_pos, pos, max_pos - 1)162 return pos163 164 165# Global Response Normalization layer (Instance Normalization ?)166 167 168class GRN(nn.Module):169 def __init__(self, dim):170 super().__init__()171 self.gamma = nn.Parameter(torch.zeros(1, 1, dim))172 self.beta = nn.Parameter(torch.zeros(1, 1, dim))173 174 def forward(self, x):175 Gx = torch.norm(x, p=2, dim=1, keepdim=True)176 Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)177 return self.gamma * (x * Nx) + self.beta + x178 179 180# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py181# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108182 183 184class ConvNeXtV2Block(nn.Module):185 def __init__(186 self,187 dim: int,188 intermediate_dim: int,189 dilation: int = 1,190 ):191 super().__init__()192 padding = (dilation * (7 - 1)) // 2193 self.dwconv = nn.Conv1d(194 dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation195 ) # depthwise conv196 self.norm = nn.LayerNorm(dim, eps=1e-6)197 self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers198 self.act = nn.GELU()199 self.grn = GRN(intermediate_dim)200 self.pwconv2 = nn.Linear(intermediate_dim, dim)201 202 def forward(self, x: torch.Tensor) -> torch.Tensor:203 residual = x204 x = x.transpose(1, 2) # b n d -> b d n205 x = self.dwconv(x)206 x = x.transpose(1, 2) # b d n -> b n d207 x = self.norm(x)208 x = self.pwconv1(x)209 x = self.act(x)210 x = self.grn(x)211 x = self.pwconv2(x)212 return residual + x213 214 215# AdaLayerNormZero216# return with modulated x for attn input, and params for later mlp modulation217 218 219class AdaLayerNormZero(nn.Module):220 def __init__(self, dim):221 super().__init__()222 223 self.silu = nn.SiLU()224 self.linear = nn.Linear(dim, dim * 6)225 226 self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)227 228 def forward(self, x, emb=None):229 emb = self.linear(self.silu(emb))230 shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)231 232 x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]233 return x, gate_msa, shift_mlp, scale_mlp, gate_mlp234 235 236# AdaLayerNormZero for final layer237# return only with modulated x for attn input, cuz no more mlp modulation238 239 240class AdaLayerNormZero_Final(nn.Module):241 def __init__(self, dim, cond_dim):242 super().__init__()243 244 self.silu = nn.SiLU()245 self.linear = nn.Linear(cond_dim, dim * 2)246 247 self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)248 249 def forward(self, x, emb):250 emb = self.linear(self.silu(emb))251 scale, shift = torch.chunk(emb, 2, dim=1)252 253 x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]254 return x255 256 257# FeedForward258 259 260class FeedForward(nn.Module):261 def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"):262 super().__init__()263 inner_dim = int(dim * mult)264 dim_out = dim_out if dim_out is not None else dim265 266 activation = nn.GELU(approximate=approximate)267 #activation = nn.SiLU()268 project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)269 self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))270 271 def forward(self, x):272 return self.ff(x)273 274 275# Attention with possible joint part276# modified from diffusers/src/diffusers/models/attention_processor.py277 278 279class Attention(nn.Module):280 def __init__(281 self,282 processor: JointAttnProcessor | AttnProcessor,283 dim: int,284 heads: int = 8,285 dim_head: int = 64,286 dropout: float = 0.0,287 context_dim: Optional[int] = None, # if not None -> joint attention288 context_pre_only=None,289 ):290 super().__init__()291 292 if not hasattr(F, "scaled_dot_product_attention"):293 raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")294 295 self.processor = processor296 297 self.dim = dim298 self.heads = heads299 self.inner_dim = dim_head * heads300 self.dropout = dropout301 302 self.context_dim = context_dim303 self.context_pre_only = context_pre_only304 305 self.to_q = nn.Linear(dim, self.inner_dim)306 self.to_k = nn.Linear(dim, self.inner_dim)307 self.to_v = nn.Linear(dim, self.inner_dim)308 309 if self.context_dim is not None:310 self.to_k_c = nn.Linear(context_dim, self.inner_dim)311 self.to_v_c = nn.Linear(context_dim, self.inner_dim)312 if self.context_pre_only is not None:313 self.to_q_c = nn.Linear(context_dim, self.inner_dim)314 315 self.to_out = nn.ModuleList([])316 self.to_out.append(nn.Linear(self.inner_dim, dim))317 self.to_out.append(nn.Dropout(dropout))318 319 if self.context_pre_only is not None and not self.context_pre_only:320 self.to_out_c = nn.Linear(self.inner_dim, dim)321 322 def forward(323 self,324 x: float["b n d"], # noised input x # noqa: F722325 c: float["b n d"] = None, # context c # noqa: F722326 mask: bool["b n"] | None = None, # noqa: F722327 rope=None, # rotary position embedding for x328 c_rope=None, # rotary position embedding for c329 ) -> torch.Tensor:330 if c is not None:331 return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope)332 else:333 return self.processor(self, x, mask=mask, rope=rope)334 335 336# Attention processor337 338 339class AttnProcessor:340 def __init__(self):341 pass342 343 def __call__(344 self,345 attn: Attention,346 x: float["b n d"], # noised input x # noqa: F722347 mask: bool["b n"] | None = None, # noqa: F722348 rope=None, # rotary position embedding349 ) -> torch.FloatTensor:350 batch_size = x.shape[0]351 352 # `sample` projections.353 query = attn.to_q(x)354 key = attn.to_k(x)355 value = attn.to_v(x)356 357 # apply rotary position embedding358 if rope is not None:359 freqs, xpos_scale = rope360 q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)361 362 query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)363 key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)364 365 # attention366 inner_dim = key.shape[-1]367 head_dim = inner_dim // attn.heads368 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)369 key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)370 value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)371 372 # mask. e.g. inference got a batch with different target durations, mask out the padding373 if mask is not None:374 attn_mask = mask375 attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n'376 attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])377 else:378 attn_mask = None379 380 x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)381 x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)382 x = x.to(query.dtype)383 384 # linear proj385 x = attn.to_out[0](x)386 # dropout387 x = attn.to_out[1](x)388 389 if mask is not None:390 mask = mask.unsqueeze(-1)391 x = x.masked_fill(~mask, 0.0)392 393 return x394 395 396# Joint Attention processor for MM-DiT397# modified from diffusers/src/diffusers/models/attention_processor.py398 399 400class JointAttnProcessor:401 def __init__(self):402 pass403 404 def __call__(405 self,406 attn: Attention,407 x: float["b n d"], # noised input x # noqa: F722408 c: float["b nt d"] = None, # context c, here text # noqa: F722409 mask: bool["b n"] | None = None, # noqa: F722410 rope=None, # rotary position embedding for x411 c_rope=None, # rotary position embedding for c412 ) -> torch.FloatTensor:413 residual = x414 415 batch_size = c.shape[0]416 417 # `sample` projections.418 query = attn.to_q(x)419 key = attn.to_k(x)420 value = attn.to_v(x)421 422 # `context` projections.423 c_query = attn.to_q_c(c)424 c_key = attn.to_k_c(c)425 c_value = attn.to_v_c(c)426 427 # apply rope for context and noised input independently428 if rope is not None:429 freqs, xpos_scale = rope430 q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)431 query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)432 key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)433 if c_rope is not None:434 freqs, xpos_scale = c_rope435 q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)436 c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale)437 c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale)438 439 # attention440 query = torch.cat([query, c_query], dim=1)441 key = torch.cat([key, c_key], dim=1)442 value = torch.cat([value, c_value], dim=1)443 444 inner_dim = key.shape[-1]445 head_dim = inner_dim // attn.heads446 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)447 key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)448 value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)449 450 # mask. e.g. inference got a batch with different target durations, mask out the padding451 if mask is not None:452 attn_mask = F.pad(mask, (0, c.shape[1]), value=True) # no mask for c (text)453 attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n'454 attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])455 else:456 attn_mask = None457 458 x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)459 x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)460 x = x.to(query.dtype)461 462 # Split the attention outputs.463 x, c = (464 x[:, : residual.shape[1]],465 x[:, residual.shape[1] :],466 )467 468 # linear proj469 x = attn.to_out[0](x)470 # dropout471 x = attn.to_out[1](x)472 if not attn.context_pre_only:473 c = attn.to_out_c(c)474 475 if mask is not None:476 mask = mask.unsqueeze(-1)477 x = x.masked_fill(~mask, 0.0)478 # c = c.masked_fill(~mask, 0.) # no mask for c (text)479 480 return x, c481 482 483# DiT Block484 485 486class DiTBlock(nn.Module):487 def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, use_style_prompt=False):488 super().__init__()489 490 self.attn_norm = AdaLayerNormZero(dim)491 self.attn = Attention(492 processor=AttnProcessor(),493 dim=dim,494 heads=heads,495 dim_head=dim_head,496 dropout=dropout,497 )498 499 self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)500 self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")501 502 self.use_style_prompt = use_style_prompt503 if use_style_prompt:504 #self.film = FiLMLayer(dim, dim)505 self.prompt_norm = AdaLayerNormZero_Final(dim)506 507 def forward(self, x, t, c=None, mask=None, rope=None): # x: noised input, t: time embedding508 if c is not None and self.use_style_prompt:509 #x = self.film(x, c)510 x = self.prompt_norm(x, c)511 512 # pre-norm & modulation for attention input513 norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)514 515 # attention516 attn_output = self.attn(x=norm, mask=mask, rope=rope)517 518 # process attention output for input x519 x = x + gate_msa.unsqueeze(1) * attn_output520 521 norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]522 ff_output = self.ff(norm)523 x = x + gate_mlp.unsqueeze(1) * ff_output524 525 return x526 527 528# MMDiT Block https://arxiv.org/abs/2403.03206529 530 531class MMDiTBlock(nn.Module):532 r"""533 modified from diffusers/src/diffusers/models/attention.py534 535 notes.536 _c: context related. text, cond, etc. (left part in sd3 fig2.b)537 _x: noised input related. (right part)538 context_pre_only: last layer only do prenorm + modulation cuz no more ffn539 """540 541 def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, context_pre_only=False):542 super().__init__()543 544 self.context_pre_only = context_pre_only545 546 self.attn_norm_c = AdaLayerNormZero_Final(dim) if context_pre_only else AdaLayerNormZero(dim)547 self.attn_norm_x = AdaLayerNormZero(dim)548 self.attn = Attention(549 processor=JointAttnProcessor(),550 dim=dim,551 heads=heads,552 dim_head=dim_head,553 dropout=dropout,554 context_dim=dim,555 context_pre_only=context_pre_only,556 )557 558 if not context_pre_only:559 self.ff_norm_c = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)560 self.ff_c = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")561 else:562 self.ff_norm_c = None563 self.ff_c = None564 self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)565 self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")566 567 def forward(self, x, c, t, mask=None, rope=None, c_rope=None): # x: noised input, c: context, t: time embedding568 # pre-norm & modulation for attention input569 if self.context_pre_only:570 norm_c = self.attn_norm_c(c, t)571 else:572 norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t)573 norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t)574 575 # attention576 x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope)577 578 # process attention output for context c579 if self.context_pre_only:580 c = None581 else: # if not last layer582 c = c + c_gate_msa.unsqueeze(1) * c_attn_output583 584 norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]585 c_ff_output = self.ff_c(norm_c)586 c = c + c_gate_mlp.unsqueeze(1) * c_ff_output587 588 # process attention output for input x589 x = x + x_gate_msa.unsqueeze(1) * x_attn_output590 591 norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None]592 x_ff_output = self.ff_x(norm_x)593 x = x + x_gate_mlp.unsqueeze(1) * x_ff_output594 595 return c, x596 597 598# time step conditioning embedding599 600 601class TimestepEmbedding(nn.Module):602 def __init__(self, dim, freq_embed_dim=256):603 super().__init__()604 self.time_embed = SinusPositionEmbedding(freq_embed_dim)605 self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))606 607 def forward(self, timestep: float["b"]): # noqa: F821608 time_hidden = self.time_embed(timestep)609 time_hidden = time_hidden.to(timestep.dtype)610 time = self.time_mlp(time_hidden) # b d611 return time612