update moss-audio-tokenizer-v2

.gitattributes

@@ -1,35 +1,47 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,125 @@
 ---
 license: apache-2.0
+library_name: transformers
+tags:
+  - audio
+  - audio-tokenizer
+  - neural-codec
+  - moss-tts-family
+  - MOSS Audio Tokenizer
+  - speech-tokenizer
+  - trust-remote-code
 ---
+
+# MossAudioTokenizer
+
+This is the code for MOSS-Audio-Tokenizer presented in [MOSS-Audio-Tokenizer: Scaling Audio Tokenizers for Future Audio Foundation Models](https://arxiv.org/abs/2602.10934). 
+
+**MOSSAudioTokenizer** is a unified discrete audio tokenizer based on the **Cat** (**C**ausal **A**udio **T**okenizer with **T**ransformer) architecture. Scaling to 1.6 billion parameters, it functions as a unified discrete interface, delivering both lossless-quality reconstruction and high-level semantic alignment.
+
+**Key Features:**
+
+*   **Extreme Compression & Variable Bitrate**: It compresses 48kHz stereo audio into a remarkably low frame rate of 12.5Hz. Utilizing a 32-layer Residual LFQ quantizer stack, it supports high-fidelity reconstruction across a wide range of bitrates.
+*   **Pure Transformer Architecture**: The model features a "CNN-free" homogeneous architecture built entirely from Causal Transformer blocks. With 1.6B combined parameters (Encoder + Decoder), it ensures exceptional scalability and supports low-latency streaming inference.
+*   **Large-Scale General Audio Training**: Trained on 3 million hours of diverse audio data, the model excels at encoding and reconstructing all audio domains, including speech, sound effects, and music.
+*   **Unified Semantic-Acoustic Representation**: While achieving state-of-the-art reconstruction quality, Cat produces discrete tokens that are "semantic-rich," making them ideal for downstream tasks like speech understanding (ASR) and generation (TTS).
+*   **Fully Trained From Scratch**: Cat does not rely on any pretrained encoders (such as HuBERT or Whisper) or distillation from teacher models. All representations are learned autonomously from raw data.
+*   **End-to-End Joint Optimization**: All components—including the encoder, quantizer, decoder, discriminator, and a decoder-only LLM for semantic alignment—are optimized jointly in a single unified training pipeline.
+
+**Summary:**
+By combining a simple, scalable architecture with massive-scale data, the Cat architecture overcomes the bottlenecks of traditional audio tokenizers. It provides a robust, high-fidelity, and semantically grounded interface for the next generation of native audio foundation models.
+
+This repository contains a lightweight remote-code implementation that mirrors the current 🤗 Transformers
+`transformers.models.moss_audio_tokenizer` module. It is intended to be uploaded to a Hugging Face Hub model repository
+and loaded with `trust_remote_code=True` when needed.
+
+## Usage
+
+### Quickstart
+
+```python
+import torch
+from transformers import AutoModel
+import torchaudio
+
+repo_id = "OpenMOSS-Team/MOSS-Audio-Tokenizer"
+model = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval()
+
+wav, sr = torchaudio.load('demo/demo_gt.wav')
+if sr != model.sampling_rate:
+    wav = torchaudio.functional.resample(wav, sr, model.sampling_rate)
+if wav.shape[0] == 1:
+    wav = wav.repeat(model.config.number_channels, 1)
+else:
+    wav = wav[: model.config.number_channels]
+wav = wav.unsqueeze(0)
+enc = model.encode(wav, return_dict=True)
+print(f"enc.audio_codes.shape: {enc.audio_codes.shape}")
+dec = model.decode(enc.audio_codes, return_dict=True)
+print(f"dec.audio.shape: {dec.audio.shape}")
+wav = dec.audio.squeeze(0)
+torchaudio.save("demo/demo_rec.wav", wav, sample_rate=model.sampling_rate)
+
+# Decode using only the first 8 layers of the RVQ
+dec_rvq8 = model.decode(enc.audio_codes[:8], return_dict=True)
+wav_rvq8 = dec_rvq8.audio.squeeze(0)
+torchaudio.save("demo/demo_rec_rvq8.wav", wav_rvq8, sample_rate=model.sampling_rate)
+```
+
+### Attention Backend And Compute Dtype
+
+`config.attention_implementation` controls whether transformer layers prefer `sdpa` or `flash_attention_2`.
+`config.compute_dtype` controls the non-quantizer autocast dtype and supports `fp32`, `bf16`, and `fp16`.
+
+```python
+model.set_attention_implementation("flash_attention_2")
+model.set_compute_dtype("fp16")
+```
+
+The quantizer always runs in fp32.
+
+### Streaming
+
+`MossAudioTokenizerModel.encode`, `decode`, `batch_encode`, and `batch_decode` all support streaming through a
+`chunk_duration` argument.
+
+- `chunk_duration` is expressed in seconds.
+- `chunk_duration * MossAudioTokenizerConfig.sampling_rate` must be divisible by `MossAudioTokenizerConfig.downsample_rate`.
+- Streaming batch inference is supported.
+- The public waveform interface expects stereo inputs shaped `(2, T)` or batched stereo inputs shaped `(B, 2, T)`.
+
+```python
+import torch
+from transformers import AutoModel
+
+repo_id = "OpenMOSS-Team/MOSS-Audio-Tokenizer"
+model = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval()
+audio = torch.randn(2, 48000 * 6)  # dummy stereo waveform
+
+# 6.0s @ 48kHz = 288000 samples, divisible by downsample_rate=3840
+enc = model.encode(audio.unsqueeze(0), return_dict=True, chunk_duration=0.08)
+dec = model.decode(enc.audio_codes, return_dict=True, chunk_duration=0.08)
+
+batch_enc = model.batch_encode([audio, audio[:, : 48000 * 3]], chunk_duration=0.08)
+codes_list = [
+    batch_enc.audio_codes[:, i, : batch_enc.audio_codes_lengths[i]]
+    for i in range(batch_enc.audio_codes.shape[1])
+]
+batch_dec = model.batch_decode(codes_list, chunk_duration=0.08)
+```
+
+## Repository layout
+
+- `configuration_moss_audio_tokenizer.py`
+- `modeling_moss_audio_tokenizer.py`
+- `__init__.py`
+- `config.json`
+- model weights
+
+
+
+## Citation
+If you use this code or result in your paper, please cite our work as:
+```tex
+
+```

__init__.py

@@ -0,0 +1 @@
+"""Remote code package for Moss audio tokenizer."""

config.json

@@ -0,0 +1,415 @@
+{
+  "architectures": [
+    "MossAudioTokenizerModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_moss_audio_tokenizer.MossAudioTokenizerConfig",
+    "AutoModel": "modeling_moss_audio_tokenizer.MossAudioTokenizerModel"
+  },
+  "model_type": "moss-audio-tokenizer",
+  "sample_rate": 48000,
+  "sampling_rate": 48000,
+  "downsample_rate": 3840,
+  "causal_transformer_context_duration": 10.0,
+  "number_channels": 2,
+  "enable_channel_interleave": true,
+  "attention_implementation": "sdpa",
+  "compute_dtype": "fp32",
+  "dtype": "float32",
+  "code_dim": 768,
+  "encoder_kwargs": [
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 240
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 240,
+      "output_dimension": 384,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 1.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 768,
+      "output_dimension": 384,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 2.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 768,
+      "output_dimension": 384,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 4.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 768,
+      "output_dimension": 384,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 8.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 768,
+      "output_dimension": 640,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 10.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 1280,
+      "output_dimension": 768,
+      "d_model": 1280,
+      "num_heads": 20,
+      "num_layers": 32,
+      "dim_feedforward": 5120,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 10.0
+    }
+  ],
+  "decoder_kwargs": [
+    {
+      "module_type": "Transformer",
+      "input_dimension": 768,
+      "output_dimension": 1280,
+      "d_model": 1280,
+      "num_heads": 20,
+      "num_layers": 32,
+      "dim_feedforward": 5120,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 10.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 640,
+      "output_dimension": 768,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 10.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 384,
+      "output_dimension": 768,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 8.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 384,
+      "output_dimension": 768,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 4.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 384,
+      "output_dimension": 768,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 2.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 384,
+      "output_dimension": 240,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 1.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 240
+    }
+  ],
+  "reversed_decoder_kwargs": [
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 240
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 240,
+      "output_dimension": 384,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 1.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 768,
+      "output_dimension": 384,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 2.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 768,
+      "output_dimension": 384,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 4.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 768,
+      "output_dimension": 384,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 8.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 768,
+      "output_dimension": 640,
+      "d_model": 768,
+      "num_heads": 12,
+      "num_layers": 12,
+      "dim_feedforward": 3072,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 10.0
+    },
+    {
+      "module_type": "PatchedPretransform",
+      "patch_size": 2
+    },
+    {
+      "module_type": "Transformer",
+      "input_dimension": 1280,
+      "output_dimension": 768,
+      "d_model": 1280,
+      "num_heads": 20,
+      "num_layers": 32,
+      "dim_feedforward": 5120,
+      "causal": true,
+      "norm": "layer_norm",
+      "positional_embedding": "rope",
+      "max_period": 10000,
+      "gating": "none",
+      "layer_scale": 0.01,
+      "conv_layout": true,
+      "context_duration": 10.0
+    }
+  ],
+  "quantizer_type": "rlfq",
+  "quantizer_kwargs": {
+    "input_dim": 768,
+    "rvq_dim": 512,
+    "output_dim": 768,
+    "num_quantizers": 32,
+    "codebook_size": 1024,
+    "codebook_dim": 8,
+    "quantizer_type": "rlfq"
+  },
+  "transformers_version": "4.56.0.dev0"
+}

configuration_moss_audio_tokenizer.py

@@ -0,0 +1,467 @@
+# coding=utf-8
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MossAudioTokenizer model configuration."""
+
+from typing import Any
+
+try:
+    from transformers.configuration_utils import PreTrainedConfig
+except ImportError:
+    from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class MossAudioTokenizerConfig(PreTrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MossAudioTokenizerModel`]. It is used to instantiate a
+    MossAudioTokenizer model according to the specified arguments, defining the model architecture.
+
+    Instantiating a configuration with the defaults will yield a similar configuration to that of the
+    [VoiceAgentGroup/moss_audio_tokenizer](https://huggingface.co/VoiceAgentGroup/moss_audio_tokenizer) architecture.
+
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
+
+    Args:
+        sampling_rate (`int`, *optional*, defaults to 48000):
+            The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz).
+        downsample_rate (`int`, *optional*, defaults to 3840):
+            Total downsampling rate from waveform to tokens.
+        causal_transformer_context_duration (`float`, *optional*, defaults to 10.0):
+            Legacy global fallback context duration in seconds for causal transformer. If an individual transformer
+            entry in `encoder_kwargs` or `decoder_kwargs` provides `context_duration`, that per-module value takes
+            precedence.
+        encoder_kwargs (`list[dict]`, *optional*):
+            List of encoder module configurations. Each dict specifies a module type and its parameters.
+        decoder_kwargs (`list[dict]`, *optional*):
+            List of decoder module configurations in execution order.
+        number_channels (`int`, *optional*, defaults to 2):
+            Number of audio channels exposed by the public waveform interface.
+        enable_channel_interleave (`bool`, *optional*, defaults to `True`):
+            Whether to flatten multi-channel waveforms into a single internal stream before codec inference.
+        attention_implementation (`str`, *optional*, defaults to `"sdpa"`):
+            Attention implementation to prefer for transformer layers. Supported values are `"sdpa"` and
+            `"flash_attention_2"`.
+        compute_dtype (`str`, *optional*, defaults to `"fp32"`):
+            Inference compute dtype for non-quantizer modules. Supported values are `"fp32"`, `"bf16"`, and `"fp16"`.
+        quantizer_type (`str`, *optional*, defaults to `"rlfq"`):
+            Quantizer type. Options include `"rvq"`, `"spec_rvq"`, `"rlfq"`, `"random_prefix_rlfq"`.
+        quantizer_kwargs (`dict`, *optional*):
+            Configuration for the quantizer including `input_dim`, `rvq_dim`, `output_dim`, `num_quantizers`,
+            `codebook_size`, and `codebook_dim`.
+
+    Example:
+
+    ```python
+    >>> from transformers import MossAudioTokenizerModel, MossAudioTokenizerConfig
+
+    >>> # Initializing a MossAudioTokenizer style configuration
+    >>> configuration = MossAudioTokenizerConfig()
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = MossAudioTokenizerModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "moss-audio-tokenizer"
+
+    # Backward-compatible alias used by some checkpoints.
+    attribute_map = {"sample_rate": "sampling_rate"}
+
+    sampling_rate: int
+    downsample_rate: int
+    causal_transformer_context_duration: float
+    encoder_kwargs: list[dict[str, Any]]
+    decoder_kwargs: list[dict[str, Any]]
+    number_channels: int
+    enable_channel_interleave: bool
+    attention_implementation: str
+    compute_dtype: str
+    quantizer_type: str
+    quantizer_kwargs: dict[str, Any]
+
+    def __init__(
+        self,
+        version: str | None = None,
+        sampling_rate: int = 48000,
+        downsample_rate: int = 3840,
+        causal_transformer_context_duration: float = 10.0,
+        encoder_kwargs: list[dict[str, Any]] | None = None,
+        decoder_kwargs: list[dict[str, Any]] | None = None,
+        number_channels: int = 2,
+        enable_channel_interleave: bool = True,
+        attention_implementation: str = "sdpa",
+        compute_dtype: str = "fp32",
+        quantizer_type: str = "rlfq",
+        quantizer_kwargs: dict[str, Any] | None = None,
+        **kwargs,
+    ):
+        # Some checkpoints might include an incorrect/legacy `model_type` (e.g. "speech_tokenizer").
+        # We drop it to avoid overriding the class-level `model_type`.
+        kwargs.pop("model_type", None)
+        if "channels_numbers" in kwargs:
+            number_channels = kwargs.pop("channels_numbers")
+        if "enable_channel_interleave" in kwargs:
+            enable_channel_interleave = kwargs.pop("enable_channel_interleave")
+        if "attention_backend" in kwargs and attention_implementation == "sdpa":
+            attention_implementation = kwargs.pop("attention_backend")
+        if "codec_compute_dtype" in kwargs and compute_dtype == "fp32":
+            compute_dtype = kwargs.pop("codec_compute_dtype")
+        reversed_decoder_kwargs = kwargs.pop("reversed_decoder_kwargs", None)
+
+        # `version` is accepted for compatibility but not used in modeling.
+        self.version = version
+        self.sampling_rate = sampling_rate
+        self.downsample_rate = downsample_rate
+        self.causal_transformer_context_duration = causal_transformer_context_duration
+        self.number_channels = number_channels
+        self.enable_channel_interleave = enable_channel_interleave
+        self.attention_implementation = attention_implementation
+        self.compute_dtype = compute_dtype
+        # Default encoder configuration
+        if encoder_kwargs is None:
+            encoder_kwargs = [
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 240,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 240,
+                    "output_dimension": 384,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 1.0,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 768,
+                    "output_dimension": 384,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 2.0,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 768,
+                    "output_dimension": 384,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 4.0,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 768,
+                    "output_dimension": 384,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 8.0,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 768,
+                    "output_dimension": 640,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 10.0,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 1280,
+                    "output_dimension": 768,
+                    "d_model": 1280,
+                    "num_heads": 20,
+                    "num_layers": 32,
+                    "dim_feedforward": 5120,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 10.0,
+                },
+            ]
+        else:
+            encoder_kwargs = [dict(module_kwargs) for module_kwargs in encoder_kwargs]
+        for module_kwargs in encoder_kwargs:
+            if module_kwargs.get("module_type") == "Transformer":
+                module_kwargs.setdefault("context_duration", causal_transformer_context_duration)
+        self.encoder_kwargs = encoder_kwargs
+
+        # Default decoder configuration (execution order)
+        if decoder_kwargs is None and reversed_decoder_kwargs is not None:
+            reversed_decoder_kwargs = [dict(module_kwargs) for module_kwargs in reversed_decoder_kwargs]
+            decoder_kwargs = []
+            for module_kwargs in reversed_decoder_kwargs[::-1]:
+                if module_kwargs.get("module_type") != "Transformer":
+                    decoder_kwargs.append(module_kwargs)
+                    continue
+                module_kwargs = dict(module_kwargs)
+                module_kwargs["input_dimension"], module_kwargs["output_dimension"] = (
+                    module_kwargs["output_dimension"],
+                    module_kwargs["input_dimension"],
+                )
+                decoder_kwargs.append(module_kwargs)
+
+        if decoder_kwargs is None:
+            decoder_kwargs = [
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 768,
+                    "output_dimension": 1280,
+                    "d_model": 1280,
+                    "num_heads": 20,
+                    "num_layers": 32,
+                    "dim_feedforward": 5120,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 10.0,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 640,
+                    "output_dimension": 768,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 10.0,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 384,
+                    "output_dimension": 768,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 8.0,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 384,
+                    "output_dimension": 768,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 4.0,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 384,
+                    "output_dimension": 768,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 2.0,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 384,
+                    "output_dimension": 240,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                    "context_duration": 1.0,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 240,
+                },
+            ]
+        else:
+            decoder_kwargs = [dict(module_kwargs) for module_kwargs in decoder_kwargs]
+        for module_kwargs in decoder_kwargs:
+            if module_kwargs.get("module_type") == "Transformer":
+                module_kwargs.setdefault("context_duration", causal_transformer_context_duration)
+        self.decoder_kwargs = decoder_kwargs
+
+        # Default quantizer configuration
+        if quantizer_kwargs is None:
+            quantizer_kwargs = {
+                "input_dim": 768,
+                "rvq_dim": 512,
+                "output_dim": 768,
+                "num_quantizers": 32,
+                "codebook_size": 1024,
+                "codebook_dim": 8,
+                "quantizer_type": "rlfq",
+            }
+
+        # Handle quantizer_type from kwargs or config
+        kw_qtype = quantizer_kwargs.get("quantizer_type", None)
+        if kw_qtype is not None:
+            self.quantizer_type = kw_qtype
+        else:
+            self.quantizer_type = quantizer_type
+            quantizer_kwargs["quantizer_type"] = quantizer_type
+
+        self.quantizer_kwargs = quantizer_kwargs
+
+        super().__init__(**kwargs)
+
+    @property
+    def num_quantizers(self) -> int:
+        """Return the number of quantizers from quantizer_kwargs."""
+        return self.quantizer_kwargs.get("num_quantizers", 32)
+
+    @property
+    def codebook_size(self) -> int:
+        """Return the codebook size from quantizer_kwargs."""
+        return self.quantizer_kwargs.get("codebook_size", 4096)
+
+    @property
+    def frame_rate(self) -> float:
+        """Return the frame rate (tokens per second)."""
+        return self.sampling_rate / self.downsample_rate
+
+
+__all__ = ["MossAudioTokenizerConfig"]

model-00001-of-00003.safetensors

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

model-00002-of-00003.safetensors

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

model-00003-of-00003.safetensors

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