Upload MOSS Audio Tokenizer v2

.gitattributes

@@ -1,49 +1 @@
-*.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
-*.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
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack 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
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* 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
-*.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
-*.png filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.wav filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,2 +1,3 @@
 dev/*
-demo/demo_rec*.wav
+*.pyc
+__pycache__/

README.md

@@ -1,151 +1,119 @@
----
-license: apache-2.0
-library_name: transformers
-tags:
-  - audio
-  - audio-tokenizer
-  - neural-codec
-  - moss-tts-family
-  - MOSS Audio Tokenizer
-  - speech-tokenizer
-  - trust-remote-code
-  - arxiv:2602.10934
----
-
-# Moss-Audio-Tokenizer-v2
-
-This is the code for the 48khz stereo version of MOSS-Audio-Tokenizer presented in [MOSS-Audio-Tokenizer: Scaling Audio Tokenizers for Future Audio Foundation Models](https://arxiv.org/abs/2602.10934). 
-
-**MOSS-Audio-Tokenizer-v2** is a unified discrete audio tokenizer based on the **Cat** (**C**ausal **A**udio **T**okenizer with **T**ransformer) architecture. Scaling to 2 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 Vector Quantization 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 2B 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 hosted as a Hugging Face Hub model repository and should be
-loaded with `trust_remote_code=True`.
-
-<br>
-<p align="center">
-    <img src="images/arch.png" width="95%"> <br>
-    Architecture of MossAudioTokenizer
-</p>
-<br>
-
-## Usage
-
-### Quickstart
-
-```python
-import torch
-from transformers import AutoModel
-import torchaudio
-
-repo_id = "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval().to(device)
-
-audio_path = "demo/demo_gt.wav"  # replace with your own 48 kHz stereo audio path if needed
-wav, sr = torchaudio.load(audio_path)
-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).to(device)
-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.cpu(), 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.cpu(), sample_rate=model.sampling_rate)
-```
-
-For production use with `trust_remote_code=True`, pin `revision` to a reviewed commit hash.
-
-### 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`.
-
-```python
-model.set_attention_implementation("flash_attention_2")
-model.set_compute_dtype("bf16")
-```
-
-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-v2"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval().to(device)
-audio = torch.randn(2, 48000 * 6).to(device)  # 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.safetensors.index.json`
-- sharded model weights: `model-00001-of-00003.safetensors`, `model-00002-of-00003.safetensors`,
-  `model-00003-of-00003.safetensors`
-- `demo/demo_gt.wav`
-
-## Citation
-If you use this code or result in your paper, please cite our work as:
-```tex
-@misc{gong2026mossaudiotokenizerscaling,
-  title={MOSS-Audio-Tokenizer: Scaling Audio Tokenizers for Future Audio Foundation Models},
-  author={Yitian Gong and Kuangwei Chen and Zhaoye Fei and Xiaogui Yang and Ke Chen and Yang Wang and Kexin Huang and Mingshu Chen and Ruixiao Li and Qingyuan Cheng and Shimin Li and Xipeng Qiu},
-  year={2026},
-  eprint={2602.10934},
-  archivePrefix={arXiv},
-  primaryClass={cs.SD},
-  url={https://arxiv.org/abs/2602.10934}
-}
-```
-
-## License
-MOSS-Audio-Tokenizer-v2 is released under the Apache 2.0 license. See `LICENSE` for the full license text.
-
+---
+license: apache-2.0
+library_name: transformers
+tags:
+  - audio
+  - audio-tokenizer
+  - neural-codec
+  - moss-tts-family
+  - trust-remote-code
+---
+
+# MOSS-Audio-Tokenizer-V2
+
+MOSS-Audio-Tokenizer-V2 is the 48 kHz stereo audio tokenizer used by the
+MOSS-TTS v1.5 local-transformer release. It exposes a Hugging Face
+`AutoModel` interface with `trust_remote_code=True` and supports both waveform
+reconstruction and chunked streaming encode/decode.
+
+## Model Details
+
+| Item | Value |
+| --- | --- |
+| Repository | `OpenMOSS-Team/MOSS-Audio-Tokenizer-V2` |
+| Interface | Hugging Face `AutoModel` remote code |
+| Audio format | 48 kHz stereo |
+| Frame rate | 12.5 Hz |
+| Quantizers | 32 |
+| Weight format | sharded `*.safetensors` |
+| Attention backends | `sdpa`, `flash_attention_2` |
+| Compute dtypes | `fp32`, `bf16`, `fp16` |
+
+The public waveform interface accepts stereo tensors shaped `(2, T)` or batched
+stereo tensors shaped `(B, 2, T)`. Mono inputs can be repeated to stereo before
+encoding.
+
+## Quickstart
+
+```python
+import torch
+import torchaudio
+from transformers import AutoModel
+
+repo_id = "OpenMOSS-Team/MOSS-Audio-Tokenizer-V2"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model = AutoModel.from_pretrained(repo_id, trust_remote_code=True).to(device).eval()
+
+wav, sr = torchaudio.load("input.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)
+elif wav.shape[0] > model.config.number_channels:
+    wav = wav[: model.config.number_channels]
+
+wav = wav.unsqueeze(0).to(device)
+
+with torch.inference_mode():
+    encoded = model.encode(wav, return_dict=True)
+    decoded = model.decode(encoded.audio_codes, return_dict=True)
+
+torchaudio.save("reconstructed.wav", decoded.audio.squeeze(0).cpu(), model.sampling_rate)
+print(encoded.audio_codes.shape)
+```
+
+## Attention Backend And Compute Dtype
+
+`config.attention_implementation` controls the transformer attention backend.
+`config.compute_dtype` controls non-quantizer inference autocast. The quantizer
+runs in fp32.
+
+```python
+model.set_attention_implementation("flash_attention_2")
+model.set_compute_dtype("bf16")
+```
+
+Use `flash_attention_2` only when FlashAttention 2 is installed and the GPU
+supports it. Otherwise use `sdpa`.
+
+## Streaming Encode/Decode
+
+`encode`, `decode`, `batch_encode`, and `batch_decode` support streaming through
+the `chunk_duration` argument.
+
+- `chunk_duration` is expressed in seconds.
+- `chunk_duration * config.sampling_rate` must be divisible by
+  `config.downsample_rate`.
+- Streaming batch inference is supported.
+
+```python
+import torch
+from transformers import AutoModel
+
+repo_id = "OpenMOSS-Team/MOSS-Audio-Tokenizer-V2"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model = AutoModel.from_pretrained(repo_id, trust_remote_code=True).to(device).eval()
+audio = torch.randn(2, 48000 * 6, device=device)
+
+with torch.inference_mode():
+    encoded = model.encode(audio.unsqueeze(0), return_dict=True, chunk_duration=0.08)
+    decoded = model.decode(encoded.audio_codes, return_dict=True, chunk_duration=0.08)
+```
+
+## Repository Files
+
+- `configuration_moss_audio_tokenizer.py`
+- `modeling_moss_audio_tokenizer.py`
+- `__init__.py`
+- `config.json`
+- `model.safetensors.index.json`
+- `model-00001-of-00003.safetensors`
+- `model-00002-of-00003.safetensors`
+- `model-00003-of-00003.safetensors`
+
+## License
+
+Apache License 2.0.

configuration_moss_audio_tokenizer.py

@@ -31,8 +31,7 @@ class MossAudioTokenizerConfig(PreTrainedConfig):
     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.
+    Instantiating a configuration with the defaults will yield a MOSS-Audio-Tokenizer-V2 style architecture.
 
     Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
     documentation from [`PreTrainedConfig`] for more information.
@@ -58,7 +57,7 @@ class MossAudioTokenizerConfig(PreTrainedConfig):
             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"`.
+            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*):

modeling_moss_audio_tokenizer.py

@@ -61,7 +61,7 @@ except ImportError:
 
 
 SUPPORTED_ATTENTION_IMPLEMENTATIONS = {"sdpa", "flash_attention_2"}
-SUPPORTED_COMPUTE_DTYPES = {"fp32": None, "bf16": torch.bfloat16}
+SUPPORTED_COMPUTE_DTYPES = {"fp32": None, "bf16": torch.bfloat16, "fp16": torch.float16}
 
 
 def resolve_compute_dtype(compute_dtype: str) -> torch.dtype | None:
@@ -1361,7 +1361,7 @@ class MossAudioTokenizerPatchedPretransform(nn.Module):
         x = x.reshape(b, d, -1, h).permute(0, 1, 3, 2).reshape(b, d * h, -1)
         # We pad the input waveform to a multiple of `downsample_rate` before applying the encoder.
         # Use a ceil division to match that padding and avoid dropping the last (partially padded) frame.
-        output_lengths = input_lengths // self.patch_size
+        output_lengths = torch.div(input_lengths + self.patch_size - 1, self.patch_size, rounding_mode="floor")
         return x, output_lengths
 
     def decode(self, x, input_lengths):
@@ -2469,7 +2469,7 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
         >>> import torch
         >>> from transformers import MossAudioTokenizerModel
 
-        >>> model = MossAudioTokenizerModel.from_pretrained("moss_audio_tokenizer-model")
+        >>> model = MossAudioTokenizerModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer-V2")
 
         >>> # Create dummy audio input
         >>> audio = torch.randn(1, 1, 24000)  # 1 second of audio at 24kHz
@@ -2485,10 +2485,38 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
         output_audio_codes_lengths: torch.Tensor | None = None
         output_audio: torch.Tensor | None = None
         output_audio_lengths: torch.Tensor | None = None
+        input_audio_lengths: torch.Tensor | None = None
         decoded_from_encoded_codes = False
 
         # Encode if input_values provided
         if input_values is not None:
+            if padding_mask is not None:
+                if padding_mask.dim() == 1:
+                    input_audio_lengths = padding_mask.sum().view(1).long()
+                elif padding_mask.dim() == 2:
+                    input_audio_lengths = padding_mask.sum(dim=-1).long()
+            elif input_values.dim() == 1:
+                input_audio_lengths = torch.full(
+                    (1,),
+                    input_values.shape[-1],
+                    device=input_values.device,
+                    dtype=torch.long,
+                )
+            elif input_values.dim() == 2:
+                batch_size = input_values.shape[0] if self.number_channels == 1 else 1
+                input_audio_lengths = torch.full(
+                    (batch_size,),
+                    input_values.shape[-1],
+                    device=input_values.device,
+                    dtype=torch.long,
+                )
+            elif input_values.dim() == 3:
+                input_audio_lengths = torch.full(
+                    (input_values.shape[0],),
+                    input_values.shape[-1],
+                    device=input_values.device,
+                    dtype=torch.long,
+                )
             encoder_output = self.encode(input_values, padding_mask, num_quantizers, return_dict=True)
             encoder_output = cast(MossAudioTokenizerEncoderOutput, encoder_output)
             output_audio_codes = encoder_output.audio_codes
@@ -2514,6 +2542,18 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
                 decoder_output = cast(MossAudioTokenizerDecoderOutput, decoder_output)
             output_audio = decoder_output.audio
             output_audio_lengths = decoder_output.audio_lengths
+            if decoded_from_encoded_codes and output_audio is not None and input_audio_lengths is not None:
+                input_audio_lengths = input_audio_lengths.to(device=output_audio.device, dtype=torch.long)
+                input_audio_lengths = input_audio_lengths.clamp(min=0, max=output_audio.shape[-1])
+                target_length = int(input_audio_lengths.max().item()) if input_audio_lengths.numel() > 0 else 0
+                if target_length != output_audio.shape[-1] or not bool((input_audio_lengths == target_length).all().item()):
+                    trimmed_audio = output_audio.new_zeros(
+                        (output_audio.shape[0], output_audio.shape[1], target_length)
+                    )
+                    for batch_index, length in enumerate(input_audio_lengths.tolist()):
+                        trimmed_audio[batch_index, :, : int(length)] = output_audio[batch_index, :, : int(length)]
+                    output_audio = trimmed_audio
+                output_audio_lengths = input_audio_lengths
 
         if not return_dict:
             return (output_audio_codes, output_audio, output_audio_lengths)