Upload folder using huggingface_hub

.gitattributes

@@ -28,7 +28,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
 *.wasm filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,279 @@
 ---
 license: apache-2.0
+tags:
+- text-to-speech
+language:
+- zh
+- yue
+- en
+- ar
+- cs
+- da
+- de
+- nl
+- es
+- fr
+- fi
+- el
+- he
+- hi
+- hu
+- ja
+- it
+- ko
+- mk
+- ms
+- ru
+- fa
+- pl
+- pt
+- sv
+- ro
+- sw
+- tl
+- th
+- tr
+- vi
 ---
+# MOSS-TTS Family
+
+
+<br>
+
+<p align="center">
+  &nbsp;&nbsp;&nbsp;&nbsp;
+  <img src="https://speech-demo.oss-cn-shanghai.aliyuncs.com/moss_tts_demo/tts_readme_imgaes_demo/openmoss_x_mosi" height="50" align="middle" />
+</p>
+
+
+
+<div align="center">
+  <a href="https://github.com/OpenMOSS/MOSS-TTS/tree/main"><img src="https://img.shields.io/badge/Project%20Page-GitHub-blue"></a>
+  <a href="https://modelscope.cn/collections/OpenMOSS-Team/MOSS-TTS"><img src="https://img.shields.io/badge/ModelScope-Models-lightgrey?logo=modelscope&amp"></a>
+  <a href="https://mosi.cn/#models"><img src="https://img.shields.io/badge/Blog-View-blue?logo=internet-explorer&amp"></a>
+  <a href="https://arxiv.org/abs/2603.18090"><img src="https://img.shields.io/badge/Arxiv-2603.18090-red?logo=Arxiv&amp"></a>
+
+  <a href="https://studio.mosi.cn"><img src="https://img.shields.io/badge/AIStudio-Try-green?logo=internet-explorer&amp"></a>
+  <a href="https://studio.mosi.cn/docs/moss-tts"><img src="https://img.shields.io/badge/API-Docs-00A3FF?logo=fastapi&amp"></a>
+  <a href="https://x.com/Open_MOSS"><img src="https://img.shields.io/badge/Twitter-Follow-black?logo=x&amp"></a>
+  <a href="https://discord.gg/fvm5TaWjU3"><img src="https://img.shields.io/badge/Discord-Join-5865F2?logo=discord&amp"></a>
+</div>
+
+
+## Overview
+MOSS‑TTS Family is an open‑source **speech and sound generation model family** from [MOSI.AI](https://mosi.cn/#hero) and the [OpenMOSS team](https://www.open-moss.com/). It is designed for **high‑fidelity**, **high‑expressiveness**, and **complex real‑world scenarios**, covering stable long‑form speech, multi‑speaker dialogue, voice/character design, environmental sound effects, and real‑time streaming TTS.
+
+
+## Introduction
+
+<p align="center">
+  <img src="https://speech-demo.oss-cn-shanghai.aliyuncs.com/moss_tts_demo/tts_readme_imgaes_demo/moss_tts_family_arch.jpeg" width="85%" />
+</p>
+
+
+When a single piece of audio needs to **sound like a real person**, **pronounce every word accurately**, **switch speaking styles across content**, **remain stable over tens of minutes**, and **support dialogue, role‑play, and real‑time interaction**, a single TTS model is often not enough. The **MOSS‑TTS Family** breaks the workflow into five production‑ready models that can be used independently or composed into a complete pipeline.
+
+- **MOSS‑TTS**: The flagship production model featuring high fidelity and optimal zero-shot voice cloning. It supports **long-speech generation**, **fine-grained control over Pinyin, phonemes, and duration**, as well as **multilingual/code-switched synthesis**.
+- **MOSS‑TTSD**: A spoken dialogue generation model for expressive, multi-speaker, and ultra-long dialogues. The new **v1.0 version** achieves **industry-leading performance on objective metrics** and **outperformed top closed-source models like Doubao and Gemini 2.5-pro** in subjective evaluations. You can visit the [MOSS-TTSD repository](https://github.com/OpenMOSS/MOSS-TTSD) for details.
+- **MOSS‑VoiceGenerator**: An open-source voice design model capable of generating diverse voices and styles directly from text prompts, **without any reference speech**. It unifies voice design, style control, and synthesis, functioning independently or as a design layer for downstream TTS. Its performance **surpasses other top-tier voice design models in arena ratings**.
+- **MOSS‑TTS‑Realtime**: A multi-turn context-aware model for real-time voice agents. It uses incremental synthesis to ensure natural and coherent replies, making it **ideal for building low-latency voice agents when paired with text models**. The TTFB (Time To First Byte) of MOSS-TTS-Realtime reaches 180 ms, and the $T_{\text{LLM-first-sentence}} + T_{\text{MOSS-TTS-Realtime-TTFB}}$ is 377 ms.
+- **MOSS‑SoundEffect**: A content creation model specialized in **sound effect generation** with wide category coverage and controllable duration. It generates audio for natural environments, urban scenes, biological sounds, human actions, and musical fragments, suitable for film, games, and interactive experiences.
+
+
+## Supported Languages
+
+MOSS-TTS-v1.5 currently supports **31 languages**. It keeps the 20 languages supported by [MOSS-TTS 1.0](https://huggingface.co/OpenMOSS-Team/MOSS-TTS) and extends multilingual continued training to additional languages including Cantonese, Dutch, Finnish, Hindi, Macedonian, Malay, Romanian, Swahili, Tagalog, Thai, and Vietnamese.
+
+| Language | Code | Flag | Language | Code | Flag | Language | Code | Flag |
+|---|---|---|---|---|---|---|---|---|
+| Chinese | zh | 🇨🇳 | Cantonese | yue | 🇭🇰 | English | en | 🇺🇸 |
+| Arabic | ar | 🇸🇦 | Czech | cs | 🇨🇿 | Danish | da | 🇩🇰 |
+| Dutch | nl | 🇳🇱 | Finnish | fi | 🇫🇮 | French | fr | 🇫🇷 |
+| German | de | 🇩🇪 | Greek | el | 🇬🇷 | Hebrew | he | 🇮🇱 |
+| Hindi | hi | 🇮🇳 | Hungarian | hu | 🇭🇺 | Italian | it | 🇮🇹 |
+| Japanese | ja | 🇯🇵 | Korean | ko | 🇰🇷 | Macedonian | mk | 🇲🇰 |
+| Malay | ms | 🇲🇾 | Persian (Farsi) | fa | 🇮🇷 | Polish | pl | 🇵🇱 |
+| Portuguese | pt | 🇵🇹 | Romanian | ro | 🇷🇴 | Russian | ru | 🇷🇺 |
+| Spanish | es | 🇪🇸 | Swahili | sw | 🇹🇿 | Swedish | sv | 🇸🇪 |
+| Tagalog | tl | 🇵🇭 | Thai | th | 🇹🇭 | Turkish | tr | 🇹🇷 |
+| Vietnamese | vi | 🇻🇳 | | | | | | |
+
+# MOSS-TTS-v1.5
+
+**MOSS-TTS-v1.5** is continued from [MOSS-TTS 1.0](https://huggingface.co/OpenMOSS-Team/MOSS-TTS). It preserves the main 1.0 capabilities, including zero-shot voice cloning, long-form speech generation, token-level duration control, Pinyin/IPA pronunciation control, multilingual synthesis, and code-switching. For the full 1.0 feature walkthrough, input schema, decoding hyperparameters, and evaluation tables, please refer to the [MOSS-TTS 1.0 README](https://huggingface.co/OpenMOSS-Team/MOSS-TTS).
+
+Compared with MOSS-TTS 1.0, v1.5 focuses on the following improvements:
+
+- **Stronger multilingual synthesis with language tags**: when the `language` field is omitted, v1.5 may improve some languages and regress slightly on others compared with 1.0. When the language is specified, v1.5 is stronger than 1.0 on almost all supported languages. Set the tag when building the user message, for example `processor.build_user_message(text=text_fr, language="French")`.
+- **More stable voice cloning**: v1.5 improves speaker similarity and reduces cloning variance, making repeated generations more consistent.
+- **Better long-reference, short-text cloning**: v1.5 handles scenarios where the reference audio is much longer than the target text more reliably than 1.0.
+- **More stable punctuation-following prosody**: v1.5 follows punctuation-driven pauses more closely, especially in long sentences.
+- **Explicit pause control**: v1.5 supports inline pause markers such as `"[pause 3.2s]"`. For example, `我今天学习了一首中国的古诗，它的名字是[pause 3.2s]静夜思！` inserts an explicit 3.2s pause before `静夜思`.
+
+## Quick Start
+
+### Environment Setup
+
+We recommend a clean, isolated Python environment with **Transformers 5.0.0** to avoid dependency conflicts.
+
+```bash
+conda create -n moss-tts python=3.12 -y
+conda activate moss-tts
+```
+
+Install all required dependencies:
+
+```bash
+git clone https://github.com/OpenMOSS/MOSS-TTS.git
+cd MOSS-TTS
+pip install --extra-index-url https://download.pytorch.org/whl/cu128 -e .
+```
+
+#### (Optional) Install FlashAttention 2
+
+For better speed and lower GPU memory usage, you can install FlashAttention 2 if your hardware supports it.
+
+```bash
+pip install --extra-index-url https://download.pytorch.org/whl/cu128 -e ".[flash-attn]"
+```
+
+If your machine has limited RAM and many CPU cores, you can cap build parallelism:
+
+```bash
+MAX_JOBS=4 pip install --extra-index-url https://download.pytorch.org/whl/cu128 -e ".[flash-attn]"
+```
+
+Notes:
+- Dependencies are managed in `pyproject.toml`, which currently pins `torch==2.9.1+cu128` and `torchaudio==2.9.1+cu128`.
+- If FlashAttention 2 fails to build on your machine, you can skip it and use the default attention backend.
+- FlashAttention 2 is only available on supported GPUs and is typically used with `torch.float16` or `torch.bfloat16`.
+
+
+### Basic Usage
+
+> Tip: MOSS-TTS-v1.5 uses the same generation API as the 1.0 **MossTTSDelay-8B** checkpoint. For multilingual inputs, set `language` whenever the language is known.
+
+MOSS-TTS provides a convenient `generate` interface for rapid usage. The examples below cover:
+1. Direct generation (Chinese / English / multilingual text with language tags / Pinyin / IPA)
+2. Voice cloning
+3. Duration control
+4. Explicit pause control with `[pause X.Ys]`
+
+```python
+from pathlib import Path
+import importlib.util
+import torch
+import torchaudio
+from transformers import AutoModel, AutoProcessor
+# Disable the broken cuDNN SDPA backend
+torch.backends.cuda.enable_cudnn_sdp(False)
+# Keep these enabled as fallbacks
+torch.backends.cuda.enable_flash_sdp(True)
+torch.backends.cuda.enable_mem_efficient_sdp(True)
+torch.backends.cuda.enable_math_sdp(True)
+
+
+pretrained_model_name_or_path = "OpenMOSS-Team/MOSS-TTS-v1.5"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+def resolve_attn_implementation() -> str:
+    # Prefer FlashAttention 2 when package + device conditions are met.
+    if (
+        device == "cuda"
+        and importlib.util.find_spec("flash_attn") is not None
+        and dtype in {torch.float16, torch.bfloat16}
+    ):
+        major, _ = torch.cuda.get_device_capability()
+        if major >= 8:
+            return "flash_attention_2"
+
+    # CUDA fallback: use PyTorch SDPA kernels.
+    if device == "cuda":
+        return "sdpa"
+
+    # CPU fallback.
+    return "eager"
+
+
+attn_implementation = resolve_attn_implementation()
+print(f"[INFO] Using attn_implementation={attn_implementation}")
+
+processor = AutoProcessor.from_pretrained(
+    pretrained_model_name_or_path,
+    trust_remote_code=True,
+)
+processor.audio_tokenizer = processor.audio_tokenizer.to(device)
+
+text_1 = "亲爱的你，\n你好呀。\n\n今天，我想用最认真、最温柔的声音，对你说一些重要的话。\n这些话，像一颗小小的星星，希望能在你的心里慢慢发光。\n\n首先，我想祝你——\n每天都能平平安安、快快乐乐。\n\n希望你早上醒来的时候，\n窗外有光，屋子里很安静，\n你的心是轻轻的，没有着急，也没有害怕。\n\n希望你吃饭的时候胃口很好，\n走路的时候脚步稳稳，\n晚上睡觉的时候，能做一个又一个甜甜的梦。\n\n我希望你能一直保持好奇心。\n对世界充满问题，\n对天空、星星、花草、书本和故事感兴趣。\n当你问“为什么”的时候，\n希望总有人愿意认真地听你说话。\n\n我也希望你学会温柔。\n温柔地对待朋友，\n温柔地对待小动物，\n也温柔地对待自己。\n\n如果有一天你犯了错，\n请不要太快责怪自己，\n因为每一个认真成长的人，\n都会在路上慢慢学会更好的方法。\n\n愿你拥有勇气。\n当你站在陌生的地方时，\n当你第一次举手发言时，\n当你遇到困难、感到害怕的时候，\n希望你能轻轻地告诉自己：\n“我可以试一试。”\n\n就算没有一次成功，也没有关系。\n失败不是坏事，\n它只是告诉你，你正在努力。\n\n我希望你学会分享快乐。\n把开心的事情告诉别人，\n把笑声送给身边的人，\n因为快乐被分享的时候，\n会变得更大、更亮。\n\n如果有一天你感到难过，\n我希望你知道——\n难过并不丢脸，\n哭泣也不是软弱。\n\n愿你能找到一个安全的地方，\n慢慢把心里的话说出来，\n然后再一次抬起头，看见希望。\n\n我还希望你能拥有梦想。\n这个梦想也许很大，\n也许很小，\n也许现在还说不清楚。\n\n没关系。\n梦想会和你一起长大，\n在时间里慢慢变得清楚。\n\n最后，我想送你一个最最重要的祝福：\n\n愿你被世界温柔对待，\n也愿你成为一个温柔的人。\n\n愿你的每一天，\n都值得被记住，\n都值得被珍惜。\n\n亲爱的你，\n请记住，\n你是独一无二的，\n你已经很棒了，\n而你的未来，\n一定会慢慢变得闪闪发光。\n\n祝你健康、勇敢、幸福，\n祝你永远带着笑容向前走。"
+text_2 = "We stand on the threshold of the AI era.\nArtificial intelligence is no longer just a concept in laboratories, but is entering every industry, every creative endeavor, and every decision. It has learned to see, hear, speak, and think, and is beginning to become an extension of human capabilities. AI is not about replacing humans, but about amplifying human creativity, making knowledge more equitable, more efficient, and allowing imagination to reach further. A new era, jointly shaped by humans and intelligent systems, has arrived."
+text_3 = "nin2 hao3，qing3 wen4 nin2 lai2 zi4 na3 zuo4 cheng2 shi4？"
+text_4 = "nin2 hao3，qing4 wen3 nin2 lai2 zi4 na4 zuo3 cheng4 shi3？"
+text_5 = "您好，请问您来自哪 zuo4 cheng2 shi4？"
+text_6 = "/həloʊ, meɪ aɪ æsk wɪtʃ sɪti juː ɑːr frʌm?/"
+text_7 = "Bonjour, je voudrais essayer une voix française naturelle et stable."
+text_8 = "我今天学习了一首中国的古诗，它的名字是[pause 3.2s]静夜思！"
+
+# Use audio from ./assets/audio to avoid downloading from the cloud.
+ref_audio_1 = "https://speech-demo.oss-cn-shanghai.aliyuncs.com/moss_tts_demo/tts_readme_demo/reference_zh.wav"
+ref_audio_2 = "https://speech-demo.oss-cn-shanghai.aliyuncs.com/moss_tts_demo/tts_readme_demo/reference_en.m4a"
+
+conversations = [
+    # Direct TTS (no reference). Language tags are recommended in v1.5.
+    [processor.build_user_message(text=text_1)],
+    [processor.build_user_message(text=text_2)],
+    # Direct TTS (no reference). For languages ​​other than Chinese and English, it is recommended to use language tags.
+    [processor.build_user_message(text=text_7, language="French")],
+    # Pinyin or IPA input
+    [processor.build_user_message(text=text_3)],
+    [processor.build_user_message(text=text_4)],
+    [processor.build_user_message(text=text_5)],
+    [processor.build_user_message(text=text_6)],
+    # Explicit pause control. Use [pause X.Ys], such as [pause 3.2s].
+    [processor.build_user_message(text=text_8)],
+    # Voice cloning (with reference)
+    [processor.build_user_message(text=text_1, reference=[ref_audio_1])],
+    [processor.build_user_message(text=text_2, reference=[ref_audio_2])],
+    # Duration control
+    [processor.build_user_message(text=text_2, tokens=325)],
+    [processor.build_user_message(text=text_2, tokens=600)],
+]
+
+model = AutoModel.from_pretrained(
+    pretrained_model_name_or_path,
+    trust_remote_code=True,
+    # If FlashAttention 2 is installed, you can set attn_implementation="flash_attention_2"
+    attn_implementation=attn_implementation,
+    torch_dtype=dtype,
+).to(device)
+model.eval()
+
+batch_size = 1
+
+save_dir = Path("inference_root")
+save_dir.mkdir(exist_ok=True, parents=True)
+sample_idx = 0
+with torch.no_grad():
+    for start in range(0, len(conversations), batch_size):
+        batch_conversations = conversations[start : start + batch_size]
+        batch = processor(batch_conversations, mode="generation")
+        input_ids = batch["input_ids"].to(device)
+        attention_mask = batch["attention_mask"].to(device)
+
+        outputs = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=4096,
+        )
+
+        for message in processor.decode(outputs):
+            audio = message.audio_codes_list[0]
+            out_path = save_dir / f"sample{sample_idx}.wav"
+            sample_idx += 1
+            torchaudio.save(out_path, audio.unsqueeze(0), processor.model_config.sampling_rate)
+
+```
+
+## More Usage
+
+MOSS-TTS-v1.5 is API-compatible with MOSS-TTS 1.0. For continuation with prefix audio, detailed `UserMessage` and `AssistantMessage` fields, generation hyperparameters, Pinyin/IPA preprocessing examples, and evaluation results, see the [MOSS-TTS 1.0 README](https://huggingface.co/OpenMOSS-Team/MOSS-TTS).

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,89 @@
+{
+  "model_type": "moss_tts_delay",
+  "architectures": [
+    "MossTTSDelayModel"
+  ],
+  "auto_map": {
+      "AutoConfig": "configuration_moss_tts.MossTTSDelayConfig",
+      "AutoModel": "modeling_moss_tts.MossTTSDelayModel"
+  },
+  "dtype": "bfloat16",
+  "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,
+    "pad_token_id": 151643,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "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",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 40960,
+    "max_window_layers": 36,
+    "model_type": "qwen3",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 1000000,
+    "sliding_window": null,
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 155648
+  },
+  "n_vq": 32,
+  "audio_vocab_size": 1024,
+  "audio_user_slot_token_id": 151654,
+  "audio_assistant_gen_slot_token_id": 151656,
+  "audio_assistant_delay_slot_token_id": 151662,
+  "audio_start_token_id": 151652,
+  "audio_end_token_id": 151653,
+  "audio_pad_code": 1024,
+  "sampling_rate": 24000,
+  "transformers_version": "4.57.1"
+}

configuration_moss_tts.py

@@ -0,0 +1,114 @@
+# 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,
+        **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
+
+        
+        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

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-00001-of-00004.safetensors

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

model-00002-of-00004.safetensors

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

model-00003-of-00004.safetensors

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

model-00004-of-00004.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,471 @@
+{
+  "metadata": {
+    "total_parameters": 8489841664,
+    "total_size": 16979683328
+  },
+  "weight_map": {
+    "emb_ext.0.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.1.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.10.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.11.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.12.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.13.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.14.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.15.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.16.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.17.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.18.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.19.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.2.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.20.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.21.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.22.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.23.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.24.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.25.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.26.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.27.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.28.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.29.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.3.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.30.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.31.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.4.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.5.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.6.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.7.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.8.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.9.weight": "model-00004-of-00004.safetensors",
+    "language_model.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.20.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.30.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.35.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.self_attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.35.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.self_attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.35.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.norm.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.0.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.1.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.10.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.11.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.12.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.13.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.14.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.15.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.16.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.17.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.18.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.19.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.2.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.20.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.21.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.22.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.23.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.24.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.25.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.26.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.27.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.28.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.29.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.3.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.30.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.31.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.32.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.4.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.5.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.6.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.7.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.8.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.9.weight": "model-00004-of-00004.safetensors"
+  }
+}

modeling_moss_tts.py

@@ -0,0 +1,525 @@
+# 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.
+""" Modeling classes for MossTTSDelay. """
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.cache_utils import Cache
+from transformers.models.qwen3 import Qwen3Model
+from transformers import initialization as init
+
+from .configuration_moss_tts import MossTTSDelayConfig
+from .inference_utils import sample_token, find_last_equal_C
+
+try:
+    from .processing_moss_tts import UserMessage, AssistantMessage, MossTTSDelayProcessor
+except Exception:
+    UserMessage = None
+    AssistantMessage = None
+    MossTTSDelayProcessor = None
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "MossTTSDelayConfig"
+
+
+@dataclass
+class MossTTSDelayOutputWithPast(ModelOutput):
+    """
+    Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Weighted sum of channel losses.
+        all_sum_losses (`torch.FloatTensor` of shape `(batch_size, n_vq + 1)`, *optional*):
+            Sum of losses for each sample and each channel before averaging.
+        all_token_nums (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Number of non-masked tokens per sample.
+        sample_losses (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
+            Loss per sample.
+        channel_losses (`torch.FloatTensor` of shape `(n_vq + 1,)`, *optional*):
+            Loss per channel (text head + vq heads).
+        logits (`List[torch.FloatTensor]`, *optional*):
+            List of prediction scores from each head.
+        past_key_values (`Cache`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks).
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed):
+            Tuple of torch.FloatTensor (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer).
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed):
+            Tuple of torch.FloatTensor (one for each layer) of the attention weights.
+    """
+    loss: Optional[torch.FloatTensor] = None
+    all_sum_losses: Optional[torch.FloatTensor] = None
+    all_token_nums: Optional[torch.LongTensor] = None
+    sample_losses: Optional[torch.FloatTensor] = None
+    channel_losses: Optional[torch.FloatTensor] = None
+    logits: Optional[List[torch.FloatTensor]] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+
+
+class MossTTSDelayPreTrainedModel(PreTrainedModel):
+    config_class = MossTTSDelayConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen3DecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    def _init_weights(self, module):
+        """
+        Transformers 5.0+ safe init:
+        - MUST use transformers.initialization helpers
+        - MUST respect param._is_hf_initialized to avoid overwriting ckpt-loaded params
+        """
+        # Let HF handle its standard modules first (LayerNorm, Linear, Embedding, etc.)
+        super()._init_weights(module)
+
+        # Pick a std consistent with HF conventions
+        # Prefer model/text config initializer_range if present.
+        std = None
+        if hasattr(self.config, "initializer_range"):
+            std = self.config.initializer_range
+        elif hasattr(self.config, "language_config") and hasattr(self.config.language_config, "initializer_range"):
+            std = self.config.language_config.initializer_range
+        else:
+            std = 0.02
+
+        # Initialize extra audio embeddings
+        if isinstance(module, nn.Embedding):
+            # Only touch our extra embeddings (avoid double touching LM's embeddings if not desired)
+            # If you prefer, you can skip this check and rely on super()._init_weights for all embeddings.
+            if getattr(module, "num_embeddings", None) == self.config.audio_vocab_size + 1:
+                init.normal_(module.weight, mean=0.0, std=std)
+                # If you later set padding_idx, you must explicitly zero it (and respect _is_hf_initialized!)
+                # init.zeros_ will internally check param flags, but slicing needs manual care.
+
+        # Initialize multi-head projections you added
+        if isinstance(module, nn.Linear):
+            # For your lm_heads, super()._init_weights already covers typical Linear.
+            # This block is only needed if you have custom Linear variants later.
+            pass
+
+
+
+MOSSTTS_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`MossTTSDelayConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The MossTTSDelay Model architecture tailored for Text-to-Speech generation with multi-head VQ prediction.",
+    MOSSTTS_START_DOCSTRING,
+)
+class MossTTSDelayModel(MossTTSDelayPreTrainedModel):
+    UserMessage = UserMessage
+    AssistantMessage = AssistantMessage
+    Processor = MossTTSDelayProcessor
+
+    def __init__(self, config: MossTTSDelayConfig):
+        super().__init__(config)
+        self.config = config
+
+        config.language_config.torch_dtype = config.torch_dtype
+        
+        self.language_model = Qwen3Model(config.language_config)
+
+        # Audio VQ Embeddings (Extra channels)
+        # Note: input_ids[..., 0] uses Qwen's embedding. 
+        # input_ids[..., 1:] use these extensions.
+        self.emb_ext = nn.ModuleList()
+        for vq_idx in range(self.config.n_vq):
+            # Add +1 for potential padding/special tokens logic if strictly required by upstream data prep
+            self.emb_ext.append(
+                nn.Embedding(self.config.audio_vocab_size + 1, config.language_config.hidden_size, padding_idx=None)
+            )
+
+        # Multi-Head Prediction Layers
+        # Head 0: Main language head
+        # Head 1..N: Audio VQ heads
+        self.lm_heads = nn.ModuleList([
+            nn.Linear(config.language_config.hidden_size, config.language_config.vocab_size, bias=False)
+        ])
+        for vq_idx in range(self.config.n_vq):
+            self.lm_heads.append(
+                nn.Linear(config.language_config.hidden_size, self.config.audio_vocab_size + 1, bias=False)
+            )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """
+        Computes the combined embeddings from text and multiple audio VQ channels.
+        
+        Args:
+            input_ids: Shape (Batch, Seq_Len, 1 + n_vq)
+        """
+        # Base Text/Content Embedding
+        # input_ids[..., 0] is standard text or semantic tokens
+        inputs_embeds = self.language_model.get_input_embeddings()(input_ids[..., 0])
+
+        # Add VQ Embeddings
+        for i, embed_layer in enumerate(self.emb_ext):
+            # i corresponds to channel i+1 in input_ids
+            # We assume the data pipeline ensures indices are within range
+            inputs_embeds = inputs_embeds + embed_layer(input_ids[..., i + 1])
+        
+        return inputs_embeds
+
+    def set_input_embeddings(self, value):
+        self.language_model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        # Returning a list of heads might break some HF utilities expecting a single head.
+        # However, for custom models, this is acceptable.
+        return self.lm_heads
+
+    @add_start_docstrings_to_model_forward(MOSSTTS_START_DOCSTRING)
+    @replace_return_docstrings(output_type=MossTTSDelayOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    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[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        hidden_out_layers: Optional[List[int]] = None,
+        channelwise_loss_weight: Optional[List[float]] = None,
+        **kwargs,
+    ) -> Union[Tuple, MossTTSDelayOutputWithPast]:
+        r"""
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length, 1 + n_vq)`):
+                Indices of input sequence tokens in the vocabulary. 
+                Dimension 2 contains: [Text/Semantics, VQ_0, VQ_1, ..., VQ_N].
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length, 1 + n_vq)`, *optional*):
+                Labels for computing the masked language modeling loss.
+            channelwise_loss_weight (`List[float]`, *optional*):
+                Manual weights for summing losses across different heads (Text vs Audio channels).
+
+        Returns:
+        """
+
+        if len(input_ids.shape) != 3 or input_ids.shape[-1] != self.config.n_vq + 1:
+            raise ValueError("`Input_ids`'s shape should be exactly (batch_size, sequence_length, 1 + n_vq).")
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
+        # 1. Prepare Embeddings
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings(input_ids)
+
+        # 2. Backbone Forward
+        # Qwen3Model outputs standard CausalLMOutputWithPast or similar
+        outputs = self.language_model(
+            input_ids=None, # Passed via inputs_embeds
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=True, # Always need hidden states for multi-head projection
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        # 3. Handle specific layer outputs if requested (Delay Pattern often requires features from specific layers)
+        last_hidden_state = outputs.last_hidden_state
+        if hidden_out_layers is None:
+            # Default to using the last layer for all heads
+            # In some architectures (like MusicGen), different codebooks come from different transformer layers.
+            # Here we default to the final layer as per original code behavior [-1] * (n + 1).
+            hidden_states_for_heads = [last_hidden_state] * (len(self.lm_heads))
+        else:
+            # If hidden_out_layers is provided (e.g. [-1, -2, -3...]), fetch them from all_hidden_states
+            # Note: outputs.hidden_states includes embedding output at index 0 usually.
+            all_hs = outputs.hidden_states
+            hidden_states_for_heads = [all_hs[idx] for idx in hidden_out_layers]
+
+        # 4. Project to Logits (Multi-Head)
+        layer_logits = []
+        for i, (hs, head) in enumerate(zip(hidden_states_for_heads, self.lm_heads)):
+            logits = head(hs)
+            # Original code logic: Mask the last token index for audio heads (indices > 0)
+            # This implies the vocab size is (N+1) but the model shouldn't predict the (N+1)-th token 
+            # (perhaps reserved for padding in the input but invalid for prediction).
+            if i > 0:
+                logits[..., -1] = float("-inf")
+            layer_logits.append(logits)
+
+        # 5. Loss Calculation
+        loss = None
+        all_sum_losses = None
+        all_token_nums = None
+        sample_losses = None
+        channel_losses = None
+
+        if labels is not None:
+            # Ensure labels match input shape rank (B, S, C)
+            if labels.dim() != 3:
+                raise ValueError(f"Labels must have rank 3 (B, S, C), got {labels.shape}")
+
+            batch_size = labels.size(0)
+            n_heads = len(layer_logits)
+            
+            # Container for per-sample, per-channel losses
+            # Shape: [Batch, n_heads]
+            all_sum_losses_list = [] 
+            
+            # Count valid tokens (not -100) per sample. 
+            # Note: Assuming mask is consistent across channels or we take sum over dim 1 (seq)
+            # Usually strict masking means checking one channel or all. 
+            # Original code: torch.sum(labels != -100, dim=1) -> [B, C]
+            all_token_nums = torch.sum(labels != -100, dim=1)
+
+            for i, logits in enumerate(layer_logits):
+                # logits: [B, S, V]
+                # cur_labels: [B, S]
+                cur_labels = labels[..., i]
+                
+                # Flatten for CrossEntropy
+                # logits: [B*S, V], labels: [B*S]
+                loss_fct = CrossEntropyLoss(reduction='none')
+                vocab_size = logits.size(-1)
+                
+                reshaped_logits = logits.view(-1, vocab_size)
+                reshaped_labels = cur_labels.contiguous().view(-1)
+                
+                # Calculate loss per token
+                per_token_loss = loss_fct(reshaped_logits, reshaped_labels)
+                
+                # Reshape back to [B, S] and sum over Sequence dimension to get per-sample loss
+                per_token_loss = per_token_loss.view(batch_size, -1)
+                per_sample_loss = torch.sum(per_token_loss, dim=-1) # [B]
+                
+                all_sum_losses_list.append(per_sample_loss)
+
+            # Stack to [B, n_heads]
+            all_sum_losses = torch.stack(all_sum_losses_list, dim=1)
+
+            # Weighted Loss Aggregation
+            if channelwise_loss_weight is not None:
+                if len(channelwise_loss_weight) != n_heads:
+                    raise ValueError(f"channelwise_loss_weight length {len(channelwise_loss_weight)} != {n_heads}")
+                
+                w_tensor = torch.tensor(channelwise_loss_weight, device=all_sum_losses.device, dtype=all_sum_losses.dtype)
+                
+                # Sample losses: Weighted sum over channels per sample / Total weight
+                # Normalize by token count per channel
+                # Avoid division by zero with epsilon or mask
+                token_counts_safe = all_token_nums.float().clamp(min=1.0)
+                
+                normalized_losses = all_sum_losses / token_counts_safe
+                sample_losses = (normalized_losses * w_tensor).sum(dim=1) / w_tensor.sum()
+                
+                # Channel losses: Sum over batch / Sum tokens over batch
+                total_loss_per_channel = all_sum_losses.sum(dim=0)
+                total_tokens_per_channel = all_token_nums.sum(dim=0).float().clamp(min=1.0)
+                channel_losses = total_loss_per_channel / total_tokens_per_channel
+                
+                # Final scalar loss
+                loss = (channel_losses * w_tensor).sum() / w_tensor.sum()
+            else:
+                # Default average if no weights provided
+                total_tokens = all_token_nums.sum().float().clamp(min=1.0)
+                loss = all_sum_losses.sum() / total_tokens
+                channel_losses = all_sum_losses.sum(dim=0) / all_token_nums.sum(dim=0).clamp(min=1.0)
+
+        return MossTTSDelayOutputWithPast(
+            loss=loss,
+            all_sum_losses=all_sum_losses,
+            all_token_nums=all_token_nums,
+            sample_losses=sample_losses,
+            channel_losses=channel_losses,
+            logits=layer_logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        max_new_tokens: int = 1000,
+        text_temperature: float = 1.5,
+        text_top_p: float = 1.0,
+        text_top_k: int = 50,
+        audio_temperature: float = 1.7,
+        audio_top_p: float = 0.8,
+        audio_top_k: int = 25,
+        audio_repetition_penalty: float = 1.0,
+    ):
+        if text_temperature > 0:
+            text_do_sample = True
+        else:
+            text_temperature = 1
+            text_do_sample = False
+        if audio_temperature > 0:
+            audio_do_sample = True
+        else:
+            audio_temperature = 1
+            audio_do_sample = False
+            
+        past_key_values = None
+        device = input_ids.device
+        current_input_ids = input_ids
+        current_attention_mask = attention_mask
+        batch_size, seq_len, n_vq = input_ids.shape
+        n_vq -= 1
+        
+        generation_ids = input_ids[:]
+        is_stopping = torch.zeros(batch_size, dtype=torch.bool, device=device)
+
+        audio_lengths = torch.zeros(batch_size, dtype=torch.int64, device=device)
+        torch_int64_max = torch.iinfo(torch.int64).max
+        delayed_lengths = torch.full((batch_size,), torch_int64_max, dtype=torch.int64, device=device)
+        
+        is_continuation = (input_ids[:, -1, 0] == self.config.audio_start_token_id) | (input_ids[:, -1, 0] == self.config.audio_assistant_gen_slot_token_id)
+        audio_start_indices = find_last_equal_C(input_ids[..., 0], self.config.audio_start_token_id)
+        audio_start_mask = is_continuation & (audio_start_indices != -1)
+        audio_lengths[audio_start_mask] = seq_len - audio_start_indices[audio_start_mask]
+        
+        is_audio = audio_start_mask.clone()
+        
+        pre_exclude_mask0 = torch.tensor([self.config.pad_token_id, self.config.audio_assistant_gen_slot_token_id, self.config.audio_assistant_delay_slot_token_id, self.config.audio_end_token_id], device=device)
+        pre_exclude_mask1 = torch.ones(self.config.language_config.vocab_size, device=device).bool()
+        pre_exclude_mask1[[self.config.audio_assistant_gen_slot_token_id, self.config.audio_assistant_delay_slot_token_id]] = False
+        
+        for time_step in tqdm(range(max_new_tokens), desc=f"Generating bs{batch_size} ..."):
+            outputs = self(
+                input_ids=current_input_ids,
+                attention_mask=current_attention_mask,
+                past_key_values=past_key_values,
+                use_cache=True,
+            )
+            past_key_values = outputs.past_key_values
+            
+            next_token_logits = [logit[:, -1, :] / text_temperature if logit_idx == 0 else logit[:, -1, :] / audio_temperature for logit_idx, logit in enumerate(outputs.logits)] # List, len=n_vq+1, [batch_size, 1, vocab_size]; 
+            next_token_logits[0] = next_token_logits[0].clone()
+            next_text_token = torch.full((batch_size,), self.config.pad_token_id, device=device)
+            next_text_token[~is_stopping & (delayed_lengths < n_vq)] = self.config.audio_assistant_delay_slot_token_id
+            is_audio_eos = ~is_stopping & (delayed_lengths == n_vq)
+            next_text_token[is_audio_eos] = self.config.audio_end_token_id
+            is_audio[is_audio_eos] = False
+            sampling_text_mask = ~is_stopping & (delayed_lengths > n_vq)
+            next_token_logits[0][~is_audio] = next_token_logits[0][~is_audio].index_fill(-1, pre_exclude_mask0, float('-inf'))
+            next_token_logits[0][is_audio] = next_token_logits[0][is_audio].masked_fill(pre_exclude_mask1, float('-inf'))
+            if time_step == 0:
+                next_token_logits[0][..., 151662] = float('-inf')
+            if time_step <= n_vq:
+                next_token_logits[0][..., self.config.im_end_token_id] = float('-inf')
+                
+            next_text_token[sampling_text_mask] = sample_token(
+                logits=next_token_logits[0][sampling_text_mask],
+                top_p=text_top_p,
+                top_k=text_top_k,
+                do_sample=text_do_sample
+            )
+            is_audio[next_text_token == self.config.audio_start_token_id] = True
+            is_stopping[next_text_token == self.config.im_end_token_id] = True
+
+            next_audio_tokens = torch.full((batch_size, n_vq), self.config.audio_pad_code, device=device)
+            
+            pre_audio_mask = audio_lengths.unsqueeze(1) > torch.arange(n_vq, dtype=int, device=device).expand(batch_size, n_vq)
+            post_audio_mask = torch.arange(n_vq, dtype=int, device=device).expand(batch_size, n_vq) > delayed_lengths.unsqueeze(1) - 1
+            post_audio_mask[delayed_lengths == torch_int64_max] = True
+            sampling_audio_mask = pre_audio_mask & post_audio_mask
+            next_audio_tokens[~sampling_audio_mask] = self.config.audio_pad_code
+            
+            if sampling_audio_mask.sum() > 0:
+                audio_ch0_logits = next_token_logits[1][sampling_audio_mask[:, 0]]
+                audio_logits = torch.stack(next_token_logits[2:], dim=1)[sampling_audio_mask[:, 1:]]
+                audio_ch0_logits[..., self.config.audio_pad_code] = float('-inf')
+                audio_logits[..., self.config.audio_pad_code] = float('-inf')
+                next_audio_tokens[:, 0][sampling_audio_mask[:, 0]] = sample_token(
+                    logits=audio_ch0_logits,
+                    prev_tokens=generation_ids[:, :, 1],
+                    repetition_penalty=audio_repetition_penalty,
+                    top_p=audio_top_p, 
+                    top_k=audio_top_k, 
+                    do_sample=audio_do_sample
+                )
+                next_audio_tokens[:, 1:][sampling_audio_mask[:, 1:]] = sample_token(
+                    logits=audio_logits,
+                    prev_tokens=generation_ids[:, :, 2:],
+                    repetition_penalty=audio_repetition_penalty,
+                    top_p=audio_top_p, 
+                    top_k=audio_top_k, 
+                    do_sample=audio_do_sample
+                )
+                
+            audio_lengths[(next_text_token == self.config.audio_start_token_id) | (next_text_token == self.config.audio_assistant_gen_slot_token_id) | (next_text_token == self.config.audio_assistant_delay_slot_token_id)] += 1
+            audio_lengths[next_text_token == self.config.audio_end_token_id] = 0
+            delayed_lengths[(delayed_lengths == torch_int64_max) & (next_text_token == self.config.audio_assistant_delay_slot_token_id)] = 0
+            delayed_lengths[delayed_lengths != torch_int64_max] += 1
+            delayed_lengths[delayed_lengths > n_vq] = torch_int64_max
+            
+            current_input_ids = torch.cat([next_text_token[:, None, None], next_audio_tokens[:, None, :]], dim=2)
+            current_attention_mask = torch.cat([current_attention_mask, (~is_stopping).unsqueeze(-1)], dim=-1)
+            generation_ids = torch.cat([generation_ids, current_input_ids], dim=1)
+            
+            if is_stopping.sum() == batch_size:
+                break
+                
+        start_indices = find_last_equal_C(input_ids[..., 0], self.config.im_start_token_id) + 3
+        start_lengths = seq_len - start_indices
+
+        output = []
+        for start_idx, start_length, cur_generation_ids in zip(start_indices, start_lengths, generation_ids):
+            output.append((start_length, cur_generation_ids[start_idx:]))
+        
+        return output

processing_moss_tts.py

@@ -0,0 +1,930 @@
+# 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", "OpenMOSS-Team/MOSS-Audio-Tokenizer"
+        )
+
+        pretrained_model_name_or_path = Path(pretrained_model_name_or_path)
+        model_config = cast(
+            MossTTSDelayConfig,
+            AutoConfig.from_pretrained(
+                pretrained_model_name_or_path,
+                *args,
+                trust_remote_code=trust_remote_code,
+                **kwargs,
+            ),
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            *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"}:
+            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)
+            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 + (delay_slot_token * (n_vq - 1))
+            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 = self.apply_delay_pattern(
+                    audio_codes, self.model_config.audio_pad_code
+                )
+                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:
+                delay_audio_codes_list[-1] = delay_audio_codes_list[-1][
+                    : -(n_vq - 1), :
+                ]
+            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]
+        audio_codes = self.apply_de_delay_pattern(audio_codes)
+
+        # 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:
+            wav, sr = torchaudio.load(wav_path)
+            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,6 @@
+{
+    "processor_class": "MossTTSDelayProcessor",
+    "auto_map": {
+        "AutoProcessor": "processing_moss_tts.MossTTSDelayProcessor"
+    }
+}

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
+}