Upload folder using huggingface_hub

config.json

@@ -0,0 +1,23 @@
+{
+  "architectures": [
+    "DashengTokenizerModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_dasheng_tokenizer.DashengTokenizerConfig",
+    "AutoModel": "modeling_dasheng_tokenizer.DashengTokenizerModel"
+  },
+  "decoder_depth": 12,
+  "decoder_embed_dim": 1280,
+  "decoder_intermediate_size": 5120,
+  "depth": 32,
+  "dtype": "float32",
+  "embed_dim": 1280,
+  "hop_length": 160,
+  "istft_hop": 320,
+  "istft_n_fft": 1280,
+  "model_type": "dashengtokenizer",
+  "n_mels_patch": 128,
+  "num_heads": 16,
+  "transformers_version": "5.1.0",
+  "upsample_tokens": 2
+}

configuration_dasheng_tokenizer.py

@@ -0,0 +1,54 @@
+"""
+Dasheng Audio Tokenizer Configuration
+"""
+
+from transformers import PretrainedConfig
+
+class DashengTokenizerConfig(PretrainedConfig):
+    """
+    Configuration class for DashEng Audio Tokenizer.
+
+    This configuration is used to initialize the DashEng model with the same
+    parameters as the original implementation in models.py.
+
+    Args:
+        target_nmels (int): Number of Mel bins for the frontend. Default: 100
+        decoder_embed_dim (int): Decoder embedding dimension. Default: 768
+        decoder_depth (int): Number of decoder layers. Default: 8
+        decoder_intermediate_size (int): Decoder intermediate size. Default: 1536
+        istft_n_fft (int): ISTFT n_fft parameter. Default: 1280
+        istft_hop (int): ISTFT hop parameter. Default: 640
+        upsample_tokens (int): Upsample factor for tokens. Default: 1
+        n_mels_patch (int): Number of Mel bins for patch embedding. Default: 100
+        hop_length (int): Hop length for Mel spectrogram. Default: 160
+    """
+
+    model_type = "dashengtokenizer"
+
+    def __init__(
+        self,
+        embed_dim: int = 1280,
+        depth:int = 32,
+        num_heads: int = 16,
+        decoder_embed_dim: int = 1280,
+        decoder_depth: int = 12,
+        decoder_intermediate_size: int = 5120,
+        istft_n_fft: int = 1280, 
+        istft_hop: int = 320, # 20ms
+        upsample_tokens: int = 2,
+        n_mels_patch: int = 128, # acoustic nmel
+        hop_length: int = 160, # acoustic hop
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.embed_dim = embed_dim
+        self.depth = depth
+        self.num_heads = num_heads
+        self.decoder_embed_dim = decoder_embed_dim
+        self.decoder_depth = decoder_depth
+        self.decoder_intermediate_size = decoder_intermediate_size
+        self.istft_n_fft = istft_n_fft
+        self.istft_hop = istft_hop
+        self.upsample_tokens = upsample_tokens
+        self.n_mels_patch = n_mels_patch
+        self.hop_length = hop_length

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c930d33c7fb5358c01f2e9e4522f6242da76c3d7eabd5e60f2af841bbe42161d
+size 3220079760

modeling_dasheng_encoder.py

@@ -0,0 +1,248 @@
+from einops import rearrange
+from einops.layers.torch import Rearrange
+import torchaudio.transforms as audio_transforms
+import torch
+import torch.nn as nn
+from typing import Optional, Type
+
+class FrontEnd(nn.Sequential):
+    def __init__(
+        self,
+        f_min: int = 0,
+        sample_rate: int = 16000,
+        win_size: int = 512,
+        center: bool = True,
+        n_fft: int = 512,
+        f_max: Optional[int] = 8000,
+        hop_size: int = 160,
+        n_mels: int = 64,
+    ):
+        self.f_min = f_min
+        self.sample_rate = sample_rate
+        self.win_size = win_size
+        self.center = center
+        self.n_fft = n_fft
+        self.f_max = f_max
+        self.hop_size = hop_size
+        self.n_mels = n_mels
+
+        with torch.device("cpu"):
+            super().__init__(
+                audio_transforms.MelSpectrogram(
+                    f_min=self.f_min,
+                    sample_rate=self.sample_rate,
+                    win_length=self.win_size,
+                    center=self.center,
+                    n_fft=self.n_fft,
+                    f_max=self.f_max,
+                    hop_length=self.hop_size,
+                    n_mels=self.n_mels,
+                ),
+                audio_transforms.AmplitudeToDB(top_db=120),
+            )
+
+    @torch.autocast(enabled=False, device_type="cuda")
+    def forward(self, x, attention_mask=None):
+        """
+        Forward pass of the frontend.
+
+        Args:
+            x: Audio tensor of shape (batch_size, num_samples)
+            attention_mask: Optional attention mask of shape (batch_size, num_samples)
+
+        Returns:
+            features: Mel spectrogram features of shape (batch_size, n_mels, num_frames)
+            attention_mask: Downsampled attention mask of shape (batch_size, num_frames)
+        """
+        features = super().forward(x)
+        if attention_mask is not None:
+            lengths = attention_mask.float().sum(-1) // self.hop_size
+            attention_mask = (torch.arange(features.shape[-1], device=features.device) < lengths.unsqueeze(-1)).int()
+        return features, attention_mask
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Type[torch.nn.Module] = nn.GELU,
+        drop: float = 0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        causal: bool = False,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.scale = (dim // num_heads) ** -0.5
+        self.causal = causal
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, mask: Optional[torch.Tensor] = None):
+        B, N, C = x.shape
+        # qkv: [3, B, heads, N, head_dim]
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+
+        # Apply Causal Mask
+        if self.causal:
+            c_mask = torch.ones(N, N, device=x.device, dtype=torch.bool).triu(1)
+            attn = attn.masked_fill(c_mask, float("-inf"))
+
+        # Apply Padding Mask (B, N) -> (B, 1, 1, N)
+        if mask is not None:
+            if mask.dtype != torch.bool:
+                padding_mask = (mask == 0) 
+            else:
+                padding_mask = mask
+            padding_mask = padding_mask.view(B, 1, 1, N)
+            attn = attn.masked_fill(padding_mask, float("-inf"))
+        attn = attn.softmax(dim=-1).nan_to_num()
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        return self.proj_drop(self.proj(x))
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
+        self.attn = Attention(dim, num_heads, qkv_bias, attn_drop, drop)
+        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
+        self.mlp = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=nn.GELU, drop=drop)
+
+    def forward(self, x, mask=None):
+        x = x + self.attn(self.norm1(x), mask=mask)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class AudioPatchEmbed(torch.nn.Module):
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__()
+        self.stride = kwargs.get('stride', [None, 4])[-1]
+        self.proj = nn.Conv2d(*args, **kwargs)
+
+    def forward(self, x:torch.Tensor, attention_mask:torch.Tensor | None =None):
+        x = self.proj(x)
+        if attention_mask is not None:
+            lengths = attention_mask.float().sum(-1) // self.stride
+            attention_mask = (torch.arange(x.shape[-1], device=x.device) < lengths.unsqueeze(-1)).int()
+        return x, attention_mask
+
+
+
+
+class DashengEncoder(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int = 1280,
+        depth: int = 32,
+        num_heads: int = 20,
+        patch_size=[64, 4],
+        patch_stride=[64, 4],
+        target_length=1008,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.time_patches = patch_stride[-1]
+        self.front_end = FrontEnd()
+        self.target_length = target_length
+        self.max_t_tokens = target_length // patch_stride[-1]
+        self.patch_embed = AudioPatchEmbed(1, embed_dim, kernel_size=patch_size, stride=patch_stride)
+        self.init_bn = nn.Sequential(
+            Rearrange("b c f t -> b f c t"),
+            torch.nn.BatchNorm2d(self.front_end.n_mels, momentum=0.01),
+            Rearrange("b f c t -> b c f t"),
+        )
+
+        self.time_pos_embed = nn.Parameter(torch.randn(1, embed_dim, 1, target_length // self.time_patches) * 0.02)
+        self.freq_pos_embed = nn.Parameter(torch.randn(1, embed_dim, 1, 1) * 0.02)
+
+        self.blocks = nn.ModuleList([Block(embed_dim, num_heads) for _ in range(depth)])
+        self.norm = nn.LayerNorm(embed_dim, eps=1e-6)
+
+    def _forward_main(self, x, attention_mask, mask_to_zero:bool = False):
+        x, attention_mask = self.patch_embed(x, attention_mask)
+        t = x.shape[-1]
+        x = x + self.time_pos_embed[:, :, :, :t] + self.freq_pos_embed
+        x = rearrange(x, "b c f t -> b (f t) c")
+        for block in self.blocks:
+            x = block(x, mask=attention_mask)
+        x = self.norm(x) 
+        if attention_mask is not None and mask_to_zero:
+            x = x * attention_mask.unsqueeze(-1) # Zero out all samples that were masked, but only after first chunk
+        return x 
+
+
+    def forward(self, x: torch.Tensor, attention_mask=None):
+        """
+        Forward pass of the AudioTransformer.
+
+        Args:
+            x: Audio tensor of shape (batch_size, num_samples)
+            attention_mask: Optional attention mask of shape (batch_size, num_samples)
+                           where True indicates valid samples and False indicates padding
+
+        Returns:
+            embeddings: Token embeddings of shape (batch_size, num_tokens, embed_dim)
+        """
+        # Process through frontend - returns features and downsampled mask
+        x, attention_mask = self.front_end(x, attention_mask)
+
+        # Rearrange features for patch embedding: (b f t) -> (b 1 f t)
+        x = rearrange(x, "b f t -> b 1 f t")
+        x = self.init_bn(x)
+
+        input_splits = x.split(self.target_length, dim=-1)
+        masks = [None for _ in range(len(input_splits))]
+        if attention_mask is not None:
+            masks = attention_mask.split(self.target_length, dim=-1)
+
+        outputs = []
+        for i, (input_split_x, mask) in enumerate(zip(input_splits, masks)):
+            output = self._forward_main(input_split_x, attention_mask=mask, mask_to_zero=i != 0)
+            outputs.append(output)
+        x = torch.cat(outputs, dim=1)
+        return x

modeling_dasheng_tokenizer.py

@@ -0,0 +1,191 @@
+from .configuration_dasheng_tokenizer import DashengTokenizerConfig
+from .modeling_dasheng_encoder import DashengEncoder
+from .vocos import VocosModel
+from typing import Optional, Tuple, Union
+import torch
+import torch.nn as nn
+from einops import rearrange
+import torchaudio
+from transformers import PreTrainedModel
+
+
+class VocosMelSpec(torch.nn.Module):
+    """MelSpectrogram frontend for Vocos."""
+    def __init__(self, sample_rate=16000, n_fft=1024, hop_length=256, n_mels=100, padding="center"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.sample_rate = sample_rate
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.n_mels = n_mels
+        with torch.device("cpu"):
+            self.mel_spec = torchaudio.transforms.MelSpectrogram(
+                    sample_rate=self.sample_rate,
+                    n_fft=self.n_fft,
+                    hop_length=self.hop_length,
+                    n_mels=self.n_mels,
+                    center=self.padding == "center",
+                    power=1,)
+
+    def forward(self, audio, **kwargs):
+        if self.padding == "same":
+            pad = self.mel_spec.win_length - self.mel_spec.hop_length
+            audio = torch.nn.functional.pad(audio, (pad // 2, pad // 2), mode="reflect")
+        mel = self.mel_spec(audio)
+        return torch.log(torch.clip(mel, min=1e-7))
+
+
+class DashengTokenizerEncoder(torch.nn.Module):
+    def __init__(
+        self,
+        embed_dim: int = 1280,
+        depth:int = 32,
+        num_heads: int = 16,
+        n_mels_patch: int = 128,
+        hop_length: int = 160,
+        **kwargs,
+    ):
+        super().__init__()
+        self.model = DashengEncoder(embed_dim=embed_dim, depth=depth, num_heads=num_heads)
+        self.embed_dim = int(self.model.embed_dim)
+        self.model.outputlayer = torch.nn.Identity()
+
+        self.front_end = VocosMelSpec(hop_length=hop_length, n_mels=n_mels_patch)
+        self.patch_embed = torch.nn.Conv2d(
+            1, self.model.embed_dim, (n_mels_patch, 4), (n_mels_patch, 4)
+        )
+        self.norm = torch.nn.LayerNorm(self.model.embed_dim)
+
+        # Store parameters for reference
+        self.n_fft = self.model.front_end.n_fft
+        self.hop_size = self.model.front_end.hop_size
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input: torch.Tensor,
+        input_attn_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Forward pass of the encoder.
+
+        Args:
+            input: Audio tensor of shape (batch_size, num_samples)
+            input_attn_mask: Optional attention mask
+
+        Returns:
+            Combined embeddings of shape (batch_size, num_tokens, embed_dim)
+        """
+        with torch.no_grad():
+            semantic_emb = self.model(input, input_attn_mask)
+
+        # acoustic part
+        mel = self.front_end(input).unsqueeze(1)
+        mel_emb = self.patch_embed(mel)
+        acoustic_emb = rearrange(mel_emb, "b c f t -> b (f t) c")
+        acoustic_emb = self.norm(acoustic_emb)
+
+        semantic_emb = semantic_emb[:, : acoustic_emb.shape[1], :]
+        emb = semantic_emb + acoustic_emb
+        return emb
+
+
+class DashengTokenizerPreTrainedModel(PreTrainedModel):
+
+    config_class = DashengTokenizerConfig
+    supports_gradient_checkpointing = True
+
+class DashengTokenizerModel(DashengTokenizerPreTrainedModel):
+    """
+    HuggingFace-compatible DashEng Tokenizer Model (Encoder + Decoder).
+
+    This model includes both the encoder and decoder for end-to-end audio processing.
+    """
+
+    def __init__(self, config: DashengTokenizerConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.encoder = DashengTokenizerEncoder(
+            embed_dim=config.embed_dim,
+            depth = config.depth,
+            num_heads=config.num_heads,
+            n_mels_patch=config.n_mels_patch,
+            hop_length=config.hop_length,
+        )
+
+        self.embed_dim = self.encoder.embed_dim
+
+        # Upsampler (if needed)
+        self.upsampler = None
+        if config.upsample_tokens > 1:
+            self.upsampler = torch.nn.ConvTranspose1d(
+                self.embed_dim, self.embed_dim,
+                kernel_size=config.upsample_tokens,
+                stride=config.upsample_tokens
+            )
+
+        # Decoder
+        self.decoder = VocosModel(
+            input_channels=self.embed_dim,
+            hidden_dim=config.decoder_embed_dim,
+            intermediate_dim=config.decoder_intermediate_size,
+            vocos_istft_hop=config.istft_hop,
+            vocos_n_fft=config.istft_n_fft,
+            num_layers=config.decoder_depth,
+        )
+
+        self.post_init()
+
+    def encode(
+        self,
+        audio: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Encode audio into embeddings."""
+        return self.encoder(audio, attention_mask)
+
+    def decode(self, embeddings: torch.Tensor) -> torch.Tensor:
+        """Decode embeddings back to audio."""
+        if self.upsampler is not None:
+            embeddings = self.upsampler(embeddings.transpose(-2, -1)).transpose(-2, -1)
+        output = self.decoder(embeddings.transpose(-2, -1))
+        return output
+
+    def forward(
+        self,
+        audio: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], dict]:
+        """
+        Forward pass of the DashEng tokenizer.
+
+        Args:
+            audio: Audio tensor of shape (batch_size, num_samples)
+            attention_mask: Optional attention mask
+            output_attentions: Whether to return attention weights
+            output_hidden_states: Whether to return hidden states
+            return_dict: Whether to return a dict
+
+        Returns:
+            Reconstructed audio of shape (batch_size, num_samples)
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Encode
+        embeddings = self.encoder(audio, attention_mask)
+
+        # Decode
+        audio_reconstructed = self.decode(embeddings)
+
+        if not return_dict:
+            return (audio_reconstructed,)
+
+        return {
+            "audio": audio_reconstructed,
+            "embeddings": embeddings,
+        }
+

vocos.py

@@ -0,0 +1,327 @@
+"""
+Standalone Vocos implementation for DashEng HuggingFace models.
+
+This is a minimal, self-contained implementation of Vocos that doesn't depend
+on external vocos libraries, making it suitable for HuggingFace Hub publication.
+"""
+
+import torch
+from torch import nn
+from typing import Optional, Tuple
+
+
+class AdaLayerNorm(nn.Module):
+    """
+    Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes
+
+    Args:
+        num_embeddings (int): Number of embeddings.
+        embedding_dim (int): Dimension of the embeddings.
+    """
+
+    def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.dim = embedding_dim
+        self.scale = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim)
+        self.shift = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim)
+        torch.nn.init.ones_(self.scale.weight)
+        torch.nn.init.zeros_(self.shift.weight)
+
+    def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor:
+        scale = self.scale(cond_embedding_id)
+        shift = self.shift(cond_embedding_id)
+        x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps)
+        x = x * scale + shift
+        return x
+
+
+class ConvNeXtBlock(nn.Module):
+    """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
+
+    Args:
+        dim (int): Number of input channels.
+        intermediate_dim (int): Dimensionality of the intermediate layer.
+        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+            Defaults to None.
+        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
+            None means non-conditional LayerNorm. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        layer_scale_init_value: float,
+        adanorm_num_embeddings: Optional[int] = None,
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
+        self.adanorm = adanorm_num_embeddings is not None
+        if adanorm_num_embeddings:
+            self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
+        else:
+            self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond_embedding_id: Optional[torch.Tensor] = None,
+        speaker_embedding: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
+        if self.adanorm:
+            assert cond_embedding_id is not None
+            x = self.norm(x, cond_embedding_id)
+        else:
+            x = self.norm(x)
+        x = self.pwconv1(x)
+        if speaker_embedding is not None:
+            x = x + speaker_embedding.unsqueeze(1)  # same speaker across all frames
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)
+
+        x = residual + x
+        return x
+
+
+class ISTFT(nn.Module):
+    """
+    Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
+    windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
+    See issue: https://github.com/pytorch/pytorch/issues/62323
+    Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
+    The NOLA constraint is met as we trim padded samples anyway.
+
+    Args:
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames.
+        win_length (int): The size of window frame and STFT filter.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        window = torch.hann_window(win_length)
+        self.register_buffer("window", window)
+
+    def forward(self, spec: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
+
+        Args:
+            spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
+                            N is the number of frequency bins, and T is the number of time frames.
+
+        Returns:
+            Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
+        """
+        if self.padding == "center":
+            # Fallback to pytorch native implementation
+            return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True)
+        elif self.padding == "same":
+            pad = (self.win_length - self.hop_length) // 2
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        assert spec.dim() == 3, "Expected a 3D tensor as input"
+        B, N, T = spec.shape
+
+        # Inverse FFT
+        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
+        ifft = ifft * self.window[None, :, None]
+
+        # Overlap and Add
+        output_size = (T - 1) * self.hop_length + self.win_length
+        y = torch.nn.functional.fold(
+            ifft,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        )[:, 0, 0, pad:-pad]
+
+        # Window envelope
+        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
+        window_envelope = torch.nn.functional.fold(
+            window_sq,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        ).squeeze()[pad:-pad]
+
+        # Normalize
+        assert (window_envelope > 1e-11).all()
+        y = y / window_envelope
+
+        return y
+
+
+class ISTFTHead(nn.Module):
+    """
+    ISTFT Head module for predicting STFT complex coefficients.
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames, which should align with
+                          the resolution of the input features.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
+        super().__init__()
+        out_dim = n_fft + 2
+        self.out = torch.nn.Linear(dim, out_dim)
+        self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
+
+    @torch.autocast(device_type="cuda", enabled=False)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the ISTFTHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x = self.out(x).transpose(1, 2)
+        mag, p = x.chunk(2, dim=1)
+        mag = torch.exp(mag)
+        mag = torch.clip(mag, max=1e2)  # safeguard to prevent excessively large magnitudes
+        # wrapping happens here. These two lines produce real and imaginary value
+        x = torch.cos(p)
+        y = torch.sin(p)
+        # recalculating phase here does not produce anything new
+        # only costs time
+        # phase = torch.atan2(y, x)
+        # S = mag * torch.exp(phase * 1j)
+        # better directly produce the complex value
+        S = mag * (x + 1j * y)
+        audio = self.istft(S)
+        return audio
+
+
+class VocosBackbone(nn.Module):
+    """
+    Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization
+
+    Args:
+        input_channels (int): Number of input features channels.
+        dim (int): Hidden dimension of the model.
+        intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
+        num_layers (int): Number of ConvNeXtBlock layers.
+        layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`.
+        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
+                                                None means non-conditional model. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        dim: int,
+        intermediate_dim: int,
+        num_layers: int,
+        layer_scale_init_value: Optional[float] = None,
+        adanorm_num_embeddings: Optional[int] = None,
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3)
+        self.adanorm = adanorm_num_embeddings is not None
+        if adanorm_num_embeddings:
+            self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
+        else:
+            self.norm = nn.LayerNorm(dim, eps=1e-6)
+        layer_scale_init_value = layer_scale_init_value or 1 / num_layers
+        self.convnext = nn.ModuleList(
+            [
+                ConvNeXtBlock(
+                    dim=dim,
+                    intermediate_dim=intermediate_dim,
+                    layer_scale_init_value=layer_scale_init_value,
+                    adanorm_num_embeddings=adanorm_num_embeddings,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv1d, nn.Linear)):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        bandwidth_id = kwargs.get("bandwidth_id", None)
+        speaker_embedding = kwargs.get("speaker_embedding", None)
+        x = self.embed(x)
+        if self.adanorm:
+            assert bandwidth_id is not None
+            x = self.norm(x.transpose(1, 2), cond_embedding_id=bandwidth_id)
+        else:
+            x = self.norm(x.transpose(1, 2))
+        x = x.transpose(1, 2)
+        for conv_block in self.convnext:
+            x = conv_block(x, cond_embedding_id=bandwidth_id, speaker_embedding=speaker_embedding)
+        x = self.final_layer_norm(x.transpose(1, 2))
+        return x
+
+
+class VocosModel(torch.nn.Module):
+    """
+    Vocos model for audio synthesis from learned representations.
+
+    Args:
+        input_channels (int): Number of input feature channels.
+        hidden_dim (int): Hidden dimension of the model.
+        intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
+        num_layers (int): Number of ConvNeXtBlock layers.
+        vocos_istft_hop (int): Hop length for ISTFT.
+        vocos_n_fft (int): FFT size for ISTFT.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(
+        self,
+        input_channels: int = 1024,
+        hidden_dim: int = 512,
+        intermediate_dim: int = 1536,
+        num_layers: int = 8,
+        vocos_istft_hop: int = 256,
+        vocos_n_fft: int = 1024,
+        padding: str = "same",
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        default_kwargs = dict(
+            input_channels=input_channels, dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=num_layers
+        )
+        self.backbone = VocosBackbone(**default_kwargs)
+        self.head = ISTFTHead(**dict(dim=hidden_dim, n_fft=vocos_n_fft, hop_length=vocos_istft_hop, padding=padding))
+
+    def forward(self, x, **kwargs):
+        x = self.backbone(x, **kwargs)
+        audio_output = self.head(x)
+        return audio_output