Add Nano-TTS CPU Gradio Space

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+generated_audio/
+.cache/
+weights/tts/.cache/
+weights/codec/.cache/

README.md

@@ -4,11 +4,19 @@ emoji: 📈
 colorFrom: red
 colorTo: green
 sdk: gradio
-sdk_version: 6.11.0
+sdk_version: 6.5.1
+python_version: "3.10"
 app_file: app.py
 pinned: false
 license: apache-2.0
-short_description: space for  MOSS-TTS-Nano
+short_description: CPU-only MOSS TTS Nano Gradio demo with local TTS and codec weights
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+This Space runs Nano-TTS on CPU using the local `weights/tts` and `weights/codec` directories.
+
+Supported modes:
+
+- `voice_clone`: upload a reference audio file or use a built-in preset voice
+- `continuation`: plain TTS, or prompt transcript plus prompt audio
+
+Realtime streaming decode is disabled in this Space. Audio is returned after full synthesis finishes.

app.py

@@ -0,0 +1,500 @@
+from __future__ import annotations
+
+import argparse
+import functools
+import logging
+import os
+import time
+from pathlib import Path
+
+import gradio as gr
+
+from nano_tts_runtime import DEFAULT_VOICE, NanoTTSService, build_default_voice_presets
+from text_normalization_pipeline import prepare_tts_request_texts
+
+APP_DIR = Path(__file__).resolve().parent
+CHECKPOINT_PATH = APP_DIR / "weights" / "tts"
+AUDIO_TOKENIZER_PATH = APP_DIR / "weights" / "codec"
+OUTPUT_DIR = Path("/tmp") / "nano-tts-space"
+PRELOAD_ENV_VAR = "NANO_TTS_PRELOAD_AT_STARTUP"
+
+MODE_VOICE_CLONE = "voice_clone"
+MODE_CONTINUATION = "continuation"
+
+_VOICE_PRESETS = build_default_voice_presets()
+
+
+def build_voice_choices() -> list[tuple[str, str]]:
+    preferred: list[tuple[str, str]] = []
+    fallback: list[tuple[str, str]] = []
+
+    for preset in _VOICE_PRESETS.values():
+        if not preset.prompt_audio_path.is_file():
+            continue
+
+        item = (f"{preset.name} - {preset.description}", preset.name)
+        fallback.append(item)
+        if preset.prompt_audio_path.suffix.lower() == ".wav":
+            preferred.append(item)
+
+    return preferred or fallback
+
+
+VOICE_CHOICES = build_voice_choices()
+DEFAULT_VOICE_VALUE = (
+    DEFAULT_VOICE
+    if any(value == DEFAULT_VOICE for _, value in VOICE_CHOICES)
+    else (VOICE_CHOICES[0][1] if VOICE_CHOICES else "")
+)
+
+
+def parse_bool_env(name: str, default: bool) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "y", "on"}
+
+
+def parse_port(value: str | None, default: int) -> int:
+    if not value:
+        return default
+    try:
+        return int(value)
+    except ValueError:
+        return default
+
+
+def maybe_delete_file(path: str | Path | None) -> None:
+    if not path:
+        return
+    try:
+        Path(path).unlink(missing_ok=True)
+    except OSError:
+        logging.warning("failed to delete temporary file: %s", path, exc_info=True)
+
+
+@functools.lru_cache(maxsize=1)
+def get_tts_service() -> NanoTTSService:
+    return NanoTTSService(
+        checkpoint_path=CHECKPOINT_PATH,
+        audio_tokenizer_path=AUDIO_TOKENIZER_PATH,
+        device="cpu",
+        dtype="float32",
+        attn_implementation="sdpa",
+        output_dir=OUTPUT_DIR,
+    )
+
+
+def preload_service() -> None:
+    started_at = time.monotonic()
+    logging.info(
+        "preloading Nano-TTS model checkpoint=%s codec=%s device=cpu",
+        CHECKPOINT_PATH,
+        AUDIO_TOKENIZER_PATH,
+    )
+    get_tts_service().get_model()
+    logging.info("Nano-TTS preload finished in %.2fs", time.monotonic() - started_at)
+
+
+def render_mode_hint(mode: str) -> str:
+    if mode == MODE_CONTINUATION:
+        return (
+            "Current mode: **Continuation**  \n"
+            "Plain TTS uses only target text. If you upload reference audio, you must also provide its transcript."
+        )
+    return (
+        "Current mode: **Voice Clone**  \n"
+        "Upload a reference audio file or use a built-in preset voice. Audio is returned only after full decoding."
+    )
+
+
+def update_mode_ui(mode: str):
+    if mode == MODE_CONTINUATION:
+        return (
+            gr.update(visible=False),
+            gr.update(
+                visible=True,
+                value="",
+                placeholder="Only for continuation with reference audio.",
+            ),
+            gr.update(label="Reference Audio Upload (optional; required if Prompt Transcript is set)"),
+            render_mode_hint(mode),
+        )
+
+    return (
+        gr.update(visible=True),
+        gr.update(visible=False, value=""),
+        gr.update(label="Reference Audio Upload (optional; overrides preset voice)"),
+        render_mode_hint(mode),
+    )
+
+
+def validate_request(
+    *,
+    text: str,
+    mode: str,
+    prompt_text: str,
+    prompt_audio_path: str | None,
+) -> tuple[str, str | None]:
+    normalized_text = str(text or "").strip()
+    normalized_prompt_text = str(prompt_text or "").strip()
+    has_prompt_audio = bool(prompt_audio_path)
+
+    if not normalized_text:
+        raise ValueError("Please enter text to synthesize.")
+
+    if mode == MODE_VOICE_CLONE:
+        if normalized_prompt_text:
+            raise ValueError("voice_clone mode does not use prompt_text. Leave Prompt Transcript empty.")
+        return normalized_text, None
+
+    if bool(normalized_prompt_text) != has_prompt_audio:
+        raise ValueError(
+            "continuation mode accepts either target text only, or prompt_text and reference audio together."
+        )
+
+    return normalized_text, (normalized_prompt_text or None)
+
+
+def build_status_text(
+    *,
+    result: dict[str, object],
+    prepared_texts: dict[str, object],
+    reference_source: str,
+) -> str:
+    text_chunks = result.get("voice_clone_text_chunks") or []
+    chunk_count = len(text_chunks) if isinstance(text_chunks, list) and text_chunks else 1
+    return (
+        f"Done | mode={result['mode']} | ref={reference_source} | elapsed={result['elapsed_seconds']:.2f}s | "
+        f"sample_rate={result['sample_rate']} | attn={result['effective_global_attn_implementation']} | "
+        f"chunks={chunk_count} | normalization={prepared_texts['normalization_method']}"
+    )
+
+
+def run_inference(
+    text: str,
+    mode: str,
+    voice: str,
+    prompt_audio_path: str | None,
+    prompt_text: str,
+    max_new_frames: int,
+    voice_clone_max_text_tokens: int,
+    do_sample: bool,
+    text_temperature: float,
+    text_top_p: float,
+    text_top_k: int,
+    audio_temperature: float,
+    audio_top_p: float,
+    audio_top_k: int,
+    audio_repetition_penalty: float,
+    seed: float | int,
+):
+    generated_audio_path: str | None = None
+    try:
+        normalized_text, normalized_prompt_text = validate_request(
+            text=text,
+            mode=mode,
+            prompt_text=prompt_text,
+            prompt_audio_path=prompt_audio_path,
+        )
+        prepared_texts = prepare_tts_request_texts(
+            text=normalized_text,
+            prompt_text=normalized_prompt_text or "",
+            voice=voice,
+            enable_wetext=False,
+            text_normalizer_manager=None,
+        )
+        reference_source = (
+            "uploaded_audio"
+            if prompt_audio_path
+            else (f"preset:{voice}" if mode == MODE_VOICE_CLONE else "none")
+        )
+        normalized_seed = None
+        if seed not in {"", None}:
+            resolved_seed = int(seed)
+            if resolved_seed != 0:
+                normalized_seed = resolved_seed
+
+        result = get_tts_service().synthesize(
+            text=str(prepared_texts["text"]),
+            mode=mode,
+            voice=voice,
+            prompt_audio_path=prompt_audio_path or None,
+            prompt_text=str(prepared_texts["prompt_text"]).strip() or None,
+            max_new_frames=int(max_new_frames),
+            voice_clone_max_text_tokens=int(voice_clone_max_text_tokens),
+            do_sample=bool(do_sample),
+            text_temperature=float(text_temperature),
+            text_top_p=float(text_top_p),
+            text_top_k=int(text_top_k),
+            audio_temperature=float(audio_temperature),
+            audio_top_p=float(audio_top_p),
+            audio_top_k=int(audio_top_k),
+            audio_repetition_penalty=float(audio_repetition_penalty),
+            seed=normalized_seed,
+            attn_implementation="sdpa",
+        )
+        generated_audio_path = str(result["audio_path"])
+        return (
+            (int(result["sample_rate"]), result["waveform_numpy"]),
+            build_status_text(
+                result=result,
+                prepared_texts=prepared_texts,
+                reference_source=reference_source,
+            ),
+            str(prepared_texts["normalized_text"]),
+            str(prepared_texts["normalized_prompt_text"]),
+        )
+    except Exception as exc:
+        logging.exception("Nano-TTS inference failed")
+        raise gr.Error(str(exc)) from exc
+    finally:
+        maybe_delete_file(generated_audio_path)
+
+
+def build_demo():
+    custom_css = """
+    :root {
+      --bg: #f5f6f0;
+      --panel: #ffffff;
+      --ink: #15221a;
+      --muted: #5a695e;
+      --line: #d9dfd6;
+      --accent: #285943;
+    }
+    .gradio-container {
+      background:
+        radial-gradient(circle at top left, rgba(162, 198, 167, 0.18), transparent 28%),
+        linear-gradient(180deg, #f5f6f0 0%, #edf1ea 100%);
+      color: var(--ink);
+    }
+    .app-card {
+      border: 1px solid var(--line);
+      border-radius: 18px;
+      background: rgba(255, 255, 255, 0.96);
+      padding: 16px;
+      box-shadow: 0 20px 40px rgba(21, 34, 26, 0.06);
+    }
+    .app-title {
+      font-size: 24px;
+      font-weight: 700;
+      letter-spacing: 0.2px;
+      margin-bottom: 6px;
+    }
+    .app-subtitle {
+      color: var(--muted);
+      font-size: 14px;
+      line-height: 1.5;
+    }
+    #run-btn {
+      background: var(--accent);
+      border: none;
+    }
+    """
+
+    with gr.Blocks(title="Nano-TTS CPU Space", css=custom_css) as demo:
+        gr.Markdown(
+            """
+            <div class="app-card">
+              <div class="app-title">Nano-TTS CPU</div>
+              <div class="app-subtitle">
+                Hugging Face Space edition backed by local <code>weights/tts</code> and <code>weights/codec</code>.
+                Realtime streaming decode is disabled; audio is returned after full synthesis.
+              </div>
+            </div>
+            """
+        )
+
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=3):
+                text = gr.Textbox(
+                    label="Target Text",
+                    lines=10,
+                    placeholder="Enter the text to synthesize.",
+                )
+                mode = gr.Radio(
+                    choices=[
+                        ("Voice Clone", MODE_VOICE_CLONE),
+                        ("Continuation", MODE_CONTINUATION),
+                    ],
+                    value=MODE_VOICE_CLONE,
+                    label="Inference Mode",
+                )
+                mode_hint = gr.Markdown(render_mode_hint(MODE_VOICE_CLONE))
+                voice = gr.Dropdown(
+                    choices=VOICE_CHOICES,
+                    value=DEFAULT_VOICE_VALUE,
+                    label="Preset Voice",
+                    info="Used only by voice_clone when no reference audio is uploaded.",
+                    visible=True,
+                )
+                prompt_audio = gr.Audio(
+                    label="Reference Audio Upload (optional; overrides preset voice)",
+                    type="filepath",
+                    sources=["upload"],
+                )
+                prompt_text = gr.Textbox(
+                    label="Prompt Transcript",
+                    lines=3,
+                    visible=False,
+                    placeholder="Only for continuation with reference audio.",
+                )
+
+                gr.Markdown(
+                    "Robust text normalization is always on. WeTextProcessing is disabled in this CPU Space for a simpler deployment path."
+                )
+
+                with gr.Accordion("Advanced Parameters", open=False):
+                    max_new_frames = gr.Slider(
+                        minimum=64,
+                        maximum=512,
+                        step=16,
+                        value=375,
+                        label="max_new_frames",
+                    )
+                    voice_clone_max_text_tokens = gr.Slider(
+                        minimum=25,
+                        maximum=200,
+                        step=5,
+                        value=75,
+                        label="voice_clone_max_text_tokens",
+                    )
+                    do_sample = gr.Checkbox(
+                        value=True,
+                        label="Enable Sampling",
+                    )
+                    seed = gr.Number(
+                        value=0,
+                        precision=0,
+                        label="Seed (0 = random)",
+                    )
+                    text_temperature = gr.Slider(
+                        minimum=0.1,
+                        maximum=2.0,
+                        step=0.05,
+                        value=1.0,
+                        label="text_temperature",
+                    )
+                    text_top_p = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        step=0.01,
+                        value=1.0,
+                        label="text_top_p",
+                    )
+                    text_top_k = gr.Slider(
+                        minimum=1,
+                        maximum=100,
+                        step=1,
+                        value=50,
+                        label="text_top_k",
+                    )
+                    audio_temperature = gr.Slider(
+                        minimum=0.1,
+                        maximum=2.0,
+                        step=0.05,
+                        value=0.8,
+                        label="audio_temperature",
+                    )
+                    audio_top_p = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        step=0.01,
+                        value=0.95,
+                        label="audio_top_p",
+                    )
+                    audio_top_k = gr.Slider(
+                        minimum=1,
+                        maximum=100,
+                        step=1,
+                        value=25,
+                        label="audio_top_k",
+                    )
+                    audio_repetition_penalty = gr.Slider(
+                        minimum=0.8,
+                        maximum=2.0,
+                        step=0.05,
+                        value=1.2,
+                        label="audio_repetition_penalty",
+                    )
+
+                run_btn = gr.Button("Generate Speech", variant="primary", elem_id="run-btn")
+
+            with gr.Column(scale=2):
+                output_audio = gr.Audio(label="Output Audio", type="numpy")
+                status = gr.Textbox(label="Status", lines=4, interactive=False)
+                normalized_text = gr.Textbox(label="Normalized Text", lines=6, interactive=False)
+                normalized_prompt_text = gr.Textbox(
+                    label="Normalized Prompt Transcript",
+                    lines=4,
+                    interactive=False,
+                )
+
+        mode.change(
+            fn=update_mode_ui,
+            inputs=[mode],
+            outputs=[voice, prompt_text, prompt_audio, mode_hint],
+        )
+
+        run_btn.click(
+            fn=run_inference,
+            inputs=[
+                text,
+                mode,
+                voice,
+                prompt_audio,
+                prompt_text,
+                max_new_frames,
+                voice_clone_max_text_tokens,
+                do_sample,
+                text_temperature,
+                text_top_p,
+                text_top_k,
+                audio_temperature,
+                audio_top_p,
+                audio_top_k,
+                audio_repetition_penalty,
+                seed,
+            ],
+            outputs=[output_audio, status, normalized_text, normalized_prompt_text],
+        )
+
+    return demo
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Nano-TTS Hugging Face Space")
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=int(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT", "7860"))),
+    )
+    parser.add_argument("--share", action="store_true")
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+        level=logging.INFO,
+    )
+
+    args.host = os.getenv("GRADIO_SERVER_NAME", args.host)
+    args.port = parse_port(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT")), args.port)
+
+    preload_enabled = parse_bool_env(PRELOAD_ENV_VAR, default=not bool(os.getenv("SPACE_ID")))
+    if preload_enabled:
+        preload_service()
+    else:
+        logging.info("Skipping model preload (set %s=1 to enable).", PRELOAD_ENV_VAR)
+
+    demo = build_demo()
+    demo.queue(max_size=4, default_concurrency_limit=1).launch(
+        server_name=args.host,
+        server_port=args.port,
+        share=args.share,
+        ssr_mode=False,
+    )
+
+
+if __name__ == "__main__":
+    main()

asserts/audio/en_1.wav

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

asserts/audio/en_2.wav

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

asserts/audio/en_3.wav

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

asserts/audio/en_4.wav

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

asserts/audio/en_5.wav

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

asserts/audio/jp_1.mp3

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

asserts/audio/jp_2.wav

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

asserts/audio/jp_3.wav

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

asserts/audio/jp_4.wav

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

asserts/audio/jp_5.wav

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

asserts/audio/zh_1.wav

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

asserts/audio/zh_2.wav

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

asserts/audio/zh_3.wav

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

asserts/audio/zh_4.wav

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

asserts/audio/zh_5.wav

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

asserts/audio/zh_6.wav

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

nano_tts_runtime.py

@@ -0,0 +1,727 @@
+from __future__ import annotations
+
+import importlib
+import logging
+import threading
+import time
+import uuid
+from dataclasses import dataclass
+from functools import lru_cache
+from pathlib import Path
+from typing import Iterator, Optional
+
+import numpy as np
+import torch
+from transformers import AutoModel, AutoModelForCausalLM
+
+MOSS_AUDIO_TOKENIZER_TYPE = "moss-audio-tokenizer-nano"
+
+APP_DIR = Path(__file__).resolve().parent
+DEFAULT_CHECKPOINT_PATH = APP_DIR / "weights" / "tts"
+DEFAULT_AUDIO_TOKENIZER_PATH = APP_DIR / "weights" / "codec"
+DEFAULT_PROMPT_AUDIO_DIR = APP_DIR / "asserts" / "audio"
+DEFAULT_OUTPUT_DIR = APP_DIR / "generated_audio"
+
+_DEFAULT_VOICE_FILES: dict[str, tuple[str, str]] = {
+    "Junhao": ("zh_1.wav", "Chinese male voice A"),
+    "Zhiming": ("zh_2.wav", "Chinese male voice B"),
+    "Weiguo": ("zh_5.wav", "Chinese male voice C"),
+    "Xiaoyu": ("zh_3.wav", "Chinese female voice A"),
+    "Yuewen": ("zh_4.wav", "Chinese female voice B"),
+    "Lingyu": ("zh_6.wav", "Chinese female voice C"),
+    "Trump": ("en_1.wav", "Trump reference voice"),
+    "Ava": ("en_2.wav", "English female voice A"),
+    "Bella": ("en_3.wav", "English female voice B"),
+    "Adam": ("en_4.wav", "English male voice A"),
+    "Nathan": ("en_5.wav", "English male voice B"),
+    "Sakura": ("jp_1.mp3", "Japanese female voice A"),
+    "Yui": ("jp_2.wav", "Japanese female voice B"),
+    "Aoi": ("jp_3.wav", "Japanese female voice C"),
+    "Hina": ("jp_4.wav", "Japanese female voice D"),
+    "Mei": ("jp_5.wav", "Japanese female voice E"),
+}
+
+DEFAULT_VOICE = "Junhao"
+FLASH_ATTENTION_DTYPES = {torch.float16, torch.bfloat16}
+
+
+@dataclass(frozen=True)
+class VoicePreset:
+    name: str
+    prompt_audio_path: Path
+    description: str
+
+
+def build_default_voice_presets() -> dict[str, VoicePreset]:
+    presets: dict[str, VoicePreset] = {}
+    for voice_name, (file_name, description) in _DEFAULT_VOICE_FILES.items():
+        prompt_path = (DEFAULT_PROMPT_AUDIO_DIR / file_name).resolve()
+        presets[voice_name] = VoicePreset(
+            name=voice_name,
+            prompt_audio_path=prompt_path,
+            description=description,
+        )
+    return presets
+
+
+def resolve_device(device_arg: str) -> torch.device:
+    if device_arg == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return torch.device(device_arg)
+
+
+def resolve_dtype(dtype_arg: str, device: torch.device) -> torch.dtype:
+    if dtype_arg == "float32":
+        return torch.float32
+    if dtype_arg == "float16":
+        return torch.float16
+    if dtype_arg == "bfloat16":
+        return torch.bfloat16
+    if device.type == "cuda":
+        if hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported():
+            return torch.bfloat16
+        return torch.float16
+    return torch.float32
+
+
+def waveform_to_numpy(waveform: torch.Tensor | np.ndarray) -> np.ndarray:
+    if torch.is_tensor(waveform):
+        array = waveform.detach().cpu().numpy()
+    else:
+        array = np.asarray(waveform)
+    if array.ndim == 1:
+        return array.astype(np.float32, copy=False)
+    if array.ndim != 2:
+        raise ValueError(f"Unsupported waveform shape: {array.shape}")
+    if array.shape[0] <= 8 and array.shape[0] < array.shape[1]:
+        array = array.T
+    return array.astype(np.float32, copy=False)
+
+
+@lru_cache(maxsize=1)
+def _has_flash_attn() -> bool:
+    try:
+        importlib.import_module("flash_attn")
+    except Exception:
+        return False
+    return True
+
+
+class NanoTTSService:
+    def __init__(
+        self,
+        *,
+        checkpoint_path: str | Path = DEFAULT_CHECKPOINT_PATH,
+        audio_tokenizer_path: str | Path = DEFAULT_AUDIO_TOKENIZER_PATH,
+        device: str = "auto",
+        dtype: str = "auto",
+        attn_implementation: str = "auto",
+        output_dir: str | Path = DEFAULT_OUTPUT_DIR,
+        voice_presets: Optional[dict[str, VoicePreset]] = None,
+    ) -> None:
+        self.checkpoint_path = Path(checkpoint_path).expanduser().resolve()
+        self.audio_tokenizer_path = Path(audio_tokenizer_path).expanduser().resolve()
+        self.output_dir = Path(output_dir).expanduser().resolve()
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        self.voice_presets = voice_presets or build_default_voice_presets()
+        self.default_voice = DEFAULT_VOICE if DEFAULT_VOICE in self.voice_presets else next(iter(self.voice_presets))
+
+        self.device = resolve_device(device)
+        self.dtype = resolve_dtype(dtype, self.device)
+        self.attn_implementation = self._resolve_attn_implementation(attn_implementation)
+
+        self._lock = threading.RLock()
+        self._model = None
+        self._audio_tokenizer = None
+        self._checkpoint_global_attn_implementation: str | None = None
+        self._checkpoint_local_attn_implementation: str | None = None
+        self._configured_global_attn_implementation: str | None = None
+        self._configured_local_attn_implementation: str | None = None
+        self._configured_audio_tokenizer_attn_implementation: str | None = None
+        self._configured_audio_tokenizer_compute_dtype: str | None = None
+
+    def _can_use_flash_attention(self) -> bool:
+        return self.device.type == "cuda" and self.dtype in FLASH_ATTENTION_DTYPES and _has_flash_attn()
+
+    def _resolve_runtime_default_attn_implementation(self) -> str:
+        return "flash_attention_2" if self._can_use_flash_attention() else "sdpa"
+
+    def _resolve_attn_implementation(self, requested: str | None) -> str | None:
+        normalized = str(requested).strip().lower() if requested is not None else "auto"
+        if not normalized or normalized in {"auto", "default", "model_default"}:
+            return None
+        if normalized not in {"sdpa", "flash_attention_2", "eager"}:
+            raise ValueError(
+                "attn_implementation must be one of: model_default/auto, sdpa, flash_attention_2, eager"
+            )
+        if normalized == "flash_attention_2":
+            if not self._can_use_flash_attention():
+                logging.warning(
+                    "flash_attention_2 requires CUDA, flash_attn, and fp16/bf16; falling back to sdpa "
+                    "(device=%s dtype=%s flash_attn=%s)",
+                    self.device,
+                    self.dtype,
+                    _has_flash_attn(),
+                )
+                return "sdpa"
+        return normalized
+
+    @staticmethod
+    def _normalize_loaded_attn_implementation(attn_implementation: object) -> str:
+        normalized = str(attn_implementation).strip().lower() if attn_implementation is not None else ""
+        if not normalized or normalized == "none":
+            return "eager"
+        return normalized
+
+    def _resolve_request_attention_implementation(
+        self,
+        requested: Optional[str],
+    ) -> tuple[str, str, str]:
+        normalized = str(requested).strip().lower() if requested is not None else ""
+        resolved = self._resolve_attn_implementation(normalized)
+        if resolved is not None:
+            return normalized, resolved, resolved
+
+        if self.attn_implementation is not None:
+            return self.attn_implementation, self.attn_implementation, self.attn_implementation
+
+        runtime_default = self._resolve_runtime_default_attn_implementation()
+        return "auto", runtime_default, runtime_default
+
+    @staticmethod
+    def _resolve_codec_attention_implementation(tts_attn_implementation: str) -> str:
+        return "flash_attention_2" if tts_attn_implementation == "flash_attention_2" else "sdpa"
+
+    def _resolve_codec_compute_dtype(self, codec_attn_implementation: str) -> str:
+        if codec_attn_implementation == "flash_attention_2":
+            return "bf16" if self.dtype == torch.bfloat16 else "fp16"
+        return "fp32"
+
+    @staticmethod
+    def _apply_model_attention_implementation(model, *, global_attn: str, local_attn: str) -> None:
+        if hasattr(model, "_set_attention_implementation"):
+            model._set_attention_implementation(global_attn, local_attn_implementation=local_attn)
+
+    def _install_stream_decode_budget_patch(self, model) -> None:
+        if self.device.type != "cuda":
+            return
+
+        model_cls = model.__class__
+        if getattr(model_cls, "_nanotts_stream_decode_budget_patch_installed", False):
+            return
+
+        compute_stream_lead = getattr(model_cls, "_compute_stream_lead_seconds", None)
+        resolve_stream_budget = getattr(model_cls, "_resolve_stream_decode_frame_budget", None)
+        if not callable(compute_stream_lead) or not callable(resolve_stream_budget):
+            return
+
+        def _patched_resolve_stream_decode_frame_budget(
+            *,
+            emitted_samples_total: int,
+            sample_rate: int,
+            first_audio_emitted_at,
+        ) -> int:
+            # The upstream streaming policy starts with one decode frame
+            # (about 80 ms audio), which makes CUDA realtime decode emit many
+            # tiny chunks and underrun browser playback on this checkpoint.
+            lead_seconds = compute_stream_lead(
+                emitted_samples_total=emitted_samples_total,
+                sample_rate=sample_rate,
+                first_audio_emitted_at=first_audio_emitted_at,
+            )
+            if first_audio_emitted_at is None or lead_seconds < 0.20:
+                return 4
+            if lead_seconds < 0.55:
+                return 6
+            if lead_seconds < 1.10:
+                return 8
+            return 12
+
+        model_cls._nanotts_original_resolve_stream_decode_frame_budget = resolve_stream_budget
+        model_cls._resolve_stream_decode_frame_budget = staticmethod(_patched_resolve_stream_decode_frame_budget)
+        model_cls._nanotts_stream_decode_budget_patch_installed = True
+        logging.info("installed Nano-TTS CUDA streaming decode budget patch")
+
+    def _discard_loaded_model_locked(self, reason: str) -> None:
+        if self._model is None:
+            return
+        logging.warning("discarding loaded Nano-TTS model state: %s", reason)
+        self._model = None
+        if self.device.type == "cuda":
+            torch.cuda.empty_cache()
+
+    def _discard_loaded_audio_tokenizer_locked(self, reason: str) -> None:
+        if self._audio_tokenizer is None:
+            return
+        logging.warning("discarding loaded Nano-TTS audio tokenizer state: %s", reason)
+        self._audio_tokenizer = None
+        self._configured_audio_tokenizer_attn_implementation = None
+        self._configured_audio_tokenizer_compute_dtype = None
+        if self.device.type == "cuda":
+            torch.cuda.empty_cache()
+
+    def _restore_model_execution_state(self, model):
+        current_parameter = next(model.parameters(), None)
+        if current_parameter is None or current_parameter.dtype == self.dtype:
+            return model
+        self._discard_loaded_model_locked(
+            f"current_dtype={current_parameter.dtype} expected_dtype={self.dtype}; reloading checkpoint"
+        )
+        return self._load_model_locked()
+
+    def _read_model_attention_implementation(self, model) -> tuple[str, str]:
+        global_attn = self._normalize_loaded_attn_implementation(
+            getattr(getattr(model, "transformer", None), "attn_implementation", None)
+        )
+        local_attn = self._normalize_loaded_attn_implementation(
+            getattr(getattr(model, "local_transformer", None), "attn_implementation", None)
+        )
+        return global_attn, local_attn
+
+    def _ensure_paths(self) -> None:
+        if not self.checkpoint_path.exists():
+            raise FileNotFoundError(f"Nano-TTS checkpoint not found: {self.checkpoint_path}")
+        if not self.audio_tokenizer_path.exists():
+            raise FileNotFoundError(f"Audio tokenizer checkpoint not found: {self.audio_tokenizer_path}")
+
+    def _load_audio_tokenizer_locked(self, *, tts_attn_implementation: str):
+        codec_attn_implementation = self._resolve_codec_attention_implementation(tts_attn_implementation)
+        codec_compute_dtype = self._resolve_codec_compute_dtype(codec_attn_implementation)
+
+        if self._audio_tokenizer is None:
+            logging.info(
+                "loading Nano-TTS audio tokenizer checkpoint=%s device=%s attn=%s compute_dtype=%s",
+                self.audio_tokenizer_path,
+                self.device,
+                codec_attn_implementation,
+                codec_compute_dtype,
+            )
+            audio_tokenizer = AutoModel.from_pretrained(
+                str(self.audio_tokenizer_path),
+                trust_remote_code=True,
+                local_files_only=True,
+            )
+            if hasattr(audio_tokenizer, "eval"):
+                audio_tokenizer.eval()
+            self._audio_tokenizer = audio_tokenizer
+
+        audio_tokenizer = self._audio_tokenizer
+        if hasattr(audio_tokenizer, "to"):
+            audio_tokenizer = audio_tokenizer.to(self.device)
+        if hasattr(audio_tokenizer, "set_attention_implementation"):
+            audio_tokenizer.set_attention_implementation(codec_attn_implementation)
+        if hasattr(audio_tokenizer, "set_compute_dtype"):
+            audio_tokenizer.set_compute_dtype(codec_compute_dtype)
+        if hasattr(audio_tokenizer, "eval"):
+            audio_tokenizer.eval()
+
+        self._audio_tokenizer = audio_tokenizer
+        self._configured_audio_tokenizer_attn_implementation = codec_attn_implementation
+        self._configured_audio_tokenizer_compute_dtype = codec_compute_dtype
+        return self._audio_tokenizer
+
+    def _load_model_locked(self):
+        if self._model is not None:
+            return self._model
+
+        self._ensure_paths()
+        logging.info(
+            "loading Nano-TTS checkpoint=%s audio_tokenizer=%s device=%s dtype=%s attn=%s",
+            self.checkpoint_path,
+            self.audio_tokenizer_path,
+            self.device,
+            self.dtype,
+            self.attn_implementation or "model_default",
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            str(self.checkpoint_path),
+            trust_remote_code=True,
+            local_files_only=True,
+        )
+        model.to(device=self.device, dtype=self.dtype)
+        self._checkpoint_global_attn_implementation, self._checkpoint_local_attn_implementation = (
+            self._read_model_attention_implementation(model)
+        )
+        _, default_global_attn, default_local_attn = self._resolve_request_attention_implementation(None)
+        self._apply_model_attention_implementation(
+            model,
+            global_attn=default_global_attn,
+            local_attn=default_local_attn,
+        )
+        self._install_stream_decode_budget_patch(model)
+        model.eval()
+        self._configured_global_attn_implementation, self._configured_local_attn_implementation = (
+            self._read_model_attention_implementation(model)
+        )
+        self._model = model
+        return self._model
+
+    def get_model(self):
+        with self._lock:
+            return self._load_model_locked()
+
+    def list_voice_names(self) -> list[str]:
+        return list(self.voice_presets.keys())
+
+    def get_voice_preset(self, voice_name: Optional[str]) -> VoicePreset:
+        if voice_name and voice_name in self.voice_presets:
+            return self.voice_presets[voice_name]
+        return self.voice_presets[self.default_voice]
+
+    def resolve_prompt_audio_path(
+        self,
+        *,
+        voice: Optional[str] = None,
+        prompt_audio_path: Optional[str | Path] = None,
+    ) -> Path:
+        if prompt_audio_path:
+            resolved = Path(prompt_audio_path).expanduser().resolve()
+            if not resolved.exists():
+                raise FileNotFoundError(f"Prompt audio not found: {resolved}")
+            return resolved
+
+        preset = self.get_voice_preset(voice)
+        if not preset.prompt_audio_path.exists():
+            raise FileNotFoundError(f"Voice preset prompt audio not found: {preset.prompt_audio_path}")
+        return preset.prompt_audio_path
+
+    def preload(self, *, voices: Optional[list[str]] = None, load_model: bool = True) -> dict[str, object]:
+        loaded_voices: list[str] = []
+        if load_model:
+            self.get_model()
+        for voice_name in voices or [self.default_voice]:
+            preset = self.get_voice_preset(voice_name)
+            if preset.prompt_audio_path.exists():
+                loaded_voices.append(preset.name)
+        return {
+            "loaded_voices": loaded_voices,
+            "device": str(self.device),
+            "dtype": str(self.dtype),
+            "attn_implementation": self.attn_implementation or "auto",
+            "checkpoint_default_attn_implementation": self._checkpoint_global_attn_implementation or "eager",
+            "checkpoint_default_local_attn_implementation": self._checkpoint_local_attn_implementation or "eager",
+            "configured_attn_implementation": self._configured_global_attn_implementation or "eager",
+            "configured_local_attn_implementation": self._configured_local_attn_implementation or "eager",
+            "configured_codec_attn_implementation": self._configured_audio_tokenizer_attn_implementation or "unknown",
+            "configured_codec_compute_dtype": self._configured_audio_tokenizer_compute_dtype or "unknown",
+        }
+
+    def _build_output_path(self, prefix: str) -> Path:
+        timestamp = time.strftime("%Y%m%d_%H%M%S")
+        random_suffix = uuid.uuid4().hex[:8]
+        return self.output_dir / f"{prefix}_{timestamp}_{random_suffix}.wav"
+
+    def synthesize(
+        self,
+        *,
+        text: str,
+        voice: Optional[str] = None,
+        mode: str = "voice_clone",
+        output_audio_path: Optional[str | Path] = None,
+        prompt_audio_path: Optional[str | Path] = None,
+        prompt_text: Optional[str] = None,
+        max_new_frames: int = 375,
+        voice_clone_max_text_tokens: int = 75,
+        voice_clone_max_memory_per_sample_gb: float = 1.0,
+        tts_max_batch_size: int = 0,
+        codec_max_batch_size: int = 0,
+        do_sample: bool = True,
+        text_temperature: float = 1.0,
+        text_top_p: float = 1.0,
+        text_top_k: int = 50,
+        audio_temperature: float = 0.8,
+        audio_top_p: float = 0.95,
+        audio_top_k: int = 25,
+        audio_repetition_penalty: float = 1.2,
+        nq: Optional[int] = None,
+        seed: Optional[int] = None,
+        attn_implementation: Optional[str] = None,
+    ) -> dict[str, object]:
+        normalized_text = str(text or "").strip()
+        if not normalized_text:
+            raise ValueError("text is required")
+
+        normalized_mode = str(mode).strip().lower()
+        if normalized_mode not in {"continuation", "voice_clone"}:
+            raise ValueError("mode must be either 'continuation' or 'voice_clone'")
+
+        effective_prompt_audio_path: Optional[Path] = None
+        resolved_voice = self.get_voice_preset(voice).name
+        if normalized_mode == "voice_clone":
+            effective_prompt_audio_path = self.resolve_prompt_audio_path(
+                voice=resolved_voice,
+                prompt_audio_path=prompt_audio_path,
+            )
+        elif prompt_audio_path is not None:
+            effective_prompt_audio_path = self.resolve_prompt_audio_path(prompt_audio_path=prompt_audio_path)
+            if not prompt_text:
+                raise ValueError("continuation mode with prompt_audio_path also requires prompt_text")
+
+        output_path = (
+            Path(output_audio_path).expanduser().resolve()
+            if output_audio_path is not None
+            else self._build_output_path(prefix=f"{resolved_voice}_{normalized_mode}")
+        )
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        started_at = time.monotonic()
+        with self._lock:
+            model = self._load_model_locked()
+            model = self._restore_model_execution_state(model)
+            requested_attn_implementation, effective_global_attn_implementation, effective_local_attn_implementation = (
+                self._resolve_request_attention_implementation(attn_implementation)
+            )
+            audio_tokenizer = self._load_audio_tokenizer_locked(
+                tts_attn_implementation=effective_global_attn_implementation
+            )
+            self._apply_model_attention_implementation(
+                model,
+                global_attn=effective_global_attn_implementation,
+                local_attn=effective_local_attn_implementation,
+            )
+            if seed is not None:
+                torch.manual_seed(seed)
+                if torch.cuda.is_available():
+                    torch.cuda.manual_seed_all(seed)
+
+            try:
+                result = model.inference(
+                    text=normalized_text,
+                    output_audio_path=str(output_path),
+                    mode=normalized_mode,
+                    prompt_text=prompt_text,
+                    prompt_audio_path=None if effective_prompt_audio_path is None else str(effective_prompt_audio_path),
+                    text_tokenizer_path=str(self.checkpoint_path),
+                    audio_tokenizer=audio_tokenizer,
+                    device=self.device,
+                    nq=nq,
+                    max_new_frames=int(max_new_frames),
+                    voice_clone_max_text_tokens=int(voice_clone_max_text_tokens),
+                    voice_clone_max_memory_per_sample_gb=float(voice_clone_max_memory_per_sample_gb),
+                    tts_max_batch_size=int(tts_max_batch_size),
+                    codec_max_batch_size=int(codec_max_batch_size),
+                    do_sample=bool(do_sample),
+                    use_kv_cache=True,
+                    text_temperature=float(text_temperature),
+                    text_top_p=float(text_top_p),
+                    text_top_k=int(text_top_k),
+                    audio_temperature=float(audio_temperature),
+                    audio_top_p=float(audio_top_p),
+                    audio_top_k=int(audio_top_k),
+                    audio_repetition_penalty=float(audio_repetition_penalty),
+                )
+            except Exception:
+                self._discard_loaded_audio_tokenizer_locked(
+                    "inference failed; reloading audio tokenizer on next request"
+                )
+                self._discard_loaded_model_locked("inference failed; reloading checkpoint on next request")
+                raise
+            effective_global_attn_implementation, effective_local_attn_implementation = (
+                self._read_model_attention_implementation(model)
+            )
+            current_parameter = next(model.parameters(), None)
+            if current_parameter is not None and current_parameter.dtype != self.dtype:
+                self._discard_loaded_model_locked(
+                    f"inference left model in dtype={current_parameter.dtype}; reloading checkpoint on next request"
+                )
+
+        waveform = result["waveform"].detach().cpu()
+        waveform_numpy = waveform_to_numpy(waveform)
+        return {
+            "audio_path": str(output_path),
+            "sample_rate": int(result["sample_rate"]),
+            "waveform": waveform,
+            "waveform_numpy": waveform_numpy,
+            "audio_token_ids": result["audio_token_ids"],
+            "reference_audio_token_ids": result["reference_audio_token_ids"],
+            "elapsed_seconds": time.monotonic() - started_at,
+            "voice": resolved_voice,
+            "mode": normalized_mode,
+            "prompt_audio_path": None if effective_prompt_audio_path is None else str(effective_prompt_audio_path),
+            "requested_attn_implementation": requested_attn_implementation,
+            "effective_global_attn_implementation": effective_global_attn_implementation,
+            "effective_local_attn_implementation": effective_local_attn_implementation,
+            "voice_clone_text_chunks": result.get("voice_clone_text_chunks"),
+            "voice_clone_chunk_batch_size": result.get("voice_clone_chunk_batch_size"),
+            "voice_clone_codec_batch_size": result.get("voice_clone_codec_batch_size"),
+        }
+
+    def synthesize_stream(
+        self,
+        *,
+        text: str,
+        voice: Optional[str] = None,
+        mode: str = "voice_clone",
+        output_audio_path: Optional[str | Path] = None,
+        prompt_audio_path: Optional[str | Path] = None,
+        prompt_text: Optional[str] = None,
+        max_new_frames: int = 375,
+        voice_clone_max_text_tokens: int = 75,
+        voice_clone_max_memory_per_sample_gb: float = 1.0,
+        tts_max_batch_size: int = 0,
+        codec_max_batch_size: int = 0,
+        do_sample: bool = True,
+        text_temperature: float = 1.0,
+        text_top_p: float = 1.0,
+        text_top_k: int = 50,
+        audio_temperature: float = 0.8,
+        audio_top_p: float = 0.95,
+        audio_top_k: int = 25,
+        audio_repetition_penalty: float = 1.2,
+        nq: Optional[int] = None,
+        seed: Optional[int] = None,
+        attn_implementation: Optional[str] = None,
+    ) -> Iterator[dict[str, object]]:
+        normalized_text = str(text or "").strip()
+        if not normalized_text:
+            raise ValueError("text is required")
+
+        normalized_mode = str(mode).strip().lower()
+        if normalized_mode not in {"continuation", "voice_clone"}:
+            raise ValueError("mode must be either 'continuation' or 'voice_clone'")
+
+        effective_prompt_audio_path: Optional[Path] = None
+        resolved_voice = self.get_voice_preset(voice).name
+        if normalized_mode == "voice_clone":
+            effective_prompt_audio_path = self.resolve_prompt_audio_path(
+                voice=resolved_voice,
+                prompt_audio_path=prompt_audio_path,
+            )
+        elif prompt_audio_path is not None:
+            effective_prompt_audio_path = self.resolve_prompt_audio_path(prompt_audio_path=prompt_audio_path)
+            if not prompt_text:
+                raise ValueError("continuation mode with prompt_audio_path also requires prompt_text")
+
+        output_path = (
+            Path(output_audio_path).expanduser().resolve()
+            if output_audio_path is not None
+            else self._build_output_path(prefix=f"{resolved_voice}_{normalized_mode}_stream")
+        )
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        started_at = time.monotonic()
+        final_result: dict[str, object] | None = None
+        with self._lock:
+            model = self._load_model_locked()
+            model = self._restore_model_execution_state(model)
+            requested_attn_implementation, effective_global_attn_implementation, effective_local_attn_implementation = (
+                self._resolve_request_attention_implementation(attn_implementation)
+            )
+            audio_tokenizer = self._load_audio_tokenizer_locked(
+                tts_attn_implementation=effective_global_attn_implementation
+            )
+            self._apply_model_attention_implementation(
+                model,
+                global_attn=effective_global_attn_implementation,
+                local_attn=effective_local_attn_implementation,
+            )
+            if seed is not None:
+                torch.manual_seed(seed)
+                if torch.cuda.is_available():
+                    torch.cuda.manual_seed_all(seed)
+
+            try:
+                for event in model.inference_stream(
+                    text=normalized_text,
+                    output_audio_path=str(output_path),
+                    mode=normalized_mode,
+                    prompt_text=prompt_text,
+                    prompt_audio_path=None if effective_prompt_audio_path is None else str(effective_prompt_audio_path),
+                    text_tokenizer_path=str(self.checkpoint_path),
+                    audio_tokenizer=audio_tokenizer,
+                    device=self.device,
+                    nq=nq,
+                    max_new_frames=int(max_new_frames),
+                    voice_clone_max_text_tokens=int(voice_clone_max_text_tokens),
+                    voice_clone_max_memory_per_sample_gb=float(voice_clone_max_memory_per_sample_gb),
+                    tts_max_batch_size=int(tts_max_batch_size),
+                    codec_max_batch_size=int(codec_max_batch_size),
+                    do_sample=bool(do_sample),
+                    use_kv_cache=True,
+                    text_temperature=float(text_temperature),
+                    text_top_p=float(text_top_p),
+                    text_top_k=int(text_top_k),
+                    audio_temperature=float(audio_temperature),
+                    audio_top_p=float(audio_top_p),
+                    audio_top_k=int(audio_top_k),
+                    audio_repetition_penalty=float(audio_repetition_penalty),
+                ):
+                    event_type = str(event.get("type", ""))
+                    if event_type == "audio":
+                        waveform = torch.as_tensor(event["waveform"], dtype=torch.float32).cpu()
+                        yield {
+                            "type": "audio",
+                            "waveform": waveform,
+                            "waveform_numpy": waveform_to_numpy(waveform),
+                            "sample_rate": int(event["sample_rate"]),
+                            "chunk_index": int(event.get("chunk_index", 0)),
+                            "is_pause": bool(event.get("is_pause", False)),
+                            "emitted_audio_seconds": float(event.get("emitted_audio_seconds", 0.0)),
+                            "lead_seconds": float(event.get("lead_seconds", 0.0)),
+                        }
+                        continue
+                    if event_type == "result":
+                        final_result = dict(event)
+            except Exception:
+                self._discard_loaded_audio_tokenizer_locked(
+                    "streaming inference failed; reloading audio tokenizer on next request"
+                )
+                self._discard_loaded_model_locked("streaming inference failed; reloading checkpoint on next request")
+                raise
+
+            effective_global_attn_implementation, effective_local_attn_implementation = (
+                self._read_model_attention_implementation(model)
+            )
+            current_parameter = next(model.parameters(), None)
+            if current_parameter is not None and current_parameter.dtype != self.dtype:
+                self._discard_loaded_model_locked(
+                    f"streaming inference left model in dtype={current_parameter.dtype}; reloading checkpoint on next request"
+                )
+
+        if final_result is None:
+            raise RuntimeError("Streaming synthesis finished without a final result.")
+
+        waveform = torch.as_tensor(final_result["waveform"], dtype=torch.float32).cpu()
+        yield {
+            "type": "result",
+            "audio_path": str(final_result["audio_path"]),
+            "sample_rate": int(final_result["sample_rate"]),
+            "waveform": waveform,
+            "waveform_numpy": waveform_to_numpy(waveform),
+            "audio_token_ids": final_result["audio_token_ids"],
+            "reference_audio_token_ids": final_result["reference_audio_token_ids"],
+            "elapsed_seconds": time.monotonic() - started_at,
+            "voice": resolved_voice,
+            "mode": normalized_mode,
+            "prompt_audio_path": None if effective_prompt_audio_path is None else str(effective_prompt_audio_path),
+            "requested_attn_implementation": requested_attn_implementation,
+            "effective_global_attn_implementation": effective_global_attn_implementation,
+            "effective_local_attn_implementation": effective_local_attn_implementation,
+            "voice_clone_text_chunks": final_result.get("voice_clone_text_chunks"),
+            "voice_clone_chunk_batch_size": final_result.get("voice_clone_chunk_batch_size"),
+            "voice_clone_codec_batch_size": final_result.get("voice_clone_codec_batch_size"),
+        }
+
+    def warmup(
+        self,
+        *,
+        text: str = "你好，欢迎使用 Nano-TTS。",
+        voice: Optional[str] = None,
+    ) -> dict[str, object]:
+        return self.synthesize(
+            text=text,
+            voice=voice or self.default_voice,
+            mode="voice_clone",
+            output_audio_path=self.output_dir / "_warmup" / "warmup.wav",
+            max_new_frames=96,
+            voice_clone_max_text_tokens=75,
+            do_sample=False,
+            text_temperature=1.0,
+            text_top_p=1.0,
+            text_top_k=50,
+            audio_temperature=0.8,
+            audio_top_p=0.95,
+            audio_top_k=25,
+            audio_repetition_penalty=1.0,
+        )

requirements.txt

@@ -0,0 +1,7 @@
+numpy>=1.24
+sentencepiece>=0.1.99
+torch==2.7.0
+torchaudio==2.7.0
+transformers==4.57.1
+safetensors>=0.4.3
+gradio==6.5.1

text_normalization_pipeline.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import logging
+import re
+import threading
+from dataclasses import dataclass
+
+from tts_robust_normalizer_single_script import normalize_tts_text
+
+ENGLISH_VOICES = frozenset({"Trump", "Ava", "Bella", "Adam", "Nathan"})
+
+
+@dataclass(frozen=True)
+class TextNormalizationSnapshot:
+    state: str
+    message: str
+    error: str | None = None
+    ready: bool = False
+    available: bool = False
+
+    @property
+    def failed(self) -> bool:
+        return self.state == "failed"
+
+
+class WeTextProcessingManager:
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._normalize_lock = threading.Lock()
+        self._thread: threading.Thread | None = None
+        self._started = False
+        self._state = "pending"
+        self._message = "Waiting for WeTextProcessing preload."
+        self._error: str | None = None
+        self._available = True
+        self._normalizers: dict[str, object] | None = None
+
+    def snapshot(self) -> TextNormalizationSnapshot:
+        with self._lock:
+            return TextNormalizationSnapshot(
+                state=self._state,
+                message=self._message,
+                error=self._error,
+                ready=self._state == "ready",
+                available=self._available,
+            )
+
+    def _set_state(self, *, state: str, message: str, error: str | None = None) -> None:
+        with self._lock:
+            self._state = state
+            self._message = message
+            self._error = error
+
+    def start(self) -> None:
+        with self._lock:
+            if self._started:
+                return
+            self._started = True
+            self._thread = threading.Thread(target=self._run, name="wetext-preload", daemon=True)
+            self._thread.start()
+
+    def ensure_ready(self) -> TextNormalizationSnapshot:
+        with self._lock:
+            if not self._started:
+                self._started = True
+                self._thread = threading.Thread(target=self._run, name="wetext-preload", daemon=True)
+                self._thread.start()
+            thread = self._thread
+        if thread is not None and thread.is_alive():
+            thread.join()
+        return self.snapshot()
+
+    def close(self) -> None:
+        return
+
+    def _run(self) -> None:
+        if not self._available:
+            self._set_state(
+                state="failed",
+                message="WeTextProcessing unavailable.",
+                error="installed WeTextProcessing modules are unavailable",
+            )
+            return
+        try:
+            self._set_state(state="running", message="Loading WeTextProcessing graphs.", error=None)
+            self._ensure_normalizers_loaded()
+            self._set_state(state="ready", message="WeTextProcessing ready. languages=zh,en", error=None)
+        except Exception as exc:
+            logging.exception("WeTextProcessing preload failed")
+            self._set_state(state="failed", message="WeTextProcessing preload failed.", error=str(exc))
+
+    def _ensure_normalizers_loaded(self) -> dict[str, object]:
+        with self._lock:
+            if self._normalizers is not None:
+                return self._normalizers
+
+            from tn.chinese.normalizer import Normalizer as ZhNormalizer
+            from tn.english.normalizer import Normalizer as EnNormalizer
+
+            logging.getLogger().setLevel(logging.INFO)
+            self._normalizers = {
+                "zh": ZhNormalizer(overwrite_cache=False),
+                "en": EnNormalizer(overwrite_cache=False),
+            }
+            return self._normalizers
+
+    def normalize(self, *, text: str, prompt_text: str, language: str) -> tuple[str, str]:
+        snapshot = self.ensure_ready()
+        if not snapshot.ready:
+            raise RuntimeError(snapshot.error or snapshot.message)
+
+        with self._normalize_lock:
+            normalizers = self._ensure_normalizers_loaded()
+            if language not in normalizers:
+                raise ValueError(f"Unsupported text normalization language: {language}")
+            normalizer = normalizers[language]
+            normalized_text = normalizer.normalize(text) if text else ""
+            normalized_prompt_text = normalizer.normalize(prompt_text) if prompt_text else ""
+            return normalized_text, normalized_prompt_text
+
+
+def resolve_text_normalization_language(*, text: str, voice: str) -> str:
+    if re.search(r"[\u3400-\u9fff]", text):
+        return "zh"
+    if re.search(r"[A-Za-z]", text):
+        return "en"
+    if voice in ENGLISH_VOICES:
+        return "en"
+    return "zh"
+
+
+def prepare_tts_request_texts(
+    *,
+    text: str,
+    prompt_text: str,
+    voice: str,
+    enable_wetext: bool,
+    text_normalizer_manager: WeTextProcessingManager | None,
+) -> dict[str, object]:
+    raw_text = str(text or "")
+    raw_prompt_text = str(prompt_text or "")
+
+    normalization_language = ""
+    intermediate_text = raw_text
+    intermediate_prompt_text = raw_prompt_text
+
+    if enable_wetext:
+        if text_normalizer_manager is None:
+            raise RuntimeError("WeTextProcessing manager is unavailable.")
+        normalization_language = resolve_text_normalization_language(text=raw_text, voice=voice)
+        intermediate_text, intermediate_prompt_text = text_normalizer_manager.normalize(
+            text=raw_text,
+            prompt_text=raw_prompt_text,
+            language=normalization_language,
+        )
+        if intermediate_text != raw_text:
+            logging.info(
+                "normalized text chars_before=%d chars_after=%d stage=wetext language=%s",
+                len(raw_text),
+                len(intermediate_text),
+                normalization_language,
+            )
+        if raw_prompt_text and intermediate_prompt_text != raw_prompt_text:
+            logging.info(
+                "normalized prompt_text chars_before=%d chars_after=%d stage=wetext language=%s",
+                len(raw_prompt_text),
+                len(intermediate_prompt_text),
+                normalization_language,
+            )
+
+    final_text = normalize_tts_text(intermediate_text)
+    final_prompt_text = normalize_tts_text(intermediate_prompt_text) if intermediate_prompt_text else ""
+
+    if final_text != intermediate_text:
+        logging.info(
+            "normalized text chars_before=%d chars_after=%d stage=robust_final",
+            len(intermediate_text),
+            len(final_text),
+        )
+    if intermediate_prompt_text and final_prompt_text != intermediate_prompt_text:
+        logging.info(
+            "normalized prompt_text chars_before=%d chars_after=%d stage=robust_final",
+            len(intermediate_prompt_text),
+            len(final_prompt_text),
+        )
+
+    return {
+        "text": final_text,
+        "prompt_text": final_prompt_text,
+        "normalized_text": final_text,
+        "normalized_prompt_text": final_prompt_text,
+        "normalization_method": (f"wetext:{normalization_language}+robust" if enable_wetext else "robust"),
+        "text_normalization_language": normalization_language,
+        "text_normalization_enabled": bool(enable_wetext),
+    }

tts_robust_normalizer_single_script.py

@@ -0,0 +1,366 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+TTS 输入鲁棒性正则化器（非语义 TN）
+
+目标
+----
+1. 只做“鲁棒性清洗”，不做数字/单位/日期/金额等语义展开。
+2. 优先保护高风险 token，避免把 `.map`、`app.js.map`、`v2.3.1`、URL、Email、@mention、#hashtag 清坏。
+3. `[]` / `{}` / `【】` / `〖〗` / `『』` / `「」` 统一转成双引号包裹内容。
+4. 对结构性符号做“替换而非删除”：
+   - `【】 / 〖〗 / 『』 / 「」` 转成双引号包裹内容。
+   - `《》` 只在“独立标题/栏目名”场景拆开；嵌入式标题保持不变。
+   - `—— / -- / ——...` 转成句边界。
+5. 对社交平台常见噪声做弱归一化：
+   - `...... / ……` -> `。`
+   - `？？？！！！` -> `？！`
+   - `！！！` -> `！`
+6. 空格按脚本类型处理：
+   - 西文片段内部：连续空格压缩为 1 个。
+   - 汉字 / 日文假名片段内部：删除空格。
+   - 汉字 / 日文假名 与“拉丁字母类 token / 受保护 token”相邻：保留或补 1 个空格。
+   - 汉字 / 日文假名 与纯数字相邻：不强行补空格。
+7. 轻量处理 Markdown 与换行：
+   - `[text](url)` -> `text url`
+   - 去掉标题 `#`、引用 `>`、列表前缀
+   - 换行转句边界 `。`
+
+非目标
+------
+1. 不决定“应该怎么读”。
+2. 不做 HTML/SSML/语义标签解释。
+"""
+
+from __future__ import annotations
+
+import re
+import unicodedata
+
+
+# ---------------------------
+# 基础常量与正则
+# ---------------------------
+
+# 不依赖空格分词的脚本：汉字 + 日文假名
+_CJK_CHARS = r"\u3400-\u4dbf\u4e00-\u9fff\u3040-\u30ff"
+_CJK = f"[{_CJK_CHARS}]"
+
+# 保护占位符
+_PROT = r"___PROT\d+___"
+
+# 需要保护的高风险 token
+_URL_RE = re.compile(r"https?://[^\s\u3000，。！？；、）】》〉」』]+")
+_EMAIL_RE = re.compile(r"(?<![\w.+-])[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}(?![\w.-])")
+_MENTION_RE = re.compile(r"(?<![A-Za-z0-9_])@[A-Za-z0-9_]{1,32}")
+_REDDIT_RE = re.compile(r"(?<![A-Za-z0-9_])(?:u|r)/[A-Za-z0-9_]+")
+_HASHTAG_RE = re.compile(r"(?<![A-Za-z0-9_])#(?!\s)[^\s#]+")
+
+# `.map` / `.env` / `.gitignore`
+_DOT_TOKEN_RE = re.compile(r"(?<![A-Za-z0-9_])\.(?=[A-Za-z0-9._-]*[A-Za-z0-9])[A-Za-z0-9._-]+")
+
+# `app.js.map` / `index.d.ts` / `v2.3.1` / `foo/bar-baz.py` 等
+_FILELIKE_RE = re.compile(
+    r"(?<![A-Za-z0-9_])"
+    r"(?=[A-Za-z0-9._/+:-]*[A-Za-z])"
+    r"(?=[A-Za-z0-9._/+:-]*[._/+:-])"
+    r"[A-Za-z0-9][A-Za-z0-9._/+:-]*"
+    r"(?![A-Za-z0-9_])"
+)
+
+# 参与“中英混排边界补空格”的 token：必须至少含 1 个拉丁字母，或本身就是受保护 token
+_LATINISH = rf"(?:{_PROT}|(?=[A-Za-z0-9._/+:-]*[A-Za-z])[A-Za-z0-9][A-Za-z0-9._/+:-]*)"
+
+# 零宽字符
+_ZERO_WIDTH_RE = re.compile(r"[\u200b-\u200d\ufeff]")
+_TRAILING_CLOSERS = set('"\')]}）】》〉」』”’')
+
+
+# ---------------------------
+# 主函数
+# ---------------------------
+
+def normalize_tts_text(text: str) -> str:
+    """对 TTS 输入做鲁棒性正则化。"""
+    text = _base_cleanup(text)
+    text = _normalize_markdown_and_lines(text)
+    text, protected = _protect_spans(text)
+
+    text = _normalize_spaces(text)
+    text = _normalize_structural_punctuation(text)
+    text = _normalize_repeated_punctuation(text)
+    text = _normalize_spaces(text)
+
+    text = _restore_spans(text, protected)
+    text = text.strip()
+    return _ensure_terminal_punctuation_by_line(text)
+
+
+# ---------------------------
+# 具体规则
+# ---------------------------
+
+def _base_cleanup(text: str) -> str:
+    text = text.replace("\r\n", "\n").replace("\r", "\n").replace("\u3000", " ")
+    text = _ZERO_WIDTH_RE.sub("", text)
+
+    cleaned = []
+    for ch in text:
+        cat = unicodedata.category(ch)
+        if ch in "\n\t " or not cat.startswith("C"):
+            cleaned.append(ch)
+    return "".join(cleaned)
+
+
+def _normalize_markdown_and_lines(text: str) -> str:
+    # Markdown 链接：[text](url) -> text url
+    text = re.sub(r"\[([^\[\]]+?)\]\((https?://[^)\s]+)\)", r"\1 \2", text)
+
+    lines = []
+    for raw in text.splitlines():
+        line = raw.strip()
+        if not line:
+            continue
+
+        line = re.sub(r"^#{1,6}\s+", "", line)   # 标题
+        line = re.sub(r"^>\s+", "", line)        # 引用
+        line = re.sub(r"^[-*+]\s+", "", line)    # 无序列表
+        line = re.sub(r"^\d+[.)]\s+", "", line)  # 有序列表
+        lines.append(line)
+
+    if not lines:
+        return ""
+
+    merged: list[str] = [lines[0]]
+    for line in lines[1:]:
+        previous = merged[-1]
+        merged[-1] = _ensure_terminal_punctuation(previous)
+        merged.append(line)
+    return "".join(merged)
+
+
+def _protect_spans(text: str) -> tuple[str, list[str]]:
+    protected: list[str] = []
+
+    def repl(match: re.Match[str]) -> str:
+        idx = len(protected)
+        protected.append(match.group(0))
+        return f"___PROT{idx}___"
+
+    for pattern in (
+        _URL_RE,
+        _EMAIL_RE,
+        _MENTION_RE,
+        _REDDIT_RE,
+        _HASHTAG_RE,
+        _DOT_TOKEN_RE,
+        _FILELIKE_RE,
+    ):
+        text = pattern.sub(repl, text)
+
+    return text, protected
+
+
+def _restore_spans(text: str, protected: list[str]) -> str:
+    for idx, original in enumerate(protected):
+        text = text.replace(f"___PROT{idx}___", original)
+    return text
+
+
+def _normalize_spaces(text: str) -> str:
+    # 统一空白
+    text = re.sub(r"[ \t\r\f\v]+", " ", text)
+
+    # 汉字 / 日文片段内部：删除空格
+    text = re.sub(rf"({_CJK})\s+(?={_CJK})", r"\1", text)
+
+    # 汉字 / 日文 与纯数字之间：删除空格（不强行保留）
+    text = re.sub(rf"({_CJK})\s+(?=\d)", r"\1", text)
+    text = re.sub(rf"(\d)\s+(?={_CJK})", r"\1", text)
+
+    # 汉字 / 日文 与拉丁字母类 token / protected token 相邻：保留或补 1 个空格
+    text = re.sub(rf"({_CJK})(?=({_LATINISH}))", r"\1 ", text)
+    text = re.sub(rf"(({_LATINISH}))(?={_CJK})", r"\1 ", text)
+
+    # 再压一遍连续空格
+    text = re.sub(r" {2,}", " ", text)
+
+    # 中文标点前后不保留空格
+    text = re.sub(r"\s+([，。！？；：、”’」』】）》])", r"\1", text)
+    text = re.sub(r"([（【「『《“‘])\s+", r"\1", text)
+    text = re.sub(r"([，。！？；：、])\s*", r"\1", text)
+
+    # ASCII 标点前不留空格；后面的英文空格不强改
+    text = re.sub(r"\s+([,.;!?])", r"\1", text)
+
+    return re.sub(r" {2,}", " ", text).strip()
+
+
+def _normalize_structural_punctuation(text: str) -> str:
+    # 各类结构性括号：统一转成双引号包裹内容
+    text = re.sub(r"\[\s*([^\[\]]+?)\s*\]", r'"\1"', text)
+    text = re.sub(r"\{\s*([^{}]+?)\s*\}", r'"\1"', text)
+    text = re.sub(r"[【〖『「]\s*([^】〗』」]+?)\s*[】〗』」]", r'"\1"', text)
+
+    # 《》只处理独立标题，不处理嵌入式标题
+    # 例：重磅。《新品发布》——现在开始！ -> 重磅。新品发布。现在开始！
+    text = re.sub(
+        r"(^|[。！？!?；;]\s*)《([^》]+)》(?=\s*(?:___PROT\d+___|[—–―-]{2,}|$|[。！？!?；;，,]))",
+        r"\1\2",
+        text,
+    )
+
+    # 长破折号 / 多连字符：转句边界
+    text = re.sub(r"\s*(?:—|–|―|-){2,}\s*", "。", text)
+
+    return text
+
+
+def _normalize_repeated_punctuation(text: str) -> str:
+    # 省略号 / 连续句点
+    text = re.sub(r"(?:\.{3,}|…{2,}|……+)", "。", text)
+
+    # 同类重复标点
+    text = re.sub(r"[。．]{2,}", "。", text)
+    text = re.sub(r"[，,]{2,}", "，", text)
+    text = re.sub(r"[!！]{2,}", "！", text)
+    text = re.sub(r"[?？]{2,}", "？", text)
+
+    # 混合问叹号：收敛到 ？！
+    def _mixed_qe(match: re.Match[str]) -> str:
+        s = match.group(0)
+        has_q = any(ch in s for ch in "?？")
+        has_e = any(ch in s for ch in "!！")
+        if has_q and has_e:
+            return "？！"
+        return "？" if has_q else "！"
+
+    text = re.sub(r"[!?！？]{2,}", _mixed_qe, text)
+    return text
+
+
+def _ensure_terminal_punctuation(text: str) -> str:
+    if not text:
+        return text
+
+    index = len(text) - 1
+    while index >= 0 and text[index].isspace():
+        index -= 1
+    while index >= 0 and text[index] in _TRAILING_CLOSERS:
+        index -= 1
+
+    if index >= 0 and unicodedata.category(text[index]).startswith("P"):
+        return text
+    return text + "。"
+
+
+def _ensure_terminal_punctuation_by_line(text: str) -> str:
+    if not text:
+        return text
+    lines = text.split("\n")
+    normalized_lines = [_ensure_terminal_punctuation(line.strip()) if line.strip() else "" for line in lines]
+    return "\n".join(normalized_lines).strip()
+
+
+# ---------------------------
+# 测试
+# ---------------------------
+
+TEST_CASES = [
+    # 1) .map / dot-leading token / 文件名 / 版本号
+    (
+        "dot_map_sentence",
+        "2026 年 3 月 31 日，安全研究员 Chaofan Shou (@Fried_rice) 发现 Anthropic 的 npm 包中暴露了 .map 文件，",
+        "2026年3月31日，安全研究员 Chaofan Shou (@Fried_rice) 发现 Anthropic 的 npm 包中暴露了 .map 文件，",
+    ),
+    ("dot_tokens", "别把 .env、.npmrc、.gitignore 提交上去。", "别把 .env、.npmrc、.gitignore 提交上去。"),
+    ("file_names", "请检查 bundle.min.js、package.json 和 processing_moss_tts.py。", "请检查 bundle.min.js、package.json 和 processing_moss_tts.py。"),
+    ("index_d_ts", "index.d.ts 里也有同样的问题。", "index.d.ts 里也有同样的问题。"),
+    ("version_build", "Bug 的讨论可以精确到 v2.3.1 (Build 15)。", "Bug 的讨论可以精确到 v2.3.1 (Build 15)。"),
+    ("version_rc", "3.0.0-rc.1 还不能上生产。", "3.0.0-rc.1 还不能上生产。"),
+    ("jar_name", "fabric-api-0.91.3+1.20.2.jar 需要单独下载。", "fabric-api-0.91.3+1.20.2.jar 需要单独下载。"),
+
+    # 2) URL / Email / mention / hashtag / Reddit
+    ("url", "仓库地址是 https://github.com/instructkr/claude-code", "仓库地址是 https://github.com/instructkr/claude-code。"),
+    ("email", "联系邮箱：ops+tts@example.ai", "联系邮箱：ops+tts@example.ai。"),
+    ("mention", "@Fried_rice 说这是 source map 暴露。", "@Fried_rice 说这是 source map 暴露。"),
+    ("reddit", "去 r/singularity 看讨论。", "去 r/singularity 看讨论。"),
+    ("hashtag_chain", "#张雪峰#张雪峰[话题]#张雪峰事件", "#张雪峰#张雪峰[话题]#张雪峰事件。"),
+    ("mention_hashtag_boundary", "关注@biscuit0228_并转发#thetime_tbs", "关注 @biscuit0228_ 并转发 #thetime_tbs。"),
+
+    # 3) bracket / 控制 token：统一转成双引号
+    ("speaker_bracket", "[S1]你好。[S2]收到。", '"S1"你好。"S2"收到。'),
+    ("event_bracket", "请模仿 {whisper} 的语气说“别出声”。", '请模仿 "whisper" 的语气说“别出声”。'),
+    ("order_bracket", "订单号：[AB-1234-XYZ]", '订单号："AB-1234-XYZ"。'),
+
+    # 4) 结构性符号：转成双引号或句边界，而不是直接删除
+    ("struct_headline", "〖重磅〗《新品发布》——现在开始！", '"重磅"《新品发布》。现在开始！'),
+    ("struct_notice", "【公告】今天 20:00 维护——预计 30 分钟。", '"公告"今天20:00维护。预计30分钟。'),
+    ("struct_quote_chain", "『特别提醒』「不要外传」", '"特别提醒""不要外传"。'),
+    ("struct_embedded_quote", "他说【重要通知】明天发布。", '他说"重要通知"明天发布。'),
+
+    # 5) 嵌入式标题：保留
+    ("embedded_title", "我喜欢《哈姆雷特》这本书。", "我喜欢《哈姆雷特》这本书。"),
+
+    # 6) 重复标点 / 社交噪声
+    ("noise_qe", "真的假的？？？！！！", "真的假的？！"),
+    ("noise_ellipsis", "这个包把 app.js.map 也发上去了......太离谱了！！！", "这个包把 app.js.map 也发上去了。太离谱了！"),
+    ("noise_ellipsis_cn", "【系统提示】请模仿{sad}低沉语气，说“今天下雨了……”", '"系统提示"请模仿"sad"低沉语气，说“今天下雨了。”'),
+
+    # 7) 空格规则：英文压缩、中文删除、中英混排保留边界
+    ("english_spaces", "This   is   a   test.", "This is a test."),
+    ("chinese_spaces", "这 是　一 段  含有多种空白的文本。", "这是一段含有多种空白的文本。"),
+    ("mixed_spaces_1", "这是Anthropic的npm包", "这是 Anthropic 的 npm 包。"),
+    ("mixed_spaces_2", "今天update到v2.3.1了", "今天 update 到 v2.3.1 了。"),
+    ("mixed_spaces_3", "处理app.js.map文件", "处理 app.js.map 文件。"),
+
+    # 8) Markdown / 列表 / 换行
+    ("markdown_link", "详情见 [release note](https://github.com/example/release)", "详情见 release note https://github.com/example/release。"),
+    ("markdown_heading", "# I made a free open source app to help with markdown files", "I made a free open source app to help with markdown files。"),
+    ("list_lines", "- 修复 .map 泄露\n- 发布 v2.3.1", "修复 .map 泄露。发布 v2.3.1。"),
+    ("numbered_lines", "1. 安装依赖\n2. 运行测试\n3. 发布 v2.3.1", "安装依赖。运行测试。发布 v2.3.1。"),
+    ("newlines", "第一行\n第二行\n第三行", "第一行。第二行。第三行。"),
+
+    # 9) 句末补标点
+    ("terminal_punct_plain", "今天发布", "今天发布。"),
+    ("terminal_punct_quoted", '他说"你好"', '他说"你好"。'),
+    ("terminal_punct_existing", "今天发布。", "今天发布。"),
+    ("terminal_punct_newlines", "第一行\n第二行。", "第一行。第二行。"),
+    ("terminal_punct_blank_lines", "第一行\n\n第二行", "第一行。第二行。"),
+
+    # 10) 零宽字符 / 幂等性
+    ("zero_width_url", "详见 https://x.com/\u200bSafety", "详见 https://x.com/Safety。"),
+]
+
+
+def run_tests(verbose: bool = True) -> None:
+    failed = []
+
+    for name, text, expected in TEST_CASES:
+        actual = normalize_tts_text(text)
+        if actual != expected:
+            failed.append((name, text, expected, actual))
+            continue
+
+        # 幂等性：第二次归一化不应继续改动结果
+        second = normalize_tts_text(actual)
+        if second != actual:
+            failed.append((name + "_idempotence", actual, actual, second))
+
+    if failed:
+        lines = ["\nTEST FAILED:\n"]
+        for name, text, expected, actual in failed:
+            lines.append(f"[{name}]")
+            lines.append(f"input   : {text}")
+            lines.append(f"expected: {expected}")
+            lines.append(f"actual  : {actual}")
+            lines.append("")
+        raise AssertionError("\n".join(lines))
+
+    if verbose:
+        print(f"All {len(TEST_CASES)} tests passed.")
+
+
+if __name__ == "__main__":
+    run_tests()

weights/codec/.gitattributes

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

weights/codec/README.md

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

weights/codec/__init__.py

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

weights/codec/config.json

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

weights/codec/configuration_moss_audio_tokenizer.py

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

weights/codec/model-00001-of-00001.safetensors

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

weights/codec/model.safetensors.index.json

@@ -0,0 +1,382 @@
+{
+  "metadata": {
+    "total_parameters": 21969664,
+    "total_size": 87878656
+  },
+  "weight_map": {
+    "encoder.1.input_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.0.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.0.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.0.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.0.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.0.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.0.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.0.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.0.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.0.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.0.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.1.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.1.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.1.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.1.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.1.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.1.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.1.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.1.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.1.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.1.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.2.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.2.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.2.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.2.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.2.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.2.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.2.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.2.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.2.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.2.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.3.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.3.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.3.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.3.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.3.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.3.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.3.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.3.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.3.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.1.transformer.layers.3.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.1.output_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.input_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.0.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.0.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.0.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.0.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.0.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.0.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.0.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.0.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.0.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.0.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.1.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.1.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.1.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.1.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.1.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.1.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.1.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.1.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.1.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.3.transformer.layers.1.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.3.output_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.input_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.0.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.0.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.0.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.0.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.0.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.0.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.0.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.0.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.0.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.0.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.1.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.1.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.1.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.1.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.1.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.1.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.1.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.1.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.1.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.5.transformer.layers.1.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.5.output_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.input_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.0.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.0.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.0.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.0.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.0.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.0.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.0.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.0.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.0.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.0.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.1.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.1.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.1.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.1.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.1.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.1.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.1.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.1.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.1.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.1.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.2.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.2.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.2.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.2.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.2.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.2.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.2.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.2.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.2.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.2.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.3.norm1.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.3.norm1.bias": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.3.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.3.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.3.norm2.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.3.norm2.bias": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.3.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.3.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.3.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "encoder.7.transformer.layers.3.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "encoder.7.output_proj.weight": "model-00001-of-00001.safetensors",
+    "quantizer.input_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.input_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.input_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.output_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.output_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.output_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.0.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.0.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.0.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.0.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.0.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.0.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.0.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.1.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.1.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.1.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.1.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.1.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.1.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.1.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.2.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.2.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.2.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.2.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.2.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.2.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.2.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.3.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.3.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.3.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.3.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.3.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.3.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.3.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.4.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.4.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.4.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.4.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.4.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.4.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.4.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.5.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.5.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.5.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.5.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.5.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.5.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.5.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.6.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.6.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.6.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.6.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.6.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.6.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.6.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.7.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.7.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.7.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.7.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.7.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.7.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.7.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.8.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.8.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.8.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.8.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.8.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.8.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.8.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.9.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.9.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.9.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.9.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.9.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.9.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.9.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.10.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.10.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.10.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.10.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.10.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.10.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.10.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.11.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.11.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.11.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.11.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.11.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.11.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.11.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.12.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.12.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.12.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.12.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.12.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.12.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.12.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.13.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.13.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.13.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.13.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.13.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.13.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.13.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.14.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.14.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.14.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.14.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.14.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.14.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.14.codebook.weight": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.15.in_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.15.in_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.15.in_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.15.out_proj.bias": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.15.out_proj.parametrizations.weight.original0": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.15.out_proj.parametrizations.weight.original1": "model-00001-of-00001.safetensors",
+    "quantizer.quantizers.15.codebook.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.input_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.0.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.0.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.0.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.0.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.0.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.0.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.0.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.0.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.0.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.0.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.1.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.1.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.1.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.1.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.1.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.1.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.1.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.1.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.1.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.1.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.2.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.2.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.2.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.2.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.2.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.2.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.2.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.2.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.2.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.2.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.3.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.3.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.3.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.3.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.3.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.3.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.3.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.3.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.3.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.1.transformer.layers.3.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.1.output_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.input_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.0.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.0.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.0.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.0.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.0.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.0.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.0.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.0.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.0.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.0.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.1.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.1.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.1.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.1.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.1.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.1.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.1.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.1.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.1.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.3.transformer.layers.1.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.3.output_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.input_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.0.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.0.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.0.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.0.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.0.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.0.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.0.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.0.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.0.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.0.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.1.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.1.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.1.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.1.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.1.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.1.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.1.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.1.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.1.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.5.transformer.layers.1.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.5.output_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.input_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.0.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.0.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.0.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.0.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.0.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.0.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.0.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.0.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.0.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.0.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.1.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.1.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.1.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.1.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.1.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.1.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.1.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.1.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.1.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.1.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.2.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.2.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.2.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.2.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.2.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.2.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.2.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.2.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.2.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.2.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.3.norm1.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.3.norm1.bias": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.3.self_attn.in_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.3.self_attn.out_proj.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.3.norm2.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.3.norm2.bias": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.3.ffn.0.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.3.ffn.2.weight": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.3.layer_scale_1.scale": "model-00001-of-00001.safetensors",
+    "decoder.7.transformer.layers.3.layer_scale_2.scale": "model-00001-of-00001.safetensors",
+    "decoder.7.output_proj.weight": "model-00001-of-00001.safetensors"
+  }
+}

weights/tts/.gitattributes

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

weights/tts/README.md

@@ -0,0 +1,3 @@
+---
+license: apache-2.0
+---

weights/tts/__init__.py

@@ -0,0 +1,31 @@
+from .configuration_nanotts import NanoTTSConfig
+from .modeling_nanotts_global_local import (
+    NanoTTSGenerationOutput,
+    NanoTTSGlobalLocalForCausalLM,
+    NanoTTSOutput,
+)
+from .tokenization_nanotts_sentencepiece import NanoTTSSentencePieceTokenizer
+
+try:
+    NanoTTSConfig.register_for_auto_class()
+except Exception:
+    pass
+
+for auto_class_name in ("AutoModel", "AutoModelForCausalLM"):
+    try:
+        NanoTTSGlobalLocalForCausalLM.register_for_auto_class(auto_class_name)
+    except Exception:
+        pass
+
+try:
+    NanoTTSSentencePieceTokenizer.register_for_auto_class("AutoTokenizer")
+except Exception:
+    pass
+
+__all__ = [
+    "NanoTTSConfig",
+    "NanoTTSGlobalLocalForCausalLM",
+    "NanoTTSSentencePieceTokenizer",
+    "NanoTTSGenerationOutput",
+    "NanoTTSOutput",
+]

weights/tts/config.json

@@ -0,0 +1,197 @@
+{
+  "add_cross_attention": false,
+  "architectures": [
+    "NanoTTSGlobalLocalForCausalLM"
+  ],
+  "attn_implementation": "sdpa",
+  "audio_assistant_slot_token_id": 9,
+  "audio_codebook_sizes": [
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024
+  ],
+  "audio_end_token_id": 7,
+  "audio_pad_token_id": 1024,
+  "audio_start_token_id": 6,
+  "audio_tokenizer_pretrained_name_or_path": "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano",
+  "audio_tokenizer_sample_rate": 48000,
+  "audio_tokenizer_type": "moss-audio-tokenizer-nano",
+  "audio_user_slot_token_id": 8,
+  "audio_vocab_size": 1024,
+  "bad_words_ids": null,
+  "begin_suppress_tokens": null,
+  "bos_token_id": null,
+  "chunk_size_feed_forward": 0,
+  "cross_attention_hidden_size": null,
+  "decoder_start_token_id": null,
+  "diversity_penalty": 0.0,
+  "do_sample": false,
+  "dtype": "float32",
+  "early_stopping": false,
+  "encoder_no_repeat_ngram_size": 0,
+  "eos_token_id": null,
+  "exponential_decay_length_penalty": null,
+  "finetuning_task": null,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "gpt2_config": {
+    "_name_or_path": "",
+    "activation_function": "gelu_new",
+    "add_cross_attention": false,
+    "architectures": null,
+    "attn_pdrop": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 1,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dtype": null,
+    "early_stopping": false,
+    "embd_pdrop": 0.0,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 2,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_epsilon": 1e-05,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "gpt2",
+    "n_ctx": 32768,
+    "n_embd": 768,
+    "n_head": 12,
+    "n_inner": 3072,
+    "n_layer": 12,
+    "n_positions": 32768,
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 3,
+    "position_embedding_type": "rope",
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "reorder_and_upcast_attn": false,
+    "repetition_penalty": 1.0,
+    "resid_pdrop": 0.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rope_base": 10000.0,
+    "scale_attn_by_inverse_layer_idx": false,
+    "scale_attn_weights": true,
+    "sep_token_id": null,
+    "summary_activation": null,
+    "summary_first_dropout": 0.1,
+    "summary_proj_to_labels": true,
+    "summary_type": "cls_index",
+    "summary_use_proj": true,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torchscript": false,
+    "transformers_version": "4.57.1",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 16384
+  },
+  "hidden_size": 768,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "im_end_token_id": 5,
+  "im_start_token_id": 4,
+  "initializer_range": 0.02,
+  "is_decoder": false,
+  "is_encoder_decoder": false,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "length_penalty": 1.0,
+  "local_transformer_attn_implementation": "sdpa",
+  "local_transformer_layers": 1,
+  "max_length": 20,
+  "max_position_embeddings": 32768,
+  "min_length": 0,
+  "model_architecture": "global_local_transformer",
+  "model_type": "nano_tts",
+  "n_vq": 16,
+  "no_repeat_ngram_size": 0,
+  "num_beam_groups": 1,
+  "num_beams": 1,
+  "num_return_sequences": 1,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "output_scores": false,
+  "pad_token_id": 3,
+  "prefix": null,
+  "problem_type": null,
+  "pruned_heads": {},
+  "remove_invalid_values": false,
+  "repetition_penalty": 1.0,
+  "return_dict": true,
+  "return_dict_in_generate": false,
+  "sep_token_id": null,
+  "suppress_tokens": null,
+  "task_specific_params": null,
+  "temperature": 1.0,
+  "tf_legacy_loss": false,
+  "tie_encoder_decoder": false,
+  "tie_word_embeddings": true,
+  "tokenizer_class": "NanoTTSSentencePieceTokenizer",
+  "tokenizer_use_fast": false,
+  "top_k": 50,
+  "top_p": 1.0,
+  "torchscript": false,
+  "transformers_version": "4.57.1",
+  "typical_p": 1.0,
+  "use_bfloat16": false,
+  "vocab_size": 16384,
+  "auto_map": {
+    "AutoConfig": "configuration_nanotts.NanoTTSConfig",
+    "AutoModel": "modeling_nanotts_global_local.NanoTTSGlobalLocalForCausalLM",
+    "AutoModelForCausalLM": "modeling_nanotts_global_local.NanoTTSGlobalLocalForCausalLM"
+  }
+}

weights/tts/configuration_nanotts.py

@@ -0,0 +1,105 @@
+# coding=utf-8
+from typing import Any, Dict, Optional, Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
+
+
+class NanoTTSConfig(PretrainedConfig):
+    model_type = "nano_tts"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        gpt2_config: Optional[Union[GPT2Config, Dict[str, Any]]] = None,
+        n_vq: int = 8,
+        audio_vocab_size: Optional[int] = 1024,
+        audio_codebook_sizes: Optional[list[int]] = None,
+        audio_pad_token_id: int = 1024,
+        pad_token_id: int = 151643,
+        im_start_token_id: int = 151644,
+        im_end_token_id: int = 151645,
+        audio_start_token_id: int = 151652,
+        audio_end_token_id: int = 151653,
+        audio_user_slot_token_id: int = 151654,
+        audio_assistant_slot_token_id: int = 151656,
+        tokenizer_use_fast: bool = False,
+        audio_tokenizer_type: str = "moss-audio-tokenizer-nano",
+        audio_tokenizer_pretrained_name_or_path: Optional[str] = "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano",
+        audio_tokenizer_sample_rate: int = 48000,
+        attn_implementation: str = "flash_attention_2",
+        initializer_range: float = 0.02,
+        model_architecture: str = "global_local_transformer",
+        local_transformer_layers: int = 4,
+        local_transformer_attn_implementation: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        if isinstance(gpt2_config, dict):
+            self.gpt2_config = GPT2Config(**gpt2_config)
+        elif gpt2_config is None:
+            self.gpt2_config = GPT2Config()
+        else:
+            self.gpt2_config = gpt2_config
+
+        self.n_vq = int(n_vq)
+        if audio_codebook_sizes is None:
+            if audio_vocab_size is None:
+                raise ValueError("audio_vocab_size must be set when audio_codebook_sizes is not provided.")
+            resolved_audio_codebook_sizes = [int(audio_vocab_size)] * self.n_vq
+        else:
+            resolved_audio_codebook_sizes = [int(codebook_size) for codebook_size in audio_codebook_sizes]
+        if len(resolved_audio_codebook_sizes) != self.n_vq:
+            raise ValueError(
+                "audio_codebook_sizes must have length n_vq "
+                f"(expected {self.n_vq}, got {len(resolved_audio_codebook_sizes)})."
+            )
+        if any(codebook_size <= 0 for codebook_size in resolved_audio_codebook_sizes):
+            raise ValueError("audio_codebook_sizes must contain positive integers.")
+
+        max_audio_codebook_size = max(resolved_audio_codebook_sizes)
+        if audio_vocab_size is not None and int(audio_vocab_size) < max_audio_codebook_size:
+            raise ValueError(
+                "audio_vocab_size must be >= max(audio_codebook_sizes) "
+                f"(got {audio_vocab_size}, expected at least {max_audio_codebook_size})."
+            )
+
+        self.audio_codebook_sizes = resolved_audio_codebook_sizes
+        self.audio_vocab_size = (
+            max_audio_codebook_size if audio_vocab_size is None else int(audio_vocab_size)
+        )
+        self.audio_pad_token_id = int(audio_pad_token_id)
+        if self.audio_pad_token_id < max_audio_codebook_size:
+            raise ValueError(
+                "audio_pad_token_id must be >= max(audio_codebook_sizes) so pad stays outside every codebook "
+                f"(got {self.audio_pad_token_id}, max codebook size {max_audio_codebook_size})."
+            )
+        self.pad_token_id = pad_token_id
+        self.im_start_token_id = im_start_token_id
+        self.im_end_token_id = im_end_token_id
+        self.audio_start_token_id = audio_start_token_id
+        self.audio_end_token_id = audio_end_token_id
+        self.audio_user_slot_token_id = audio_user_slot_token_id
+        self.audio_assistant_slot_token_id = audio_assistant_slot_token_id
+        self.tokenizer_use_fast = tokenizer_use_fast
+        self.audio_tokenizer_type = audio_tokenizer_type
+        self.audio_tokenizer_pretrained_name_or_path = audio_tokenizer_pretrained_name_or_path
+        self.audio_tokenizer_sample_rate = audio_tokenizer_sample_rate
+        self.attn_implementation = attn_implementation
+        self.initializer_range = initializer_range
+        self.model_architecture = model_architecture
+        self.local_transformer_layers = local_transformer_layers
+        self.local_transformer_attn_implementation = (
+            attn_implementation
+            if local_transformer_attn_implementation is None
+            else local_transformer_attn_implementation
+        )
+        self.vocab_size = self.gpt2_config.vocab_size
+        self.hidden_size = self.gpt2_config.hidden_size
+        self.max_position_embeddings = self.gpt2_config.n_positions
+
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+    def to_dict(self) -> Dict[str, Any]:
+        output = super().to_dict()
+        output["gpt2_config"] = self.gpt2_config.to_dict()
+        return output

weights/tts/gpt2_decoder.py

@@ -0,0 +1,618 @@
+# coding=utf-8
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
+
+try:
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import pad_input, unpad_input
+
+    _FLASH_ATTN_AVAILABLE = True
+except Exception:
+    flash_attn_func = None
+    flash_attn_varlen_func = None
+    pad_input = None
+    unpad_input = None
+    _FLASH_ATTN_AVAILABLE = False
+
+
+@dataclass
+class PackedSequenceMetadata:
+    cu_seqlens: torch.Tensor
+    max_seqlen: int
+    indices: Optional[torch.Tensor] = None
+    batch_size: Optional[int] = None
+    seq_len: Optional[int] = None
+
+
+class NanoGPT2RotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, base: float = 10000.0) -> None:
+        super().__init__()
+        if dim % 2 != 0:
+            raise ValueError(f"RoPE head_dim must be even, got {dim}")
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(
+        self,
+        position_ids: torch.LongTensor,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if position_ids.ndim == 1:
+            position_ids = position_ids.unsqueeze(0)
+        freqs = torch.einsum("bs,d->bsd", position_ids.to(device=device, dtype=self.inv_freq.dtype), self.inv_freq)
+        cos = freqs.cos().repeat_interleave(2, dim=-1).unsqueeze(2).to(dtype=dtype)
+        sin = freqs.sin().repeat_interleave(2, dim=-1).unsqueeze(2).to(dtype=dtype)
+        return cos, sin
+
+
+def rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
+    even = hidden_states[..., ::2]
+    odd = hidden_states[..., 1::2]
+    return torch.stack((-odd, even), dim=-1).reshape_as(hidden_states)
+
+
+def apply_rotary_pos_emb(
+    hidden_states: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    return (hidden_states * cos) + (rotate_half(hidden_states) * sin)
+
+
+class NanoGPT2MLP(nn.Module):
+    def __init__(self, config: GPT2Config) -> None:
+        super().__init__()
+        hidden_size = int(config.hidden_size)
+        inner_size = int(config.n_inner or 4 * hidden_size)
+        self.fc_in = nn.Linear(hidden_size, inner_size)
+        self.fc_out = nn.Linear(inner_size, hidden_size)
+        self.act = ACT2FN[config.activation_function]
+        self.dropout = nn.Dropout(config.resid_pdrop)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc_in(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.fc_out(hidden_states)
+        return self.dropout(hidden_states)
+
+
+class NanoGPT2Attention(nn.Module):
+    def __init__(self, config: GPT2Config, layer_idx: int, attn_implementation: str) -> None:
+        super().__init__()
+        hidden_size = int(config.hidden_size)
+        num_heads = int(config.num_attention_heads)
+        if hidden_size % num_heads != 0:
+            raise ValueError(f"hidden_size={hidden_size} must be divisible by num_attention_heads={num_heads}")
+
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.embed_dim = hidden_size
+        self.layer_idx = layer_idx
+        self.attn_implementation = attn_implementation
+        self.attn_dropout = float(config.attn_pdrop)
+        self.resid_dropout = nn.Dropout(config.resid_pdrop)
+        self.scale_attn_weights = bool(getattr(config, "scale_attn_weights", True))
+        self.scale_attn_by_inverse_layer_idx = bool(getattr(config, "scale_attn_by_inverse_layer_idx", False))
+        self.position_embedding_type = str(getattr(config, "position_embedding_type", "absolute")).lower()
+        if self.position_embedding_type not in {"absolute", "rope"}:
+            raise ValueError(f"Unsupported position_embedding_type={self.position_embedding_type!r}")
+
+        self.c_attn = nn.Linear(hidden_size, 3 * hidden_size)
+        self.c_proj = nn.Linear(hidden_size, hidden_size)
+        self.rotary_emb = None
+        if self.position_embedding_type == "rope":
+            self.rotary_emb = NanoGPT2RotaryEmbedding(
+                self.head_dim,
+                base=float(getattr(config, "rope_base", 10000.0)),
+            )
+
+    def _split_heads(self, tensor: torch.Tensor) -> torch.Tensor:
+        if tensor.ndim == 3:
+            batch_size, seq_len, _ = tensor.shape
+            return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim)
+        if tensor.ndim == 2:
+            total_tokens, _ = tensor.shape
+            return tensor.view(total_tokens, self.num_heads, self.head_dim)
+        raise ValueError(f"Unsupported tensor rank for attention split: {tensor.ndim}")
+
+    def _merge_heads(self, tensor: torch.Tensor) -> torch.Tensor:
+        if tensor.ndim == 4:
+            batch_size, seq_len, _, _ = tensor.shape
+            return tensor.reshape(batch_size, seq_len, self.embed_dim)
+        if tensor.ndim == 3:
+            total_tokens, _, _ = tensor.shape
+            return tensor.reshape(total_tokens, self.embed_dim)
+        raise ValueError(f"Unsupported tensor rank for attention merge: {tensor.ndim}")
+
+    def _causal_attention_mask(
+        self,
+        attention_mask: Optional[torch.Tensor],
+        query_length: int,
+        key_length: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        query_positions = torch.arange(query_length, device=device, dtype=torch.long)
+        query_positions = query_positions + max(key_length - query_length, 0)
+        key_positions = torch.arange(key_length, device=device, dtype=torch.long)
+        causal = key_positions.unsqueeze(0) <= query_positions.unsqueeze(1)
+        causal = causal.unsqueeze(0).unsqueeze(0)
+        if attention_mask is None:
+            return causal
+        key_mask = attention_mask[:, None, None, :].to(dtype=torch.bool)
+        return causal & key_mask
+
+    def _eager_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+
+        scale = 1.0
+        if self.scale_attn_weights:
+            scale /= self.head_dim ** 0.5
+        if self.scale_attn_by_inverse_layer_idx:
+            scale /= float(self.layer_idx + 1)
+
+        scores = torch.matmul(query, key.transpose(-1, -2)) * scale
+        causal_mask = self._causal_attention_mask(
+            attention_mask=attention_mask,
+            query_length=query.shape[-2],
+            key_length=key.shape[-2],
+            device=query.device,
+        )
+        scores = scores.masked_fill(~causal_mask, torch.finfo(scores.dtype).min)
+        probs = torch.softmax(scores, dim=-1)
+        if self.training and self.attn_dropout > 0:
+            probs = torch.dropout(probs, self.attn_dropout, train=True)
+        output = torch.matmul(probs, value)
+        return output.transpose(1, 2).contiguous()
+
+    def _sdpa_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        mask = None
+        query_attention_mask = None
+        if attention_mask is not None:
+            query_length = query.shape[-2]
+            key_length = key.shape[-2]
+            mask = self._causal_attention_mask(
+                attention_mask=attention_mask,
+                query_length=query_length,
+                key_length=key_length,
+                device=query.device,
+            )
+            query_attention_mask = attention_mask[:, -query_length:].to(dtype=torch.bool, device=query.device)
+            if not bool(query_attention_mask.all()):
+                # SDPA can produce NaNs when a query row is fully masked. For padded query positions,
+                # keep a single aligned key visible, then zero the query output after attention.
+                mask = mask.expand(query.shape[0], -1, -1, -1).clone()
+                invalid_batch, invalid_query = torch.nonzero(~query_attention_mask, as_tuple=True)
+                aligned_key = invalid_query + max(key_length - query_length, 0)
+                mask[invalid_batch, :, invalid_query, aligned_key] = True
+        output = torch.nn.functional.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=mask,
+            dropout_p=self.attn_dropout if self.training else 0.0,
+            is_causal=mask is None,
+        )
+        if query_attention_mask is not None and not bool(query_attention_mask.all()):
+            output = output.masked_fill(~query_attention_mask[:, None, :, None], 0.0)
+        return output.transpose(1, 2).contiguous()
+
+    def _flash_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        packed_metadata: Optional[PackedSequenceMetadata],
+    ) -> torch.Tensor:
+        if not _FLASH_ATTN_AVAILABLE:
+            raise ImportError("flash_attn is not installed, but attn_implementation='flash_attention_2' was requested.")
+        if query.device.type != "cuda":
+            raise ValueError("flash_attention_2 requires CUDA tensors.")
+        if query.dtype not in (torch.float16, torch.bfloat16):
+            raise ValueError(
+                f"flash_attention_2 requires fp16/bf16 tensors, but received dtype={query.dtype}."
+            )
+
+        dropout_p = self.attn_dropout if self.training else 0.0
+        if packed_metadata is not None:
+            if packed_metadata.indices is not None:
+                query = query.reshape(-1, self.num_heads, self.head_dim).index_select(0, packed_metadata.indices)
+                key = key.reshape(-1, self.num_heads, self.head_dim).index_select(0, packed_metadata.indices)
+                value = value.reshape(-1, self.num_heads, self.head_dim).index_select(0, packed_metadata.indices)
+            output = flash_attn_varlen_func(
+                query,
+                key,
+                value,
+                packed_metadata.cu_seqlens,
+                packed_metadata.cu_seqlens,
+                packed_metadata.max_seqlen,
+                packed_metadata.max_seqlen,
+                dropout_p=dropout_p,
+                causal=True,
+            )
+            if packed_metadata.indices is None:
+                return output
+            return pad_input(
+                output,
+                packed_metadata.indices,
+                packed_metadata.batch_size,
+                packed_metadata.seq_len,
+            )
+
+        if attention_mask is None or bool(attention_mask.all()):
+            return flash_attn_func(
+                query,
+                key,
+                value,
+                dropout_p=dropout_p,
+                causal=True,
+            )
+
+        unpadded_query, indices, cu_seqlens, max_seqlen, _ = unpad_input(query, attention_mask)
+        unpadded_key, _, _, _, _ = unpad_input(key, attention_mask)
+        unpadded_value, _, _, _, _ = unpad_input(value, attention_mask)
+        output = flash_attn_varlen_func(
+            unpadded_query,
+            unpadded_key,
+            unpadded_value,
+            cu_seqlens,
+            cu_seqlens,
+            max_seqlen,
+            max_seqlen,
+            dropout_p=dropout_p,
+            causal=True,
+        )
+        return pad_input(output, indices, query.shape[0], query.shape[1])
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        packed_metadata: Optional[PackedSequenceMetadata] = None,
+        layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
+        qkv = self.c_attn(hidden_states)
+        query, key, value = qkv.split(self.embed_dim, dim=-1)
+        query = self._split_heads(query)
+        key = self._split_heads(key)
+        value = self._split_heads(value)
+
+        if self.rotary_emb is not None:
+            if position_ids is None:
+                raise ValueError("position_ids must be provided when position_embedding_type='rope'.")
+            cos, sin = self.rotary_emb(
+                position_ids.to(device=query.device),
+                device=query.device,
+                dtype=query.dtype,
+            )
+            query = apply_rotary_pos_emb(query, cos, sin)
+            key = apply_rotary_pos_emb(key, cos, sin)
+
+        if layer_past is not None:
+            past_key, past_value = layer_past
+            key = torch.cat([past_key.to(device=key.device, dtype=key.dtype), key], dim=1)
+            value = torch.cat([past_value.to(device=value.device, dtype=value.dtype), value], dim=1)
+
+        present = (key, value) if use_cache else None
+
+        if self.attn_implementation == "flash_attention_2" and layer_past is None:
+            attn_output = self._flash_attention(
+                query=query,
+                key=key,
+                value=value,
+                attention_mask=attention_mask,
+                packed_metadata=packed_metadata,
+            )
+        elif self.attn_implementation == "sdpa":
+            attn_output = self._sdpa_attention(
+                query=query,
+                key=key,
+                value=value,
+                attention_mask=attention_mask,
+            )
+        else:
+            attn_output = self._eager_attention(
+                query=query,
+                key=key,
+                value=value,
+                attention_mask=attention_mask,
+            )
+
+        attn_output = self._merge_heads(attn_output)
+        attn_output = self.c_proj(attn_output)
+        return self.resid_dropout(attn_output), present
+
+
+class NanoGPT2Block(nn.Module):
+    def __init__(self, config: GPT2Config, layer_idx: int, attn_implementation: str) -> None:
+        super().__init__()
+        hidden_size = int(config.hidden_size)
+        self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.attn = NanoGPT2Attention(config, layer_idx=layer_idx, attn_implementation=attn_implementation)
+        self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.mlp = NanoGPT2MLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        packed_metadata: Optional[PackedSequenceMetadata] = None,
+        layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
+        attn_output, present = self.attn(
+            self.ln_1(hidden_states),
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            packed_metadata=packed_metadata,
+            layer_past=layer_past,
+            use_cache=use_cache,
+        )
+        hidden_states = hidden_states + attn_output
+        hidden_states = hidden_states + self.mlp(self.ln_2(hidden_states))
+        return hidden_states, present
+
+
+class NanoGPT2Model(nn.Module):
+    def __init__(self, config: GPT2Config, attn_implementation: str = "eager") -> None:
+        super().__init__()
+        self.config = config
+        self.attn_implementation = attn_implementation
+        self.position_embedding_type = str(getattr(config, "position_embedding_type", "absolute")).lower()
+        if self.position_embedding_type not in {"absolute", "rope"}:
+            raise ValueError(f"Unsupported position_embedding_type={self.position_embedding_type!r}")
+        hidden_size = int(config.hidden_size)
+        self.wte = nn.Embedding(config.vocab_size, hidden_size)
+        self.wpe = nn.Embedding(config.n_positions, hidden_size) if self.position_embedding_type == "absolute" else nn.Identity()
+        self.drop = nn.Dropout(config.embd_pdrop)
+        self.h = nn.ModuleList(
+            [NanoGPT2Block(config, layer_idx=index, attn_implementation=attn_implementation) for index in range(config.n_layer)]
+        )
+        self.ln_f = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.gradient_checkpointing = False
+        self._reset_parameters()
+
+    def _reset_parameters(self) -> None:
+        init_std = float(self.config.initializer_range)
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, mean=0.0, std=init_std)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, mean=0.0, std=init_std)
+            elif isinstance(module, nn.LayerNorm):
+                nn.init.ones_(module.weight)
+                nn.init.zeros_(module.bias)
+
+    @staticmethod
+    def _normalize_num_sequences(
+        cu_seqlens: torch.Tensor,
+        num_sequences: Optional[torch.Tensor],
+        device: torch.device,
+    ) -> torch.Tensor:
+        if cu_seqlens.ndim == 1:
+            cu_seqlens = cu_seqlens.unsqueeze(0)
+        if num_sequences is None:
+            counts = []
+            for boundary in cu_seqlens:
+                diffs = boundary[1:] - boundary[:-1]
+                counts.append(int((diffs > 0).sum().item()))
+            return torch.tensor(counts, dtype=torch.int32, device=device)
+        if num_sequences.ndim == 0:
+            return num_sequences.unsqueeze(0)
+        return num_sequences
+
+    @staticmethod
+    def build_packed_position_ids(
+        attention_mask: Optional[torch.Tensor],
+        cu_seqlens: torch.Tensor,
+        num_sequences: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        if cu_seqlens.ndim == 1:
+            cu_seqlens = cu_seqlens.unsqueeze(0)
+        batch_size, seq_len = cu_seqlens.shape[0], cu_seqlens.shape[1] - 1
+        device = cu_seqlens.device
+        position_ids = torch.zeros((batch_size, seq_len), dtype=torch.long, device=device)
+        counts = NanoGPT2Model._normalize_num_sequences(cu_seqlens, num_sequences, device=device)
+        for batch_index in range(batch_size):
+            sequence_count = int(counts[batch_index].item())
+            boundaries = cu_seqlens[batch_index, : sequence_count + 1].tolist()
+            for start, end in zip(boundaries[:-1], boundaries[1:]):
+                start = int(start)
+                end = int(end)
+                if end > start:
+                    position_ids[batch_index, start:end] = torch.arange(end - start, device=device)
+        if attention_mask is not None:
+            position_ids = position_ids * attention_mask.to(dtype=position_ids.dtype)
+        return position_ids
+
+    @staticmethod
+    def build_packed_metadata(
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        num_sequences: Optional[torch.Tensor],
+    ) -> PackedSequenceMetadata:
+        if cu_seqlens.ndim == 1:
+            cu_seqlens = cu_seqlens.unsqueeze(0)
+        device = hidden_states.device
+        counts = NanoGPT2Model._normalize_num_sequences(cu_seqlens, num_sequences, device=device)
+        flat_indices = []
+        cumulative = [0]
+        max_seqlen = 0
+        seq_len = hidden_states.shape[1]
+
+        for batch_index in range(hidden_states.shape[0]):
+            sequence_count = int(counts[batch_index].item())
+            boundaries = cu_seqlens[batch_index, : sequence_count + 1].tolist()
+            for start, end in zip(boundaries[:-1], boundaries[1:]):
+                start = int(start)
+                end = int(end)
+                if end <= start:
+                    continue
+                segment_indices = batch_index * seq_len + torch.arange(start, end, device=device)
+                flat_indices.append(segment_indices)
+                cumulative.append(cumulative[-1] + (end - start))
+                max_seqlen = max(max_seqlen, end - start)
+
+        if not flat_indices:
+            raise ValueError("cu_seqlens did not describe any non-empty packed sequences.")
+
+        indices = torch.cat(flat_indices, dim=0)
+        return PackedSequenceMetadata(
+            cu_seqlens=torch.tensor(cumulative, dtype=torch.int32, device=device),
+            max_seqlen=max_seqlen,
+            indices=indices,
+            batch_size=hidden_states.shape[0],
+            seq_len=hidden_states.shape[1],
+        )
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[tuple[tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: bool = True,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        num_sequences: Optional[torch.Tensor] = None,
+    ) -> BaseModelOutputWithPast:
+        del input_ids, output_attentions
+
+        if inputs_embeds is None:
+            raise ValueError("inputs_embeds must be provided.")
+
+        use_cache = bool(use_cache)
+        if use_cache and cu_seqlens is not None:
+            raise ValueError("use_cache=True is not supported together with cu_seqlens packing.")
+
+        hidden_states = inputs_embeds
+        if attention_mask is None:
+            attention_mask = torch.ones(hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
+        else:
+            attention_mask = attention_mask.to(dtype=torch.bool, device=hidden_states.device)
+        query_attention_mask = attention_mask[:, -hidden_states.shape[1] :]
+
+        packed_metadata = None
+        if position_ids is None:
+            if cu_seqlens is not None:
+                position_ids = self.build_packed_position_ids(
+                    attention_mask=attention_mask,
+                    cu_seqlens=cu_seqlens.to(device=hidden_states.device),
+                    num_sequences=num_sequences.to(device=hidden_states.device) if num_sequences is not None else None,
+                )
+            elif attention_mask is not None:
+                position_ids = attention_mask.long().cumsum(dim=-1) - 1
+                position_ids = position_ids.masked_fill(~attention_mask, 0)
+                position_ids = position_ids[:, -hidden_states.shape[1] :]
+            else:
+                past_length = 0
+                if past_key_values is not None and len(past_key_values) > 0:
+                    past_length = past_key_values[0][0].shape[1]
+                position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device, dtype=torch.long)
+                position_ids = position_ids + past_length
+                position_ids = position_ids.unsqueeze(0).expand(hidden_states.shape[0], -1)
+
+        if cu_seqlens is not None and self.attn_implementation == "flash_attention_2":
+            packed_metadata = self.build_packed_metadata(
+                hidden_states=hidden_states,
+                cu_seqlens=cu_seqlens.to(device=hidden_states.device),
+                num_sequences=num_sequences.to(device=hidden_states.device) if num_sequences is not None else None,
+            )
+
+        if self.position_embedding_type == "absolute":
+            hidden_states = hidden_states + self.wpe(position_ids)
+        hidden_states = self.drop(hidden_states)
+        hidden_states = hidden_states * query_attention_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
+
+        all_hidden_states = () if output_hidden_states else None
+        presents = [] if use_cache else None
+        for layer_index, block in enumerate(self.h):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                if use_cache:
+                    raise ValueError("use_cache=True is not supported when gradient checkpointing is enabled during training.")
+
+                def custom_forward(*inputs):
+                    output, _ = block(
+                        inputs[0],
+                        attention_mask=inputs[1],
+                        position_ids=inputs[2],
+                        packed_metadata=packed_metadata,
+                        layer_past=None,
+                        use_cache=False,
+                    )
+                    return output
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    custom_forward,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    use_reentrant=False,
+                )
+                present = None
+            else:
+                hidden_states, present = block(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    packed_metadata=packed_metadata,
+                    layer_past=None if past_key_values is None else past_key_values[layer_index],
+                    use_cache=use_cache,
+                )
+            hidden_states = hidden_states * query_attention_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
+            if presents is not None:
+                presents.append(present)
+
+        hidden_states = self.ln_f(hidden_states)
+        hidden_states = hidden_states * query_attention_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return (hidden_states, tuple(presents) if presents is not None else None, all_hidden_states, None)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=tuple(presents) if presents is not None else None,
+            hidden_states=all_hidden_states,
+            attentions=None,
+        )

weights/tts/prompting.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from typing import List, Sequence
+
+from .configuration_nanotts import NanoTTSConfig
+
+
+USER_ROLE_PREFIX = "user\n"
+USER_TEMPLATE_REFERENCE_PREFIX = (
+    "<user_inst>\n"
+    "- Reference(s):\n"
+)
+USER_TEMPLATE_AFTER_REFERENCE = (
+    "\n- Instruction:\nNone\n"
+    "- Tokens:\nNone\n"
+    "- Quality:\nNone\n"
+    "- Sound Event:\nNone\n"
+    "- Ambient Sound:\nNone\n"
+    "- Language:\nNone\n"
+    "- Text:\n"
+)
+USER_TEMPLATE_PREFIX = USER_TEMPLATE_REFERENCE_PREFIX + "None" + USER_TEMPLATE_AFTER_REFERENCE
+USER_TEMPLATE_SUFFIX = "\n</user_inst>"
+ASSISTANT_TURN_PREFIX = "\n"
+ASSISTANT_ROLE_PREFIX = "assistant\n"
+
+
+def encode_text(tokenizer, text: str) -> List[int]:
+    try:
+        return list(tokenizer.encode(text, add_special_tokens=False))
+    except TypeError:
+        return list(tokenizer.encode(text))
+
+
+def decode_text(tokenizer, token_ids: Sequence[int]) -> str:
+    try:
+        return str(
+            tokenizer.decode(
+                list(token_ids),
+                skip_special_tokens=False,
+                clean_up_tokenization_spaces=False,
+            )
+        )
+    except TypeError:
+        try:
+            return str(tokenizer.decode(list(token_ids), skip_special_tokens=False))
+        except TypeError:
+            return str(tokenizer.decode(list(token_ids)))
+
+
+def build_user_prompt_prefix(tokenizer, config: NanoTTSConfig) -> List[int]:
+    return [config.im_start_token_id] + encode_text(tokenizer, USER_ROLE_PREFIX) + encode_text(
+        tokenizer,
+        USER_TEMPLATE_REFERENCE_PREFIX,
+    )
+
+
+def build_user_prompt_after_reference(tokenizer) -> List[int]:
+    return encode_text(tokenizer, USER_TEMPLATE_AFTER_REFERENCE)
+
+
+def build_assistant_prompt_prefix(tokenizer, config: NanoTTSConfig) -> List[int]:
+    return encode_text(tokenizer, USER_TEMPLATE_SUFFIX) + [config.im_end_token_id] + encode_text(
+        tokenizer,
+        ASSISTANT_TURN_PREFIX,
+    ) + [config.im_start_token_id] + encode_text(
+        tokenizer,
+        ASSISTANT_ROLE_PREFIX,
+    )
+
+
+def build_prompt_prefix(tokenizer, config: NanoTTSConfig) -> List[int]:
+    return (
+        build_user_prompt_prefix(tokenizer, config)
+        + encode_text(tokenizer, "None")
+        + build_user_prompt_after_reference(tokenizer)
+    )
+
+
+def build_prompt_suffix(tokenizer, config: NanoTTSConfig) -> List[int]:
+    return build_assistant_prompt_prefix(tokenizer, config)
+
+
+def build_prompt_token_ids(
+    tokenizer,
+    config: NanoTTSConfig,
+    text_token_ids: Sequence[int],
+) -> List[int]:
+    return build_prompt_prefix(tokenizer, config) + [int(token_id) for token_id in text_token_ids] + build_prompt_suffix(
+        tokenizer,
+        config,
+    )

weights/tts/pytorch_model.bin

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

weights/tts/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

weights/tts/tokenization_nanotts_sentencepiece.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+from typing import Any
+
+import sentencepiece as spm
+from transformers import PreTrainedTokenizer
+
+
+VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}
+
+
+class NanoTTSSentencePieceTokenizer(PreTrainedTokenizer):
+    vocab_files_names = VOCAB_FILES_NAMES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file: str,
+        unk_token: str = "<unk>",
+        bos_token: str = "<s>",
+        eos_token: str = "</s>",
+        pad_token: str = "<pad>",
+        sp_model_kwargs: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> None:
+        self.vocab_file = str(vocab_file)
+        self.sp_model_kwargs = {} if sp_model_kwargs is None else dict(sp_model_kwargs)
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(self.vocab_file)
+        super().__init__(
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        return int(self.sp_model.get_piece_size())
+
+    def get_vocab(self) -> dict[str, int]:
+        vocab = {self.sp_model.id_to_piece(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text: str) -> list[str]:
+        return list(self.sp_model.encode(text, out_type=str))
+
+    def _convert_token_to_id(self, token: str) -> int:
+        token_id = int(self.sp_model.piece_to_id(token))
+        return token_id
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return str(self.sp_model.id_to_piece(int(index)))
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        return str(self.sp_model.decode(tokens))
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+        save_dir = Path(save_directory)
+        save_dir.mkdir(parents=True, exist_ok=True)
+        out_name = "tokenizer.model" if filename_prefix is None else f"{filename_prefix}-tokenizer.model"
+        out_path = save_dir / out_name
+        if Path(self.vocab_file).resolve() != out_path.resolve():
+            shutil.copyfile(self.vocab_file, out_path)
+        return (str(out_path),)
+
+    def build_inputs_with_special_tokens(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+    ) -> list[int]:
+        if token_ids_1 is None:
+            return list(token_ids_0)
+        return list(token_ids_0) + list(token_ids_1)
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+        already_has_special_tokens: bool = False,
+    ) -> list[int]:
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0,
+                token_ids_1=token_ids_1,
+                already_has_special_tokens=True,
+            )
+        if token_ids_1 is None:
+            return [0] * len(token_ids_0)
+        return [0] * (len(token_ids_0) + len(token_ids_1))
+
+    def create_token_type_ids_from_sequences(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+    ) -> list[int]:
+        if token_ids_1 is None:
+            return [0] * len(token_ids_0)
+        return [0] * (len(token_ids_0) + len(token_ids_1))

weights/tts/tokenizer.model

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

weights/tts/tokenizer_config.json

@@ -0,0 +1,52 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [],
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_nanotts_sentencepiece.NanoTTSSentencePieceTokenizer",
+      null
+    ]
+  },
+  "backend": "custom",
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "model_max_length": 16384,
+  "pad_token": "<pad>",
+  "tokenizer_class": "NanoTTSSentencePieceTokenizer",
+  "unk_token": "<unk>"
+}