choritogabriel/echo-tts-2
2
1# SPDX-License-Identifier: Apache-2.02 3# This file contains portions adapted from:4# • Descript Audio Codec (DAC) — MIT License (full text appended below)5# • Fish-Speech S1 DAC Autoencoder — reference implementation (Apache-2.0 / CC-BY-NC),6# rewritten here in a single-file Torch module for interoperability and transparency.7#8# OVERALL LICENSE (this file): Apache-2.0, except where explicitly marked:9# # SPDX-License-Identifier: MIT10# Keep these notices and the embedded MIT text if you redistribute this file.11 12# NOTE13# Self-contained autoencoder implementation of Fish-S1-DAC (inlining DAC code to avoid dependencies).14# Code in this module has been largely copy-and-pasted from the Fish-S1-DAC and DAC repositories,15# and refactored with help from ChatGPT/Claude (these models also helped with licensing).16# Thus, it differs stylistically from the rest of the codebase (and is likely internally inconsistent as well).17 18from __future__ import annotations19 20import math21from dataclasses import dataclass22from typing import List, Optional, Tuple, Union23 24import numpy as np25import torch26from torch import Tensor, nn27from torch.nn import functional as F28from torch.nn.utils.parametrizations import weight_norm29from torch.nn.utils.parametrize import remove_parametrizations30 31from einops import rearrange32 33 34# --------------------------------------------------------------------35# Shared helpers36# --------------------------------------------------------------------37 38def find_multiple(n: int, k: int) -> int:39 return n if n % k == 0 else n + k - (n % k)40 41def unpad1d(x: Tensor, paddings: Tuple[int, int]) -> Tensor:42 """Remove padding from x, handling properly zero padding. Only for 1d!"""43 padding_left, padding_right = paddings44 assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)45 assert (padding_left + padding_right) <= x.shape[-1]46 end = x.shape[-1] - padding_right47 return x[..., padding_left:end]48 49def get_extra_padding_for_conv1d(50 x: Tensor, kernel_size: int, stride: int, padding_total: int = 051) -> int:52 """See pad_for_conv1d; enough right pad so striding evenly covers length."""53 length = x.shape[-1]54 n_frames = (length - kernel_size + padding_total) / stride + 155 ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)56 return ideal_length - length57 58def pad1d(59 x: Tensor,60 paddings: Tuple[int, int],61 mode: str = "zeros",62 value: float = 0.0,63) -> Tensor:64 """65 Reflect‑safe 1D pad: if reflect would underflow on small inputs, insert66 temporary right zero-pad before reflecting.67 """68 length = x.shape[-1]69 padding_left, padding_right = paddings70 assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)71 if mode == "reflect":72 max_pad = max(padding_left, padding_right)73 extra_pad = 074 if length <= max_pad:75 extra_pad = max_pad - length + 176 x = F.pad(x, (0, extra_pad))77 padded = F.pad(x, (padding_left, padding_right), mode, value)78 end = padded.shape[-1] - extra_pad79 return padded[..., :end]80 else:81 return F.pad(x, (padding_left, padding_right), mode, value)82 83 84# --------------------------------------------------------------------85# DAC Layers (adapted) — MIT86# Original: https://github.com/descriptinc/descript-audio-codec/blob/main/dac/nn/layers.py87# SPDX-License-Identifier: MIT88# --------------------------------------------------------------------89 90def WNConv1d(*args, **kwargs):91 return weight_norm(nn.Conv1d(*args, **kwargs))92 93def WNConvTranspose1d(*args, **kwargs):94 return weight_norm(nn.ConvTranspose1d(*args, **kwargs))95 96@torch.jit.script97def snake(x: Tensor, alpha: Tensor) -> Tensor:98 shape = x.shape99 x = x.reshape(shape[0], shape[1], -1)100 x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)101 x = x.reshape(shape)102 return x103 104class Snake1d(nn.Module):105 def __init__(self, channels: int):106 super().__init__()107 self.alpha = nn.Parameter(torch.ones(1, channels, 1))108 def forward(self, x: Tensor) -> Tensor:109 return snake(x, self.alpha)110 111# --------------------------------------------------------------------112# DAC Vector Quantize (adapted) — MIT113# Original: https://github.com/descriptinc/descript-audio-codec/blob/main/dac/nn/quantize.py114# SPDX-License-Identifier: MIT115# --------------------------------------------------------------------116 117class VectorQuantize(nn.Module):118 """119 VQ with factorized, l2-normalized codes (ViT‑VQGAN style).120 I/O in (B, D, T).121 """122 def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):123 super().__init__()124 self.codebook_size = codebook_size125 self.codebook_dim = codebook_dim126 self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)127 self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)128 self.codebook = nn.Embedding(codebook_size, codebook_dim)129 130 def forward(self, z: Tensor):131 z_e = self.in_proj(z) # (B, D, T)132 z_q, indices = self.decode_latents(z_e)133 commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])134 codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])135 z_q = z_e + (z_q - z_e).detach() # straight‑through136 z_q = self.out_proj(z_q)137 return z_q, commitment_loss, codebook_loss, indices, z_e138 139 def embed_code(self, embed_id: Tensor) -> Tensor:140 return F.embedding(embed_id, self.codebook.weight)141 142 def decode_code(self, embed_id: Tensor) -> Tensor:143 return self.embed_code(embed_id).transpose(1, 2)144 145 def decode_latents(self, latents: Tensor) -> Tuple[Tensor, Tensor]:146 encodings = rearrange(latents, "b d t -> (b t) d")147 codebook = self.codebook.weight148 encodings = F.normalize(encodings)149 codebook = F.normalize(codebook)150 dist = (151 encodings.pow(2).sum(1, keepdim=True)152 - 2 * encodings @ codebook.t()153 + codebook.pow(2).sum(1, keepdim=True).t()154 )155 indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))156 z_q = self.decode_code(indices)157 return z_q, indices158 159 160class ResidualVectorQuantize(nn.Module):161 """SoundStream-style residual VQ stack."""162 def __init__(163 self,164 input_dim: int = 512,165 n_codebooks: int = 9,166 codebook_size: int = 1024,167 codebook_dim: Union[int, List[int]] = 8,168 quantizer_dropout: float = 0.0,169 ):170 super().__init__()171 if isinstance(codebook_dim, int):172 codebook_dim = [codebook_dim for _ in range(n_codebooks)]173 174 self.n_codebooks = n_codebooks175 self.codebook_dim = codebook_dim176 self.codebook_size = codebook_size177 178 self.quantizers = nn.ModuleList([179 VectorQuantize(input_dim, codebook_size, codebook_dim[i])180 for i in range(n_codebooks)181 ])182 self.quantizer_dropout = quantizer_dropout183 184 def forward(self, z: Tensor, n_quantizers: Optional[int] = None):185 z_q = 0186 residual = z187 commitment_loss = 0188 codebook_loss = 0189 190 codebook_indices = []191 latents = []192 193 if n_quantizers is None:194 n_quantizers = self.n_codebooks195 if self.training:196 n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1197 dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],))198 n_dropout = int(z.shape[0] * self.quantizer_dropout)199 n_quantizers[:n_dropout] = dropout[:n_dropout]200 n_quantizers = n_quantizers.to(z.device)201 202 for i, quantizer in enumerate(self.quantizers):203 if self.training is False and i >= n_quantizers:204 break205 206 z_q_i, commit_i, codebk_i, indices_i, z_e_i = quantizer(residual)207 208 mask = (torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers)209 z_q = z_q + z_q_i * mask[:, None, None]210 residual = residual - z_q_i211 212 commitment_loss += (commit_i * mask).mean()213 codebook_loss += (codebk_i * mask).mean()214 215 codebook_indices.append(indices_i)216 latents.append(z_e_i)217 218 codes = torch.stack(codebook_indices, dim=1)219 latents = torch.cat(latents, dim=1)220 221 return z_q, codes, latents, commitment_loss, codebook_loss222 223 def from_codes(self, codes: Tensor) -> Tuple[Tensor, Tensor, Tensor]:224 z_q = 0.0225 z_p = []226 n_codebooks = codes.shape[1]227 for i in range(n_codebooks):228 z_p_i = self.quantizers[i].decode_code(codes[:, i, :])229 z_p.append(z_p_i)230 z_q_i = self.quantizers[i].out_proj(z_p_i)231 z_q = z_q + z_q_i232 return z_q, torch.cat(z_p, dim=1), codes233 234 def from_latents(self, latents: Tensor) -> Tuple[Tensor, Tensor, Tensor]:235 z_q = 0236 z_p = []237 codes = []238 dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers])239 n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[0]240 for i in range(n_codebooks):241 j, k = dims[i], dims[i + 1]242 z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])243 z_p.append(z_p_i)244 codes.append(codes_i)245 z_q_i = self.quantizers[i].out_proj(z_p_i)246 z_q = z_q + z_q_i247 return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)248 249 250# --------------------------------------------------------------------251# S1 DAC rvq252# --------------------------------------------------------------------253 254@dataclass255class VQResult:256 z: Tensor257 codes: Tensor258 latents: Tensor259 codebook_loss: Tensor260 commitment_loss: Tensor261 semantic_distill_z: Optional[Tensor] = None262 263 264class CausalConvNet(nn.Module):265 def __init__(266 self,267 in_channels,268 out_channels,269 kernel_size,270 dilation=1,271 stride=1,272 groups=1,273 padding=None,274 ):275 super().__init__()276 self.conv = nn.Conv1d(277 in_channels, out_channels, kernel_size,278 stride=stride, dilation=dilation, groups=groups,279 )280 self.stride = stride281 self.kernel_size = (kernel_size - 1) * dilation + 1282 self.dilation = dilation283 self.padding = self.kernel_size - self.stride284 285 def forward(self, x: Tensor) -> Tensor:286 pad = self.padding287 extra = get_extra_padding_for_conv1d(x, self.kernel_size, self.stride, pad)288 x = pad1d(x, (pad, extra), mode="constant", value=0)289 return self.conv(x).contiguous()290 291 def weight_norm(self, name="weight", dim=0):292 self.conv = weight_norm(self.conv, name=name, dim=dim)293 return self294 295 def remove_weight_norm(self):296 self.conv = remove_parametrizations(self.conv)297 return self298 299 300class CausalTransConvNet(nn.Module):301 def __init__(self, in_channels, out_channels, kernel_size, dilation=1, stride=1, padding=None):302 super().__init__()303 self.conv = nn.ConvTranspose1d(304 in_channels, out_channels, kernel_size,305 stride=stride, dilation=dilation306 )307 self.stride = stride308 self.kernel_size = kernel_size309 310 def forward(self, x: Tensor) -> Tensor:311 x = self.conv(x)312 pad = self.kernel_size - self.stride313 padding_right = math.ceil(pad)314 padding_left = pad - padding_right315 x = unpad1d(x, (padding_left, padding_right))316 return x.contiguous()317 318 def weight_norm(self, name="weight", dim=0):319 self.conv = weight_norm(self.conv, name=name, dim=dim)320 return self321 322 def remove_weight_norm(self):323 self.conv = remove_parametrizations(self.conv)324 return self325 326 327def CausalWNConv1d(*args, **kwargs):328 return CausalConvNet(*args, **kwargs).weight_norm()329 330def CausalWNConvTranspose1d(*args, **kwargs):331 return CausalTransConvNet(*args, **kwargs).weight_norm()332 333class ConvNeXtBlock(nn.Module):334 r"""ConvNeXt Block (1D).335 DwConv -> (N, C, L) → (N, L, C) -> LN -> Linear -> GELU -> Linear -> (N, C, L) with residual336 """337 def __init__(338 self,339 dim: int,340 layer_scale_init_value: float = 1e-6,341 mlp_ratio: float = 4.0,342 kernel_size: int = 7,343 dilation: int = 1,344 ):345 super().__init__()346 convnet_type = CausalConvNet347 self.dwconv = convnet_type(348 dim, dim, kernel_size=kernel_size,349 groups=dim, dilation=dilation,350 ) # depthwise conv351 self.norm = nn.LayerNorm(dim, eps=1e-6)352 self.pwconv1 = nn.Linear(dim, int(mlp_ratio * dim))353 self.act = nn.GELU()354 self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim)355 self.gamma = (356 nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)357 if layer_scale_init_value > 0 else None358 )359 360 def forward(self, x: Tensor, apply_residual: bool = True) -> Tensor:361 inp = x362 x = self.dwconv(x)363 x = x.permute(0, 2, 1) # (N, C, L) -> (N, L, C)364 x = self.norm(x)365 x = self.pwconv1(x)366 x = self.act(x)367 x = self.pwconv2(x)368 if self.gamma is not None:369 x = self.gamma * x370 x = x.permute(0, 2, 1) # (N, L, C) -> (N, C, L)371 if apply_residual:372 x = inp + x373 return x374 375 376class DownsampleResidualVectorQuantize(nn.Module):377 def __init__(378 self,379 input_dim: int = 1024,380 n_codebooks: int = 9,381 codebook_dim: int = 8,382 quantizer_dropout: float = 0.5,383 codebook_size: int = 1024,384 semantic_codebook_size: int = 4096,385 downsample_factor: Tuple[int, ...] = (2, 2),386 downsample_dims: Optional[Tuple[int, ...]] = None,387 pre_module: Optional[nn.Module] = None,388 post_module: Optional[nn.Module] = None,389 semantic_predictor_module: Optional[nn.Module] = None,390 ):391 super().__init__()392 393 if downsample_dims is None:394 downsample_dims = tuple(input_dim for _ in range(len(downsample_factor)))395 396 all_dims = (input_dim,) + tuple(downsample_dims)397 398 self.semantic_quantizer = ResidualVectorQuantize(399 input_dim=input_dim,400 n_codebooks=1,401 codebook_size=semantic_codebook_size,402 codebook_dim=codebook_dim,403 quantizer_dropout=0.0,404 )405 406 self.quantizer = ResidualVectorQuantize(407 input_dim=input_dim,408 n_codebooks=n_codebooks,409 codebook_size=codebook_size,410 codebook_dim=codebook_dim,411 quantizer_dropout=quantizer_dropout,412 )413 414 convnet_type = CausalConvNet415 transconvnet_type = CausalTransConvNet416 417 self.downsample = nn.Sequential(418 *[419 nn.Sequential(420 convnet_type(all_dims[idx], all_dims[idx + 1], kernel_size=factor, stride=factor),421 ConvNeXtBlock(dim=all_dims[idx + 1]),422 )423 for idx, factor in enumerate(downsample_factor)424 ]425 )426 427 self.upsample = nn.Sequential(428 *[429 nn.Sequential(430 transconvnet_type(all_dims[idx + 1], all_dims[idx], kernel_size=factor, stride=factor),431 ConvNeXtBlock(dim=all_dims[idx]),432 )433 for idx, factor in reversed(list(enumerate(downsample_factor)))434 ]435 )436 437 self.apply(self._init_weights)438 self.pre_module = pre_module if pre_module is not None else nn.Identity()439 self.post_module = post_module if post_module is not None else nn.Identity()440 self.semantic_predictor_module = (441 semantic_predictor_module if semantic_predictor_module is not None else nn.Identity()442 )443 444 @staticmethod445 def _init_weights(m):446 if isinstance(m, (nn.Conv1d, nn.Linear)):447 nn.init.trunc_normal_(m.weight, std=0.02)448 if getattr(m, "bias", None) is not None:449 nn.init.constant_(m.bias, 0)450 451 def forward(self, z: Tensor, n_quantizers: Optional[int] = None, semantic_len: Optional[Tensor] = None, **kwargs):452 # z: (B, D, T)453 original_shape = z.shape454 if semantic_len is None:455 semantic_len = torch.LongTensor([z.shape[-1]])456 457 z = self.downsample(z)458 z = self.pre_module(z) # (B, D, T) or (B, T, D) depending on module; original uses channels-first in/out459 460 semantic_z, semantic_codes, semantic_latents, semantic_commitment_loss, semantic_codebook_loss = \461 self.semantic_quantizer(z)462 residual_z = z - semantic_z463 residual_z, codes, latents, commitment_loss, codebook_loss = self.quantizer(residual_z, n_quantizers=n_quantizers)464 z = semantic_z + residual_z465 commitment_loss = commitment_loss + semantic_commitment_loss466 codebook_loss = codebook_loss + semantic_codebook_loss467 codes = torch.cat([semantic_codes, codes], dim=1)468 latents = torch.cat([semantic_latents, latents], dim=1)469 z = self.post_module(z)470 z = self.upsample(z)471 472 # Pad or crop z to match original shape (time dimension)473 diff = original_shape[-1] - z.shape[-1]474 right = 0475 left = abs(diff) - right476 if diff > 0:477 z = F.pad(z, (left, right))478 elif diff < 0:479 z = z[..., left:]480 481 return VQResult(482 z=z, codes=codes, latents=latents,483 commitment_loss=commitment_loss, codebook_loss=codebook_loss,484 )485 486 def decode(self, indices: Tensor) -> Tensor:487 new_indices = torch.zeros_like(indices)488 new_indices[:, 0] = torch.clamp(indices[:, 0], max=self.semantic_quantizer.codebook_size - 1)489 new_indices[:, 1:] = torch.clamp(indices[:, 1:], max=self.quantizer.codebook_size - 1)490 491 z_q_semantic = self.semantic_quantizer.from_codes(new_indices[:, :1])[0]492 z_q_residual = self.quantizer.from_codes(new_indices[:, 1:])[0]493 z_q = z_q_semantic + z_q_residual494 z_q = self.post_module(z_q)495 z_q = self.upsample(z_q)496 return z_q497 498 499# --------------------------------------------------------------------500# Transformer stack501# --------------------------------------------------------------------502 503@dataclass504class ModelArgs:505 block_size: int = 2048506 n_layer: int = 8507 n_head: int = 8508 dim: int = 512509 intermediate_size: int = 1536510 n_local_heads: int = -1511 head_dim: int = 64512 rope_base: float = 10000513 norm_eps: float = 1e-5514 dropout_rate: float = 0.1515 attn_dropout_rate: float = 0.1516 channels_first: bool = True # to be compatible with conv1d input/output517 pos_embed_type: str = "rope" # "rope" or "conformer"518 max_relative_position: int = 128519 520 def __post_init__(self):521 if self.n_local_heads == -1:522 self.n_local_heads = self.n_head523 if self.intermediate_size is None:524 hidden_dim = 4 * self.dim525 n_hidden = int(2 * hidden_dim / 3)526 self.intermediate_size = find_multiple(n_hidden, 256)527 assert self.pos_embed_type in ["rope", "conformer"]528 529 530class KVCache(nn.Module):531 def __init__(self, max_batch_size, max_seq_length, n_heads, head_dim, dtype=torch.bfloat16):532 super().__init__()533 cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim)534 self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype))535 self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype))536 537 def update(self, input_pos: Tensor, k_val: Tensor, v_val: Tensor):538 # input_pos: [S], k_val: [B, H, S, D]539 assert input_pos.shape[0] == k_val.shape[2]540 k_out = self.k_cache541 v_out = self.v_cache542 k_out[:, :, input_pos] = k_val543 v_out[:, :, input_pos] = v_val544 return (545 k_out[:, :, : input_pos.max() + 1, :],546 v_out[:, :, : input_pos.max() + 1, :],547 )548 549 def clear_cache(self, prompt_len: int):550 self.k_cache[:, :, prompt_len:, :].fill_(0)551 self.v_cache[:, :, prompt_len:, :].fill_(0)552 553 554class Transformer(nn.Module):555 def __init__(self, config: ModelArgs) -> None:556 super().__init__()557 self.config = config558 559 self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))560 self.norm = RMSNorm(config.dim, eps=config.norm_eps)561 562 if config.pos_embed_type == "rope":563 freqs_cis = precompute_freqs_cis(self.config.block_size, self.config.head_dim, self.config.rope_base)564 self.register_buffer("freqs_cis", freqs_cis)565 else:566 self.register_buffer("freqs_cis", None)567 568 causal_mask = torch.tril(torch.ones(self.config.block_size, self.config.block_size, dtype=torch.bool))569 self.register_buffer("causal_mask", causal_mask)570 571 self.max_batch_size = -1572 self.max_seq_length = -1573 self.use_kv_cache = False574 575 def setup_caches(self, max_batch_size, max_seq_length):576 head_dim = self.config.dim // self.config.n_head577 max_seq_length = find_multiple(max_seq_length, 8)578 self.max_seq_length = max_seq_length579 self.max_batch_size = max_batch_size580 dtype = self.norm.weight.dtype581 device = self.norm.weight.device582 583 for b in self.layers:584 b.attention.kv_cache = KVCache(585 max_batch_size, max_seq_length, self.config.n_local_heads, head_dim, dtype586 ).to(device)587 588 self.use_kv_cache = True589 590 def forward(self, x: Tensor, input_pos: Optional[Tensor] = None, mask: Optional[Tensor] = None) -> Tensor:591 if self.config.pos_embed_type == "rope":592 assert self.freqs_cis is not None593 freqs_cis = self.freqs_cis[input_pos]594 else:595 freqs_cis = None596 597 if mask is None:598 if not self.training and self.use_kv_cache:599 mask = self.causal_mask[None, None, input_pos]600 mask = mask[..., : input_pos.max() + 1]601 else:602 mask = self.causal_mask[None, None, input_pos]603 mask = mask[..., input_pos]604 605 for layer in self.layers:606 x = layer(x, input_pos, freqs_cis, mask)607 x = self.norm(x)608 return x609 610 611class TransformerBlock(nn.Module):612 def __init__(self, config: ModelArgs) -> None:613 super().__init__()614 self.attention = Attention(config)615 self.feed_forward = FeedForward(config)616 self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)617 self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)618 self.attention_layer_scale = LayerScale(config.dim, inplace=True)619 self.ffn_layer_scale = LayerScale(config.dim, inplace=True)620 621 def forward(self, x: Tensor, input_pos: Tensor, freqs_cis: Tensor, mask: Tensor) -> Tensor:622 h = x + self.attention_layer_scale(623 self.attention(self.attention_norm(x), freqs_cis, mask, input_pos)624 )625 out = h + self.ffn_layer_scale(self.feed_forward(self.ffn_norm(h)))626 return out627 628 629class Attention(nn.Module):630 def __init__(self, config: ModelArgs):631 super().__init__()632 assert config.dim % config.n_head == 0633 634 total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim635 self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)636 self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)637 self.kv_cache = None638 639 self.n_head = config.n_head640 self.head_dim = config.head_dim641 self.n_local_heads = config.n_local_heads642 self.dim = config.dim643 self.attn_dropout_rate = config.attn_dropout_rate644 self.pos_embed_type = config.pos_embed_type645 646 if self.pos_embed_type == "conformer":647 self.max_relative_position = config.max_relative_position648 num_pos_embeddings = 2 * config.max_relative_position + 1649 self.rel_pos_embeddings = nn.Parameter(torch.zeros(num_pos_embeddings, self.head_dim))650 nn.init.normal_(self.rel_pos_embeddings, mean=0.0, std=0.02)651 652 def _compute_conformer_pos_scores(self, q: Tensor, seqlen: int) -> Tensor:653 positions = torch.arange(seqlen, device=q.device)654 relative_positions = positions.unsqueeze(1) - positions.unsqueeze(0) # [S, S]655 relative_positions = torch.clamp(relative_positions + self.max_relative_position,656 0, 2 * self.max_relative_position)657 rel_embeddings = self.rel_pos_embeddings[relative_positions] # [S, S, D]658 q = q.transpose(1, 2) # [B, S, H, D]659 rel_logits = torch.matmul(q, rel_embeddings.transpose(-2, -1)) # [B, S, H, S]660 rel_logits = rel_logits.transpose(1, 2) # [B, H, S, S]661 return rel_logits662 663 def forward(self, x: Tensor, freqs_cis: Tensor, mask: Tensor, input_pos: Optional[Tensor] = None) -> Tensor:664 bsz, seqlen, _ = x.shape665 666 kv_size = self.n_local_heads * self.head_dim667 q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1)668 context_seqlen = seqlen669 670 q = q.view(bsz, seqlen, self.n_head, self.head_dim)671 k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)672 v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)673 674 if self.pos_embed_type == "rope":675 q = apply_rotary_emb(q, freqs_cis)676 k = apply_rotary_emb(k, freqs_cis)677 678 q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))679 680 if self.kv_cache is not None:681 k, v = self.kv_cache.update(input_pos, k, v)682 683 k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)684 v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)685 686 if self.pos_embed_type == "conformer":687 scale = 1.0 / math.sqrt(self.head_dim)688 scores = torch.matmul(q, k.transpose(-2, -1)) * scale689 rel_scores = self._compute_conformer_pos_scores(q, seqlen)690 scores = scores + rel_scores691 if mask is not None:692 scores = scores.masked_fill(~mask, float("-inf"))693 attn = F.softmax(scores, dim=-1)694 if self.attn_dropout_rate > 0 and self.training:695 attn = F.dropout(attn, p=self.attn_dropout_rate)696 y = torch.matmul(attn, v)697 else:698 y = F.scaled_dot_product_attention(699 q, k, v,700 dropout_p=self.attn_dropout_rate if self.training else 0.0,701 attn_mask=mask,702 )703 y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.head_dim * self.n_head)704 y = self.wo(y)705 return y706 707 708class FeedForward(nn.Module):709 def __init__(self, config: ModelArgs) -> None:710 super().__init__()711 self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)712 self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)713 self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)714 self.dropout = nn.Dropout(config.dropout_rate)715 716 def forward(self, x: Tensor) -> Tensor:717 return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))718 719 720class RMSNorm(nn.Module):721 def __init__(self, dim: int, eps: float = 1e-5):722 super().__init__()723 self.eps = eps724 self.weight = nn.Parameter(torch.ones(dim))725 726 def _norm(self, x):727 return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)728 729 def forward(self, x: Tensor) -> Tensor:730 output = self._norm(x.float()).type_as(x)731 return output * self.weight732 733 734class LayerScale(nn.Module):735 def __init__(self, dim: int, init_values: Union[float, Tensor] = 1e-2, inplace: bool = False) -> None:736 super().__init__()737 self.inplace = inplace738 self.gamma = nn.Parameter(init_values * torch.ones(dim))739 740 def forward(self, x: Tensor) -> Tensor:741 return x.mul_(self.gamma) if self.inplace else x * self.gamma742 743 744class WindowLimitedTransformer(Transformer):745 """Transformer with window-limited causal attention."""746 def __init__(747 self,748 config: ModelArgs,749 input_dim: int = 512,750 window_size: Optional[int] = None,751 causal: bool = True,752 look_ahead_conv: Optional[nn.Module] = None,753 ):754 super().__init__(config)755 self.window_size = window_size756 self.causal = causal757 self.channels_first = config.channels_first758 self.look_ahead_conv = look_ahead_conv if look_ahead_conv is not None else nn.Identity()759 self.input_proj = nn.Linear(input_dim, config.dim) if input_dim != config.dim else nn.Identity()760 self.output_proj = nn.Linear(config.dim, input_dim) if input_dim != config.dim else nn.Identity()761 762 def make_window_limited_mask(self, max_length: int, x_lens: Optional[Tensor] = None) -> Tensor:763 if self.causal:764 mask = torch.tril(torch.ones(max_length, max_length))765 row_indices = torch.arange(max_length).view(-1, 1)766 window_size = self.window_size or max_length767 valid_range = (row_indices - window_size + 1).clamp(min=0)768 column_indices = torch.arange(max_length)769 mask = (column_indices >= valid_range) & mask.bool()770 else:771 raise NotImplementedError772 mask = mask.bool()[None, None]773 return mask774 775 def make_mask(self, max_length: int, x_lens: Optional[Tensor] = None) -> Tensor:776 if self.causal:777 mask = torch.tril(torch.ones(max_length, max_length))778 else:779 mask = torch.ones(max_length, max_length)780 mask = mask.bool()[None, None]781 for i, x_len in enumerate(x_lens):782 mask[:x_len, i] = 0783 mask = mask.bool()[None, None]784 return mask785 786 def forward(self, x: Tensor, x_lens: Optional[Tensor] = None) -> Tensor:787 if self.channels_first:788 x = x.transpose(1, 2)789 x = self.input_proj(x)790 x = self.look_ahead_conv(x)791 input_pos = torch.arange(x.shape[1], device=x.device)792 max_length = x.shape[1]793 if self.window_size is not None:794 mask = self.make_window_limited_mask(max_length, x_lens)795 else:796 mask = self.make_mask(max_length, x_lens)797 mask = mask.to(x.device)798 x = super().forward(x, input_pos, mask)799 x = self.output_proj(x)800 if self.channels_first:801 x = x.transpose(1, 2)802 return x803 804 805def precompute_freqs_cis(806 seq_len: int, n_elem: int, base: int = 10000, dtype: torch.dtype = torch.bfloat16807) -> Tensor:808 freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))809 t = torch.arange(seq_len, device=freqs.device)810 freqs = torch.outer(t, freqs)811 freqs_cis = torch.polar(torch.ones_like(freqs), freqs)812 cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)813 return cache.to(dtype=dtype)814 815def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:816 xshaped = x.float().reshape(*x.shape[:-1], -1, 2)817 freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)818 x_out2 = torch.stack(819 [820 xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],821 xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],822 ],823 -1,824 )825 x_out2 = x_out2.flatten(3)826 return x_out2.type_as(x)827 828 829def init_weights(m):830 if isinstance(m, nn.Conv1d):831 nn.init.trunc_normal_(m.weight, std=0.02)832 nn.init.constant_(m.bias, 0)833 834 835# --------------------------------------------------------------------836# Top-level AE837# --------------------------------------------------------------------838 839class EncoderBlock(nn.Module):840 def __init__(841 self,842 dim: int = 16,843 stride: int = 1,844 causal: bool = False,845 n_t_layer: int = 0,846 transformer_general_config=None,847 ):848 super().__init__()849 conv_class = CausalWNConv1d if causal else WNConv1d850 transformer_module = (851 nn.Identity()852 if n_t_layer == 0853 else WindowLimitedTransformer(854 causal=causal,855 input_dim=dim,856 window_size=512,857 config=transformer_general_config(858 n_layer=n_t_layer,859 n_head=dim // 64,860 dim=dim,861 intermediate_size=dim * 3,862 ),863 )864 )865 self.block = nn.Sequential(866 # three multi‑receptive‑field residual units867 ResidualUnit(dim // 2, dilation=1, causal=causal),868 ResidualUnit(dim // 2, dilation=3, causal=causal),869 ResidualUnit(dim // 2, dilation=9, causal=causal),870 Snake1d(dim // 2),871 conv_class(dim // 2, dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)),872 transformer_module,873 )874 875 def forward(self, x: Tensor) -> Tensor:876 return self.block(x)877 878 879class ResidualUnit(nn.Module):880 def __init__(self, dim: int = 16, dilation: int = 1, causal: bool = False):881 super().__init__()882 conv_class = CausalWNConv1d if causal else WNConv1d883 pad = ((7 - 1) * dilation) // 2884 self.block = nn.Sequential(885 Snake1d(dim),886 conv_class(dim, dim, kernel_size=7, dilation=dilation, padding=pad),887 Snake1d(dim),888 conv_class(dim, dim, kernel_size=1),889 )890 self.causal = causal891 892 def forward(self, x: Tensor) -> Tensor:893 y = self.block(x)894 pad = x.shape[-1] - y.shape[-1]895 if pad > 0:896 if self.causal:897 x = x[..., :-pad]898 else:899 x = x[..., pad // 2 : -pad // 2]900 return x + y901 902 903class Encoder(nn.Module):904 def __init__(905 self,906 d_model: int = 64,907 strides: List[int] = [2, 4, 8, 8],908 d_latent: int = 64,909 n_transformer_layers: List[int] = [0, 0, 4, 4],910 transformer_general_config: Optional[ModelArgs] = None,911 causal: bool = False,912 ):913 super().__init__()914 conv_class = CausalWNConv1d if causal else WNConv1d915 layers: List[nn.Module] = [conv_class(1, d_model, kernel_size=7, padding=3)]916 for stride, n_t_layer in zip(strides, n_transformer_layers):917 d_model *= 2918 layers.append(919 EncoderBlock(920 d_model, stride=stride, causal=causal,921 n_t_layer=n_t_layer, transformer_general_config=transformer_general_config,922 )923 )924 layers += [Snake1d(d_model), conv_class(d_model, d_latent, kernel_size=3, padding=1)]925 self.block = nn.Sequential(*layers)926 self.enc_dim = d_model927 928 def forward(self, x: Tensor) -> Tensor:929 return self.block(x)930 931 932class DecoderBlock(nn.Module):933 def __init__(934 self,935 input_dim: int = 16,936 output_dim: int = 8,937 stride: int = 1,938 causal: bool = False,939 n_t_layer: int = 0,940 transformer_general_config=None,941 ):942 super().__init__()943 conv_trans_class = CausalWNConvTranspose1d if causal else WNConvTranspose1d944 transformer_module = (945 nn.Identity()946 if n_t_layer == 0947 else WindowLimitedTransformer(948 causal=causal,949 input_dim=input_dim,950 window_size=None,951 config=transformer_general_config(952 n_layer=n_t_layer,953 n_head=input_dim // 64,954 dim=input_dim,955 intermediate_size=input_dim * 3,956 ),957 )958 )959 self.block = nn.Sequential(960 Snake1d(input_dim),961 conv_trans_class(input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)),962 ResidualUnit(output_dim, dilation=1, causal=causal),963 ResidualUnit(output_dim, dilation=3, causal=causal),964 ResidualUnit(output_dim, dilation=9, causal=causal),965 )966 967 def forward(self, x: Tensor) -> Tensor:968 return self.block(x)969 970 971class Decoder(nn.Module):972 def __init__(973 self,974 input_channel: int,975 channels: int,976 rates: List[int],977 d_out: int = 1,978 causal: bool = False,979 n_transformer_layers: List[int] = [0, 0, 0, 0],980 transformer_general_config=None,981 ):982 super().__init__()983 conv_class = CausalWNConv1d if causal else WNConv1d984 layers: List[nn.Module] = [conv_class(input_channel, channels, kernel_size=7, padding=3)]985 for i, (stride, n_t_layer) in enumerate(zip(rates, n_transformer_layers)):986 input_dim = channels // 2**i987 output_dim = channels // 2 ** (i + 1)988 layers.append(989 DecoderBlock(990 input_dim, output_dim, stride, causal=causal,991 n_t_layer=n_t_layer, transformer_general_config=transformer_general_config,992 )993 )994 layers += [Snake1d(output_dim), conv_class(output_dim, d_out, kernel_size=7, padding=3), nn.Tanh()]995 self.model = nn.Sequential(*layers)996 997 def forward(self, x: Tensor) -> Tensor:998 return self.model(x)999 1000 1001class DAC(nn.Module):1002 def __init__(1003 self,1004 encoder_dim: int = 64,1005 encoder_rates: List[int] = [2, 4, 8, 8],1006 latent_dim: Optional[int] = None,1007 decoder_dim: int = 1536,1008 decoder_rates: List[int] = [8, 8, 4, 2],1009 quantizer: Optional[nn.Module] = None,1010 sample_rate: int = 44100,1011 causal: bool = True,1012 encoder_transformer_layers: List[int] = [0, 0, 0, 0],1013 decoder_transformer_layers: List[int] = [0, 0, 0, 0],1014 transformer_general_config=None,1015 ):1016 super().__init__()1017 1018 self.encoder_dim = encoder_dim1019 self.encoder_rates = encoder_rates1020 self.decoder_dim = decoder_dim1021 self.decoder_rates = decoder_rates1022 self.sample_rate = sample_rate1023 1024 if latent_dim is None:1025 latent_dim = encoder_dim * (2 ** len(encoder_rates))1026 self.latent_dim = latent_dim1027 1028 self.hop_length = int(np.prod(encoder_rates))1029 self.encoder = Encoder(1030 encoder_dim, encoder_rates, latent_dim, causal=causal,1031 n_transformer_layers=encoder_transformer_layers,1032 transformer_general_config=transformer_general_config,1033 )1034 self.quantizer = quantizer1035 self.decoder = Decoder(1036 latent_dim, decoder_dim, decoder_rates, causal=causal,1037 n_transformer_layers=decoder_transformer_layers,1038 transformer_general_config=transformer_general_config,1039 )1040 self.sample_rate = sample_rate1041 self.apply(init_weights)1042 1043 self.delay = self.get_delay()1044 self.frame_length = self.hop_length * 41045 1046 def get_output_length(self, input_length: int) -> int:1047 length = input_length1048 for stride in self.encoder_rates:1049 length = math.ceil(length / stride)1050 return length1051 1052 def get_delay(self) -> int:1053 l_out = self.get_output_length(0)1054 L = l_out1055 1056 layers = [layer for layer in self.modules() if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d))]1057 for layer in reversed(layers):1058 d = layer.dilation[0]1059 k = layer.kernel_size[0]1060 s = layer.stride[0]1061 if isinstance(layer, nn.ConvTranspose1d):1062 L = ((L - d * (k - 1) - 1) / s) + 11063 elif isinstance(layer, nn.Conv1d):1064 L = (L - 1) * s + d * (k - 1) + 11065 L = math.ceil(L)1066 1067 l_in = L1068 return (l_in - l_out) // 21069 1070 def preprocess(self, audio_data: Tensor, sample_rate: Optional[int]) -> Tensor:1071 if sample_rate is None:1072 sample_rate = self.sample_rate1073 assert sample_rate == self.sample_rate1074 1075 length = audio_data.shape[-1]1076 right_pad = math.ceil(length / self.hop_length) * self.hop_length - length1077 audio_data = F.pad(audio_data, (0, right_pad))1078 return audio_data1079 1080 def encode(1081 self,1082 audio_data: Tensor,1083 audio_lengths: Optional[Tensor] = None,1084 n_quantizers: Optional[int] = None,1085 **kwargs,1086 ):1087 """Encode audio to quantized code indices."""1088 if audio_data.ndim == 2:1089 audio_data = audio_data.unsqueeze(1)1090 length = audio_data.shape[-1]1091 right_pad = math.ceil(length / self.frame_length) * self.frame_length - length1092 audio_data = F.pad(audio_data, (0, right_pad))1093 if audio_lengths is None:1094 audio_lengths = torch.LongTensor([length + right_pad]).to(audio_data.device)1095 1096 z = self.encoder(audio_data)1097 vq_results = self.quantizer(z, n_quantizers, **kwargs)1098 indices = vq_results.codes1099 indices_lens = torch.ceil(audio_lengths / self.frame_length).long()1100 return indices, indices_lens1101 1102 def decode(self, indices: Tensor, feature_lengths: Tensor):1103 """Decode code indices to audio."""1104 if indices.ndim == 2:1105 indices = indices[None]1106 z = self.quantizer.decode(indices)1107 audio_lengths = feature_lengths * self.frame_length1108 return self.decoder(z), audio_lengths1109 1110 def encode_to_codes(self, audio: Tensor, audio_lengths: Optional[Tensor] = None, n_quantizers: Optional[int] = None, **kw):1111 return self.encode(audio, audio_lengths, n_quantizers, **kw)1112 1113 def decode_codes(self, indices: Tensor, feature_lengths: Tensor):1114 return self.decode(indices, feature_lengths)1115 1116 @torch.no_grad()1117 def encode_zq(self, audio_data: Tensor) -> Tensor:1118 indices, _ = self.encode(audio_data)1119 new_indices = torch.zeros_like(indices)1120 new_indices[:, 0] = torch.clamp(indices[:, 0], max=self.quantizer.semantic_quantizer.codebook_size - 1)1121 new_indices[:, 1:] = torch.clamp(indices[:, 1:], max=self.quantizer.quantizer.codebook_size - 1)1122 1123 z_q_semantic = self.quantizer.semantic_quantizer.from_codes(new_indices[:, :1])[0]1124 z_q_residual = self.quantizer.quantizer.from_codes(new_indices[:, 1:])[0]1125 z_q = z_q_semantic + z_q_residual1126 return z_q1127 1128 @torch.no_grad()1129 def decode_zq(self, z_q: Tensor) -> Tensor:1130 z_q = self.quantizer.post_module(z_q)1131 z_q = self.quantizer.upsample(z_q)1132 return self.decoder(z_q)1133 1134 @property1135 def device(self) -> torch.device: return next(self.parameters()).device1136 1137 @property1138 def dtype(self) -> torch.dtype: return next(self.parameters()).dtype1139 1140# --------------------------------------------------------------------1141# Build helpers1142# --------------------------------------------------------------------1143 1144def build_ae(**cfg) -> DAC:1145 """1146 Factory used by external loaders1147 """1148 # Shared transformer config for the RVQ pre/post modules1149 q_config = ModelArgs(1150 block_size=4096, n_layer=8, n_head=16, dim=1024,1151 intermediate_size=3072, head_dim=64, norm_eps=1e-5,1152 dropout_rate=0.1, attn_dropout_rate=0.1, channels_first=True1153 )1154 1155 def make_transformer():1156 return WindowLimitedTransformer(1157 causal=True, window_size=128, input_dim=1024, config=q_config1158 )1159 1160 quantizer = DownsampleResidualVectorQuantize(1161 input_dim=1024, n_codebooks=9, codebook_size=1024, codebook_dim=8,1162 quantizer_dropout=0.5, downsample_factor=(2, 2),1163 semantic_codebook_size=4096,1164 pre_module=make_transformer(),1165 post_module=make_transformer(),1166 )1167 1168 def transformer_general_config(**kw):1169 return ModelArgs(1170 block_size=kw.get("block_size", 16384),1171 n_layer=kw.get("n_layer", 8),1172 n_head=kw.get("n_head", 8),1173 dim=kw.get("dim", 512),1174 intermediate_size=kw.get("intermediate_size", 1536),1175 n_local_heads=kw.get("n_local_heads", -1),1176 head_dim=kw.get("head_dim", 64),1177 rope_base=kw.get("rope_base", 10000),1178 norm_eps=kw.get("norm_eps", 1e-5),1179 dropout_rate=kw.get("dropout_rate", 0.1),1180 attn_dropout_rate=kw.get("attn_dropout_rate", 0.1),1181 channels_first=kw.get("channels_first", True),1182 )1183 1184 dac = DAC(1185 encoder_dim=64, encoder_rates=[2, 4, 8, 8], latent_dim=1024,1186 decoder_dim=1536, decoder_rates=[8, 8, 4, 2],1187 quantizer=quantizer, sample_rate=44100, causal=True,1188 encoder_transformer_layers=[0, 0, 0, 4],1189 decoder_transformer_layers=[4, 0, 0, 0],1190 transformer_general_config=transformer_general_config,1191 )1192 return dac1193 1194__all__ = [1195 "DAC",1196 "build_ae",1197 "VectorQuantize",1198 "ResidualVectorQuantize",1199 "DownsampleResidualVectorQuantize",1200]