| 1 | """ |
| 2 | ein notation: |
| 3 | b - batch |
| 4 | n - sequence |
| 5 | nt - text sequence |
| 6 | nw - raw wave length |
| 7 | d - dimension |
| 8 | """ |
| 9 | # ruff: noqa: F722 F821 |
| 10 | |
| 11 | from __future__ import annotations |
| 12 | |
| 13 | import math |
| 14 | import warnings |
| 15 | from typing import Optional |
| 16 | |
| 17 | import torch |
| 18 | import torch.nn.functional as F |
| 19 | import torchaudio |
| 20 | from librosa.filters import mel as librosa_mel_fn |
| 21 | from torch import nn |
| 22 | from x_transformers.x_transformers import apply_rotary_pos_emb |
| 23 | |
| 24 | from f5_tts.model.utils import is_package_available |
| 25 | |
| 26 | |
| 27 | # raw wav to mel spec |
| 28 | |
| 29 | |
| 30 | mel_basis_cache = {} |
| 31 | hann_window_cache = {} |
| 32 | vocos_mel_stft_cache = {} |
| 33 | |
| 34 | |
| 35 | def get_bigvgan_mel_spectrogram( |
| 36 | waveform, |
| 37 | n_fft=1024, |
| 38 | n_mel_channels=100, |
| 39 | target_sample_rate=24000, |
| 40 | hop_length=256, |
| 41 | win_length=1024, |
| 42 | fmin=0, |
| 43 | fmax=None, |
| 44 | center=False, |
| 45 | ): # Copy from https://github.com/NVIDIA/BigVGAN/tree/main |
| 46 | device = waveform.device |
| 47 | key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}" |
| 48 | |
| 49 | if key not in mel_basis_cache: |
| 50 | mel = librosa_mel_fn(sr=target_sample_rate, n_fft=n_fft, n_mels=n_mel_channels, fmin=fmin, fmax=fmax) |
| 51 | mel_basis_cache[key] = torch.from_numpy(mel).float().to(device) # TODO: why they need .float()? |
| 52 | hann_window_cache[key] = torch.hann_window(win_length).to(device) |
| 53 | |
| 54 | mel_basis = mel_basis_cache[key] |
| 55 | hann_window = hann_window_cache[key] |
| 56 | |
| 57 | padding = (n_fft - hop_length) // 2 |
| 58 | waveform = torch.nn.functional.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1) |
| 59 | |
| 60 | spec = torch.stft( |
| 61 | waveform, |
| 62 | n_fft, |
| 63 | hop_length=hop_length, |
| 64 | win_length=win_length, |
| 65 | window=hann_window, |
| 66 | center=center, |
| 67 | pad_mode="reflect", |
| 68 | normalized=False, |
| 69 | onesided=True, |
| 70 | return_complex=True, |
| 71 | ) |
| 72 | spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9) |
| 73 | |
| 74 | mel_spec = torch.matmul(mel_basis, spec) |
| 75 | mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5)) |
| 76 | |
| 77 | return mel_spec |
| 78 | |
| 79 | |
| 80 | def get_vocos_mel_spectrogram( |
| 81 | waveform, |
| 82 | n_fft=1024, |
| 83 | n_mel_channels=100, |
| 84 | target_sample_rate=24000, |
| 85 | hop_length=256, |
| 86 | win_length=1024, |
| 87 | ): |
| 88 | device = waveform.device |
| 89 | key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{device}" |
| 90 | if key not in vocos_mel_stft_cache: |
| 91 | vocos_mel_stft_cache[key] = torchaudio.transforms.MelSpectrogram( |
| 92 | sample_rate=target_sample_rate, |
| 93 | n_fft=n_fft, |
| 94 | win_length=win_length, |
| 95 | hop_length=hop_length, |
| 96 | n_mels=n_mel_channels, |
| 97 | power=1, |
| 98 | center=True, |
| 99 | normalized=False, |
| 100 | norm=None, |
| 101 | ).to(device) |
| 102 | if len(waveform.shape) == 3: |
| 103 | waveform = waveform.squeeze(1) # 'b 1 nw -> b nw' |
| 104 | |
| 105 | assert len(waveform.shape) == 2 |
| 106 | |
| 107 | mel = vocos_mel_stft_cache[key](waveform) |
| 108 | mel = mel.clamp(min=1e-5).log() |
| 109 | return mel |
| 110 | |
| 111 | |
| 112 | class MelSpec(nn.Module): |
| 113 | def __init__( |
| 114 | self, |
| 115 | n_fft=1024, |
| 116 | hop_length=256, |
| 117 | win_length=1024, |
| 118 | n_mel_channels=100, |
| 119 | target_sample_rate=24_000, |
| 120 | mel_spec_type="vocos", |
| 121 | ): |
| 122 | super().__init__() |
| 123 | assert mel_spec_type in ["vocos", "bigvgan"], print("We only support two extract mel backend: vocos or bigvgan") |
| 124 | |
| 125 | self.n_fft = n_fft |
| 126 | self.hop_length = hop_length |
| 127 | self.win_length = win_length |
| 128 | self.n_mel_channels = n_mel_channels |
| 129 | self.target_sample_rate = target_sample_rate |
| 130 | |
| 131 | if mel_spec_type == "vocos": |
| 132 | self.extractor = get_vocos_mel_spectrogram |
| 133 | elif mel_spec_type == "bigvgan": |
| 134 | self.extractor = get_bigvgan_mel_spectrogram |
| 135 | |
| 136 | self.register_buffer("dummy", torch.tensor(0), persistent=False) |
| 137 | |
| 138 | def forward(self, wav): |
| 139 | if self.dummy.device != wav.device: |
| 140 | self.to(wav.device) |
| 141 | |
| 142 | mel = self.extractor( |
| 143 | waveform=wav, |
| 144 | n_fft=self.n_fft, |
| 145 | n_mel_channels=self.n_mel_channels, |
| 146 | target_sample_rate=self.target_sample_rate, |
| 147 | hop_length=self.hop_length, |
| 148 | win_length=self.win_length, |
| 149 | ) |
| 150 | |
| 151 | return mel |
| 152 | |
| 153 | |
| 154 | # sinusoidal position embedding |
| 155 | |
| 156 | |
| 157 | class SinusPositionEmbedding(nn.Module): |
| 158 | def __init__(self, dim): |
| 159 | super().__init__() |
| 160 | self.dim = dim |
| 161 | |
| 162 | def forward(self, x, scale=1000): |
| 163 | device = x.device |
| 164 | half_dim = self.dim // 2 |
| 165 | emb = math.log(10000) / (half_dim - 1) |
| 166 | emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) |
| 167 | emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) |
| 168 | emb = torch.cat((emb.sin(), emb.cos()), dim=-1) |
| 169 | return emb |
| 170 | |
| 171 | |
| 172 | # convolutional position embedding |
| 173 | |
| 174 | |
| 175 | class ConvPositionEmbedding(nn.Module): |
| 176 | def __init__(self, dim, kernel_size=31, groups=16): |
| 177 | super().__init__() |
| 178 | assert kernel_size % 2 != 0 |
| 179 | self.conv1d = nn.Sequential( |
| 180 | nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2), |
| 181 | nn.Mish(), |
| 182 | nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2), |
| 183 | nn.Mish(), |
| 184 | ) |
| 185 | self.layer_need_mask_idx = [i for i, layer in enumerate(self.conv1d) if isinstance(layer, nn.Conv1d)] |
| 186 | |
| 187 | def forward(self, x: float["b n d"], mask: bool["b n"] | None = None): |
| 188 | if mask is not None: |
| 189 | mask = mask.unsqueeze(1) # [B 1 N] |
| 190 | x = x.permute(0, 2, 1) # [B D N] |
| 191 | |
| 192 | if mask is not None: |
| 193 | x = x.masked_fill(~mask, 0.0) |
| 194 | for i, block in enumerate(self.conv1d): |
| 195 | x = block(x) |
| 196 | if mask is not None and i in self.layer_need_mask_idx: |
| 197 | x = x.masked_fill(~mask, 0.0) |
| 198 | |
| 199 | x = x.permute(0, 2, 1) # [B N D] |
| 200 | |
| 201 | return x |
| 202 | |
| 203 | |
| 204 | # rotary positional embedding related |
| 205 | |
| 206 | |
| 207 | def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0): |
| 208 | # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning |
| 209 | # has some connection to NTK literature |
| 210 | # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/ |
| 211 | # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py |
| 212 | theta *= theta_rescale_factor ** (dim / (dim - 2)) |
| 213 | freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) |
| 214 | t = torch.arange(end, device=freqs.device) # type: ignore |
| 215 | freqs = torch.outer(t, freqs).float() # type: ignore |
| 216 | freqs_cos = torch.cos(freqs) # real part |
| 217 | freqs_sin = torch.sin(freqs) # imaginary part |
| 218 | return torch.cat([freqs_cos, freqs_sin], dim=-1) |
| 219 | |
| 220 | |
| 221 | def get_pos_embed_indices(start, length, max_pos, scale=1.0): |
| 222 | # length = length if isinstance(length, int) else length.max() |
| 223 | scale = scale * torch.ones_like(start, dtype=torch.float32) # in case scale is a scalar |
| 224 | pos = ( |
| 225 | start.unsqueeze(1) |
| 226 | + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long() |
| 227 | ) |
| 228 | # avoid extra long error. |
| 229 | pos = torch.where(pos < max_pos, pos, max_pos - 1) |
| 230 | return pos |
| 231 | |
| 232 | |
| 233 | # Global Response Normalization layer (Instance Normalization ?) |
| 234 | |
| 235 | |
| 236 | class GRN(nn.Module): |
| 237 | def __init__(self, dim): |
| 238 | super().__init__() |
| 239 | self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) |
| 240 | self.beta = nn.Parameter(torch.zeros(1, 1, dim)) |
| 241 | |
| 242 | def forward(self, x): |
| 243 | Gx = torch.norm(x, p=2, dim=1, keepdim=True) |
| 244 | Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) |
| 245 | return self.gamma * (x * Nx) + self.beta + x |
| 246 | |
| 247 | |
| 248 | # ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py |
| 249 | # ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108 |
| 250 | |
| 251 | |
| 252 | class ConvNeXtV2Block(nn.Module): |
| 253 | def __init__( |
| 254 | self, |
| 255 | dim: int, |
| 256 | intermediate_dim: int, |
| 257 | dilation: int = 1, |
| 258 | ): |
| 259 | super().__init__() |
| 260 | padding = (dilation * (7 - 1)) // 2 |
| 261 | self.dwconv = nn.Conv1d( |
| 262 | dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation |
| 263 | ) # depthwise conv |
| 264 | self.norm = nn.LayerNorm(dim, eps=1e-6) |
| 265 | self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers |
| 266 | self.act = nn.GELU() |
| 267 | self.grn = GRN(intermediate_dim) |
| 268 | self.pwconv2 = nn.Linear(intermediate_dim, dim) |
| 269 | |
| 270 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 271 | residual = x |
| 272 | x = x.transpose(1, 2) # b n d -> b d n |
| 273 | x = self.dwconv(x) |
| 274 | x = x.transpose(1, 2) # b d n -> b n d |
| 275 | x = self.norm(x) |
| 276 | x = self.pwconv1(x) |
| 277 | x = self.act(x) |
| 278 | x = self.grn(x) |
| 279 | x = self.pwconv2(x) |
| 280 | return residual + x |
| 281 | |
| 282 | |
| 283 | # RMSNorm |
| 284 | |
| 285 | |
| 286 | class RMSNorm(nn.Module): |
| 287 | def __init__(self, dim: int, eps: float): |
| 288 | super().__init__() |
| 289 | self.eps = eps |
| 290 | self.weight = nn.Parameter(torch.ones(dim)) |
| 291 | self.native_rms_norm = float(torch.__version__[:3]) >= 2.4 |
| 292 | |
| 293 | def forward(self, x): |
| 294 | if self.native_rms_norm: |
| 295 | if self.weight.dtype in [torch.float16, torch.bfloat16]: |
| 296 | x = x.to(self.weight.dtype) |
| 297 | x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps) |
| 298 | else: |
| 299 | variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True) |
| 300 | x = x * torch.rsqrt(variance + self.eps) |
| 301 | if self.weight.dtype in [torch.float16, torch.bfloat16]: |
| 302 | x = x.to(self.weight.dtype) |
| 303 | x = x * self.weight |
| 304 | |
| 305 | return x |
| 306 | |
| 307 | |
| 308 | # AdaLayerNorm |
| 309 | # return with modulated x for attn input, and params for later mlp modulation |
| 310 | |
| 311 | |
| 312 | class AdaLayerNorm(nn.Module): |
| 313 | def __init__(self, dim): |
| 314 | super().__init__() |
| 315 | |
| 316 | self.silu = nn.SiLU() |
| 317 | self.linear = nn.Linear(dim, dim * 6) |
| 318 | |
| 319 | self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 320 | |
| 321 | def forward(self, x, emb=None): |
| 322 | emb = self.linear(self.silu(emb)) |
| 323 | shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1) |
| 324 | |
| 325 | x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] |
| 326 | return x, gate_msa, shift_mlp, scale_mlp, gate_mlp |
| 327 | |
| 328 | |
| 329 | # AdaLayerNorm for final layer |
| 330 | # return only with modulated x for attn input, cuz no more mlp modulation |
| 331 | |
| 332 | |
| 333 | class AdaLayerNorm_Final(nn.Module): |
| 334 | def __init__(self, dim): |
| 335 | super().__init__() |
| 336 | |
| 337 | self.silu = nn.SiLU() |
| 338 | self.linear = nn.Linear(dim, dim * 2) |
| 339 | |
| 340 | self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 341 | |
| 342 | def forward(self, x, emb): |
| 343 | emb = self.linear(self.silu(emb)) |
| 344 | scale, shift = torch.chunk(emb, 2, dim=1) |
| 345 | |
| 346 | x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] |
| 347 | return x |
| 348 | |
| 349 | |
| 350 | # FeedForward |
| 351 | |
| 352 | |
| 353 | class FeedForward(nn.Module): |
| 354 | def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"): |
| 355 | super().__init__() |
| 356 | inner_dim = int(dim * mult) |
| 357 | dim_out = dim_out if dim_out is not None else dim |
| 358 | |
| 359 | activation = nn.GELU(approximate=approximate) |
| 360 | project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) |
| 361 | self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)) |
| 362 | |
| 363 | def forward(self, x): |
| 364 | return self.ff(x) |
| 365 | |
| 366 | |
| 367 | # Attention with possible joint part |
| 368 | # modified from diffusers/src/diffusers/models/attention_processor.py |
| 369 | |
| 370 | |
| 371 | class Attention(nn.Module): |
| 372 | def __init__( |
| 373 | self, |
| 374 | processor: JointAttnProcessor | AttnProcessor, |
| 375 | dim: int, |
| 376 | heads: int = 8, |
| 377 | dim_head: int = 64, |
| 378 | dropout: float = 0.0, |
| 379 | context_dim: Optional[int] = None, # if not None -> joint attention |
| 380 | context_pre_only: bool = False, |
| 381 | qk_norm: Optional[str] = None, |
| 382 | ): |
| 383 | super().__init__() |
| 384 | |
| 385 | if not hasattr(F, "scaled_dot_product_attention"): |
| 386 | raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") |
| 387 | |
| 388 | self.processor = processor |
| 389 | |
| 390 | self.dim = dim |
| 391 | self.heads = heads |
| 392 | self.inner_dim = dim_head * heads |
| 393 | self.dropout = dropout |
| 394 | |
| 395 | self.context_dim = context_dim |
| 396 | self.context_pre_only = context_pre_only |
| 397 | |
| 398 | self.to_q = nn.Linear(dim, self.inner_dim) |
| 399 | self.to_k = nn.Linear(dim, self.inner_dim) |
| 400 | self.to_v = nn.Linear(dim, self.inner_dim) |
| 401 | |
| 402 | if qk_norm is None: |
| 403 | self.q_norm = None |
| 404 | self.k_norm = None |
| 405 | elif qk_norm == "rms_norm": |
| 406 | self.q_norm = RMSNorm(dim_head, eps=1e-6) |
| 407 | self.k_norm = RMSNorm(dim_head, eps=1e-6) |
| 408 | else: |
| 409 | raise ValueError(f"Unimplemented qk_norm: {qk_norm}") |
| 410 | |
| 411 | if self.context_dim is not None: |
| 412 | self.to_q_c = nn.Linear(context_dim, self.inner_dim) |
| 413 | self.to_k_c = nn.Linear(context_dim, self.inner_dim) |
| 414 | self.to_v_c = nn.Linear(context_dim, self.inner_dim) |
| 415 | if qk_norm is None: |
| 416 | self.c_q_norm = None |
| 417 | self.c_k_norm = None |
| 418 | elif qk_norm == "rms_norm": |
| 419 | self.c_q_norm = RMSNorm(dim_head, eps=1e-6) |
| 420 | self.c_k_norm = RMSNorm(dim_head, eps=1e-6) |
| 421 | |
| 422 | self.to_out = nn.ModuleList([]) |
| 423 | self.to_out.append(nn.Linear(self.inner_dim, dim)) |
| 424 | self.to_out.append(nn.Dropout(dropout)) |
| 425 | |
| 426 | if self.context_dim is not None and not self.context_pre_only: |
| 427 | self.to_out_c = nn.Linear(self.inner_dim, context_dim) |
| 428 | |
| 429 | def forward( |
| 430 | self, |
| 431 | x: float["b n d"], # noised input x |
| 432 | c: float["b n d"] = None, # context c |
| 433 | mask: bool["b n"] | None = None, |
| 434 | rope=None, # rotary position embedding for x |
| 435 | c_rope=None, # rotary position embedding for c |
| 436 | c_mask: bool["b nt"] | None = None, # text mask |
| 437 | ) -> torch.Tensor: |
| 438 | if c is not None: |
| 439 | return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask) |
| 440 | else: |
| 441 | return self.processor(self, x, mask=mask, rope=rope) |
| 442 | |
| 443 | |
| 444 | # Attention processor |
| 445 | |
| 446 | if is_package_available("flash_attn"): |
| 447 | from flash_attn import flash_attn_func, flash_attn_varlen_func |
| 448 | from flash_attn.bert_padding import pad_input, unpad_input |
| 449 | |
| 450 | |
| 451 | class AttnProcessor: |
| 452 | def __init__( |
| 453 | self, |
| 454 | pe_attn_head: int | None = None, # number of attention head to apply rope, None for all |
| 455 | attn_backend: str = "torch", # "torch" or "flash_attn" |
| 456 | attn_mask_enabled: bool = True, |
| 457 | ): |
| 458 | if attn_backend == "flash_attn": |
| 459 | assert is_package_available("flash_attn"), "Please install flash-attn first." |
| 460 | if attn_backend == "torch" and attn_mask_enabled: |
| 461 | warnings.warn( |
| 462 | "attn_mask_enabled=True with attn_backend='torch' can consume large GPU memory. " |
| 463 | "Please switch attn_backend to 'flash_attn'.", |
| 464 | UserWarning, |
| 465 | ) |
| 466 | |
| 467 | self.pe_attn_head = pe_attn_head |
| 468 | self.attn_backend = attn_backend |
| 469 | self.attn_mask_enabled = attn_mask_enabled |
| 470 | |
| 471 | def __call__( |
| 472 | self, |
| 473 | attn: Attention, |
| 474 | x: float["b n d"], # noised input x |
| 475 | mask: bool["b n"] | None = None, |
| 476 | rope=None, # rotary position embedding |
| 477 | ) -> torch.FloatTensor: |
| 478 | batch_size = x.shape[0] |
| 479 | |
| 480 | # `sample` projections |
| 481 | query = attn.to_q(x) |
| 482 | key = attn.to_k(x) |
| 483 | value = attn.to_v(x) |
| 484 | |
| 485 | # attention |
| 486 | inner_dim = key.shape[-1] |
| 487 | head_dim = inner_dim // attn.heads |
| 488 | query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 489 | key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 490 | value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 491 | |
| 492 | # qk norm |
| 493 | if attn.q_norm is not None: |
| 494 | query = attn.q_norm(query) |
| 495 | if attn.k_norm is not None: |
| 496 | key = attn.k_norm(key) |
| 497 | |
| 498 | # apply rotary position embedding |
| 499 | if rope is not None: |
| 500 | freqs, xpos_scale = rope |
| 501 | q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) |
| 502 | |
| 503 | if self.pe_attn_head is not None: |
| 504 | pn = self.pe_attn_head |
| 505 | query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale) |
| 506 | key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale) |
| 507 | else: |
| 508 | query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) |
| 509 | key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) |
| 510 | |
| 511 | if self.attn_backend == "torch": |
| 512 | # mask. e.g. inference got a batch with different target durations, mask out the padding |
| 513 | if self.attn_mask_enabled and mask is not None: |
| 514 | attn_mask = mask |
| 515 | attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n' |
| 516 | attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2]) |
| 517 | else: |
| 518 | attn_mask = None |
| 519 | x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) |
| 520 | x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) |
| 521 | |
| 522 | elif self.attn_backend == "flash_attn": |
| 523 | query = query.transpose(1, 2) # [b, h, n, d] -> [b, n, h, d] |
| 524 | key = key.transpose(1, 2) |
| 525 | value = value.transpose(1, 2) |
| 526 | if self.attn_mask_enabled and mask is not None: |
| 527 | query, indices, q_cu_seqlens, q_max_seqlen_in_batch, _ = unpad_input(query, mask) |
| 528 | key, _, k_cu_seqlens, k_max_seqlen_in_batch, _ = unpad_input(key, mask) |
| 529 | value, _, _, _, _ = unpad_input(value, mask) |
| 530 | x = flash_attn_varlen_func( |
| 531 | query, |
| 532 | key, |
| 533 | value, |
| 534 | q_cu_seqlens, |
| 535 | k_cu_seqlens, |
| 536 | q_max_seqlen_in_batch, |
| 537 | k_max_seqlen_in_batch, |
| 538 | ) |
| 539 | x = pad_input(x, indices, batch_size, q_max_seqlen_in_batch) |
| 540 | x = x.reshape(batch_size, -1, attn.heads * head_dim) |
| 541 | else: |
| 542 | x = flash_attn_func(query, key, value, dropout_p=0.0, causal=False) |
| 543 | x = x.reshape(batch_size, -1, attn.heads * head_dim) |
| 544 | |
| 545 | x = x.to(query.dtype) |
| 546 | |
| 547 | # linear proj |
| 548 | x = attn.to_out[0](x) |
| 549 | # dropout |
| 550 | x = attn.to_out[1](x) |
| 551 | |
| 552 | if mask is not None: |
| 553 | mask = mask.unsqueeze(-1) |
| 554 | x = x.masked_fill(~mask, 0.0) |
| 555 | |
| 556 | return x |
| 557 | |
| 558 | |
| 559 | # Joint Attention processor for MM-DiT |
| 560 | # modified from diffusers/src/diffusers/models/attention_processor.py |
| 561 | |
| 562 | |
| 563 | class JointAttnProcessor: |
| 564 | def __init__( |
| 565 | self, |
| 566 | attn_backend: str = "torch", # "torch" or "flash_attn" |
| 567 | attn_mask_enabled: bool = True, |
| 568 | ): |
| 569 | if attn_backend == "flash_attn": |
| 570 | assert is_package_available("flash_attn"), "Please install flash-attn first." |
| 571 | if attn_backend == "torch" and attn_mask_enabled: |
| 572 | warnings.warn( |
| 573 | "attn_mask_enabled=True with attn_backend='torch' can consume large GPU memory. " |
| 574 | "Please switch attn_backend to 'flash_attn'.", |
| 575 | UserWarning, |
| 576 | ) |
| 577 | |
| 578 | self.attn_backend = attn_backend |
| 579 | self.attn_mask_enabled = attn_mask_enabled |
| 580 | |
| 581 | def __call__( |
| 582 | self, |
| 583 | attn: Attention, |
| 584 | x: float["b n d"], # noised input x |
| 585 | c: float["b nt d"] = None, # context c, here text |
| 586 | mask: bool["b n"] | None = None, |
| 587 | rope=None, # rotary position embedding for x |
| 588 | c_rope=None, # rotary position embedding for c |
| 589 | c_mask: bool["b nt"] | None = None, # text mask |
| 590 | ) -> torch.FloatTensor: |
| 591 | residual = x |
| 592 | audio_mask = mask |
| 593 | |
| 594 | batch_size = c.shape[0] |
| 595 | |
| 596 | # `sample` projections |
| 597 | query = attn.to_q(x) |
| 598 | key = attn.to_k(x) |
| 599 | value = attn.to_v(x) |
| 600 | |
| 601 | # `context` projections |
| 602 | c_query = attn.to_q_c(c) |
| 603 | c_key = attn.to_k_c(c) |
| 604 | c_value = attn.to_v_c(c) |
| 605 | |
| 606 | # attention |
| 607 | inner_dim = key.shape[-1] |
| 608 | head_dim = inner_dim // attn.heads |
| 609 | query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 610 | key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 611 | value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 612 | c_query = c_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 613 | c_key = c_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 614 | c_value = c_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) |
| 615 | |
| 616 | # qk norm |
| 617 | if attn.q_norm is not None: |
| 618 | query = attn.q_norm(query) |
| 619 | if attn.k_norm is not None: |
| 620 | key = attn.k_norm(key) |
| 621 | if attn.c_q_norm is not None: |
| 622 | c_query = attn.c_q_norm(c_query) |
| 623 | if attn.c_k_norm is not None: |
| 624 | c_key = attn.c_k_norm(c_key) |
| 625 | |
| 626 | # apply rope for context and noised input independently |
| 627 | if rope is not None: |
| 628 | freqs, xpos_scale = rope |
| 629 | q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) |
| 630 | query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) |
| 631 | key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) |
| 632 | if c_rope is not None: |
| 633 | freqs, xpos_scale = c_rope |
| 634 | q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) |
| 635 | c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale) |
| 636 | c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale) |
| 637 | |
| 638 | # joint attention |
| 639 | query = torch.cat([query, c_query], dim=2) |
| 640 | key = torch.cat([key, c_key], dim=2) |
| 641 | value = torch.cat([value, c_value], dim=2) |
| 642 | |
| 643 | # build combined mask for joint attention: audio mask + text mask |
| 644 | if self.attn_mask_enabled and mask is not None: |
| 645 | if c_mask is not None: |
| 646 | mask = torch.cat([mask, c_mask], dim=1) |
| 647 | else: |
| 648 | mask = F.pad(mask, (0, c.shape[1]), value=True) |
| 649 | |
| 650 | if self.attn_backend == "torch": |
| 651 | # mask. e.g. inference got a batch with different target durations, mask out the padding |
| 652 | if self.attn_mask_enabled and mask is not None: |
| 653 | attn_mask = mask |
| 654 | attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n' |
| 655 | attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2]) |
| 656 | else: |
| 657 | attn_mask = None |
| 658 | x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) |
| 659 | x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) |
| 660 | |
| 661 | elif self.attn_backend == "flash_attn": |
| 662 | query = query.transpose(1, 2) # [b, h, n, d] -> [b, n, h, d] |
| 663 | key = key.transpose(1, 2) |
| 664 | value = value.transpose(1, 2) |
| 665 | if self.attn_mask_enabled and mask is not None: |
| 666 | total_seq_len = query.shape[1] |
| 667 | query, indices, q_cu_seqlens, q_max_seqlen_in_batch, _ = unpad_input(query, mask) |
| 668 | key, _, k_cu_seqlens, k_max_seqlen_in_batch, _ = unpad_input(key, mask) |
| 669 | value, _, _, _, _ = unpad_input(value, mask) |
| 670 | x = flash_attn_varlen_func( |
| 671 | query, |
| 672 | key, |
| 673 | value, |
| 674 | q_cu_seqlens, |
| 675 | k_cu_seqlens, |
| 676 | q_max_seqlen_in_batch, |
| 677 | k_max_seqlen_in_batch, |
| 678 | ) |
| 679 | x = pad_input(x, indices, batch_size, total_seq_len) |
| 680 | x = x.reshape(batch_size, -1, attn.heads * head_dim) |
| 681 | else: |
| 682 | x = flash_attn_func(query, key, value, dropout_p=0.0, causal=False) |
| 683 | x = x.reshape(batch_size, -1, attn.heads * head_dim) |
| 684 | |
| 685 | x = x.to(query.dtype) |
| 686 | |
| 687 | # Split the attention outputs. |
| 688 | x, c = ( |
| 689 | x[:, : residual.shape[1]], |
| 690 | x[:, residual.shape[1] :], |
| 691 | ) |
| 692 | |
| 693 | # linear proj |
| 694 | x = attn.to_out[0](x) |
| 695 | # dropout |
| 696 | x = attn.to_out[1](x) |
| 697 | if not attn.context_pre_only: |
| 698 | c = attn.to_out_c(c) |
| 699 | |
| 700 | if audio_mask is not None: |
| 701 | x = x.masked_fill(~audio_mask.unsqueeze(-1), 0.0) |
| 702 | if c_mask is not None: |
| 703 | c = c.masked_fill(~c_mask.unsqueeze(-1), 0.0) |
| 704 | |
| 705 | return x, c |
| 706 | |
| 707 | |
| 708 | # DiT Block |
| 709 | |
| 710 | |
| 711 | class DiTBlock(nn.Module): |
| 712 | def __init__( |
| 713 | self, |
| 714 | dim, |
| 715 | heads, |
| 716 | dim_head, |
| 717 | ff_mult=4, |
| 718 | dropout=0.1, |
| 719 | qk_norm=None, |
| 720 | pe_attn_head=None, |
| 721 | attn_backend="torch", # "torch" or "flash_attn" |
| 722 | attn_mask_enabled=True, |
| 723 | ): |
| 724 | super().__init__() |
| 725 | |
| 726 | self.attn_norm = AdaLayerNorm(dim) |
| 727 | self.attn = Attention( |
| 728 | processor=AttnProcessor( |
| 729 | pe_attn_head=pe_attn_head, |
| 730 | attn_backend=attn_backend, |
| 731 | attn_mask_enabled=attn_mask_enabled, |
| 732 | ), |
| 733 | dim=dim, |
| 734 | heads=heads, |
| 735 | dim_head=dim_head, |
| 736 | dropout=dropout, |
| 737 | qk_norm=qk_norm, |
| 738 | ) |
| 739 | |
| 740 | self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 741 | self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") |
| 742 | |
| 743 | def forward(self, x, t, mask=None, rope=None): # x: noised input, t: time embedding |
| 744 | # pre-norm & modulation for attention input |
| 745 | norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) |
| 746 | |
| 747 | # attention |
| 748 | attn_output = self.attn(x=norm, mask=mask, rope=rope) |
| 749 | |
| 750 | # process attention output for input x |
| 751 | x = x + gate_msa.unsqueeze(1) * attn_output |
| 752 | |
| 753 | norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] |
| 754 | ff_output = self.ff(norm) |
| 755 | x = x + gate_mlp.unsqueeze(1) * ff_output |
| 756 | |
| 757 | return x |
| 758 | |
| 759 | |
| 760 | # MMDiT Block https://arxiv.org/abs/2403.03206 |
| 761 | |
| 762 | |
| 763 | class MMDiTBlock(nn.Module): |
| 764 | r""" |
| 765 | modified from diffusers/src/diffusers/models/attention.py |
| 766 | |
| 767 | notes. |
| 768 | _c: context related. text, cond, etc. (left part in sd3 fig2.b) |
| 769 | _x: noised input related. (right part) |
| 770 | context_pre_only: last layer only do prenorm + modulation cuz no more ffn |
| 771 | """ |
| 772 | |
| 773 | def __init__( |
| 774 | self, |
| 775 | dim, |
| 776 | heads, |
| 777 | dim_head, |
| 778 | ff_mult=4, |
| 779 | dropout=0.1, |
| 780 | context_dim=None, |
| 781 | context_pre_only=False, |
| 782 | qk_norm=None, |
| 783 | attn_backend="torch", |
| 784 | attn_mask_enabled=False, |
| 785 | ): |
| 786 | super().__init__() |
| 787 | if context_dim is None: |
| 788 | context_dim = dim |
| 789 | self.context_pre_only = context_pre_only |
| 790 | |
| 791 | self.attn_norm_c = AdaLayerNorm_Final(context_dim) if context_pre_only else AdaLayerNorm(context_dim) |
| 792 | self.attn_norm_x = AdaLayerNorm(dim) |
| 793 | self.attn = Attention( |
| 794 | processor=JointAttnProcessor( |
| 795 | attn_backend=attn_backend, |
| 796 | attn_mask_enabled=attn_mask_enabled, |
| 797 | ), |
| 798 | dim=dim, |
| 799 | heads=heads, |
| 800 | dim_head=dim_head, |
| 801 | dropout=dropout, |
| 802 | context_dim=context_dim, |
| 803 | context_pre_only=context_pre_only, |
| 804 | qk_norm=qk_norm, |
| 805 | ) |
| 806 | |
| 807 | if not context_pre_only: |
| 808 | self.ff_norm_c = nn.LayerNorm(context_dim, elementwise_affine=False, eps=1e-6) |
| 809 | self.ff_c = FeedForward(dim=context_dim, mult=ff_mult, dropout=dropout, approximate="tanh") |
| 810 | else: |
| 811 | self.ff_norm_c = None |
| 812 | self.ff_c = None |
| 813 | self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 814 | self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") |
| 815 | |
| 816 | def forward( |
| 817 | self, x, c, t, mask=None, rope=None, c_rope=None, c_mask=None |
| 818 | ): # x: noised input, c: context, t: time embedding |
| 819 | # pre-norm & modulation for attention input |
| 820 | if self.context_pre_only: |
| 821 | norm_c = self.attn_norm_c(c, t) |
| 822 | else: |
| 823 | norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t) |
| 824 | norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t) |
| 825 | |
| 826 | # attention |
| 827 | x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask) |
| 828 | |
| 829 | # process attention output for context c |
| 830 | if self.context_pre_only: |
| 831 | c = None |
| 832 | else: # if not last layer |
| 833 | c = c + c_gate_msa.unsqueeze(1) * c_attn_output |
| 834 | |
| 835 | norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] |
| 836 | c_ff_output = self.ff_c(norm_c) |
| 837 | c = c + c_gate_mlp.unsqueeze(1) * c_ff_output |
| 838 | |
| 839 | # process attention output for input x |
| 840 | x = x + x_gate_msa.unsqueeze(1) * x_attn_output |
| 841 | |
| 842 | norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None] |
| 843 | x_ff_output = self.ff_x(norm_x) |
| 844 | x = x + x_gate_mlp.unsqueeze(1) * x_ff_output |
| 845 | |
| 846 | return c, x |
| 847 | |
| 848 | |
| 849 | # time step conditioning embedding |
| 850 | |
| 851 | |
| 852 | class TimestepEmbedding(nn.Module): |
| 853 | def __init__(self, dim, freq_embed_dim=256): |
| 854 | super().__init__() |
| 855 | self.time_embed = SinusPositionEmbedding(freq_embed_dim) |
| 856 | self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) |
| 857 | |
| 858 | def forward(self, timestep: float["b"]): |
| 859 | time_hidden = self.time_embed(timestep) |
| 860 | time_hidden = time_hidden.to(timestep.dtype) |
| 861 | time = self.time_mlp(time_hidden) # b d |
| 862 | return time |
| 863 |