Document architecture, prompts, code, and full run statistics

.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
+assets/sft3/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+stats/training_steps.jsonl filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,124 @@
+---
+license: cc-by-4.0
+language:
+- en
+- de
+tags:
+- text-to-speech
+- voice-acting
+- qwen3
+- moss-audio-tokenizer-v2
+- audio-generation
+datasets:
+- laion/tts-scaling-ladder-de-en
+- laion/humaneness-voice-small-acting-eval
+model-index:
+- name: Humaneness Voice Small
+  results: []
+---
+
+# Humaneness Voice Small
+
+Humaneness Voice Small is an experimental English/German text-to-speech and voice-acting model. This repository contains **ten stage-local checkpoints from one continuous S1→S10 training run**, the training and inference source, run statistics, and a guide to the exact prompt surfaces. It is not a standard `AutoModel.from_pretrained()` checkpoint: it uses a Qwen3-0.6B semantic backbone, a MOSS-style local Talker, and the separately loaded [MOSS Audio Tokenizer v2](https://huggingface.co/OpenMOSS-Team/MOSS-Audio-Tokenizer-v2). Read the inference instructions before loading weights.
+
+The semantic backbone was **initialized from pretrained Qwen3-0.6B**, not from random weights. The **112,764,928-parameter Talker and bridge were initialized fresh** for this run. This corrects the loose description “trained from scratch” sometimes applied to earlier experiments. The Qwen hidden width is 1,024 and the local Talker width is 2,560. One semantic frame spans 80 ms; the Talker predicts **12 RVQ codebook indices per frame**, or 150 codebook indices per second. It does not predict 32 codebooks per frame. The codec reconstructs 48-kHz audio. The full state is approximately 715 million parameters, including the frozen, unused score conditioner; the model-only exports are BF16 PyTorch state dictionaries.
+
+## Choose a checkpoint
+
+The `checkpoints/S*/model_bf16.pt` files are the **lowest independent held-out validation-loss checkpoint within each stage**. The common validation set has 50 held-out prompts in each of two conditioning modes (100 rows/checkpoint), and all 49 complete candidates were swept. The acting benchmark did not select the checkpoints. A lower validation loss does not necessarily mean a more natural voice.
+
+| Stage | Ladder tier | Approx. measured presentation-hours¹ | Best global step | Held-out loss | Shared-cell acting reward² |
+|---|---|---:|---:|---:|---:|
+| S1 | x100 | 188,354 | 16,265 | 4.8054 | 0.4292 |
+| S2 | x50 | 116,838 | 26,284 | 4.7232 | 0.4489 |
+| S3 | tier7 | 72,314 | 32,337 | 4.7045 | **0.4628** |
+| S4 | tier6 | 57,419 | 37,182 | **4.6966** | 0.4581 |
+| S5 | tier5 | 28,716 | 38,183 | 4.6987 | 0.4543 |
+| S6 | tier4 | 14,361 | 40,935 | 4.7005 | 0.4539 |
+| S7 | tier3 | 5,750 | 41,466 | 4.7056 | 0.4540 |
+| S8 | tier2 | 2,877 | 41,739 | 4.7109 | 0.4586 |
+| S9 | tier1 | 1,439 | 41,878 | 4.7198 | 0.4540 |
+| S10 | tier0 | 578 | 41,936 | 4.7278 | 0.4572 |
+
+¹ These are **training presentations**, not unique source-audio hours. Multi-format prompts repeat some recordings. The ladder tiers are nested and paired with sidecar material; do not add these hours to estimate unique audio. ² Target-aware automatic reward on 1,651 shared valid acting cells per stage, not a human MOS. The S3–S10 paired difference is +0.00557 with a 95% challenge-bootstrap CI of −0.00701 to +0.01899: **the benchmark does not establish S3 as significantly better than S10**. S3 is a reasonable provisional default; inspect the same prompts in S3/S10 before choosing for a production voice. S4 has the best held-out validation loss. All underlying values and intervals are in `stats/acting_findings.json` and `stats/acting_analysis.json.gz`.
+
+The run used 32 JUPITER Booster nodes × 4 GPUs/node, batch 4,096, **one** 2,097-step linear warmup followed by one cosine decay across all 41,936 updates—**no stage LR restarts**. Peak LRs were 8e-5 for the semantic backbone and 2.4e-4 for the Talker, ending at 5% of peak. The source `code/plan_caption_production_32nodes.json` records the exact historical manifests and hashes; its absolute JUPITER paths are provenance, not portable defaults.
+
+## Prompt formats actually trained
+
+The outer model template is `<user_inst>` with `Reference(s)`, `Instruction`, `Tokens`, `Quality`, `Sound Event`, `Ambient Sound`, `Language`, and `Text` fields. The examples below show the **Instruction** field. `Text` separately repeats the SCRIPT after `SCRIPT:` for structured prompts, or otherwise follows the inherited packer behavior. `Tokens` is a **codec-frame budget**, not a literal transcript token count. Reference audio, when supplied, is codec-encoded through the dedicated reference channel. Do not paste a reference filename into the instruction.
+
+**Structured GENERAL/SCRIPT (C01).** Use when a specific delivery per sentence, explicit durations, pauses or vocal bursts matter:
+
+```text
+GENERAL: frightened, breath-held, intimate, increasingly panicked, adult, low-register
+SCRIPT:
+(quiet alarm, searching hesitation) [3.7 seconds duration] I locked the back door... didn't I?
+(fear rising, clipped self-correction) [4.2 seconds duration] I heard the latch click, but now it is standing open.
+(whispered panic, sharp inhale) [5.1 seconds duration] Wait—please don't move; there is someone breathing in the hallway.
+```
+
+The [prompting guide](https://laion-humaneness-voice-small-acting-eval-space.static.hf.space/index.html) and the historical training code distinguish `GENERAL:` global cues from the `SCRIPT:` line-by-line performance and quoted text of other formats. Sentence durations and pauses are optional training variations, not promises of exact timing. If a vocal burst is requested in the structured SCRIPT, put it at its intended position in the spoken sequence, outside the literal speech text.
+
+**Caption + exact transcript (C02–C10).** The description outside quotes is English; the spoken text inside `TRANSCRIPT` may be English or German. The model saw BUD-E/Body Whisper v1.0 and v1.1, Timbre Whisper, Voice Tagging Whisper and one validated Gemma-4-E4B paraphrase per enriched UID. The six possible Gemma presentation modes were *not* all generated for each UID.
+
+```text
+CAPTION: Use a close, low-register voice that begins in contained fear and tightens into whispered panic.
+TRANSCRIPT: "I locked the back door... didn't I? I heard the latch click, but now it is standing open."
+```
+
+For Timbre Whisper and Voice Tagging Whisper, the training loader split comma-separated tags from prose and presented **tags only, prose only, or both**, each one-third of that family's examples. Tag order was shuffled deterministically. The combined form is literal:
+
+```text
+CAPTION: Tags: increasingly_panicked, tense, close-mic, breath-held, low-register. Caption: Use a close, low-register voice that tightens into whispered panic.
+TRANSCRIPT: "I locked the back door... didn't I?"
+```
+
+The model also saw **transcript only with a distinct reference recording** (C11):
+
+```text
+TRANSCRIPT: "I locked the back door... didn't I?"
+```
+
+For ordinary intelligible speech, start with a concise `CAPTION:` and an exact quoted `TRANSCRIPT:`. For voice identity, supply a short, clean, distinct reference audio. The S3 benchmark's Timbre-prose surface (C05) had mean WER 0.0437 across its mixed reference assignments; S3 Gemma-style C10 had the highest *descriptive* caption-style reward (0.505), while S10's Voice-Tagging tags+prose C09 had WER 0.0473 and the highest caption-style emotion-band hit there. These are not causal prompt-format rankings: reference assignments and authored captions differ by condition. The structured C01 format is currently much less reliable for words (S3 mean WER ≈0.348; S10 ≈0.313 across reference modes), but it is the surface that explicitly expresses sentence timing and bursts. It has not shown reliable exact burst-onset control.
+
+## Evaluation: what the scores do and do not say
+
+The [public listening Space](https://huggingface.co/spaces/laion/humaneness-voice-small-acting-eval-space) and [direct site](https://laion-humaneness-voice-small-acting-eval-space.static.hf.space/index.html) contain **1,872 takes per checkpoint × ten checkpoints = 18,720 takes**: 50 English/German acting challenges, 11 prompt surfaces, three paired seeds, plus controlled reference and duration variations. Every stage has a [dedicated listening page](https://laion-humaneness-voice-small-acting-eval-space.static.hf.space/S3.html); replace `S3` in that URL with `S1`–`S10`. Separate [extreme S3 prompting](https://laion-humaneness-voice-small-acting-eval-space.static.hf.space/extreme_s3.html) and [S10 reference/CFG exploration](https://laion-humaneness-voice-small-acting-eval-space.static.hf.space/extreme_s10_cfg.html) are exploratory, **not** the frozen 18,720-take benchmark. The audio/score dataset is [here](https://huggingface.co/datasets/laion/humaneness-voice-small-acting-eval).
+
+| Stage | Mean WER ↓ | Genuineness ↑ | Burst blend ↑ | Requested-emotion hit ↑ | VoiceNet target hit ↑ | Requested burst type+onset F1 ↑ | Sentence duration MAE, s ↓ |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| S1 | 0.152 | 1.368 | 4.042 | 0.192 | 0.323 | 0.017 | 0.870 |
+| S2 | 0.128 | 1.533 | 3.981 | 0.199 | 0.338 | **0.053** | **0.850** |
+| S3 | 0.121 | **1.636** | 4.470 | 0.193 | 0.329 | 0.011 | 0.884 |
+| S4 | 0.120 | 1.615 | 4.354 | 0.191 | 0.333 | 0.030 | 0.934 |
+| S5 | **0.113** | 1.571 | 4.139 | 0.198 | 0.336 | 0.022 | 0.949 |
+| S6 | 0.122 | 1.609 | 4.577 | 0.201 | 0.326 | 0.031 | 0.924 |
+| S7 | 0.127 | 1.588 | **4.606** | 0.195 | 0.323 | 0.035 | 1.022 |
+| S8 | 0.115 | 1.591 | 4.593 | 0.192 | 0.326 | 0.015 | 1.007 |
+| S9 | 0.124 | 1.565 | 4.561 | 0.196 | 0.324 | 0.022 | 1.024 |
+| S10 | 0.121 | 1.561 | 4.556 | **0.207** | 0.320 | 0.030 | 0.966 |
+
+WER comes from Parakeet v3. Genuineness, burst blend, emotion and VoiceNet are learned proxy scores, **not** human listener ratings. Burst realization and duration error are restricted to eligible prompts; duration MAE excludes failed alignments, so consult alignment coverage before comparing it. A per-stage optimum in one column is not an all-round win. The balanced caption-condition comparison suggests reference audio often reduces *descriptive* WER, but the paired C01 reference/no-reference WER differences for S3 and S10 have confidence intervals crossing zero. No universal causal reference benefit is claimed. The benchmark's C01 prompts did carry parenthesized sentence directions, but lacked explicit pause tags, requested 0.7-second bursts only at sentence onset, and used durations for only half the challenges. These are limitations, not successful fine-grained-control results.
+
+## Files and inference
+
+- `checkpoints/S1` through `S10`: selected BF16 model-only state plus `metadata.json` with SHA-256 and original checkpoint contract. **Optimizer/RNG state is not uploaded.** Full 8.5-GB-per-checkpoint resume states remain in the original JUPITER scratch run; the public weights alone do not resume the optimizer exactly.
+- `code/continuous_train.py`, `code/train.sbatch`, `code/plan_caption_production_32nodes.json` and adjacent modules: the historical JUPITER source and training contract. The raw plan refers to local manifests and intentionally fails its hash checks if edited; regenerate plans and manifests for new data. `code/infer.py` is the standalone example for the exported weights. The code bundles MOSS export modules, tokenizer configuration and score schema; the audio codec is fetched separately from its upstream Hugging Face repository.
+- `stats/training_steps.jsonl`: every logged training-step scalar line across all four Slurm invocations. `stats/invocation_summaries.json`, `stats/run_configuration.json`, `stats/stage_winners.json`, `stats/validation_loss_full.tar.gz`, `stats/acting_findings.json`, and `stats/acting_analysis.json.gz` hold the remaining scalar summaries, all held-out checkpoint-sweep records, and the public acting aggregate. `stats/CONTENTS.json` gives hashes.
+
+Example (requires a CUDA environment with PyTorch, Transformers, NumPy, SoundFile, Torchaudio, Safetensors, and permission to download the separate codec):
+
+```bash
+python code/infer.py --stage S3 --language en --frames 75 \
+  --prompt $'CAPTION: Warm, softly amused, conversational.\nTRANSCRIPT: "I cannot believe we made it."' \
+  --text 'I cannot believe we made it.' --output example.wav
+```
+
+Add `--reference-wav reference.wav` for a distinct, clean recording under about 3 seconds. Do not use the target recording as its own reference. The example is a source-level starting point; the public listening Space is the already-audited playback path. This repository does **not** contain the large original SFT-3 model weights or original Qwen checkpoint separately, because the exported Humaneness Voice Small states include every trained model weight required for inference.
+
+## Training data, license and responsible use
+
+The core nested scaling ladder and associated annotation sidecars are documented at [laion/tts-scaling-ladder-de-en](https://huggingface.co/datasets/laion/tts-scaling-ladder-de-en); inspect each linked dataset's license/provenance before a derivative training run. The DPO corpora are a **separate future-adaptation resource**, not evidence that these base checkpoints were DPO-trained: [real-speech DPO](https://huggingface.co/datasets/laion/tts-realspeech-dpo-en-de), [voice-profile DPO](https://huggingface.co/datasets/laion/laion-voice-profiles-dpo), [CFG DPO](https://huggingface.co/datasets/laion/laion-voice-profiles-dpo-cfg), and [S5 pitch DPO](https://huggingface.co/datasets/laion/humaneness-voice-s5-pitch-dpo-v1). The pitch set's chosen/rejected quality gap is substantial; do not treat it as a clean voice-identity-only preference corpus.
+
+Model weights and original release materials are provided under **CC BY 4.0**; attribute LAION / Humaneness Voice Small and preserve attribution to the source projects. Qwen3 and MOSS Audio Tokenizer v2 are separate upstream dependencies under their own stated licenses. Respect speakers' rights, consent, provenance and applicable law. Synthetic voice acting can be mistaken for real speech; disclose generated audio and avoid impersonation or deceptive use. This is a research release with known shakiness, prompt sensitivity and imperfect word/timing/burst control, not a safety-certified production voice.

assets/qwen3/config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "max_position_embeddings": 40960,
+  "max_window_layers": 28,
+  "model_type": "qwen3",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

assets/score_schema.json

@@ -0,0 +1,1043 @@
+{
+  "version": 1,
+  "partition": "existing vn_time19 plus the other 38 VoiceNet axes",
+  "global_dimensions": [
+    {
+      "name": "vn:AGEV",
+      "label": "Voice Age",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_AGEV_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:GEND",
+      "label": "Perceived Gender",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_GEND_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:REGS",
+      "label": "Register",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_REGS_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:WARM",
+      "label": "Warmth",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_WARM_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:BRGT",
+      "label": "Brightness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_BRGT_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:FULL",
+      "label": "Fullness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_FULL_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:ROUG",
+      "label": "Roughness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_ROUG_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:METL",
+      "label": "Metallic Character",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_METL_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:HARM",
+      "label": "Harmonicity",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_HARM_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:R_CHST",
+      "label": "Chest Resonance",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_R_CHST_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:R_HEAD",
+      "label": "Head Resonance",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_R_HEAD_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:R_MASK",
+      "label": "Mask Resonance",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_R_MASK_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:R_MIXD",
+      "label": "Mixed Resonance",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_R_MIXD_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:R_NASL",
+      "label": "Nasal Resonance",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_R_NASL_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:R_ORAL",
+      "label": "Oral Resonance",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_R_ORAL_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:R_THRT",
+      "label": "Throat Resonance",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_R_THRT_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:RCQL",
+      "label": "Recording Quality",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_RCQL_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:BKGN",
+      "label": "Background Noise",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_BKGN_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:ESTH",
+      "label": "Esthetics",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_ESTH_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "genu",
+      "label": "Genuineness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "genuineness_0_6",
+      "head": "Genuineness regression"
+    },
+    {
+      "name": "blend",
+      "label": "Vocal-burst blend",
+      "normalization_bounds": [
+        0,
+        10
+      ],
+      "source_global_column": "blend_0_10",
+      "head": "Blend regression"
+    }
+  ],
+  "local_dimensions": [
+    {
+      "name": "vn:AROU",
+      "label": "Arousal",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_AROU_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:ARSH",
+      "label": "Arousal Shift",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_ARSH_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:ATCK",
+      "label": "Attack",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_ATCK_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:CHNK",
+      "label": "Chunking",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_CHNK_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:CLRT",
+      "label": "Articulation Clarity",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_CLRT_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:COGL",
+      "label": "Cognitive Load",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_COGL_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:DARC",
+      "label": "Dynamic Arc",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_DARC_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:DFLU",
+      "label": "Disfluency",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_DFLU_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:EMPH",
+      "label": "Emphasis",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_EMPH_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:EXPL",
+      "label": "Content Appropriateness (3-point Scale)",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_EXPL_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:FOCS",
+      "label": "Focus",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_FOCS_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:RANG",
+      "label": "Pitch Range",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_RANG_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:RESP",
+      "label": "Respiration",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_RESP_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:SMTH",
+      "label": "Smoothness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_SMTH_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:STNC",
+      "label": "Stance",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_STNC_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:STRU",
+      "label": "Structure",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_STRU_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_ASMR",
+      "label": "ASMR Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_ASMR_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_AUTH",
+      "label": "Authoritative Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_AUTH_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_CART",
+      "label": "Cartoonish Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_CART_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_CASU",
+      "label": "Casual Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_CASU_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_CONV",
+      "label": "Conversational Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_CONV_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_DRAM",
+      "label": "Dramatic Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_DRAM_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_FORM",
+      "label": "Formal Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_FORM_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_MONO",
+      "label": "Monologue Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_MONO_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_NARR",
+      "label": "Narrator Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_NARR_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_NEWS",
+      "label": "Newsreader Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_NEWS_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_PLAY",
+      "label": "Playful Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_PLAY_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_RANT",
+      "label": "Ranting/Angry Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_RANT_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_STRY",
+      "label": "Storytelling Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_STRY_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_TECH",
+      "label": "Teacher/Didactic Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_TECH_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:S_WHIS",
+      "label": "Whisper-Talk Style",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_S_WHIS_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:TEMP",
+      "label": "Tempo",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_TEMP_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:TENS",
+      "label": "Tension",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_TENS_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:VALN",
+      "label": "Valence",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_VALN_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:VALS",
+      "label": "Valence Shift",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_VALS_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:VFLX",
+      "label": "Velocity Flux",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_VFLX_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:VOLT",
+      "label": "Volatility",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_VOLT_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "vn:VULN",
+      "label": "Vulnerability",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "vn_VULN_reg",
+      "head": "VoiceNet regression"
+    },
+    {
+      "name": "emo:Affection",
+      "label": "Affection",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Affection"
+    },
+    {
+      "name": "emo:Amusement",
+      "label": "Amusement",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Amusement"
+    },
+    {
+      "name": "emo:Anger",
+      "label": "Anger",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Anger"
+    },
+    {
+      "name": "emo:Astonishment_Surprise",
+      "label": "Astonishment_Surprise",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Astonishment_Surprise"
+    },
+    {
+      "name": "emo:Awe",
+      "label": "Awe",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Awe"
+    },
+    {
+      "name": "emo:Bitterness",
+      "label": "Bitterness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Bitterness"
+    },
+    {
+      "name": "emo:Concentration",
+      "label": "Concentration",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Concentration"
+    },
+    {
+      "name": "emo:Confusion",
+      "label": "Confusion",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Confusion"
+    },
+    {
+      "name": "emo:Contemplation",
+      "label": "Contemplation",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Contemplation"
+    },
+    {
+      "name": "emo:Contempt",
+      "label": "Contempt",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Contempt"
+    },
+    {
+      "name": "emo:Contentment",
+      "label": "Contentment",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Contentment"
+    },
+    {
+      "name": "emo:Disappointment",
+      "label": "Disappointment",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Disappointment"
+    },
+    {
+      "name": "emo:Disgust",
+      "label": "Disgust",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Disgust"
+    },
+    {
+      "name": "emo:Distress",
+      "label": "Distress",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Distress"
+    },
+    {
+      "name": "emo:Doubt",
+      "label": "Doubt",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Doubt"
+    },
+    {
+      "name": "emo:Elation",
+      "label": "Elation",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Elation"
+    },
+    {
+      "name": "emo:Embarrassment",
+      "label": "Embarrassment",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Embarrassment"
+    },
+    {
+      "name": "emo:Emotional_Numbness",
+      "label": "Emotional_Numbness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Emotional_Numbness"
+    },
+    {
+      "name": "emo:Fatigue_Exhaustion",
+      "label": "Fatigue_Exhaustion",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Fatigue_Exhaustion"
+    },
+    {
+      "name": "emo:Fear",
+      "label": "Fear",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Fear"
+    },
+    {
+      "name": "emo:Helplessness",
+      "label": "Helplessness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Helplessness"
+    },
+    {
+      "name": "emo:Hope_Enthusiasm_Optimism",
+      "label": "Hope_Enthusiasm_Optimism",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Hope_Enthusiasm_Optimism"
+    },
+    {
+      "name": "emo:Impatience_and_Irritability",
+      "label": "Impatience_and_Irritability",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Impatience_and_Irritability"
+    },
+    {
+      "name": "emo:Infatuation",
+      "label": "Infatuation",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Infatuation"
+    },
+    {
+      "name": "emo:Interest",
+      "label": "Interest",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Interest"
+    },
+    {
+      "name": "emo:Intoxication_Altered_States_of_Consciousness",
+      "label": "Intoxication_Altered_States_of_Consciousness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Intoxication_Altered_States_of_Consciousness"
+    },
+    {
+      "name": "emo:Jealousy_and_Envy",
+      "label": "Jealousy_&_Envy",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Jealousy_and_Envy"
+    },
+    {
+      "name": "emo:Longing",
+      "label": "Longing",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Longing"
+    },
+    {
+      "name": "emo:Malevolence_Malice",
+      "label": "Malevolence_Malice",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Malevolence_Malice"
+    },
+    {
+      "name": "emo:Pain",
+      "label": "Pain",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Pain"
+    },
+    {
+      "name": "emo:Pleasure_Ecstasy",
+      "label": "Pleasure_Ecstasy",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Pleasure_Ecstasy"
+    },
+    {
+      "name": "emo:Pride",
+      "label": "Pride",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Pride"
+    },
+    {
+      "name": "emo:Relief",
+      "label": "Relief",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Relief"
+    },
+    {
+      "name": "emo:Sadness",
+      "label": "Sadness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Sadness"
+    },
+    {
+      "name": "emo:Sexual_Lust",
+      "label": "Sexual_Lust",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Sexual_Lust"
+    },
+    {
+      "name": "emo:Shame",
+      "label": "Shame",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Shame"
+    },
+    {
+      "name": "emo:Sourness",
+      "label": "Sourness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Sourness"
+    },
+    {
+      "name": "emo:Teasing",
+      "label": "Teasing",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Teasing"
+    },
+    {
+      "name": "emo:Thankfulness_Gratitude",
+      "label": "Thankfulness_Gratitude",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Thankfulness_Gratitude"
+    },
+    {
+      "name": "emo:Triumph",
+      "label": "Triumph",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "head": "EmoNet raw regression",
+      "source_global_column": "emo_Triumph"
+    },
+    {
+      "name": "genu",
+      "label": "Genuineness",
+      "normalization_bounds": [
+        0,
+        6
+      ],
+      "source_global_column": "genuineness_0_6",
+      "head": "Genuineness regression"
+    },
+    {
+      "name": "blend",
+      "label": "Vocal-burst blend",
+      "normalization_bounds": [
+        0,
+        10
+      ],
+      "source_global_column": "blend_0_10",
+      "head": "Blend regression"
+    }
+  ],
+  "global_tokens": 5,
+  "tokens_per_sentence": 10,
+  "missing_values": "zero normalized input with explicit mask=0; never represented as a measured zero",
+  "normalization": "Affine scale by nominal bounds; do not clip raw regressions outside the nominal range.",
+  "global_scope": "Measured whole-clip features, stability-oriented proxies, not proven speaker constants.",
+  "local_scope": "Measured sentence/span audio only, with pinned source/time/text alignment. Never copy whole-clip values to a sub-sentence.",
+  "quality_scope": "BKGN/RCQL/ESTH are the VoiceNet background/recording/aesthetic axes. They are not interchangeable with separate Empathic-Insight Quality heads.",
+  "sources": [
+    {
+      "path": "/e/home/jusers/schuhmann1/jupiter/layer_forensics/w2/actf_w2.py",
+      "sha256": "103f04aaed55cb4e0d8de6886f3c96813f18d1633431a1963611f79bbde3a0b2"
+    },
+    {
+      "path": "/e/home/jusers/schuhmann1/jupiter/layer_forensics/w2/README.md",
+      "sha256": "c5c70cc20bc1e2381af143bfaf3a7cc621179724309995aad8e3d715f6903fad"
+    },
+    {
+      "path": "/e/data1/datasets/staging/mmlaion/dramabox_top3_rebuild/hf/schema.json",
+      "sha256": "3ba48124f225e31374212707bf949ce4df8dcdf9043100d78674c4e51eb315e3"
+    }
+  ],
+  "historical_partition_lines": [
+    70,
+    73
+  ]
+}

assets/sft3/__init__.py

@@ -0,0 +1,9 @@
+from .configuration_moss_tts import MossTTSLocalConfig
+from .modeling_moss_tts import MossTTSLocalModel
+from .processing_moss_tts import MossTTSLocalProcessor
+
+__all__ = [
+    "MossTTSLocalConfig",
+    "MossTTSLocalModel",
+    "MossTTSLocalProcessor",
+]

assets/sft3/added_tokens.json

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

assets/sft3/config.json

@@ -0,0 +1,267 @@
+{
+  "architectures": [
+    "MossTTSLocalModel"
+  ],
+  "attn_implementation": "sdpa",
+  "audio_assistant_gen_slot_token_id": 151656,
+  "audio_assistant_slot_token_id": 151656,
+  "audio_codebook_sizes": [
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024
+  ],
+  "audio_end_token_id": 151670,
+  "audio_pad_code": 1024,
+  "audio_pad_token_id": 1024,
+  "audio_start_token_id": 151669,
+  "audio_tokenizer_name_or_path": "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2",
+  "audio_user_slot_token_id": 151654,
+  "audio_vocab_size": 1024,
+  "auto_map": {
+    "AutoConfig": "configuration_moss_tts.MossTTSLocalConfig",
+    "AutoModel": "modeling_moss_tts.MossTTSLocalModel",
+    "AutoProcessor": "processing_moss_tts.MossTTSLocalProcessor"
+  },
+  "dtype": "bfloat16",
+  "gpt2_config": {
+    "activation_function": "silu",
+    "add_cross_attention": false,
+    "attn_pdrop": 0.0,
+    "bos_token_id": null,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "embd_pdrop": 0.0,
+    "eos_token_id": 151645,
+    "finetuning_task": null,
+    "initializer_range": 0.02,
+    "is_decoder": false,
+    "layer_norm_epsilon": 1e-06,
+    "model_type": "gpt2",
+    "n_ctx": 10240,
+    "n_embd": 2560,
+    "n_head": 32,
+    "n_inner": 9728,
+    "n_layer": 1,
+    "n_positions": 10240,
+    "pad_token_id": null,
+    "position_embedding_type": "rope",
+    "prefix": null,
+    "pruned_heads": {},
+    "reorder_and_upcast_attn": false,
+    "resid_pdrop": 0.0,
+    "rope_base": 1000000.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,
+    "task_specific_params": null,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "torchscript": false,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 151936
+  },
+  "hidden_size": 2560,
+  "im_end_token_id": 151645,
+  "im_start_token_id": 151644,
+  "initializer_range": 0.02,
+  "language_config": {
+    "add_cross_attention": false,
+    "architectures": [
+      "Qwen3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "dtype": "bfloat16",
+    "eos_token_id": 151643,
+    "finetuning_task": null,
+    "gradient_checkpointing_use_reentrant": false,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 2560,
+    "initializer_range": 0.02,
+    "intermediate_size": 9728,
+    "is_decoder": false,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 32768,
+    "max_window_layers": 36,
+    "model_type": "qwen3",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "pad_token_id": 151643,
+    "prefix": null,
+    "pruned_heads": {},
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "rope_theta": 1000000,
+      "rope_type": "default"
+    },
+    "sep_token_id": null,
+    "sliding_window": null,
+    "task_specific_params": null,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "torchscript": false,
+    "use_bfloat16": false,
+    "use_cache": false,
+    "use_sliding_window": false,
+    "vocab_size": 151936
+  },
+  "local_hidden_size": 2560,
+  "local_text_head_mode": "binary",
+  "local_transformer_attn_implementation": "sdpa",
+  "local_transformer_layers": 1,
+  "model_type": "moss_tts_local",
+  "n_vq": 12,
+  "pad_token_id": 151643,
+  "processor_class": "MossTTSLocalProcessor",
+  "qwen3_config": {
+    "add_cross_attention": false,
+    "architectures": [
+      "Qwen3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "dtype": "bfloat16",
+    "eos_token_id": 151643,
+    "finetuning_task": null,
+    "gradient_checkpointing_use_reentrant": false,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 2560,
+    "initializer_range": 0.02,
+    "intermediate_size": 9728,
+    "is_decoder": false,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 32768,
+    "max_window_layers": 36,
+    "model_type": "qwen3",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "pad_token_id": 151643,
+    "prefix": null,
+    "pruned_heads": {},
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "rope_theta": 1000000,
+      "rope_type": "default"
+    },
+    "sep_token_id": null,
+    "sliding_window": null,
+    "task_specific_params": null,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "torchscript": false,
+    "use_bfloat16": false,
+    "use_cache": false,
+    "use_sliding_window": false,
+    "vocab_size": 151936
+  },
+  "sampling_rate": 48000,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.14.1",
+  "use_static_local_kv_cache": true,
+  "vocab_size": 151936
+}

assets/sft3/configuration_moss_tts.py

@@ -0,0 +1,158 @@
+# coding=utf-8
+"""Configuration for the MOSS-TTS-Local-Transformer-v1.5 release."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+
+
+SUPPORTED_ATTENTION_IMPLEMENTATIONS = {"flash_attention_2", "sdpa", "eager"}
+
+
+def _normalize_attention_implementation(value: Optional[str], default: str = "flash_attention_2") -> str:
+    normalized = str(value or default).strip().lower()
+    if normalized in {"flash", "flash_attn", "flash-attn", "flash_attention"}:
+        normalized = "flash_attention_2"
+    if normalized not in SUPPORTED_ATTENTION_IMPLEMENTATIONS:
+        raise ValueError(
+            "attn_implementation must be one of "
+            f"{sorted(SUPPORTED_ATTENTION_IMPLEMENTATIONS)}, got {value!r}."
+        )
+    return normalized
+
+
+class MossTTSLocalConfig(PretrainedConfig):
+    model_type = "moss_tts_local"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        qwen3_config: Optional[Union[Qwen3Config, Dict[str, Any]]] = None,
+        gpt2_config: Optional[Union[GPT2Config, Dict[str, Any]]] = None,
+        language_config: Optional[Union[Qwen3Config, Dict[str, Any]]] = None,
+        n_vq: int = 12,
+        audio_vocab_size: int = 1024,
+        audio_codebook_sizes: Optional[list[int]] = None,
+        audio_pad_token_id: int = 1024,
+        audio_pad_code: Optional[int] = None,
+        pad_token_id: int = 151643,
+        im_start_token_id: int = 151644,
+        im_end_token_id: int = 151645,
+        audio_start_token_id: int = 151669,
+        audio_end_token_id: int = 151670,
+        audio_user_slot_token_id: int = 151654,
+        audio_assistant_slot_token_id: int = 151656,
+        audio_assistant_gen_slot_token_id: Optional[int] = None,
+        sampling_rate: int = 48000,
+        audio_tokenizer_name_or_path: Optional[str] = None,
+        attn_implementation: str = "flash_attention_2",
+        local_transformer_attn_implementation: Optional[str] = None,
+        local_text_head_mode: str = "binary",
+        initializer_range: float = 0.02,
+        **kwargs: Any,
+    ) -> None:
+        if qwen3_config is None and language_config is not None:
+            qwen3_config = language_config
+        if isinstance(qwen3_config, dict):
+            self.qwen3_config = Qwen3Config(**qwen3_config)
+        elif qwen3_config is None:
+            self.qwen3_config = Qwen3Config()
+        else:
+            self.qwen3_config = qwen3_config
+
+        if isinstance(gpt2_config, dict):
+            self.gpt2_config = GPT2Config(**gpt2_config)
+        elif gpt2_config is None:
+            self.gpt2_config = GPT2Config(
+                vocab_size=int(self.qwen3_config.vocab_size),
+                n_embd=int(self.qwen3_config.hidden_size),
+                n_layer=1,
+                n_head=max(1, int(self.qwen3_config.hidden_size) // 80),
+                n_positions=int(n_vq) + 1,
+                n_ctx=int(n_vq) + 1,
+                activation_function="silu",
+                layer_norm_epsilon=1e-6,
+                resid_pdrop=0.0,
+                embd_pdrop=0.0,
+                attn_pdrop=0.0,
+            )
+        else:
+            self.gpt2_config = gpt2_config
+
+        self.n_vq = int(n_vq)
+        if self.n_vq <= 0:
+            raise ValueError("n_vq must be positive.")
+        if audio_codebook_sizes is None:
+            self.audio_codebook_sizes = [int(audio_vocab_size)] * self.n_vq
+        else:
+            self.audio_codebook_sizes = [int(size) for size in audio_codebook_sizes]
+        if len(self.audio_codebook_sizes) != self.n_vq:
+            raise ValueError(
+                f"audio_codebook_sizes must have length n_vq={self.n_vq}, "
+                f"got {len(self.audio_codebook_sizes)}."
+            )
+        if any(size <= 0 for size in self.audio_codebook_sizes):
+            raise ValueError("audio_codebook_sizes must contain positive integers.")
+        self.audio_vocab_size = int(max(int(audio_vocab_size), max(self.audio_codebook_sizes)))
+        self.audio_pad_token_id = int(audio_pad_code if audio_pad_code is not None else audio_pad_token_id)
+        self.audio_pad_code = self.audio_pad_token_id
+        if self.audio_pad_token_id < self.audio_vocab_size:
+            raise ValueError("audio_pad_token_id/audio_pad_code must be outside the audio vocab.")
+
+        self.pad_token_id = int(pad_token_id)
+        self.im_start_token_id = int(im_start_token_id)
+        self.im_end_token_id = int(im_end_token_id)
+        self.audio_start_token_id = int(audio_start_token_id)
+        self.audio_end_token_id = int(audio_end_token_id)
+        self.audio_user_slot_token_id = int(audio_user_slot_token_id)
+        self.audio_assistant_slot_token_id = int(
+            audio_assistant_slot_token_id
+            if audio_assistant_gen_slot_token_id is None
+            else audio_assistant_gen_slot_token_id
+        )
+        self.audio_assistant_gen_slot_token_id = self.audio_assistant_slot_token_id
+
+        self.sampling_rate = int(sampling_rate)
+        self.audio_tokenizer_name_or_path = audio_tokenizer_name_or_path
+        self.attn_implementation = _normalize_attention_implementation(attn_implementation)
+        self.local_transformer_attn_implementation = _normalize_attention_implementation(
+            local_transformer_attn_implementation,
+            default=self.attn_implementation,
+        )
+        self.initializer_range = float(initializer_range)
+
+        self.hidden_size = int(self.qwen3_config.hidden_size)
+        self.vocab_size = int(self.qwen3_config.vocab_size)
+        self.local_hidden_size = int(self.gpt2_config.hidden_size)
+        if self.local_hidden_size != self.hidden_size:
+            raise ValueError(
+                "This MOSS-TTS-Local-Transformer-v1.5 release expects local hidden size to "
+                "match Qwen3 hidden size so audio embeddings and heads are tied."
+            )
+
+        normalized_text_head_mode = str(local_text_head_mode or "full_vocab").strip().lower()
+        if normalized_text_head_mode in {"full", "full-vocab", "vocab"}:
+            normalized_text_head_mode = "full_vocab"
+        if normalized_text_head_mode not in {"full_vocab", "binary"}:
+            raise ValueError("local_text_head_mode must be 'full_vocab' or 'binary'.")
+        self.local_text_head_mode = normalized_text_head_mode
+
+        kwargs.setdefault("tie_word_embeddings", True)
+        super().__init__(pad_token_id=self.pad_token_id, **kwargs)
+
+    @property
+    def language_config(self) -> Qwen3Config:
+        return self.qwen3_config
+
+    def to_dict(self) -> Dict[str, Any]:
+        output = super().to_dict()
+        output["qwen3_config"] = self.qwen3_config.to_dict()
+        output["language_config"] = self.qwen3_config.to_dict()
+        output["gpt2_config"] = self.gpt2_config.to_dict()
+        output["audio_pad_code"] = self.audio_pad_token_id
+        output["audio_assistant_gen_slot_token_id"] = self.audio_assistant_slot_token_id
+        return output

assets/sft3/gpt2_decoder.py

@@ -0,0 +1,721 @@
+# 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
+
+
+def _is_static_kv_cache_layer(layer_past: object) -> bool:
+    return isinstance(layer_past, dict) and bool(layer_past.get("static_kv_cache", False))
+
+
+class MossTTSNanoGPT2RotaryEmbedding(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}")
+        self.dim = int(dim)
+        self.base = float(base)
+        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)
+
+    def _compute_inv_freq(self, device: Optional[torch.device] = None) -> torch.Tensor:
+        return 1.0 / (
+            self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim)
+        )
+
+    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)
+        inv_freq = self._compute_inv_freq(device=device)
+        freqs = torch.einsum("bs,d->bsd", position_ids.to(device=device, dtype=inv_freq.dtype), 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 MossTTSNanoGPT2MLP(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 MossTTSNanoGPT2Attention(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 = MossTTSNanoGPT2RotaryEmbedding(
+                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
+        if attention_mask is not None or query.shape[-2] != key.shape[-2]:
+            mask = self._causal_attention_mask(
+                attention_mask=attention_mask,
+                query_length=query.shape[-2],
+                key_length=key.shape[-2],
+                device=query.device,
+            )
+        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,
+        )
+        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,
+            )
+
+        if query.shape[1] != key.shape[1]:
+            query_attention_mask = attention_mask[:, -query.shape[1] :]
+            unpadded_query, query_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
+                query,
+                query_attention_mask,
+            )
+            unpadded_key, _, cu_seqlens_k, max_seqlen_k, _ = 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_q,
+                cu_seqlens_k,
+                max_seqlen_q,
+                max_seqlen_k,
+                dropout_p=dropout_p,
+                causal=True,
+            )
+            return pad_input(output, query_indices, query.shape[0], query.shape[1])
+
+        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)
+
+        static_layer_past = layer_past is not None and _is_static_kv_cache_layer(layer_past)
+        if static_layer_past:
+            past_length = int(layer_past.get("length", 0))
+            new_length = past_length + int(key.shape[1])
+            key_cache = layer_past["key"]
+            value_cache = layer_past["value"]
+            if new_length > int(key_cache.shape[1]):
+                raise ValueError(
+                    f"Static KV cache is too short: need {new_length}, capacity={int(key_cache.shape[1])}."
+                )
+            key_cache[:, past_length:new_length].copy_(key)
+            value_cache[:, past_length:new_length].copy_(value)
+            key = key_cache[:, :new_length]
+            value = value_cache[:, :new_length]
+            layer_past["length"] = new_length
+        elif 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 = layer_past if (use_cache and static_layer_past) else ((key, value) if use_cache else None)
+
+        if self.attn_implementation == "flash_attention_2":
+            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 MossTTSNanoGPT2Block(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 = MossTTSNanoGPT2Attention(config, layer_idx=layer_idx, attn_implementation=attn_implementation)
+        self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.mlp = MossTTSNanoGPT2MLP(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 MossTTSNanoGPT2Model(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(
+            [MossTTSNanoGPT2Block(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:
+            diffs = cu_seqlens[:, 1:] - cu_seqlens[:, :-1]
+            return diffs.gt(0).sum(dim=-1).to(device=device, dtype=torch.long)
+        if num_sequences.ndim == 0:
+            num_sequences = num_sequences.unsqueeze(0)
+        return num_sequences.to(device=device, dtype=torch.long)
+
+    @staticmethod
+    def _packed_segments_from_cu_seqlens(
+        cu_seqlens: torch.Tensor,
+        num_sequences: Optional[torch.Tensor],
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if cu_seqlens.ndim == 1:
+            cu_seqlens = cu_seqlens.unsqueeze(0)
+        cu_seqlens = cu_seqlens.to(device=device)
+        batch_size, boundary_count = cu_seqlens.shape
+        segment_slots = boundary_count - 1
+        if segment_slots <= 0:
+            empty = torch.empty(0, dtype=torch.long, device=device)
+            return empty, empty, empty
+
+        counts = MossTTSNanoGPT2Model._normalize_num_sequences(cu_seqlens, num_sequences, device=device)
+        counts = counts.clamp(min=0, max=segment_slots)
+        segment_slots = int(counts.max().item()) if counts.numel() > 0 else 0
+        if segment_slots <= 0:
+            empty = torch.empty(0, dtype=torch.long, device=device)
+            return empty, empty, empty
+        cu_seqlens = cu_seqlens[:, : segment_slots + 1]
+        slot_ids = torch.arange(segment_slots, device=device).unsqueeze(0)
+        valid_slots = slot_ids < counts.unsqueeze(1)
+
+        starts = cu_seqlens[:, :-1].to(dtype=torch.long)
+        ends = cu_seqlens[:, 1:].to(dtype=torch.long)
+        lengths = (ends - starts).clamp_min(0)
+        lengths = torch.where(valid_slots, lengths, torch.zeros((), dtype=torch.long, device=device))
+
+        batch_ids = torch.arange(batch_size, device=device, dtype=torch.long).unsqueeze(1).expand(batch_size, segment_slots)
+        batch_ids = batch_ids.reshape(-1)
+        starts = starts.reshape(-1)
+        lengths = lengths.reshape(-1)
+
+        valid_segments = lengths.gt(0)
+        valid_count = int(valid_segments.to(dtype=torch.long).sum().item())
+        if valid_count <= 0:
+            empty = torch.empty(0, dtype=torch.long, device=device)
+            return empty, empty, empty
+        if valid_count == lengths.numel():
+            return batch_ids, starts, lengths
+
+        valid_order = torch.argsort(valid_segments.to(dtype=torch.long), descending=True, stable=True)[:valid_count]
+        return (
+            batch_ids.index_select(0, valid_order),
+            starts.index_select(0, valid_order),
+            lengths.index_select(0, valid_order),
+        )
+
+    @staticmethod
+    def _packed_token_indices(
+        batch_ids: torch.Tensor,
+        starts: torch.Tensor,
+        lengths: torch.Tensor,
+        seq_len: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        total_tokens = int(lengths.sum().item())
+        if total_tokens <= 0:
+            empty = torch.empty(0, dtype=torch.long, device=lengths.device)
+            return empty, empty
+
+        segment_ids = torch.repeat_interleave(
+            torch.arange(lengths.numel(), device=lengths.device, dtype=torch.long),
+            lengths,
+            output_size=total_tokens,
+        )
+        segment_starts = torch.cumsum(lengths, dim=0) - lengths
+        positions = torch.arange(total_tokens, device=lengths.device, dtype=torch.long) - segment_starts[segment_ids]
+        indices = batch_ids[segment_ids] * seq_len + starts[segment_ids] + positions
+        return indices, positions
+
+    @staticmethod
+    def build_packed_position_ids(
+        attention_mask: Optional[torch.Tensor],
+        cu_seqlens: torch.Tensor,
+        num_sequences: Optional[torch.Tensor],
+        sequence_length: Optional[int] = None,
+    ) -> torch.Tensor:
+        if cu_seqlens.ndim == 1:
+            cu_seqlens = cu_seqlens.unsqueeze(0)
+        batch_size = cu_seqlens.shape[0]
+        seq_len = int(sequence_length or (cu_seqlens.shape[1] - 1))
+        device = cu_seqlens.device
+        position_ids = torch.zeros((batch_size, seq_len), dtype=torch.long, device=device)
+        batch_ids, starts, lengths = MossTTSNanoGPT2Model._packed_segments_from_cu_seqlens(
+            cu_seqlens,
+            num_sequences,
+            device,
+        )
+        if lengths.numel() > 0:
+            indices, positions = MossTTSNanoGPT2Model._packed_token_indices(batch_ids, starts, lengths, seq_len)
+            position_ids.view(-1).scatter_(0, indices, positions)
+        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
+        seq_len = hidden_states.shape[1]
+        batch_ids, starts, lengths = MossTTSNanoGPT2Model._packed_segments_from_cu_seqlens(
+            cu_seqlens,
+            num_sequences,
+            device,
+        )
+        if lengths.numel() == 0:
+            raise ValueError("cu_seqlens did not describe any non-empty packed sequences.")
+
+        indices, _ = MossTTSNanoGPT2Model._packed_token_indices(batch_ids, starts, lengths, seq_len)
+        cumulative = torch.empty(lengths.numel() + 1, dtype=torch.int32, device=device)
+        cumulative[0] = 0
+        cumulative[1:] = lengths.to(dtype=torch.int32).cumsum(dim=0)
+        return PackedSequenceMetadata(
+            cu_seqlens=cumulative,
+            max_seqlen=int(lengths.max().item()),
+            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
+        query_attention_mask = None
+        if attention_mask is not None:
+            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:
+                if attention_mask is None:
+                    raise ValueError("attention_mask must be provided with cu_seqlens packing.")
+                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,
+                    sequence_length=hidden_states.shape[1],
+                )
+            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:
+                    first_layer_past = past_key_values[0]
+                    if _is_static_kv_cache_layer(first_layer_past):
+                        past_length = int(first_layer_past.get("length", 0))
+                    else:
+                        past_length = first_layer_past[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)
+        if query_attention_mask is not None:
+            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,
+                )
+            if query_attention_mask is not None:
+                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)
+        if query_attention_mask is not None:
+            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,
+        )

assets/sft3/modeling_moss_tts.py

@@ -0,0 +1,623 @@
+# coding=utf-8
+"""Modeling code for the MOSS-TTS-Local-Transformer-v1.5 HuggingFace release."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+import torch
+import torch.nn as nn
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
+from transformers.utils import ModelOutput
+
+from .configuration_moss_tts import MossTTSLocalConfig
+from .gpt2_decoder import MossTTSNanoGPT2Model
+from .qwen3_decoder import MossQwen3Model
+
+
+@dataclass
+class MossTTSLocalOutput(ModelOutput):
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[tuple[tuple[torch.Tensor, torch.Tensor], ...]] = None
+    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+
+
+def _find_last_equal(input_ids: torch.LongTensor, value: int) -> torch.LongTensor:
+    matches = input_ids.eq(int(value))
+    if not bool(matches.any(dim=1).all().item()):
+        raise ValueError(f"Every sample must contain token id {int(value)}.")
+    positions = torch.arange(input_ids.shape[1], device=input_ids.device, dtype=torch.long)
+    masked_positions = positions.unsqueeze(0).masked_fill(~matches, -1)
+    return masked_positions.max(dim=1).values
+
+
+class MossTTSLocalPreTrainedModel(PreTrainedModel):
+    config_class = MossTTSLocalConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MossTTSNanoGPT2Block", "MossQwen3DecoderLayer"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False) -> None:
+        if isinstance(module, MossTTSNanoGPT2Model) or isinstance(module, MossQwen3Model):
+            module.gradient_checkpointing = value
+
+
+class MossTTSLocalModel(MossTTSLocalPreTrainedModel):
+    _tied_weights_keys = None
+
+    def __init__(self, config: MossTTSLocalConfig) -> None:
+        super().__init__(config)
+        self._tied_weights_keys = self._build_tied_weights_keys(config)
+
+        config.qwen3_config.pad_token_id = config.pad_token_id
+        config.qwen3_config._attn_implementation = config.attn_implementation
+        local_gpt2_config = config.gpt2_config.to_dict()
+        local_gpt2_config["n_layer"] = int(getattr(config, "local_transformer_layers", config.gpt2_config.n_layer))
+        local_gpt2_config["n_positions"] = int(config.n_vq) + 1
+        local_gpt2_config["n_ctx"] = int(config.n_vq) + 1
+        local_gpt2_config = GPT2Config(**local_gpt2_config)
+        local_gpt2_config.pad_token_id = config.pad_token_id
+        local_gpt2_config._attn_implementation = config.local_transformer_attn_implementation
+
+        self.transformer = MossQwen3Model(config.qwen3_config)
+        self.local_transformer = MossTTSNanoGPT2Model(
+            local_gpt2_config,
+            attn_implementation=config.local_transformer_attn_implementation,
+        )
+        self.local_transformer.wte = nn.Identity()
+
+        hidden_size = int(config.hidden_size)
+        self.audio_embeddings = nn.ModuleList(
+            [
+                nn.Embedding(int(config.audio_codebook_sizes[index]), hidden_size)
+                for index in range(config.n_vq)
+            ]
+        )
+        self.text_lm_head = nn.Linear(hidden_size, int(config.vocab_size), bias=False)
+        self.audio_lm_heads = nn.ModuleList(
+            [
+                nn.Linear(hidden_size, int(config.audio_codebook_sizes[index]), bias=False)
+                for index in range(config.n_vq)
+            ]
+        )
+        self.local_text_lm_head = (
+            nn.Linear(hidden_size, 2, bias=False)
+            if self._use_binary_local_text_head()
+            else None
+        )
+
+        self.post_init()
+        self.tie_weights()
+        self.initialize_local_text_lm_head_from_text_lm_head()
+
+    def can_generate(self) -> bool:
+        return True
+
+    @staticmethod
+    def _build_tied_weights_keys(config: MossTTSLocalConfig) -> dict[str, str]:
+        tied_weights = {"text_lm_head.weight": "transformer.embed_tokens.weight"}
+        tied_weights.update(
+            {
+                f"audio_lm_heads.{index}.weight": f"audio_embeddings.{index}.weight"
+                for index in range(config.n_vq)
+            }
+        )
+        return tied_weights
+
+    def tie_weights(self, *args, **kwargs) -> None:
+        del args, kwargs
+        self.text_lm_head.weight = self.transformer.embed_tokens.weight
+        for embedding, head in zip(self.audio_embeddings, self.audio_lm_heads):
+            head.weight = embedding.weight
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.transformer.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.transformer.embed_tokens = value
+        self.tie_weights()
+        self.initialize_local_text_lm_head_from_text_lm_head()
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.text_lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.text_lm_head = new_embeddings
+        self.tie_weights()
+        self.initialize_local_text_lm_head_from_text_lm_head()
+
+    def _use_binary_local_text_head(self) -> bool:
+        return str(getattr(self.config, "local_text_head_mode", "full_vocab")).strip().lower() == "binary"
+
+    def _local_text_candidate_ids(self, device: torch.device) -> torch.LongTensor:
+        return torch.tensor(
+            [
+                int(self.config.audio_assistant_slot_token_id),
+                int(self.config.audio_end_token_id),
+            ],
+            dtype=torch.long,
+            device=device,
+        )
+
+    def initialize_local_text_lm_head_from_text_lm_head(self) -> None:
+        if not self._use_binary_local_text_head() or self.local_text_lm_head is None:
+            return
+        candidate_ids = self._local_text_candidate_ids(self.text_lm_head.weight.device)
+        with torch.no_grad():
+            source_weight = self.text_lm_head.weight.index_select(0, candidate_ids)
+            if tuple(source_weight.shape) == tuple(self.local_text_lm_head.weight.shape):
+                self.local_text_lm_head.weight.copy_(
+                    source_weight.to(
+                        device=self.local_text_lm_head.weight.device,
+                        dtype=self.local_text_lm_head.weight.dtype,
+                    )
+                )
+
+    def _resolve_fixed_nq(
+        self,
+        n_vq_for_inference: Optional[int] = None,
+        nq: Optional[int] = None,
+    ) -> int:
+        requested = n_vq_for_inference if n_vq_for_inference is not None else nq
+        config_nq = int(self.config.n_vq)
+        if requested is not None and int(requested) != config_nq:
+            raise ValueError(
+                "This MOSS-TTS-Local-Transformer-v1.5 release is trained with a fixed RVQ depth. "
+                f"Expected n_vq={config_nq}, got {int(requested)}."
+            )
+        return config_nq
+
+    def _build_inputs_embeds(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
+        if input_ids.ndim != 3 or input_ids.shape[-1] != self.config.n_vq + 1:
+            raise ValueError(
+                f"Expected input_ids shape [batch, seq, {self.config.n_vq + 1}], "
+                f"got {tuple(input_ids.shape)}."
+            )
+        text_ids = input_ids[..., 0]
+        inputs_embeds = self.transformer.embed_tokens(text_ids)
+        for channel_index, embedding in enumerate(self.audio_embeddings):
+            channel_ids = input_ids[..., channel_index + 1]
+            valid_mask = channel_ids.ne(self.config.audio_pad_token_id)
+            safe_ids = channel_ids.masked_fill(~valid_mask, 0)
+            audio_embeds = embedding(safe_ids) * valid_mask.unsqueeze(-1)
+            inputs_embeds = inputs_embeds + audio_embeds
+        return inputs_embeds
+
+    def _global_hidden_to_local(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+        return hidden_states
+
+    @staticmethod
+    def _local_past_length(past_key_values: Optional[tuple[Any, ...]]) -> int:
+        if past_key_values is None or len(past_key_values) == 0:
+            return 0
+        first_layer_past = past_key_values[0]
+        if isinstance(first_layer_past, dict) and bool(first_layer_past.get("static_kv_cache", False)):
+            return int(first_layer_past.get("length", 0))
+        return int(first_layer_past[0].shape[1])
+
+    def _new_static_local_past_key_values(
+        self,
+        batch_size: int,
+        max_length: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[dict[str, Any], ...]:
+        layers = []
+        for block in self.local_transformer.h:
+            attn = block.attn
+            cache_shape = (
+                int(batch_size),
+                int(max_length),
+                int(attn.num_heads),
+                int(attn.head_dim),
+            )
+            layers.append(
+                {
+                    "static_kv_cache": True,
+                    "key": torch.empty(cache_shape, device=device, dtype=dtype),
+                    "value": torch.empty(cache_shape, device=device, dtype=dtype),
+                    "length": 0,
+                }
+            )
+        return tuple(layers)
+
+    def _decode_local_hidden_states_with_cache(
+        self,
+        local_inputs_embeds: torch.FloatTensor,
+        past_key_values: Optional[tuple[Any, ...]] = None,
+    ) -> tuple[torch.FloatTensor, Optional[tuple[Any, ...]]]:
+        if (
+            past_key_values is None
+            and not self.training
+            and bool(getattr(self.config, "use_static_local_kv_cache", True))
+        ):
+            max_length = max(int(getattr(self.config, "n_vq", 0)) + 1, int(local_inputs_embeds.shape[1]))
+            past_key_values = self._new_static_local_past_key_values(
+                batch_size=int(local_inputs_embeds.shape[0]),
+                max_length=max_length,
+                device=local_inputs_embeds.device,
+                dtype=local_inputs_embeds.dtype,
+            )
+        past_length = self._local_past_length(past_key_values)
+        local_seq_len = int(local_inputs_embeds.shape[1])
+        local_position_ids = torch.arange(
+            past_length,
+            past_length + local_seq_len,
+            device=local_inputs_embeds.device,
+            dtype=torch.long,
+        ).unsqueeze(0)
+        if int(local_inputs_embeds.shape[0]) != 1:
+            local_position_ids = local_position_ids.expand(int(local_inputs_embeds.shape[0]), -1)
+        local_outputs = self.local_transformer(
+            input_ids=None,
+            past_key_values=past_key_values,
+            attention_mask=None,
+            position_ids=local_position_ids,
+            inputs_embeds=local_inputs_embeds,
+            use_cache=True,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+            cu_seqlens=None,
+            num_sequences=None,
+        )
+        return local_outputs.last_hidden_state, local_outputs.past_key_values
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[tuple[tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = True,
+        **kwargs,
+    ) -> Union[tuple, MossTTSLocalOutput]:
+        del kwargs
+        if inputs_embeds is None:
+            if input_ids is None:
+                raise ValueError("Either input_ids or inputs_embeds must be provided.")
+            inputs_embeds = self._build_inputs_embeds(input_ids)
+        outputs = self.transformer(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cu_seqlens=None,
+            num_sequences=None,
+        )
+        if not return_dict:
+            return (
+                outputs.last_hidden_state,
+                outputs.past_key_values,
+                outputs.hidden_states,
+                outputs.attentions,
+            )
+        return MossTTSLocalOutput(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def _decode_local_last_hidden_state(
+        self,
+        local_inputs_embeds: torch.FloatTensor,
+    ) -> torch.FloatTensor:
+        local_seq_len = int(local_inputs_embeds.shape[1])
+        local_position_ids = torch.arange(
+            0,
+            local_seq_len,
+            device=local_inputs_embeds.device,
+            dtype=torch.long,
+        ).unsqueeze(0)
+        if int(local_inputs_embeds.shape[0]) != 1:
+            local_position_ids = local_position_ids.expand(int(local_inputs_embeds.shape[0]), -1)
+        local_outputs = self.local_transformer(
+            input_ids=None,
+            attention_mask=None,
+            position_ids=local_position_ids,
+            inputs_embeds=local_inputs_embeds,
+            use_cache=False,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+            cu_seqlens=None,
+            num_sequences=None,
+        )
+        return local_outputs.last_hidden_state[:, -1, :]
+
+    def _filter_logits(
+        self,
+        logits: torch.FloatTensor,
+        top_k: Optional[int],
+        top_p: Optional[float],
+    ) -> torch.FloatTensor:
+        scores = logits
+        if top_k is not None and int(top_k) > 0 and int(top_k) < scores.shape[-1]:
+            kth = torch.topk(scores, int(top_k), dim=-1).values[..., -1, None]
+            scores = scores.masked_fill(scores < kth, -torch.inf)
+        if top_p is not None and 0.0 < float(top_p) < 1.0:
+            sorted_scores, sorted_indices = torch.sort(scores, descending=True, dim=-1)
+            sorted_probs = torch.softmax(sorted_scores, dim=-1)
+            cumulative_probs = sorted_probs.cumsum(dim=-1)
+            sorted_mask = cumulative_probs > float(top_p)
+            sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
+            sorted_mask[..., 0] = False
+            remove_mask = torch.zeros_like(scores, dtype=torch.bool)
+            remove_mask.scatter_(dim=-1, index=sorted_indices, src=sorted_mask)
+            scores = scores.masked_fill(remove_mask, -torch.inf)
+        return scores
+
+    def _apply_repetition_penalty(
+        self,
+        scores: torch.FloatTensor,
+        previous_token_ids: Optional[torch.LongTensor],
+        penalty: float,
+    ) -> torch.FloatTensor:
+        if previous_token_ids is None or float(penalty) == 1.0:
+            return scores
+        if previous_token_ids.ndim == 1:
+            previous_token_ids = previous_token_ids.unsqueeze(0)
+        updated = scores.clone()
+        for batch_index in range(updated.shape[0]):
+            unique_token_ids = torch.unique(previous_token_ids[batch_index])
+            unique_token_ids = unique_token_ids[
+                (unique_token_ids >= 0) & (unique_token_ids < updated.shape[-1])
+            ]
+            if unique_token_ids.numel() == 0:
+                continue
+            token_scores = updated[batch_index].index_select(0, unique_token_ids)
+            token_scores = torch.where(
+                token_scores < 0,
+                token_scores * float(penalty),
+                token_scores / float(penalty),
+            )
+            updated[batch_index].scatter_(0, unique_token_ids, token_scores)
+        return updated
+
+    def _sample_next_token(
+        self,
+        logits: torch.FloatTensor,
+        do_sample: bool,
+        temperature: float,
+        top_k: Optional[int],
+        top_p: Optional[float],
+        previous_token_ids: Optional[torch.LongTensor] = None,
+        repetition_penalty: float = 1.0,
+    ) -> torch.LongTensor:
+        scores = logits.float()
+        scores = self._apply_repetition_penalty(scores, previous_token_ids, repetition_penalty)
+        if not do_sample:
+            return torch.argmax(scores, dim=-1)
+        if float(temperature) <= 0:
+            raise ValueError("temperature must be positive when do_sample=True.")
+        scores = scores / float(temperature)
+        scores = self._filter_logits(scores, top_k=top_k, top_p=top_p)
+        probs = torch.softmax(scores, dim=-1)
+        return torch.multinomial(probs, num_samples=1).squeeze(-1)
+
+    def _sample_next_assistant_text_token(
+        self,
+        local_hidden_states: torch.FloatTensor,
+        do_sample: bool,
+        temperature: float,
+        top_k: Optional[int],
+        top_p: Optional[float],
+    ) -> torch.LongTensor:
+        if self._use_binary_local_text_head() and self.local_text_lm_head is not None:
+            logits = self.local_text_lm_head(local_hidden_states)
+            sampled_indices = self._sample_next_token(
+                logits=logits,
+                do_sample=do_sample,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+            )
+            candidate_ids = self._local_text_candidate_ids(logits.device)
+            return candidate_ids[sampled_indices]
+
+        candidate_ids = self._local_text_candidate_ids(local_hidden_states.device)
+        logits = self.text_lm_head(local_hidden_states).index_select(dim=-1, index=candidate_ids)
+        sampled_indices = self._sample_next_token(
+            logits=logits,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+        )
+        return candidate_ids[sampled_indices]
+
+    def _build_generation_row(
+        self,
+        batch_size: int,
+        device: torch.device,
+        audio_token_ids: torch.LongTensor,
+    ) -> torch.LongTensor:
+        row = torch.full(
+            (batch_size, 1, self.config.n_vq + 1),
+            int(self.config.audio_pad_token_id),
+            dtype=torch.long,
+            device=device,
+        )
+        row[:, :, 0] = int(self.config.audio_assistant_slot_token_id)
+        row[:, :, 1:] = audio_token_ids.unsqueeze(1)
+        return row
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        max_new_tokens: Optional[int] = None,
+        max_new_frames: Optional[int] = None,
+        do_sample: bool = True,
+        text_temperature: float = 1.0,
+        text_top_p: float = 1.0,
+        text_top_k: int = 50,
+        audio_temperature: Optional[float] = None,
+        audio_top_p: Optional[float] = None,
+        audio_top_k: Optional[int] = None,
+        audio_repetition_penalty: Optional[float] = None,
+        temperature: float = 1.0,
+        top_p: float = 0.95,
+        top_k: int = 50,
+        repetition_penalty: float = 1.0,
+        use_kv_cache: bool = True,
+        n_vq_for_inference: Optional[int] = None,
+        nq: Optional[int] = None,
+        **kwargs,
+    ) -> list[tuple[int, torch.LongTensor]]:
+        del kwargs
+        self._resolve_fixed_nq(n_vq_for_inference=n_vq_for_inference, nq=nq)
+
+        if input_ids.ndim == 2:
+            input_ids = input_ids.unsqueeze(0)
+        if input_ids.ndim != 3:
+            raise ValueError(f"Expected input_ids with 3 dims, got {tuple(input_ids.shape)}.")
+        if input_ids.shape[-1] != self.config.n_vq + 1:
+            raise ValueError(
+                f"Expected {self.config.n_vq + 1} channels from config.n_vq, got {input_ids.shape[-1]}."
+            )
+        if attention_mask is None:
+            attention_mask = torch.ones(input_ids.shape[:2], dtype=torch.bool, device=input_ids.device)
+        elif attention_mask.ndim == 1:
+            attention_mask = attention_mask.unsqueeze(0)
+        attention_mask = attention_mask.to(device=input_ids.device, dtype=torch.bool)
+
+        frame_budget = max_new_frames if max_new_frames is not None else max_new_tokens
+        if frame_budget is None:
+            frame_budget = 4096
+        frame_budget = int(frame_budget)
+
+        audio_temperature = float(temperature if audio_temperature is None else audio_temperature)
+        audio_top_p = float(top_p if audio_top_p is None else audio_top_p)
+        audio_top_k = int(top_k if audio_top_k is None else audio_top_k)
+        audio_repetition_penalty = float(
+            repetition_penalty if audio_repetition_penalty is None else audio_repetition_penalty
+        )
+
+        batch_size = input_ids.shape[0]
+        input_ids_length = input_ids.shape[1]
+        current_input_ids = input_ids
+        current_attention_mask = attention_mask
+        current_model_input_ids = current_input_ids
+        generated_frames: list[torch.LongTensor] = []
+        finished = torch.zeros(batch_size, dtype=torch.bool, device=input_ids.device)
+        past_key_values = None
+        local_dtype = self.local_transformer.ln_f.weight.dtype
+
+        for _ in range(frame_budget):
+            generated_audio_history = torch.stack(generated_frames, dim=1) if generated_frames else None
+            global_inputs_embeds = self._build_inputs_embeds(current_model_input_ids)
+            global_outputs = self.transformer(
+                input_ids=None,
+                past_key_values=past_key_values,
+                attention_mask=current_attention_mask,
+                position_ids=None,
+                inputs_embeds=global_inputs_embeds,
+                use_cache=use_kv_cache,
+                output_attentions=False,
+                output_hidden_states=False,
+                return_dict=True,
+                cu_seqlens=None,
+                num_sequences=None,
+            )
+            global_hidden_states = global_outputs.last_hidden_state[:, -1, :]
+            local_global_hidden_states = self._global_hidden_to_local(global_hidden_states).to(dtype=local_dtype)
+
+            local_prefix_hidden_states, local_prefix_past_key_values = self._decode_local_hidden_states_with_cache(
+                local_global_hidden_states.unsqueeze(1)
+            )
+            local_hidden_states = local_prefix_hidden_states[:, -1, :]
+            next_text_tokens = self._sample_next_assistant_text_token(
+                local_hidden_states=local_hidden_states,
+                do_sample=do_sample,
+                temperature=text_temperature,
+                top_k=text_top_k,
+                top_p=text_top_p,
+            )
+            should_continue = next_text_tokens.eq(int(self.config.audio_assistant_slot_token_id)) & ~finished
+            finished = finished | next_text_tokens.eq(int(self.config.audio_end_token_id))
+            if not bool(should_continue.any().item()):
+                break
+
+            next_frame_tokens = []
+            for channel_index in range(int(self.config.n_vq)):
+                channel_logits = self.audio_lm_heads[channel_index](local_hidden_states)
+                channel_token = self._sample_next_token(
+                    logits=channel_logits,
+                    do_sample=do_sample,
+                    temperature=audio_temperature,
+                    top_k=audio_top_k,
+                    top_p=audio_top_p,
+                    previous_token_ids=(
+                        None
+                        if generated_audio_history is None
+                        else generated_audio_history[:, :, channel_index]
+                    ),
+                    repetition_penalty=audio_repetition_penalty,
+                )
+                next_frame_tokens.append(channel_token)
+                if channel_index + 1 < int(self.config.n_vq):
+                    current_local_input = self.audio_embeddings[channel_index](channel_token).to(dtype=local_dtype)
+                    local_token_hidden_states, local_prefix_past_key_values = (
+                        self._decode_local_hidden_states_with_cache(
+                            current_local_input.unsqueeze(1),
+                            past_key_values=local_prefix_past_key_values,
+                        )
+                    )
+                    local_hidden_states = local_token_hidden_states[:, -1, :]
+
+            next_frame = torch.stack(next_frame_tokens, dim=-1)
+            next_frame = next_frame.masked_fill(
+                ~should_continue.unsqueeze(-1),
+                int(self.config.audio_pad_token_id),
+            )
+            generated_frames.append(next_frame)
+
+            next_row = self._build_generation_row(
+                batch_size=batch_size,
+                device=input_ids.device,
+                audio_token_ids=next_frame,
+            )
+            if bool((~should_continue).any().item()):
+                next_row[~should_continue, 0, 0] = int(self.config.pad_token_id)
+                next_row[~should_continue, 0, 1:] = int(self.config.audio_pad_token_id)
+
+            current_input_ids = torch.cat([current_input_ids, next_row], dim=1)
+            current_attention_mask = torch.cat(
+                [current_attention_mask, should_continue.unsqueeze(1)],
+                dim=1,
+            )
+            if use_kv_cache:
+                current_model_input_ids = next_row
+                past_key_values = global_outputs.past_key_values
+            else:
+                current_model_input_ids = current_input_ids
+
+        start_indices = _find_last_equal(input_ids[..., 0], int(self.config.audio_start_token_id))
+        start_lengths = input_ids_length - start_indices - 1
+        outputs: list[tuple[int, torch.LongTensor]] = []
+        for start_index, start_length, generation_ids in zip(
+            start_indices.tolist(),
+            start_lengths.tolist(),
+            current_input_ids,
+        ):
+            outputs.append((int(start_length), generation_ids[int(start_index):].detach().cpu()))
+        return outputs

assets/sft3/processing_moss_tts.py

@@ -0,0 +1,914 @@
+# coding=utf-8
+"""Processor for the MOSS-TTS-Local-Transformer-v1.5 HuggingFace release."""
+
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
+
+import torch
+import torchaudio
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+    BatchFeature,
+    PreTrainedTokenizerBase,
+    ProcessorMixin,
+    logging,
+    processing_utils,
+)
+
+from .configuration_moss_tts import MossTTSLocalConfig
+
+
+if hasattr(processing_utils, "MODALITY_TO_BASE_CLASS_MAPPING"):
+    processing_utils.MODALITY_TO_BASE_CLASS_MAPPING["audio_tokenizer"] = "PreTrainedModel"
+else:
+    processing_utils.AUTO_TO_BASE_CLASS_MAPPING["AutoModel"] = "PreTrainedModel"
+logger = logging.get_logger(__name__)
+
+AUDIO_PLACEHOLDER = "<|audio|>"
+USER_ROLE_PREFIX = "user\n"
+USER_TEMPLATE_REFERENCE_PREFIX = (
+    "<user_inst>\n"
+    "- Reference(s):\n"
+)
+USER_TEMPLATE_AFTER_REFERENCE_SUFFIX = (
+    "\n"
+    "- Text:\n"
+)
+USER_TEMPLATE_SUFFIX = "\n</user_inst>"
+ASSISTANT_TURN_PREFIX = "\n"
+ASSISTANT_ROLE_PREFIX = "assistant\n"
+USER_MESSAGE_FIELDS = (
+    "text",
+    "reference",
+    "instruction",
+    "tokens",
+    "quality",
+    "sound_event",
+    "ambient_sound",
+    "language",
+)
+
+
+def _normalize_template_value(value: Any) -> str:
+    if value is None:
+        return "None"
+    resolved = str(value).strip()
+    return resolved or "None"
+
+
+def _render_user_prompt_after_reference(
+    language_code: object | None = None,
+    prompt_fields: Optional[Dict[str, Any]] = None,
+) -> str:
+    fields = dict(prompt_fields or {})
+    return (
+        "\n- Instruction:\n"
+        + _normalize_template_value(fields.get("instruction"))
+        + "\n- Tokens:\n"
+        + _normalize_template_value(fields.get("tokens"))
+        + "\n- Quality:\n"
+        + _normalize_template_value(fields.get("quality"))
+        + "\n- Sound Event:\n"
+        + _normalize_template_value(fields.get("sound_event"))
+        + "\n- Ambient Sound:\n"
+        + _normalize_template_value(fields.get("ambient_sound"))
+        + "\n- Language:\n"
+        + _normalize_template_value(fields.get("language", language_code))
+        + USER_TEMPLATE_AFTER_REFERENCE_SUFFIX
+    )
+
+
+@dataclass
+class Message:
+    def to_dict(self) -> Dict[str, Any]:
+        raise NotImplementedError
+
+
+@dataclass
+class UserMessage(Message):
+    text: Optional[str] = None
+    reference: Optional[List[Optional[Union[str, os.PathLike, torch.Tensor]]]] = None
+    instruction: Optional[str] = None
+    tokens: Optional[int] = None
+    quality: Optional[str] = None
+    sound_event: Optional[str] = None
+    ambient_sound: Optional[str] = None
+    language: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        template = """<user_inst>
+- Reference(s):
+{reference}
+- Instruction:
+{instruction}
+- Tokens:
+{tokens}
+- Quality:
+{quality}
+- Sound Event:
+{sound_event}
+- Ambient Sound:
+{ambient_sound}
+- Language:
+{language}
+- Text:
+{text}
+</user_inst>"""
+
+        audio_codes_list: list[Union[str, os.PathLike, torch.Tensor]] = []
+        if self.reference is None:
+            reference = "None"
+        else:
+            reference_items: list[str] = []
+            for speaker_idx, speaker_reference in enumerate(self.reference):
+                if speaker_reference is None:
+                    continue
+                # Keep raw audio placeholders directly under "- Reference(s):".
+                # Speaker labels such as "[S1]:" change the token sequence and
+                # can affect voice-clone conditioning.
+                reference_items.append(AUDIO_PLACEHOLDER)
+                audio_codes_list.append(speaker_reference)
+            reference = "\n".join(reference_items) if reference_items else "None"
+
+        self._content = (
+            template.replace("{reference}", str(reference))
+            .replace("{instruction}", str(self.instruction))
+            .replace("{tokens}", str(self.tokens))
+            .replace("{quality}", str(self.quality))
+            .replace("{sound_event}", str(self.sound_event))
+            .replace("{ambient_sound}", str(self.ambient_sound))
+            .replace("{language}", str(self.language))
+            .replace("{text}", str(self.text))
+        )
+        self._audio_codes_list = audio_codes_list
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "role": "user",
+            "content": self._content,
+            "audio_codes_list": self._audio_codes_list,
+            "text": self.text,
+            "instruction": self.instruction,
+            "tokens": self.tokens,
+            "quality": self.quality,
+            "sound_event": self.sound_event,
+            "ambient_sound": self.ambient_sound,
+            "language": self.language,
+        }
+
+
+@dataclass
+class AssistantMessage(Message):
+    audio_codes_list: List[Union[str, os.PathLike, torch.Tensor]]
+    content: str = AUDIO_PLACEHOLDER
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "role": "assistant",
+            "content": self.content,
+            "audio_codes_list": self.audio_codes_list,
+        }
+
+
+class MossTTSLocalProcessor(ProcessorMixin):
+    attributes = ["tokenizer"]
+    tokenizer_class = "AutoTokenizer"
+    audio_tokenizer_class = "AutoModel"
+
+    tokenizer: PreTrainedTokenizerBase
+    audio_tokenizer: Any
+
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        audio_tokenizer: Any = None,
+        model_config: Optional[MossTTSLocalConfig] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(tokenizer=tokenizer, audio_tokenizer=audio_tokenizer, **kwargs)
+        self.tokenizer = tokenizer
+        self.audio_tokenizer = audio_tokenizer
+        self.model_config = model_config or MossTTSLocalConfig()
+
+        def _id_to_token(token_id: int) -> str:
+            token = tokenizer.convert_ids_to_tokens(int(token_id))
+            if isinstance(token, list):
+                return token[0] if token else ""
+            return cast(str, token)
+
+        self.audio_user_slot_token = _id_to_token(self.model_config.audio_user_slot_token_id)
+        self.audio_assistant_slot_token = _id_to_token(self.model_config.audio_assistant_slot_token_id)
+        self.audio_start_token = _id_to_token(self.model_config.audio_start_token_id)
+        self.audio_end_token = _id_to_token(self.model_config.audio_end_token_id)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        trust_remote_code = kwargs.pop("trust_remote_code", True)
+        kwargs.pop("_from_auto", None)
+        codec_path = kwargs.pop("codec_path", None)
+        codec_weight_dtype = kwargs.pop("codec_weight_dtype", "fp32")
+        codec_compute_dtype = kwargs.pop("codec_compute_dtype", None)
+        codec_attention_implementation = kwargs.pop("codec_attention_implementation", None)
+
+        codec_kwargs: Dict[str, Any] = {}
+        if codec_weight_dtype is not None:
+            # Default to fp32 codec weights; callers can pass codec_weight_dtype="bf16" to reduce memory.
+            codec_kwargs["codec_weight_dtype"] = codec_weight_dtype
+
+        model_ref = Path(str(pretrained_model_name_or_path))
+        model_ref_or_name = model_ref if model_ref.exists() else pretrained_model_name_or_path
+        model_config = cast(
+            MossTTSLocalConfig,
+            AutoConfig.from_pretrained(
+                model_ref_or_name,
+                *args,
+                trust_remote_code=trust_remote_code,
+                **kwargs,
+            ),
+        )
+
+        if codec_path is None:
+            try:
+                processor_dict, _ = cls.get_processor_dict(
+                    pretrained_model_name_or_path,
+                    **dict(kwargs),
+                )
+                codec_path = processor_dict.get("audio_tokenizer_name_or_path")
+                audio_tokenizer_dict = processor_dict.get("audio_tokenizer", {})
+                if isinstance(audio_tokenizer_dict, dict):
+                    codec_path = audio_tokenizer_dict.get("audio_tokenizer_name_or_path") or codec_path
+            except Exception:
+                codec_path = None
+        if codec_path is None:
+            codec_path = getattr(model_config, "audio_tokenizer_name_or_path", None)
+        if codec_path is None:
+            codec_path = "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2"
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_ref_or_name,
+            *args,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+        audio_tokenizer = AutoModel.from_pretrained(
+            codec_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+            **codec_kwargs,
+        )
+        if codec_compute_dtype is not None and hasattr(audio_tokenizer, "set_compute_dtype"):
+            audio_tokenizer.set_compute_dtype(codec_compute_dtype)
+        if codec_attention_implementation is not None and hasattr(
+            audio_tokenizer, "set_attention_implementation"
+        ):
+            audio_tokenizer.set_attention_implementation(codec_attention_implementation)
+        return cls(
+            tokenizer=tokenizer,
+            audio_tokenizer=audio_tokenizer,
+            model_config=model_config,
+            **kwargs,
+        )
+
+    @staticmethod
+    def build_user_message(
+        text: Optional[str] = None,
+        reference: Optional[List[Optional[Union[str, os.PathLike, torch.Tensor]]]] = None,
+        instruction: Optional[str] = None,
+        tokens: Optional[int] = None,
+        quality: Optional[str] = None,
+        sound_event: Optional[str] = None,
+        ambient_sound: Optional[str] = None,
+        language: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        if reference is not None and not isinstance(reference, list):
+            reference = [reference]
+        return UserMessage(
+            text=text,
+            reference=reference,
+            instruction=instruction,
+            tokens=tokens,
+            quality=quality,
+            sound_event=sound_event,
+            ambient_sound=ambient_sound,
+            language=language,
+        ).to_dict()
+
+    @staticmethod
+    def build_assistant_message(
+        audio_codes_list: List[Union[str, os.PathLike, torch.Tensor]],
+        content: str = AUDIO_PLACEHOLDER,
+    ) -> Dict[str, Any]:
+        return AssistantMessage(audio_codes_list=audio_codes_list, content=content).to_dict()
+
+    def _assert_fixed_nq(self, n_vq: Optional[int]) -> int:
+        config_nq = int(self.model_config.n_vq)
+        if n_vq is not None and int(n_vq) != config_nq:
+            raise ValueError(
+                "This MOSS-TTS-Local-Transformer-v1.5 release uses the RVQ depth stored in the model config. "
+                f"Expected n_vq={config_nq}, got {int(n_vq)}."
+            )
+        return config_nq
+
+    def _encode_text(self, text: str) -> list[int]:
+        try:
+            return list(self.tokenizer.encode(text, add_special_tokens=False))
+        except TypeError:
+            return list(self.tokenizer.encode(text))
+
+    def _build_text_rows(self, token_ids: Sequence[int], *, device: Optional[torch.device] = None) -> torch.Tensor:
+        rows = torch.full(
+            (len(token_ids), int(self.model_config.n_vq) + 1),
+            int(self.model_config.audio_pad_token_id),
+            dtype=torch.long,
+            device=device,
+        )
+        if token_ids:
+            rows[:, 0] = torch.tensor([int(token_id) for token_id in token_ids], dtype=torch.long, device=rows.device)
+        return rows
+
+    def _build_audio_rows(self, audio_tokens: torch.Tensor, slot_token_id: int) -> torch.Tensor:
+        rows = torch.full(
+            (int(audio_tokens.shape[0]), int(self.model_config.n_vq) + 1),
+            int(self.model_config.audio_pad_token_id),
+            dtype=torch.long,
+            device=audio_tokens.device,
+        )
+        if rows.shape[0] > 0:
+            rows[:, 0] = int(slot_token_id)
+            rows[:, 1:] = audio_tokens.to(dtype=torch.long)
+        return rows
+
+    def _user_prompt_prefix_ids(self) -> list[int]:
+        return (
+            [int(self.model_config.im_start_token_id)]
+            + self._encode_text(USER_ROLE_PREFIX)
+            + self._encode_text(USER_TEMPLATE_REFERENCE_PREFIX)
+        )
+
+    def _user_prompt_after_reference_ids(
+        self,
+        language_code: object | None,
+        prompt_fields: Optional[Dict[str, Any]],
+    ) -> list[int]:
+        return self._encode_text(
+            _render_user_prompt_after_reference(
+                language_code=language_code,
+                prompt_fields=prompt_fields,
+            )
+        )
+
+    def _assistant_prompt_prefix_ids(self) -> list[int]:
+        return (
+            self._encode_text(USER_TEMPLATE_SUFFIX)
+            + [int(self.model_config.im_end_token_id)]
+            + self._encode_text(ASSISTANT_TURN_PREFIX)
+            + [int(self.model_config.im_start_token_id)]
+            + self._encode_text(ASSISTANT_ROLE_PREFIX)
+        )
+
+    def _prompt_fields_from_user_message(self, message: Dict[str, Any]) -> dict[str, Any]:
+        fields = {}
+        for key in ("instruction", "tokens", "quality", "sound_event", "ambient_sound"):
+            if key in message and message.get(key) is not None:
+                fields[key] = message.get(key)
+        if "language" in message and message.get("language") is not None:
+            fields["language"] = message.get("language")
+        return fields
+
+    def _build_generation_or_voice_clone_codes(
+        self,
+        message: Dict[str, Any],
+        n_vq: int,
+    ) -> torch.Tensor:
+        if "text" not in message:
+            raise ValueError("Direct MOSS-TTS-Local-Transformer-v1.5 generation requires messages built by build_user_message(...).")
+        text = "" if message.get("text") is None else str(message.get("text"))
+        prompt_fields = self._prompt_fields_from_user_message(message)
+        language_code = message.get("language")
+        audio_codes_list = self._resolve_audio_items(message.get("audio_codes_list", []), n_vq)
+        text_token_ids = self._encode_text(text)
+
+        if audio_codes_list:
+            parts: list[torch.Tensor] = [self._build_text_rows(
+                self._user_prompt_prefix_ids(),
+                device=audio_codes_list[0].device,
+            )]
+            for reference_codes in audio_codes_list:
+                parts.append(self._build_text_rows([int(self.model_config.audio_start_token_id)], device=reference_codes.device))
+                parts.append(self._build_audio_rows(reference_codes, int(self.model_config.audio_user_slot_token_id)))
+                parts.append(self._build_text_rows([int(self.model_config.audio_end_token_id)], device=reference_codes.device))
+            parts.append(
+                self._build_text_rows(
+                    self._user_prompt_after_reference_ids(language_code, prompt_fields)
+                    + text_token_ids
+                    + self._assistant_prompt_prefix_ids()
+                    + [int(self.model_config.audio_start_token_id)],
+                    device=audio_codes_list[0].device,
+                )
+            )
+            return torch.cat(parts, dim=0)
+
+        prompt_token_ids = (
+            self._user_prompt_prefix_ids()
+            + self._encode_text("None")
+            + self._user_prompt_after_reference_ids(language_code, prompt_fields)
+            + text_token_ids
+            + self._assistant_prompt_prefix_ids()
+            + [int(self.model_config.audio_start_token_id)]
+        )
+        return self._build_text_rows(prompt_token_ids)
+
+    def _build_continuation_codes(
+        self,
+        conversation: list[Dict[str, Any]],
+        n_vq: int,
+    ) -> torch.Tensor:
+        if len(conversation) < 2:
+            raise ValueError("continuation mode requires a user message followed by an assistant audio message.")
+        user_message = conversation[-2]
+        assistant_message = conversation[-1]
+        if user_message.get("role") != "user" or assistant_message.get("role") != "assistant":
+            raise ValueError("continuation mode requires the last two messages to be user, assistant.")
+        if "text" not in user_message:
+            raise ValueError("Direct MOSS-TTS-Local-Transformer-v1.5 continuation requires user messages built by build_user_message(...).")
+
+        text = "" if user_message.get("text") is None else str(user_message.get("text"))
+        prompt_fields = self._prompt_fields_from_user_message(user_message)
+        language_code = user_message.get("language")
+        prompt_token_ids = (
+            self._user_prompt_prefix_ids()
+            + self._encode_text("None")
+            + self._user_prompt_after_reference_ids(language_code, prompt_fields)
+            + self._encode_text(text)
+            + self._assistant_prompt_prefix_ids()
+            + [int(self.model_config.audio_start_token_id)]
+        )
+        audio_codes_list = self._resolve_audio_items(assistant_message.get("audio_codes_list", []), n_vq)
+        if not audio_codes_list:
+            return self._build_text_rows(prompt_token_ids)
+        if len(audio_codes_list) != 1:
+            raise ValueError("The MOSS-TTS-Local-Transformer-v1.5 continuation path expects exactly one prompt audio item.")
+        prompt_audio_codes = audio_codes_list[0]
+        return torch.cat(
+            [
+                self._build_text_rows(prompt_token_ids, device=prompt_audio_codes.device),
+                self._build_audio_rows(prompt_audio_codes, int(self.model_config.audio_assistant_slot_token_id)),
+            ],
+            dim=0,
+        )
+
+    def _try_build_direct_codes(
+        self,
+        conversation: list[Dict[str, Any]],
+        mode: str,
+        n_vq: int,
+    ) -> Optional[torch.Tensor]:
+        if mode == "generation" and len(conversation) == 1 and conversation[-1].get("role") == "user":
+            if "text" in conversation[-1]:
+                return self._build_generation_or_voice_clone_codes(conversation[-1], n_vq)
+            return None
+        if mode == "continuation" and len(conversation) >= 2:
+            if "text" in conversation[-2]:
+                return self._build_continuation_codes(conversation, n_vq)
+            return None
+        return None
+
+    def __call__(self, *args, **kwargs) -> BatchFeature:
+        conversations = args[0] if args else kwargs.pop("conversations")
+        mode: str = kwargs.pop("mode", "generation")
+        apply_chat_template: bool = kwargs.pop("apply_chat_template", True)
+        n_vq = self._assert_fixed_nq(kwargs.pop("n_vq", None))
+
+        kwargs.pop("return_tensors", None)
+        kwargs.pop("padding", None)
+        kwargs.pop("truncation", None)
+
+        if mode not in {"generation", "continuation", "computing_loss"}:
+            raise ValueError(f"Unsupported mode: {mode}")
+        if isinstance(conversations, (Message, dict)):
+            conversations = [conversations]
+        elif isinstance(conversations, list) and conversations and all(
+            isinstance(item, (Message, dict)) for item in conversations
+        ):
+            conversations = [conversations]
+
+        input_ids_list: list[torch.Tensor] = []
+        for conversation in conversations:
+            if isinstance(conversation, (Message, dict)):
+                conversation = [conversation]
+            conversation = [self._normalize_message(message) for message in conversation]
+
+            if (mode == "generation") ^ (conversation[-1]["role"] == "user"):
+                raise ValueError("generation mode must end with a user message.")
+            if mode == "continuation" and conversation[-1]["role"] != "assistant":
+                raise ValueError("continuation mode must end with an assistant message.")
+
+            direct_codes = self._try_build_direct_codes(conversation, mode, n_vq)
+            if direct_codes is not None:
+                input_ids_list.append(direct_codes)
+                continue
+
+            unified_parts = []
+            for message_idx, message in enumerate(conversation):
+                content = str(message["content"])
+                if apply_chat_template:
+                    add_generation_prompt = mode == "generation" and message_idx == len(conversation) - 1
+                    try:
+                        content = self.tokenizer.apply_chat_template(
+                            [{"role": message["role"], "content": content}],
+                            add_generation_prompt=add_generation_prompt,
+                            tokenize=False,
+                        )
+                    except Exception:
+                        logger.warning("apply_chat_template failed; falling back to raw message content.")
+
+                raw_audio_items = message.get("audio_codes_list", [])
+                audio_codes_list = self._resolve_audio_items(raw_audio_items, n_vq)
+                unified_parts.append(
+                    self._get_unified_codes(
+                        role=message["role"],
+                        content=content,
+                        audio_codes_list=audio_codes_list,
+                        truncation=(mode == "continuation"),
+                    )
+                )
+
+            unified_codes = torch.cat(unified_parts, dim=0)
+            if mode == "generation":
+                audio_start_row = torch.full(
+                    (1, n_vq + 1),
+                    int(self.model_config.audio_pad_token_id),
+                    dtype=unified_codes.dtype,
+                    device=unified_codes.device,
+                )
+                audio_start_row[:, 0] = int(self.model_config.audio_start_token_id)
+                unified_codes = torch.cat([unified_codes, audio_start_row], dim=0)
+            input_ids_list.append(unified_codes)
+
+        return BatchFeature(data=self._pad(input_ids_list))
+
+    def _normalize_message(self, message: Union[Message, Dict[str, Any]]) -> Dict[str, Any]:
+        if isinstance(message, Message):
+            return message.to_dict()
+        if not isinstance(message, dict):
+            raise TypeError("Each message must be a Message or dict.")
+        if "content" in message and "audio_codes_list" in message:
+            return message
+        role = message.get("role")
+        if role == "user":
+            return self.build_user_message(**{key: message.get(key) for key in USER_MESSAGE_FIELDS})
+        if role == "assistant":
+            return self.build_assistant_message(
+                audio_codes_list=message.get("audio_codes_list", []),
+                content=message.get("content", AUDIO_PLACEHOLDER),
+            )
+        raise ValueError(f"Unsupported role: {role}")
+
+    def _resolve_audio_items(
+        self,
+        raw_audio_items: list[Any],
+        n_vq: int,
+    ) -> list[torch.Tensor]:
+        if not raw_audio_items:
+            return []
+        resolved: list[Optional[torch.Tensor]] = [None] * len(raw_audio_items)
+        paths: list[str] = []
+        path_positions: list[int] = []
+        for index, item in enumerate(raw_audio_items):
+            if isinstance(item, torch.Tensor):
+                if item.ndim != 2 or int(item.shape[1]) != n_vq:
+                    raise ValueError(f"audio code tensor must have shape [T, {n_vq}], got {tuple(item.shape)}.")
+                resolved[index] = item.to(dtype=torch.long).cpu()
+            elif isinstance(item, (str, os.PathLike)):
+                paths.append(str(item))
+                path_positions.append(index)
+            else:
+                raise TypeError("Audio items must be tensors or path-like values.")
+        if paths:
+            encoded = self.encode_audios_from_path(paths, n_vq=n_vq)
+            for position, codes in zip(path_positions, encoded):
+                resolved[position] = codes
+        return [cast(torch.Tensor, item) for item in resolved]
+
+    def _pad(self, input_ids_list: list[torch.Tensor]) -> Dict[str, torch.Tensor]:
+        device = input_ids_list[0].device
+        lengths = torch.tensor([item.shape[0] for item in input_ids_list], device=device)
+        padded = torch.nn.utils.rnn.pad_sequence(
+            input_ids_list,
+            batch_first=True,
+            padding_value=int(self.model_config.audio_pad_token_id),
+            padding_side="left",
+        )
+        left_pad_mask = (padded.shape[1] - lengths).unsqueeze(1) > torch.arange(
+            padded.shape[1],
+            device=device,
+        ).unsqueeze(0)
+        padded[..., 0][left_pad_mask] = int(self.model_config.pad_token_id)
+        attention_mask = torch.zeros(padded.shape[:2], dtype=torch.bool, device=device)
+        attention_mask[~left_pad_mask] = True
+        return {"input_ids": padded, "attention_mask": attention_mask}
+
+    @staticmethod
+    def _replace_audio_placeholders(
+        content: str,
+        lengths: list[int],
+        slot_token: str,
+        audio_start_token: str,
+        audio_end_token: str,
+    ) -> str:
+        placeholder_count = content.count(AUDIO_PLACEHOLDER)
+        if placeholder_count != len(lengths):
+            raise ValueError(
+                f"Number of {AUDIO_PLACEHOLDER} ({placeholder_count}) does not match "
+                f"audio item count ({len(lengths)})."
+            )
+        lengths_iter = iter(lengths)
+
+        def replacer(_: re.Match) -> str:
+            length = int(next(lengths_iter))
+            if length <= 0:
+                return f"{audio_start_token}{audio_end_token}"
+            return f"{audio_start_token}{slot_token * length}{audio_end_token}"
+
+        return re.sub(re.escape(AUDIO_PLACEHOLDER), replacer, content)
+
+    def _get_unified_codes(
+        self,
+        role: str,
+        content: str,
+        audio_codes_list: list[torch.Tensor],
+        truncation: bool,
+    ) -> torch.Tensor:
+        n_vq = int(self.model_config.n_vq)
+        slot_token = self.audio_user_slot_token if role == "user" else self.audio_assistant_slot_token
+        content = self._replace_audio_placeholders(
+            content=content,
+            lengths=[int(codes.shape[0]) for codes in audio_codes_list],
+            slot_token=slot_token,
+            audio_start_token=self.audio_start_token,
+            audio_end_token=self.audio_end_token,
+        )
+        text_codes = torch.tensor(
+            self.tokenizer.encode(content),
+            dtype=torch.long,
+            device=audio_codes_list[0].device if audio_codes_list else None,
+        )
+
+        audio_start_indices = torch.where(text_codes == int(self.model_config.audio_start_token_id))[0]
+        audio_end_indices = torch.where(text_codes == int(self.model_config.audio_end_token_id))[0]
+        if len(audio_start_indices) != len(audio_codes_list) or len(audio_end_indices) != len(audio_codes_list):
+            raise ValueError("Audio placeholders do not match the encoded audio spans.")
+
+        if not audio_codes_list:
+            audio_codes = torch.full(
+                (len(text_codes), n_vq),
+                int(self.model_config.audio_pad_token_id),
+                dtype=torch.long,
+                device=text_codes.device,
+            )
+        else:
+            pieces: list[torch.Tensor] = []
+            prefix_idx = 0
+            for start_t, end_t, codes in zip(audio_start_indices, audio_end_indices, audio_codes_list):
+                start_idx = int(start_t.item())
+                end_idx = int(end_t.item())
+                pad_before = torch.full(
+                    (start_idx - prefix_idx + 1, n_vq),
+                    int(self.model_config.audio_pad_token_id),
+                    dtype=torch.long,
+                    device=codes.device,
+                )
+                pieces.extend([pad_before, codes.to(dtype=torch.long)])
+                prefix_idx = end_idx
+            if truncation:
+                trailing = torch.zeros(
+                    (0, n_vq),
+                    dtype=torch.long,
+                    device=audio_codes_list[0].device,
+                )
+            else:
+                last_end = int(audio_end_indices[-1].item())
+                trailing = torch.full(
+                    (len(text_codes) - last_end, n_vq),
+                    int(self.model_config.audio_pad_token_id),
+                    dtype=torch.long,
+                    device=audio_codes_list[0].device,
+                )
+            pieces.append(trailing)
+            audio_codes = torch.cat(pieces, dim=0)
+
+        if text_codes.shape[0] != audio_codes.shape[0]:
+            min_len = min(text_codes.shape[0], audio_codes.shape[0])
+            text_codes = text_codes[:min_len]
+            audio_codes = audio_codes[:min_len]
+        return torch.cat([text_codes.unsqueeze(1), audio_codes], dim=1)
+
+    def _parse_text_codes(self, start_length: int, text_codes: torch.LongTensor) -> str:
+        text = cast(str, self.tokenizer.decode(text_codes))
+        prefix = cast(str, self.tokenizer.decode(text_codes[:start_length]))
+        text = text[len(prefix):]
+        audio_pattern = re.compile(
+            rf"(?:{re.escape(self.audio_start_token)})?"
+            rf"(?:{re.escape(self.audio_assistant_slot_token)})*"
+            rf"{re.escape(self.audio_end_token)}"
+        )
+        return audio_pattern.sub(
+            lambda match: AUDIO_PLACEHOLDER if self.audio_assistant_slot_token in match.group(0) else "",
+            text,
+        )
+
+    def _parse_audio_codes(
+        self,
+        start_length: int,
+        audio_codes: torch.LongTensor,
+        *,
+        return_stereo: bool = True,
+    ) -> list[torch.Tensor]:
+        is_pad = audio_codes.eq(int(self.model_config.audio_pad_token_id)).all(dim=1)
+        non_pad = ~is_pad
+        if not bool(non_pad.any().item()):
+            return []
+        idx = torch.nonzero(non_pad).squeeze(1)
+        breaks = torch.where(idx[1:] != idx[:-1] + 1)[0] + 1
+        segment_indices = [idx] if breaks.numel() == 0 else list(torch.tensor_split(idx, breaks.cpu().tolist()))
+        code_segments = [audio_codes[segment] for segment in segment_indices]
+        decoded = self.decode_audio_codes(code_segments, return_stereo=return_stereo)
+
+        if start_length > 0 and code_segments and decoded:
+            first_code_length = int(code_segments[0].shape[0])
+            if first_code_length > 0:
+                trim_ratio = max(0.0, min(float(start_length) / float(first_code_length), 1.0))
+                if trim_ratio >= 1.0:
+                    decoded = decoded[1:]
+                elif trim_ratio > 0.0:
+                    trim_samples = int(decoded[0].shape[-1] * trim_ratio)
+                    decoded[0] = decoded[0][..., trim_samples:]
+        return decoded
+
+    def decode(self, output: Any, *, return_stereo: bool = True) -> list[Optional[AssistantMessage]]:
+        generated_messages: list[Optional[AssistantMessage]] = []
+        for start_length, generation_ids in output:
+            content = self._parse_text_codes(int(start_length), generation_ids[:, 0])
+            audio_codes_list = self._parse_audio_codes(
+                int(start_length),
+                generation_ids[:, 1:],
+                return_stereo=return_stereo,
+            )
+            if content == "":
+                generated_messages.append(None)
+            else:
+                generated_messages.append(
+                    AssistantMessage(
+                        content=content,
+                        audio_codes_list=cast(list[Union[str, torch.Tensor]], audio_codes_list),
+                    )
+                )
+        return generated_messages
+
+    @staticmethod
+    def loudness_normalize(
+        wav: torch.Tensor,
+        target_dbfs: float = -20.0,
+        gain_range: tuple[float, float] = (-3.0, 3.0),
+    ) -> torch.Tensor:
+        wav = wav.to(torch.float32)
+        if wav.numel() == 0:
+            return wav
+        current_dbfs = 10.0 * torch.log10(torch.mean(wav**2) + 1e-9)
+        gain = max(gain_range[0], min(float(target_dbfs - current_dbfs), gain_range[1]))
+        return wav * (10.0 ** (gain / 20.0))
+
+    def _get_audio_tokenizer_device(self) -> torch.device:
+        audio_tokenizer = getattr(self, "audio_tokenizer", None)
+        if audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set.")
+        try:
+            return next(audio_tokenizer.parameters()).device
+        except StopIteration:
+            return torch.device("cpu")
+
+    def encode_audios_from_wav(
+        self,
+        wav_list: Union[torch.Tensor, list[torch.Tensor]],
+        sampling_rate: int,
+        n_vq: Optional[int] = None,
+    ) -> list[torch.Tensor]:
+        n_vq = self._assert_fixed_nq(n_vq)
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set.")
+        if isinstance(wav_list, torch.Tensor):
+            wav_list = [wav_list]
+        target_sr = int(self.model_config.sampling_rate)
+        device = self._get_audio_tokenizer_device()
+        prepared = []
+        for wav in wav_list:
+            if wav.ndim == 1:
+                wav = wav.unsqueeze(0)
+            if wav.shape[0] == 1:
+                wav = wav.repeat(2, 1)
+            elif wav.shape[0] > 2:
+                wav = wav[:2]
+            if int(sampling_rate) != target_sr:
+                wav = torchaudio.functional.resample(wav, int(sampling_rate), target_sr)
+            prepared.append(self.loudness_normalize(wav).to(device))
+
+        if hasattr(self.audio_tokenizer, "batch_encode"):
+            encoded = self.audio_tokenizer.batch_encode(prepared, num_quantizers=n_vq)
+            audio_codes = encoded.audio_codes
+            audio_lengths = encoded.audio_codes_lengths
+        else:
+            max_len = max(int(wav.shape[-1]) for wav in prepared)
+            input_values = torch.zeros(len(prepared), 1, max_len, dtype=torch.float32, device=device)
+            padding_mask = torch.zeros(len(prepared), max_len, dtype=torch.bool, device=device)
+            for index, wav in enumerate(prepared):
+                input_values[index, 0, : wav.shape[-1]] = wav
+                padding_mask[index, : wav.shape[-1]] = True
+            encoded = self.audio_tokenizer.encode(
+                input_values,
+                padding_mask=padding_mask,
+                num_quantizers=n_vq,
+                return_dict=True,
+            )
+            audio_codes = encoded.audio_codes
+            audio_lengths = encoded.audio_codes_lengths
+
+        if audio_codes is None or audio_lengths is None:
+            raise RuntimeError("audio_tokenizer did not return audio_codes/audio_codes_lengths.")
+        result = []
+        for index in range(int(audio_codes.shape[1])):
+            length = int(audio_lengths[index].item())
+            result.append(audio_codes[:, index, :length].transpose(0, 1).contiguous().cpu().long())
+        return result
+
+    def encode_audios_from_path(
+        self,
+        wav_path_list: Union[str, os.PathLike, list[Union[str, os.PathLike]]],
+        n_vq: Optional[int] = None,
+    ) -> list[torch.Tensor]:
+        if isinstance(wav_path_list, (str, os.PathLike)):
+            wav_path_list = [wav_path_list]
+        wavs = []
+        target_sr = int(self.model_config.sampling_rate)
+        for wav_path in wav_path_list:
+            wav, sr = torchaudio.load(str(wav_path))
+            if int(sr) != target_sr:
+                wav = torchaudio.functional.resample(wav, int(sr), target_sr)
+            wavs.append(wav)
+        return self.encode_audios_from_wav(wavs, target_sr, n_vq=n_vq)
+
+    def decode_audio_codes(
+        self,
+        audio_tokens_list: Union[torch.Tensor, list[torch.Tensor]],
+        *,
+        return_stereo: bool = True,
+    ) -> list[torch.Tensor]:
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set.")
+        if isinstance(audio_tokens_list, torch.Tensor):
+            audio_tokens_list = [audio_tokens_list]
+        if not audio_tokens_list:
+            return []
+
+        n_vq = int(self.model_config.n_vq)
+        device = self._get_audio_tokenizer_device()
+        codes_list = [
+            codes[:, :n_vq].transpose(0, 1).contiguous().to(device=device, dtype=torch.long)
+            for codes in audio_tokens_list
+        ]
+        max_len = max(int(codes.shape[1]) for codes in codes_list)
+        audio_codes = torch.zeros(n_vq, len(codes_list), max_len, device=device, dtype=torch.long)
+        padding_mask = torch.zeros(len(codes_list), max_len, device=device, dtype=torch.bool)
+        for index, codes in enumerate(codes_list):
+            length = int(codes.shape[1])
+            audio_codes[:, index, :length] = codes
+            padding_mask[index, :length] = True
+
+        decoded = self.audio_tokenizer.decode(
+            audio_codes,
+            padding_mask=padding_mask,
+            num_quantizers=n_vq,
+            return_dict=True,
+            chunk_duration=8,
+        )
+        audio = decoded.audio
+        audio_lengths = decoded.audio_lengths
+        if audio is None or audio_lengths is None:
+            raise RuntimeError("audio_tokenizer.decode did not return audio/audio_lengths.")
+        wavs = []
+        for index in range(int(audio.shape[0])):
+            length = int(audio_lengths[index].item())
+            wav = audio[index, :, :length].contiguous().cpu().to(torch.float32)
+            if not return_stereo:
+                if wav.shape[0] == 1:
+                    wav = wav.squeeze(0)
+                else:
+                    wav = wav.mean(dim=0)
+            wavs.append(wav)
+        return wavs

assets/sft3/processor_config.json

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

assets/sft3/qwen3_decoder.py

@@ -0,0 +1,582 @@
+# coding=utf-8
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from safetensors.torch import load_file
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutputWithPast
+
+from .gpt2_decoder import PackedSequenceMetadata, MossTTSNanoGPT2Model
+
+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
+
+
+class MossQwen3RMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class MossQwen3RotaryEmbedding(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        head_dim = int(getattr(config, "head_dim", config.hidden_size // config.num_attention_heads))
+        rope_theta = getattr(config, "rope_theta", None)
+        if rope_theta is None:
+            rope_scaling = getattr(config, "rope_scaling", None)
+            if isinstance(rope_scaling, dict):
+                rope_theta = rope_scaling.get("rope_theta")
+        if rope_theta is None:
+            rope_theta = 1000000.0
+        rope_theta = float(rope_theta)
+        self.head_dim = head_dim
+        self.rope_theta = rope_theta
+        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)
+
+    def _compute_inv_freq(self, device: Optional[torch.device] = None) -> torch.Tensor:
+        return 1.0 / (
+            self.rope_theta ** (torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) / self.head_dim)
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.LongTensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        inv_freq = self._compute_inv_freq(device=hidden_states.device)
+        freqs = torch.einsum(
+            "bs,d->bsd",
+            position_ids.to(device=hidden_states.device, dtype=inv_freq.dtype),
+            inv_freq,
+        )
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb.cos().to(dtype=hidden_states.dtype), emb.sin().to(dtype=hidden_states.dtype)
+
+
+def rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
+    first_half = hidden_states[..., : hidden_states.shape[-1] // 2]
+    second_half = hidden_states[..., hidden_states.shape[-1] // 2 :]
+    return torch.cat((-second_half, first_half), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    cos = cos.unsqueeze(-2)
+    sin = sin.unsqueeze(-2)
+    query = (query * cos) + (rotate_half(query) * sin)
+    key = (key * cos) + (rotate_half(key) * sin)
+    return query, key
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    if n_rep == 1:
+        return hidden_states
+    batch, seq_len, num_key_value_heads, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:, :, :, None, :].expand(batch, seq_len, num_key_value_heads, n_rep, head_dim)
+    return hidden_states.reshape(batch, seq_len, num_key_value_heads * n_rep, head_dim)
+
+
+class MossQwen3MLP(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.act_fn = ACT2FN[getattr(config, "hidden_act", "silu")]
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+
+
+class MossQwen3Attention(nn.Module):
+    def __init__(self, config, layer_idx: int) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = int(layer_idx)
+        self.hidden_size = int(config.hidden_size)
+        self.num_heads = int(config.num_attention_heads)
+        self.num_key_value_heads = int(config.num_key_value_heads)
+        self.head_dim = int(getattr(config, "head_dim", self.hidden_size // self.num_heads))
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.scaling = self.head_dim ** -0.5
+        self.attention_dropout = float(getattr(config, "attention_dropout", 0.0))
+        self.attn_implementation = str(getattr(config, "_attn_implementation", "eager"))
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=bool(config.attention_bias))
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=bool(config.attention_bias),
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=bool(config.attention_bias),
+        )
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=bool(config.attention_bias))
+        self.q_norm = MossQwen3RMSNorm(self.head_dim, eps=float(config.rms_norm_eps))
+        self.k_norm = MossQwen3RMSNorm(self.head_dim, eps=float(config.rms_norm_eps))
+
+    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:
+        key = repeat_kv(key, self.num_key_value_groups)
+        value = repeat_kv(value, self.num_key_value_groups)
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        scores = torch.matmul(query, key.transpose(-1, -2)) * self.scaling
+        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(~mask, torch.finfo(scores.dtype).min)
+        probs = torch.softmax(scores, dim=-1)
+        if self.training and self.attention_dropout > 0:
+            probs = torch.dropout(probs, self.attention_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:
+        key = repeat_kv(key, self.num_key_value_groups)
+        value = repeat_kv(value, self.num_key_value_groups)
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        mask = None
+        if attention_mask is not None or query.shape[-2] != key.shape[-2]:
+            mask = self._causal_attention_mask(
+                attention_mask=attention_mask,
+                query_length=query.shape[-2],
+                key_length=key.shape[-2],
+                device=query.device,
+            )
+        output = torch.nn.functional.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=mask is None,
+            scale=self.scaling,
+        )
+        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, got dtype={query.dtype}.")
+
+        dropout_p = self.attention_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_key_value_heads, self.head_dim).index_select(0, packed_metadata.indices)
+                value = value.reshape(-1, self.num_key_value_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)
+
+        if query.shape[1] != key.shape[1]:
+            query_attention_mask = attention_mask[:, -query.shape[1] :]
+            unpadded_query, query_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
+                query,
+                query_attention_mask,
+            )
+            unpadded_key, _, cu_seqlens_k, max_seqlen_k, _ = 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_q,
+                cu_seqlens_k,
+                max_seqlen_q,
+                max_seqlen_k,
+                dropout_p=dropout_p,
+                causal=True,
+            )
+            return pad_input(output, query_indices, query.shape[0], query.shape[1])
+
+        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,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = 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]]]:
+        input_shape = hidden_states.shape[:-1]
+        query_states = self.q_norm(
+            self.q_proj(hidden_states).view(*input_shape, self.num_heads, self.head_dim)
+        )
+        key_states = self.k_norm(
+            self.k_proj(hidden_states).view(*input_shape, self.num_key_value_heads, self.head_dim)
+        )
+        value_states = self.v_proj(hidden_states).view(*input_shape, self.num_key_value_heads, self.head_dim)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if layer_past is not None:
+            past_key, past_value = layer_past
+            key_states = torch.cat([past_key.to(device=key_states.device, dtype=key_states.dtype), key_states], dim=1)
+            value_states = torch.cat(
+                [past_value.to(device=value_states.device, dtype=value_states.dtype), value_states],
+                dim=1,
+            )
+
+        present = (key_states, value_states) if use_cache else None
+        if self.attn_implementation == "flash_attention_2":
+            attn_output = self._flash_attention(
+                query=query_states,
+                key=key_states,
+                value=value_states,
+                attention_mask=attention_mask,
+                packed_metadata=packed_metadata,
+            )
+        elif self.attn_implementation == "sdpa":
+            attn_output = self._sdpa_attention(
+                query=query_states,
+                key=key_states,
+                value=value_states,
+                attention_mask=attention_mask,
+            )
+        else:
+            attn_output = self._eager_attention(
+                query=query_states,
+                key=key_states,
+                value=value_states,
+                attention_mask=attention_mask,
+            )
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        return self.o_proj(attn_output), present
+
+
+class MossQwen3DecoderLayer(nn.Module):
+    def __init__(self, config, layer_idx: int) -> None:
+        super().__init__()
+        self.self_attn = MossQwen3Attention(config=config, layer_idx=layer_idx)
+        self.mlp = MossQwen3MLP(config)
+        self.input_layernorm = MossQwen3RMSNorm(config.hidden_size, eps=float(config.rms_norm_eps))
+        self.post_attention_layernorm = MossQwen3RMSNorm(config.hidden_size, eps=float(config.rms_norm_eps))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        packed_metadata: Optional[PackedSequenceMetadata] = None,
+        layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        attn_output, present = self.self_attn(
+            hidden_states=hidden_states,
+            position_embeddings=position_embeddings,
+            attention_mask=attention_mask,
+            packed_metadata=packed_metadata,
+            layer_past=layer_past,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + attn_output
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + self.mlp(hidden_states)
+        return hidden_states, present
+
+
+class MossQwen3Model(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.config = config
+        self.attn_implementation = str(getattr(config, "_attn_implementation", "eager"))
+        self.padding_idx = getattr(config, "pad_token_id", None)
+        self.vocab_size = int(config.vocab_size)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [MossQwen3DecoderLayer(config, layer_idx=index) for index in range(config.num_hidden_layers)]
+        )
+        self.norm = MossQwen3RMSNorm(config.hidden_size, eps=float(config.rms_norm_eps))
+        self.rotary_emb = MossQwen3RotaryEmbedding(config)
+        self.gradient_checkpointing = False
+        self.gradient_checkpointing_use_reentrant = bool(
+            getattr(config, "gradient_checkpointing_use_reentrant", False)
+        )
+        self._reset_parameters()
+
+    def _reset_parameters(self) -> None:
+        init_std = float(getattr(self.config, "initializer_range", 0.02))
+        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)
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def load_qwen3_pretrained_weights(self, pretrained_path: str) -> None:
+        model_dir = Path(pretrained_path)
+        index_path = model_dir / "model.safetensors.index.json"
+        if not index_path.exists():
+            raise FileNotFoundError(f"Missing Qwen3 safetensors index: {index_path}")
+        with index_path.open("r", encoding="utf-8") as handle:
+            index = json.load(handle)
+        weight_map = index.get("weight_map", {})
+        shard_to_keys: dict[str, list[str]] = {}
+        for key, shard in weight_map.items():
+            if not key.startswith("model."):
+                continue
+            shard_to_keys.setdefault(str(shard), []).append(key)
+
+        state_dict = self.state_dict()
+        loaded_state = {}
+        for shard, keys in sorted(shard_to_keys.items()):
+            shard_tensors = load_file(str(model_dir / shard), device="cpu")
+            for key in keys:
+                target_key = key[len("model.") :]
+                if target_key not in state_dict:
+                    continue
+                tensor = shard_tensors[key]
+                if tuple(tensor.shape) != tuple(state_dict[target_key].shape):
+                    raise ValueError(
+                        f"Shape mismatch while loading Qwen3 weight {key}: "
+                        f"checkpoint={tuple(tensor.shape)} model={tuple(state_dict[target_key].shape)}"
+                    )
+                loaded_state[target_key] = tensor
+
+        missing, unexpected = self.load_state_dict(loaded_state, strict=False)
+        unexpected = [key for key in unexpected if key]
+        if unexpected:
+            raise RuntimeError(f"Unexpected Qwen3 pretrained keys after load: {unexpected[:10]}")
+        missing = [key for key in missing if key not in loaded_state]
+        if missing:
+            raise RuntimeError(f"Missing Qwen3 pretrained keys after load: {missing[:10]}")
+
+    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 = MossTTSNanoGPT2Model.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,
+                    sequence_length=hidden_states.shape[1],
+                )
+            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 = MossTTSNanoGPT2Model.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,
+            )
+
+        hidden_states = hidden_states * query_attention_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        all_hidden_states = () if output_hidden_states else None
+        presents = [] if use_cache else None
+        for layer_index, decoder_layer in enumerate(self.layers):
+            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, _ = decoder_layer(
+                        hidden_states=inputs[0],
+                        attention_mask=inputs[1],
+                        packed_metadata=packed_metadata,
+                        layer_past=None,
+                        use_cache=False,
+                        position_embeddings=position_embeddings,
+                    )
+                    return output
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    custom_forward,
+                    hidden_states,
+                    attention_mask,
+                    use_reentrant=self.gradient_checkpointing_use_reentrant,
+                )
+                present = None
+            else:
+                hidden_states, present = decoder_layer(
+                    hidden_states=hidden_states,
+                    attention_mask=attention_mask,
+                    packed_metadata=packed_metadata,
+                    layer_past=None if past_key_values is None else past_key_values[layer_index],
+                    use_cache=use_cache,
+                    position_embeddings=position_embeddings,
+                )
+
+            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.norm(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,
+        )

assets/sft3/special_tokens_map.json

@@ -0,0 +1,21 @@
+{
+  "additional_special_tokens": [
+    "<|audio_start|>",
+    "<|audio_end|>",
+    "<|audio_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

assets/sft3/tokenizer.json

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

assets/sft3/tokenizer_config.json

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

checkpoints/S1/metadata.json

@@ -0,0 +1,17 @@
+{
+  "stage": "S1",
+  "bytes": 1431076525,
+  "sha256": "89cf9ee92e83ea8fda7661c53c2346e68b08f407ac359248dabdf42a8c01c602",
+  "step": 16265,
+  "loss": 4.805444754959801,
+  "original_training_state_sha256": "e33af9a8bfcace52d36072dc90db4847a764da709fcf781775c4e40db3de4b96",
+  "contract": {
+    "plan_sha256": "277eb0027da710e3b0141c131090792fa6c2da8797ab5a4483bfe8e4f846ddd9",
+    "schema_sha256": "8eee17be1be56dcc4c3b5e8ac95b5494dd147451051f2d9a0c34d418f354b9e7",
+    "manifest_sha256": "b7d691148613bc04bbcf942d7e63152fe5669c922aff2400c1efdbe5b815227c",
+    "world": 128,
+    "objective": "global_token_mean_continuous_ladder_v1",
+    "seed": 20260923
+  },
+  "format": "torch.save of base.state_dict() in bfloat16; weights_only=True"
+}

checkpoints/S10/metadata.json

@@ -0,0 +1,17 @@
+{
+  "stage": "S10",
+  "bytes": 1431076525,
+  "sha256": "d92965a2656b35a38bbf76aee545bcd64b813235549b29bab4a64b09a256acb4",
+  "step": 41936,
+  "loss": 4.727833971921586,
+  "original_training_state_sha256": "1fbb313660f71bf74aca97237dc0a77e0852193875a4ba32d56e6dfb1ae03c55",
+  "contract": {
+    "plan_sha256": "277eb0027da710e3b0141c131090792fa6c2da8797ab5a4483bfe8e4f846ddd9",
+    "schema_sha256": "8eee17be1be56dcc4c3b5e8ac95b5494dd147451051f2d9a0c34d418f354b9e7",
+    "manifest_sha256": "b7d691148613bc04bbcf942d7e63152fe5669c922aff2400c1efdbe5b815227c",
+    "world": 128,
+    "objective": "global_token_mean_continuous_ladder_v1",
+    "seed": 20260923
+  },
+  "format": "torch.save of base.state_dict() in bfloat16; weights_only=True"
+}

checkpoints/S2/metadata.json

@@ -0,0 +1,17 @@
+{
+  "stage": "S2",
+  "bytes": 1431076525,
+  "sha256": "b05ee43210cb35629368229adb00bdb45546aad47c0a640757e82be480162d0c",
+  "step": 26284,
+  "loss": 4.723223350145455,
+  "original_training_state_sha256": "bed312f0116c75bbbd2ee87af2c3d799b07dce97bb0f47a2e184c3d1d7adf0d3",
+  "contract": {
+    "plan_sha256": "277eb0027da710e3b0141c131090792fa6c2da8797ab5a4483bfe8e4f846ddd9",
+    "schema_sha256": "8eee17be1be56dcc4c3b5e8ac95b5494dd147451051f2d9a0c34d418f354b9e7",
+    "manifest_sha256": "b7d691148613bc04bbcf942d7e63152fe5669c922aff2400c1efdbe5b815227c",
+    "world": 128,
+    "objective": "global_token_mean_continuous_ladder_v1",
+    "seed": 20260923
+  },
+  "format": "torch.save of base.state_dict() in bfloat16; weights_only=True"
+}

checkpoints/S3/metadata.json

@@ -0,0 +1,17 @@
+{
+  "stage": "S3",
+  "bytes": 1431076525,
+  "sha256": "60079723bf797a81107def1e6fc479e5353b3cc51f3780f98ea65a9f5c109f97",
+  "step": 32337,
+  "loss": 4.704502838977402,
+  "original_training_state_sha256": "02a62d17867616531a0af1a8663841e6bd638d0cb39e526f0b7283aa65440dd9",
+  "contract": {
+    "plan_sha256": "277eb0027da710e3b0141c131090792fa6c2da8797ab5a4483bfe8e4f846ddd9",
+    "schema_sha256": "8eee17be1be56dcc4c3b5e8ac95b5494dd147451051f2d9a0c34d418f354b9e7",
+    "manifest_sha256": "b7d691148613bc04bbcf942d7e63152fe5669c922aff2400c1efdbe5b815227c",
+    "world": 128,
+    "objective": "global_token_mean_continuous_ladder_v1",
+    "seed": 20260923
+  },
+  "format": "torch.save of base.state_dict() in bfloat16; weights_only=True"
+}

checkpoints/S4/metadata.json

@@ -0,0 +1,17 @@
+{
+  "stage": "S4",
+  "bytes": 1431076525,
+  "sha256": "c6d07d86ec057844cbbcfd5f6d67b152ddff92df9fcd0063e25de032ed33d6ce",
+  "step": 37182,
+  "loss": 4.69656242843785,
+  "original_training_state_sha256": "59b9f2cbea9cf14b9651bdf587bb9d6622872087d5c6e9997975074987d16273",
+  "contract": {
+    "plan_sha256": "277eb0027da710e3b0141c131090792fa6c2da8797ab5a4483bfe8e4f846ddd9",
+    "schema_sha256": "8eee17be1be56dcc4c3b5e8ac95b5494dd147451051f2d9a0c34d418f354b9e7",
+    "manifest_sha256": "b7d691148613bc04bbcf942d7e63152fe5669c922aff2400c1efdbe5b815227c",
+    "world": 128,
+    "objective": "global_token_mean_continuous_ladder_v1",
+    "seed": 20260923
+  },
+  "format": "torch.save of base.state_dict() in bfloat16; weights_only=True"
+}

checkpoints/S5/metadata.json

@@ -0,0 +1,17 @@
+{
+  "stage": "S5",
+  "bytes": 1431077587,
+  "sha256": "b90981030618ea64f2b0e5cdc4167de222182be4ad1f3dc8b358ebb59d778fba",
+  "step": 38183,
+  "loss": 4.698665114248258,
+  "original_training_state_sha256": "0b844c61ccd4a1f6f052d749dcbb9d5d7922c49c832d8a81bb18a06e4cdfc2fc",
+  "contract": {
+    "plan_sha256": "277eb0027da710e3b0141c131090792fa6c2da8797ab5a4483bfe8e4f846ddd9",
+    "schema_sha256": "8eee17be1be56dcc4c3b5e8ac95b5494dd147451051f2d9a0c34d418f354b9e7",
+    "manifest_sha256": "b7d691148613bc04bbcf942d7e63152fe5669c922aff2400c1efdbe5b815227c",
+    "world": 128,
+    "objective": "global_token_mean_continuous_ladder_v1",
+    "seed": 20260923
+  },
+  "format": "torch.save of base.state_dict() in bfloat16; weights_only=True"
+}

checkpoints/S6/metadata.json

@@ -0,0 +1,17 @@
+{
+  "stage": "S6",
+  "bytes": 1431076525,
+  "sha256": "49b4bbe028faaf99d2501e234db7fd8e5faa1e004b7acbf9a96f1f0cb054c65c",
+  "step": 40935,
+  "loss": 4.700463458320118,
+  "original_training_state_sha256": "ff4ef3c3f0559b38351251f951f8476f60402ee61f6a942834a03f90386dcf80",
+  "contract": {
+    "plan_sha256": "277eb0027da710e3b0141c131090792fa6c2da8797ab5a4483bfe8e4f846ddd9",
+    "schema_sha256": "8eee17be1be56dcc4c3b5e8ac95b5494dd147451051f2d9a0c34d418f354b9e7",
+    "manifest_sha256": "b7d691148613bc04bbcf942d7e63152fe5669c922aff2400c1efdbe5b815227c",
+    "world": 128,
+    "objective": "global_token_mean_continuous_ladder_v1",
+    "seed": 20260923
+  },
+  "format": "torch.save of base.state_dict() in bfloat16; weights_only=True"
+}

checkpoints/S7/metadata.json

@@ -0,0 +1,17 @@
+{
+  "stage": "S7",
+  "bytes": 1431076525,
+  "sha256": "ec2dc60f6df2b8b0a842b35048a9b6c6bb6bedf87640f5affdeea49899009b6d",
+  "step": 41466,
+  "loss": 4.705551250746067,
+  "original_training_state_sha256": "977654975fc889d3c726477b0159740b9c1d4b579434a82b7ca382f4e5d710ff",
+  "contract": {
+    "plan_sha256": "277eb0027da710e3b0141c131090792fa6c2da8797ab5a4483bfe8e4f846ddd9",
+    "schema_sha256": "8eee17be1be56dcc4c3b5e8ac95b5494dd147451051f2d9a0c34d418f354b9e7",
+    "manifest_sha256": "b7d691148613bc04bbcf942d7e63152fe5669c922aff2400c1efdbe5b815227c",
+    "world": 128,
+    "objective": "global_token_mean_continuous_ladder_v1",
+    "seed": 20260923
+  },
+  "format": "torch.save of base.state_dict() in bfloat16; weights_only=True"
+}

checkpoints/S8/metadata.json

@@ -0,0 +1,17 @@
+{
+  "stage": "S8",
+  "bytes": 1431076525,
+  "sha256": "7ea0868438d1fbae7cafc4aad366199cfb5f535c4a7ed17f250d48683ae94dad",
+  "step": 41739,
+  "loss": 4.710860909975666,
+  "original_training_state_sha256": "86285dc461cefb437c0d29a75100d3232982a215f73a3ab23bfb471a3a8f1fde",
+  "contract": {
+    "plan_sha256": "277eb0027da710e3b0141c131090792fa6c2da8797ab5a4483bfe8e4f846ddd9",
+    "schema_sha256": "8eee17be1be56dcc4c3b5e8ac95b5494dd147451051f2d9a0c34d418f354b9e7",
+    "manifest_sha256": "b7d691148613bc04bbcf942d7e63152fe5669c922aff2400c1efdbe5b815227c",
+    "world": 128,
+    "objective": "global_token_mean_continuous_ladder_v1",
+    "seed": 20260923
+  },
+  "format": "torch.save of base.state_dict() in bfloat16; weights_only=True"
+}

checkpoints/S9/metadata.json

@@ -0,0 +1,17 @@
+{
+  "stage": "S9",
+  "bytes": 1431076525,
+  "sha256": "4109129bc881584e27b8ffd0f073f7f3536ba472042fd77fe84a42f9cdb1c008",
+  "step": 41878,
+  "loss": 4.719842956252719,
+  "original_training_state_sha256": "d9ed815939a42d2cfe16ee212a435cc38ed134c0d28f0345fbfd5105f823aafd",
+  "contract": {
+    "plan_sha256": "277eb0027da710e3b0141c131090792fa6c2da8797ab5a4483bfe8e4f846ddd9",
+    "schema_sha256": "8eee17be1be56dcc4c3b5e8ac95b5494dd147451051f2d9a0c34d418f354b9e7",
+    "manifest_sha256": "b7d691148613bc04bbcf942d7e63152fe5669c922aff2400c1efdbe5b815227c",
+    "world": 128,
+    "objective": "global_token_mean_continuous_ladder_v1",
+    "seed": 20260923
+  },
+  "format": "torch.save of base.state_dict() in bfloat16; weights_only=True"
+}

code/anno_timed.py

@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+"""Timed-script annotation engine (v2).
+
+WHY.  Listening to the SFT+DPO model surfaced two failure modes: vocal bursts that run far too
+long and sound unnatural, and occasional hallucinated content that was never prompted.  Both are
+BUDGET failures -- nothing in the v1 prompt told the model how long anything was allowed to take.
+v1 carried at most one optional `[7.3 seconds duration]` prefix for the whole clip and bare
+`(sigh)` cues with no duration at all.
+
+v2 renders the target text as a TIMED SCRIPT built from the word-level timestamps and the
+DETECTED burst spans, so every second of the clip is accounted for:
+
+    [0.9 seconds pause] [3.4 seconds duration] Die Handelskammer betreibt diesen Test.
+    [0.6 seconds pause] (contented sigh, 0.4 seconds) [1.2 seconds pause]
+    [4.1 seconds duration] Die entwickeln und fuehren ihn durch, wissen Sie.
+
+Rules, each asserted by `verify()` on the EMITTED string:
+  1. `[D seconds duration]` precedes a speech segment and equals that SEGMENT's last-word offset
+     minus its first-word onset -- never the clip's.
+  2. Segments break at sentence-final punctuation.  A segment still longer than MAX_SEG_S is
+     split again at its largest internal word gap, recursively.  A 20 s two-sentence utterance
+     therefore renders as `[12.0 ...] <s1> [8.0 ...] <s2>`.
+  3. Every gap >= PAUSE_MIN between consecutive events -- including before the first and after
+     the last -- is emitted as `[G seconds pause]`.
+  4. Every detected burst is emitted at its position, as `(label, D seconds)`.
+  5. The emitted numbers sum to the clip duration.
+
+FOUR THINGS A NAIVE IMPLEMENTATION GETS WRONG, all measured on this corpus:
+
+  * The transcript must not be rebuilt by joining word tokens.  The forced aligner drops
+    numerals ("154 Euro" aligns as "Euro"), so a join silently loses text on 9.9 % of
+    real-speech rows.  Words are ALIGNED to the original string by character offset and the
+    original string is SLICED, so output text is byte-identical apart from inserted tags.
+  * The script is built from DETECTED bursts, not from the parenthetical cues already sitting in
+    `text_with_bursts`.  On the voice-profile corpus those parentheticals are SYNTHESIS
+    DIRECTIONS -- 39.4 % were never confirmed by the burst detector.  A target text carrying a
+    cue that is not in the audio teaches precisely the behaviour being removed.
+  * 49.0 % of detected bursts OVERLAP an aligned word span.  Adding such a burst's full length
+    to the budget double-counts time -- it drove the sum 2.1 s past the clip length on the worst
+    row of the first test.  A burst's printed duration is its length MINUS its overlap with
+    speech; when nothing is left it is printed WITHOUT a number, which is the honest statement
+    that it is co-articulated and has no slot of its own.
+  * A gap SHORTER than PAUSE_MIN is not worth a tag, but skipping it breaks the sum and the
+    residual then reads as an error in the duration values.  Such a gap is absorbed into the
+    adjacent speech segment (never into a burst -- a burst tag must state the burst's own length).
+"""
+import re
+
+PAUSE_MIN   = 0.20     # gaps below this are absorbed rather than tagged
+MAX_SEG_S   = 12.0     # a speech segment longer than this is split again at its largest gap
+MIN_SPLIT_S = 2.0      # never create a segment shorter than this by splitting
+BURST_MIN_S = 0.05     # below this a burst prints without a duration
+
+_SENT_END = re.compile(r"[.!?…]+[\"'“”’\)\]]*$")
+
+
+def _fmt(x):
+    return f"{x:.1f}"
+
+
+def align_words(text, words):
+    """-> [(char_start, char_end, t_start, t_end)] for the word tokens locatable in `text`."""
+    out, pos = [], 0
+    for w in words or []:
+        tok = (w.get("w") or "").strip()
+        if not tok:
+            continue
+        try:
+            s, e = float(w["s"]), float(w["e"])
+        except Exception:
+            continue
+        if not (e >= s >= -0.001):
+            continue
+        i = text.find(tok, pos)
+        if i < 0:
+            continue
+        j = i + len(tok)
+        out.append((i, j, s, e))
+        pos = j
+    return out
+
+
+def _sentence_breaks(text, aw):
+    br = {k for k, (i, j, s, e) in enumerate(aw) if _SENT_END.search(text[i:j])}
+    br.add(len(aw) - 1)
+    return br
+
+
+def _split_long(aw, lo, hi):
+    dur = aw[hi][3] - aw[lo][2]
+    if dur <= MAX_SEG_S or hi - lo < 1:
+        return [(lo, hi)]
+    best, bestscore = None, -1e9
+    for k in range(lo, hi):
+        if (aw[k][3] - aw[lo][2]) < MIN_SPLIT_S or (aw[hi][3] - aw[k + 1][2]) < MIN_SPLIT_S:
+            continue
+        gap = aw[k + 1][2] - aw[k][3]
+        score = gap - 0.02 * abs((aw[k][3] - aw[lo][2]) - dur / 2.0)
+        if score > bestscore:
+            bestscore, best = score, k
+    if best is None:
+        return [(lo, hi)]
+    return _split_long(aw, lo, best) + _split_long(aw, best + 1, hi)
+
+
+def _free_len(s, e, spans):
+    """length of [s,e) not covered by any interval in `spans` (sorted, disjoint)."""
+    free = e - s
+    for (a, b) in spans:
+        if b <= s:
+            continue
+        if a >= e:
+            break
+        free -= min(e, b) - max(s, a)
+    return max(free, 0.0)
+
+
+def build_events(text, words, bursts, dur_s):
+    aw = align_words(text, words)
+    bs = []
+    for t in bursts or []:
+        try:
+            s, e, lab = float(t[0]), float(t[1]), str(t[2] or "").strip()
+        except Exception:
+            continue
+        if e <= s:
+            continue
+        bs.append((s, e, lab))
+    bs.sort()
+
+    if not aw:
+        ev = [{"kind": "burst", "s": s, "e": e, "label": l, "dur": e - s, "timed": True}
+              for (s, e, l) in bs]
+        _absorb(ev, dur_s)
+        return ev, aw
+
+    # 1. sentence segmentation
+    br = _sentence_breaks(text, aw)
+    segs, lo = [], 0
+    for k in range(len(aw)):
+        if k in br:
+            segs.append((lo, k))
+            lo = k + 1
+
+    # 2. every burst splits the segment it lands in, at the WORD GAP nearest its midpoint.
+    #    (Requiring the burst to sit entirely inside a gap would fail for the 49 % that overlap
+    #    a word, and those would then be emitted out of order.)
+    cuts = set()
+    for (s, e, l) in bs:
+        mid = 0.5 * (s + e)
+        for (lo, hi) in segs:
+            if aw[lo][2] <= mid <= aw[hi][3] and hi > lo:
+                k = min(range(lo, hi), key=lambda k: abs(0.5 * (aw[k][3] + aw[k + 1][2]) - mid))
+                cuts.add(k)
+                break
+    out, = [[]]
+    for (lo, hi) in segs:
+        a = lo
+        for k in range(lo, hi):
+            if k in cuts:
+                out.append((a, k))
+                a = k + 1
+        out.append((a, hi))
+    segs = [x for x in out if x[0] <= x[1]]
+
+    # 3. length cap
+    out = []
+    for (lo, hi) in segs:
+        out.extend(_split_long(aw, lo, hi))
+    segs = out
+
+    spans = sorted((aw[lo][2], aw[hi][3]) for (lo, hi) in segs)
+    ev = [{"kind": "speech", "s": aw[lo][2], "e": aw[hi][3], "lo": lo, "hi": hi} for (lo, hi) in segs]
+    for (s, e, l) in bs:
+        f = _free_len(s, e, spans)
+        ev.append({"kind": "burst", "s": s, "e": e, "label": l,
+                   "dur": f, "timed": f >= BURST_MIN_S})
+    ev.sort(key=lambda d: (d["s"], 0 if d["kind"] == "speech" else 1))
+    _absorb(ev, dur_s)
+    return ev, aw
+
+
+def _absorb(ev, dur_s):
+    prev = 0.0
+    for k, e in enumerate(ev):
+        if e["kind"] == "burst" and not e.get("timed"):
+            continue                                  # occupies no budget of its own
+        gap = e["s"] - prev
+        if 0 < gap < PAUSE_MIN:
+            if e["kind"] == "speech":
+                e["s"] = prev
+            elif k > 0 and ev[k - 1]["kind"] == "speech":
+                ev[k - 1]["e"] = e["s"]
+        prev = max(prev, e["e"] if e["kind"] == "speech" else e["s"] + e["dur"])
+    if dur_s and ev:
+        last = None
+        for e in reversed(ev):
+            if e["kind"] == "speech":
+                last = e
+                break
+        if last is not None:
+            tail = float(dur_s) - last["e"]
+            if 0 < tail < PAUSE_MIN and last is ev[-1]:
+                last["e"] = float(dur_s)
+
+
+def render(text, words, bursts, dur_s, timed=True, drop_bursts=False, trailing=True,
+           blank_text=False):
+    """-> (script, meta).  `timed=False` keeps burst labels but omits every number.
+
+    `blank_text=True` emits `...` in place of every SPEECH CHUNK and leaves every tag -- pauses,
+    per-segment durations, burst labels and their lengths -- exactly where it was.  It exists for
+    the classifier-free-guidance DPO pairs, where the two clips of a pair say different words and
+    the preference must therefore not be decidable from the words.  Nothing else about the
+    rendering changes: the segmentation, the numbers and their sum are computed from the real text
+    and the real word timings, so the emitted budget still describes the real clip.
+    """
+    text = text or ""
+    ev, aw = build_events(text, words, bursts, dur_s)
+    # `not aw` used to be part of this guard, which made the wordless branch of build_events
+    # (line ~131, "no aligned words -> the events ARE the bursts") unreachable: a clip that is
+    # nothing but a vocal burst rendered as its empty text instead of as `(chuckle, 0.5 seconds)`.
+    # `not ev` is equivalent for every other input -- with no words and no bursts build_events
+    # returns an empty list anyway -- so this only changes rows that are wordless AND carry a
+    # burst: 709 of the 3,144,739 corpus rows, plus the standalone-burst training set.
+    if not ev:
+        head = ("..." if text.strip() else "") if blank_text else text.strip()
+        return head, {"n_seg": 0, "n_burst": 0, "n_pause": 0, "n_burst_untimed": 0,
+                      "timed": False, "acc": 0.0}
+
+    parts, t, cut = [], 0.0, 0
+    n_seg = n_burst = n_pause = n_untimed = 0
+    acc = 0.0
+    for k, e in enumerate(ev):
+        if e["kind"] == "burst" and drop_bursts:
+            continue                                  # its span folds into the surrounding pause
+        if e["kind"] == "burst" and not e["timed"]:
+            parts.append(f"({e['label'].lower() or 'vocal burst'})")
+            n_burst += 1
+            n_untimed += 1
+            continue
+        gap = e["s"] - t
+        if timed and gap >= PAUSE_MIN:
+            parts.append(f"[{_fmt(gap)} seconds pause]")
+            acc += round(gap, 1)
+            n_pause += 1
+        if e["kind"] == "speech":
+            is_last_speech = not any(x["kind"] == "speech" for x in ev[k + 1:])
+            end_char = len(text) if is_last_speech else aw[e["hi"]][1]
+            chunk = text[cut:end_char].strip()
+            cut = end_char
+            if blank_text:
+                chunk = "..." if chunk else ""
+            d = e["e"] - e["s"]
+            if timed:
+                parts.append(f"[{_fmt(d)} seconds duration]")
+                acc += round(d, 1)
+            if chunk:
+                parts.append(chunk)
+            n_seg += 1
+            t = e["e"]
+        else:
+            lab = e["label"].lower() or "vocal burst"
+            parts.append(f"({lab}, {_fmt(e['dur'])} seconds)" if timed else f"({lab})")
+            if timed:
+                acc += round(e["dur"], 1)
+            n_burst += 1
+            t = max(t, e["s"] + e["dur"])
+    if timed and trailing and dur_s and (float(dur_s) - t) >= PAUSE_MIN:
+        parts.append(f"[{_fmt(float(dur_s) - t)} seconds pause]")
+        acc += round(float(dur_s) - t, 1)
+        n_pause += 1
+
+    script = re.sub(r"\s{2,}", " ", " ".join(p for p in parts if p)).strip()
+    return script, {"n_seg": n_seg, "n_burst": n_burst, "n_pause": n_pause,
+                    "n_burst_untimed": n_untimed, "timed": bool(timed), "acc": round(acc, 3)}
+
+
+# ---------------------------------------------------------------- independent verification
+_TAG = re.compile(r"\[(\d+\.\d) seconds (?:pause|duration)\]")
+_BUR = re.compile(r"\([^,()]*, (\d+\.\d) seconds\)")
+
+
+def verify(script, dur_s, text=None):
+    """Re-parses the EMITTED string.  Deliberately does not reuse the renderer's accumulator, so
+    a bookkeeping error in the renderer cannot hide behind its own arithmetic."""
+    tot = sum(float(m.group(1)) for m in _TAG.finditer(script))
+    tot += sum(float(m.group(1)) for m in _BUR.finditer(script))
+    stripped = re.sub(r"\([^)]*\)", " ", re.sub(r"\[[^\]]*\]", " ", script))
+    stripped = re.sub(r"\s+", " ", stripped).strip()
+    return {"sum": round(tot, 2), "dur": round(float(dur_s or 0), 2),
+            "residual": round(float(dur_s or 0) - tot, 2), "text_out": stripped}
+
+
+# ---------------------------------------------------------------- script -> plan
+_ITEM = re.compile(r"\[(\d+\.\d) seconds (pause|duration)\]|\(([^()]*?), (\d+\.\d) seconds\)|\(([^()]*?)\)")
+
+
+def parse_script(script):
+    """Read an emitted timed script back into a plan.
+
+    Used by the GRPO reward: the prompt states WHERE each burst should fall and HOW LONG it and
+    every utterance should be, and the reward has to compare that against what the detector found
+    in the generated audio.  Parsing the emitted string rather than keeping a parallel structure
+    means the reward can only ever score what the model was actually shown.
+
+    -> (items, total_s) where each item is
+       {"kind": "pause"|"speech"|"burst", "t": start_s, "dur": s, "label": str|None}
+    """
+    items, t = [], 0.0
+    for m in _ITEM.finditer(script or ""):
+        if m.group(1) is not None:
+            d = float(m.group(1))
+            items.append({"kind": m.group(2), "t": t, "dur": d, "label": None})
+            t += d
+        elif m.group(3) is not None:
+            d = float(m.group(4))
+            items.append({"kind": "burst", "t": t, "dur": d, "label": m.group(3).strip()})
+            t += d
+        else:
+            items.append({"kind": "burst", "t": t, "dur": None, "label": m.group(5).strip()})
+    return items, round(t, 2)
+
+
+def prompted_bursts(script):
+    items, _ = parse_script(script)
+    return [(i["label"], i["t"], i["dur"]) for i in items
+            if i["kind"] == "burst" and i["dur"] is not None]

code/caption_curriculum_dataset.py

@@ -0,0 +1,172 @@
+#!/usr/bin/env python3
+"""Lazy container-backed loader for the multi-caption S1->S10 curriculum."""
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from pathlib import Path
+import random
+
+import numpy as np
+
+from timed_prompt import repair_prompt
+
+
+TAR_MEMBER_MAP = "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/tar_member_map/manifest.json"
+FAMILIES = ("general_script", "bude_v1", "bude_v1_1", "timbre", "voice_tagging", "gemma", "transcript_only")
+DESC_DTYPE = np.dtype([
+    ("base", "<u4"), ("reference", "<u4"), ("annotation", "<u4"),
+    ("family", "u1"), ("surface", "u1"), ("mode", "u1"), ("reserved", "u1"),
+])
+MISSING = np.iinfo(np.uint32).max
+
+
+class JsonlOffsets:
+    def __init__(self, records: str | Path, offsets: str | Path):
+        self.path = Path(records)
+        self.offsets = np.load(offsets, mmap_mode="r", allow_pickle=False)
+        self.fd = os.open(self.path, os.O_RDONLY)
+        assert self.offsets.dtype == np.int64 and int(self.offsets[0]) == 0
+        assert int(self.offsets[-1]) == self.path.stat().st_size
+
+    def row(self, index: int) -> dict:
+        start, end = int(self.offsets[index]), int(self.offsets[index + 1])
+        return json.loads(os.pread(self.fd, end - start, start))
+
+
+class PhysicalBase(JsonlOffsets):
+    def __init__(self, manifest: str | Path):
+        path = Path(manifest)
+        spec = json.loads(path.read_text())
+        root = path.parent
+        super().__init__(root / spec["records"], root / spec["offsets"])
+
+
+def shuffled(tags: list[str], uid: str, stage: str, presentation: int, family: str) -> list[str]:
+    result = list(tags)
+    seed = int.from_bytes(hashlib.blake2b(
+        f"20260923|{stage}|{uid}|{family}|{presentation}".encode(), digest_size=8).digest(), "little")
+    random.Random(seed).shuffle(result)
+    return result
+
+
+def caption_prompt(caption: str, transcript: str) -> str:
+    caption = " ".join(str(caption).replace('"', "'").split())
+    assert caption
+    return f'CAPTION: {caption}\nTRANSCRIPT: "{transcript}"'
+
+
+class CaptionCurriculumRecords:
+    def __init__(self, manifest: str | Path):
+        self.manifest_path = Path(manifest)
+        self.spec = json.loads(self.manifest_path.read_text())
+        assert self.spec["status"] == "complete"
+        assert self.spec["format"] == "m2-caption-curriculum-v1"
+        root = self.manifest_path.parent
+        self.phase = self.spec["stage"]
+        self.descriptors = np.load(root / self.spec["descriptors"], mmap_mode="r", allow_pickle=False)
+        self.order = np.load(root / self.spec["order"], mmap_mode="r", allow_pickle=False)
+        assert self.descriptors.dtype == DESC_DTYPE
+        assert len(self.descriptors) == len(self.order) == self.spec["presentations"]
+        self.base = PhysicalBase(self.spec["base_manifest"])
+        assets = json.loads(Path(self.spec["assets_manifest"]).read_text())
+        assert assets["status"] == "complete"
+        asset_root = Path(self.spec["assets_manifest"]).parent
+        self.annotations = JsonlOffsets(asset_root / assets["annotations"]["records"],
+                                        asset_root / assets["annotations"]["offsets"])
+        self.references = JsonlOffsets(asset_root / assets["references"]["records"],
+                                       asset_root / assets["references"]["offsets"])
+
+        import sys
+        sys.path.insert(0, "/e/scratch/reformo/schuhmann1_moss/code/m2_20k_score")
+        from dataset import ShardReader
+        from tar_member_map import TarMemberMap
+
+        class MappedShardReader(ShardReader):
+            def __init__(self):
+                super().__init__(capacity=16)
+                self.packed_members = TarMemberMap(TAR_MEMBER_MAP)
+
+            def member(self, path, name):
+                offset, size = self.packed_members.locate(path, name)
+                return self.bytes_at(path, offset, size)
+
+        self.reader = MappedShardReader()
+
+    def __len__(self):
+        return len(self.order)
+
+    def example(self, index: int, packer):
+        descriptor = self.descriptors[int(self.order[index])]
+        base = self.base.row(int(descriptor["base"]))
+        uid = base["uid"]
+        family = FAMILIES[int(descriptor["family"])]
+        mode = "reference" if int(descriptor["mode"]) else "instruction"
+        meta = dict(self.reader.record(base["record_locator"]))
+        assert meta["uid"] == uid
+        meta["frames"] = int(meta.get("moss_frames", base["target_locator"]["full_frames"]))
+        target = self.reader.codes(base["target_locator"])
+        assert len(target) == meta["frames"]
+        transcript = str(meta["text"])
+
+        annotation = None
+        if int(descriptor["annotation"]) != MISSING:
+            annotation = self.annotations.row(int(descriptor["annotation"]))
+            assert annotation["uid"] == uid
+
+        prompt_info = {"timed": False, "eligible": False, "duration_tags": 0,
+                       "format_id": "m2-multiformat-caption-v1-20260923"}
+        form = "B"
+        if family == "general_script":
+            meta["prompt"], prompt_info = repair_prompt(
+                meta, caption=None, form="A", mode=mode,
+                presentation=int(index), phase=self.phase)
+            form = "A"
+        elif family == "bude_v1":
+            assert base.get("bude_caption")
+            meta["prompt"] = caption_prompt(base["bude_caption"], transcript)
+        elif family == "bude_v1_1":
+            meta["prompt"] = caption_prompt(annotation["bude"], transcript)
+        elif family in ("timbre", "voice_tagging"):
+            prefix = "timbre" if family == "timbre" else "voice"
+            tags = shuffled(annotation[prefix + "_tags"], uid, self.phase, int(index), family)
+            prose = annotation[prefix + "_prose"]
+            surface = int(descriptor["surface"])
+            if surface == 0:
+                assert tags
+                body = ", ".join(tags)
+            elif surface == 1:
+                body = prose
+            else:
+                assert tags and prose
+                body = "Tags: " + ", ".join(tags) + ". Caption: " + prose
+            meta["prompt"] = caption_prompt(body, transcript)
+        elif family == "gemma":
+            meta["prompt"] = caption_prompt(annotation["gemma_caption"], transcript)
+        elif family == "transcript_only":
+            meta["prompt"] = f'TRANSCRIPT: "{transcript}"'
+            form = "A"
+        else:
+            raise AssertionError(family)
+
+        reference = None
+        if mode == "reference":
+            assert int(descriptor["reference"]) != MISSING
+            reference_row = self.references.row(int(descriptor["reference"]))
+            assert reference_row["uid"] == uid and reference_row["reference_uid"] != uid
+            reference = self.reader.codes(reference_row["locator"])
+            assert 0 < len(reference) <= 37
+
+        result = packer.pack_mode(meta, target, mode, reference)
+        result["accounting"] = {
+            "uid": uid, "mode": mode, "form": form, "family": family,
+            "frames": len(target), "target_audio_hours": float(meta["dur_s"]) / 3600,
+            "reference_frames": len(reference) if reference is not None else 0,
+            "prompt_timed": bool(prompt_info["timed"]),
+            "prompt_eligible": bool(prompt_info["eligible"]),
+            "duration_tags": int(prompt_info["duration_tags"]),
+            "prompt_format_id": (prompt_info["format_id"] if family == "general_script"
+                                 else "m2-multiformat-caption-v1-20260923"),
+        }
+        return result

code/cascade_dataset.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""M2-600M cascade adapter with deterministic timed-prompt repair.
+
+Audio/code targets and manifests remain immutable. Every presentation passes
+through ``timed_prompt.repair_prompt``; repair facts are returned in accounting
+so the trainer can prove the observed timed/untimed mixture.
+"""
+import json
+import os
+import numpy as np
+from timed_prompt import repair_prompt
+
+MODES = ('instruction', 'reference')
+TAR_MEMBER_MAP = '/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/tar_member_map/manifest.json'
+
+class CascadeRecords:
+    def __init__(self, manifest):
+        from pathlib import Path
+        self.manifest_path = Path(manifest)
+        spec = json.loads(self.manifest_path.read_text())
+        assert spec['status'] == 'complete' and spec['presentations'] > 0
+        root = self.manifest_path.parent
+        self.spec = spec
+        self.phase = str(spec.get('phase') or spec.get('ladder_stage') or self.manifest_path.parent.name)
+        self.offsets = np.load(root / spec['offsets'], mmap_mode='r', allow_pickle=False)
+        assert len(self.offsets) == spec['presentations'] + 1
+        assert self.offsets.dtype == np.int64 and int(self.offsets[0]) == 0
+        self.order = np.load(root / spec['order'], mmap_mode='r', allow_pickle=False)
+        assert len(self.order) == spec['presentations']
+        assert self.order.dtype in (np.dtype('uint32'), np.dtype('int64'))
+        # Full uniqueness is established once by the manifest verifier.  Never
+        # materialize 48M Python ints/sets independently on all 128 DDP ranks.
+        assert int(self.order.min()) == 0 and int(self.order.max()) == spec['presentations'] - 1
+        self.path = root / spec['records']
+        self.fd = os.open(self.path, os.O_RDONLY)
+        assert int(self.offsets[-1]) == self.path.stat().st_size
+        if 'records_bytes' in spec:
+            assert self.path.stat().st_size == spec['records_bytes']
+            assert (root / spec['offsets']).stat().st_size == spec['offsets_bytes']
+            assert (root / spec['order']).stat().st_size == spec['order_bytes']
+        import sys
+        sys.path.insert(0, '/e/scratch/reformo/schuhmann1_moss/code/m2_20k_score')
+        from dataset import ShardReader
+        from tar_member_map import TarMemberMap
+
+        class MappedShardReader(ShardReader):
+            def __init__(self):
+                # ShardReader uses the same bounded capacity for Parquet
+                # reference tables as for TAR member dictionaries.  TAR
+                # dictionaries are overridden below, but Parquet tables still
+                # need a non-zero LRU (capacity=0 caused an immediate eviction
+                # followed by move_to_end(KeyError) in S1 smoke 1892952).
+                super().__init__(capacity=16)
+                self.packed_members = TarMemberMap(TAR_MEMBER_MAP)
+
+            def member(self, path, name):
+                offset, size = self.packed_members.locate(path, name)
+                return self.bytes_at(path, offset, size)
+
+        self.reader = MappedShardReader()
+
+    def __len__(self):
+        return len(self.order)
+
+    def row(self, index):
+        j = int(self.order[index])
+        start, end = int(self.offsets[j]), int(self.offsets[j + 1])
+        return json.loads(os.pread(self.fd, end - start, start))
+
+    def example(self, index, packer):
+        row = self.row(index)
+        mode = row['mode']
+        form = row.get('form', 'A')
+        assert mode in MODES, row['uid']
+        assert form in ('A', 'B'), row['uid']
+        meta = self.reader.record(row['record_locator'])
+        assert meta['uid'] == row['uid']
+        if 'moss_frames' in meta:
+            meta['frames'] = int(meta['moss_frames'])
+        else:
+            # Sidecar element JSONs carry no moss_frames key (pre-upload state);
+            # frame count comes from the manifest locator (validated in codes()).
+            meta['frames'] = int(row['target_locator']['full_frames'])
+        codes = self.reader.codes(row['target_locator'])
+        assert len(codes) == meta['frames']
+        meta = dict(meta)
+        meta['prompt'], prompt_info = repair_prompt(
+            meta, caption=row.get('bude_caption'), form=form, mode=mode,
+            presentation=int(index), phase=self.phase)
+        reference = None
+        if mode == 'reference':
+            assert row.get('reference_uid') and row['reference_uid'] != row['uid']
+            assert row.get('reference_locator')
+            reference = self.reader.codes(row['reference_locator'])
+            assert 0 < len(reference) <= 37
+        result = packer.pack_mode(meta, codes, mode, reference)
+        result['accounting'] = {'uid': row['uid'], 'mode': mode, 'form': form,
+                                'frames': len(codes),
+                                'target_audio_hours': float(meta['dur_s']) / 3600,
+                                'reference_frames': len(reference) if reference is not None else 0,
+                                'prompt_timed': bool(prompt_info['timed']),
+                                'prompt_eligible': bool(prompt_info['eligible']),
+                                'duration_tags': int(prompt_info['duration_tags']),
+                                'prompt_format_id': prompt_info['format_id']}
+        return result

code/conditioning.py

@@ -0,0 +1,153 @@
+"""Projected score tokens. Missing measurements are masked, never imputed scores.
+
+The schema file pins names, units and measurement scopes. A tokenizer reserves
+six textual marker IDs inside M2's existing padded vocabulary; its original
+embedding rows and all original model weights remain loadable without resizing.
+"""
+from contextlib import contextmanager
+import json
+from pathlib import Path
+import sys
+import torch
+from torch import nn
+
+SMALL = Path('/e/scratch/reformo/schuhmann1_moss/code/small_tts')
+sys.path.insert(0, str(SMALL))
+MARKERS = ('<|score_global_start|>', '<|score_global_slot|>', '<|score_global_end|>',
+           '<|score_sentence_start|>', '<|score_sentence_slot|>', '<|score_sentence_end|>')
+GLOBAL_TOKENS = 5
+LOCAL_TOKENS = 10
+
+
+def install_markers(tokenizer, vocab_size):
+    import inspect
+    parameters = inspect.signature(tokenizer.add_special_tokens).parameters
+    option = ('replace_extra_special_tokens' if 'replace_extra_special_tokens' in parameters
+              else 'replace_additional_special_tokens')
+    tokenizer.add_special_tokens({'additional_special_tokens': list(MARKERS)}, **{option: False})
+    ids = dict(zip(MARKERS, tokenizer.convert_tokens_to_ids(list(MARKERS))))
+    assert len(set(ids.values())) == 6 and all(0 <= x < vocab_size for x in ids.values()), ids
+    for marker, token in ids.items():
+        assert tokenizer.encode(marker, add_special_tokens=False) == [token]
+    return ids
+
+
+def feature_vector(values, dimensions):
+    """Return values scaled by documented bounds and an explicit validity mask.
+
+    Missing values become numerical zero only alongside mask=0. Normalization
+    bounds are scale anchors, not clipping bounds: actual regression outputs can
+    exceed nominal score ranges. Only non-finite measurements (or violations of
+    explicitly configured hard bounds) are rejected.
+    """
+    import math
+    scaled, mask = [], []
+    for dim in dimensions:
+        value = values.get(dim['name'])
+        valid = value is not None
+        if valid:
+            value = float(value)
+            low, high = map(float, dim['normalization_bounds'])
+            legal = dim.get('hard_bounds')
+            if not math.isfinite(value) or high <= low or (legal is not None and not legal[0] <= value <= legal[1]):
+                raise ValueError(f'Invalid measured feature {dim["name"]}={value}')
+            value = 2 * (value - low) / (high - low) - 1
+        scaled.append(value if valid else 0.0)
+        mask.append(float(valid))
+    return scaled + mask
+
+
+class ScoreConditioner(nn.Module):
+    def __init__(self, global_dims, local_dims, hidden=1024):
+        super().__init__()
+        self.global_dims, self.local_dims = global_dims, local_dims
+        self.hidden = hidden
+        self.global_projection = nn.Sequential(nn.Linear(2 * global_dims, 256), nn.SiLU(),
+                                              nn.Linear(256, GLOBAL_TOKENS * hidden))
+        self.local_projection = nn.Sequential(nn.Linear(2 * local_dims, 512), nn.SiLU(),
+                                             nn.Linear(512, LOCAL_TOKENS * hidden))
+        self.boundaries = nn.Embedding(4, hidden)
+        for module in self.modules():
+            if isinstance(module, (nn.Linear, nn.Embedding)):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+                if isinstance(module, nn.Linear):
+                    nn.init.zeros_(module.bias)
+
+    def forward(self, global_values, local_values):
+        # At least one padded sentence keeps the same DDP graph in every mode.
+        B, S, _ = local_values.shape
+        g = self.global_projection(global_values).reshape(B, GLOBAL_TOKENS, self.hidden)
+        local = self.local_projection(local_values).reshape(B, S, LOCAL_TOKENS, self.hidden)
+        boundary = self.boundaries.weight.to(g.dtype)
+        gs = boundary[0].reshape(1, 1, -1).expand(B, 1, -1)
+        ge = boundary[1].reshape(1, 1, -1).expand(B, 1, -1)
+        ls = boundary[2].reshape(1, 1, 1, -1).expand(B, S, 1, -1)
+        le = boundary[3].reshape(1, 1, 1, -1).expand(B, S, 1, -1)
+        local = torch.cat([ls, local, le], dim=2).reshape(B, S * (LOCAL_TOKENS + 2), self.hidden)
+        # Bank index zero is not a score token. Gradients still traverse the
+        # complete bank with zero contribution when a batch has no score mode.
+        zero = torch.zeros(B, 1, self.hidden, device=g.device, dtype=g.dtype)
+        return torch.cat([zero, gs, g, ge, local], dim=1)
+
+
+def scored_class():
+    import moss_small
+    parent = moss_small.MossTTSSmallModel.cls()
+
+    class ScoreM2(parent):
+        def __init__(self, config, schema):
+            super().__init__(config)
+            self.score_conditioner = ScoreConditioner(len(schema['global_dimensions']), len(schema['local_dimensions']),
+                                                     hidden=int(config.hidden_size))
+            self._generation_conditioning = None
+
+        def score_embeddings(self, input_ids, conditioning):
+            embeds = super()._build_inputs_embeds(input_ids)
+            if conditioning is None:
+                return embeds
+            values, local, index = conditioning
+            bank = self.score_conditioner(values, local).to(embeds.dtype)
+            if index.shape[1] != input_ids.shape[1]:
+                # Cached generation after prefill has only new audio rows.
+                if input_ids.shape[1] != 1:
+                    raise ValueError('Conditioning positions do not match the prompt; KV-cache generation is required')
+                return embeds
+            scores = bank.gather(1, index.unsqueeze(-1).expand(-1, -1, bank.shape[-1]))
+            return torch.where(index.unsqueeze(-1).gt(0), scores, embeds)
+
+        def _build_inputs_embeds(self, input_ids):
+            return self.score_embeddings(input_ids, self._generation_conditioning)
+
+        def forward(self, input_ids=None, attention_mask=None, score_conditioning=None, **kwargs):
+            if score_conditioning is not None:
+                kwargs['inputs_embeds'] = self.score_embeddings(input_ids, score_conditioning)
+                input_ids = None
+            return super().forward(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
+
+        @contextmanager
+        def generation_scores(self, conditioning):
+            if self._generation_conditioning is not None:
+                raise RuntimeError('Overlapping generation on one score-conditioned model is unsupported')
+            self._generation_conditioning = conditioning
+            try:
+                yield
+            finally:
+                self._generation_conditioning = None
+
+    return ScoreM2
+
+
+def load_start_model(schema_path, checkpoint, log=print):
+    import moss_small
+    schema = json.loads(Path(schema_path).read_text())
+    config = moss_small.make_config('M2')
+    model = scored_class()(config, schema)
+    state = torch.load(Path(checkpoint) / 'model_bf16.pt', map_location='cpu', weights_only=True)
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    expected = {'score_conditioner.' + name for name in model.score_conditioner.state_dict()}
+    if set(missing) != expected or unexpected:
+        raise RuntimeError(f'Initial M2 checkpoint mismatch: missing={missing}, unexpected={unexpected}')
+    del state
+    model.tie_weights()
+    log(f'Loaded final M2 checkpoint {checkpoint}; new score parameters={sum(p.numel() for p in model.score_conditioner.parameters())}')
+    return model, config, schema

code/continuation_control.py

@@ -0,0 +1,51 @@
+"""Pure control rules shared by the real trainer and its recovery verifier."""
+import hashlib
+import json
+import math
+from pathlib import Path
+
+
+def sha(path):
+    result=hashlib.sha256()
+    with Path(path).open('rb') as stream:
+        for block in iter(lambda:stream.read(8<<20),b''):result.update(block)
+    return result.hexdigest()
+
+
+def complete_checkpoint(directory,contract=None,world=None,content=False):
+    directory=Path(directory);meta=json.loads((directory/'COMPLETE.json').read_text())
+    assert meta['step']>=1
+    if contract is not None:assert meta['contract']==contract
+    if world is not None:assert meta['world']==world
+    assert (directory/'training_state.pt').stat().st_size==meta['state_bytes']>0
+    if content:assert sha(directory/'training_state.pt')==meta['state_sha256']
+    if 'export_bytes' in meta:
+        assert (directory/'model_bf16.pt').stat().st_size==meta['export_bytes']>0
+        if content:assert sha(directory/'model_bf16.pt')==meta['export_sha256']
+    return meta
+
+
+def latest_checkpoint(root,contract,world):
+    """A crash after directory rename but before pointer update is recoverable."""
+    found=[]
+    for directory in Path(root).glob('step-*'):
+        # A named committed directory with invalid content is a real error;
+        # never silently rewind a possibly corrupted training history.
+        meta=complete_checkpoint(directory,contract,world)
+        assert directory.name==f'step-{meta["step"]:08d}'
+        found.append((meta['step'],directory))
+    return max(found)[1] if found else None
+
+
+def epoch_plan(samples,world,per_gpu):
+    assert samples>0 and world>0 and per_gpu>0
+    batch=world*per_gpu;steps=math.ceil(samples/batch);tail=samples-(steps-1)*batch
+    assert tail>=world,'Every final rank must receive a real target; no silent tail drop/padding'
+    quarters=sorted({math.ceil(steps*q/4) for q in (1,2,3,4)})
+    return {'samples':samples,'world':world,'global_batch':batch,'steps':steps,'tail_samples':tail,
+        'quarter_steps':quarters,'epochs':1,'presentations_per_target':1,'dropped_targets':0,'duplicated_targets':0}
+
+
+def contract(plan_path,plan,world):
+    return {'plan_sha256':sha(plan_path),'schema_sha256':plan['schema_sha256'],'manifest_sha256':plan['manifest_sha256'],
+        'world':world,'objective':'global_token_mean_v1','seed':plan['seed']}

code/continuous_train.py

@@ -0,0 +1,428 @@
+#!/usr/bin/env python3
+"""One optimizer and one LR schedule across the complete S1->S10 ladder."""
+from __future__ import annotations
+
+import argparse
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import nullcontext
+from datetime import datetime, timedelta
+import hashlib
+import json
+import math
+import os
+from pathlib import Path
+import random
+import statistics
+import sys
+import time
+
+SC = Path("/e/scratch/reformo/schuhmann1_moss")
+HERE = Path(__file__).resolve().parent
+V3 = Path("/e/home/jusers/schuhmann1/jupiter/m2_600m_tts_training/cascade_v3_timed")
+M2 = SC / "code/m2_20k_score"
+SMALL = SC / "code/small_tts"
+for candidate in (str(HERE), str(V3), str(M2), str(SMALL)):
+    while candidate in sys.path:
+        sys.path.remove(candidate)
+sys.path[:0] = [str(HERE), str(V3), str(M2), str(SMALL)]
+
+
+def sha(path: str | Path) -> str:
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as stream:
+        for block in iter(lambda: stream.read(8 << 20), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def validate(plan: dict, plan_path: Path) -> list[dict]:
+    assert plan["kind"] == "m2_continuous_large_talker_v1"
+    assert plan["architecture"] == "qwen3_0.6b_plus_fresh_sft3_width_talker"
+    assert plan["scheduler"] == "single_linear_warmup_then_cosine_no_stage_restarts"
+    assert plan["nodes"] >= 1 and plan["ranks_per_node"] in (1, 4)
+    assert plan["world_size"] == plan["nodes"] * plan["ranks_per_node"]
+    assert plan["global_batch"] == plan["world_size"] * plan["samples_per_gpu"]
+    assert plan["lr_backbone_peak"] == 8e-5
+    assert plan["lr_talker_peak"] == 2.4e-4
+    assert plan["warmup_steps"] >= 1 and 0 <= plan["lr_floor_factor"] <= 1
+    assert plan["initialization"]["semantic"] == "original_pretrained_Qwen3-0.6B"
+    assert plan["initialization"]["talker"] == "random_fresh"
+    assert plan["checkpoint_policy"] == "all_stage_boundaries_plus_interval"
+    expected_names = [f"S{i}" for i in range(1, 11)]
+    assert [stage["name"] for stage in plan["stages"]] == expected_names
+    specs = []
+    total_updates = 0
+    for stage in plan["stages"]:
+        spec = json.loads(Path(stage["manifest"]).read_text())
+        assert spec["status"] == "complete" and spec["presentations"] > 0
+        assert sha(stage["manifest"]) == stage["manifest_sha256"]
+        natural = math.ceil(spec["presentations"] / plan["global_batch"])
+        updates = int(stage.get("max_updates", natural))
+        assert 1 <= updates <= natural
+        assert updates == stage["updates"]
+        # A deliberately truncated smoke never reaches the natural final
+        # partial batch. Validate that tail only when this plan consumes the
+        # complete stage; every earlier update is a full global batch.
+        if updates == natural:
+            tail = spec["presentations"] - (natural - 1) * plan["global_batch"]
+            assert tail >= plan["world_size"]
+        total_updates += updates
+        specs.append(spec)
+    assert total_updates == plan["total_updates"]
+    assert plan["warmup_steps"] < total_updates
+    for name, expected in plan["code_sha256"].items():
+        source = HERE / name if name in ("continuous_train.py", "large_talker.py",
+                                          "caption_curriculum_dataset.py") else V3 / name
+        assert sha(source) == expected, (name, source)
+    return specs
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("plan")
+    args = parser.parse_args()
+    plan_path = Path(args.plan).resolve()
+    plan = json.loads(plan_path.read_text())
+
+    import numpy as np
+    import torch
+    import torch.distributed as dist
+    from torch.nn.parallel import DistributedDataParallel as DDP
+    from transformers import AutoTokenizer
+    from packing import ScorePacker
+    from plan import step_indices, microbatches, dynamic_weights
+    from continuation_control import latest_checkpoint, complete_checkpoint
+    from cascade_dataset import CascadeRecords
+    from caption_curriculum_dataset import CaptionCurriculumRecords
+    import moss_small
+    import state_io
+    import va_loss
+    from large_talker import build_fresh, parameter_counts
+
+    specs = validate(plan, plan_path)
+    rank, world, local = (int(os.environ[key]) for key in
+                          ("SLURM_PROCID", "SLURM_NTASKS", "SLURM_LOCALID"))
+    assert world == int(plan["world_size"])
+    device_index = 0 if torch.cuda.device_count() == 1 else local
+    torch.cuda.set_device(device_index)
+    device = torch.device("cuda", device_index)
+    os.environ.update(RANK=str(rank), WORLD_SIZE=str(world), LOCAL_RANK=str(device_index))
+    dist.init_process_group("nccl", timeout=timedelta(minutes=30), device_id=device)
+    torch.set_num_threads(4)
+    if plan.get("deterministic", False):
+        torch.use_deterministic_algorithms(True)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+    def log(message: str) -> None:
+        if rank == 0:
+            print(datetime.now().astimezone().isoformat(timespec="seconds"), message, flush=True)
+
+    random.seed(plan["seed"])
+    np.random.seed(plan["seed"])
+    torch.manual_seed(plan["seed"])
+    torch.cuda.manual_seed_all(plan["seed"])
+    invocation = os.environ["SLURM_JOB_ID"] + os.environ.get("M2_INVOCATION_SUFFIX", "")
+    invocation_dir = Path(plan["output"]) / "invocations" / invocation
+    invocation_dir.mkdir(parents=True, exist_ok=True)
+    if rank == 0:
+        assert not (invocation_dir / "steps.jsonl").exists()
+
+    schema_path = SC / "out/m2_20k_score/score_schema.json"
+    schema = json.loads(schema_path.read_text())
+    model, config = build_fresh(schema, log=log)
+    counts = parameter_counts(model)
+    log("Fresh large-Talker model: " + json.dumps(counts, sort_keys=True))
+
+    from functools import partial
+    from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+        apply_activation_checkpointing, checkpoint_wrapper, CheckpointImpl)
+    for module in model.modules():
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = False
+    apply_activation_checkpointing(
+        model,
+        check_fn=lambda module: type(module).__name__ == "MossQwen3DecoderLayer",
+        checkpoint_wrapper_fn=partial(checkpoint_wrapper,
+                                      checkpoint_impl=CheckpointImpl.NO_REENTRANT),
+    )
+
+    class Step(torch.nn.Module):
+        def __init__(self, base):
+            super().__init__()
+            self.base = base
+
+        def forward(self, input_ids, attention_mask, labels, score_conditioning,
+                    weights, scale):
+            hidden = self.base(input_ids=input_ids, attention_mask=attention_mask,
+                               score_conditioning=score_conditioning,
+                               use_cache=False).last_hidden_state
+            loss, per = va_loss.compute_supervised_loss_from_hidden(
+                self.base, global_hidden_states=hidden, labels=labels,
+                channelwise_loss_weight=weights, return_per_channel=True)
+            return loss * scale, per
+
+    step_model = DDP(Step(model.to(device)).train(), device_ids=[device_index],
+                     broadcast_buffers=False, find_unused_parameters=False,
+                     gradient_as_bucket_view=True)
+    if plan["gradient_communication"] == "bfloat16":
+        from torch.distributed.algorithms.ddp_comm_hooks.default_hooks import bf16_compress_hook
+        step_model.register_comm_hook(None, bf16_compress_hook)
+    else:
+        assert plan["gradient_communication"] == "float32"
+
+    groups = {"backbone": [], "talker": []}
+    for name, parameter in model.named_parameters():
+        if not parameter.requires_grad:
+            continue
+        key = "backbone" if name.startswith(("transformer.", "text_lm_head.")) else "talker"
+        groups[key].append(parameter)
+    optimizer = torch.optim.AdamW([
+        {"params": groups["backbone"], "lr": plan["lr_backbone_peak"], "name": "backbone"},
+        {"params": groups["talker"], "lr": plan["lr_talker_peak"], "name": "talker"},
+    ], betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1, foreach=False)
+
+    total_updates = int(plan["total_updates"])
+    warmup = int(plan["warmup_steps"])
+    floor = float(plan["lr_floor_factor"])
+
+    def lr_factor(step: int) -> float:
+        if step < warmup:
+            return (step + 1) / warmup
+        progress = min(1.0, (step - warmup) / max(1, total_updates - warmup))
+        return floor + (1.0 - floor) * 0.5 * (1.0 + math.cos(math.pi * progress))
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_factor)
+    _, _, Processor = moss_small.export_classes()
+    processor = Processor(tokenizer=AutoTokenizer.from_pretrained(
+        moss_small.SFT3, trust_remote_code=True, local_files_only=True),
+        audio_tokenizer=None, model_config=config)
+    packer = ScorePacker(processor, config, schema)
+    datasets = []
+    for stage, spec in zip(plan["stages"], specs):
+        cls = CaptionCurriculumRecords if spec.get("format") == "m2-caption-curriculum-v1" else CascadeRecords
+        datasets.append(cls(stage["manifest"]))
+    assert all(len(dataset) == spec["presentations"]
+               for dataset, spec in zip(datasets, specs))
+
+    manifest_digest = hashlib.sha256("".join(
+        stage["manifest_sha256"] for stage in plan["stages"]).encode()).hexdigest()
+    contract = {
+        "plan_sha256": sha(plan_path), "schema_sha256": sha(schema_path),
+        "manifest_sha256": manifest_digest, "world": world,
+        "objective": "global_token_mean_continuous_ladder_v1", "seed": plan["seed"],
+    }
+    start_step = 0
+    selection = [None]
+    if rank == 0:
+        try:
+            directory = latest_checkpoint(Path(plan["output"]) / "checkpoints", contract, world)
+            if directory:
+                complete_checkpoint(directory, contract, world, content=True)
+            selection[0] = {"path": str(directory) if directory else None}
+        except Exception as error:
+            selection[0] = {"error": repr(error)}
+    dist.broadcast_object_list(selection, src=0)
+    assert "error" not in selection[0], selection[0]
+    if selection[0]["path"]:
+        start_step = state_io.load(selection[0]["path"], step_model, optimizer,
+                                   scheduler, contract)
+    log(f"Continuous resume point={start_step}/{total_updates}; one scheduler, no stage restarts")
+    if start_step == total_updates:
+        log(f"ALREADY_COMPLETE at {start_step}/{total_updates}")
+        dist.barrier(); dist.destroy_process_group(); return
+    assert start_step < total_updates
+
+    def prepare(dataset, phase, phase_step):
+        began = time.monotonic()
+        indices = step_indices(phase, phase_step, rank, world, len(dataset))
+        examples = [dataset.example(index, packer) for index in indices]
+        assert examples
+        assert max(len(example["input_ids"]) for example in examples) <= int(
+            getattr(config, "max_position_embeddings", 32768))
+        batches = microbatches(examples, plan["max_padded_tokens"], plan["max_examples"])
+        prepared = []
+        for examples_batch in batches:
+            batch = packer.collate(examples_batch)
+            frames = sum(example["accounting"]["frames"] for example in examples_batch)
+            assert int(batch["labels"][:, :, 1].ge(0).sum()) == frames
+            prepared.append((batch, frames, len(examples_batch)))
+        stats = [len(examples), sum(x["accounting"]["frames"] for x in examples),
+                 sum(x["accounting"]["target_audio_hours"] for x in examples),
+                 sum(x["accounting"]["reference_frames"] for x in examples),
+                 sum(len(x["input_ids"]) for x in examples),
+                 sum(batch["attention_mask"].numel() for batch, _, _ in prepared)]
+        stats += [sum(x["accounting"]["mode"] == mode for x in examples)
+                  for mode in ("instruction", "reference")]
+        stats += [sum(x["accounting"]["form"] == form for x in examples)
+                  for form in ("A", "B")]
+        allowed_prompt_ids = set(plan.get("allowed_prompt_format_ids", [plan["prompt_contract"]]))
+        assert all(x["accounting"]["prompt_format_id"] in allowed_prompt_ids for x in examples)
+        stats += [sum(x["accounting"]["prompt_timed"] for x in examples),
+                  sum(x["accounting"]["prompt_eligible"] for x in examples),
+                  sum(x["accounting"]["duration_tags"] for x in examples)]
+        return prepared, stats, time.monotonic() - began
+
+    def update(prepared, local_stats):
+        stats = torch.tensor(local_stats, dtype=torch.float64, device=device)
+        dist.all_reduce(stats)
+        values = stats.tolist()
+        total_samples, total_frames = int(values[0]), int(values[1])
+        optimizer.zero_grad(set_to_none=True)
+        local_loss = torch.zeros((), device=device)
+        per_sums = torch.zeros(13, device=device)
+        for index, (batch, frames, samples) in enumerate(prepared):
+            batch = {name: tuple(value.to(device, non_blocking=True) for value in content)
+                     if name == "score_conditioning" else content.to(device, non_blocking=True)
+                     for name, content in batch.items() if name not in ("modes", "meta")}
+            weights, scale = dynamic_weights(frames, samples, total_frames, total_samples, world)
+            context = step_model.no_sync() if index + 1 < len(prepared) else nullcontext()
+            with context:
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    loss, per = step_model(**batch, weights=weights, scale=scale)
+                assert torch.isfinite(loss)
+                loss.backward()
+            local_loss += loss.detach() / world
+            per_sums += per * torch.tensor([frames + samples] + [frames] * 12,
+                                           device=device)
+        gradient_norm = torch.nn.utils.clip_grad_norm_(step_model.parameters(), 1.0)
+        assert torch.isfinite(gradient_norm)
+        optimizer.step()
+        scheduler.step()
+        dist.all_reduce(local_loss)
+        dist.all_reduce(per_sums)
+        per_sums /= torch.tensor([total_frames + total_samples] + [total_frames] * 12,
+                                 device=device)
+        torch.cuda.synchronize()
+        return {
+            "loss": float(local_loss), "per_channel": per_sums.tolist(),
+            "global_samples": total_samples, "target_frames": total_frames,
+            "target_audio_hours": values[2], "reference_frames": int(values[3]),
+            "input_rows": int(values[4]), "padded_rows": int(values[5]),
+            "mode_counts": dict(zip(("instruction", "reference"), map(int, values[6:8]))),
+            "form_counts": dict(zip(("A", "B"), map(int, values[8:10]))),
+            "prompt_timed": int(values[10]), "prompt_eligible": int(values[11]),
+            "duration_tags": int(values[12]),
+            "prompt_format_id": plan["prompt_contract"],
+            "gradient_norm": float(gradient_norm),
+            "lr_backbone": float(optimizer.param_groups[0]["lr"]),
+            "lr_talker": float(optimizer.param_groups[1]["lr"]),
+        }
+
+    metrics = (invocation_dir / "steps.jsonl").open("a", buffering=1) if rank == 0 else None
+    stage_offsets = []
+    cursor = 0
+    for stage in plan["stages"]:
+        stage_offsets.append(cursor)
+        cursor += stage["updates"]
+    assert cursor == total_updates
+    global_quarters = {math.ceil(total_updates * q / 4) for q in (1, 2, 3, 4)}
+    stage_ends = {offset + stage["updates"] for offset, stage in zip(stage_offsets, plan["stages"])}
+    if plan.get("smoke_mode"):
+        # A full optimizer state is several GiB.  The smoke proves both an early
+        # boundary checkpoint and the final checkpoint without writing ten
+        # redundant copies during its deliberately tiny stage transitions.
+        global_quarters = set()
+        stage_ends = set(map(int, plan["smoke_checkpoint_steps"]))
+    last_save = time.monotonic()
+    invocation_begin = time.monotonic()
+    final_step = start_step
+    best_loss = None
+    best_step = None
+    stopped = False
+    observed_times = []
+
+    with ThreadPoolExecutor(max_workers=1) as pool:
+        for stage_index, (stage, dataset, offset) in enumerate(
+                zip(plan["stages"], datasets, stage_offsets)):
+            phase = {"name": stage["name"], "offset": 0,
+                     "global_batch": plan["global_batch"], "updates": stage["updates"]}
+            if start_step >= offset + stage["updates"]:
+                continue
+            pending = None
+            local_start = max(0, start_step - offset)
+            log(f"Entering {stage['name']} at local step {local_start}/{stage['updates']} "
+                f"without optimizer/scheduler reset")
+            for phase_step in range(local_start, stage["updates"]):
+                if pending is None:
+                    pending = pool.submit(prepare, dataset, phase, phase_step)
+                began = time.monotonic()
+                ready, local_stats, preparation_seconds = pending.result()
+                data_wait_seconds = time.monotonic() - began
+                pending = None
+                if phase_step + 1 < stage["updates"]:
+                    pending = pool.submit(prepare, dataset, phase, phase_step + 1)
+                result = update(ready, local_stats)
+                wall_seconds = time.monotonic() - began
+                timing = torch.tensor([wall_seconds, data_wait_seconds, preparation_seconds],
+                                      dtype=torch.float64, device=device)
+                dist.all_reduce(timing, op=dist.ReduceOp.MAX)
+                final_step = offset + phase_step + 1
+                observed_times.append(float(timing[0]))
+                result.update(
+                    stage=stage["name"], stage_index=stage_index + 1,
+                    stage_step=phase_step + 1, stage_updates=stage["updates"],
+                    step=final_step, total_updates=total_updates,
+                    steady=final_step > warmup, wall_seconds=float(timing[0]),
+                    data_wait_seconds=float(timing[1]), preparation_seconds=float(timing[2]),
+                    job=os.environ["SLURM_JOB_ID"], invocation=invocation,
+                )
+                if rank == 0:
+                    metrics.write(json.dumps(result) + "\n")
+                    log(f"{stage['name']} {phase_step+1}/{stage['updates']} global="
+                        f"{final_step}/{total_updates} loss={result['loss']:.5f} "
+                        f"wall={result['wall_seconds']:.3f}s")
+                    if best_loss is None or result["loss"] < best_loss:
+                        best_loss, best_step = result["loss"], final_step
+
+                interval_due = time.monotonic() - last_save >= plan["checkpoint_interval_seconds"]
+                deadline_due = time.monotonic() - invocation_begin >= plan["max_invocation_seconds"]
+                control = torch.tensor([int(interval_due), int(deadline_due)] if rank == 0 else [0, 0],
+                                       device=device)
+                dist.broadcast(control, src=0)
+                stopped = bool(control[1]) and final_step < total_updates
+                export = final_step in stage_ends or final_step in global_quarters
+                if export or bool(control[0]) or stopped or final_step == total_updates:
+                    saved = state_io.save(Path(plan["output"]) / "checkpoints", step_model,
+                                          optimizer, scheduler, final_step, contract,
+                                          export=export or final_step == total_updates)
+                    last_save = time.monotonic()
+                    if rank == 0:
+                        log("Committed retained checkpoint " + str(saved))
+                    dist.barrier()
+                if stopped:
+                    log(f"STOPPED_RESUMABLE at global step {final_step}")
+                    break
+            if stopped:
+                break
+
+    peak = torch.tensor(torch.cuda.max_memory_allocated() / 2**30, device=device)
+    dist.all_reduce(peak, op=dist.ReduceOp.MAX)
+    median_step = statistics.median(observed_times)
+    threshold = plan.get("max_acceptable_median_step_seconds")
+    throughput_failed = bool(threshold is not None and median_step > float(threshold))
+    if rank == 0:
+        status = ("FAIL_THROUGHPUT" if throughput_failed else
+                  ("PASS" if final_step == total_updates else "STOPPED_RESUMABLE"))
+        summary = {
+            "status": status, "job": os.environ["SLURM_JOB_ID"],
+            "invocation": invocation, "final_step": final_step,
+            "total_updates": total_updates, "one_continuous_schedule": True,
+            "optimizer_or_scheduler_restarts_at_stage_boundaries": 0,
+            "best_observed_loss_this_invocation": best_loss,
+            "best_observed_step_this_invocation": best_step,
+            "median_step_seconds_this_invocation": median_step,
+            "p90_step_seconds_this_invocation": float(np.quantile(observed_times, 0.9)),
+            "peak_allocated_gib_max_rank": float(peak), **counts,
+            "plan_sha256": sha(plan_path), "contract": contract,
+        }
+        (invocation_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
+        log(status + " " + str(invocation_dir / "summary.json"))
+    dist.barrier()
+    dist.destroy_process_group()
+    if throughput_failed:
+        raise SystemExit(42)
+
+
+if __name__ == "__main__":
+    main()

code/dataset.py

@@ -0,0 +1,167 @@
+"""Bounded random access to indexed JSONL metadata and existing codec shards.
+
+Only offsets are memory-mapped. No corpus-sized Python list or concatenated
+audio/code tensor is copied into each GPU process. TAR member indices and
+Parquet reference tables are kept in small per-thread LRU caches.
+"""
+from collections import OrderedDict
+import hashlib
+import io
+import json
+import os
+from pathlib import Path
+import tarfile
+import threading
+import numpy as np
+
+
+class ShardReader:
+    def __init__(self, capacity=16, file_capacity=None):
+        import resource
+        self.capacity = capacity
+        soft_limit = resource.getrlimit(resource.RLIMIT_NOFILE)[0]
+        self.file_capacity = file_capacity or max(32, min(2048, soft_limit // 2))
+        self.local = threading.local()
+
+    def cache(self):
+        if not hasattr(self.local, 'files'):
+            self.local.files = OrderedDict()
+            self.local.tables = OrderedDict()
+            self.local.members = OrderedDict()
+        return self.local
+
+    def bytes_at(self, path, offset, size):
+        state = self.cache()
+        if path not in state.files:
+            state.files[path] = os.open(path, os.O_RDONLY)
+            while len(state.files) > self.file_capacity:
+                _, fd = state.files.popitem(last=False)
+                os.close(fd)
+        fd = state.files[path]
+        state.files.move_to_end(path)
+        value = os.pread(fd, int(size), int(offset))
+        if len(value) != size:
+            raise IOError(f'Short shard read: {path}@{offset}+{size}')
+        return value
+
+    def member(self, path, name):
+        state = self.cache()
+        if path not in state.members:
+            with tarfile.open(path, 'r:') as archive:
+                state.members[path] = {m.name: (m.offset_data, m.size) for m in archive if m.isfile()}
+            while len(state.members) > self.capacity:
+                state.members.popitem(last=False)
+        state.members.move_to_end(path)
+        offset, size = state.members[path][name]
+        return self.bytes_at(path, offset, size)
+
+    def codes(self, locator):
+        kind = locator.get('kind', 'tar_moss_npy')
+        if kind == 'tar_moss_npy':
+            blob = (self.bytes_at(locator['path'], locator['offset'], locator['size'])
+                    if 'offset' in locator else self.member(locator['path'], locator['member']))
+            try:
+                result = np.load(io.BytesIO(blob), allow_pickle=False)
+            except ValueError:
+                # Raw little-endian uint16 dump without NumPy header
+                # (sidecar packer wrote tobytes()); shape from byte count.
+                result = np.frombuffer(blob, dtype=np.uint16).reshape(-1, 12)
+            assert result.dtype == np.uint16
+        elif kind == 'local_moss_parquet':
+            import pyarrow.parquet as pq
+            state = self.cache()
+            key = (locator['path'], locator['codes_field'], locator['frames_field'])
+            if key not in state.tables:
+                state.tables[key] = pq.read_table(locator['path'], columns=[locator['codes_field'], locator['frames_field']])
+                while len(state.tables) > self.capacity:
+                    state.tables.popitem(last=False)
+            state.tables.move_to_end(key)
+            row = state.tables[key].slice(locator['row_index'], 1).to_pylist()[0]
+            frames = int(row[locator['frames_field']])
+            result = np.frombuffer(row[locator['codes_field']], dtype='<i2').reshape(frames, 12)
+        else:
+            raise ValueError(f'Unsupported verified codec locator: {kind}')
+        assert result.ndim == 2 and result.shape[1] == 12 and len(result) > 0
+        assert int(result.min()) >= 0 and int(result.max()) < 1024
+        expected = locator.get('full_frames')
+        if expected is not None:
+            assert len(result) == expected
+        if 'crop_frames' in locator:
+            start, length = int(locator.get('crop_start_frame', 0)), int(locator['crop_frames'])
+            assert start >= 0 and length > 0 and start + length <= len(result)
+            result = result[start:start + length]
+        return result
+
+    def record(self, locator):
+        blob = (self.bytes_at(locator['path'], locator['offset'], locator['size'])
+                if 'offset' in locator else self.member(locator['path'], locator['member']))
+        return json.loads(blob)
+
+
+class IndexedRecords:
+    def __init__(self, manifest):
+        self.manifest_path = Path(manifest)
+        self.spec = json.loads(self.manifest_path.read_text())
+        root = self.manifest_path.parent
+        self.offsets = np.load(root / self.spec['offsets'], mmap_mode='r', allow_pickle=False)
+        self.path = root / self.spec['records']
+        self.fd = os.open(self.path, os.O_RDONLY)
+        assert len(self.offsets) == self.spec['samples'] + 1
+        assert int(self.offsets[0]) == 0 and int(self.offsets[-1]) == self.path.stat().st_size
+        self.reader = ShardReader()
+
+    def __len__(self):
+        return len(self.offsets) - 1
+
+    def row(self, index):
+        start, end = int(self.offsets[index]), int(self.offsets[index + 1])
+        return json.loads(os.pread(self.fd, end - start, start))
+
+    def example(self, index, packer):
+        from modes import verify_record_mode
+        row = self.row(index)
+        verify_record_mode(row)
+        mode = row['mode']
+        meta = self.reader.record(row['record_locator'])
+        assert meta['uid'] == row['uid']
+        meta['frames'] = int(meta['moss_frames'])
+        codes = self.reader.codes(row['target_locator'])
+        assert len(codes) == meta['frames']
+        reference = None
+        if mode == 'reference':
+            assert row['reference_uid'] != row['uid'] and row['reference_evidence']
+            reference = self.reader.codes(row['reference_locator'])
+            assert 0 < len(reference) <= 37
+        elif mode == 'score':
+            meta['conditioning_measurements'] = row['conditioning_measurements']
+        result = packer.pack_mode(meta, codes, mode, reference)
+        result['accounting'] = {'uid': row['uid'], 'mode': mode, 'frames': len(codes),
+                                'target_audio_hours': float(meta['dur_s']) / 3600,
+                                'reference_frames': len(reference) if reference is not None else 0}
+        return result
+
+
+def write_manifest(directory, rows, contract):
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    offsets, n, modes, duration = [0], 0, {}, 0.
+    path = directory / 'records.jsonl'
+    digest = hashlib.sha256()
+    seen = set()
+    with path.with_suffix('.jsonl.tmp').open('wb') as stream:
+        for row in rows:
+            assert row['uid'] not in seen, row['uid']
+            seen.add(row['uid'])
+            assert row['mode'] in ('instruction', 'reference', 'score')
+            blob = (json.dumps(row, ensure_ascii=False, separators=(',', ':')) + '\n').encode()
+            stream.write(blob); digest.update(blob); offsets.append(offsets[-1] + len(blob)); n += 1
+            modes[row['mode']] = modes.get(row['mode'], 0) + 1
+            duration += row.get('dur_s', 0.) / 3600
+    path.with_suffix('.jsonl.tmp').replace(path)
+    np.save(directory / 'offsets.npy', np.asarray(offsets, dtype=np.int64), allow_pickle=False)
+    result = dict(contract, records=path.name, offsets='offsets.npy', samples=n, modes=modes,
+                  target_audio_hours=duration, records_sha256=digest.hexdigest())
+    temporary = directory / 'manifest.json.tmp'
+    temporary.write_text(json.dumps(result, indent=2) + '\n')
+    temporary.replace(directory / 'manifest.json')
+    return directory / 'manifest.json'

code/infer.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+"""Single-clip Humaneness Voice Small inference; use on a CUDA GPU.
+
+Example (from the model repository root):
+  python code/infer.py --stage S3 --prompt 'CAPTION: warm, amused narration\nTRANSCRIPT: "Hello there."' \
+      --text 'Hello there.' --frames 60 --output hello.wav
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+import sys
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / 'code'))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--stage', choices=[f'S{i}' for i in range(1, 11)], default='S3')
+    parser.add_argument('--prompt', required=True, help='Literal GENERAL/SCRIPT, CAPTION/TRANSCRIPT or TRANSCRIPT text')
+    parser.add_argument('--text', required=True, help='The exact spoken transcript')
+    parser.add_argument('--frames', type=int, required=True, help='Frame budget at 12.5 frames/second')
+    parser.add_argument('--language', choices=('en', 'de'), default='en')
+    parser.add_argument('--reference-wav', type=Path, help='Optional distinct reference recording')
+    parser.add_argument('--seed', type=int, default=777)
+    parser.add_argument('--output', type=Path, required=True)
+    args = parser.parse_args()
+    assert args.frames > 0
+
+    import numpy as np
+    import soundfile as sf
+    import torch
+    import moss_small
+    from large_talker import build_fresh
+    from packing import ScorePacker, generated_audio
+
+    if not torch.cuda.is_available():
+        raise RuntimeError('This inference example requires CUDA')
+    device = torch.device('cuda:0')
+    torch.cuda.set_device(device)
+    moss_small.SFT3 = str(ROOT / 'assets/sft3')
+    moss_small.QWEN = str(ROOT / 'assets/qwen3')
+    # Full published model state includes the Qwen3 backbone weights.  Do not
+    # download the separate original pretraining file merely to overwrite it.
+    moss_small.load_qwen_backbone = lambda model, log=print: None
+    schema = json.loads((ROOT / 'assets/score_schema.json').read_text())
+    model, config = build_fresh(schema, log=lambda _: None)
+    state = torch.load(ROOT / 'checkpoints' / args.stage / 'model_bf16.pt',
+                       map_location='cpu', weights_only=True)
+    model.load_state_dict(state, strict=True)
+    model.tie_weights()
+    model = model.to(device, dtype=torch.bfloat16).eval()
+    del state
+
+    _, _, Processor = moss_small.export_classes()
+    processor = Processor.from_pretrained(
+        moss_small.SFT3, codec_path='OpenMOSS-Team/MOSS-Audio-Tokenizer-v2',
+        codec_weight_dtype='fp32', codec_compute_dtype='bf16')
+    processor.audio_tokenizer = processor.audio_tokenizer.to(device).eval()
+    packer = ScorePacker(processor, config, schema)
+    reference = None
+    if args.reference_wav:
+        reference = processor.encode_audios_from_path(args.reference_wav, n_vq=12)[0]
+        if not 0 < len(reference) <= 37:
+            raise ValueError('Reference must be 1–37 codec frames (at most 2.96 seconds)')
+    mode = 'reference' if reference is not None else 'instruction'
+    record = {'prompt': args.prompt, 'text': args.text, 'frames': args.frames,
+              'lang': args.language}
+    example = packer.pack_mode(record, [], mode, reference, generation=True)
+    batch = packer.collate([example])
+    ids = batch['input_ids'].to(device)
+    mask = batch['attention_mask'].to(device)
+    conditioning = tuple(t.to(device) for t in batch['score_conditioning'])
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+    with (torch.inference_mode(), model.generation_scores(conditioning),
+          torch.autocast('cuda', dtype=torch.bfloat16)):
+        result = model.generate(input_ids=ids, attention_mask=mask,
+            max_new_frames=args.frames + 60, do_sample=True,
+            audio_temperature=1.0, audio_top_p=0.95, audio_top_k=50,
+            audio_repetition_penalty=1.0, use_kv_cache=True)
+    codes = generated_audio(result, config).cpu()
+    if len(codes) < 2:
+        raise RuntimeError('Generated fewer than two codec frames')
+    waveform = processor.decode_audio_codes([codes.to(device)], return_stereo=False)[0]
+    wave = np.asarray(waveform.float().cpu().numpy()).reshape(-1)
+    sample_rate = int(processor.model_config.sampling_rate)
+    assert np.isfinite(wave).all()
+    assert abs(len(wave) / len(codes) - sample_rate / 12.5) <= 0.01 * sample_rate / 12.5
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    sf.write(args.output, wave, sample_rate, subtype='PCM_16')
+    print(json.dumps({'output': str(args.output), 'frames': len(codes),
+                      'duration_seconds': len(wave) / sample_rate, 'stage': args.stage}))
+
+
+if __name__ == '__main__':
+    main()

code/large_talker.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+"""Qwen3-0.6B semantic path plus a fresh SFT-3-width local Talker.
+
+This is the existing M1 bridge architecture without loading the SFT-3 Talker
+weights and without the experimental K4 memory module.  Qwen starts from the
+original pretrained Qwen3-0.6B weights; every local/audio/bridge parameter is
+fresh and trainable.
+"""
+from __future__ import annotations
+
+from torch import nn
+
+
+EXPECTED_LOCAL_TALKER_PARAMETERS = 112_764_928
+
+
+def build_fresh(schema: dict, log=print):
+    import moss_small
+    from conditioning import scored_class
+
+    config = moss_small.make_config("M1")
+    model = scored_class()(config, schema)
+    moss_small.load_qwen_backbone(model, log=log)
+
+    # The score adapter is not part of this caption-conditioned run.  It stays
+    # frozen and receives no score-token prompts, exactly as in the old ladder.
+    for parameter in model.score_conditioner.parameters():
+        parameter.requires_grad_(False)
+
+    names = (
+        "local_transformer.", "audio_embeddings.", "local_text_lm_head.",
+        "proj_in.", "proj_out.",
+    )
+    local = sum(parameter.numel() for name, parameter in model.named_parameters()
+                if name.startswith(names))
+    assert local == EXPECTED_LOCAL_TALKER_PARAMETERS, local
+    assert int(config.hidden_size) == 1024
+    assert int(config.local_hidden_size) == 2560
+    assert int(config.n_vq) == 12
+    for index in range(int(config.n_vq)):
+        assert model.audio_lm_heads[index].weight.data_ptr() == \
+               model.audio_embeddings[index].weight.data_ptr()
+    return model, config
+
+
+def parameter_counts(model: nn.Module) -> dict:
+    backbone_prefixes = ("transformer.", "text_lm_head.")
+    backbone = sum(parameter.numel() for name, parameter in model.named_parameters()
+                   if name.startswith(backbone_prefixes) and parameter.requires_grad)
+    head = sum(parameter.numel() for name, parameter in model.named_parameters()
+               if not name.startswith(backbone_prefixes) and parameter.requires_grad)
+    frozen = sum(parameter.numel() for parameter in model.parameters()
+                 if not parameter.requires_grad)
+    assert head == EXPECTED_LOCAL_TALKER_PARAMETERS, head
+    return {
+        "trainable_backbone_parameters": backbone,
+        "trainable_talker_parameters": head,
+        "frozen_score_conditioner_parameters": frozen,
+        "total_parameters": sum(parameter.numel() for parameter in model.parameters()),
+    }

code/moss_export/__init__.py

@@ -0,0 +1,9 @@
+from .configuration_moss_tts import MossTTSLocalConfig
+from .modeling_moss_tts import MossTTSLocalModel
+from .processing_moss_tts import MossTTSLocalProcessor
+
+__all__ = [
+    "MossTTSLocalConfig",
+    "MossTTSLocalModel",
+    "MossTTSLocalProcessor",
+]

code/moss_export/configuration_moss_tts.py

@@ -0,0 +1,158 @@
+# coding=utf-8
+"""Configuration for the MOSS-TTS-Local-Transformer-v1.5 release."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+
+
+SUPPORTED_ATTENTION_IMPLEMENTATIONS = {"flash_attention_2", "sdpa", "eager"}
+
+
+def _normalize_attention_implementation(value: Optional[str], default: str = "flash_attention_2") -> str:
+    normalized = str(value or default).strip().lower()
+    if normalized in {"flash", "flash_attn", "flash-attn", "flash_attention"}:
+        normalized = "flash_attention_2"
+    if normalized not in SUPPORTED_ATTENTION_IMPLEMENTATIONS:
+        raise ValueError(
+            "attn_implementation must be one of "
+            f"{sorted(SUPPORTED_ATTENTION_IMPLEMENTATIONS)}, got {value!r}."
+        )
+    return normalized
+
+
+class MossTTSLocalConfig(PretrainedConfig):
+    model_type = "moss_tts_local"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        qwen3_config: Optional[Union[Qwen3Config, Dict[str, Any]]] = None,
+        gpt2_config: Optional[Union[GPT2Config, Dict[str, Any]]] = None,
+        language_config: Optional[Union[Qwen3Config, Dict[str, Any]]] = None,
+        n_vq: int = 12,
+        audio_vocab_size: int = 1024,
+        audio_codebook_sizes: Optional[list[int]] = None,
+        audio_pad_token_id: int = 1024,
+        audio_pad_code: Optional[int] = None,
+        pad_token_id: int = 151643,
+        im_start_token_id: int = 151644,
+        im_end_token_id: int = 151645,
+        audio_start_token_id: int = 151669,
+        audio_end_token_id: int = 151670,
+        audio_user_slot_token_id: int = 151654,
+        audio_assistant_slot_token_id: int = 151656,
+        audio_assistant_gen_slot_token_id: Optional[int] = None,
+        sampling_rate: int = 48000,
+        audio_tokenizer_name_or_path: Optional[str] = None,
+        attn_implementation: str = "flash_attention_2",
+        local_transformer_attn_implementation: Optional[str] = None,
+        local_text_head_mode: str = "binary",
+        initializer_range: float = 0.02,
+        **kwargs: Any,
+    ) -> None:
+        if qwen3_config is None and language_config is not None:
+            qwen3_config = language_config
+        if isinstance(qwen3_config, dict):
+            self.qwen3_config = Qwen3Config(**qwen3_config)
+        elif qwen3_config is None:
+            self.qwen3_config = Qwen3Config()
+        else:
+            self.qwen3_config = qwen3_config
+
+        if isinstance(gpt2_config, dict):
+            self.gpt2_config = GPT2Config(**gpt2_config)
+        elif gpt2_config is None:
+            self.gpt2_config = GPT2Config(
+                vocab_size=int(self.qwen3_config.vocab_size),
+                n_embd=int(self.qwen3_config.hidden_size),
+                n_layer=1,
+                n_head=max(1, int(self.qwen3_config.hidden_size) // 80),
+                n_positions=int(n_vq) + 1,
+                n_ctx=int(n_vq) + 1,
+                activation_function="silu",
+                layer_norm_epsilon=1e-6,
+                resid_pdrop=0.0,
+                embd_pdrop=0.0,
+                attn_pdrop=0.0,
+            )
+        else:
+            self.gpt2_config = gpt2_config
+
+        self.n_vq = int(n_vq)
+        if self.n_vq <= 0:
+            raise ValueError("n_vq must be positive.")
+        if audio_codebook_sizes is None:
+            self.audio_codebook_sizes = [int(audio_vocab_size)] * self.n_vq
+        else:
+            self.audio_codebook_sizes = [int(size) for size in audio_codebook_sizes]
+        if len(self.audio_codebook_sizes) != self.n_vq:
+            raise ValueError(
+                f"audio_codebook_sizes must have length n_vq={self.n_vq}, "
+                f"got {len(self.audio_codebook_sizes)}."
+            )
+        if any(size <= 0 for size in self.audio_codebook_sizes):
+            raise ValueError("audio_codebook_sizes must contain positive integers.")
+        self.audio_vocab_size = int(max(int(audio_vocab_size), max(self.audio_codebook_sizes)))
+        self.audio_pad_token_id = int(audio_pad_code if audio_pad_code is not None else audio_pad_token_id)
+        self.audio_pad_code = self.audio_pad_token_id
+        if self.audio_pad_token_id < self.audio_vocab_size:
+            raise ValueError("audio_pad_token_id/audio_pad_code must be outside the audio vocab.")
+
+        self.pad_token_id = int(pad_token_id)
+        self.im_start_token_id = int(im_start_token_id)
+        self.im_end_token_id = int(im_end_token_id)
+        self.audio_start_token_id = int(audio_start_token_id)
+        self.audio_end_token_id = int(audio_end_token_id)
+        self.audio_user_slot_token_id = int(audio_user_slot_token_id)
+        self.audio_assistant_slot_token_id = int(
+            audio_assistant_slot_token_id
+            if audio_assistant_gen_slot_token_id is None
+            else audio_assistant_gen_slot_token_id
+        )
+        self.audio_assistant_gen_slot_token_id = self.audio_assistant_slot_token_id
+
+        self.sampling_rate = int(sampling_rate)
+        self.audio_tokenizer_name_or_path = audio_tokenizer_name_or_path
+        self.attn_implementation = _normalize_attention_implementation(attn_implementation)
+        self.local_transformer_attn_implementation = _normalize_attention_implementation(
+            local_transformer_attn_implementation,
+            default=self.attn_implementation,
+        )
+        self.initializer_range = float(initializer_range)
+
+        self.hidden_size = int(self.qwen3_config.hidden_size)
+        self.vocab_size = int(self.qwen3_config.vocab_size)
+        self.local_hidden_size = int(self.gpt2_config.hidden_size)
+        if self.local_hidden_size != self.hidden_size:
+            raise ValueError(
+                "This MOSS-TTS-Local-Transformer-v1.5 release expects local hidden size to "
+                "match Qwen3 hidden size so audio embeddings and heads are tied."
+            )
+
+        normalized_text_head_mode = str(local_text_head_mode or "full_vocab").strip().lower()
+        if normalized_text_head_mode in {"full", "full-vocab", "vocab"}:
+            normalized_text_head_mode = "full_vocab"
+        if normalized_text_head_mode not in {"full_vocab", "binary"}:
+            raise ValueError("local_text_head_mode must be 'full_vocab' or 'binary'.")
+        self.local_text_head_mode = normalized_text_head_mode
+
+        kwargs.setdefault("tie_word_embeddings", True)
+        super().__init__(pad_token_id=self.pad_token_id, **kwargs)
+
+    @property
+    def language_config(self) -> Qwen3Config:
+        return self.qwen3_config
+
+    def to_dict(self) -> Dict[str, Any]:
+        output = super().to_dict()
+        output["qwen3_config"] = self.qwen3_config.to_dict()
+        output["language_config"] = self.qwen3_config.to_dict()
+        output["gpt2_config"] = self.gpt2_config.to_dict()
+        output["audio_pad_code"] = self.audio_pad_token_id
+        output["audio_assistant_gen_slot_token_id"] = self.audio_assistant_slot_token_id
+        return output

code/moss_export/gpt2_decoder.py

@@ -0,0 +1,721 @@
+# 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
+
+
+def _is_static_kv_cache_layer(layer_past: object) -> bool:
+    return isinstance(layer_past, dict) and bool(layer_past.get("static_kv_cache", False))
+
+
+class MossTTSNanoGPT2RotaryEmbedding(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}")
+        self.dim = int(dim)
+        self.base = float(base)
+        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)
+
+    def _compute_inv_freq(self, device: Optional[torch.device] = None) -> torch.Tensor:
+        return 1.0 / (
+            self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim)
+        )
+
+    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)
+        inv_freq = self._compute_inv_freq(device=device)
+        freqs = torch.einsum("bs,d->bsd", position_ids.to(device=device, dtype=inv_freq.dtype), 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 MossTTSNanoGPT2MLP(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 MossTTSNanoGPT2Attention(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 = MossTTSNanoGPT2RotaryEmbedding(
+                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
+        if attention_mask is not None or query.shape[-2] != key.shape[-2]:
+            mask = self._causal_attention_mask(
+                attention_mask=attention_mask,
+                query_length=query.shape[-2],
+                key_length=key.shape[-2],
+                device=query.device,
+            )
+        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,
+        )
+        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,
+            )
+
+        if query.shape[1] != key.shape[1]:
+            query_attention_mask = attention_mask[:, -query.shape[1] :]
+            unpadded_query, query_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
+                query,
+                query_attention_mask,
+            )
+            unpadded_key, _, cu_seqlens_k, max_seqlen_k, _ = 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_q,
+                cu_seqlens_k,
+                max_seqlen_q,
+                max_seqlen_k,
+                dropout_p=dropout_p,
+                causal=True,
+            )
+            return pad_input(output, query_indices, query.shape[0], query.shape[1])
+
+        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)
+
+        static_layer_past = layer_past is not None and _is_static_kv_cache_layer(layer_past)
+        if static_layer_past:
+            past_length = int(layer_past.get("length", 0))
+            new_length = past_length + int(key.shape[1])
+            key_cache = layer_past["key"]
+            value_cache = layer_past["value"]
+            if new_length > int(key_cache.shape[1]):
+                raise ValueError(
+                    f"Static KV cache is too short: need {new_length}, capacity={int(key_cache.shape[1])}."
+                )
+            key_cache[:, past_length:new_length].copy_(key)
+            value_cache[:, past_length:new_length].copy_(value)
+            key = key_cache[:, :new_length]
+            value = value_cache[:, :new_length]
+            layer_past["length"] = new_length
+        elif 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 = layer_past if (use_cache and static_layer_past) else ((key, value) if use_cache else None)
+
+        if self.attn_implementation == "flash_attention_2":
+            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 MossTTSNanoGPT2Block(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 = MossTTSNanoGPT2Attention(config, layer_idx=layer_idx, attn_implementation=attn_implementation)
+        self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.mlp = MossTTSNanoGPT2MLP(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 MossTTSNanoGPT2Model(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(
+            [MossTTSNanoGPT2Block(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:
+            diffs = cu_seqlens[:, 1:] - cu_seqlens[:, :-1]
+            return diffs.gt(0).sum(dim=-1).to(device=device, dtype=torch.long)
+        if num_sequences.ndim == 0:
+            num_sequences = num_sequences.unsqueeze(0)
+        return num_sequences.to(device=device, dtype=torch.long)
+
+    @staticmethod
+    def _packed_segments_from_cu_seqlens(
+        cu_seqlens: torch.Tensor,
+        num_sequences: Optional[torch.Tensor],
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if cu_seqlens.ndim == 1:
+            cu_seqlens = cu_seqlens.unsqueeze(0)
+        cu_seqlens = cu_seqlens.to(device=device)
+        batch_size, boundary_count = cu_seqlens.shape
+        segment_slots = boundary_count - 1
+        if segment_slots <= 0:
+            empty = torch.empty(0, dtype=torch.long, device=device)
+            return empty, empty, empty
+
+        counts = MossTTSNanoGPT2Model._normalize_num_sequences(cu_seqlens, num_sequences, device=device)
+        counts = counts.clamp(min=0, max=segment_slots)
+        segment_slots = int(counts.max().item()) if counts.numel() > 0 else 0
+        if segment_slots <= 0:
+            empty = torch.empty(0, dtype=torch.long, device=device)
+            return empty, empty, empty
+        cu_seqlens = cu_seqlens[:, : segment_slots + 1]
+        slot_ids = torch.arange(segment_slots, device=device).unsqueeze(0)
+        valid_slots = slot_ids < counts.unsqueeze(1)
+
+        starts = cu_seqlens[:, :-1].to(dtype=torch.long)
+        ends = cu_seqlens[:, 1:].to(dtype=torch.long)
+        lengths = (ends - starts).clamp_min(0)
+        lengths = torch.where(valid_slots, lengths, torch.zeros((), dtype=torch.long, device=device))
+
+        batch_ids = torch.arange(batch_size, device=device, dtype=torch.long).unsqueeze(1).expand(batch_size, segment_slots)
+        batch_ids = batch_ids.reshape(-1)
+        starts = starts.reshape(-1)
+        lengths = lengths.reshape(-1)
+
+        valid_segments = lengths.gt(0)
+        valid_count = int(valid_segments.to(dtype=torch.long).sum().item())
+        if valid_count <= 0:
+            empty = torch.empty(0, dtype=torch.long, device=device)
+            return empty, empty, empty
+        if valid_count == lengths.numel():
+            return batch_ids, starts, lengths
+
+        valid_order = torch.argsort(valid_segments.to(dtype=torch.long), descending=True, stable=True)[:valid_count]
+        return (
+            batch_ids.index_select(0, valid_order),
+            starts.index_select(0, valid_order),
+            lengths.index_select(0, valid_order),
+        )
+
+    @staticmethod
+    def _packed_token_indices(
+        batch_ids: torch.Tensor,
+        starts: torch.Tensor,
+        lengths: torch.Tensor,
+        seq_len: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        total_tokens = int(lengths.sum().item())
+        if total_tokens <= 0:
+            empty = torch.empty(0, dtype=torch.long, device=lengths.device)
+            return empty, empty
+
+        segment_ids = torch.repeat_interleave(
+            torch.arange(lengths.numel(), device=lengths.device, dtype=torch.long),
+            lengths,
+            output_size=total_tokens,
+        )
+        segment_starts = torch.cumsum(lengths, dim=0) - lengths
+        positions = torch.arange(total_tokens, device=lengths.device, dtype=torch.long) - segment_starts[segment_ids]
+        indices = batch_ids[segment_ids] * seq_len + starts[segment_ids] + positions
+        return indices, positions
+
+    @staticmethod
+    def build_packed_position_ids(
+        attention_mask: Optional[torch.Tensor],
+        cu_seqlens: torch.Tensor,
+        num_sequences: Optional[torch.Tensor],
+        sequence_length: Optional[int] = None,
+    ) -> torch.Tensor:
+        if cu_seqlens.ndim == 1:
+            cu_seqlens = cu_seqlens.unsqueeze(0)
+        batch_size = cu_seqlens.shape[0]
+        seq_len = int(sequence_length or (cu_seqlens.shape[1] - 1))
+        device = cu_seqlens.device
+        position_ids = torch.zeros((batch_size, seq_len), dtype=torch.long, device=device)
+        batch_ids, starts, lengths = MossTTSNanoGPT2Model._packed_segments_from_cu_seqlens(
+            cu_seqlens,
+            num_sequences,
+            device,
+        )
+        if lengths.numel() > 0:
+            indices, positions = MossTTSNanoGPT2Model._packed_token_indices(batch_ids, starts, lengths, seq_len)
+            position_ids.view(-1).scatter_(0, indices, positions)
+        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
+        seq_len = hidden_states.shape[1]
+        batch_ids, starts, lengths = MossTTSNanoGPT2Model._packed_segments_from_cu_seqlens(
+            cu_seqlens,
+            num_sequences,
+            device,
+        )
+        if lengths.numel() == 0:
+            raise ValueError("cu_seqlens did not describe any non-empty packed sequences.")
+
+        indices, _ = MossTTSNanoGPT2Model._packed_token_indices(batch_ids, starts, lengths, seq_len)
+        cumulative = torch.empty(lengths.numel() + 1, dtype=torch.int32, device=device)
+        cumulative[0] = 0
+        cumulative[1:] = lengths.to(dtype=torch.int32).cumsum(dim=0)
+        return PackedSequenceMetadata(
+            cu_seqlens=cumulative,
+            max_seqlen=int(lengths.max().item()),
+            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
+        query_attention_mask = None
+        if attention_mask is not None:
+            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:
+                if attention_mask is None:
+                    raise ValueError("attention_mask must be provided with cu_seqlens packing.")
+                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,
+                    sequence_length=hidden_states.shape[1],
+                )
+            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:
+                    first_layer_past = past_key_values[0]
+                    if _is_static_kv_cache_layer(first_layer_past):
+                        past_length = int(first_layer_past.get("length", 0))
+                    else:
+                        past_length = first_layer_past[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)
+        if query_attention_mask is not None:
+            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,
+                )
+            if query_attention_mask is not None:
+                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)
+        if query_attention_mask is not None:
+            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,
+        )

code/moss_export/modeling_moss_tts.py

@@ -0,0 +1,623 @@
+# coding=utf-8
+"""Modeling code for the MOSS-TTS-Local-Transformer-v1.5 HuggingFace release."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+import torch
+import torch.nn as nn
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
+from transformers.utils import ModelOutput
+
+from .configuration_moss_tts import MossTTSLocalConfig
+from .gpt2_decoder import MossTTSNanoGPT2Model
+from .qwen3_decoder import MossQwen3Model
+
+
+@dataclass
+class MossTTSLocalOutput(ModelOutput):
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[tuple[tuple[torch.Tensor, torch.Tensor], ...]] = None
+    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+
+
+def _find_last_equal(input_ids: torch.LongTensor, value: int) -> torch.LongTensor:
+    matches = input_ids.eq(int(value))
+    if not bool(matches.any(dim=1).all().item()):
+        raise ValueError(f"Every sample must contain token id {int(value)}.")
+    positions = torch.arange(input_ids.shape[1], device=input_ids.device, dtype=torch.long)
+    masked_positions = positions.unsqueeze(0).masked_fill(~matches, -1)
+    return masked_positions.max(dim=1).values
+
+
+class MossTTSLocalPreTrainedModel(PreTrainedModel):
+    config_class = MossTTSLocalConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MossTTSNanoGPT2Block", "MossQwen3DecoderLayer"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False) -> None:
+        if isinstance(module, MossTTSNanoGPT2Model) or isinstance(module, MossQwen3Model):
+            module.gradient_checkpointing = value
+
+
+class MossTTSLocalModel(MossTTSLocalPreTrainedModel):
+    _tied_weights_keys = None
+
+    def __init__(self, config: MossTTSLocalConfig) -> None:
+        super().__init__(config)
+        self._tied_weights_keys = self._build_tied_weights_keys(config)
+
+        config.qwen3_config.pad_token_id = config.pad_token_id
+        config.qwen3_config._attn_implementation = config.attn_implementation
+        local_gpt2_config = config.gpt2_config.to_dict()
+        local_gpt2_config["n_layer"] = int(getattr(config, "local_transformer_layers", config.gpt2_config.n_layer))
+        local_gpt2_config["n_positions"] = int(config.n_vq) + 1
+        local_gpt2_config["n_ctx"] = int(config.n_vq) + 1
+        local_gpt2_config = GPT2Config(**local_gpt2_config)
+        local_gpt2_config.pad_token_id = config.pad_token_id
+        local_gpt2_config._attn_implementation = config.local_transformer_attn_implementation
+
+        self.transformer = MossQwen3Model(config.qwen3_config)
+        self.local_transformer = MossTTSNanoGPT2Model(
+            local_gpt2_config,
+            attn_implementation=config.local_transformer_attn_implementation,
+        )
+        self.local_transformer.wte = nn.Identity()
+
+        hidden_size = int(config.hidden_size)
+        self.audio_embeddings = nn.ModuleList(
+            [
+                nn.Embedding(int(config.audio_codebook_sizes[index]), hidden_size)
+                for index in range(config.n_vq)
+            ]
+        )
+        self.text_lm_head = nn.Linear(hidden_size, int(config.vocab_size), bias=False)
+        self.audio_lm_heads = nn.ModuleList(
+            [
+                nn.Linear(hidden_size, int(config.audio_codebook_sizes[index]), bias=False)
+                for index in range(config.n_vq)
+            ]
+        )
+        self.local_text_lm_head = (
+            nn.Linear(hidden_size, 2, bias=False)
+            if self._use_binary_local_text_head()
+            else None
+        )
+
+        self.post_init()
+        self.tie_weights()
+        self.initialize_local_text_lm_head_from_text_lm_head()
+
+    def can_generate(self) -> bool:
+        return True
+
+    @staticmethod
+    def _build_tied_weights_keys(config: MossTTSLocalConfig) -> dict[str, str]:
+        tied_weights = {"text_lm_head.weight": "transformer.embed_tokens.weight"}
+        tied_weights.update(
+            {
+                f"audio_lm_heads.{index}.weight": f"audio_embeddings.{index}.weight"
+                for index in range(config.n_vq)
+            }
+        )
+        return tied_weights
+
+    def tie_weights(self, *args, **kwargs) -> None:
+        del args, kwargs
+        self.text_lm_head.weight = self.transformer.embed_tokens.weight
+        for embedding, head in zip(self.audio_embeddings, self.audio_lm_heads):
+            head.weight = embedding.weight
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.transformer.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.transformer.embed_tokens = value
+        self.tie_weights()
+        self.initialize_local_text_lm_head_from_text_lm_head()
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.text_lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.text_lm_head = new_embeddings
+        self.tie_weights()
+        self.initialize_local_text_lm_head_from_text_lm_head()
+
+    def _use_binary_local_text_head(self) -> bool:
+        return str(getattr(self.config, "local_text_head_mode", "full_vocab")).strip().lower() == "binary"
+
+    def _local_text_candidate_ids(self, device: torch.device) -> torch.LongTensor:
+        return torch.tensor(
+            [
+                int(self.config.audio_assistant_slot_token_id),
+                int(self.config.audio_end_token_id),
+            ],
+            dtype=torch.long,
+            device=device,
+        )
+
+    def initialize_local_text_lm_head_from_text_lm_head(self) -> None:
+        if not self._use_binary_local_text_head() or self.local_text_lm_head is None:
+            return
+        candidate_ids = self._local_text_candidate_ids(self.text_lm_head.weight.device)
+        with torch.no_grad():
+            source_weight = self.text_lm_head.weight.index_select(0, candidate_ids)
+            if tuple(source_weight.shape) == tuple(self.local_text_lm_head.weight.shape):
+                self.local_text_lm_head.weight.copy_(
+                    source_weight.to(
+                        device=self.local_text_lm_head.weight.device,
+                        dtype=self.local_text_lm_head.weight.dtype,
+                    )
+                )
+
+    def _resolve_fixed_nq(
+        self,
+        n_vq_for_inference: Optional[int] = None,
+        nq: Optional[int] = None,
+    ) -> int:
+        requested = n_vq_for_inference if n_vq_for_inference is not None else nq
+        config_nq = int(self.config.n_vq)
+        if requested is not None and int(requested) != config_nq:
+            raise ValueError(
+                "This MOSS-TTS-Local-Transformer-v1.5 release is trained with a fixed RVQ depth. "
+                f"Expected n_vq={config_nq}, got {int(requested)}."
+            )
+        return config_nq
+
+    def _build_inputs_embeds(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
+        if input_ids.ndim != 3 or input_ids.shape[-1] != self.config.n_vq + 1:
+            raise ValueError(
+                f"Expected input_ids shape [batch, seq, {self.config.n_vq + 1}], "
+                f"got {tuple(input_ids.shape)}."
+            )
+        text_ids = input_ids[..., 0]
+        inputs_embeds = self.transformer.embed_tokens(text_ids)
+        for channel_index, embedding in enumerate(self.audio_embeddings):
+            channel_ids = input_ids[..., channel_index + 1]
+            valid_mask = channel_ids.ne(self.config.audio_pad_token_id)
+            safe_ids = channel_ids.masked_fill(~valid_mask, 0)
+            audio_embeds = embedding(safe_ids) * valid_mask.unsqueeze(-1)
+            inputs_embeds = inputs_embeds + audio_embeds
+        return inputs_embeds
+
+    def _global_hidden_to_local(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+        return hidden_states
+
+    @staticmethod
+    def _local_past_length(past_key_values: Optional[tuple[Any, ...]]) -> int:
+        if past_key_values is None or len(past_key_values) == 0:
+            return 0
+        first_layer_past = past_key_values[0]
+        if isinstance(first_layer_past, dict) and bool(first_layer_past.get("static_kv_cache", False)):
+            return int(first_layer_past.get("length", 0))
+        return int(first_layer_past[0].shape[1])
+
+    def _new_static_local_past_key_values(
+        self,
+        batch_size: int,
+        max_length: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[dict[str, Any], ...]:
+        layers = []
+        for block in self.local_transformer.h:
+            attn = block.attn
+            cache_shape = (
+                int(batch_size),
+                int(max_length),
+                int(attn.num_heads),
+                int(attn.head_dim),
+            )
+            layers.append(
+                {
+                    "static_kv_cache": True,
+                    "key": torch.empty(cache_shape, device=device, dtype=dtype),
+                    "value": torch.empty(cache_shape, device=device, dtype=dtype),
+                    "length": 0,
+                }
+            )
+        return tuple(layers)
+
+    def _decode_local_hidden_states_with_cache(
+        self,
+        local_inputs_embeds: torch.FloatTensor,
+        past_key_values: Optional[tuple[Any, ...]] = None,
+    ) -> tuple[torch.FloatTensor, Optional[tuple[Any, ...]]]:
+        if (
+            past_key_values is None
+            and not self.training
+            and bool(getattr(self.config, "use_static_local_kv_cache", True))
+        ):
+            max_length = max(int(getattr(self.config, "n_vq", 0)) + 1, int(local_inputs_embeds.shape[1]))
+            past_key_values = self._new_static_local_past_key_values(
+                batch_size=int(local_inputs_embeds.shape[0]),
+                max_length=max_length,
+                device=local_inputs_embeds.device,
+                dtype=local_inputs_embeds.dtype,
+            )
+        past_length = self._local_past_length(past_key_values)
+        local_seq_len = int(local_inputs_embeds.shape[1])
+        local_position_ids = torch.arange(
+            past_length,
+            past_length + local_seq_len,
+            device=local_inputs_embeds.device,
+            dtype=torch.long,
+        ).unsqueeze(0)
+        if int(local_inputs_embeds.shape[0]) != 1:
+            local_position_ids = local_position_ids.expand(int(local_inputs_embeds.shape[0]), -1)
+        local_outputs = self.local_transformer(
+            input_ids=None,
+            past_key_values=past_key_values,
+            attention_mask=None,
+            position_ids=local_position_ids,
+            inputs_embeds=local_inputs_embeds,
+            use_cache=True,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+            cu_seqlens=None,
+            num_sequences=None,
+        )
+        return local_outputs.last_hidden_state, local_outputs.past_key_values
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[tuple[tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = True,
+        **kwargs,
+    ) -> Union[tuple, MossTTSLocalOutput]:
+        del kwargs
+        if inputs_embeds is None:
+            if input_ids is None:
+                raise ValueError("Either input_ids or inputs_embeds must be provided.")
+            inputs_embeds = self._build_inputs_embeds(input_ids)
+        outputs = self.transformer(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cu_seqlens=None,
+            num_sequences=None,
+        )
+        if not return_dict:
+            return (
+                outputs.last_hidden_state,
+                outputs.past_key_values,
+                outputs.hidden_states,
+                outputs.attentions,
+            )
+        return MossTTSLocalOutput(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def _decode_local_last_hidden_state(
+        self,
+        local_inputs_embeds: torch.FloatTensor,
+    ) -> torch.FloatTensor:
+        local_seq_len = int(local_inputs_embeds.shape[1])
+        local_position_ids = torch.arange(
+            0,
+            local_seq_len,
+            device=local_inputs_embeds.device,
+            dtype=torch.long,
+        ).unsqueeze(0)
+        if int(local_inputs_embeds.shape[0]) != 1:
+            local_position_ids = local_position_ids.expand(int(local_inputs_embeds.shape[0]), -1)
+        local_outputs = self.local_transformer(
+            input_ids=None,
+            attention_mask=None,
+            position_ids=local_position_ids,
+            inputs_embeds=local_inputs_embeds,
+            use_cache=False,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+            cu_seqlens=None,
+            num_sequences=None,
+        )
+        return local_outputs.last_hidden_state[:, -1, :]
+
+    def _filter_logits(
+        self,
+        logits: torch.FloatTensor,
+        top_k: Optional[int],
+        top_p: Optional[float],
+    ) -> torch.FloatTensor:
+        scores = logits
+        if top_k is not None and int(top_k) > 0 and int(top_k) < scores.shape[-1]:
+            kth = torch.topk(scores, int(top_k), dim=-1).values[..., -1, None]
+            scores = scores.masked_fill(scores < kth, -torch.inf)
+        if top_p is not None and 0.0 < float(top_p) < 1.0:
+            sorted_scores, sorted_indices = torch.sort(scores, descending=True, dim=-1)
+            sorted_probs = torch.softmax(sorted_scores, dim=-1)
+            cumulative_probs = sorted_probs.cumsum(dim=-1)
+            sorted_mask = cumulative_probs > float(top_p)
+            sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
+            sorted_mask[..., 0] = False
+            remove_mask = torch.zeros_like(scores, dtype=torch.bool)
+            remove_mask.scatter_(dim=-1, index=sorted_indices, src=sorted_mask)
+            scores = scores.masked_fill(remove_mask, -torch.inf)
+        return scores
+
+    def _apply_repetition_penalty(
+        self,
+        scores: torch.FloatTensor,
+        previous_token_ids: Optional[torch.LongTensor],
+        penalty: float,
+    ) -> torch.FloatTensor:
+        if previous_token_ids is None or float(penalty) == 1.0:
+            return scores
+        if previous_token_ids.ndim == 1:
+            previous_token_ids = previous_token_ids.unsqueeze(0)
+        updated = scores.clone()
+        for batch_index in range(updated.shape[0]):
+            unique_token_ids = torch.unique(previous_token_ids[batch_index])
+            unique_token_ids = unique_token_ids[
+                (unique_token_ids >= 0) & (unique_token_ids < updated.shape[-1])
+            ]
+            if unique_token_ids.numel() == 0:
+                continue
+            token_scores = updated[batch_index].index_select(0, unique_token_ids)
+            token_scores = torch.where(
+                token_scores < 0,
+                token_scores * float(penalty),
+                token_scores / float(penalty),
+            )
+            updated[batch_index].scatter_(0, unique_token_ids, token_scores)
+        return updated
+
+    def _sample_next_token(
+        self,
+        logits: torch.FloatTensor,
+        do_sample: bool,
+        temperature: float,
+        top_k: Optional[int],
+        top_p: Optional[float],
+        previous_token_ids: Optional[torch.LongTensor] = None,
+        repetition_penalty: float = 1.0,
+    ) -> torch.LongTensor:
+        scores = logits.float()
+        scores = self._apply_repetition_penalty(scores, previous_token_ids, repetition_penalty)
+        if not do_sample:
+            return torch.argmax(scores, dim=-1)
+        if float(temperature) <= 0:
+            raise ValueError("temperature must be positive when do_sample=True.")
+        scores = scores / float(temperature)
+        scores = self._filter_logits(scores, top_k=top_k, top_p=top_p)
+        probs = torch.softmax(scores, dim=-1)
+        return torch.multinomial(probs, num_samples=1).squeeze(-1)
+
+    def _sample_next_assistant_text_token(
+        self,
+        local_hidden_states: torch.FloatTensor,
+        do_sample: bool,
+        temperature: float,
+        top_k: Optional[int],
+        top_p: Optional[float],
+    ) -> torch.LongTensor:
+        if self._use_binary_local_text_head() and self.local_text_lm_head is not None:
+            logits = self.local_text_lm_head(local_hidden_states)
+            sampled_indices = self._sample_next_token(
+                logits=logits,
+                do_sample=do_sample,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+            )
+            candidate_ids = self._local_text_candidate_ids(logits.device)
+            return candidate_ids[sampled_indices]
+
+        candidate_ids = self._local_text_candidate_ids(local_hidden_states.device)
+        logits = self.text_lm_head(local_hidden_states).index_select(dim=-1, index=candidate_ids)
+        sampled_indices = self._sample_next_token(
+            logits=logits,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+        )
+        return candidate_ids[sampled_indices]
+
+    def _build_generation_row(
+        self,
+        batch_size: int,
+        device: torch.device,
+        audio_token_ids: torch.LongTensor,
+    ) -> torch.LongTensor:
+        row = torch.full(
+            (batch_size, 1, self.config.n_vq + 1),
+            int(self.config.audio_pad_token_id),
+            dtype=torch.long,
+            device=device,
+        )
+        row[:, :, 0] = int(self.config.audio_assistant_slot_token_id)
+        row[:, :, 1:] = audio_token_ids.unsqueeze(1)
+        return row
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        max_new_tokens: Optional[int] = None,
+        max_new_frames: Optional[int] = None,
+        do_sample: bool = True,
+        text_temperature: float = 1.0,
+        text_top_p: float = 1.0,
+        text_top_k: int = 50,
+        audio_temperature: Optional[float] = None,
+        audio_top_p: Optional[float] = None,
+        audio_top_k: Optional[int] = None,
+        audio_repetition_penalty: Optional[float] = None,
+        temperature: float = 1.0,
+        top_p: float = 0.95,
+        top_k: int = 50,
+        repetition_penalty: float = 1.0,
+        use_kv_cache: bool = True,
+        n_vq_for_inference: Optional[int] = None,
+        nq: Optional[int] = None,
+        **kwargs,
+    ) -> list[tuple[int, torch.LongTensor]]:
+        del kwargs
+        self._resolve_fixed_nq(n_vq_for_inference=n_vq_for_inference, nq=nq)
+
+        if input_ids.ndim == 2:
+            input_ids = input_ids.unsqueeze(0)
+        if input_ids.ndim != 3:
+            raise ValueError(f"Expected input_ids with 3 dims, got {tuple(input_ids.shape)}.")
+        if input_ids.shape[-1] != self.config.n_vq + 1:
+            raise ValueError(
+                f"Expected {self.config.n_vq + 1} channels from config.n_vq, got {input_ids.shape[-1]}."
+            )
+        if attention_mask is None:
+            attention_mask = torch.ones(input_ids.shape[:2], dtype=torch.bool, device=input_ids.device)
+        elif attention_mask.ndim == 1:
+            attention_mask = attention_mask.unsqueeze(0)
+        attention_mask = attention_mask.to(device=input_ids.device, dtype=torch.bool)
+
+        frame_budget = max_new_frames if max_new_frames is not None else max_new_tokens
+        if frame_budget is None:
+            frame_budget = 4096
+        frame_budget = int(frame_budget)
+
+        audio_temperature = float(temperature if audio_temperature is None else audio_temperature)
+        audio_top_p = float(top_p if audio_top_p is None else audio_top_p)
+        audio_top_k = int(top_k if audio_top_k is None else audio_top_k)
+        audio_repetition_penalty = float(
+            repetition_penalty if audio_repetition_penalty is None else audio_repetition_penalty
+        )
+
+        batch_size = input_ids.shape[0]
+        input_ids_length = input_ids.shape[1]
+        current_input_ids = input_ids
+        current_attention_mask = attention_mask
+        current_model_input_ids = current_input_ids
+        generated_frames: list[torch.LongTensor] = []
+        finished = torch.zeros(batch_size, dtype=torch.bool, device=input_ids.device)
+        past_key_values = None
+        local_dtype = self.local_transformer.ln_f.weight.dtype
+
+        for _ in range(frame_budget):
+            generated_audio_history = torch.stack(generated_frames, dim=1) if generated_frames else None
+            global_inputs_embeds = self._build_inputs_embeds(current_model_input_ids)
+            global_outputs = self.transformer(
+                input_ids=None,
+                past_key_values=past_key_values,
+                attention_mask=current_attention_mask,
+                position_ids=None,
+                inputs_embeds=global_inputs_embeds,
+                use_cache=use_kv_cache,
+                output_attentions=False,
+                output_hidden_states=False,
+                return_dict=True,
+                cu_seqlens=None,
+                num_sequences=None,
+            )
+            global_hidden_states = global_outputs.last_hidden_state[:, -1, :]
+            local_global_hidden_states = self._global_hidden_to_local(global_hidden_states).to(dtype=local_dtype)
+
+            local_prefix_hidden_states, local_prefix_past_key_values = self._decode_local_hidden_states_with_cache(
+                local_global_hidden_states.unsqueeze(1)
+            )
+            local_hidden_states = local_prefix_hidden_states[:, -1, :]
+            next_text_tokens = self._sample_next_assistant_text_token(
+                local_hidden_states=local_hidden_states,
+                do_sample=do_sample,
+                temperature=text_temperature,
+                top_k=text_top_k,
+                top_p=text_top_p,
+            )
+            should_continue = next_text_tokens.eq(int(self.config.audio_assistant_slot_token_id)) & ~finished
+            finished = finished | next_text_tokens.eq(int(self.config.audio_end_token_id))
+            if not bool(should_continue.any().item()):
+                break
+
+            next_frame_tokens = []
+            for channel_index in range(int(self.config.n_vq)):
+                channel_logits = self.audio_lm_heads[channel_index](local_hidden_states)
+                channel_token = self._sample_next_token(
+                    logits=channel_logits,
+                    do_sample=do_sample,
+                    temperature=audio_temperature,
+                    top_k=audio_top_k,
+                    top_p=audio_top_p,
+                    previous_token_ids=(
+                        None
+                        if generated_audio_history is None
+                        else generated_audio_history[:, :, channel_index]
+                    ),
+                    repetition_penalty=audio_repetition_penalty,
+                )
+                next_frame_tokens.append(channel_token)
+                if channel_index + 1 < int(self.config.n_vq):
+                    current_local_input = self.audio_embeddings[channel_index](channel_token).to(dtype=local_dtype)
+                    local_token_hidden_states, local_prefix_past_key_values = (
+                        self._decode_local_hidden_states_with_cache(
+                            current_local_input.unsqueeze(1),
+                            past_key_values=local_prefix_past_key_values,
+                        )
+                    )
+                    local_hidden_states = local_token_hidden_states[:, -1, :]
+
+            next_frame = torch.stack(next_frame_tokens, dim=-1)
+            next_frame = next_frame.masked_fill(
+                ~should_continue.unsqueeze(-1),
+                int(self.config.audio_pad_token_id),
+            )
+            generated_frames.append(next_frame)
+
+            next_row = self._build_generation_row(
+                batch_size=batch_size,
+                device=input_ids.device,
+                audio_token_ids=next_frame,
+            )
+            if bool((~should_continue).any().item()):
+                next_row[~should_continue, 0, 0] = int(self.config.pad_token_id)
+                next_row[~should_continue, 0, 1:] = int(self.config.audio_pad_token_id)
+
+            current_input_ids = torch.cat([current_input_ids, next_row], dim=1)
+            current_attention_mask = torch.cat(
+                [current_attention_mask, should_continue.unsqueeze(1)],
+                dim=1,
+            )
+            if use_kv_cache:
+                current_model_input_ids = next_row
+                past_key_values = global_outputs.past_key_values
+            else:
+                current_model_input_ids = current_input_ids
+
+        start_indices = _find_last_equal(input_ids[..., 0], int(self.config.audio_start_token_id))
+        start_lengths = input_ids_length - start_indices - 1
+        outputs: list[tuple[int, torch.LongTensor]] = []
+        for start_index, start_length, generation_ids in zip(
+            start_indices.tolist(),
+            start_lengths.tolist(),
+            current_input_ids,
+        ):
+            outputs.append((int(start_length), generation_ids[int(start_index):].detach().cpu()))
+        return outputs

code/moss_export/processing_moss_tts.py

@@ -0,0 +1,914 @@
+# coding=utf-8
+"""Processor for the MOSS-TTS-Local-Transformer-v1.5 HuggingFace release."""
+
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
+
+import torch
+import torchaudio
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+    BatchFeature,
+    PreTrainedTokenizerBase,
+    ProcessorMixin,
+    logging,
+    processing_utils,
+)
+
+from .configuration_moss_tts import MossTTSLocalConfig
+
+
+if hasattr(processing_utils, "MODALITY_TO_BASE_CLASS_MAPPING"):
+    processing_utils.MODALITY_TO_BASE_CLASS_MAPPING["audio_tokenizer"] = "PreTrainedModel"
+else:
+    processing_utils.AUTO_TO_BASE_CLASS_MAPPING["AutoModel"] = "PreTrainedModel"
+logger = logging.get_logger(__name__)
+
+AUDIO_PLACEHOLDER = "<|audio|>"
+USER_ROLE_PREFIX = "user\n"
+USER_TEMPLATE_REFERENCE_PREFIX = (
+    "<user_inst>\n"
+    "- Reference(s):\n"
+)
+USER_TEMPLATE_AFTER_REFERENCE_SUFFIX = (
+    "\n"
+    "- Text:\n"
+)
+USER_TEMPLATE_SUFFIX = "\n</user_inst>"
+ASSISTANT_TURN_PREFIX = "\n"
+ASSISTANT_ROLE_PREFIX = "assistant\n"
+USER_MESSAGE_FIELDS = (
+    "text",
+    "reference",
+    "instruction",
+    "tokens",
+    "quality",
+    "sound_event",
+    "ambient_sound",
+    "language",
+)
+
+
+def _normalize_template_value(value: Any) -> str:
+    if value is None:
+        return "None"
+    resolved = str(value).strip()
+    return resolved or "None"
+
+
+def _render_user_prompt_after_reference(
+    language_code: object | None = None,
+    prompt_fields: Optional[Dict[str, Any]] = None,
+) -> str:
+    fields = dict(prompt_fields or {})
+    return (
+        "\n- Instruction:\n"
+        + _normalize_template_value(fields.get("instruction"))
+        + "\n- Tokens:\n"
+        + _normalize_template_value(fields.get("tokens"))
+        + "\n- Quality:\n"
+        + _normalize_template_value(fields.get("quality"))
+        + "\n- Sound Event:\n"
+        + _normalize_template_value(fields.get("sound_event"))
+        + "\n- Ambient Sound:\n"
+        + _normalize_template_value(fields.get("ambient_sound"))
+        + "\n- Language:\n"
+        + _normalize_template_value(fields.get("language", language_code))
+        + USER_TEMPLATE_AFTER_REFERENCE_SUFFIX
+    )
+
+
+@dataclass
+class Message:
+    def to_dict(self) -> Dict[str, Any]:
+        raise NotImplementedError
+
+
+@dataclass
+class UserMessage(Message):
+    text: Optional[str] = None
+    reference: Optional[List[Optional[Union[str, os.PathLike, torch.Tensor]]]] = None
+    instruction: Optional[str] = None
+    tokens: Optional[int] = None
+    quality: Optional[str] = None
+    sound_event: Optional[str] = None
+    ambient_sound: Optional[str] = None
+    language: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        template = """<user_inst>
+- Reference(s):
+{reference}
+- Instruction:
+{instruction}
+- Tokens:
+{tokens}
+- Quality:
+{quality}
+- Sound Event:
+{sound_event}
+- Ambient Sound:
+{ambient_sound}
+- Language:
+{language}
+- Text:
+{text}
+</user_inst>"""
+
+        audio_codes_list: list[Union[str, os.PathLike, torch.Tensor]] = []
+        if self.reference is None:
+            reference = "None"
+        else:
+            reference_items: list[str] = []
+            for speaker_idx, speaker_reference in enumerate(self.reference):
+                if speaker_reference is None:
+                    continue
+                # Keep raw audio placeholders directly under "- Reference(s):".
+                # Speaker labels such as "[S1]:" change the token sequence and
+                # can affect voice-clone conditioning.
+                reference_items.append(AUDIO_PLACEHOLDER)
+                audio_codes_list.append(speaker_reference)
+            reference = "\n".join(reference_items) if reference_items else "None"
+
+        self._content = (
+            template.replace("{reference}", str(reference))
+            .replace("{instruction}", str(self.instruction))
+            .replace("{tokens}", str(self.tokens))
+            .replace("{quality}", str(self.quality))
+            .replace("{sound_event}", str(self.sound_event))
+            .replace("{ambient_sound}", str(self.ambient_sound))
+            .replace("{language}", str(self.language))
+            .replace("{text}", str(self.text))
+        )
+        self._audio_codes_list = audio_codes_list
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "role": "user",
+            "content": self._content,
+            "audio_codes_list": self._audio_codes_list,
+            "text": self.text,
+            "instruction": self.instruction,
+            "tokens": self.tokens,
+            "quality": self.quality,
+            "sound_event": self.sound_event,
+            "ambient_sound": self.ambient_sound,
+            "language": self.language,
+        }
+
+
+@dataclass
+class AssistantMessage(Message):
+    audio_codes_list: List[Union[str, os.PathLike, torch.Tensor]]
+    content: str = AUDIO_PLACEHOLDER
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "role": "assistant",
+            "content": self.content,
+            "audio_codes_list": self.audio_codes_list,
+        }
+
+
+class MossTTSLocalProcessor(ProcessorMixin):
+    attributes = ["tokenizer"]
+    tokenizer_class = "AutoTokenizer"
+    audio_tokenizer_class = "AutoModel"
+
+    tokenizer: PreTrainedTokenizerBase
+    audio_tokenizer: Any
+
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        audio_tokenizer: Any = None,
+        model_config: Optional[MossTTSLocalConfig] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(tokenizer=tokenizer, audio_tokenizer=audio_tokenizer, **kwargs)
+        self.tokenizer = tokenizer
+        self.audio_tokenizer = audio_tokenizer
+        self.model_config = model_config or MossTTSLocalConfig()
+
+        def _id_to_token(token_id: int) -> str:
+            token = tokenizer.convert_ids_to_tokens(int(token_id))
+            if isinstance(token, list):
+                return token[0] if token else ""
+            return cast(str, token)
+
+        self.audio_user_slot_token = _id_to_token(self.model_config.audio_user_slot_token_id)
+        self.audio_assistant_slot_token = _id_to_token(self.model_config.audio_assistant_slot_token_id)
+        self.audio_start_token = _id_to_token(self.model_config.audio_start_token_id)
+        self.audio_end_token = _id_to_token(self.model_config.audio_end_token_id)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        trust_remote_code = kwargs.pop("trust_remote_code", True)
+        kwargs.pop("_from_auto", None)
+        codec_path = kwargs.pop("codec_path", None)
+        codec_weight_dtype = kwargs.pop("codec_weight_dtype", "fp32")
+        codec_compute_dtype = kwargs.pop("codec_compute_dtype", None)
+        codec_attention_implementation = kwargs.pop("codec_attention_implementation", None)
+
+        codec_kwargs: Dict[str, Any] = {}
+        if codec_weight_dtype is not None:
+            # Default to fp32 codec weights; callers can pass codec_weight_dtype="bf16" to reduce memory.
+            codec_kwargs["codec_weight_dtype"] = codec_weight_dtype
+
+        model_ref = Path(str(pretrained_model_name_or_path))
+        model_ref_or_name = model_ref if model_ref.exists() else pretrained_model_name_or_path
+        model_config = cast(
+            MossTTSLocalConfig,
+            AutoConfig.from_pretrained(
+                model_ref_or_name,
+                *args,
+                trust_remote_code=trust_remote_code,
+                **kwargs,
+            ),
+        )
+
+        if codec_path is None:
+            try:
+                processor_dict, _ = cls.get_processor_dict(
+                    pretrained_model_name_or_path,
+                    **dict(kwargs),
+                )
+                codec_path = processor_dict.get("audio_tokenizer_name_or_path")
+                audio_tokenizer_dict = processor_dict.get("audio_tokenizer", {})
+                if isinstance(audio_tokenizer_dict, dict):
+                    codec_path = audio_tokenizer_dict.get("audio_tokenizer_name_or_path") or codec_path
+            except Exception:
+                codec_path = None
+        if codec_path is None:
+            codec_path = getattr(model_config, "audio_tokenizer_name_or_path", None)
+        if codec_path is None:
+            codec_path = "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2"
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_ref_or_name,
+            *args,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+        audio_tokenizer = AutoModel.from_pretrained(
+            codec_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+            **codec_kwargs,
+        )
+        if codec_compute_dtype is not None and hasattr(audio_tokenizer, "set_compute_dtype"):
+            audio_tokenizer.set_compute_dtype(codec_compute_dtype)
+        if codec_attention_implementation is not None and hasattr(
+            audio_tokenizer, "set_attention_implementation"
+        ):
+            audio_tokenizer.set_attention_implementation(codec_attention_implementation)
+        return cls(
+            tokenizer=tokenizer,
+            audio_tokenizer=audio_tokenizer,
+            model_config=model_config,
+            **kwargs,
+        )
+
+    @staticmethod
+    def build_user_message(
+        text: Optional[str] = None,
+        reference: Optional[List[Optional[Union[str, os.PathLike, torch.Tensor]]]] = None,
+        instruction: Optional[str] = None,
+        tokens: Optional[int] = None,
+        quality: Optional[str] = None,
+        sound_event: Optional[str] = None,
+        ambient_sound: Optional[str] = None,
+        language: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        if reference is not None and not isinstance(reference, list):
+            reference = [reference]
+        return UserMessage(
+            text=text,
+            reference=reference,
+            instruction=instruction,
+            tokens=tokens,
+            quality=quality,
+            sound_event=sound_event,
+            ambient_sound=ambient_sound,
+            language=language,
+        ).to_dict()
+
+    @staticmethod
+    def build_assistant_message(
+        audio_codes_list: List[Union[str, os.PathLike, torch.Tensor]],
+        content: str = AUDIO_PLACEHOLDER,
+    ) -> Dict[str, Any]:
+        return AssistantMessage(audio_codes_list=audio_codes_list, content=content).to_dict()
+
+    def _assert_fixed_nq(self, n_vq: Optional[int]) -> int:
+        config_nq = int(self.model_config.n_vq)
+        if n_vq is not None and int(n_vq) != config_nq:
+            raise ValueError(
+                "This MOSS-TTS-Local-Transformer-v1.5 release uses the RVQ depth stored in the model config. "
+                f"Expected n_vq={config_nq}, got {int(n_vq)}."
+            )
+        return config_nq
+
+    def _encode_text(self, text: str) -> list[int]:
+        try:
+            return list(self.tokenizer.encode(text, add_special_tokens=False))
+        except TypeError:
+            return list(self.tokenizer.encode(text))
+
+    def _build_text_rows(self, token_ids: Sequence[int], *, device: Optional[torch.device] = None) -> torch.Tensor:
+        rows = torch.full(
+            (len(token_ids), int(self.model_config.n_vq) + 1),
+            int(self.model_config.audio_pad_token_id),
+            dtype=torch.long,
+            device=device,
+        )
+        if token_ids:
+            rows[:, 0] = torch.tensor([int(token_id) for token_id in token_ids], dtype=torch.long, device=rows.device)
+        return rows
+
+    def _build_audio_rows(self, audio_tokens: torch.Tensor, slot_token_id: int) -> torch.Tensor:
+        rows = torch.full(
+            (int(audio_tokens.shape[0]), int(self.model_config.n_vq) + 1),
+            int(self.model_config.audio_pad_token_id),
+            dtype=torch.long,
+            device=audio_tokens.device,
+        )
+        if rows.shape[0] > 0:
+            rows[:, 0] = int(slot_token_id)
+            rows[:, 1:] = audio_tokens.to(dtype=torch.long)
+        return rows
+
+    def _user_prompt_prefix_ids(self) -> list[int]:
+        return (
+            [int(self.model_config.im_start_token_id)]
+            + self._encode_text(USER_ROLE_PREFIX)
+            + self._encode_text(USER_TEMPLATE_REFERENCE_PREFIX)
+        )
+
+    def _user_prompt_after_reference_ids(
+        self,
+        language_code: object | None,
+        prompt_fields: Optional[Dict[str, Any]],
+    ) -> list[int]:
+        return self._encode_text(
+            _render_user_prompt_after_reference(
+                language_code=language_code,
+                prompt_fields=prompt_fields,
+            )
+        )
+
+    def _assistant_prompt_prefix_ids(self) -> list[int]:
+        return (
+            self._encode_text(USER_TEMPLATE_SUFFIX)
+            + [int(self.model_config.im_end_token_id)]
+            + self._encode_text(ASSISTANT_TURN_PREFIX)
+            + [int(self.model_config.im_start_token_id)]
+            + self._encode_text(ASSISTANT_ROLE_PREFIX)
+        )
+
+    def _prompt_fields_from_user_message(self, message: Dict[str, Any]) -> dict[str, Any]:
+        fields = {}
+        for key in ("instruction", "tokens", "quality", "sound_event", "ambient_sound"):
+            if key in message and message.get(key) is not None:
+                fields[key] = message.get(key)
+        if "language" in message and message.get("language") is not None:
+            fields["language"] = message.get("language")
+        return fields
+
+    def _build_generation_or_voice_clone_codes(
+        self,
+        message: Dict[str, Any],
+        n_vq: int,
+    ) -> torch.Tensor:
+        if "text" not in message:
+            raise ValueError("Direct MOSS-TTS-Local-Transformer-v1.5 generation requires messages built by build_user_message(...).")
+        text = "" if message.get("text") is None else str(message.get("text"))
+        prompt_fields = self._prompt_fields_from_user_message(message)
+        language_code = message.get("language")
+        audio_codes_list = self._resolve_audio_items(message.get("audio_codes_list", []), n_vq)
+        text_token_ids = self._encode_text(text)
+
+        if audio_codes_list:
+            parts: list[torch.Tensor] = [self._build_text_rows(
+                self._user_prompt_prefix_ids(),
+                device=audio_codes_list[0].device,
+            )]
+            for reference_codes in audio_codes_list:
+                parts.append(self._build_text_rows([int(self.model_config.audio_start_token_id)], device=reference_codes.device))
+                parts.append(self._build_audio_rows(reference_codes, int(self.model_config.audio_user_slot_token_id)))
+                parts.append(self._build_text_rows([int(self.model_config.audio_end_token_id)], device=reference_codes.device))
+            parts.append(
+                self._build_text_rows(
+                    self._user_prompt_after_reference_ids(language_code, prompt_fields)
+                    + text_token_ids
+                    + self._assistant_prompt_prefix_ids()
+                    + [int(self.model_config.audio_start_token_id)],
+                    device=audio_codes_list[0].device,
+                )
+            )
+            return torch.cat(parts, dim=0)
+
+        prompt_token_ids = (
+            self._user_prompt_prefix_ids()
+            + self._encode_text("None")
+            + self._user_prompt_after_reference_ids(language_code, prompt_fields)
+            + text_token_ids
+            + self._assistant_prompt_prefix_ids()
+            + [int(self.model_config.audio_start_token_id)]
+        )
+        return self._build_text_rows(prompt_token_ids)
+
+    def _build_continuation_codes(
+        self,
+        conversation: list[Dict[str, Any]],
+        n_vq: int,
+    ) -> torch.Tensor:
+        if len(conversation) < 2:
+            raise ValueError("continuation mode requires a user message followed by an assistant audio message.")
+        user_message = conversation[-2]
+        assistant_message = conversation[-1]
+        if user_message.get("role") != "user" or assistant_message.get("role") != "assistant":
+            raise ValueError("continuation mode requires the last two messages to be user, assistant.")
+        if "text" not in user_message:
+            raise ValueError("Direct MOSS-TTS-Local-Transformer-v1.5 continuation requires user messages built by build_user_message(...).")
+
+        text = "" if user_message.get("text") is None else str(user_message.get("text"))
+        prompt_fields = self._prompt_fields_from_user_message(user_message)
+        language_code = user_message.get("language")
+        prompt_token_ids = (
+            self._user_prompt_prefix_ids()
+            + self._encode_text("None")
+            + self._user_prompt_after_reference_ids(language_code, prompt_fields)
+            + self._encode_text(text)
+            + self._assistant_prompt_prefix_ids()
+            + [int(self.model_config.audio_start_token_id)]
+        )
+        audio_codes_list = self._resolve_audio_items(assistant_message.get("audio_codes_list", []), n_vq)
+        if not audio_codes_list:
+            return self._build_text_rows(prompt_token_ids)
+        if len(audio_codes_list) != 1:
+            raise ValueError("The MOSS-TTS-Local-Transformer-v1.5 continuation path expects exactly one prompt audio item.")
+        prompt_audio_codes = audio_codes_list[0]
+        return torch.cat(
+            [
+                self._build_text_rows(prompt_token_ids, device=prompt_audio_codes.device),
+                self._build_audio_rows(prompt_audio_codes, int(self.model_config.audio_assistant_slot_token_id)),
+            ],
+            dim=0,
+        )
+
+    def _try_build_direct_codes(
+        self,
+        conversation: list[Dict[str, Any]],
+        mode: str,
+        n_vq: int,
+    ) -> Optional[torch.Tensor]:
+        if mode == "generation" and len(conversation) == 1 and conversation[-1].get("role") == "user":
+            if "text" in conversation[-1]:
+                return self._build_generation_or_voice_clone_codes(conversation[-1], n_vq)
+            return None
+        if mode == "continuation" and len(conversation) >= 2:
+            if "text" in conversation[-2]:
+                return self._build_continuation_codes(conversation, n_vq)
+            return None
+        return None
+
+    def __call__(self, *args, **kwargs) -> BatchFeature:
+        conversations = args[0] if args else kwargs.pop("conversations")
+        mode: str = kwargs.pop("mode", "generation")
+        apply_chat_template: bool = kwargs.pop("apply_chat_template", True)
+        n_vq = self._assert_fixed_nq(kwargs.pop("n_vq", None))
+
+        kwargs.pop("return_tensors", None)
+        kwargs.pop("padding", None)
+        kwargs.pop("truncation", None)
+
+        if mode not in {"generation", "continuation", "computing_loss"}:
+            raise ValueError(f"Unsupported mode: {mode}")
+        if isinstance(conversations, (Message, dict)):
+            conversations = [conversations]
+        elif isinstance(conversations, list) and conversations and all(
+            isinstance(item, (Message, dict)) for item in conversations
+        ):
+            conversations = [conversations]
+
+        input_ids_list: list[torch.Tensor] = []
+        for conversation in conversations:
+            if isinstance(conversation, (Message, dict)):
+                conversation = [conversation]
+            conversation = [self._normalize_message(message) for message in conversation]
+
+            if (mode == "generation") ^ (conversation[-1]["role"] == "user"):
+                raise ValueError("generation mode must end with a user message.")
+            if mode == "continuation" and conversation[-1]["role"] != "assistant":
+                raise ValueError("continuation mode must end with an assistant message.")
+
+            direct_codes = self._try_build_direct_codes(conversation, mode, n_vq)
+            if direct_codes is not None:
+                input_ids_list.append(direct_codes)
+                continue
+
+            unified_parts = []
+            for message_idx, message in enumerate(conversation):
+                content = str(message["content"])
+                if apply_chat_template:
+                    add_generation_prompt = mode == "generation" and message_idx == len(conversation) - 1
+                    try:
+                        content = self.tokenizer.apply_chat_template(
+                            [{"role": message["role"], "content": content}],
+                            add_generation_prompt=add_generation_prompt,
+                            tokenize=False,
+                        )
+                    except Exception:
+                        logger.warning("apply_chat_template failed; falling back to raw message content.")
+
+                raw_audio_items = message.get("audio_codes_list", [])
+                audio_codes_list = self._resolve_audio_items(raw_audio_items, n_vq)
+                unified_parts.append(
+                    self._get_unified_codes(
+                        role=message["role"],
+                        content=content,
+                        audio_codes_list=audio_codes_list,
+                        truncation=(mode == "continuation"),
+                    )
+                )
+
+            unified_codes = torch.cat(unified_parts, dim=0)
+            if mode == "generation":
+                audio_start_row = torch.full(
+                    (1, n_vq + 1),
+                    int(self.model_config.audio_pad_token_id),
+                    dtype=unified_codes.dtype,
+                    device=unified_codes.device,
+                )
+                audio_start_row[:, 0] = int(self.model_config.audio_start_token_id)
+                unified_codes = torch.cat([unified_codes, audio_start_row], dim=0)
+            input_ids_list.append(unified_codes)
+
+        return BatchFeature(data=self._pad(input_ids_list))
+
+    def _normalize_message(self, message: Union[Message, Dict[str, Any]]) -> Dict[str, Any]:
+        if isinstance(message, Message):
+            return message.to_dict()
+        if not isinstance(message, dict):
+            raise TypeError("Each message must be a Message or dict.")
+        if "content" in message and "audio_codes_list" in message:
+            return message
+        role = message.get("role")
+        if role == "user":
+            return self.build_user_message(**{key: message.get(key) for key in USER_MESSAGE_FIELDS})
+        if role == "assistant":
+            return self.build_assistant_message(
+                audio_codes_list=message.get("audio_codes_list", []),
+                content=message.get("content", AUDIO_PLACEHOLDER),
+            )
+        raise ValueError(f"Unsupported role: {role}")
+
+    def _resolve_audio_items(
+        self,
+        raw_audio_items: list[Any],
+        n_vq: int,
+    ) -> list[torch.Tensor]:
+        if not raw_audio_items:
+            return []
+        resolved: list[Optional[torch.Tensor]] = [None] * len(raw_audio_items)
+        paths: list[str] = []
+        path_positions: list[int] = []
+        for index, item in enumerate(raw_audio_items):
+            if isinstance(item, torch.Tensor):
+                if item.ndim != 2 or int(item.shape[1]) != n_vq:
+                    raise ValueError(f"audio code tensor must have shape [T, {n_vq}], got {tuple(item.shape)}.")
+                resolved[index] = item.to(dtype=torch.long).cpu()
+            elif isinstance(item, (str, os.PathLike)):
+                paths.append(str(item))
+                path_positions.append(index)
+            else:
+                raise TypeError("Audio items must be tensors or path-like values.")
+        if paths:
+            encoded = self.encode_audios_from_path(paths, n_vq=n_vq)
+            for position, codes in zip(path_positions, encoded):
+                resolved[position] = codes
+        return [cast(torch.Tensor, item) for item in resolved]
+
+    def _pad(self, input_ids_list: list[torch.Tensor]) -> Dict[str, torch.Tensor]:
+        device = input_ids_list[0].device
+        lengths = torch.tensor([item.shape[0] for item in input_ids_list], device=device)
+        padded = torch.nn.utils.rnn.pad_sequence(
+            input_ids_list,
+            batch_first=True,
+            padding_value=int(self.model_config.audio_pad_token_id),
+            padding_side="left",
+        )
+        left_pad_mask = (padded.shape[1] - lengths).unsqueeze(1) > torch.arange(
+            padded.shape[1],
+            device=device,
+        ).unsqueeze(0)
+        padded[..., 0][left_pad_mask] = int(self.model_config.pad_token_id)
+        attention_mask = torch.zeros(padded.shape[:2], dtype=torch.bool, device=device)
+        attention_mask[~left_pad_mask] = True
+        return {"input_ids": padded, "attention_mask": attention_mask}
+
+    @staticmethod
+    def _replace_audio_placeholders(
+        content: str,
+        lengths: list[int],
+        slot_token: str,
+        audio_start_token: str,
+        audio_end_token: str,
+    ) -> str:
+        placeholder_count = content.count(AUDIO_PLACEHOLDER)
+        if placeholder_count != len(lengths):
+            raise ValueError(
+                f"Number of {AUDIO_PLACEHOLDER} ({placeholder_count}) does not match "
+                f"audio item count ({len(lengths)})."
+            )
+        lengths_iter = iter(lengths)
+
+        def replacer(_: re.Match) -> str:
+            length = int(next(lengths_iter))
+            if length <= 0:
+                return f"{audio_start_token}{audio_end_token}"
+            return f"{audio_start_token}{slot_token * length}{audio_end_token}"
+
+        return re.sub(re.escape(AUDIO_PLACEHOLDER), replacer, content)
+
+    def _get_unified_codes(
+        self,
+        role: str,
+        content: str,
+        audio_codes_list: list[torch.Tensor],
+        truncation: bool,
+    ) -> torch.Tensor:
+        n_vq = int(self.model_config.n_vq)
+        slot_token = self.audio_user_slot_token if role == "user" else self.audio_assistant_slot_token
+        content = self._replace_audio_placeholders(
+            content=content,
+            lengths=[int(codes.shape[0]) for codes in audio_codes_list],
+            slot_token=slot_token,
+            audio_start_token=self.audio_start_token,
+            audio_end_token=self.audio_end_token,
+        )
+        text_codes = torch.tensor(
+            self.tokenizer.encode(content),
+            dtype=torch.long,
+            device=audio_codes_list[0].device if audio_codes_list else None,
+        )
+
+        audio_start_indices = torch.where(text_codes == int(self.model_config.audio_start_token_id))[0]
+        audio_end_indices = torch.where(text_codes == int(self.model_config.audio_end_token_id))[0]
+        if len(audio_start_indices) != len(audio_codes_list) or len(audio_end_indices) != len(audio_codes_list):
+            raise ValueError("Audio placeholders do not match the encoded audio spans.")
+
+        if not audio_codes_list:
+            audio_codes = torch.full(
+                (len(text_codes), n_vq),
+                int(self.model_config.audio_pad_token_id),
+                dtype=torch.long,
+                device=text_codes.device,
+            )
+        else:
+            pieces: list[torch.Tensor] = []
+            prefix_idx = 0
+            for start_t, end_t, codes in zip(audio_start_indices, audio_end_indices, audio_codes_list):
+                start_idx = int(start_t.item())
+                end_idx = int(end_t.item())
+                pad_before = torch.full(
+                    (start_idx - prefix_idx + 1, n_vq),
+                    int(self.model_config.audio_pad_token_id),
+                    dtype=torch.long,
+                    device=codes.device,
+                )
+                pieces.extend([pad_before, codes.to(dtype=torch.long)])
+                prefix_idx = end_idx
+            if truncation:
+                trailing = torch.zeros(
+                    (0, n_vq),
+                    dtype=torch.long,
+                    device=audio_codes_list[0].device,
+                )
+            else:
+                last_end = int(audio_end_indices[-1].item())
+                trailing = torch.full(
+                    (len(text_codes) - last_end, n_vq),
+                    int(self.model_config.audio_pad_token_id),
+                    dtype=torch.long,
+                    device=audio_codes_list[0].device,
+                )
+            pieces.append(trailing)
+            audio_codes = torch.cat(pieces, dim=0)
+
+        if text_codes.shape[0] != audio_codes.shape[0]:
+            min_len = min(text_codes.shape[0], audio_codes.shape[0])
+            text_codes = text_codes[:min_len]
+            audio_codes = audio_codes[:min_len]
+        return torch.cat([text_codes.unsqueeze(1), audio_codes], dim=1)
+
+    def _parse_text_codes(self, start_length: int, text_codes: torch.LongTensor) -> str:
+        text = cast(str, self.tokenizer.decode(text_codes))
+        prefix = cast(str, self.tokenizer.decode(text_codes[:start_length]))
+        text = text[len(prefix):]
+        audio_pattern = re.compile(
+            rf"(?:{re.escape(self.audio_start_token)})?"
+            rf"(?:{re.escape(self.audio_assistant_slot_token)})*"
+            rf"{re.escape(self.audio_end_token)}"
+        )
+        return audio_pattern.sub(
+            lambda match: AUDIO_PLACEHOLDER if self.audio_assistant_slot_token in match.group(0) else "",
+            text,
+        )
+
+    def _parse_audio_codes(
+        self,
+        start_length: int,
+        audio_codes: torch.LongTensor,
+        *,
+        return_stereo: bool = True,
+    ) -> list[torch.Tensor]:
+        is_pad = audio_codes.eq(int(self.model_config.audio_pad_token_id)).all(dim=1)
+        non_pad = ~is_pad
+        if not bool(non_pad.any().item()):
+            return []
+        idx = torch.nonzero(non_pad).squeeze(1)
+        breaks = torch.where(idx[1:] != idx[:-1] + 1)[0] + 1
+        segment_indices = [idx] if breaks.numel() == 0 else list(torch.tensor_split(idx, breaks.cpu().tolist()))
+        code_segments = [audio_codes[segment] for segment in segment_indices]
+        decoded = self.decode_audio_codes(code_segments, return_stereo=return_stereo)
+
+        if start_length > 0 and code_segments and decoded:
+            first_code_length = int(code_segments[0].shape[0])
+            if first_code_length > 0:
+                trim_ratio = max(0.0, min(float(start_length) / float(first_code_length), 1.0))
+                if trim_ratio >= 1.0:
+                    decoded = decoded[1:]
+                elif trim_ratio > 0.0:
+                    trim_samples = int(decoded[0].shape[-1] * trim_ratio)
+                    decoded[0] = decoded[0][..., trim_samples:]
+        return decoded
+
+    def decode(self, output: Any, *, return_stereo: bool = True) -> list[Optional[AssistantMessage]]:
+        generated_messages: list[Optional[AssistantMessage]] = []
+        for start_length, generation_ids in output:
+            content = self._parse_text_codes(int(start_length), generation_ids[:, 0])
+            audio_codes_list = self._parse_audio_codes(
+                int(start_length),
+                generation_ids[:, 1:],
+                return_stereo=return_stereo,
+            )
+            if content == "":
+                generated_messages.append(None)
+            else:
+                generated_messages.append(
+                    AssistantMessage(
+                        content=content,
+                        audio_codes_list=cast(list[Union[str, torch.Tensor]], audio_codes_list),
+                    )
+                )
+        return generated_messages
+
+    @staticmethod
+    def loudness_normalize(
+        wav: torch.Tensor,
+        target_dbfs: float = -20.0,
+        gain_range: tuple[float, float] = (-3.0, 3.0),
+    ) -> torch.Tensor:
+        wav = wav.to(torch.float32)
+        if wav.numel() == 0:
+            return wav
+        current_dbfs = 10.0 * torch.log10(torch.mean(wav**2) + 1e-9)
+        gain = max(gain_range[0], min(float(target_dbfs - current_dbfs), gain_range[1]))
+        return wav * (10.0 ** (gain / 20.0))
+
+    def _get_audio_tokenizer_device(self) -> torch.device:
+        audio_tokenizer = getattr(self, "audio_tokenizer", None)
+        if audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set.")
+        try:
+            return next(audio_tokenizer.parameters()).device
+        except StopIteration:
+            return torch.device("cpu")
+
+    def encode_audios_from_wav(
+        self,
+        wav_list: Union[torch.Tensor, list[torch.Tensor]],
+        sampling_rate: int,
+        n_vq: Optional[int] = None,
+    ) -> list[torch.Tensor]:
+        n_vq = self._assert_fixed_nq(n_vq)
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set.")
+        if isinstance(wav_list, torch.Tensor):
+            wav_list = [wav_list]
+        target_sr = int(self.model_config.sampling_rate)
+        device = self._get_audio_tokenizer_device()
+        prepared = []
+        for wav in wav_list:
+            if wav.ndim == 1:
+                wav = wav.unsqueeze(0)
+            if wav.shape[0] == 1:
+                wav = wav.repeat(2, 1)
+            elif wav.shape[0] > 2:
+                wav = wav[:2]
+            if int(sampling_rate) != target_sr:
+                wav = torchaudio.functional.resample(wav, int(sampling_rate), target_sr)
+            prepared.append(self.loudness_normalize(wav).to(device))
+
+        if hasattr(self.audio_tokenizer, "batch_encode"):
+            encoded = self.audio_tokenizer.batch_encode(prepared, num_quantizers=n_vq)
+            audio_codes = encoded.audio_codes
+            audio_lengths = encoded.audio_codes_lengths
+        else:
+            max_len = max(int(wav.shape[-1]) for wav in prepared)
+            input_values = torch.zeros(len(prepared), 1, max_len, dtype=torch.float32, device=device)
+            padding_mask = torch.zeros(len(prepared), max_len, dtype=torch.bool, device=device)
+            for index, wav in enumerate(prepared):
+                input_values[index, 0, : wav.shape[-1]] = wav
+                padding_mask[index, : wav.shape[-1]] = True
+            encoded = self.audio_tokenizer.encode(
+                input_values,
+                padding_mask=padding_mask,
+                num_quantizers=n_vq,
+                return_dict=True,
+            )
+            audio_codes = encoded.audio_codes
+            audio_lengths = encoded.audio_codes_lengths
+
+        if audio_codes is None or audio_lengths is None:
+            raise RuntimeError("audio_tokenizer did not return audio_codes/audio_codes_lengths.")
+        result = []
+        for index in range(int(audio_codes.shape[1])):
+            length = int(audio_lengths[index].item())
+            result.append(audio_codes[:, index, :length].transpose(0, 1).contiguous().cpu().long())
+        return result
+
+    def encode_audios_from_path(
+        self,
+        wav_path_list: Union[str, os.PathLike, list[Union[str, os.PathLike]]],
+        n_vq: Optional[int] = None,
+    ) -> list[torch.Tensor]:
+        if isinstance(wav_path_list, (str, os.PathLike)):
+            wav_path_list = [wav_path_list]
+        wavs = []
+        target_sr = int(self.model_config.sampling_rate)
+        for wav_path in wav_path_list:
+            wav, sr = torchaudio.load(str(wav_path))
+            if int(sr) != target_sr:
+                wav = torchaudio.functional.resample(wav, int(sr), target_sr)
+            wavs.append(wav)
+        return self.encode_audios_from_wav(wavs, target_sr, n_vq=n_vq)
+
+    def decode_audio_codes(
+        self,
+        audio_tokens_list: Union[torch.Tensor, list[torch.Tensor]],
+        *,
+        return_stereo: bool = True,
+    ) -> list[torch.Tensor]:
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set.")
+        if isinstance(audio_tokens_list, torch.Tensor):
+            audio_tokens_list = [audio_tokens_list]
+        if not audio_tokens_list:
+            return []
+
+        n_vq = int(self.model_config.n_vq)
+        device = self._get_audio_tokenizer_device()
+        codes_list = [
+            codes[:, :n_vq].transpose(0, 1).contiguous().to(device=device, dtype=torch.long)
+            for codes in audio_tokens_list
+        ]
+        max_len = max(int(codes.shape[1]) for codes in codes_list)
+        audio_codes = torch.zeros(n_vq, len(codes_list), max_len, device=device, dtype=torch.long)
+        padding_mask = torch.zeros(len(codes_list), max_len, device=device, dtype=torch.bool)
+        for index, codes in enumerate(codes_list):
+            length = int(codes.shape[1])
+            audio_codes[:, index, :length] = codes
+            padding_mask[index, :length] = True
+
+        decoded = self.audio_tokenizer.decode(
+            audio_codes,
+            padding_mask=padding_mask,
+            num_quantizers=n_vq,
+            return_dict=True,
+            chunk_duration=8,
+        )
+        audio = decoded.audio
+        audio_lengths = decoded.audio_lengths
+        if audio is None or audio_lengths is None:
+            raise RuntimeError("audio_tokenizer.decode did not return audio/audio_lengths.")
+        wavs = []
+        for index in range(int(audio.shape[0])):
+            length = int(audio_lengths[index].item())
+            wav = audio[index, :, :length].contiguous().cpu().to(torch.float32)
+            if not return_stereo:
+                if wav.shape[0] == 1:
+                    wav = wav.squeeze(0)
+                else:
+                    wav = wav.mean(dim=0)
+            wavs.append(wav)
+        return wavs

code/moss_export/qwen3_decoder.py

@@ -0,0 +1,582 @@
+# coding=utf-8
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from safetensors.torch import load_file
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutputWithPast
+
+from .gpt2_decoder import PackedSequenceMetadata, MossTTSNanoGPT2Model
+
+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
+
+
+class MossQwen3RMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class MossQwen3RotaryEmbedding(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        head_dim = int(getattr(config, "head_dim", config.hidden_size // config.num_attention_heads))
+        rope_theta = getattr(config, "rope_theta", None)
+        if rope_theta is None:
+            rope_scaling = getattr(config, "rope_scaling", None)
+            if isinstance(rope_scaling, dict):
+                rope_theta = rope_scaling.get("rope_theta")
+        if rope_theta is None:
+            rope_theta = 1000000.0
+        rope_theta = float(rope_theta)
+        self.head_dim = head_dim
+        self.rope_theta = rope_theta
+        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)
+
+    def _compute_inv_freq(self, device: Optional[torch.device] = None) -> torch.Tensor:
+        return 1.0 / (
+            self.rope_theta ** (torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) / self.head_dim)
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.LongTensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        inv_freq = self._compute_inv_freq(device=hidden_states.device)
+        freqs = torch.einsum(
+            "bs,d->bsd",
+            position_ids.to(device=hidden_states.device, dtype=inv_freq.dtype),
+            inv_freq,
+        )
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb.cos().to(dtype=hidden_states.dtype), emb.sin().to(dtype=hidden_states.dtype)
+
+
+def rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
+    first_half = hidden_states[..., : hidden_states.shape[-1] // 2]
+    second_half = hidden_states[..., hidden_states.shape[-1] // 2 :]
+    return torch.cat((-second_half, first_half), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    cos = cos.unsqueeze(-2)
+    sin = sin.unsqueeze(-2)
+    query = (query * cos) + (rotate_half(query) * sin)
+    key = (key * cos) + (rotate_half(key) * sin)
+    return query, key
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    if n_rep == 1:
+        return hidden_states
+    batch, seq_len, num_key_value_heads, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:, :, :, None, :].expand(batch, seq_len, num_key_value_heads, n_rep, head_dim)
+    return hidden_states.reshape(batch, seq_len, num_key_value_heads * n_rep, head_dim)
+
+
+class MossQwen3MLP(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.act_fn = ACT2FN[getattr(config, "hidden_act", "silu")]
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+
+
+class MossQwen3Attention(nn.Module):
+    def __init__(self, config, layer_idx: int) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = int(layer_idx)
+        self.hidden_size = int(config.hidden_size)
+        self.num_heads = int(config.num_attention_heads)
+        self.num_key_value_heads = int(config.num_key_value_heads)
+        self.head_dim = int(getattr(config, "head_dim", self.hidden_size // self.num_heads))
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.scaling = self.head_dim ** -0.5
+        self.attention_dropout = float(getattr(config, "attention_dropout", 0.0))
+        self.attn_implementation = str(getattr(config, "_attn_implementation", "eager"))
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=bool(config.attention_bias))
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=bool(config.attention_bias),
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=bool(config.attention_bias),
+        )
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=bool(config.attention_bias))
+        self.q_norm = MossQwen3RMSNorm(self.head_dim, eps=float(config.rms_norm_eps))
+        self.k_norm = MossQwen3RMSNorm(self.head_dim, eps=float(config.rms_norm_eps))
+
+    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:
+        key = repeat_kv(key, self.num_key_value_groups)
+        value = repeat_kv(value, self.num_key_value_groups)
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        scores = torch.matmul(query, key.transpose(-1, -2)) * self.scaling
+        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(~mask, torch.finfo(scores.dtype).min)
+        probs = torch.softmax(scores, dim=-1)
+        if self.training and self.attention_dropout > 0:
+            probs = torch.dropout(probs, self.attention_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:
+        key = repeat_kv(key, self.num_key_value_groups)
+        value = repeat_kv(value, self.num_key_value_groups)
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        mask = None
+        if attention_mask is not None or query.shape[-2] != key.shape[-2]:
+            mask = self._causal_attention_mask(
+                attention_mask=attention_mask,
+                query_length=query.shape[-2],
+                key_length=key.shape[-2],
+                device=query.device,
+            )
+        output = torch.nn.functional.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=mask is None,
+            scale=self.scaling,
+        )
+        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, got dtype={query.dtype}.")
+
+        dropout_p = self.attention_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_key_value_heads, self.head_dim).index_select(0, packed_metadata.indices)
+                value = value.reshape(-1, self.num_key_value_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)
+
+        if query.shape[1] != key.shape[1]:
+            query_attention_mask = attention_mask[:, -query.shape[1] :]
+            unpadded_query, query_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
+                query,
+                query_attention_mask,
+            )
+            unpadded_key, _, cu_seqlens_k, max_seqlen_k, _ = 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_q,
+                cu_seqlens_k,
+                max_seqlen_q,
+                max_seqlen_k,
+                dropout_p=dropout_p,
+                causal=True,
+            )
+            return pad_input(output, query_indices, query.shape[0], query.shape[1])
+
+        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,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = 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]]]:
+        input_shape = hidden_states.shape[:-1]
+        query_states = self.q_norm(
+            self.q_proj(hidden_states).view(*input_shape, self.num_heads, self.head_dim)
+        )
+        key_states = self.k_norm(
+            self.k_proj(hidden_states).view(*input_shape, self.num_key_value_heads, self.head_dim)
+        )
+        value_states = self.v_proj(hidden_states).view(*input_shape, self.num_key_value_heads, self.head_dim)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if layer_past is not None:
+            past_key, past_value = layer_past
+            key_states = torch.cat([past_key.to(device=key_states.device, dtype=key_states.dtype), key_states], dim=1)
+            value_states = torch.cat(
+                [past_value.to(device=value_states.device, dtype=value_states.dtype), value_states],
+                dim=1,
+            )
+
+        present = (key_states, value_states) if use_cache else None
+        if self.attn_implementation == "flash_attention_2":
+            attn_output = self._flash_attention(
+                query=query_states,
+                key=key_states,
+                value=value_states,
+                attention_mask=attention_mask,
+                packed_metadata=packed_metadata,
+            )
+        elif self.attn_implementation == "sdpa":
+            attn_output = self._sdpa_attention(
+                query=query_states,
+                key=key_states,
+                value=value_states,
+                attention_mask=attention_mask,
+            )
+        else:
+            attn_output = self._eager_attention(
+                query=query_states,
+                key=key_states,
+                value=value_states,
+                attention_mask=attention_mask,
+            )
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        return self.o_proj(attn_output), present
+
+
+class MossQwen3DecoderLayer(nn.Module):
+    def __init__(self, config, layer_idx: int) -> None:
+        super().__init__()
+        self.self_attn = MossQwen3Attention(config=config, layer_idx=layer_idx)
+        self.mlp = MossQwen3MLP(config)
+        self.input_layernorm = MossQwen3RMSNorm(config.hidden_size, eps=float(config.rms_norm_eps))
+        self.post_attention_layernorm = MossQwen3RMSNorm(config.hidden_size, eps=float(config.rms_norm_eps))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        packed_metadata: Optional[PackedSequenceMetadata] = None,
+        layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        attn_output, present = self.self_attn(
+            hidden_states=hidden_states,
+            position_embeddings=position_embeddings,
+            attention_mask=attention_mask,
+            packed_metadata=packed_metadata,
+            layer_past=layer_past,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + attn_output
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + self.mlp(hidden_states)
+        return hidden_states, present
+
+
+class MossQwen3Model(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.config = config
+        self.attn_implementation = str(getattr(config, "_attn_implementation", "eager"))
+        self.padding_idx = getattr(config, "pad_token_id", None)
+        self.vocab_size = int(config.vocab_size)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [MossQwen3DecoderLayer(config, layer_idx=index) for index in range(config.num_hidden_layers)]
+        )
+        self.norm = MossQwen3RMSNorm(config.hidden_size, eps=float(config.rms_norm_eps))
+        self.rotary_emb = MossQwen3RotaryEmbedding(config)
+        self.gradient_checkpointing = False
+        self.gradient_checkpointing_use_reentrant = bool(
+            getattr(config, "gradient_checkpointing_use_reentrant", False)
+        )
+        self._reset_parameters()
+
+    def _reset_parameters(self) -> None:
+        init_std = float(getattr(self.config, "initializer_range", 0.02))
+        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)
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def load_qwen3_pretrained_weights(self, pretrained_path: str) -> None:
+        model_dir = Path(pretrained_path)
+        index_path = model_dir / "model.safetensors.index.json"
+        if not index_path.exists():
+            raise FileNotFoundError(f"Missing Qwen3 safetensors index: {index_path}")
+        with index_path.open("r", encoding="utf-8") as handle:
+            index = json.load(handle)
+        weight_map = index.get("weight_map", {})
+        shard_to_keys: dict[str, list[str]] = {}
+        for key, shard in weight_map.items():
+            if not key.startswith("model."):
+                continue
+            shard_to_keys.setdefault(str(shard), []).append(key)
+
+        state_dict = self.state_dict()
+        loaded_state = {}
+        for shard, keys in sorted(shard_to_keys.items()):
+            shard_tensors = load_file(str(model_dir / shard), device="cpu")
+            for key in keys:
+                target_key = key[len("model.") :]
+                if target_key not in state_dict:
+                    continue
+                tensor = shard_tensors[key]
+                if tuple(tensor.shape) != tuple(state_dict[target_key].shape):
+                    raise ValueError(
+                        f"Shape mismatch while loading Qwen3 weight {key}: "
+                        f"checkpoint={tuple(tensor.shape)} model={tuple(state_dict[target_key].shape)}"
+                    )
+                loaded_state[target_key] = tensor
+
+        missing, unexpected = self.load_state_dict(loaded_state, strict=False)
+        unexpected = [key for key in unexpected if key]
+        if unexpected:
+            raise RuntimeError(f"Unexpected Qwen3 pretrained keys after load: {unexpected[:10]}")
+        missing = [key for key in missing if key not in loaded_state]
+        if missing:
+            raise RuntimeError(f"Missing Qwen3 pretrained keys after load: {missing[:10]}")
+
+    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 = MossTTSNanoGPT2Model.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,
+                    sequence_length=hidden_states.shape[1],
+                )
+            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 = MossTTSNanoGPT2Model.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,
+            )
+
+        hidden_states = hidden_states * query_attention_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        all_hidden_states = () if output_hidden_states else None
+        presents = [] if use_cache else None
+        for layer_index, decoder_layer in enumerate(self.layers):
+            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, _ = decoder_layer(
+                        hidden_states=inputs[0],
+                        attention_mask=inputs[1],
+                        packed_metadata=packed_metadata,
+                        layer_past=None,
+                        use_cache=False,
+                        position_embeddings=position_embeddings,
+                    )
+                    return output
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    custom_forward,
+                    hidden_states,
+                    attention_mask,
+                    use_reentrant=self.gradient_checkpointing_use_reentrant,
+                )
+                present = None
+            else:
+                hidden_states, present = decoder_layer(
+                    hidden_states=hidden_states,
+                    attention_mask=attention_mask,
+                    packed_metadata=packed_metadata,
+                    layer_past=None if past_key_values is None else past_key_values[layer_index],
+                    use_cache=use_cache,
+                    position_embeddings=position_embeddings,
+                )
+
+            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.norm(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,
+        )

code/moss_pack.py

@@ -0,0 +1,95 @@
+"""Packing of (instruction, text, tokens, reference codes, target codes) into MOSS-TTS tensors.
+
+Derived from $NB/vblora/code/vbpack.py (which itself replaced the private MOSS-TTS repo's
+`MossTTSLocalV15SFTDataset._pack_record`, that repo not being on this filesystem).  Extended here
+with **reference-audio conditioning**: the processor accepts reference audio as a LongTensor
+[T, n_vq] of codec codes directly (see `_resolve_audio_items` in processing_moss_tts.py), so no
+audio tokenizer is needed at train time.
+
+Label derivation is pinned to the loss consumer (`va_loss.compute_supervised_loss_from_hidden`):
+  * labels[t] = input_ids[t+1]     (hidden at t decodes frame t+1)
+  * supervised iff the TARGET text channel is audio_assistant_slot (continue) or audio_end (stop)
+  * audio channels >= audio_pad_code -> -100
+
+Crucially the reference audio rows sit inside the USER turn, whose target text channel is neither
+assistant_slot nor audio_end, so they are automatically unsupervised.  Verified in verify_pack.py.
+"""
+import numpy as np
+import torch
+
+
+def codes_from_blob(blob, frames, n_vq=12):
+    """target_codes/ref_codes/chosen_codes/... are int16 little-endian [frames, n_vq]."""
+    if blob is None or not frames:
+        return None
+    a = np.frombuffer(blob, dtype=np.int16)
+    n = int(frames) * int(n_vq)
+    if a.size < n:
+        return None
+    return torch.from_numpy(a[:n].reshape(int(frames), int(n_vq)).astype(np.int64))
+
+
+class Packer:
+    def __init__(self, processor, model_config):
+        self.proc = processor
+        c = model_config
+        self.n_vq = int(c.n_vq)
+        self.assist = int(c.audio_assistant_slot_token_id)
+        self.aend = int(c.audio_end_token_id)
+        self.audio_pad = int(getattr(c, "audio_pad_token_id", 1024))
+        self.text_pad = int(c.pad_token_id)
+
+    def pack(self, instruction, text, tokens, codes, ref_codes=None, language="English"):
+        """codes: [T, n_vq] long.  ref_codes: [R, n_vq] long or None.
+        -> dict(input_ids [T,13], labels [T,13], n_sup)"""
+        codes = torch.as_tensor(np.asarray(codes), dtype=torch.long)
+        ref = None
+        if ref_codes is not None:
+            ref = [torch.as_tensor(np.asarray(ref_codes), dtype=torch.long)]
+        um = self.proc.build_user_message(text=text, instruction=instruction, language=language,
+                                          tokens=tokens, reference=ref)
+        am = self.proc.build_assistant_message(audio_codes_list=[codes])
+        b = self.proc([[um, am]], mode="computing_loss")
+        ids = b["input_ids"][0].long()                     # [T, 13]
+        lab = torch.full_like(ids, -100)
+        lab[:-1] = ids[1:]
+        tt = lab[:, 0]
+        sup = (tt == self.assist) | (tt == self.aend)
+        # --- reference-audio contamination guard ---------------------------------------
+        # The reference clip inside the USER turn is terminated by <|audio_end|>, which is the
+        # SAME token id the model uses for the assistant's STOP decision.  Without this guard the
+        # position preceding the reference's terminator is supervised as "stop", teaching the
+        # model to stop at the end of the user's reference audio.  Measured: exactly 1 spurious
+        # supervised position per referenced sample (n_sup 154 vs 153 at target_frames=152).
+        # Supervise only from one position before the first assistant audio row.
+        fa = (ids[:, 0] == self.assist).nonzero()
+        if fa.numel() > 0:
+            k = int(fa[0])
+            sup[: max(0, k - 1)] = False
+        lab[~sup] = -100
+        au = lab[:, 1:]
+        au[(au >= self.audio_pad) | (au < 0)] = -100
+        lab[:, 1:] = au
+        return {"input_ids": ids, "labels": lab, "n_sup": int(sup.sum())}
+
+
+def collate(items, text_pad, audio_pad):
+    """Right-pad a list of pack() outputs.  Padding is fully masked in `labels`."""
+    T = max(it["input_ids"].shape[0] for it in items)
+    C = items[0]["input_ids"].shape[1]
+    B = len(items)
+    ids = torch.empty((B, T, C), dtype=torch.long)
+    ids[:, :, 0] = text_pad
+    ids[:, :, 1:] = audio_pad
+    lab = torch.full((B, T, C), -100, dtype=torch.long)
+    am = torch.zeros((B, T), dtype=torch.long)
+    for i, it in enumerate(items):
+        t = it["input_ids"].shape[0]
+        ids[i, :t] = it["input_ids"]
+        lab[i, :t] = it["labels"]
+        am[i, :t] = 1
+    out = {"input_ids": ids, "labels": lab, "attention_mask": am}
+    for k in ("meta",):
+        if k in items[0]:
+            out[k] = [it[k] for it in items]
+    return out

code/moss_small.py

@@ -0,0 +1,220 @@
+#!/usr/bin/env python3
+"""M1 (Arm P) und M2 (Arm S): Qwen3-0.6B-Backbone + MOSS-Talker-Head.
+
+M1: Head (lokaler Transformer, 12 Audio-Embeddings = 12 Audio-Köpfe, binärer Stop-Kopf) aus $SC/out/sft3/export,
+    Hidden 2560, EINGEFROREN; trainierbar: Backbone (1024) + proj_out 1024->2560 (vor dem Head) + proj_in 2560->1024
+    (Summe der Audio-Embeddings am Backbone-Eingang).  Der Head wird NIE in ein Backbone gemergt;
+    audio_lm_heads.N.weight IST audio_embeddings.N.weight (gewichtsgeteilt).
+M2: gleiche Architektur, Head mit Hidden 1024 zufällig initialisiert, alles trainierbar.
+
+`hidden_size != local_hidden_size` ist im Original nicht vorgesehen (Config wirft, _global_hidden_to_local ist Identität,
+Audio-Embeddings haben hidden_size) -> Unterklasse mit eigenem __init__ und Projektionen."""
+import os, sys, json, importlib.util
+import torch, torch.nn as nn
+from transformers import Qwen3Config
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
+from safetensors import safe_open
+from safetensors.torch import load_file
+
+SC = '/e/scratch/reformo/schuhmann1_moss'
+SFT3 = f'{SC}/out/sft3/export'
+QWEN = f'{SC}/models/qwen3-0.6b'
+
+
+def _load_export_modules():
+    """modeling_moss_tts aus dem SFT3-Export als Paket importieren (relative Imports)."""
+    import importlib
+    if 'moss_export' in sys.modules:
+        return sys.modules['moss_export']
+    spec = importlib.util.spec_from_file_location('moss_export', f'{SFT3}/__init__.py',
+                                                  submodule_search_locations=[SFT3])
+    mod = importlib.util.module_from_spec(spec)
+    sys.modules['moss_export'] = mod
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def export_classes():
+    _load_export_modules()
+    import importlib
+    cfgm = importlib.import_module('moss_export.configuration_moss_tts')
+    modm = importlib.import_module('moss_export.modeling_moss_tts')
+    procm = importlib.import_module('moss_export.processing_moss_tts')
+    return cfgm.MossTTSLocalConfig, modm.MossTTSLocalModel, procm.MossTTSLocalProcessor
+
+
+def make_config(kind):
+    """kind: 'M1' (local 2560, SFT3-Head) | 'M2' (local 1024, von Null)."""
+    MossCfg, _, _ = export_classes()
+    base = json.load(open(f'{SFT3}/config.json'))
+    q = json.load(open(f'{QWEN}/config.json'))
+    for k in ('architectures', 'transformers_version'):
+        q.pop(k, None)
+    q['use_cache'] = False
+    q['gradient_checkpointing_use_reentrant'] = False
+    g = dict(base['gpt2_config'])
+    if kind == 'M1':
+        local = 2560                         # n_embd 2560, n_head 32, n_inner 9728 bleiben (SFT3-Head)
+    elif kind == 'M2':
+        local = 1024
+        g['n_embd'] = 1024; g['n_head'] = 16; g['n_inner'] = 4096
+    else:
+        raise ValueError(kind)
+    g_tmp = dict(g); g_tmp['n_embd'] = int(q['hidden_size'])    # Config-Prüfung umgehen
+    kw = {k: base[k] for k in ('audio_codebook_sizes', 'audio_pad_code', 'audio_pad_token_id', 'audio_start_token_id',
+                                'audio_end_token_id', 'audio_user_slot_token_id', 'audio_assistant_slot_token_id',
+                                'audio_vocab_size', 'im_end_token_id', 'im_start_token_id', 'n_vq', 'pad_token_id',
+                                'local_text_head_mode', 'local_transformer_layers', 'sampling_rate',
+                                'audio_tokenizer_name_or_path', 'use_static_local_kv_cache')}
+    cfg = MossCfg(qwen3_config=q, gpt2_config=g_tmp, attn_implementation='sdpa',
+                  local_transformer_attn_implementation='sdpa', **kw)
+    cfg.gpt2_config = GPT2Config(**g)
+    cfg.local_hidden_size = local
+    cfg.small_kind = kind
+    cfg.dtype = 'bfloat16'
+    return cfg
+
+
+class MossTTSSmallModel:
+    """Fabrik: erzeugt die Unterklasse erst, wenn das Export-Modul geladen ist."""
+    _cls = None
+
+    @classmethod
+    def cls(klass):
+        if klass._cls is not None:
+            return klass._cls
+        _, MossTTSLocalModel, _ = export_classes()
+        import importlib
+        gpt2m = importlib.import_module('moss_export.gpt2_decoder')
+        qw = importlib.import_module('moss_export.qwen3_decoder')
+
+        class _Small(MossTTSLocalModel):
+            def __init__(self, config):
+                # Elternklasse überspringen (sie baut Embeddings mit hidden_size); PreTrainedModel-Init direkt.
+                super(MossTTSLocalModel, self).__init__(config)
+                self._tied_weights_keys = self._build_tied_weights_keys(config)
+                config.qwen3_config.pad_token_id = config.pad_token_id
+                config.qwen3_config._attn_implementation = config.attn_implementation
+                lg = config.gpt2_config.to_dict()
+                lg['n_layer'] = int(getattr(config, 'local_transformer_layers', config.gpt2_config.n_layer))
+                lg['n_positions'] = int(config.n_vq) + 1
+                lg['n_ctx'] = int(config.n_vq) + 1
+                lg = GPT2Config(**lg)
+                lg.pad_token_id = config.pad_token_id
+                lg._attn_implementation = config.local_transformer_attn_implementation
+                self.transformer = qw.MossQwen3Model(config.qwen3_config)
+                self.local_transformer = gpt2m.MossTTSNanoGPT2Model(lg, attn_implementation=config.local_transformer_attn_implementation)
+                self.local_transformer.wte = nn.Identity()
+                H = int(config.hidden_size); L = int(config.local_hidden_size)
+                self.audio_embeddings = nn.ModuleList([nn.Embedding(int(config.audio_codebook_sizes[i]), L) for i in range(config.n_vq)])
+                self.text_lm_head = nn.Linear(H, int(config.vocab_size), bias=False)
+                self.audio_lm_heads = nn.ModuleList([nn.Linear(L, int(config.audio_codebook_sizes[i]), bias=False) for i in range(config.n_vq)])
+                self.local_text_lm_head = nn.Linear(L, 2, bias=False) if self._use_binary_local_text_head() else None
+                if H != L:
+                    self.proj_in = nn.Linear(L, H, bias=False)      # Audio-Embedding-Summe -> Backbone
+                    self.proj_out = nn.Linear(H, L, bias=False)     # Backbone-Hidden -> Head
+                else:
+                    self.proj_in = nn.Identity(); self.proj_out = nn.Identity()
+                self.post_init()
+                self.tie_weights()
+                self.initialize_local_text_lm_head_from_text_lm_head()
+
+            def _build_inputs_embeds(self, input_ids):
+                if input_ids.ndim != 3 or input_ids.shape[-1] != self.config.n_vq + 1:
+                    raise ValueError(f'Expected [B,T,{self.config.n_vq + 1}], got {tuple(input_ids.shape)}')
+                text_ids = input_ids[..., 0]
+                emb = self.transformer.embed_tokens(text_ids)
+                acc = None
+                for ci, e in enumerate(self.audio_embeddings):
+                    ids = input_ids[..., ci + 1]
+                    valid = ids.ne(self.config.audio_pad_token_id)
+                    a = e(ids.masked_fill(~valid, 0)) * valid.unsqueeze(-1)
+                    acc = a if acc is None else acc + a
+                return emb + self.proj_in(acc.to(emb.dtype))
+
+            def _global_hidden_to_local(self, h):
+                return self.proj_out(h)
+
+        klass._cls = _Small
+        return _Small
+
+
+def load_qwen_backbone(model, log=print):
+    sd = load_file(f'{QWEN}/model.safetensors', device='cpu')
+    tgt = model.transformer.state_dict()
+    loaded = {}
+    for k, v in sd.items():
+        if k.startswith('model.'):
+            kk = k[len('model.'):]
+            if kk in tgt:
+                assert tuple(v.shape) == tuple(tgt[kk].shape), (k, v.shape, tgt[kk].shape)
+                loaded[kk] = v.to(tgt[kk].dtype)
+    missing, unexpected = model.transformer.load_state_dict(loaded, strict=False)
+    missing = [m for m in missing if 'rotary' not in m]
+    if missing or unexpected:
+        raise RuntimeError(f'Qwen3 load: missing={missing[:10]} unexpected={unexpected[:10]}')
+    extra = [k for k in sd if not k.startswith('model.')]
+    log(f'[qwen3] loaded {len(loaded)} tensors into backbone; non-backbone keys in file: {extra}')
+    model.tie_weights()
+
+
+def load_sft3_head(model, log=print):
+    """Lokaler Transformer + 12 Audio-Embeddings (+ Köpfe, gewichtsgeteilt) + binärer Stop-Kopf aus dem SFT3-Export."""
+    want_prefix = ('local_transformer.', 'audio_embeddings.', 'local_text_lm_head.')
+    loaded = {}
+    with safe_open(f'{SFT3}/model.safetensors', framework='pt', device='cpu') as f:
+        keys = list(f.keys())
+        for k in keys:
+            if k.startswith(want_prefix):
+                loaded[k] = f.get_tensor(k)
+    sd = model.state_dict()
+    for k, v in loaded.items():
+        assert k in sd, k
+        assert tuple(v.shape) == tuple(sd[k].shape), (k, v.shape, sd[k].shape)
+    missing, unexpected = model.load_state_dict({k: v.to(sd[k].dtype) for k, v in loaded.items()}, strict=False)
+    assert not unexpected, unexpected
+    model.tie_weights()
+    # Kontrolle: Head-Gewichte sind Embeddings
+    for i in range(model.config.n_vq):
+        assert model.audio_lm_heads[i].weight.data_ptr() == model.audio_embeddings[i].weight.data_ptr()
+    n = sum(v.numel() for v in loaded.values())
+    log(f'[sft3-head] loaded {len(loaded)} tensors ({n/1e6:.1f}M params) — local_transformer + audio_embeddings + stop head')
+    return loaded.keys()
+
+
+def freeze_head(model, log=print):
+    n = 0
+    for mod in (model.local_transformer, model.audio_embeddings, model.local_text_lm_head):
+        for p in mod.parameters():
+            p.requires_grad_(False); n += p.numel()
+    log(f'[freeze] head frozen: {n/1e6:.1f}M params')
+
+
+def build(kind, log=print, dtype=torch.float32):
+    cfg = make_config(kind)
+    cls = MossTTSSmallModel.cls()
+    model = cls(cfg)
+    load_qwen_backbone(model, log)
+    if kind == 'M1':
+        load_sft3_head(model, log)
+        freeze_head(model, log)
+    ntr = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    ntot = sum(p.numel() for p in model.parameters())
+    log(f'[{kind}] params total={ntot/1e6:.1f}M trainable={ntr/1e6:.1f}M hidden={cfg.hidden_size} local={cfg.local_hidden_size}')
+    return model.to(dtype), cfg
+
+
+def param_groups(model, kind, lr_backbone, lr_new):
+    """Backbone (vortrainiert) vs. neue Teile (Projektionen / Head von Null)."""
+    bb, new = [], []
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        (bb if n.startswith('transformer.') or n.startswith('text_lm_head.') else new).append(p)
+    return [{'params': bb, 'lr': lr_backbone, 'name': 'backbone'}, {'params': new, 'lr': lr_new, 'name': 'new'}]
+
+
+if __name__ == '__main__':
+    kind = sys.argv[1] if len(sys.argv) > 1 else 'M2'
+    m, cfg = build(kind)
+    print(m.__class__.__name__, cfg.small_kind)

code/pack_moss.py

@@ -0,0 +1,40 @@
+"""Packen für M1/M2 mit dem SFT3-Prozessor (identisch zu sft_train2.build_examples, ohne Prompt-Zufall)."""
+import sys, torch
+sys.path.insert(0, '/e/data1/datasets/playground/mmlaion/schuhmann1/dramabox/train')
+from moss_pack import collate          # Right-Padding, Labels maskiert
+from prompt_fmt import user_content, AUDIO_PLACEHOLDER
+
+
+class MossPacker:
+    def __init__(self, proc, cfg):
+        self.proc = proc
+        self.assist = int(cfg.audio_assistant_slot_token_id); self.aend = int(cfg.audio_end_token_id)
+        self.audio_pad = int(cfg.audio_pad_token_id); self.text_pad = int(cfg.pad_token_id)
+
+    def pack(self, meta, codes, ref_codes=None):
+        use_ref = ref_codes is not None
+        content = user_content(meta, use_ref)
+        um = {"role": "user", "content": content, "audio_codes_list": [torch.as_tensor(ref_codes)] if use_ref else []}
+        am = self.proc.build_assistant_message(audio_codes_list=[torch.as_tensor(codes)])
+        b = self.proc([[um, am]], mode="computing_loss")
+        ids = b["input_ids"][0].long()
+        lab = torch.full_like(ids, -100)
+        lab[:-1] = ids[1:]
+        tt = lab[:, 0]
+        sup = (tt == self.assist) | (tt == self.aend)
+        fa = (ids[:, 0] == self.assist).nonzero()
+        if fa.numel() > 0:
+            sup[: max(0, int(fa[0]) - 1)] = False       # Referenz-<|audio_end|> nicht als Stop lernen
+        lab[~sup] = -100
+        au = lab[:, 1:]
+        au[(au >= self.audio_pad) | (au < 0)] = -100
+        lab[:, 1:] = au
+        return {"input_ids": ids, "labels": lab, "n_sup": int(sup.sum())}
+
+    def prompt_ids(self, meta, ref_codes=None, frames=None):
+        """Generierungs-Eingabe [T,13] (endet mit <|audio_start|>)."""
+        use_ref = ref_codes is not None
+        content = user_content(meta, use_ref, frames=frames)
+        um = {"role": "user", "content": content, "audio_codes_list": [torch.as_tensor(ref_codes)] if use_ref else []}
+        b = self.proc([[um]], mode="generation")
+        return b["input_ids"][0].long()

code/packing.py

@@ -0,0 +1,162 @@
+"""Preserve the inherited prompt skeleton and inject measured score slots.
+
+Only the global caption and the leading sentence-cue parentheses are replaced
+in score mode. The inherited template repeats the SCRIPT in Instruction and
+Text; both copies map to the same measured local feature bank.
+"""
+import re
+import sys
+from pathlib import Path
+import numpy as np
+import torch
+from conditioning import MARKERS, feature_vector, install_markers
+
+sys.path.insert(0, '/e/scratch/reformo/schuhmann1_moss/code/small_tts')
+from pack_moss import MossPacker
+from prompt_fmt import user_content
+
+GLOBAL_BLOCK = MARKERS[0] + MARKERS[1] * 5 + MARKERS[2]
+LOCAL_BLOCK = MARKERS[3] + MARKERS[4] * 10 + MARKERS[5]
+
+
+def score_prompt(prompt, sentence_count):
+    if not prompt.startswith('GENERAL: ') or '\nSCRIPT:\n' not in prompt:
+        raise ValueError('Score mode requires the verified GENERAL/SCRIPT prompt format')
+    _, script = prompt.split('\nSCRIPT:\n', 1)
+    lines = script.splitlines()
+    if len(lines) != sentence_count or not lines:
+        raise ValueError(f'Prompt cues and measured sentence rows disagree: {len(lines)} vs {sentence_count}')
+    edited, bodies = [], []
+    for line in lines:
+        match = re.match(r'^\([^)]*\)(.*)$', line)
+        if not match:
+            raise ValueError('Expected one leading parenthesized sentence cue')
+        body = match.group(1)
+        bodies.append(body)
+        edited.append(LOCAL_BLOCK + body)
+    changed = 'GENERAL: ' + GLOBAL_BLOCK + '\nSCRIPT:\n' + '\n'.join(edited)
+    # Everything following the cue stays byte-identical, including burst tags,
+    # pause spellings, numbers and punctuation. Durations live in the template.
+    assert all(new[len(LOCAL_BLOCK):] == old for new, old in zip(edited, bodies))
+    return changed
+
+
+def labels_for(ids, assist, aend, audio_pad):
+    labels = torch.full_like(ids, -100)
+    labels[:-1] = ids[1:]
+    targets = labels[:, 0]
+    supervised = targets.eq(assist) | targets.eq(aend)
+    first = ids[:, 0].eq(assist).nonzero()
+    if not first.numel():
+        raise ValueError('No assistant target frames in the packed training example')
+    supervised[:max(0, int(first[0]) - 1)] = False
+    labels[~supervised] = -100
+    audio = labels[:, 1:]
+    audio[(audio >= audio_pad) | (audio < 0)] = -100
+    return labels
+
+
+def generated_audio(result, config):
+    """MOSS generate returns BOA + 13-channel rows, not a [frames,12] array.
+
+    Follow the already validated M1/M2 evaluator's extraction. Any unexpected
+    remaining code value is an error rather than silently clipping corruption.
+    """
+    start_length, sequence = result[0]
+    assert start_length == 0, 'This continuation path requires an empty target prefix'
+    assert sequence.ndim == 2 and sequence.shape[1] == int(config.n_vq) + 1
+    assert int(sequence[0, 0]) == int(config.audio_start_token_id)
+    audio = sequence[1:, 1:]
+    keep = (audio[:, 0] != int(config.audio_pad_token_id)) & (audio[:, 0] >= 0)
+    audio = audio[keep]
+    if audio.numel():
+        assert int(audio.min()) >= 0 and int(audio.max()) < int(config.audio_pad_token_id)
+    return audio
+
+
+class ScorePacker(MossPacker):
+    def __init__(self, proc, config, schema):
+        super().__init__(proc, config)
+        self.schema = schema
+        self.markers = install_markers(proc.tokenizer, int(config.vocab_size))
+
+    def _prepare(self, meta, mode, ref_codes):
+        if mode not in ('instruction', 'reference', 'score'):
+            raise ValueError(mode)
+        if (ref_codes is not None) != (mode == 'reference'):
+            raise ValueError('Reference mode must have a genuine non-target reference')
+        record = dict(meta)
+        global_values = [0.] * (2 * len(self.schema['global_dimensions']))
+        local_values = []
+        if mode == 'score':
+            measurements = meta['conditioning_measurements']
+            if not measurements.get('validated'):
+                raise ValueError('Score mode requires independently validated measurement provenance')
+            local = measurements['sentences']
+            record['prompt'] = score_prompt(meta['prompt'], len(local))
+            global_values = feature_vector(measurements['global'], self.schema['global_dimensions'])
+            for row in local:
+                if row['scope'] not in ('sentence_audio', 'whole_target_single_cue_rescored', 'windowed_sentence_audio') or not row.get('provenance'):
+                    raise ValueError('Whole-clip or unknown measurements cannot populate local score slots')
+                local_values.append(feature_vector(row['values'], self.schema['local_dimensions']))
+            if not local_values:
+                raise ValueError('No local measurements for score mode')
+        return record, global_values, local_values
+
+    def _index(self, ids, mode, sentence_count):
+        index = torch.zeros(ids.shape[0], dtype=torch.long)
+        text = ids[:, 0].tolist()
+        global_start = self.markers[MARKERS[0]]
+        local_start = self.markers[MARKERS[3]]
+        gs = [i for i, token in enumerate(text) if token == global_start]
+        ls = [i for i, token in enumerate(text) if token == local_start]
+        if mode != 'score':
+            if gs or ls:
+                raise ValueError('Score markers unexpectedly present outside score mode')
+            return index
+        if len(gs) != 1 or len(ls) != 2 * sentence_count:
+            raise ValueError(f'Unexpected template score-block coverage: global={len(gs)}, sentence={len(ls)}')
+        expect_global = [global_start] + [self.markers[MARKERS[1]]] * 5 + [self.markers[MARKERS[2]]]
+        expect_local = [local_start] + [self.markers[MARKERS[4]]] * 10 + [self.markers[MARKERS[5]]]
+        i = gs[0]
+        assert text[i:i + 7] == expect_global
+        index[i:i + 7] = torch.arange(1, 8)
+        for repeat, i in enumerate(ls):
+            assert text[i:i + 12] == expect_local
+            base = 8 + (repeat % sentence_count) * 12
+            index[i:i + 12] = torch.arange(base, base + 12)
+        return index
+
+    def pack_mode(self, meta, codes, mode, ref_codes=None, generation=False):
+        record, glob, local = self._prepare(meta, mode, ref_codes)
+        content = user_content(record, mode == 'reference')
+        user = {'role': 'user', 'content': content,
+                'audio_codes_list': [torch.as_tensor(np.asarray(ref_codes).copy(), dtype=torch.long)] if ref_codes is not None else []}
+        conversation = [user]
+        if not generation:
+            conversation.append(self.proc.build_assistant_message(audio_codes_list=[torch.as_tensor(np.asarray(codes), dtype=torch.long)]))
+        packed = self.proc([conversation], mode='generation' if generation else 'computing_loss')
+        ids = packed['input_ids'][0].long()
+        index = self._index(ids, mode, len(local))
+        labels = torch.full_like(ids, -100) if generation else labels_for(ids, self.assist, self.aend, self.audio_pad)
+        if not generation:
+            assert int((labels[:, 1] >= 0).sum()) == len(codes), 'Target frame supervision changed'
+            assert int((labels[:, 0] == self.aend).sum()) == 1, 'Reference end leaked into the loss'
+        return {'input_ids': ids, 'labels': labels, 'score_index': index,
+                'global_values': glob, 'local_values': local, 'mode': mode}
+
+    def collate(self, examples):
+        from pack_moss import collate
+        batch = collate(examples, self.text_pad, self.audio_pad)
+        B, T = batch['attention_mask'].shape
+        S = max(1, max(len(row['local_values']) for row in examples))
+        glob = torch.tensor([row['global_values'] for row in examples], dtype=torch.float32)
+        local = torch.zeros(B, S, 2 * len(self.schema['local_dimensions']), dtype=torch.float32)
+        index = torch.zeros(B, T, dtype=torch.long)
+        for i, row in enumerate(examples):
+            index[i, :len(row['score_index'])] = row['score_index']
+            if row['local_values']:
+                local[i, :len(row['local_values'])] = torch.tensor(row['local_values'])
+        batch['score_conditioning'] = (glob, local, index)
+        batch['modes'] = [row['mode'] for row in examples]
+        return batch

code/plan.py

@@ -0,0 +1,96 @@
+"""Same global sample stream across ranks, without dropped tail targets."""
+import hashlib
+import json
+import numpy as np
+
+
+def benchmark_phases(world):
+    assert world in (8, 16)
+    return [{'name': 'strong', 'offset': 0, 'global_batch': 256, 'updates': 40, 'warmup': 10},
+            {'name': 'weak', 'offset': 10240, 'global_batch': world * 32, 'updates': 12, 'warmup': 4, 'replicated_base_world': 8}]
+
+
+def step_indices(phase, step, rank, world, samples):
+    if phase.get('replicated_base_world'):
+        # Benchmark only: duplicate the same eight real rank streams on the
+        # second pair of nodes, giving exactly equal per-GPU input/token load.
+        base = phase['replicated_base_world']
+        local_count = phase['global_batch'] // world
+        start = phase['offset'] + step * base * local_count
+        return list(range(start + rank % base, start + base * local_count, base))
+    start = phase['offset'] + step * phase['global_batch']
+    end = min(samples, start + phase['global_batch'])
+    if start >= samples: return []
+    return list(range(start + rank, end, world))
+
+
+def microbatches(examples, max_padded_tokens=12000, max_examples=16):
+    """Length sort only within one already fixed optimizer batch."""
+    examples = sorted(examples, key=lambda ex: len(ex['input_ids']))
+    batches, current = [], []
+    for row in examples:
+        n = len(row['input_ids'])
+        if current and (n * (len(current) + 1) > max_padded_tokens or len(current) >= max_examples):
+            batches.append(current); current = []
+        current.append(row)
+    if current: batches.append(current)
+    return batches
+
+
+def dynamic_weights(local_frames, local_samples, global_frames, global_samples, world):
+    """Preserve one global per-channel token mean across microbatches/ranks.
+
+    The inherited loss normalizes by sum(weights), so the caller multiplies
+    its return value by returned scale. DDP subsequently averages gradients.
+    This is algebraically the ordinary inherited objective on the global batch.
+    """
+    weights = [world * (local_frames + local_samples) / (global_frames + global_samples)]
+    weights += [(32. / 12) * world * local_frames / global_frames] * 12
+    return weights, sum(weights) / 33.
+
+
+def check_plan(samples=16384):
+    unions = {}
+    for world in (8, 16):
+        phases = benchmark_phases(world)
+        for phase in phases:
+            seen = []
+            for step in range(phase['updates']):
+                global_ids = sorted(i for rank in range(world) for i in step_indices(phase, step, rank, world, samples))
+                if phase['name'] == 'weak':
+                    start = phase['offset'] + step * 256
+                    expect = sorted(list(range(start, start + 256)) * (world // 8))
+                else:
+                    expect = list(range(phase['offset'] + step * phase['global_batch'], phase['offset'] + (step + 1) * phase['global_batch']))
+                assert global_ids == expect
+                seen.extend(global_ids)
+            assert len(seen) == len(set(seen)) * (world // 8 if phase['name'] == 'weak' else 1)
+            unions[(world, phase['name'])] = seen
+    assert unions[(8, 'strong')] == unions[(16, 'strong')]
+    for step in range(12):
+        for rank in range(16):
+            assert step_indices(benchmark_phases(16)[1], step, rank, 16, samples) == step_indices(benchmark_phases(8)[1], step, rank % 8, 8, samples)
+    # Test mean/gradient coefficients independently, including unequal lengths.
+    rng = np.random.default_rng(71)
+    channels = rng.normal(size=(263, 13))
+    frame_groups = [13, 49, 81, 120]
+    sample_groups = [1, 3, 5, 7]
+    expected = channels.mean(axis=0)
+    actual = np.zeros(13)
+    start = 0
+    for frames, sample_count in zip(frame_groups, sample_groups):
+        weights, scale = dynamic_weights(frames, sample_count, sum(frame_groups), sum(sample_groups), 4)
+        # Audio-channel linear coefficients recover the exact global mean after DDP.
+        audio_coefficient = weights[1] / sum(weights) * scale / 4
+        actual[1:] += channels[start:start + frames, 1:].mean(axis=0) * audio_coefficient
+        start += frames
+    assert np.allclose(actual[1:], expected[1:] * (32 / 12) / 33)
+    return {'status': 'PASS', 'strong_unique_samples': len(unions[(8, 'strong')]),
+            'strong_union_sha256': hashlib.sha256(np.asarray(unions[(8, 'strong')], dtype='<i8').tobytes()).hexdigest(),
+            'worlds': [8, 16], 'global_batch_strong': 256, 'per_gpu_weak': 32,
+            'strong_warmup_updates': 10, 'strong_steady_updates': 30, 'weak_warmup_updates': 4, 'weak_steady_updates': 8,
+            'weak_input_rule': 'Exact same eight real per-GPU streams, replicated once across the additional eight GPUs at four nodes. Benchmark-only repetitions, explicitly counted as presentations.',
+            'loss_weighting': 'Exact global per-channel supervised-token mean, independent of rank and microbatch partition.'}
+
+
+if __name__ == '__main__': print(json.dumps(check_plan(), indent=2))

code/plan_caption_production_32nodes.json

@@ -0,0 +1,137 @@
+{
+  "kind": "m2_continuous_large_talker_v1",
+  "architecture": "qwen3_0.6b_plus_fresh_sft3_width_talker",
+  "initialization": {
+    "semantic": "original_pretrained_Qwen3-0.6B",
+    "talker": "random_fresh"
+  },
+  "nodes": 32,
+  "ranks_per_node": 4,
+  "world_size": 128,
+  "samples_per_gpu": 32,
+  "global_batch": 4096,
+  "seed": 20260923,
+  "max_padded_tokens": 12000,
+  "max_examples": 16,
+  "lr_backbone_peak": 8e-05,
+  "lr_talker_peak": 0.00024,
+  "lr_floor_factor": 0.05,
+  "scheduler": "single_linear_warmup_then_cosine_no_stage_restarts",
+  "gradient_communication": "bfloat16",
+  "deterministic": false,
+  "checkpoint_policy": "all_stage_boundaries_plus_interval",
+  "checkpoint_interval_seconds": 3600,
+  "max_invocation_seconds": 40000,
+  "resume_policy": "latest_complete_same_plan_world",
+  "prompt_contract": "m2-multiformat-caption-v1-20260923",
+  "allowed_prompt_format_ids": [
+    "m2-cascade-timed-repair-v1-20260920",
+    "m2-multiformat-caption-v1-20260923"
+  ],
+  "codec_contract": "MOSS-Audio-Tokenizer-v2; 12 codebooks/frame; 12.5 frames/s; 150 codebook tokens/s",
+  "code_sha256": {
+    "continuous_train.py": "58b6c82beaaf6ff8e2b65fea6312a28afd75079153c13dfdb9d432de4c42befc",
+    "large_talker.py": "3a1d0ed44467a6933dfac2a21a143a1c802482bafb2e992a1e6d8793e3fcdae4",
+    "caption_curriculum_dataset.py": "6c01a3e26e5908d01391bfd4aa89815746a6a26aedca397ba1299e5d688a2a64",
+    "cascade_dataset.py": "784967ce6d1a545d565a95a16c940ee2743085362576d1566c2b37aa2f40b351",
+    "timed_prompt.py": "4d06d0447573284f09804ba516bb2121f65d4cb26ec5a2c01c5bd8d425c8a642",
+    "anno_timed.py": "06df8ed3cb7f67d59d1bfca2931e26dd4489344925c01e4412836d2fcf946ec6"
+  },
+  "created_at": "2026-09-25T08:43:16.244121+02:00",
+  "stages": [
+    {
+      "name": "S1",
+      "tier": "x100",
+      "manifest": "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1/manifests/S1/manifest.json",
+      "manifest_sha256": "447cc8c91fa507dad7e4c30fbe36e1f5510efcc13e49b2e267e4fdbc04cad56f",
+      "presentations": 66621417,
+      "measured_hours": 188353.851,
+      "updates": 16265
+    },
+    {
+      "name": "S2",
+      "tier": "x50",
+      "manifest": "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1/manifests/S2/manifest.json",
+      "manifest_sha256": "1db622e32c210ad3bf258251852e4f4fd0095d521da8aa08c91431e1ac175ade",
+      "presentations": 41036973,
+      "measured_hours": 116837.568,
+      "updates": 10019
+    },
+    {
+      "name": "S3",
+      "tier": "tier7",
+      "manifest": "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1/manifests/S3/manifest.json",
+      "manifest_sha256": "7625d9fb0a74709af1f6ad7698c77369c32e131271b8f88f8714d12e821eee50",
+      "presentations": 24789493,
+      "measured_hours": 72314.446,
+      "updates": 6053
+    },
+    {
+      "name": "S4",
+      "tier": "tier6",
+      "manifest": "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1/manifests/S4/manifest.json",
+      "manifest_sha256": "255cb00bc50183ceadf7b1dec23b65a3fb6f575daa06002365a9c658e21a6226",
+      "presentations": 19842627,
+      "measured_hours": 57419.06,
+      "updates": 4845
+    },
+    {
+      "name": "S5",
+      "tier": "tier5",
+      "manifest": "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1/manifests/S5/manifest.json",
+      "manifest_sha256": "eb3b6e0d004f6460adff2274d3041925d5b417a853010dad3c39106b30074a24",
+      "presentations": 10152521,
+      "measured_hours": 28716.381,
+      "updates": 2479
+    },
+    {
+      "name": "S6",
+      "tier": "tier4",
+      "manifest": "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1/manifests/S6/manifest.json",
+      "manifest_sha256": "90094e2cb294791ff7e24df7f36acddac39e95cca25a6c2fc885bc2b9f2f5bb2",
+      "presentations": 5217313,
+      "measured_hours": 14361.366,
+      "updates": 1274
+    },
+    {
+      "name": "S7",
+      "tier": "tier3",
+      "manifest": "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1/manifests/S7/manifest.json",
+      "manifest_sha256": "b3d94f4b4b10e78fc4b4a33186b86611f7d1ad69a88a84175f9d9c0a69f98df7",
+      "presentations": 2171373,
+      "measured_hours": 5749.569,
+      "updates": 531
+    },
+    {
+      "name": "S8",
+      "tier": "tier2",
+      "manifest": "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1/manifests/S8/manifest.json",
+      "manifest_sha256": "0187e2809cee89f6f4431b3827d8046ce2f8d4f059c37e3086e38d42197db865",
+      "presentations": 1116702,
+      "measured_hours": 2877.426,
+      "updates": 273
+    },
+    {
+      "name": "S9",
+      "tier": "tier1",
+      "manifest": "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1/manifests/S9/manifest.json",
+      "manifest_sha256": "edf6e8fff72d6d6e8128a0c2f06e7a0ee2aa1064c420d85dbc17fd7f0871372f",
+      "presentations": 568338,
+      "measured_hours": 1439.466,
+      "updates": 139
+    },
+    {
+      "name": "S10",
+      "tier": "tier0",
+      "manifest": "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1/manifests/S10/manifest.json",
+      "manifest_sha256": "24d9d3f500c8132fc07168ad07470f7cafdecec4dde7c7b064d395d49909ec13",
+      "presentations": 234815,
+      "measured_hours": 577.881,
+      "updates": 58
+    }
+  ],
+  "total_updates": 41936,
+  "warmup_steps": 2097,
+  "output": "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1/training",
+  "purpose": "authorized continuous S1-to-S10 multi-caption production"
+}

code/prompt_fmt.py

@@ -0,0 +1,39 @@
+"""Prompt-Oberfläche für alle drei Modelle: das SFT3-Nutzer-Template (prompt_lib2.USER_INST_TEMPLATE, byte-gleich),
+Instruction = Hausformat-Prompt der Leiter (GENERAL + SCRIPT), Text = SCRIPT-Teil, Tokens = Frames, Referenz = <|audio|> oder None."""
+AUDIO_PLACEHOLDER = "<|audio|>"
+USER_INST_TEMPLATE = """<user_inst>
+- Reference(s):
+{reference}
+- Instruction:
+{instruction}
+- Tokens:
+{tokens}
+- Quality:
+{quality}
+- Sound Event:
+{sound_event}
+- Ambient Sound:
+{ambient_sound}
+- Language:
+{language}
+- Text:
+{text}
+</user_inst>"""
+
+
+def lang_name(lang):
+    return "German" if str(lang or "en").lower().startswith("de") else "English"
+
+
+def split_prompt(prompt):
+    if 'SCRIPT:\n' in prompt:
+        return prompt, prompt.split('SCRIPT:\n', 1)[1].strip()
+    return prompt, prompt.strip()
+
+
+def user_content(meta, use_ref, frames=None):
+    instruction, script = split_prompt(meta['prompt'])
+    return USER_INST_TEMPLATE.format(
+        reference=AUDIO_PLACEHOLDER if use_ref else "None", instruction=instruction.strip() or "None",
+        tokens=int(frames if frames is not None else meta['frames']), quality="None", sound_event="None",
+        ambient_sound="None", language=lang_name(meta['lang']), text=script or "None")

code/state_io.py

@@ -0,0 +1,86 @@
+"""Atomic full-state DDP checkpoints with per-rank RNG and pinned contracts."""
+from datetime import datetime
+import json
+import os
+from pathlib import Path
+import random
+import numpy as np
+import torch
+import torch.distributed as dist
+from continuation_control import sha
+
+
+def ranks():
+    return (dist.get_rank(), dist.get_world_size()) if dist.is_initialized() else (0, 1)
+
+
+def barrier():
+    if dist.is_initialized():
+        dist.barrier()
+
+
+def unwrap(model):
+    return model.module if hasattr(model, 'module') else model
+
+
+def save(root, model, optimizer, scheduler, step, contract, export=False):
+    rank, world = ranks()
+    root = Path(root)
+    rng = {'torch': torch.get_rng_state(), 'cuda': torch.cuda.get_rng_state(),
+           'numpy': np.random.get_state(), 'python': random.getstate()}
+    all_rng = [None] * world if rank == 0 else None
+    if world > 1:
+        dist.gather_object(rng, all_rng, dst=0)
+    else:
+        all_rng = [rng]
+    directory = root / f'step-{step:08d}'
+    if rank == 0:
+        temporary = root / f'.step-{step:08d}-{os.environ.get("SLURM_JOB_ID", "local")}.partial'
+        temporary.mkdir(parents=True, exist_ok=True)
+        torch.save({'model': unwrap(model).state_dict(), 'optimizer': optimizer.state_dict(),
+                    'scheduler': scheduler.state_dict(), 'rng': all_rng, 'step': step, 'world': world,
+                    'contract': contract}, temporary / 'training_state.pt')
+        metadata = {'step': step, 'world': world, 'contract': contract,
+                    'created_at': datetime.now().astimezone().isoformat(),
+                    'state_bytes': (temporary / 'training_state.pt').stat().st_size,
+                    'state_sha256': sha(temporary / 'training_state.pt')}
+        if export:
+            base = unwrap(model).base
+            cache, state = {}, {}
+            for name, value in base.state_dict().items():
+                key = (value.data_ptr(), tuple(value.shape), value.dtype)
+                if key not in cache:
+                    cache[key] = value.detach().to(dtype=torch.bfloat16, device='cpu')
+                state[name] = cache[key]
+            torch.save(state, temporary / 'model_bf16.pt')
+            metadata['export_bytes'] = (temporary / 'model_bf16.pt').stat().st_size
+            metadata['export_sha256'] = sha(temporary / 'model_bf16.pt')
+        (temporary / 'COMPLETE.json').write_text(json.dumps(metadata, indent=2) + '\n')
+        if directory.exists():
+            raise RuntimeError(f'Refusing to overwrite existing checkpoint {directory}')
+        temporary.replace(directory)
+        pointer = root / 'current.json.tmp'
+        pointer.write_text(json.dumps({'directory': directory.name, 'step': step}) + '\n')
+        pointer.replace(root / 'current.json')
+    barrier()
+    return directory
+
+
+def load(directory, model, optimizer, scheduler, contract):
+    directory = Path(directory)
+    rank, world = ranks()
+    metadata = json.loads((directory / 'COMPLETE.json').read_text())
+    assert metadata['contract'] == contract and metadata['world'] == world
+    assert (directory / 'training_state.pt').stat().st_size == metadata['state_bytes']
+    state = torch.load(directory / 'training_state.pt', map_location='cpu', weights_only=False)
+    assert state['contract'] == contract and state['world'] == world and len(state['rng']) == world
+    unwrap(model).load_state_dict(state['model'], strict=True)
+    optimizer.load_state_dict(state['optimizer']); scheduler.load_state_dict(state['scheduler'])
+    rng = state['rng'][rank]
+    torch.set_rng_state(rng['torch']); torch.cuda.set_rng_state(rng['cuda'])
+    np.random.set_state(rng['numpy']); random.setstate(rng['python'])
+    step = int(state['step'])
+    assert step==metadata['step']
+    del state
+    barrier()
+    return step

code/tar_member_map.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+"""Packed, read-only map from TAR path/member name to byte offset and size.
+
+The corpus stays in the existing large TAR containers.  A single compressed
+record per TAR replaces repeated full TAR-directory scans on every DDP rank.
+"""
+from collections import OrderedDict
+import json
+import os
+from pathlib import Path
+import threading
+import zlib
+
+
+class TarMemberMap:
+    def __init__(self, manifest, capacity=256):
+        self.manifest_path = Path(manifest)
+        spec = json.loads(self.manifest_path.read_text())
+        assert spec['status'] == 'complete' and spec['format'] == 'zlib-json-tar-member-map-v1'
+        root = self.manifest_path.parent
+        index_path = root / spec['index']
+        data_path = root / spec['data']
+        assert index_path.stat().st_size == spec['index_bytes']
+        assert data_path.stat().st_size == spec['data_bytes']
+        index = json.loads(index_path.read_text())
+        assert len(index['paths']) == spec['tars']
+        assert all(len(index[k]) == spec['tars'] for k in ('offsets', 'sizes', 'members', 'tar_sizes'))
+        self.rows = {path: (int(index['offsets'][i]), int(index['sizes'][i]),
+                            int(index['members'][i]), int(index['tar_sizes'][i]))
+                     for i, path in enumerate(index['paths'])}
+        self.fd = os.open(data_path, os.O_RDONLY)
+        self.capacity = int(capacity)
+        self.local = threading.local()
+
+    def _cache(self):
+        if not hasattr(self.local, 'members'):
+            self.local.members = OrderedDict()
+        return self.local.members
+
+    def locate(self, path, member):
+        cache = self._cache()
+        if path not in cache:
+            assert path in self.rows, f'TAR absent from packed member map: {path}'
+            offset, size, expected_members, expected_tar_size = self.rows[path]
+            assert os.path.getsize(path) == expected_tar_size, f'TAR size changed after index build: {path}'
+            compressed = os.pread(self.fd, size, offset)
+            assert len(compressed) == size
+            decoded = json.loads(zlib.decompress(compressed))
+            assert len(decoded) == expected_members
+            cache[path] = decoded
+            while len(cache) > self.capacity:
+                cache.popitem(last=False)
+        cache.move_to_end(path)
+        value = cache[path].get(member)
+        assert value is not None, f'Member absent from packed map: {path}::{member}'
+        return int(value[0]), int(value[1])
+