vicentemovil/echo-tts-preview
0
1# SPDX-FileCopyrightText: 2025 Jordan Darefsky2# SPDX-License-Identifier: Apache-2.03#4# This file contains portions adapted from:5# • Descript Audio Codec (DAC) — MIT License (full text appended below)6# • Fish-Speech S1 DAC Autoencoder — reference implementation (Apache-2.0 / CC-BY-NC),7# rewritten here in a single-file Torch module for interoperability and transparency.8#9# OVERALL LICENSE (this file): Apache-2.0, except where explicitly marked:10# # SPDX-License-Identifier: MIT11# Keep these notices and the embedded MIT text if you redistribute this file.12 13# NOTE (style/provenance):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 stylistically differs from the rest of the codebase (I'm not even sure about internal consistency)17# and is likely much messier than it would have been had it been written from scratch.18 19 20from __future__ import annotations21 22import math23from dataclasses import dataclass24from typing import List, Optional, Tuple, Union25 26import numpy as np27import torch28from torch import Tensor, nn29from torch.nn import functional as F30from torch.nn.utils.parametrizations import weight_norm31from torch.nn.utils.parametrize import remove_parametrizations32 33from einops import rearrange34 35 36# --------------------------------------------------------------------37# Shared helpers38# --------------------------------------------------------------------39 40def find_multiple(n: int, k: int) -> int:41 return n if n % k == 0 else n + k - (n % k)42 43def unpad1d(x: Tensor, paddings: Tuple[int, int]) -> Tensor:44 """Remove padding from x, handling properly zero padding. Only for 1d!"""45 padding_left, padding_right = paddings46 assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)47 assert (padding_left + padding_right) <= x.shape[-1]48 end = x.shape[-1] - padding_right49 return x[..., padding_left:end]50 51def get_extra_padding_for_conv1d(52 x: Tensor, kernel_size: int, stride: int, padding_total: int = 053) -> int:54 """See pad_for_conv1d; enough right pad so striding evenly covers length."""55 length = x.shape[-1]56 n_frames = (length - kernel_size + padding_total) / stride + 157 ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)58 return ideal_length - length59 60def pad1d(61 x: Tensor,62 paddings: Tuple[int, int],63 mode: str = "zeros",64 value: float = 0.0,65) -> Tensor:66 """67 Reflect‑safe 1D pad: if reflect would underflow on small inputs, insert68 temporary right zero-pad before reflecting.69 """70 length = x.shape[-1]71 padding_left, padding_right = paddings72 assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)73 if mode == "reflect":74 max_pad = max(padding_left, padding_right)75 extra_pad = 076 if length <= max_pad:77 extra_pad = max_pad - length + 178 x = F.pad(x, (0, extra_pad))79 padded = F.pad(x, (padding_left, padding_right), mode, value)80 end = padded.shape[-1] - extra_pad81 return padded[..., :end]82 else:83 return F.pad(x, (padding_left, padding_right), mode, value)84 85 86# --------------------------------------------------------------------87# DAC Layers (adapted) — MIT88# Original: https://github.com/descriptinc/descript-audio-codec/blob/main/dac/nn/layers.py89# SPDX-License-Identifier: MIT90# --------------------------------------------------------------------91 92def WNConv1d(*args, **kwargs):93 return weight_norm(nn.Conv1d(*args, **kwargs))94 95def WNConvTranspose1d(*args, **kwargs):96 return weight_norm(nn.ConvTranspose1d(*args, **kwargs))97 98@torch.jit.script99def snake(x: Tensor, alpha: Tensor) -> Tensor:100 shape = x.shape101 x = x.reshape(shape[0], shape[1], -1)102 x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)103 x = x.reshape(shape)104 return x105 106class Snake1d(nn.Module):107 def __init__(self, channels: int):108 super().__init__()109 self.alpha = nn.Parameter(torch.ones(1, channels, 1))110 def forward(self, x: Tensor) -> Tensor:111 return snake(x, self.alpha)112 113# --------------------------------------------------------------------114# DAC Vector Quantize (adapted) — MIT115# Original: https://github.com/descriptinc/descript-audio-codec/blob/main/dac/nn/quantize.py116# SPDX-License-Identifier: MIT117# --------------------------------------------------------------------118 119class VectorQuantize(nn.Module):120 """121 VQ with factorized, l2-normalized codes (ViT‑VQGAN style).122 I/O in (B, D, T).123 """124 def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):125 super().__init__()126 self.codebook_size = codebook_size127 self.codebook_dim = codebook_dim128 self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)129 self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)130 self.codebook = nn.Embedding(codebook_size, codebook_dim)131 132 def forward(self, z: Tensor):133 z_e = self.in_proj(z) # (B, D, T)134 z_q, indices = self.decode_latents(z_e)135 commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])136 codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])137 z_q = z_e + (z_q - z_e).detach() # straight‑through138 z_q = self.out_proj(z_q)139 return z_q, commitment_loss, codebook_loss, indices, z_e140 141 def embed_code(self, embed_id: Tensor) -> Tensor:142 return F.embedding(embed_id, self.codebook.weight)143 144 def decode_code(self, embed_id: Tensor) -> Tensor:145 return self.embed_code(embed_id).transpose(1, 2)146 147 def decode_latents(self, latents: Tensor) -> Tuple[Tensor, Tensor]:148 encodings = rearrange(latents, "b d t -> (b t) d")149 codebook = self.codebook.weight150 encodings = F.normalize(encodings)151 codebook = F.normalize(codebook)152 dist = (153 encodings.pow(2).sum(1, keepdim=True)154 - 2 * encodings @ codebook.t()155 + codebook.pow(2).sum(1, keepdim=True).t()156 )157 indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))158 z_q = self.decode_code(indices)159 return z_q, indices160 161 162class ResidualVectorQuantize(nn.Module):163 """SoundStream-style residual VQ stack."""164 def __init__(165 self,166 input_dim: int = 512,167 n_codebooks: int = 9,168 codebook_size: int = 1024,169 codebook_dim: Union[int, List[int]] = 8,170 quantizer_dropout: float = 0.0,171 ):172 super().__init__()173 if isinstance(codebook_dim, int):174 codebook_dim = [codebook_dim for _ in range(n_codebooks)]175 176 self.n_codebooks = n_codebooks177 self.codebook_dim = codebook_dim178 self.codebook_size = codebook_size179 180 self.quantizers = nn.ModuleList([181 VectorQuantize(input_dim, codebook_size, codebook_dim[i])182 for i in range(n_codebooks)183 ])184 self.quantizer_dropout = quantizer_dropout185 186 def forward(self, z: Tensor, n_quantizers: Optional[int] = None):187 z_q = 0188 residual = z189 commitment_loss = 0190 codebook_loss = 0191 192 codebook_indices = []193 latents = []194 195 if n_quantizers is None:196 n_quantizers = self.n_codebooks197 if self.training:198 n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1199 dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],))200 n_dropout = int(z.shape[0] * self.quantizer_dropout)201 n_quantizers[:n_dropout] = dropout[:n_dropout]202 n_quantizers = n_quantizers.to(z.device)203 204 for i, quantizer in enumerate(self.quantizers):205 if self.training is False and i >= n_quantizers:206 break207 208 z_q_i, commit_i, codebk_i, indices_i, z_e_i = quantizer(residual)209 210 mask = (torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers)211 z_q = z_q + z_q_i * mask[:, None, None]212 residual = residual - z_q_i213 214 commitment_loss += (commit_i * mask).mean()215 codebook_loss += (codebk_i * mask).mean()216 217 codebook_indices.append(indices_i)218 latents.append(z_e_i)219 220 codes = torch.stack(codebook_indices, dim=1)221 latents = torch.cat(latents, dim=1)222 223 return z_q, codes, latents, commitment_loss, codebook_loss224 225 def from_codes(self, codes: Tensor) -> Tuple[Tensor, Tensor, Tensor]:226 z_q = 0.0227 z_p = []228 n_codebooks = codes.shape[1]229 for i in range(n_codebooks):230 z_p_i = self.quantizers[i].decode_code(codes[:, i, :])231 z_p.append(z_p_i)232 z_q_i = self.quantizers[i].out_proj(z_p_i)233 z_q = z_q + z_q_i234 return z_q, torch.cat(z_p, dim=1), codes235 236 def from_latents(self, latents: Tensor) -> Tuple[Tensor, Tensor, Tensor]:237 z_q = 0238 z_p = []239 codes = []240 dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers])241 n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[0]242 for i in range(n_codebooks):243 j, k = dims[i], dims[i + 1]244 z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])245 z_p.append(z_p_i)246 codes.append(codes_i)247 z_q_i = self.quantizers[i].out_proj(z_p_i)248 z_q = z_q + z_q_i249 return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)250 251 252# --------------------------------------------------------------------253# S1 DAC rvq254# --------------------------------------------------------------------255 256@dataclass257class VQResult:258 z: Tensor259 codes: Tensor260 latents: Tensor261 codebook_loss: Tensor262 commitment_loss: Tensor263 semantic_distill_z: Optional[Tensor] = None264 265 266class CausalConvNet(nn.Module):267 def __init__(268 self,269 in_channels,270 out_channels,271 kernel_size,272 dilation=1,273 stride=1,274 groups=1,275 padding=None,276 ):277 super().__init__()278 self.conv = nn.Conv1d(279 in_channels, out_channels, kernel_size,280 stride=stride, dilation=dilation, groups=groups,281 )282 self.stride = stride283 self.kernel_size = (kernel_size - 1) * dilation + 1284 self.dilation = dilation285 self.padding = self.kernel_size - self.stride286 287 def forward(self, x: Tensor) -> Tensor:288 pad = self.padding289 extra = get_extra_padding_for_conv1d(x, self.kernel_size, self.stride, pad)290 x = pad1d(x, (pad, extra), mode="constant", value=0)291 return self.conv(x).contiguous()292 293 def weight_norm(self, name="weight", dim=0):294 self.conv = weight_norm(self.conv, name=name, dim=dim)295 return self296 297 def remove_weight_norm(self):298 self.conv = remove_parametrizations(self.conv)299 return self300 301 302class CausalTransConvNet(nn.Module):303 def __init__(self, in_channels, out_channels, kernel_size, dilation=1, stride=1, padding=None):304 super().__init__()305 self.conv = nn.ConvTranspose1d(306 in_channels, out_channels, kernel_size,307 stride=stride, dilation=dilation308 )309 self.stride = stride310 self.kernel_size = kernel_size311 312 def forward(self, x: Tensor) -> Tensor:313 x = self.conv(x)314 pad = self.kernel_size - self.stride315 padding_right = math.ceil(pad)316 padding_left = pad - padding_right317 x = unpad1d(x, (padding_left, padding_right))318 return x.contiguous()319 320 def weight_norm(self, name="weight", dim=0):321 self.conv = weight_norm(self.conv, name=name, dim=dim)322 return self323 324 def remove_weight_norm(self):325 self.conv = remove_parametrizations(self.conv)326 return self327 328 329def CausalWNConv1d(*args, **kwargs):330 return CausalConvNet(*args, **kwargs).weight_norm()331 332def CausalWNConvTranspose1d(*args, **kwargs):333 return CausalTransConvNet(*args, **kwargs).weight_norm()334 335class ConvNeXtBlock(nn.Module):336 r"""ConvNeXt Block (1D).337 DwConv -> (N, C, L) → (N, L, C) -> LN -> Linear -> GELU -> Linear -> (N, C, L) with residual338 """339 def __init__(340 self,341 dim: int,342 layer_scale_init_value: float = 1e-6,343 mlp_ratio: float = 4.0,344 kernel_size: int = 7,345 dilation: int = 1,346 ):347 super().__init__()348 convnet_type = CausalConvNet349 self.dwconv = convnet_type(350 dim, dim, kernel_size=kernel_size,351 groups=dim, dilation=dilation,352 ) # depthwise conv353 self.norm = nn.LayerNorm(dim, eps=1e-6)354 self.pwconv1 = nn.Linear(dim, int(mlp_ratio * dim))355 self.act = nn.GELU()356 self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim)357 self.gamma = (358 nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)359 if layer_scale_init_value > 0 else None360 )361 362 def forward(self, x: Tensor, apply_residual: bool = True) -> Tensor:363 inp = x364 x = self.dwconv(x)365 x = x.permute(0, 2, 1) # (N, C, L) -> (N, L, C)366 x = self.norm(x)367 x = self.pwconv1(x)368 x = self.act(x)369 x = self.pwconv2(x)370 if self.gamma is not None:371 x = self.gamma * x372 x = x.permute(0, 2, 1) # (N, L, C) -> (N, C, L)373 if apply_residual:374 x = inp + x375 return x376 377 378class DownsampleResidualVectorQuantize(nn.Module):379 def __init__(380 self,381 input_dim: int = 1024,382 n_codebooks: int = 9,383 codebook_dim: int = 8,384 quantizer_dropout: float = 0.5,385 codebook_size: int = 1024,386 semantic_codebook_size: int = 4096,387 downsample_factor: Tuple[int, ...] = (2, 2),388 downsample_dims: Optional[Tuple[int, ...]] = None,389 pre_module: Optional[nn.Module] = None,390 post_module: Optional[nn.Module] = None,391 semantic_predictor_module: Optional[nn.Module] = None,392 ):393 super().__init__()394 395 if downsample_dims is None:396 downsample_dims = tuple(input_dim for _ in range(len(downsample_factor)))397 398 all_dims = (input_dim,) + tuple(downsample_dims)399 400 self.semantic_quantizer = ResidualVectorQuantize(401 input_dim=input_dim,402 n_codebooks=1,403 codebook_size=semantic_codebook_size,404 codebook_dim=codebook_dim,405 quantizer_dropout=0.0,406 )407 408 self.quantizer = ResidualVectorQuantize(409 input_dim=input_dim,410 n_codebooks=n_codebooks,411 codebook_size=codebook_size,412 codebook_dim=codebook_dim,413 quantizer_dropout=quantizer_dropout,414 )415 416 convnet_type = CausalConvNet417 transconvnet_type = CausalTransConvNet418 419 self.downsample = nn.Sequential(420 *[421 nn.Sequential(422 convnet_type(all_dims[idx], all_dims[idx + 1], kernel_size=factor, stride=factor),423 ConvNeXtBlock(dim=all_dims[idx + 1]),424 )425 for idx, factor in enumerate(downsample_factor)426 ]427 )428 429 self.upsample = nn.Sequential(430 *[431 nn.Sequential(432 transconvnet_type(all_dims[idx + 1], all_dims[idx], kernel_size=factor, stride=factor),433 ConvNeXtBlock(dim=all_dims[idx]),434 )435 for idx, factor in reversed(list(enumerate(downsample_factor)))436 ]437 )438 439 self.apply(self._init_weights)440 self.pre_module = pre_module if pre_module is not None else nn.Identity()441 self.post_module = post_module if post_module is not None else nn.Identity()442 self.semantic_predictor_module = (443 semantic_predictor_module if semantic_predictor_module is not None else nn.Identity()444 )445 446 @staticmethod447 def _init_weights(m):448 if isinstance(m, (nn.Conv1d, nn.Linear)):449 nn.init.trunc_normal_(m.weight, std=0.02)450 if getattr(m, "bias", None) is not None:451 nn.init.constant_(m.bias, 0)452 453 def forward(self, z: Tensor, n_quantizers: Optional[int] = None, semantic_len: Optional[Tensor] = None, **kwargs):454 # z: (B, D, T)455 original_shape = z.shape456 if semantic_len is None:457 semantic_len = torch.LongTensor([z.shape[-1]])458 459 z = self.downsample(z)460 z = self.pre_module(z) # (B, D, T) or (B, T, D) depending on module; original uses channels-first in/out461 462 semantic_z, semantic_codes, semantic_latents, semantic_commitment_loss, semantic_codebook_loss = \463 self.semantic_quantizer(z)464 residual_z = z - semantic_z465 residual_z, codes, latents, commitment_loss, codebook_loss = self.quantizer(residual_z, n_quantizers=n_quantizers)466 z = semantic_z + residual_z467 commitment_loss = commitment_loss + semantic_commitment_loss468 codebook_loss = codebook_loss + semantic_codebook_loss469 codes = torch.cat([semantic_codes, codes], dim=1)470 latents = torch.cat([semantic_latents, latents], dim=1)471 z = self.post_module(z)472 z = self.upsample(z)473 474 # Pad or crop z to match original shape (time dimension)475 diff = original_shape[-1] - z.shape[-1]476 right = 0477 left = abs(diff) - right478 if diff > 0:479 z = F.pad(z, (left, right))480 elif diff < 0:481 z = z[..., left:]482 483 return VQResult(484 z=z, codes=codes, latents=latents,485 commitment_loss=commitment_loss, codebook_loss=codebook_loss,486 )487 488 def decode(self, indices: Tensor) -> Tensor:489 new_indices = torch.zeros_like(indices)490 new_indices[:, 0] = torch.clamp(indices[:, 0], max=self.semantic_quantizer.codebook_size - 1)491 new_indices[:, 1:] = torch.clamp(indices[:, 1:], max=self.quantizer.codebook_size - 1)492 493 z_q_semantic = self.semantic_quantizer.from_codes(new_indices[:, :1])[0]494 z_q_residual = self.quantizer.from_codes(new_indices[:, 1:])[0]495 z_q = z_q_semantic + z_q_residual496 z_q = self.post_module(z_q)497 z_q = self.upsample(z_q)498 return z_q499 500 501# --------------------------------------------------------------------502# Transformer stack503# --------------------------------------------------------------------504 505@dataclass506class ModelArgs:507 block_size: int = 2048508 n_layer: int = 8509 n_head: int = 8510 dim: int = 512511 intermediate_size: int = 1536512 n_local_heads: int = -1513 head_dim: int = 64514 rope_base: float = 10000515 norm_eps: float = 1e-5516 dropout_rate: float = 0.1517 attn_dropout_rate: float = 0.1518 channels_first: bool = True # to be compatible with conv1d input/output519 pos_embed_type: str = "rope" # "rope" or "conformer"520 max_relative_position: int = 128521 522 def __post_init__(self):523 if self.n_local_heads == -1:524 self.n_local_heads = self.n_head525 if self.intermediate_size is None:526 hidden_dim = 4 * self.dim527 n_hidden = int(2 * hidden_dim / 3)528 self.intermediate_size = find_multiple(n_hidden, 256)529 assert self.pos_embed_type in ["rope", "conformer"]530 531 532class KVCache(nn.Module):533 def __init__(self, max_batch_size, max_seq_length, n_heads, head_dim, dtype=torch.bfloat16):534 super().__init__()535 cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim)536 self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype))537 self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype))538 539 def update(self, input_pos: Tensor, k_val: Tensor, v_val: Tensor):540 # input_pos: [S], k_val: [B, H, S, D]541 assert input_pos.shape[0] == k_val.shape[2]542 k_out = self.k_cache543 v_out = self.v_cache544 k_out[:, :, input_pos] = k_val545 v_out[:, :, input_pos] = v_val546 return (547 k_out[:, :, : input_pos.max() + 1, :],548 v_out[:, :, : input_pos.max() + 1, :],549 )550 551 def clear_cache(self, prompt_len: int):552 self.k_cache[:, :, prompt_len:, :].fill_(0)553 self.v_cache[:, :, prompt_len:, :].fill_(0)554 555 556class Transformer(nn.Module):557 def __init__(self, config: ModelArgs) -> None:558 super().__init__()559 self.config = config560 561 self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))562 self.norm = RMSNorm(config.dim, eps=config.norm_eps)563 564 if config.pos_embed_type == "rope":565 freqs_cis = precompute_freqs_cis(self.config.block_size, self.config.head_dim, self.config.rope_base)566 self.register_buffer("freqs_cis", freqs_cis)567 else:568 self.register_buffer("freqs_cis", None)569 570 causal_mask = torch.tril(torch.ones(self.config.block_size, self.config.block_size, dtype=torch.bool))571 self.register_buffer("causal_mask", causal_mask)572 573 self.max_batch_size = -1574 self.max_seq_length = -1575 self.use_kv_cache = False576 577 def setup_caches(self, max_batch_size, max_seq_length):578 head_dim = self.config.dim // self.config.n_head579 max_seq_length = find_multiple(max_seq_length, 8)580 self.max_seq_length = max_seq_length581 self.max_batch_size = max_batch_size582 dtype = self.norm.weight.dtype583 device = self.norm.weight.device584 585 for b in self.layers:586 b.attention.kv_cache = KVCache(587 max_batch_size, max_seq_length, self.config.n_local_heads, head_dim, dtype588 ).to(device)589 590 self.use_kv_cache = True591 592 def forward(self, x: Tensor, input_pos: Optional[Tensor] = None, mask: Optional[Tensor] = None) -> Tensor:593 if self.config.pos_embed_type == "rope":594 assert self.freqs_cis is not None595 freqs_cis = self.freqs_cis[input_pos]596 else:597 freqs_cis = None598 599 if mask is None:600 if not self.training and self.use_kv_cache:601 mask = self.causal_mask[None, None, input_pos]602 mask = mask[..., : input_pos.max() + 1]603 else:604 mask = self.causal_mask[None, None, input_pos]605 mask = mask[..., input_pos]606 607 for layer in self.layers:608 x = layer(x, input_pos, freqs_cis, mask)609 x = self.norm(x)610 return x611 612 613class TransformerBlock(nn.Module):614 def __init__(self, config: ModelArgs) -> None:615 super().__init__()616 self.attention = Attention(config)617 self.feed_forward = FeedForward(config)618 self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)619 self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)620 self.attention_layer_scale = LayerScale(config.dim, inplace=True)621 self.ffn_layer_scale = LayerScale(config.dim, inplace=True)622 623 def forward(self, x: Tensor, input_pos: Tensor, freqs_cis: Tensor, mask: Tensor) -> Tensor:624 h = x + self.attention_layer_scale(625 self.attention(self.attention_norm(x), freqs_cis, mask, input_pos)626 )627 out = h + self.ffn_layer_scale(self.feed_forward(self.ffn_norm(h)))628 return out629 630 631class Attention(nn.Module):632 def __init__(self, config: ModelArgs):633 super().__init__()634 assert config.dim % config.n_head == 0635 636 total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim637 self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)638 self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)639 self.kv_cache = None640 641 self.n_head = config.n_head642 self.head_dim = config.head_dim643 self.n_local_heads = config.n_local_heads644 self.dim = config.dim645 self.attn_dropout_rate = config.attn_dropout_rate646 self.pos_embed_type = config.pos_embed_type647 648 if self.pos_embed_type == "conformer":649 self.max_relative_position = config.max_relative_position650 num_pos_embeddings = 2 * config.max_relative_position + 1651 self.rel_pos_embeddings = nn.Parameter(torch.zeros(num_pos_embeddings, self.head_dim))652 nn.init.normal_(self.rel_pos_embeddings, mean=0.0, std=0.02)653 654 def _compute_conformer_pos_scores(self, q: Tensor, seqlen: int) -> Tensor:655 positions = torch.arange(seqlen, device=q.device)656 relative_positions = positions.unsqueeze(1) - positions.unsqueeze(0) # [S, S]657 relative_positions = torch.clamp(relative_positions + self.max_relative_position,658 0, 2 * self.max_relative_position)659 rel_embeddings = self.rel_pos_embeddings[relative_positions] # [S, S, D]660 q = q.transpose(1, 2) # [B, S, H, D]661 rel_logits = torch.matmul(q, rel_embeddings.transpose(-2, -1)) # [B, S, H, S]662 rel_logits = rel_logits.transpose(1, 2) # [B, H, S, S]663 return rel_logits664 665 def forward(self, x: Tensor, freqs_cis: Tensor, mask: Tensor, input_pos: Optional[Tensor] = None) -> Tensor:666 bsz, seqlen, _ = x.shape667 668 kv_size = self.n_local_heads * self.head_dim669 q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1)670 context_seqlen = seqlen671 672 q = q.view(bsz, seqlen, self.n_head, self.head_dim)673 k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)674 v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)675 676 if self.pos_embed_type == "rope":677 q = apply_rotary_emb(q, freqs_cis)678 k = apply_rotary_emb(k, freqs_cis)679 680 q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))681 682 if self.kv_cache is not None:683 k, v = self.kv_cache.update(input_pos, k, v)684 685 k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)686 v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)687 688 if self.pos_embed_type == "conformer":689 scale = 1.0 / math.sqrt(self.head_dim)690 scores = torch.matmul(q, k.transpose(-2, -1)) * scale691 rel_scores = self._compute_conformer_pos_scores(q, seqlen)692 scores = scores + rel_scores693 if mask is not None:694 scores = scores.masked_fill(~mask, float("-inf"))695 attn = F.softmax(scores, dim=-1)696 if self.attn_dropout_rate > 0 and self.training:697 attn = F.dropout(attn, p=self.attn_dropout_rate)698 y = torch.matmul(attn, v)699 else:700 y = F.scaled_dot_product_attention(701 q, k, v,702 dropout_p=self.attn_dropout_rate if self.training else 0.0,703 attn_mask=mask,704 )705 y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.head_dim * self.n_head)706 y = self.wo(y)707 return y708 709 710class FeedForward(nn.Module):711 def __init__(self, config: ModelArgs) -> None:712 super().__init__()713 self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)714 self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)715 self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)716 self.dropout = nn.Dropout(config.dropout_rate)717 718 def forward(self, x: Tensor) -> Tensor:719 return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))720 721 722class RMSNorm(nn.Module):723 def __init__(self, dim: int, eps: float = 1e-5):724 super().__init__()725 self.eps = eps726 self.weight = nn.Parameter(torch.ones(dim))727 728 def _norm(self, x):729 return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)730 731 def forward(self, x: Tensor) -> Tensor:732 output = self._norm(x.float()).type_as(x)733 return output * self.weight734 735 736class LayerScale(nn.Module):737 def __init__(self, dim: int, init_values: Union[float, Tensor] = 1e-2, inplace: bool = False) -> None:738 super().__init__()739 self.inplace = inplace740 self.gamma = nn.Parameter(init_values * torch.ones(dim))741 742 def forward(self, x: Tensor) -> Tensor:743 return x.mul_(self.gamma) if self.inplace else x * self.gamma744 745 746class WindowLimitedTransformer(Transformer):747 """Transformer with window-limited causal attention."""748 def __init__(749 self,750 config: ModelArgs,751 input_dim: int = 512,752 window_size: Optional[int] = None,753 causal: bool = True,754 look_ahead_conv: Optional[nn.Module] = None,755 ):756 super().__init__(config)757 self.window_size = window_size758 self.causal = causal759 self.channels_first = config.channels_first760 self.look_ahead_conv = look_ahead_conv if look_ahead_conv is not None else nn.Identity()761 self.input_proj = nn.Linear(input_dim, config.dim) if input_dim != config.dim else nn.Identity()762 self.output_proj = nn.Linear(config.dim, input_dim) if input_dim != config.dim else nn.Identity()763 764 def make_window_limited_mask(self, max_length: int, x_lens: Optional[Tensor] = None) -> Tensor:765 if self.causal:766 mask = torch.tril(torch.ones(max_length, max_length))767 row_indices = torch.arange(max_length).view(-1, 1)768 window_size = self.window_size or max_length769 valid_range = (row_indices - window_size + 1).clamp(min=0)770 column_indices = torch.arange(max_length)771 mask = (column_indices >= valid_range) & mask.bool()772 else:773 raise NotImplementedError774 mask = mask.bool()[None, None]775 return mask776 777 def make_mask(self, max_length: int, x_lens: Optional[Tensor] = None) -> Tensor:778 if self.causal:779 mask = torch.tril(torch.ones(max_length, max_length))780 else:781 mask = torch.ones(max_length, max_length)782 mask = mask.bool()[None, None]783 for i, x_len in enumerate(x_lens):784 mask[:x_len, i] = 0785 mask = mask.bool()[None, None]786 return mask787 788 def forward(self, x: Tensor, x_lens: Optional[Tensor] = None) -> Tensor:789 if self.channels_first:790 x = x.transpose(1, 2)791 x = self.input_proj(x)792 x = self.look_ahead_conv(x)793 input_pos = torch.arange(x.shape[1], device=x.device)794 max_length = x.shape[1]795 if self.window_size is not None:796 mask = self.make_window_limited_mask(max_length, x_lens)797 else:798 mask = self.make_mask(max_length, x_lens)799 mask = mask.to(x.device)800 x = super().forward(x, input_pos, mask)801 x = self.output_proj(x)802 if self.channels_first:803 x = x.transpose(1, 2)804 return x805 806 807def precompute_freqs_cis(808 seq_len: int, n_elem: int, base: int = 10000, dtype: torch.dtype = torch.bfloat16809) -> Tensor:810 freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))811 t = torch.arange(seq_len, device=freqs.device)812 freqs = torch.outer(t, freqs)813 freqs_cis = torch.polar(torch.ones_like(freqs), freqs)814 cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)815 return cache.to(dtype=dtype)816 817def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:818 xshaped = x.float().reshape(*x.shape[:-1], -1, 2)819 freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)820 x_out2 = torch.stack(821 [822 xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],823 xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],824 ],825 -1,826 )827 x_out2 = x_out2.flatten(3)828 return x_out2.type_as(x)829 830 831def init_weights(m):832 if isinstance(m, nn.Conv1d):833 nn.init.trunc_normal_(m.weight, std=0.02)834 nn.init.constant_(m.bias, 0)835 836 837# --------------------------------------------------------------------838# Top-level AE839# --------------------------------------------------------------------840 841class EncoderBlock(nn.Module):842 def __init__(843 self,844 dim: int = 16,845 stride: int = 1,846 causal: bool = False,847 n_t_layer: int = 0,848 transformer_general_config=None,849 ):850 super().__init__()851 conv_class = CausalWNConv1d if causal else WNConv1d852 transformer_module = (853 nn.Identity()854 if n_t_layer == 0855 else WindowLimitedTransformer(856 causal=causal,857 input_dim=dim,858 window_size=512,859 config=transformer_general_config(860 n_layer=n_t_layer,861 n_head=dim // 64,862 dim=dim,863 intermediate_size=dim * 3,864 ),865 )866 )867 self.block = nn.Sequential(868 # three multi‑receptive‑field residual units869 ResidualUnit(dim // 2, dilation=1, causal=causal),870 ResidualUnit(dim // 2, dilation=3, causal=causal),871 ResidualUnit(dim // 2, dilation=9, causal=causal),872 Snake1d(dim // 2),873 conv_class(dim // 2, dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)),874 transformer_module,875 )876 877 def forward(self, x: Tensor) -> Tensor:878 return self.block(x)879 880 881class ResidualUnit(nn.Module):882 def __init__(self, dim: int = 16, dilation: int = 1, causal: bool = False):883 super().__init__()884 conv_class = CausalWNConv1d if causal else WNConv1d885 pad = ((7 - 1) * dilation) // 2886 self.block = nn.Sequential(887 Snake1d(dim),888 conv_class(dim, dim, kernel_size=7, dilation=dilation, padding=pad),889 Snake1d(dim),890 conv_class(dim, dim, kernel_size=1),891 )892 self.causal = causal893 894 def forward(self, x: Tensor) -> Tensor:895 y = self.block(x)896 pad = x.shape[-1] - y.shape[-1]897 if pad > 0:898 if self.causal:899 x = x[..., :-pad]900 else:901 x = x[..., pad // 2 : -pad // 2]902 return x + y903 904 905class Encoder(nn.Module):906 def __init__(907 self,908 d_model: int = 64,909 strides: List[int] = [2, 4, 8, 8],910 d_latent: int = 64,911 n_transformer_layers: List[int] = [0, 0, 4, 4],912 transformer_general_config: Optional[ModelArgs] = None,913 causal: bool = False,914 ):915 super().__init__()916 conv_class = CausalWNConv1d if causal else WNConv1d917 layers: List[nn.Module] = [conv_class(1, d_model, kernel_size=7, padding=3)]918 for stride, n_t_layer in zip(strides, n_transformer_layers):919 d_model *= 2920 layers.append(921 EncoderBlock(922 d_model, stride=stride, causal=causal,923 n_t_layer=n_t_layer, transformer_general_config=transformer_general_config,924 )925 )926 layers += [Snake1d(d_model), conv_class(d_model, d_latent, kernel_size=3, padding=1)]927 self.block = nn.Sequential(*layers)928 self.enc_dim = d_model929 930 def forward(self, x: Tensor) -> Tensor:931 return self.block(x)932 933 934class DecoderBlock(nn.Module):935 def __init__(936 self,937 input_dim: int = 16,938 output_dim: int = 8,939 stride: int = 1,940 causal: bool = False,941 n_t_layer: int = 0,942 transformer_general_config=None,943 ):944 super().__init__()945 conv_trans_class = CausalWNConvTranspose1d if causal else WNConvTranspose1d946 transformer_module = (947 nn.Identity()948 if n_t_layer == 0949 else WindowLimitedTransformer(950 causal=causal,951 input_dim=input_dim,952 window_size=None,953 config=transformer_general_config(954 n_layer=n_t_layer,955 n_head=input_dim // 64,956 dim=input_dim,957 intermediate_size=input_dim * 3,958 ),959 )960 )961 self.block = nn.Sequential(962 Snake1d(input_dim),963 conv_trans_class(input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)),964 ResidualUnit(output_dim, dilation=1, causal=causal),965 ResidualUnit(output_dim, dilation=3, causal=causal),966 ResidualUnit(output_dim, dilation=9, causal=causal),967 )968 969 def forward(self, x: Tensor) -> Tensor:970 return self.block(x)971 972 973class Decoder(nn.Module):974 def __init__(975 self,976 input_channel: int,977 channels: int,978 rates: List[int],979 d_out: int = 1,980 causal: bool = False,981 n_transformer_layers: List[int] = [0, 0, 0, 0],982 transformer_general_config=None,983 ):984 super().__init__()985 conv_class = CausalWNConv1d if causal else WNConv1d986 layers: List[nn.Module] = [conv_class(input_channel, channels, kernel_size=7, padding=3)]987 for i, (stride, n_t_layer) in enumerate(zip(rates, n_transformer_layers)):988 input_dim = channels // 2**i989 output_dim = channels // 2 ** (i + 1)990 layers.append(991 DecoderBlock(992 input_dim, output_dim, stride, causal=causal,993 n_t_layer=n_t_layer, transformer_general_config=transformer_general_config,994 )995 )996 layers += [Snake1d(output_dim), conv_class(output_dim, d_out, kernel_size=7, padding=3), nn.Tanh()]997 self.model = nn.Sequential(*layers)998 999 def forward(self, x: Tensor) -> Tensor:1000 return self.model(x)1001 1002 1003class DAC(nn.Module):1004 def __init__(1005 self,1006 encoder_dim: int = 64,1007 encoder_rates: List[int] = [2, 4, 8, 8],1008 latent_dim: Optional[int] = None,1009 decoder_dim: int = 1536,1010 decoder_rates: List[int] = [8, 8, 4, 2],1011 quantizer: Optional[nn.Module] = None,1012 sample_rate: int = 44100,1013 causal: bool = True,1014 encoder_transformer_layers: List[int] = [0, 0, 0, 0],1015 decoder_transformer_layers: List[int] = [0, 0, 0, 0],1016 transformer_general_config=None,1017 ):1018 super().__init__()1019 1020 self.encoder_dim = encoder_dim1021 self.encoder_rates = encoder_rates1022 self.decoder_dim = decoder_dim1023 self.decoder_rates = decoder_rates1024 self.sample_rate = sample_rate1025 1026 if latent_dim is None:1027 latent_dim = encoder_dim * (2 ** len(encoder_rates))1028 self.latent_dim = latent_dim1029 1030 self.hop_length = int(np.prod(encoder_rates))1031 self.encoder = Encoder(1032 encoder_dim, encoder_rates, latent_dim, causal=causal,1033 n_transformer_layers=encoder_transformer_layers,1034 transformer_general_config=transformer_general_config,1035 )1036 self.quantizer = quantizer1037 self.decoder = Decoder(1038 latent_dim, decoder_dim, decoder_rates, causal=causal,1039 n_transformer_layers=decoder_transformer_layers,1040 transformer_general_config=transformer_general_config,1041 )1042 self.sample_rate = sample_rate1043 self.apply(init_weights)1044 1045 self.delay = self.get_delay()1046 self.frame_length = self.hop_length * 41047 1048 def get_output_length(self, input_length: int) -> int:1049 length = input_length1050 for stride in self.encoder_rates:1051 length = math.ceil(length / stride)1052 return length1053 1054 def get_delay(self) -> int:1055 l_out = self.get_output_length(0)1056 L = l_out1057 1058 layers = [layer for layer in self.modules() if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d))]1059 for layer in reversed(layers):1060 d = layer.dilation[0]1061 k = layer.kernel_size[0]1062 s = layer.stride[0]1063 if isinstance(layer, nn.ConvTranspose1d):1064 L = ((L - d * (k - 1) - 1) / s) + 11065 elif isinstance(layer, nn.Conv1d):1066 L = (L - 1) * s + d * (k - 1) + 11067 L = math.ceil(L)1068 1069 l_in = L1070 return (l_in - l_out) // 21071 1072 def preprocess(self, audio_data: Tensor, sample_rate: Optional[int]) -> Tensor:1073 if sample_rate is None:1074 sample_rate = self.sample_rate1075 assert sample_rate == self.sample_rate1076 1077 length = audio_data.shape[-1]1078 right_pad = math.ceil(length / self.hop_length) * self.hop_length - length1079 audio_data = F.pad(audio_data, (0, right_pad))1080 return audio_data1081 1082 def encode(1083 self,1084 audio_data: Tensor,1085 audio_lengths: Optional[Tensor] = None,1086 n_quantizers: Optional[int] = None,1087 **kwargs,1088 ):1089 """Encode audio to quantized code indices."""1090 if audio_data.ndim == 2:1091 audio_data = audio_data.unsqueeze(1)1092 length = audio_data.shape[-1]1093 right_pad = math.ceil(length / self.frame_length) * self.frame_length - length1094 audio_data = F.pad(audio_data, (0, right_pad))1095 if audio_lengths is None:1096 audio_lengths = torch.LongTensor([length + right_pad]).to(audio_data.device)1097 1098 z = self.encoder(audio_data)1099 vq_results = self.quantizer(z, n_quantizers, **kwargs)1100 indices = vq_results.codes1101 indices_lens = torch.ceil(audio_lengths / self.frame_length).long()1102 return indices, indices_lens1103 1104 def decode(self, indices: Tensor, feature_lengths: Tensor):1105 """Decode code indices to audio."""1106 if indices.ndim == 2:1107 indices = indices[None]1108 z = self.quantizer.decode(indices)1109 audio_lengths = feature_lengths * self.frame_length1110 return self.decoder(z), audio_lengths1111 1112 def encode_to_codes(self, audio: Tensor, audio_lengths: Optional[Tensor] = None, n_quantizers: Optional[int] = None, **kw):1113 return self.encode(audio, audio_lengths, n_quantizers, **kw)1114 1115 def decode_codes(self, indices: Tensor, feature_lengths: Tensor):1116 return self.decode(indices, feature_lengths)1117 1118 @torch.no_grad()1119 def encode_zq(self, audio_data: Tensor) -> Tensor:1120 indices, _ = self.encode(audio_data)1121 new_indices = torch.zeros_like(indices)1122 new_indices[:, 0] = torch.clamp(indices[:, 0], max=self.quantizer.semantic_quantizer.codebook_size - 1)1123 new_indices[:, 1:] = torch.clamp(indices[:, 1:], max=self.quantizer.quantizer.codebook_size - 1)1124 1125 z_q_semantic = self.quantizer.semantic_quantizer.from_codes(new_indices[:, :1])[0]1126 z_q_residual = self.quantizer.quantizer.from_codes(new_indices[:, 1:])[0]1127 z_q = z_q_semantic + z_q_residual1128 return z_q1129 1130 @torch.no_grad()1131 def decode_zq(self, z_q: Tensor) -> Tensor:1132 z_q = self.quantizer.post_module(z_q)1133 z_q = self.quantizer.upsample(z_q)1134 return self.decoder(z_q)1135 1136 @property1137 def device(self) -> torch.device: return next(self.parameters()).device1138 1139 @property1140 def dtype(self) -> torch.dtype: return next(self.parameters()).dtype1141 1142# --------------------------------------------------------------------1143# Build helpers1144# --------------------------------------------------------------------1145 1146def build_ae(**cfg) -> DAC:1147 """1148 Factory used by external loaders1149 """1150 # Shared transformer config for the RVQ pre/post modules1151 q_config = ModelArgs(1152 block_size=4096, n_layer=8, n_head=16, dim=1024,1153 intermediate_size=3072, head_dim=64, norm_eps=1e-5,1154 dropout_rate=0.1, attn_dropout_rate=0.1, channels_first=True1155 )1156 1157 def make_transformer():1158 return WindowLimitedTransformer(1159 causal=True, window_size=128, input_dim=1024, config=q_config1160 )1161 1162 quantizer = DownsampleResidualVectorQuantize(1163 input_dim=1024, n_codebooks=9, codebook_size=1024, codebook_dim=8,1164 quantizer_dropout=0.5, downsample_factor=(2, 2),1165 semantic_codebook_size=4096,1166 pre_module=make_transformer(),1167 post_module=make_transformer(),1168 )1169 1170 def transformer_general_config(**kw):1171 return ModelArgs(1172 block_size=kw.get("block_size", 16384),1173 n_layer=kw.get("n_layer", 8),1174 n_head=kw.get("n_head", 8),1175 dim=kw.get("dim", 512),1176 intermediate_size=kw.get("intermediate_size", 1536),1177 n_local_heads=kw.get("n_local_heads", -1),1178 head_dim=kw.get("head_dim", 64),1179 rope_base=kw.get("rope_base", 10000),1180 norm_eps=kw.get("norm_eps", 1e-5),1181 dropout_rate=kw.get("dropout_rate", 0.1),1182 attn_dropout_rate=kw.get("attn_dropout_rate", 0.1),1183 channels_first=kw.get("channels_first", True),1184 )1185 1186 dac = DAC(1187 encoder_dim=64, encoder_rates=[2, 4, 8, 8], latent_dim=1024,1188 decoder_dim=1536, decoder_rates=[8, 8, 4, 2],1189 quantizer=quantizer, sample_rate=44100, causal=True,1190 encoder_transformer_layers=[0, 0, 0, 4],1191 decoder_transformer_layers=[4, 0, 0, 0],1192 transformer_general_config=transformer_general_config,1193 )1194 return dac1195 1196__all__ = [1197 "DAC",1198 "build_ae",1199 "VectorQuantize",1200 "ResidualVectorQuantize",