emotion-conditioned Indian-English TTS (non-commercial demo, Skit-AI BY-NC)

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,127 @@
+---
+license: cc-by-nc-4.0
+language:
+- en
+library_name: transformers
+pipeline_tag: text-to-speech
+base_model: OpenMOSS-Team/MOSS-TTS-Local-Transformer
+tags:
+- text-to-speech
+- tts
+- emotion
+- expressive
+- indian-english
+- lora
+- non-commercial
+---
+
+# Emotion TTS, Indian-English (proof-of-concept)
+
+A non-commercial research demo. An emotion-conditioned Indian-English voice: one consistent
+female speaker with steerable emotion (neutral, happy, sad, angry, excited, calm, apologetic,
+fear), selected per sentence through a text instruction. It is a LoRA fine-tune of the 1.7B
+MOSS-TTS-Local-Transformer.
+
+> NON-COMMERCIAL. This model was trained on the Skit-AI Emotional TTS dataset, which is
+> licensed CC BY-NC 4.0. The model is therefore released under CC BY-NC 4.0 and must not be
+> used for commercial purposes. For commercial use, obtain a license from Skit-AI.
+
+## Why this exists
+
+Part of the Roxi TTS research line. The commercial Roxi voices are neutral read-speech,
+because their training speaker had little emotional range, so emotion instructions had no
+effect. This demo shows that with an emotionally varied speaker, the same pipeline produces
+genuinely steerable emotion in a single consistent voice, controlled by instruction.
+
+## Emotions and how to steer
+
+Set the `instruction` field to one of the following:
+
+| Emotion | Instruction |
+|---|---|
+| neutral | Speak in a neutral, clear, conversational Indian-English style. |
+| happy | Speak in a happy, cheerful, warm tone, in a clear Indian-English style. |
+| sad | Speak in a sad, downcast, sorrowful tone, in a clear Indian-English style. |
+| angry | Speak in an angry, irritated, forceful tone, in a clear Indian-English style. |
+| excited | Speak in an excited, high-energy, enthusiastic tone, in a clear Indian-English style. |
+| calm | Speak in a calm, relaxed, soothing tone, in a clear Indian-English style. |
+| apologetic | Speak in an apologetic, regretful, gentle tone, in a clear Indian-English style. |
+| fear | Speak in a fearful, anxious, worried tone, in a clear Indian-English style. |
+
+## Model at a glance
+
+| Field | Value |
+|---|---|
+| Base model | OpenMOSS-Team/MOSS-TTS-Local-Transformer (1.7B, Apache-2.0) |
+| Audio tokenizer | OpenMOSS-Team/MOSS-Audio-Tokenizer (Apache-2.0) |
+| Method | LoRA (PEFT), r=16, alpha=32, merged into the base weights, bf16 |
+| Training data | Skit-AI Emotional TTS, single Indian-English female speaker, 8 emotions, about 2825 clips |
+| Output | 24 kHz mono |
+| Voice consistency | 0.958 speaker-similarity across emotions (WavLM-SV) |
+| Emotion range | 2.1x pitch-variation spread across the eight emotions |
+| Speed | Real-time factor about 2.5 on a 16 GB GPU, not real-time |
+
+## Requirements
+
+Built for transformers 4.57.1. Install the MOSS-TTS repository so the model class is importable.
+
+```bash
+pip install "transformers==4.57.1" torch torchaudio soundfile librosa peft
+git clone https://github.com/OpenMOSS/MOSS-TTS.git
+```
+
+## Usage
+
+```python
+import sys, torch, soundfile as sf
+sys.path.insert(0, "MOSS-TTS")
+from transformers import AutoProcessor
+from moss_tts_local.modeling_moss_tts import MossTTSDelayModel
+
+repo = "IOTEverythin/roxi-tts-emotion-demo"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+processor = AutoProcessor.from_pretrained(repo, trust_remote_code=True)
+processor.audio_tokenizer = processor.audio_tokenizer.to(device)
+model = MossTTSDelayModel.from_pretrained(repo, torch_dtype=dtype, attn_implementation="sdpa").to(device).eval()
+
+text = "I just heard the news about the meeting tomorrow."
+instruction = "Speak in an excited, high-energy, enthusiastic tone, in a clear Indian-English style."
+
+conv = [[processor.build_user_message(text=text, instruction=instruction)]]
+batch = processor(conv, mode="generation")
+out = model.generate(
+    input_ids=batch["input_ids"].to(device),
+    attention_mask=batch["attention_mask"].to(device),
+    max_new_tokens=4096, do_sample=True, temperature=0.9,
+)
+audio = processor.decode(out)[0].audio_codes_list[0]
+sf.write("out.wav", audio.float().cpu().numpy(), processor.model_config.sampling_rate)
+```
+
+Generation is autoregressive and can under-generate. If a clip is short, generate a few
+times and keep the longest, then trim silence. Keep sentences to about twelve words.
+
+## Limitations
+
+- Non-commercial license (see below).
+- Not real-time on a single consumer GPU.
+- Eight emotions. Surprise is excluded because the source had no audio for it.
+- Emotion is set per sentence through the instruction, not through inline tags inside a sentence.
+- Requires transformers 4.57.1.
+
+## Attribution and license
+
+This model is released under CC BY-NC 4.0 (non-commercial). It is built on
+MOSS-TTS-Local-Transformer (Apache-2.0) and the MOSS Audio Tokenizer (Apache-2.0). The
+emotion control is learned from the Skit-AI Emotional TTS dataset
+(https://github.com/skit-ai/emotion-tts-dataset), which is licensed CC BY-NC 4.0,
+copyright Skit.ai. Because the training data is non-commercial, this derivative model is
+non-commercial as well.
+
+## Responsible use
+
+This voice is derived from a real dataset speaker. Do not use it to impersonate real people
+or for fraud, social engineering, or deception. Disclose AI-generated audio where required by
+law or policy. Provided as is, without warranty.

added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|audio_end|>": 151653,
+  "<|audio_pad|>": 151654,
+  "<|audio_start|>": 151652,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656
+}

chat_template.jinja

@@ -0,0 +1,4 @@
+{% for message in messages %}<|im_start|>{{ message['role'] }}
+{% if message['content'] is string %}{{ message['content'] }}{% else %}{% for content in message['content'] %}{% if content.get('type') == 'text' %}{{ content['text'] }}{% endif %}{% endfor %}{% endif %}<|im_end|>
+{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
+{% endif %}

config.json

@@ -0,0 +1,93 @@
+{
+  "additional_mlp_ffn_hidden_size": 2048,
+  "architectures": [
+    "MossTTSDelayModel"
+  ],
+  "audio_assistant_delay_slot_token_id": 151662,
+  "audio_assistant_gen_slot_token_id": 151656,
+  "audio_end_token_id": 151653,
+  "audio_pad_code": 1024,
+  "audio_start_token_id": 151652,
+  "audio_user_slot_token_id": 151654,
+  "audio_vocab_size": 1024,
+  "auto_map": {
+    "AutoConfig": "configuration_moss_tts.MossTTSDelayConfig",
+    "AutoModel": "modeling_moss_tts.MossTTSDelayModel"
+  },
+  "dtype": "bfloat16",
+  "hidden_size": 2048,
+  "im_end_token_id": 151645,
+  "im_start_token_id": 151644,
+  "initializer_range": 0.02,
+  "language_config": {
+    "_name_or_path": "Qwen/Qwen3-8B",
+    "architectures": [
+      "Qwen3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "eos_token_id": 151645,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": 6144,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 40960,
+    "max_window_layers": 28,
+    "model_type": "qwen3",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 28,
+    "num_key_value_heads": 8,
+    "pad_token_id": 151643,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "rope_theta": 1000000,
+      "rope_type": "default"
+    },
+    "sliding_window": null,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 155648
+  },
+  "local_ffn_hidden_size": 8960,
+  "local_hidden_size": 1536,
+  "local_num_layers": 4,
+  "model_type": "moss_tts_delay",
+  "n_vq": 32,
+  "pad_token_id": 151643,
+  "sampling_rate": 24000,
+  "transformers_version": "5.0.0",
+  "vocab_size": 155648
+}

configuration_moss_tts.py

@@ -0,0 +1,122 @@
+# 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.
+""" MossTTSDelay model configuration """
+
+from typing import Optional, Union
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+from transformers.models.qwen3 import Qwen3Config
+
+logger = logging.get_logger(__name__)
+
+
+class MossTTSDelayConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MossTTSDelayModel`]. It is used to instantiate an
+    MossTTSDelay 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 MossTTSDelay [MossTTSDelay-8B](https://huggingface.co/OpenMOSS/mosstts-8b) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        language_config (`Union[Qwen3Config, dict]`, *optional*):
+            Configuration for the backbone language model (Qwen3).
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        n_vq (`int`, *optional*, defaults to 32):
+            Number of additional VQ (Vector Quantization) heads/channels for audio. 
+            Determines the number of codebooks used in the audio representation.
+        audio_vocab_size (`int`, *optional*, defaults to 1024):
+            Vocabulary size for the audio tokens (codebooks 1 to N).
+        audio_user_slot_token_id (`int`, *optional*, defaults to 151654):
+            The specific token ID used as a placeholder/slot for user-side audio inputs in the prompt.
+        audio_assistant_gen_slot_token_id (`int`, *optional*, defaults to 151656):
+            The specific token ID representing the generation slot for the assistant's audio output. 
+            Acting as the trigger for the TTS generation process.
+        audio_assistant_delay_slot_token_id (`int`, *optional*, defaults to 151662):
+            The token ID used in the 'Delay Pattern' paradigm to represent the delayed/offset positions 
+            between different VQ channels.
+        audio_start_token_id (`int`, *optional*, defaults to 151652):
+            Special token ID used to denote the start of an audio sequence in the stream.
+        audio_end_token_id (`int`, *optional*, defaults to 151653):
+            Special token ID used to denote the end of an audio sequence (EOS for audio).
+        audio_pad_code (`int`, *optional*, defaults to 1024):
+            The padding value used within the audio VQ codebooks. Typically equals `audio_vocab_size`.
+    """
+    model_type = "moss_tts_delay"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        language_config: Optional[Union[Qwen3Config, dict]] = None,
+        initializer_range: float = 0.02,
+        n_vq: int = 32,
+        pad_token_id: int = 151643,
+        im_start_token_id: int = 151644,
+        im_end_token_id: int = 151645,
+        audio_vocab_size: int = 1024,
+        audio_user_slot_token_id: int = 151654,
+        audio_assistant_gen_slot_token_id: int = 151656,
+        audio_assistant_delay_slot_token_id: int = 151662,
+        audio_start_token_id: int = 151652,
+        audio_end_token_id: int = 151653,
+        audio_pad_code: int = 1024,
+        sampling_rate: int = 24000,
+        additional_mlp_ffn_hidden_size: int = 2048,
+        local_ffn_hidden_size: int = 8960,
+        local_hidden_size: int = 1536,
+        local_num_layers: int = 4,
+        **kwargs,
+    ):
+        if isinstance(language_config, dict):
+            self.language_config = Qwen3Config(**language_config)
+        elif language_config is None:
+            self.language_config = Qwen3Config()
+        else:
+            self.language_config = language_config
+
+        self.initializer_range = initializer_range
+        self.n_vq = n_vq
+        self.audio_vocab_size = audio_vocab_size
+        self.audio_user_slot_token_id = audio_user_slot_token_id
+        self.audio_assistant_gen_slot_token_id = audio_assistant_gen_slot_token_id
+        self.audio_assistant_delay_slot_token_id = audio_assistant_delay_slot_token_id
+        self.audio_start_token_id = audio_start_token_id
+        self.audio_end_token_id = audio_end_token_id
+        self.audio_pad_code = audio_pad_code
+        self.sampling_rate = sampling_rate
+
+        self.hidden_size = self.language_config.hidden_size
+        self.vocab_size = self.language_config.vocab_size
+        self.im_start_token_id = self.language_config
+        self.pad_token_id = pad_token_id
+        self.im_start_token_id = im_start_token_id
+        self.im_end_token_id = im_end_token_id
+        
+        self.additional_mlp_ffn_hidden_size = additional_mlp_ffn_hidden_size
+        self.local_ffn_hidden_size = local_ffn_hidden_size
+        self.local_hidden_size = local_hidden_size
+        self.local_num_layers = local_num_layers
+        
+        super().__init__(**kwargs)
+
+    def to_dict(self):
+        output = super().to_dict()
+        if hasattr(self.language_config, "to_dict"):
+            output["language_config"] = self.language_config.to_dict()
+        else:
+            output["language_config"] = self.language_config
+        return output

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "transformers_version": "5.0.0"
+}

inference_utils.py

@@ -0,0 +1,154 @@
+import torch
+import torchaudio
+import torch.nn.functional as F
+from typing import Optional, List, Tuple
+from tqdm import tqdm
+
+
+def apply_top_k(logits, top_k):
+    batch_size, vocab_size = logits.shape
+    top_k = min(top_k, vocab_size)
+    top_k_values, top_k_indices = torch.topk(logits, top_k, dim=-1)
+    filtered_logits = torch.full_like(logits, float("-inf"))
+    batch_indices = torch.arange(batch_size).unsqueeze(-1)
+    filtered_logits[batch_indices, top_k_indices] = top_k_values
+    return filtered_logits
+
+
+def apply_top_p(logits, top_p):
+    probs = F.softmax(logits, dim=-1)
+    sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
+    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+    sorted_indices_to_remove = cumulative_probs > top_p
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = False
+    batch_size = logits.shape[0]
+    filtered_logits = logits.clone()
+    for i in range(batch_size):
+        indices_to_remove = sorted_indices[i][sorted_indices_to_remove[i]]
+        filtered_logits[i, indices_to_remove] = float("-inf")
+    return filtered_logits
+
+
+def apply_top_p_optimized(logits, top_p):
+    probs = F.softmax(logits, dim=-1)
+    sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
+
+    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+
+    sorted_indices_to_remove = cumulative_probs > top_p
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = False
+
+    indices_to_remove = torch.zeros_like(logits, dtype=torch.bool).scatter_(
+        dim=-1, index=sorted_indices, src=sorted_indices_to_remove
+    )
+
+    logits[indices_to_remove] = float("-inf")
+    return logits
+
+
+def apply_repetition_penalty_delay_pattern(
+    logits: torch.Tensor,
+    prev_tokens: torch.LongTensor,
+    penalty: float,
+):
+    """
+    logits: [B, H, V]  or [N, V]
+    prev_tokens: [B, T, H] or [N, T] or [B, H]
+
+    Apply the repetition penalty independently for each H (VQ head).
+    """
+    if penalty == 1.0 or prev_tokens is None:
+        return logits
+
+    vocab_size = logits.size(-1)
+
+    # Case 1: regular [N, V] (text layer)
+    if logits.dim() == 2:
+        prev_tokens_flat = prev_tokens.reshape(-1)
+        unique_tokens = torch.unique(prev_tokens_flat)
+
+        token_logits = logits[:, unique_tokens]
+        pos_mask = token_logits > 0
+        token_logits[pos_mask] /= penalty
+        token_logits[~pos_mask] *= penalty
+        logits[:, unique_tokens] = token_logits
+        return logits
+
+    # Case 2: Delay Pattern audio [B, H, V]
+    assert logits.dim() == 3, "Delay Pattern audio logits must be [B, H, V]"
+    B, H, V = logits.shape
+
+    for h in range(H):
+        # prev_tokens_h: [B, T] or [B]
+        prev_tokens_h = prev_tokens[..., h].reshape(-1)
+        unique_tokens = torch.unique(prev_tokens_h)
+
+        if unique_tokens.numel() == 0:
+            continue
+
+        token_logits = logits[:, h, unique_tokens]
+        pos_mask = token_logits > 0
+        token_logits[pos_mask] /= penalty
+        token_logits[~pos_mask] *= penalty
+        logits[:, h, unique_tokens] = token_logits
+
+    return logits
+
+
+def sample_token(
+    logits,
+    prev_tokens: Optional[torch.LongTensor] = None,
+    repetition_penalty: float = 1.0,
+    top_p=None,
+    top_k=None,
+    do_sample=True,
+):
+    vocab_size = logits.size(-1)
+
+    # ===== Repetition Penalty (before reshaping!) =====
+    if prev_tokens is not None and repetition_penalty != 1.0:
+        logits = apply_repetition_penalty_delay_pattern(
+            logits,
+            prev_tokens,
+            repetition_penalty,
+        )
+
+    if not do_sample:
+        return torch.argmax(logits, dim=-1)
+
+    # ===== Only flatten after this, for top-k / top-p / multinomial =====
+    original_shape = logits.shape
+    reshaped_logits = logits.view(-1, vocab_size)
+
+    if top_k is not None and top_k > 0:
+        reshaped_logits = apply_top_k(reshaped_logits, top_k)
+
+    if top_p is not None and top_p < 1.0:
+        reshaped_logits = apply_top_p_optimized(reshaped_logits, top_p)
+
+    probs = F.softmax(reshaped_logits, dim=-1)
+    next_tokens = torch.multinomial(probs, num_samples=1)
+
+    return next_tokens.view(original_shape[:-1])
+
+
+def find_last_equal_C(tensor, C):
+    """
+    tensor: torch.Tensor of shape [batch_size, seq_len]
+    C: scalar value to match
+    Returns: torch.Tensor of shape [batch_size] with last indices
+    """
+    mask = (tensor == C).int()  # Shape: [batch_size, seq_len], bool tensor
+    flipped_mask = mask.flip(dims=[1])  # Flip along sequence dimension
+    flipped_indices = flipped_mask.argmax(dim=1)  # First True in flipped
+    seq_len = tensor.shape[1]
+    last_indices = (seq_len - 1) - flipped_indices  # Convert to original indices
+
+    # Optional: Handle cases with no C (set to -1), though problem assumes existence
+    actual_values = tensor[torch.arange(tensor.shape[0]), last_indices]
+    no_match = actual_values != C
+    last_indices[no_match] = -1
+
+    return last_indices

model.safetensors

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

modeling_moss_tts.py

@@ -0,0 +1,880 @@
+import os
+import copy
+import torch
+import torch.nn as nn
+import logging
+import sys
+
+from tqdm import tqdm
+from dataclasses import dataclass
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.utils import ModelOutput
+from transformers.cache_utils import Cache
+from typing import Optional, List, Tuple, Union
+from transformers.loss.loss_utils import ForCausalLMLoss
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.generation.streamers import BaseStreamer
+from transformers.models.qwen3.modeling_qwen3 import Qwen3Model, Qwen3Attention, eager_attention_forward
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.generation.stopping_criteria import StoppingCriteriaList
+from transformers.generation.logits_process import LogitsProcessorList, RepetitionPenaltyLogitsProcessor, TopKLogitsWarper, TopPLogitsWarper, TemperatureLogitsWarper
+from transformers.masking_utils import create_causal_mask
+
+from .inference_utils import find_last_equal_C
+from .configuration_moss_tts import MossTTSDelayConfig
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MossTTSRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [..., dim]
+        norm = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(norm + self.eps)
+        return x * self.weight
+
+
+class MossTTSMLP(nn.Module):
+    """
+    HF-style MLP adapter equivalent to Megatron's SwiGLU FFN:
+      in:  input_size
+      mid: ffn_hidden_size
+      out: output_size
+
+    Computes:
+      y = down( silu(gate(x)) * up(x) )
+
+    Optionally includes a pre-norm on input (common in Megatron blocks).
+    """
+    def __init__(
+        self,
+        input_size: int,
+        ffn_hidden_size: int,
+        output_size: int,
+        bias: bool = False,
+        prenorm: bool = False,
+        norm_eps: float = 1e-6,
+        use_rmsnorm: bool = True,
+    ):
+        super().__init__()
+
+        self.prenorm = prenorm
+        if prenorm:
+            if use_rmsnorm:
+                self.norm = MossTTSRMSNorm(input_size, eps=norm_eps)
+            else:
+                self.norm = nn.LayerNorm(input_size, eps=norm_eps)
+        else:
+            self.norm = None
+
+        # SwiGLU uses two projections to ffn_hidden_size: gate and up
+        self.gate_proj = nn.Linear(input_size, ffn_hidden_size, bias=bias)
+        self.up_proj   = nn.Linear(input_size, ffn_hidden_size, bias=bias)
+
+        # down projection to output_size (note: output can differ from input)
+        self.down_proj = nn.Linear(ffn_hidden_size, output_size, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.norm is not None:
+            x = self.norm(x)
+
+        gate = self.gate_proj(x)
+        up   = self.up_proj(x)
+        h = F.silu(gate) * up
+        y = self.down_proj(h)
+        return y
+
+def moss_tts_masked_embedding(embedding: nn.Embedding,
+                     input_ids: torch.LongTensor,
+                     ignore_index: int = -100) -> torch.Tensor:
+    """
+    对 input_ids 中 != ignore_index 的位置做 embedding，ignore_index 的位置输出全 0 向量。
+
+    Args:
+        embedding: 一个 nn.Embedding 层
+        input_ids: 任意形状的 LongTensor，里面允许出现 ignore_index
+        ignore_index: 需要被忽略的位置标记（默认 -100）
+
+    Returns:
+        embeddings: 形状为 (*input_ids.shape, embedding.embedding_dim) 的张量
+    """
+    # mask: True 表示需要正常 embedding，False 表示输出 0
+    mask = (input_ids != ignore_index)  # shape: [...]
+
+    # 为了避免 -100 这种非法 index 传进 embedding，这里先临时替换掉
+    safe_ids = input_ids.clone()
+    safe_ids[~mask] = 0
+
+    # 正常过 embedding
+    out = embedding(safe_ids)  # shape: [..., dim]
+
+    # 把 ignore_index 对应的位置置 0
+    # NOTE: out-of-place mask (was `out[~mask] = 0.0`) — the in-place variant raises
+    # "a leaf Variable that requires grad is being used in an in-place operation" once
+    # LoRA + gradient-checkpointing make the embedding output track grad.
+    out = out * mask.unsqueeze(-1).to(out.dtype)
+
+    return out
+
+class MossTTSAttentionWithoutPositionalEmbedding(Qwen3Attention):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: MossTTSDelayConfig, layer_idx: int):
+        super().__init__(config, layer_idx)
+
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        assert past_key_value is None
+
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
+                print(
+                    "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
+                    'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+                )
+            else:
+                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            is_causal=True,
+            attention_mask=None,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,  # diff with Llama
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+class MossTTSLocalTransformer(Qwen3Model):
+    def __init__(self, config: MossTTSDelayConfig):
+        super().__init__(config)
+        del self.rotary_emb
+        del self.embed_tokens
+        for layer_idx in range(config.num_hidden_layers):
+            self.layers[layer_idx].self_attn = MossTTSAttentionWithoutPositionalEmbedding(config, layer_idx)
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **flash_attn_kwargs,
+    ) -> BaseModelOutputWithPast:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        use_cache = False
+        assert not use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            print(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache
+        if not isinstance(past_key_values, (type(None), Cache)):
+            raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            assert False
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # causal_mask = self._update_causal_mask( # ???
+        #     attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+        # )
+        mask_kwargs = {
+            "config": self.config,
+            "input_embeds": inputs_embeds,
+            "attention_mask": attention_mask,
+            "cache_position": cache_position,
+            "past_key_values": past_key_values,
+            "position_ids": position_ids,
+        }
+        causal_mask = create_causal_mask(**mask_kwargs)
+
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        # position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=None,
+                past_key_value=None,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=None,
+                position_embeddings=None,
+                **flash_attn_kwargs,
+            )
+
+            hidden_states = layer_outputs
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+        
+@dataclass
+class MosiTTSOutputWithPast(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    loss_all: Optional[Tuple[torch.FloatTensor]] = None
+    logits_all: Optional[Tuple[torch.FloatTensor]] = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class MossTTSGenerateDecoderOnlyOutput(ModelOutput):
+    sequences: torch.LongTensor = None
+    scores: Optional[Tuple[torch.FloatTensor]] = None
+    logits: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    past_key_values: Optional[Tuple[Tuple[Tuple[torch.FloatTensor]]]] = None
+
+
+class CustomMixin(GenerationMixin): # TODO 待检查正确性
+    def _sample(
+        self,
+        input_ids: torch.LongTensor, # (B, T, 1+Nq)
+        logits_processor: LogitsProcessorList,
+        stopping_criteria: StoppingCriteriaList,
+        generation_config: GenerationConfig,
+        synced_gpus: bool,
+        streamer: Optional["BaseStreamer"] = None,
+        **model_kwargs,
+    ) -> Union[MossTTSGenerateDecoderOnlyOutput, torch.LongTensor]:
+        # 提取配置参数
+        # assert False
+        speech_pad_idx = self.config.audio_pad_code
+        device = input_ids.device
+        eos_token_id = generation_config.eos_token_id
+        output_attentions = generation_config.output_attentions
+        output_hidden_states = generation_config.output_hidden_states
+        output_scores = generation_config.output_scores
+        output_logits = generation_config.output_logits
+        return_dict_in_generate = generation_config.return_dict_in_generate
+        max_length = generation_config.max_length
+        has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria)
+        do_sample = generation_config.do_sample
+
+        # 初始化输出元组
+        scores = () if (return_dict_in_generate and output_scores) else None
+        raw_logits = () if (return_dict_in_generate and output_logits) else None
+        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+        decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+
+        # 初始化跟踪变量
+        batch_size, cur_len, channels = input_ids.shape  # channels = 8
+        input_ids_length = cur_len
+        # assert batch_size == 1
+        this_peer_finished = False
+        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) # (B, )
+        base_length = input_ids.shape[1]
+        model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs)
+        # model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
+
+        # 定义logits processor
+        if generation_config.do_samples is not None:
+            do_samples = generation_config.do_samples
+            realprocessor = [LogitsProcessorList() for _ in range(channels)]
+            for i, layer_config in enumerate(generation_config.layers):
+                if not do_samples[i]:
+                    continue
+                if layer_config.get("repetition_penalty") is not None and i != 0: # 文本层不用重复惩罚
+                    realprocessor[i].append(RepetitionPenaltyLogitsProcessor(penalty=layer_config.get("repetition_penalty")))
+                if layer_config.get("temperature") is not None:
+                    realprocessor[i].append(TemperatureLogitsWarper(temperature=layer_config.get("temperature")))
+                if layer_config.get("top_k") is not None:
+                    realprocessor[i].append(TopKLogitsWarper(top_k=layer_config.get("top_k")))
+                if layer_config.get("top_p") is not None:
+                    realprocessor[i].append(TopPLogitsWarper(top_p=layer_config.get("top_p")))
+        else:
+            assert False
+            do_samples = [do_sample for _ in range(channels)]
+            realprocessor = [logits_processor for _ in range(channels)]
+        
+        pbar = tqdm()
+        while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+             # 准备模型输入
+            pbar.update()
+            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+            model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
+            model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
+            # 前向传递
+            outputs = self(**model_inputs, n_vq_for_inference=generation_config.n_vq_for_inference, return_dict=True, output_hidden_states=True)
+            model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs)
+
+            if synced_gpus and this_peer_finished:
+                continue
+
+            global_trm_output_hidden_states = outputs.hidden_states[-1][:, -1, :] # (B, D)
+            dtype = global_trm_output_hidden_states.dtype
+
+            local_trm_dim = self.local_transformer_config.hidden_size
+            local_transformer_inputs = torch.zeros(batch_size, 0, local_trm_dim).to(device).to(dtype) # (B, 0 <= t <= Nq, D), 维护当前 local trm 的输入
+            current_local_transformer_input = self.speech_embedding_to_local_mlp(global_trm_output_hidden_states) # (B, D) 维护当前 timestamp 的 local trm 的输入，
+
+            next_tokens = [] # 1+Nq * (B, )
+            # n_vq_for_inference = int(os.environ['N_VQ_FOR_INFERENCE'])
+            n_vq_for_inference = generation_config.n_vq_for_inference
+            for layer_index in range(min(channels, 1 + n_vq_for_inference)):
+                local_transformer_inputs = torch.cat([local_transformer_inputs, current_local_transformer_input.unsqueeze(1)], dim=1) # (B, t, D)
+                local_transformer_outputs = self.local_transformer(
+                    input_ids=None,
+                    attention_mask=None,
+                    inputs_embeds=local_transformer_inputs # (B, t=1+Nq, D)
+                )[0] # (B, t=1+Nq, D)
+                local_transformer_outputs = self.layer_norm_before_lm_heads[layer_index](
+                    self.local_to_speech_embedding_mlps[layer_index](local_transformer_outputs) # (B, t=1+Nq, D)
+                ) # (B, t=1+Nq, D)
+
+                next_token_logit = self.lm_heads[layer_index](local_transformer_outputs[:, -1, :]) # (B, V)
+                if layer_index != 0:
+                    next_token_logit[:, speech_pad_idx] = -torch.inf
+                next_token_score = realprocessor[layer_index](input_ids[..., layer_index], next_token_logit) # (B, V)
+
+                if do_samples[layer_index]:
+                    channel_ntk = torch.multinomial(nn.functional.softmax(next_token_score, dim=-1), num_samples=1).squeeze(1) # (B, )
+                else:
+                    channel_ntk = torch.argmax(next_token_score, dim=-1) # (B, )
+
+                next_tokens.append(channel_ntk) # 1+Nq * (B, )
+                current_local_transformer_input = self.model.embedding_list[layer_index](channel_ntk) # (B, D)
+                current_local_transformer_input = self.speech_embedding_to_local_mlp(current_local_transformer_input) # (B, D)
+
+            for layer_index in range(1 + n_vq_for_inference, channels):
+                next_tokens.append(torch.zeros((batch_size, )).to(torch.int).to(device))
+            next_tokens = torch.stack(next_tokens, dim=-1)  # (B, 1+Nq)
+
+            if has_eos_stopping_criteria:
+                for i in range(channels):
+                    pddp = eos_token_id if i == 0 else speech_pad_idx
+                    next_tokens[:, i] = next_tokens[:, i] * unfinished_sequences + pddp * (1 - unfinished_sequences)
+            
+            input_ids = torch.cat([input_ids, next_tokens[:, None, :]], dim=1) # (B, T, 1+Nq)
+            if streamer is not None:
+                streamer.put(next_tokens[:, 0].cpu())
+
+            stopping = stopping_criteria(input_ids[..., 0], scores)
+            # stopping = stopping_criteria(input_ids[..., 0], scores)
+            unfinished_sequences = unfinished_sequences & ~stopping
+            this_peer_finished = unfinished_sequences.max() == 0
+
+            if return_dict_in_generate:
+                if output_scores:
+                    assert False
+                    scores += (next_token_scores,)
+                if output_logits:
+                    assert False
+                    raw_logits += (next_token_logits,)
+                if output_attentions:
+                    decoder_attentions += (outputs.attentions,)
+                if output_hidden_states:
+                    decoder_hidden_states += (outputs.hidden_states,)
+
+            cur_len += 1
+            del outputs
+
+        if streamer is not None:
+            streamer.end()
+
+        if return_dict_in_generate:
+            return MossTTSGenerateDecoderOnlyOutput(
+                sequences=input_ids,
+                scores=scores,
+                logits=raw_logits,
+                attentions=decoder_attentions,
+                hidden_states=decoder_hidden_states,
+                past_key_values=model_kwargs.get("past_key_values"),
+            )
+        else:
+            start_indices = find_last_equal_C(input_ids[..., 0], self.config.audio_start_token_id)
+            start_lengths = input_ids_length - start_indices - 1 # voice clone 下是 0，续写情况下是 prompt 音频的长度，不包括 audio_start_token
+            output = []
+            for start_idx, start_length, cur_generation_ids in zip(start_indices, start_lengths, input_ids):
+                output.append((start_length, cur_generation_ids[start_idx:]))
+            
+            return output
+
+
+class MosiTTSPretrainedModel(PreTrainedModel):
+    config_class = MossTTSDelayConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen2DecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_cache_class = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+
+class MosiTTSModel(MosiTTSPretrainedModel):
+    def __init__(self, config: MossTTSDelayConfig):
+        super().__init__(config)
+        self.text_pad_idx = config.pad_token_id
+        self.speech_pad_idx = config.audio_pad_code
+        self.embedding_list = nn.ModuleList([])
+        self.embedding_list.append(nn.Embedding(config.vocab_size, config.hidden_size, self.text_pad_idx))
+        self.channels = 1 + config.n_vq
+        for _ in range(1, self.channels):
+            self.embedding_list.append(nn.Embedding(config.audio_vocab_size + 1, config.hidden_size, self.speech_pad_idx))
+
+        self.language_model = Qwen3Model(config.language_config)
+        self.post_init()
+        self.language_model.embed_tokens.requires_grad_(False)
+
+    def get_input_embeddings(self):
+        return self.embedding_list[0]
+
+    def set_input_embeddings(self, value: nn.Embedding):
+        self.embedding_list[0] = value
+
+    def _prepare_multi_modal_inputs(
+        self,
+        input_ids: torch.LongTensor,
+        n_vq_for_inference: Optional[int] = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        """
+        Prepares multi-modal embeddings from input_ids of shape (batch_size, channels, sequence_length).
+        For channel 0: text + speech tokens, for channels 1 to channels-1: speech tokens padded with speech_pad_token.
+        """
+        batch_size, seq_length, channels = input_ids.shape
+        if channels != self.channels:
+            raise ValueError(f"Expected {self.channels} channels, got {channels}")
+
+        if n_vq_for_inference is None:
+            n_vq_for_inference = self.channels - 1
+        n_vq_for_inference = max(1, min(self.channels - 1, int(n_vq_for_inference)))
+
+        inputs_embeds = torch.zeros(batch_size, seq_length, self.config.hidden_size, device=input_ids.device, dtype=self.embedding_list[0].weight.dtype)
+        for i in range(min(channels, 1 + n_vq_for_inference)):
+            embed_layer = self.embedding_list[i]
+            channel_input = input_ids[...,i]
+            inputs_embeds += embed_layer(channel_input)
+
+        return inputs_embeds # (B, T, D)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,  # Shape: (batch_size, channels, sequence_length)
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if input_ids is not None:
+            inputs_embeds = self._prepare_multi_modal_inputs(input_ids, **kwargs) # (B, T, D)
+
+        outputs = self.language_model(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+        return outputs
+
+
+class MossTTSDelayModel(MosiTTSPretrainedModel, CustomMixin):
+    _tied_weights_keys = None
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config: MossTTSDelayConfig):
+        super().__init__(config)
+        self.model = MosiTTSModel(config)
+        self.channels = 1 + config.n_vq
+        self.weights = [1 for _ in range(self.channels)]
+        self.vocab_size = config.vocab_size
+
+        local_transformer_config = copy.deepcopy(config.language_config)
+        local_transformer_config.num_hidden_layers = config.local_num_layers
+        local_transformer_config.hidden_size = config.local_hidden_size
+        local_transformer_config.intermediate_size = config.local_ffn_hidden_size
+        self.local_transformer_config = local_transformer_config
+        self.local_transformer = MossTTSLocalTransformer(self.local_transformer_config)
+
+        self.speech_embedding_to_local_mlp = MossTTSMLP(
+            input_size=config.hidden_size,
+            ffn_hidden_size=config.additional_mlp_ffn_hidden_size,
+            output_size=config.local_hidden_size
+        )
+        self.local_to_speech_embedding_mlps = nn.ModuleList([
+            MossTTSMLP(
+                input_size=config.local_hidden_size,
+                ffn_hidden_size=config.additional_mlp_ffn_hidden_size,
+                output_size=config.hidden_size
+            )
+            for _ in range(self.channels)
+        ])
+
+        self.layer_norm_before_lm_heads = nn.ModuleList([
+            MossTTSRMSNorm(config.hidden_size)
+            for _ in range(self.channels)
+        ])
+
+        self.lm_heads = nn.ModuleList([])
+        self.lm_heads.append(nn.Linear(config.hidden_size, config.vocab_size, bias=False))
+        for _ in range(1, self.channels):
+            self.lm_heads.append(nn.Linear(config.hidden_size, 1 + config.audio_vocab_size, bias=False))
+        self.post_init()
+        self._freeze_unused_qwen_embeddings()
+        self.register_load_state_dict_post_hook(self._freeze_unused_qwen_embeddings_post_load)
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        model = super().from_pretrained(*args, **kwargs)
+        model._freeze_unused_qwen_embeddings()
+        return model
+
+    def get_input_embeddings(self):
+        return self.model.embedding_list[0]
+
+    def _freeze_unused_qwen_embeddings(self) -> None:
+        self.model.language_model.embed_tokens.requires_grad_(False)
+
+    def _freeze_unused_qwen_embeddings_post_load(self, module, incompatible_keys) -> None:
+        module._freeze_unused_qwen_embeddings()
+
+    def can_generate(self):
+        return True
+
+    def _build_generation_config(
+        self,
+        generation_config: Optional[GenerationConfig] = None,
+        max_new_tokens: Optional[int] = None,
+        text_temperature: Optional[float] = None,
+        text_top_p: Optional[float] = None,
+        text_top_k: Optional[int] = None,
+        text_repetition_penalty: Optional[float] = None,
+        audio_temperature: Optional[float] = None,
+        audio_top_p: Optional[float] = None,
+        audio_top_k: Optional[int] = None,
+        audio_repetition_penalty: Optional[float] = None,
+        n_vq_for_inference: Optional[int] = None,
+    ) -> GenerationConfig:
+        config = copy.deepcopy(generation_config or self.generation_config)
+
+        text_temperature = 1.5 if text_temperature is None else float(text_temperature)
+        text_top_p = 1.0 if text_top_p is None else float(text_top_p)
+        text_top_k = 50 if text_top_k is None else int(text_top_k)
+        text_repetition_penalty = 1.0 if text_repetition_penalty is None else float(text_repetition_penalty)
+        audio_temperature = 1.0 if audio_temperature is None else float(audio_temperature)
+        audio_top_p = 0.95 if audio_top_p is None else float(audio_top_p)
+        audio_top_k = 50 if audio_top_k is None else int(audio_top_k)
+        audio_repetition_penalty = 1.1 if audio_repetition_penalty is None else float(audio_repetition_penalty)
+
+        text_do_sample = text_temperature > 0
+        if not text_do_sample:
+            text_temperature = 1.0
+        audio_do_sample = audio_temperature > 0
+        if not audio_do_sample:
+            audio_temperature = 1.0
+
+        if max_new_tokens is not None:
+            config.max_new_tokens = int(max_new_tokens)
+        elif getattr(config, "max_new_tokens", None) is None:
+            config.max_new_tokens = 100000 # about 2.2 hours , can be overridden by user input, you can set to a smaller value for faster generation during debugging
+
+        if getattr(config, "pad_token_id", None) is None:
+            config.pad_token_id = self.config.pad_token_id
+        config.eos_token_id = self.config.audio_end_token_id
+        config.use_cache = True
+        config.do_sample = text_do_sample or audio_do_sample
+
+        resolved_n_vq = self.channels - 1 if n_vq_for_inference is None else int(n_vq_for_inference)
+        resolved_n_vq = max(1, min(self.channels - 1, resolved_n_vq))
+        config.n_vq_for_inference = resolved_n_vq
+        config.do_samples = [text_do_sample] + [audio_do_sample] * (self.channels - 1)
+        config.layers = [
+            {
+                "repetition_penalty": text_repetition_penalty,
+                "temperature": text_temperature,
+                "top_p": text_top_p,
+                "top_k": text_top_k,
+            }
+        ] + [
+            {
+                "repetition_penalty": audio_repetition_penalty,
+                "temperature": audio_temperature,
+                "top_p": audio_top_p,
+                "top_k": audio_top_k,
+            }
+            for _ in range(self.channels - 1)
+        ]
+        return config
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        max_new_tokens: Optional[int] = None,
+        text_temperature: Optional[float] = None,
+        text_top_p: Optional[float] = None,
+        text_top_k: Optional[int] = None,
+        text_repetition_penalty: Optional[int] = None,
+        audio_temperature: Optional[float] = None,
+        audio_top_p: Optional[float] = None,
+        audio_top_k: Optional[int] = None,
+        audio_repetition_penalty: Optional[float] = None,
+        n_vq_for_inference: Optional[int] = None,
+        **kwargs,
+    ):
+        resolved_generation_config = self._build_generation_config(
+            generation_config=generation_config,
+            max_new_tokens=max_new_tokens,
+            text_temperature=text_temperature,
+            text_top_p=text_top_p,
+            text_top_k=text_top_k,
+            text_repetition_penalty=text_repetition_penalty,
+            audio_temperature=audio_temperature,
+            audio_top_p=audio_top_p,
+            audio_top_k=audio_top_k,
+            audio_repetition_penalty=audio_repetition_penalty,
+            n_vq_for_inference=n_vq_for_inference,
+        )
+        return super().generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            generation_config=resolved_generation_config,
+            **kwargs,
+        )
+
+    # def tie_weights(self):
+    #     ...
+        # for i in range(self.config.channels):
+        #     self._tie_or_clone_weights(self.lm_heads[i], self.model.embedding_list[i])
+
+    def set_input_embeddings(self, value):
+        self.model.embedding_list[0] = value
+
+    def get_output_embeddings(self):
+        return self.lm_heads[0]
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_heads[0] = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def set_weights(self, weights):
+        self.weights = weights
+
+    def _prepare_shifted_audio_inputs(self, label_ids): # (B, T, 1 + Nq) 可能有 -100
+        text_and_audio_label_embed_list = [] # Nq * (1, T, B, D)
+        for i in range(self.channels - 1):
+            text_and_audio_label_embed_list.append(
+                moss_tts_masked_embedding(self.model.embedding_list[i], label_ids[:, :, i]).unsqueeze(0).transpose(1, 2) # (B, T) -> (B, T, D) -> (1, B, T, D) -> (1, T, B, D)
+            ) # (1, T, B, D)
+        audio_label_embeds = torch.stack(text_and_audio_label_embed_list, dim=0) # (Nq, 1, T, B, D)
+        audio_label_embeds = audio_label_embeds.contiguous()[:, 0, :, :, :].transpose(1, 2) # (Nq, B, T, D)
+        return audio_label_embeds # (Nq, B, T, D)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None, # (B, T, 1 + Nq)
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None, # (B, T, 1 + Nq), TODO labels 为 input_ids shift 一位的结果
+        channelwise_loss_weight: Optional[List[float]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, MosiTTSOutputWithPast]:
+        device = input_ids.device if not input_ids is None else inputs_embeds.device
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids, # (B, T, 1 + Nq)
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        if labels is not None:
+            local_transformer_inputs_from_global = outputs[0].unsqueeze(0) # (1, B, T, D)
+            D_global= local_transformer_inputs_from_global.shape[-1]
+            local_transformer_inputs_from_speech_embeddings = self._prepare_shifted_audio_inputs(labels) # (B, T, 1 + Nq) -> (Nq, B, T, D)
+            local_transformer_input_hidden_states = torch.cat([local_transformer_inputs_from_global, local_transformer_inputs_from_speech_embeddings], dim=0).contiguous() # (1 + Nq, B, T, D)
+            local_transformer_input_hidden_states = self.speech_embedding_to_local_mlp(local_transformer_input_hidden_states) # (1 + Nq, B, T, D)
+            N_channels, B, T, D_local = local_transformer_input_hidden_states.shape
+            local_transformer_input_hidden_states = local_transformer_input_hidden_states.permute(1, 2, 0, 3) # (B, T, 1 + Nq, D)
+            local_transformer_input_hidden_states = local_transformer_input_hidden_states.reshape(B * T, N_channels, D_local) # (batch_size=B * T, time=1+Nq, D)
+            local_transformer_output_hidden_states = self.local_transformer( # TODO 没有开位置编码
+                input_ids=None,
+                attention_mask=None,
+                inputs_embeds=local_transformer_input_hidden_states # (batch_size=B * T, time=1+Nq, D)
+            )[0] # (batch_size=B * T, time=1+Nq, D)
+            after_lm_head_mlp_hidden_states = [] # Nq+1 * (B*T, D) TODO ???
+            for i in range(self.channels):
+                after_lm_head_mlp_hidden_states.append(
+                    self.layer_norm_before_lm_heads[i](
+                        self.local_to_speech_embedding_mlps[i](
+                            local_transformer_output_hidden_states[:, i, :] # (B*T, D)
+                        )
+                    )
+                )  # Nq+1 * (B*T, D)
+
+            after_lm_head_mlp_hidden_states = torch.stack(after_lm_head_mlp_hidden_states, dim=0)  # (1 + Nq, B*T, D)
+            after_lm_head_mlp_hidden_states = after_lm_head_mlp_hidden_states.reshape(N_channels, B, T, D_global) # (1 + Nq, B, T, D)
+            logits_all = [lm_head(h_i) for lm_head, h_i in zip(self.lm_heads, after_lm_head_mlp_hidden_states)] # 1+Nq * (B, T, V)
+
+            loss_all = torch.empty(self.channels, device=device) # (1 + Nq)
+
+            for i in range(self.channels):
+                vocab_size = self.config.vocab_size if i == 0 else (1 + self.config.audio_vocab_size)
+                loss_all[i] = ForCausalLMLoss(logits_all[i], labels[..., i], vocab_size, shift_labels=labels[..., i]) # (B, T, V), (B, T) => (1, )
+            effective_weights = channelwise_loss_weight if channelwise_loss_weight is not None else self.weights
+            if len(effective_weights) != self.channels:
+                raise ValueError(
+                    f"`channelwise_loss_weight` length {len(effective_weights)} != {self.channels}."
+                )
+            total_weight = float(sum(effective_weights))
+            if total_weight <= 0:
+                raise ValueError("`channelwise_loss_weight` must sum to a positive value.")
+            normalized_weights = [float(weight_i) / total_weight for weight_i in effective_weights] # (1+Nq, )
+
+            total_loss = 0
+            for w, loss in zip(normalized_weights, loss_all):
+                total_loss += w * loss
+        else:
+            total_loss = None
+            loss_all = None
+            logits_all = [None] * self.channels
+
+        if not return_dict:
+            output = (logits_all,) + outputs[1:]
+            return (total_loss, loss_all) + output if total_loss is not None else output
+
+        return MosiTTSOutputWithPast(
+            loss=total_loss,
+            logits=logits_all[0],
+            loss_all=loss_all,
+            logits_all=logits_all, # 1+Nq * (B, T, V)
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states, # L * (B, T, D)
+            attentions=outputs.attentions,
+        )

processing_moss_tts.py

@@ -0,0 +1,950 @@
+# 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.
+
+import os
+from typing import Any, Dict, List, Optional, Tuple, Type, Union, Literal, Final, cast
+from dataclasses import dataclass
+from pathlib import Path
+import re
+import torchaudio
+
+from transformers import processing_utils
+
+processing_utils.MODALITY_TO_BASE_CLASS_MAPPING["audio_tokenizer"] = "PreTrainedModel"
+
+import torch
+from transformers import (
+    PreTrainedTokenizerBase,
+    BatchFeature,
+    ProcessorMixin,
+    logging,
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+)
+
+from .configuration_moss_tts import MossTTSDelayConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+AUDIO_PLACEHOLDER = "<|audio|>"
+
+
+@dataclass
+class Message:
+    def to_dict(self) -> Dict[str, Any]:
+        raise NotImplementedError
+
+
+@dataclass
+class UserMessage(Message):
+    text: Optional[str] = None
+    reference: Optional[List[Optional[Union[str, torch.Tensor]]]] = None
+    instruction: Optional[str] = None
+    tokens: Optional[int] = None
+    quality: Optional[str] = None
+    sound_event: Optional[str] = None
+    ambient_sound: Optional[str] = None
+    language: Optional[str] = None
+
+    def __post_init__(self):
+        template = """<user_inst>
+- Reference(s):
+{reference}
+- Instruction:
+{instruction}
+- Tokens:
+{tokens}
+- Quality:
+{quality}
+- Sound Event:
+{sound_event}
+- Ambient Sound:
+{ambient_sound}
+- Language:
+{language}
+- Text:
+{text}
+</user_inst>"""
+
+        audio_codes_list = []
+        if self.reference is None:
+            reference = "None"
+        elif isinstance(self.reference, List):
+            reference = []
+            for speaker_idx, speaker_reference in enumerate(self.reference):
+                if speaker_reference is not None:
+                    reference.append(f"[S{speaker_idx+1}]:\n{AUDIO_PLACEHOLDER}")
+            reference = "\n".join(reference)
+            audio_codes_list = [
+                speaker_reference
+                for speaker_reference in self.reference
+                if speaker_reference is not None
+            ]
+        else:
+            raise TypeError("`reference` should be exactly a list when it is not None.")
+
+        content = (
+            template.replace("{reference}", str(reference))
+            .replace("{instruction}", str(self.instruction))
+            .replace("{tokens}", str(self.tokens))
+            .replace("{quality}", str(self.quality))
+            .replace("{sound_event}", str(self.sound_event))
+            .replace("{ambient_sound}", str(self.ambient_sound))
+            .replace("{language}", str(self.language))
+            .replace("{text}", str(self.text))
+        )
+
+        self._content = content
+        self._audio_codes_list = audio_codes_list
+
+    def to_dict(self):
+        return {
+            "role": "user",
+            "content": self._content,
+            "audio_codes_list": self._audio_codes_list,
+        }
+
+
+@dataclass
+class AssistantMessage(Message):
+    audio_codes_list: List[Union[str, torch.Tensor]]
+    content: str = AUDIO_PLACEHOLDER
+
+    def to_dict(self):
+        return {
+            "role": "assistant",
+            "content": self.content,
+            "audio_codes_list": self.audio_codes_list,
+        }
+
+
+USER_MESSAGE_FIELDS = (
+    "text",
+    "reference",
+    "instruction",
+    "tokens",
+    "quality",
+    "sound_event",
+    "ambient_sound",
+    "language",
+)
+
+
+class MossTTSDelayProcessor(ProcessorMixin):
+    tokenizer_class = "AutoTokenizer"
+    audio_tokenizer_class = "AutoModel"
+
+    tokenizer: PreTrainedTokenizerBase
+    audio_tokenizer: Any
+
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        audio_tokenizer: Any = None,
+        model_config: Optional[MossTTSDelayConfig] = None,
+        **kwargs,
+    ):
+        super().__init__(tokenizer=tokenizer, audio_tokenizer=audio_tokenizer, **kwargs)
+
+        # Explicit assignments for type-checkers; ProcessorMixin sets these too.
+        self.tokenizer = tokenizer
+        self.audio_tokenizer = audio_tokenizer
+        if model_config is None:
+            model_config = MossTTSDelayConfig()
+        self.model_config = model_config
+
+        self.imstart_token_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
+        self.imend_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
+        self.newline_token_id = 198
+
+        def _id_to_token(token_id: int) -> str:
+            tok = tokenizer.convert_ids_to_tokens(int(token_id))
+            if isinstance(tok, list):
+                return tok[0] if len(tok) > 0 else ""
+            return cast(str, tok)
+
+        self.audio_user_slot_token = _id_to_token(
+            self.model_config.audio_user_slot_token_id
+        )
+        self.audio_assistant_gen_slot_token = _id_to_token(
+            self.model_config.audio_assistant_gen_slot_token_id
+        )
+        self.audio_assistant_delay_slot_token = _id_to_token(
+            self.model_config.audio_assistant_delay_slot_token_id
+        )
+        self.audio_start_token = _id_to_token(self.model_config.audio_start_token_id)
+        self.audio_end_token = _id_to_token(self.model_config.audio_end_token_id)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        trust_remote_code = kwargs.pop("trust_remote_code", True)
+        kwargs.pop("_from_auto", None)
+
+        audio_tokenizer_name_or_path = kwargs.pop("codec_path", None)
+        if audio_tokenizer_name_or_path is None:
+            processor_lookup_kwargs = dict(kwargs)
+            try:
+                processor_dict, _ = cls.get_processor_dict(
+                    pretrained_model_name_or_path,
+                    **processor_lookup_kwargs,
+                )
+                audio_tokenizer_name_or_path = processor_dict.get(
+                    "audio_tokenizer_name_or_path"
+                )
+                audio_tokenizer_dict = processor_dict.get("audio_tokenizer", {})
+                if isinstance(audio_tokenizer_dict, dict):
+                    audio_tokenizer_name_or_path = audio_tokenizer_dict.get(
+                        "audio_tokenizer_name_or_path"
+                    ) or audio_tokenizer_name_or_path
+            except Exception:
+                audio_tokenizer_name_or_path = None
+        if audio_tokenizer_name_or_path is None:
+            audio_tokenizer_name_or_path = "OpenMOSS-Team/MOSS-Audio-Tokenizer"
+
+        pretrained_model_name_or_path = str(pretrained_model_name_or_path)
+        local_model_path = Path(pretrained_model_name_or_path)
+        pretrained_model_ref = local_model_path if local_model_path.exists() else pretrained_model_name_or_path
+        model_config = cast(
+            MossTTSDelayConfig,
+            AutoConfig.from_pretrained(
+                pretrained_model_ref,
+                *args,
+                trust_remote_code=trust_remote_code,
+                **kwargs,
+            ),
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_ref,
+            *args,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+        audio_tokenizer = AutoModel.from_pretrained(
+            audio_tokenizer_name_or_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+
+        return cls(
+            tokenizer=tokenizer,
+            audio_tokenizer=audio_tokenizer,
+            model_config=model_config,
+            **kwargs,
+        )
+
+    def __call__(self, *args, **kwargs) -> BatchFeature:
+        conversations = args[0] if len(args) > 0 else kwargs.pop("conversations")
+        mode: str = kwargs.pop("mode", "generation")
+        apply_chat_template: bool = kwargs.pop("apply_chat_template", True)
+        n_vq: Optional[int] = kwargs.pop("n_vq", None)
+
+        # Common ProcessorMixin kwargs that we ignore because we always return torch tensors.
+        kwargs.pop("return_tensors", None)
+        kwargs.pop("padding", None)
+        kwargs.pop("truncation", None)
+
+        """
+        mode only works when a Message is converted to a dict.
+        """
+
+        if mode not in {"generation", "continuation", "computing_loss"}:
+            raise RuntimeError
+
+        if isinstance(conversations, (Message, Dict)):
+            conversations = [conversations]
+
+        truncation = False
+        if mode == "continuation":
+            truncation = True
+
+        input_ids_list = []
+        for conversation in conversations:
+            if isinstance(conversation, (Message, Dict)):
+                conversation = [conversation]
+
+            # Normalize early so downstream logic always deals with dict messages.
+            conversation = [self._normalize_message(m) for m in conversation]
+
+            if (mode == "generation") ^ (len(conversation) % 2 != 0):
+                raise ValueError
+
+            if (mode == "generation") ^ (conversation[-1]["role"] == "user"):
+                raise ValueError
+
+            unified_codes = []
+            for message_idx, message in enumerate(conversation):
+                if apply_chat_template:
+                    add_generation_prompt = (
+                        mode == "generation" and message_idx == len(conversation) - 1
+                    )
+                    try:
+                        content = self.tokenizer.apply_chat_template(
+                            [{"role": message["role"], "content": message["content"]}],
+                            add_generation_prompt=add_generation_prompt,
+                            tokenize=False,
+                        )
+                    except TypeError:
+                        try:
+                            content = self.tokenizer.apply_chat_template(
+                                [
+                                    {
+                                        "role": message["role"],
+                                        "content": message["content"],
+                                    }
+                                ],
+                                add_generation_prompt=add_generation_prompt,
+                            )
+                        except Exception:
+                            logger.warning(
+                                "apply_chat_template failed; fallback to raw content."
+                            )
+                            content = message["content"]
+                else:
+                    content = message["content"]
+
+                if not isinstance(content, str):
+                    content = str(content)
+
+                # Batch-encode all path-based references in one call when possible.
+                # This ensures we actually exercise audio_tokenizer.batch_encode for multi-reference prompts,
+                # instead of repeatedly calling it with batch=1.
+                raw_audio_items = message.get("audio_codes_list", [])
+
+                audio_codes_list: List[torch.Tensor] = []
+                if len(raw_audio_items) > 0:
+                    encoded_items: List[Optional[torch.Tensor]] = [None] * len(
+                        raw_audio_items
+                    )
+                    paths: List[str] = []
+                    path_positions: List[int] = []
+
+                    for idx, item in enumerate(raw_audio_items):
+                        if isinstance(item, torch.Tensor):
+                            if n_vq is not None and item.shape[1] != n_vq:
+                                raise RuntimeError(
+                                    "audio_codes's n_vq is not equal to the parameter `n_vq`. Your can set the parameter `n_vq` as None if you have already tokenzied the wavs."
+                                )
+                            encoded_items[idx] = item
+                            continue
+
+                        if isinstance(item, (str, os.PathLike)):
+                            paths.append(str(item))
+                            path_positions.append(idx)
+                            continue
+
+                        raise TypeError(
+                            "Each audio item must be a torch.Tensor of codes or a path-like string."
+                        )
+
+                    if len(paths) > 0:
+                        encoded_from_paths = self.encode_audios_from_path(paths, n_vq)
+                        if len(encoded_from_paths) != len(paths):
+                            raise RuntimeError(
+                                "encode_audios_from_path returned an unexpected number of items."
+                            )
+                        for pos, codes in zip(path_positions, encoded_from_paths):
+                            encoded_items[pos] = codes
+
+                    audio_codes_list = [cast(torch.Tensor, t) for t in encoded_items]
+                unified_codes.append(
+                    self._get_unified_codes(
+                        message["role"], content, audio_codes_list, truncation
+                    )
+                )
+
+            unified_codes = torch.cat(unified_codes) # (T, Nq)
+            if mode == "generation":
+                audio_start_position_tokens = torch.zeros((1, unified_codes.shape[-1])).to(unified_codes.dtype).to(unified_codes.device) # (1, Nq)
+                audio_start_position_tokens[:, 0] = self.tokenizer.encode(self.audio_start_token)[0]
+                audio_start_position_tokens[:, 1:] = self.model_config.audio_pad_code
+                unified_codes = torch.cat([unified_codes, audio_start_position_tokens], dim=0) # (T, Nq)
+            input_ids_list.append(unified_codes)
+
+        return BatchFeature(data=self._pad(input_ids_list))
+
+    @staticmethod
+    def build_user_message(
+        text: Optional[str] = None,
+        reference: Optional[List[Optional[Union[str, torch.Tensor]]]] = None,
+        instruction: Optional[str] = None,
+        tokens: Optional[int] = None,
+        quality: Optional[str] = None,
+        sound_event: Optional[str] = None,
+        ambient_sound: Optional[str] = None,
+        language: Optional[str] = None,
+    ) -> Dict:
+        if reference is not None and not isinstance(reference, list):
+            reference = [reference]
+        return UserMessage(
+            text=text,
+            reference=reference,
+            instruction=instruction,
+            tokens=tokens,
+            quality=quality,
+            sound_event=sound_event,
+            ambient_sound=ambient_sound,
+            language=language,
+        ).to_dict()
+
+    @staticmethod
+    def build_assistant_message(
+        audio_codes_list: List[Union[str, torch.Tensor]],
+        content: str = AUDIO_PLACEHOLDER,
+    ) -> Dict:
+        return AssistantMessage(
+            audio_codes_list=audio_codes_list,
+            content=content,
+        ).to_dict()
+
+    def _normalize_message(self, message: Union[Message, Dict]) -> Dict:
+        if isinstance(message, Message):
+            return message.to_dict()
+        if not isinstance(message, dict):
+            raise TypeError("Each message must be a Message or dict.")
+        if "role" not in message:
+            raise ValueError("Message dict must include a 'role' field.")
+        if "content" in message and "audio_codes_list" in message:
+            return message
+        role = message["role"]
+        if role == "user":
+            kwargs = {key: message.get(key) for key in USER_MESSAGE_FIELDS}
+            return self.build_user_message(**kwargs)
+        if role == "assistant":
+            return self.build_assistant_message(
+                audio_codes_list=message.get("audio_codes_list", []),
+                content=message.get("content", AUDIO_PLACEHOLDER),
+            )
+        raise ValueError(f"Unsupported role: {role}")
+
+    def _pad(self, input_ids_list: List[torch.Tensor]):
+        device = input_ids_list[0].device
+        lengths = torch.tensor([w.shape[0] for w in input_ids_list], device=device)
+        pad_input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids_list,
+            batch_first=True,
+            padding_value=self.model_config.audio_pad_code,
+            padding_side="left",
+        )
+        other_channel_mask = (pad_input_ids.shape[1] - lengths).unsqueeze(
+            1
+        ) > torch.arange(pad_input_ids.shape[1], device=device).unsqueeze(0)
+        pad_input_ids[..., 0][other_channel_mask] = self.model_config.pad_token_id
+        attention_mask = torch.zeros(
+            pad_input_ids.shape[0], pad_input_ids.shape[1], device=device
+        )
+        attention_mask[~other_channel_mask] = 1
+        attention_mask = attention_mask.bool()
+        return {
+            "input_ids": pad_input_ids,  # [batch_size, seqlen, n_vq]
+            "attention_mask": attention_mask,
+        }
+
+    @staticmethod
+    def _replace_audio_placeholders(
+        content: str,
+        lengths: List[int],
+        n_vq: int,
+        gen_slot_token: str,
+        delay_slot_token: str,
+        audio_start_token: str,
+        audio_end_token: str,
+    ) -> str:
+        if n_vq < 1:
+            raise ValueError(f"n_vq must be >= 1, got {n_vq}")
+
+        num_placeholders = content.count(AUDIO_PLACEHOLDER)
+        if num_placeholders != len(lengths):
+            raise ValueError(
+                f"Number of {AUDIO_PLACEHOLDER} ({num_placeholders}) "
+                f"does not match lengths ({len(lengths)})"
+            )
+
+        def build_audio_block(length: int) -> str:
+            if length < 0:
+                raise ValueError(f"length must be >= 0, got {length}")
+
+            if length == 0:
+                return f"{audio_start_token}{audio_end_token}"
+
+            step_tokens = gen_slot_token * length
+            return f"{audio_start_token}{step_tokens}{audio_end_token}"
+
+        lengths_iter = iter(lengths)
+
+        def replacer(match: re.Match) -> str:
+            length = next(lengths_iter)
+            return build_audio_block(length)
+
+        result = re.sub(re.escape(AUDIO_PLACEHOLDER), replacer, content)
+
+        return result
+
+    @staticmethod
+    def _merge_consecutive_audio_placeholders(
+        content: str,
+        audio_codes_list: List[torch.Tensor],
+    ) -> Tuple[str, List[torch.Tensor]]:
+        matches = list(re.finditer(re.escape(AUDIO_PLACEHOLDER), content))
+        if len(matches) <= 1:
+            return content, audio_codes_list
+
+        if len(matches) != len(audio_codes_list):
+            raise ValueError(
+                "Audio placeholders do not match the provided audio codes list."
+            )
+
+        new_audio_codes_list = []
+        new_parts = []
+        last_pos = 0
+        i = 0
+        while i < len(matches):
+            j = i
+            while (
+                j + 1 < len(matches)
+                and content[matches[j].end() : matches[j + 1].start()].strip() == ""
+            ):
+                j += 1
+
+            new_parts.append(content[last_pos : matches[i].start()])
+            new_parts.append(AUDIO_PLACEHOLDER)
+            last_pos = matches[j].end()
+
+            if j == i:
+                new_audio_codes_list.append(audio_codes_list[i])
+            else:
+                new_audio_codes_list.append(
+                    torch.cat(audio_codes_list[i : j + 1], dim=0)
+                )
+
+            i = j + 1
+
+        new_parts.append(content[last_pos:])
+        return "".join(new_parts), new_audio_codes_list
+
+    @staticmethod
+    def apply_delay_pattern(codes: torch.Tensor, pad_code: int) -> torch.Tensor:
+        delayed_tokens = torch.full(
+            (codes.shape[0] + codes.shape[1] - 1, codes.shape[1]),
+            pad_code,
+            device=codes.device,
+            dtype=codes.dtype,
+        )
+        for i in range(codes.shape[1]):
+            delayed_tokens[i : i + codes.shape[0], i] = codes[:, i]
+        return delayed_tokens
+
+    @staticmethod
+    def apply_de_delay_pattern(delay_codes: torch.Tensor) -> torch.Tensor:
+        tokens = torch.full(
+            (delay_codes.shape[0] - delay_codes.shape[1] + 1, delay_codes.shape[1]),
+            0,
+            device=delay_codes.device,
+            dtype=delay_codes.dtype,
+        )
+        for i in range(delay_codes.shape[1]):
+            tokens[:, i] = delay_codes[i : i + tokens.shape[0], i]
+        return tokens
+
+    def _get_unified_codes(
+        self,
+        role: str,
+        content: str,
+        audio_codes_list: List[torch.Tensor],
+        truncation: bool,
+    ) -> torch.Tensor:
+        """
+        此时的 content 已经是带上了对话格式
+        """
+        if role == "user":
+            audio_gen_slot_token = audio_delay_slot_token = self.audio_user_slot_token
+            truncation = False
+        else:
+            audio_gen_slot_token = self.audio_assistant_gen_slot_token
+            audio_delay_slot_token = self.audio_assistant_delay_slot_token
+
+        if len(audio_codes_list):
+            n_vq = audio_codes_list[0].shape[1]
+        else:
+            n_vq = self.model_config.n_vq
+
+        if len(audio_codes_list) > 1 and AUDIO_PLACEHOLDER in content:
+            content, audio_codes_list = self._merge_consecutive_audio_placeholders(
+                content, audio_codes_list
+            )
+        content = self._replace_audio_placeholders(
+            content=content,
+            lengths=[len(audio_codes) for audio_codes in audio_codes_list],
+            n_vq=n_vq,
+            gen_slot_token=audio_gen_slot_token,
+            delay_slot_token=audio_delay_slot_token,
+            audio_start_token=self.audio_start_token,
+            audio_end_token=self.audio_end_token,
+        )
+        text_codes = torch.tensor(
+            self.tokenizer.encode(content),
+            device=audio_codes_list[0].device if audio_codes_list else None,
+        )
+
+        audio_start_indices = torch.where(
+            text_codes == self.model_config.audio_start_token_id
+        )[0]
+        audio_end_indices = torch.where(
+            text_codes == self.model_config.audio_end_token_id
+        )[0]
+        if len(audio_start_indices) != len(audio_codes_list) or len(
+            audio_end_indices
+        ) != len(audio_codes_list):
+            raise ValueError(
+                "Audio placeholders do not match the provided audio codes list."
+            )
+
+        delay_audio_codes_list = []
+        if len(audio_codes_list) == 0:
+            delay_audio_codes_list = torch.full(
+                (len(text_codes), n_vq),
+                self.model_config.audio_pad_code,
+                device=text_codes.device,
+                dtype=text_codes.dtype,
+            )
+        else:
+            prefix_idx = 0
+            for audio_start_idx_t, audio_end_idx_t, audio_codes in zip(
+                audio_start_indices, audio_end_indices, audio_codes_list
+            ):
+                audio_start_idx = int(audio_start_idx_t.item())
+                audio_end_idx = int(audio_end_idx_t.item())
+                delay_audio_codes = audio_codes # not delay 
+                pad_codes = torch.full(
+                    (audio_start_idx - prefix_idx + 1, n_vq),
+                    self.model_config.audio_pad_code,
+                    device=audio_codes.device,
+                    dtype=audio_codes.dtype,
+                )
+                delay_audio_codes_list.extend([pad_codes, delay_audio_codes])
+                prefix_idx = audio_end_idx
+
+            if truncation:
+                ...
+            else:
+                last_audio_end_idx = int(audio_end_indices[-1].item())
+                pad_codes = torch.full(
+                    (len(text_codes) - last_audio_end_idx, n_vq),
+                    self.model_config.audio_pad_code,
+                    device=audio_codes_list[0].device,
+                    dtype=audio_codes_list[0].dtype,
+                )
+                delay_audio_codes_list.append(pad_codes)
+
+            delay_audio_codes_list = torch.cat(delay_audio_codes_list)
+
+        if text_codes.shape[0] != delay_audio_codes_list.shape[0]:
+            text_codes = text_codes[: delay_audio_codes_list.shape[0]]
+
+        unified_codes = torch.cat(
+            [text_codes.unsqueeze(1), delay_audio_codes_list], dim=1
+        )
+        return unified_codes
+
+    def _parse_text_codes(self, start_length, text_codes):
+        text = cast(str, self.tokenizer.decode(text_codes))
+        prefix = cast(str, self.tokenizer.decode(text_codes[:start_length]))
+        text = text[len(prefix) :]
+
+        AUDIO_PATTERN = re.compile(
+            rf"(?:{self.audio_start_token})?"
+            rf"(?:{self.audio_assistant_gen_slot_token})*"
+            rf"(?:{self.audio_assistant_delay_slot_token})*"
+            rf"{self.audio_end_token}"
+        )
+
+        def normalize_audio_segments(text: str) -> str:
+            def repl(match: re.Match) -> str:
+                seg = match.group(0)
+                # Replace with <|audio|> if gen_slot is present in the segment;
+                if self.audio_assistant_gen_slot_token in seg:
+                    return AUDIO_PLACEHOLDER
+                # Otherwise, remove it.
+                return ""
+
+            return AUDIO_PATTERN.sub(repl, text)
+
+        return normalize_audio_segments(text)
+
+    def _parse_audio_codes(self, start_length, audio_codes):
+        # De-delay back to [T', n_vq]
+        # Rows that are all pad are separators between real audio segments.
+        is_pad = (audio_codes == self.model_config.audio_pad_code).all(dim=1)
+        non_pad = ~is_pad
+        if not non_pad.any():
+            return []
+
+        idx = torch.nonzero(non_pad).squeeze(1)
+        breaks = torch.where(idx[1:] != idx[:-1] + 1)[0] + 1
+        if breaks.numel() == 0:
+            segments_idx = [idx]
+        else:
+            segments_idx = torch.split(idx, breaks.tolist())
+
+        audio_codes_list = [audio_codes[s] for s in segments_idx]
+
+        # Batch-decode all audio segments together.
+        decoded_audio_list = self.decode_audio_codes(audio_codes_list)
+
+        # Keep codec causal context by decoding the whole first segment first,
+        # then trim at waveform level according to start_length ratio.
+        if (
+            start_length > 0
+            and len(audio_codes_list) > 0
+            and len(decoded_audio_list) > 0
+        ):
+            first_codes_length = audio_codes_list[0].shape[0]
+            if first_codes_length > 0:
+                trim_ratio = max(
+                    0.0, min(float(start_length) / float(first_codes_length), 1.0)
+                )
+                first_audio = decoded_audio_list[0]
+                if trim_ratio >= 1.0:
+                    decoded_audio_list = decoded_audio_list[1:]
+                elif trim_ratio > 0.0:
+                    trim_samples = int(first_audio.shape[-1] * trim_ratio)
+                    decoded_audio_list[0] = first_audio[..., trim_samples:]
+
+        return decoded_audio_list
+
+    def decode(self, output: List[Tuple[int, torch.Tensor]]):
+        """
+        1. 这里不管怎样，都需要一个完整的 assistant generation ids;
+        2. 支持从任意位置进行截断；
+        """
+
+        genearted_messages = []
+        for start_length, generation_ids in output:
+            content = self._parse_text_codes(start_length, generation_ids[:, 0])
+            audio_codes_list = self._parse_audio_codes(
+                start_length, generation_ids[:, 1:]
+            )
+            if content == "":
+                message = None
+            else:
+                message = AssistantMessage(
+                    content=content,
+                    audio_codes_list=cast(
+                        List[Union[str, torch.Tensor]], audio_codes_list
+                    ),
+                )
+            genearted_messages.append(message)
+        return genearted_messages
+
+    @staticmethod
+    def loudness_normalize(
+        wav: torch.Tensor,
+        target_dbfs: float = -20,
+        gain_range: tuple[float, float] = (-3.0, 3.0),
+    ) -> torch.Tensor:
+        wav = wav.to(torch.float32)
+        if wav.numel() == 0:
+            return wav
+        current_dbfs = 10.0 * torch.log10(torch.mean(wav**2) + 1e-9)
+        gain = float(target_dbfs - current_dbfs)
+        gain = max(gain_range[0], min(gain, gain_range[1]))
+        factor = 10.0 ** (gain / 20.0)
+        return wav * factor
+
+    def _get_audio_tokenizer_device(self) -> torch.device:
+        """Best-effort device inference for `self.audio_tokenizer`.
+
+        Notes:
+        - Old TAC wrapper exposed `.device`, but standard `torch.nn.Module` does not.
+        - New MossAudioTokenizerModel is a `PreTrainedModel`; parameters define its device.
+        """
+
+        audio_tokenizer = getattr(self, "audio_tokenizer", None)
+        if audio_tokenizer is None:
+            logger.warning(
+                "audio_tokenizer is not set on processor. Using CPU as default."
+            )
+            return torch.device("cpu")
+
+        device_attr = getattr(audio_tokenizer, "device", None)
+        if isinstance(device_attr, torch.device):
+            return device_attr
+
+        try:
+            return next(audio_tokenizer.parameters()).device
+        except StopIteration:
+            # No parameters (shouldn't happen for real models); default to CPU.
+            logger.warning(
+                "No parameters found on audio_tokenizer. Using CPU as default."
+            )
+            return torch.device("cpu")
+
+    def encode_audios_from_wav(
+        self,
+        wav_list: List[torch.Tensor],
+        sampling_rate: int,
+        n_vq: Optional[int] = None,
+    ):
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set on processor.")
+        audio_tokenizer = self.audio_tokenizer
+
+        if isinstance(wav_list, torch.Tensor):
+            wav_list = [wav_list]
+        wav_list_ = []
+        resample = False
+        if sampling_rate != self.model_config.sampling_rate:
+            resample = True
+        device = self._get_audio_tokenizer_device()
+        for wav in wav_list:
+            if wav.shape[0] > 1:
+                wav = torch.mean(wav, dim=0, keepdim=True)
+            if resample:
+                wav = torchaudio.functional.resample(
+                    waveform=wav,
+                    orig_freq=sampling_rate,
+                    new_freq=self.model_config.sampling_rate,
+                )
+            wav = wav.to(device)
+            wav_list_.append(self.loudness_normalize(wav.squeeze(0)))
+
+        # New MossAudioTokenizerModel API: prefer batch_encode(list[wav])
+        if hasattr(audio_tokenizer, "batch_encode"):
+            enc = audio_tokenizer.batch_encode(wav_list_, num_quantizers=n_vq)
+            audio_codes = enc.audio_codes  # (NQ, B, T)
+            audio_codes_lengths = enc.audio_codes_lengths  # (B,)
+        else:
+            # Fallback: use encode() with explicit padding.
+            max_len = max(int(wav.shape[-1]) for wav in wav_list_)
+            input_values = torch.zeros(
+                len(wav_list_), 1, max_len, device=device, dtype=torch.float32
+            )
+            padding_mask = torch.zeros(
+                len(wav_list_), max_len, device=device, dtype=torch.bool
+            )
+            for i, wav in enumerate(wav_list_):
+                this_len = int(wav.shape[-1])
+                input_values[i, 0, :this_len] = wav
+                padding_mask[i, :this_len] = True
+            enc = audio_tokenizer.encode(
+                input_values,
+                padding_mask=padding_mask,
+                num_quantizers=n_vq,
+                return_dict=True,
+            )
+            audio_codes = enc.audio_codes
+            audio_codes_lengths = enc.audio_codes_lengths
+
+        if audio_codes is None or audio_codes_lengths is None:
+            raise RuntimeError(
+                "audio_tokenizer.encode() returned empty outputs (audio_codes/audio_codes_lengths)."
+            )
+
+        # Keep processor's historical contract: list[Tensor] with shape (T, NQ)
+        # and on CPU (so downstream text/audio packing remains device-agnostic).
+        codes_list: List[torch.Tensor] = []
+        for i in range(int(audio_codes.shape[1])):
+            length_i = int(audio_codes_lengths[i].item())
+            codes_i = (
+                audio_codes[:, i, :length_i]
+                .transpose(0, 1)
+                .contiguous()
+                .to(torch.long)
+                .cpu()
+            )
+            codes_list.append(codes_i)
+        return codes_list
+
+    def encode_audios_from_path(
+        self, wav_path_list: Union[str, List[str]], n_vq: Optional[int] = None
+    ):
+        if isinstance(wav_path_list, str):
+            wav_path_list = [wav_path_list]
+
+        if len(wav_path_list) == 0:
+            raise ValueError("Empty wav_path_list")
+
+        # Load + (if needed) resample each wav independently, so callers can
+        # pass a heterogeneous batch of files while still benefiting from
+        # audio_tokenizer.batch_encode.
+        target_sr = int(self.model_config.sampling_rate)
+        wav_list: List[torch.Tensor] = []
+        for wav_path in wav_path_list:
+            import soundfile as _sf  # bypass torchcodec (Windows FFmpeg-shared issue)
+            _d, sr = _sf.read(str(wav_path), always_2d=True, dtype="float32")
+            wav = torch.from_numpy(_d.T).contiguous()  # [channels, frames]
+            if int(sr) != target_sr:
+                wav = torchaudio.functional.resample(
+                    waveform=wav,
+                    orig_freq=int(sr),
+                    new_freq=target_sr,
+                )
+            wav_list.append(wav)
+
+        return self.encode_audios_from_wav(wav_list, target_sr, n_vq)
+
+    def decode_audio_codes(
+        self, audio_tokens_list: Union[torch.Tensor, List[torch.Tensor]]
+    ):
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set on processor.")
+        audio_tokenizer = self.audio_tokenizer
+
+        if isinstance(audio_tokens_list, torch.Tensor):
+            audio_tokens_list = [audio_tokens_list]
+        if len(audio_tokens_list) == 0:
+            return []
+
+        device = self._get_audio_tokenizer_device()
+
+        # Processor uses (T, NQ); MossAudioTokenizer expects (NQ, T) (or (NQ, B, T)).
+        codes_list = [
+            codes.transpose(0, 1).contiguous().to(device=device, dtype=torch.long)
+            for codes in audio_tokens_list
+        ]
+    
+        # Fallback: pad to (NQ, B, T) + mask, then decode.
+        nq = int(codes_list[0].shape[0])
+        max_t = max(int(c.shape[1]) for c in codes_list)
+        audio_codes = torch.zeros(
+            nq, len(codes_list), max_t, device=device, dtype=torch.long
+        )
+        padding_mask = torch.zeros(
+            len(codes_list), max_t, device=device, dtype=torch.bool
+        )
+        for i, c in enumerate(codes_list):
+            t = int(c.shape[1])
+            audio_codes[:, i, :t] = c
+            padding_mask[i, :t] = True
+        dec = audio_tokenizer.decode(
+            audio_codes, padding_mask=padding_mask, return_dict=True, chunk_duration=8
+        )
+        audio = dec.audio
+        audio_lengths = dec.audio_lengths
+
+        if audio is None or audio_lengths is None:
+            raise RuntimeError(
+                "audio_tokenizer.decode() returned empty outputs (audio/audio_lengths)."
+            )
+
+        # Return historical contract: list of 1D waveforms (T,)
+        wav_list: List[torch.Tensor] = []
+        for i in range(int(audio.shape[0])):
+            length_i = int(audio_lengths[i].item())
+            wav = audio[i, 0, :length_i].contiguous().to(torch.float32).cpu()
+            wav_list.append(wav)
+        return wav_list

processor_config.json

@@ -0,0 +1,7 @@
+{
+  "processor_class": "MossTTSDelayProcessor",
+  "auto_map": {
+    "AutoProcessor": "processing_moss_tts.MossTTSDelayProcessor"
+  },
+  "audio_tokenizer_name_or_path": "OpenMOSS-Team/MOSS-Audio-Tokenizer"
+}

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|audio_start|>",
+    "<|audio_end|>",
+    "<|audio_user_slot|>",
+    "<|image_pad|>",
+    "<|audio_assistant_gen_slot|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,240 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|audio_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|audio_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|audio_user_slot|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|audio_assistant_gen_slot|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|audio_assistant_delay_slot|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|audio_start|>",
+    "<|audio_end|>",
+    "<|audio_user_slot|>",
+    "<|image_pad|>",
+    "<|audio_assistant_gen_slot|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "processor_class": "AsteroidProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}