[Model] Add Audio8 TTS Preview 0.6B (DualAR, 44.1 kHz codec) (#6157)

Signed-off-by: NancyFyong <NancyFyong@users.noreply.github.com>
Co-authored-by: NancyFyong <NancyFyong@users.noreply.github.com>

benchmarks/tts/README.md

@@ -85,6 +85,25 @@ vllm bench serve --omni \
     --save-result --result-dir ./results
 ```
 
+#### Set the sample rate for models that are not 24 kHz
+
+Streamed PCM carries no header, so the harness derives `audio_duration` from the
+byte count and a **default 24000 Hz** assumption
+(`vllm_omni/metrics/definitions.py:DEFAULT_AUDIO_SAMPLE_RATE`). For a model that
+emits another rate, export the real one or every duration-derived metric is
+wrong by `real_sr / 24000`:
+
+```bash
+# Audio8 TTS, Fish Speech S2 Pro, Ming-omni-tts, ... all decode at 44.1 kHz
+export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100
+```
+
+Without it, a 44.1 kHz model reports `audio_duration` 1.84x too long, so
+`audio_rtf` looks 1.84x *better* than reality and `audio_underrun` /
+`Streaming continuity OK rate` are measured against inflated chunk budgets.
+A quick sanity check: compare `Mean AUDIO_DURATION (s)` against the actual
+length of a saved WAV for the same prompt.
+
 #### Streaming continuity (`audio_underrun`)
 
 `audio_underrun` (seconds) is the per-request worst-case buffer deficit

benchmarks/tts/model_configs.yaml

@@ -17,6 +17,17 @@ models:
         extra_params:
           lang: en
           text_normalization: true
+  # Audio8 TTS Preview: DualAR, 44.1 kHz. trust_remote_code stays false.
+  # Bench at 44.1 kHz: export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100, otherwise
+  # the streamed-PCM duration defaults to 24 kHz and RTF looks 1.84x too good.
+  Audio8/Audio8-TTS-Preview-0.6b:
+    supported_tasks: [default_voice, voice_clone]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    trust_remote_code: false
+    task_extra_body:
+      default_voice: {}
+      voice_clone: {}
 
   # -CustomVoice checkpoints lack speaker_encoder weights, so voice_clone is
   # NOT supported (an attempt raises ValueError from _extract_speaker_embedding

docs/models/supported_models.md

@@ -97,6 +97,7 @@ th {
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
 | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
 | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `Audio8TTSSlowARForConditionalGeneration` | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✅︎ | | | | — |
 | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) |
 | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |

examples/offline_inference/text_to_speech/README.md

@@ -14,6 +14,7 @@ list of supported architectures across all modalities, see
 
 | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
 |---|---|---|---|---|---|---|
+| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz |
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
 | Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz |
@@ -44,6 +45,57 @@ python examples/offline_inference/text_to_speech/<model>/end2end.py \
 
 ---
 
+## Audio8 TTS Preview
+
+0.6B DualAR text-to-speech model (Audio8) with its own 44.1 kHz neural audio
+codec: a slow AR transformer predicts one semantic token per codec frame, and a
+4-layer fast AR predicts that frame's remaining 9 codebooks.
+
+### Prerequisites
+No extra packages: the codec architecture ships in tree and its weights
+(`codec.pth`) come with the checkpoint.
+
+### Quick start
+```bash
+python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
+    --text "Welcome to Audio8 TTS, a compact multilingual text to speech model."
+```
+
+### Voice cloning
+```bash
+python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
+    --text "Welcome to Audio8 TTS." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "The exact transcript of the reference recording."
+```
+The reference transcript must match what is actually spoken in the clip;
+a mismatch degrades both stability and speaker similarity.
+
+### Streaming
+```bash
+python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
+    --text "Welcome to Audio8 TTS." --streaming
+```
+
+### Notes
+- Output: 44.1 kHz mono WAV; the codec emits ~21.5 frames/s (2048 samples per frame).
+- Context is 2048 packed text+audio positions, so `max_model_len` is 2048.
+- Greedy decoding (`temperature: 0`) is degenerate for this checkpoint — it
+  ends the utterance almost immediately. Keep the deploy defaults
+  (`temperature 0.7`, `top_k 50`, `top_p 0.9`), which also match
+  `generation_config.json`.
+- Repetition-Aware Sampling is applied to the semantic token: a repeat inside
+  the last 10 frames is re-drawn from a flatter distribution instead of being
+  masked out.
+- `--num-prompts 4` smoke-tests concurrency (stage 0 runs `max_num_seqs: 4`).
+
+### Measured behaviour (1x H20, shared GPU)
+`--streaming` on a 58-character English sentence: TTFA ~0.74 s (4-frame first
+chunk), 7 chunks, 1.55 s wall for 3.16 s of audio (RTF ~0.49). Treat these as
+smoke-test magnitudes, not a benchmark.
+
+---
+
 ## CosyVoice3
 
 2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz.

examples/offline_inference/text_to_speech/audio8_tts/end2end.py

@@ -0,0 +1,237 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Offline inference demo for Audio8 TTS Preview 0.6B via vLLM Omni.
+
+Usage:
+    # Text-only synthesis
+    python end2end.py --text "Welcome to Audio8 TTS."
+
+    # Zero-shot voice cloning (the transcript must match the reference audio)
+    python end2end.py --text "Welcome to Audio8 TTS." \
+        --ref-audio reference.wav --ref-text "The exact transcript of the reference recording."
+
+    # Streaming (per-chunk TTFA / inter-chunk timings)
+    python end2end.py --text "Welcome to Audio8 TTS." --streaming
+"""
+
+import asyncio
+import os
+import time
+
+import numpy as np
+import soundfile as sf
+import torch
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+from vllm_omni import AsyncOmni, Omni
+from vllm_omni.model_executor.models.audio8_tts.codec_utils import estimate_reference_code_frames
+from vllm_omni.model_executor.models.audio8_tts.prompt_utils import (
+    build_text_only_prompt_ids,
+    estimate_voice_clone_prompt_len,
+    normalize_text,
+)
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+
+DEFAULT_MODEL = "Audio8/Audio8-TTS-Preview-0.6b"
+
+
+def build_prompt(
+    text: str,
+    ref_audio_path: str | None = None,
+    ref_text: str | None = None,
+    model_name: str = DEFAULT_MODEL,
+) -> dict:
+    """Build an engine prompt, using the same protocol as online serving."""
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    if ref_audio_path is None and ref_text is None:
+        prompt_ids, normalized_text = build_text_only_prompt_ids(tokenizer, text)
+        return {
+            "prompt_token_ids": prompt_ids,
+            "additional_information": {"text": [normalized_text]},
+        }
+
+    if not ref_audio_path or not ref_text:
+        raise ValueError("Audio8 TTS voice cloning requires both --ref-audio and --ref-text")
+
+    normalized_text = normalize_text(text)
+    normalized_ref_text = normalize_text(ref_text, add_default_speaker=True)
+    ref_audio_wav, ref_audio_sr = sf.read(ref_audio_path, dtype="float32", always_2d=False)
+    if ref_audio_wav.ndim > 1:
+        ref_audio_wav = ref_audio_wav.mean(axis=-1)
+    ref_frames = estimate_reference_code_frames(int(ref_audio_wav.shape[0]), int(ref_audio_sr))
+    placeholder_len = estimate_voice_clone_prompt_len(tokenizer, normalized_text, normalized_ref_text, ref_frames)
+
+    # The clone prompt embeds the reference audio's own codec codes, so the
+    # model builds it in preprocess(); we only reserve the exact length here.
+    return {
+        "prompt_token_ids": [1] * placeholder_len,
+        "additional_information": {
+            "text": normalized_text,
+            "ref_text": normalized_ref_text,
+            "ref_audio_wav": torch.from_numpy(np.asarray(ref_audio_wav, dtype=np.float32)),
+            "ref_audio_sr": int(ref_audio_sr),
+            "audio8_structured_voice_clone": True,
+        },
+    }
+
+
+def extract_audio(multimodal_output: dict) -> tuple[torch.Tensor, int]:
+    """Pull the waveform and sample rate out of a multimodal output.
+
+    The output processor concatenates the per-step delta tensors under
+    ``model_outputs``; ``audio`` is the older alias. Never use ``a or b`` on
+    these values -- truthiness on a tensor raises.
+    """
+    audio = multimodal_output.get("model_outputs")
+    if audio is None:
+        audio = multimodal_output.get("audio")
+    if audio is None:
+        raise ValueError(f"No audio in multimodal_output: {sorted(multimodal_output)}")
+    if isinstance(audio, list):
+        audio = torch.cat([torch.as_tensor(chunk).reshape(-1) for chunk in audio], dim=0)
+    audio = torch.as_tensor(audio).reshape(-1)
+
+    sr_raw = multimodal_output.get("sr")
+    if isinstance(sr_raw, list):
+        sr_raw = sr_raw[-1] if sr_raw else None
+    if sr_raw is None:
+        raise ValueError("Missing sample rate in multimodal_output")
+    sample_rate = int(sr_raw.item()) if hasattr(sr_raw, "item") else int(sr_raw)
+    return audio, sample_rate
+
+
+def _save_wav(output_dir: str, request_id: str, audio: torch.Tensor, sample_rate: int) -> None:
+    out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+    sf.write(out_wav, audio.float().cpu().numpy(), samplerate=sample_rate, format="WAV")
+    duration = audio.numel() / sample_rate
+    # print(), not logging: vLLM reconfigures logging on engine start and
+    # disables loggers created before that, which would swallow these lines.
+    print(f"Request {request_id}: saved {out_wav} (sr={sample_rate}, {duration:.2f}s)")
+
+
+def _omni_kwargs(args, model_name: str) -> dict:
+    kwargs = {
+        "model": model_name,
+        "log_stats": args.log_stats,
+        "stage_init_timeout": args.stage_init_timeout,
+    }
+    if args.deploy_config:
+        kwargs["deploy_config"] = args.deploy_config
+    return kwargs
+
+
+def main(args) -> None:
+    os.makedirs(args.output_dir, exist_ok=True)
+    model_name = args.model or DEFAULT_MODEL
+    prompts = [
+        build_prompt(text, args.ref_audio, args.ref_text, model_name) for text in ([args.text] * args.num_prompts)
+    ]
+
+    omni = Omni(**_omni_kwargs(args, model_name))
+    t_start = time.perf_counter()
+    total_audio_seconds = 0.0
+    for stage_outputs in omni.generate(prompts):
+        request_output = stage_outputs
+        if request_output is None or not request_output.outputs:
+            continue
+        audio, sample_rate = extract_audio(request_output.outputs[0].multimodal_output)
+        assert audio.numel() > 0, "Audio8 TTS produced an empty waveform"
+        total_audio_seconds += audio.numel() / sample_rate
+        _save_wav(args.output_dir, request_output.request_id, audio, sample_rate)
+    elapsed = time.perf_counter() - t_start
+    rtf = elapsed / max(total_audio_seconds, 1e-6)
+    print(f"Total {elapsed * 1000:.1f} ms for {total_audio_seconds:.2f} s of audio (RTF={rtf:.3f})")
+
+
+async def main_streaming(args) -> None:
+    os.makedirs(args.output_dir, exist_ok=True)
+    model_name = args.model or DEFAULT_MODEL
+    prompt = build_prompt(args.text, args.ref_audio, args.ref_text, model_name)
+
+    omni = AsyncOmni(**_omni_kwargs(args, model_name))
+    request_id = "0"
+    chunks: list[torch.Tensor] = []
+    sample_rate = None
+    t_start = time.perf_counter()
+    t_prev = t_start
+    chunk_index = 0
+
+    async for stage_output in omni.generate(prompt, request_id=request_id):
+        multimodal_output = stage_output.outputs[0].multimodal_output
+        if not multimodal_output:
+            continue
+        audio = multimodal_output.get("model_outputs")
+        if audio is None:
+            audio = multimodal_output.get("audio")
+        if audio is not None:
+            if isinstance(audio, list):
+                chunks.extend(torch.as_tensor(chunk).reshape(-1) for chunk in audio)
+            else:
+                chunks.append(torch.as_tensor(audio).reshape(-1))
+        sr_raw = multimodal_output.get("sr")
+        if isinstance(sr_raw, list):
+            sr_raw = sr_raw[-1] if sr_raw else None
+        if sr_raw is not None:
+            sample_rate = int(sr_raw.item()) if hasattr(sr_raw, "item") else int(sr_raw)
+
+        if stage_output.finished:
+            break
+        now = time.perf_counter()
+        if chunk_index == 0:
+            print(f"chunk 0: TTFA={(now - t_start) * 1000:.1f} ms")
+        else:
+            print(f"chunk {chunk_index}: inter_chunk={(now - t_prev) * 1000:.1f} ms")
+        t_prev = now
+        chunk_index += 1
+
+    if not chunks or sample_rate is None:
+        raise RuntimeError("Audio8 TTS streaming produced no audio")
+    _save_wav(args.output_dir, request_id, torch.cat(chunks, dim=0), sample_rate)
+    print(f"Streaming done: total={(time.perf_counter() - t_start) * 1000:.1f} ms chunks={chunk_index}")
+
+
+def parse_args():
+    parser = TrackingArgumentParser(description="Audio8 TTS Preview offline inference with vLLM Omni")
+    parser.add_argument("--model", type=str, default=DEFAULT_MODEL, help="Model path or HuggingFace repo ID.")
+    parser.add_argument(
+        "--text",
+        type=str,
+        default="Welcome to Audio8 TTS, a compact multilingual text to speech model.",
+        help="Text to synthesize.",
+    )
+    parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.")
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcript of --ref-audio; required for voice cloning and must match the recording.",
+    )
+    parser.add_argument("--output-dir", type=str, default="output_audio", help="Directory for output WAV files.")
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Duplicate the prompt N times to smoke-test concurrent requests.",
+    )
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help="Override the deploy config path; defaults to vllm_omni/deploy/audio8_tts.yaml via model_type.",
+    )
+    parser.add_argument("--streaming", action="store_true", help="Use AsyncOmni and report per-chunk latency.")
+    parser.add_argument("--log-stats", action="store_true", help="Enable engine stats logging.")
+    parser.add_argument("--stage-init-timeout", type=int, default=600, help="Per-stage init timeout in seconds.")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    parsed_args = parse_args()
+    if parsed_args.streaming:
+        asyncio.run(main_streaming(parsed_args))
+    else:
+        main(parsed_args)

examples/online_serving/text_to_speech/README.md

@@ -14,6 +14,7 @@ For the full list of supported architectures across all modalities, see
 
 | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
 | --- | --- | --- | --- | --- | --- |
+| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | uploaded audio voice only; no presets | ✓ |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
 | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
 | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — |
@@ -253,6 +254,138 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
 
 ---
 
+## Audio8 TTS Preview
+
+0.6B DualAR TTS at 44.1 kHz, 11 languages, zero-shot voice cloning.
+
+### Prerequisites
+None beyond the base install: the neural audio codec is implemented in tree and
+its weights (`codec.pth`) ship with the checkpoint.
+
+### Launch
+```bash
+vllm serve Audio8/Audio8-TTS-Preview-0.6b --omni --port 8092
+# or:
+./audio8_tts/run_server.sh
+```
+The deploy config auto-loads from `vllm_omni/deploy/audio8_tts.yaml` (HF
+`model_type` is `arktts`). Do **not** pass `--trust-remote-code`: vllm-omni
+registers its own `arktts` config, and transformers would otherwise prefer the
+checkpoint's remote code and bypass it.
+
+### Text-only synthesis
+```bash
+curl -X POST http://localhost:8092/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Welcome to Audio8 TTS.",
+        "response_format": "wav"
+    }' --output output.wav
+```
+
+### Voice cloning
+```bash
+curl -X POST http://localhost:8092/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Welcome to Audio8 TTS.",
+        "ref_audio": "https://example.com/reference.wav",
+        "ref_text": "The exact transcript of the reference recording."
+    }' --output cloned.wav
+```
+`ref_audio` also accepts a `data:audio/wav;base64,...` URL. `ref_text` is
+mandatory whenever `ref_audio` is present and must match the recording.
+Uploading a voice via `POST /v1/audio/voices` lets the server reuse the encoded
+reference codes across requests (`voice: "<name>"`).
+
+### Streaming
+```bash
+python audio8_tts/speech_client.py --text "Welcome to Audio8 TTS." --stream --output out.pcm
+ffplay -f s16le -ar 44100 -ac 1 out.pcm
+```
+
+### Gradio demo
+```bash
+./audio8_tts/run_gradio_demo.sh          # server + demo
+python audio8_tts/gradio_demo.py --api-base http://localhost:8092   # demo only
+```
+
+### Benchmark
+
+```bash
+export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100   # required: harness defaults to 24 kHz
+vllm bench serve --omni --host 127.0.0.1 --port 8092 \
+    --model Audio8/Audio8-TTS-Preview-0.6b \
+    --backend openai-audio-speech --endpoint /v1/audio/speech \
+    --dataset-name seed-tts-text --dataset-path benchmarks/build_dataset/seed_tts_smoke \
+    --seed-tts-locale en --num-prompts 20 --num-warmups 2 \
+    --max-concurrency 1 --request-rate inf \
+    --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun
+```
+
+Measured on 1x H20 **shared with an unrelated job holding ~80% SM**, so treat
+these as a lower bound. 20 prompts from `seed_tts_smoke/en`, mean audio 4.5 s,
+deploy defaults:
+
+| | HF reference (batch=1) | vllm-omni c=1 | c=4 | c=8 |
+|---|---|---|---|---|
+| RTF (mean) | 1.008 | **0.19** | 0.30 | 0.54 |
+| Throughput (req/s) | 0.23 | 1.18 | 2.84 | **2.93** |
+| E2E latency (mean, ms) | 4375 | 850 | 1349 | 2361 |
+| Time to first packet (mean, ms) | n/a (no streaming) | 63 | 105 | 1124 |
+| Underrun p99 / continuity OK | n/a | 0 s / 100% | 0 s / 100% | 0 s / 100% |
+
+~5.3x lower per-request RTF than the upstream reference implementation at
+concurrency 1, ~12.8x aggregate throughput at concurrency 8.
+
+#### Stage-0 `max_num_seqs`: first-packet vs throughput
+
+This subsection is operator tuning guidance, not a claim this PR optimizes
+anything: the shipped deploy default is unchanged (`max_num_seqs: 4`). The
+numbers below are a single run on the same contended H20 as above (one unrelated
+job holding ~80% SM, no repeats), so read them as directional, not as measured
+speedups. Same setup, varying stage 0's `max_num_seqs` (stage 1 stays at 1):
+
+| stage-0 `max_num_seqs` | client concurrency | req/s | RTF | TTFP mean / p99 (ms) | underrun p99 (s) |
+|---|---|---|---|---|---|
+| 4 (default) | 1 | 1.18 | 0.19 | 63 / 67 | 0.00 |
+| 4 | 4 | 2.84 | 0.30 | 105 / 148 | 0.00 |
+| 4 | 8 | 2.93 | 0.54 | 1124 / 1642 | 0.00 |
+| 8 | 1 | 1.16 | 0.19 | 64 / 70 | 0.00 |
+| 8 | 4 | 2.88 | 0.30 | 107 / 148 | 0.00 |
+| 8 | 8 | **3.72** | 0.40 | **247** / 733 | 0.22 |
+| 8 | 16 | 3.91 | 0.65 | 1366 / 2530 | 0.29 |
+
+Reading of this:
+
+- Below the limit (`concurrency <= max_num_seqs`) the setting is irrelevant, as
+  expected — c=1 and c=4 are identical for both.
+- At c=8, raising the limit to 8 admits every request instead of queueing half
+  of them: in this single run, roughly +27% throughput and ~4.5x lower
+  first-packet latency.
+- Throughput saturates around 3.7-3.9 req/s (~17 s of audio per second). Past
+  c=8 you only buy queueing delay.
+- The cost of admitting more streams is a 0.22-0.29 s worst-case buffer deficit
+  in 1 of 20 requests (`Streaming continuity OK rate` 95%), because stage 1 runs
+  `max_num_seqs: 1` and the codec window round-robins across streams. Client-side
+  prebuffering (the bundled Gradio demo holds 2 chunks) absorbs that; raising
+  stage 1's `max_num_seqs` is the untested knob if you need both.
+
+The shipped default stays at 4 because every measured point there had zero
+underrun. Raise it to 8 if first-packet latency under load matters more than a
+rare sub-300 ms gap.
+
+### Notes
+- Output: 44.1 kHz mono; ~21.5 codec frames per second.
+- No built-in speaker presets. Omit `voice` for a random timbre, or clone one.
+- `max_new_tokens` caps generated codec frames (1 frame ~= 46 ms).
+- Context is 2048 packed text+audio positions (`max_model_len: 2048`).
+- Do not send `temperature: 0`: greedy decoding is degenerate for this
+  checkpoint and truncates the utterance. The deploy defaults
+  (0.7 / top_k 50 / top_p 0.9) mirror `generation_config.json`.
+
+---
+
 ## Fish Speech S2 Pro
 
 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec, which is vendored in

examples/online_serving/text_to_speech/audio8_tts/gradio_demo.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Gradio demo for Audio8 TTS Preview via the /v1/audio/speech API.
+
+Supports:
+  - text-to-speech synthesis in 11 languages
+  - zero-shot voice cloning from reference audio (upload, microphone, or URL)
+  - streaming (progressive PCM chunks) and non-streaming modes
+
+Usage:
+    # Start the server first (see run_server.sh), then:
+    python gradio_demo.py --api-base http://localhost:8092
+"""
+
+import argparse
+import base64
+import io
+
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import httpx
+import numpy as np
+import soundfile as sf
+
+#: Audio8 TTS decodes at 44.1 kHz; raw PCM must be interpreted at that rate.
+PCM_SAMPLE_RATE = 44100
+DEFAULT_API_BASE = "http://localhost:8092"
+
+
+def encode_audio_to_base64(audio_data: tuple) -> str:
+    """Encode a Gradio ``(sample_rate, numpy_array)`` input as a data URL."""
+    sample_rate, audio_np = audio_data
+    if audio_np.dtype != np.int16:
+        if audio_np.dtype in (np.float32, np.float64):
+            audio_np = (np.clip(audio_np, -1.0, 1.0) * 32767).astype(np.int16)
+        else:
+            audio_np = audio_np.astype(np.int16)
+    buf = io.BytesIO()
+    sf.write(buf, audio_np, sample_rate, format="WAV")
+    return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}"
+
+
+def build_payload(
+    text: str,
+    ref_audio: tuple | None,
+    ref_audio_url: str,
+    ref_text: str,
+    response_format: str = "wav",
+    stream: bool = False,
+) -> dict:
+    """Build the ``/v1/audio/speech`` payload.
+
+    Audio8 TTS has no built-in speakers: omit ``voice`` for a random timbre, or
+    supply ``ref_audio`` + ``ref_text`` to clone a voice.
+    """
+    if not text or not text.strip():
+        raise gr.Error("Please enter text to synthesize.")
+
+    payload: dict = {
+        "input": text.strip(),
+        "response_format": "pcm" if stream else response_format,
+        "stream": stream,
+    }
+    if stream:
+        payload["stream_format"] = "audio"
+
+    ref_url = ref_audio_url.strip() if ref_audio_url else ""
+    if ref_audio is not None:
+        payload["ref_audio"] = encode_audio_to_base64(ref_audio)
+    elif ref_url:
+        payload["ref_audio"] = ref_url
+
+    if "ref_audio" in payload:
+        if not ref_text or not ref_text.strip():
+            raise gr.Error("Voice cloning requires the transcript of the reference audio.")
+        payload["ref_text"] = ref_text.strip()
+
+    return payload
+
+
+def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format):
+    """Non-streaming request; returns the whole waveform."""
+    payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=False)
+    try:
+        with httpx.Client(timeout=300.0) as client:
+            resp = client.post(
+                f"{api_base}/v1/audio/speech",
+                json=payload,
+                headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
+            )
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.") from None
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") from None
+
+    if resp.status_code != 200:
+        raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+    if "application/json" in resp.headers.get("content-type", ""):
+        raise gr.Error(f"Server error: {resp.text}")
+
+    try:
+        if response_format == "pcm":
+            audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0
+            return (PCM_SAMPLE_RATE, audio_np)
+        audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
+        if audio_np.ndim > 1:
+            audio_np = audio_np[:, 0]
+        return (sample_rate, audio_np.astype(np.float32))
+    except Exception as e:
+        raise gr.Error(f"Failed to decode audio: {e}") from None
+
+
+def generate_speech_stream(api_base, text, ref_audio, ref_audio_url, ref_text, response_format, chunk_seconds=0.5):
+    """Streaming request; yields PCM chunks for progressive playback."""
+    payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=True)
+    min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds)
+    prebuffer_chunks = 2
+    pending: list[np.ndarray] = []
+    pending_len = 0
+    chunks_yielded = 0
+    prebuffer: list[np.ndarray] = []
+    leftover = b""
+
+    try:
+        with (
+            httpx.Client(timeout=300.0) as client,
+            client.stream(
+                "POST",
+                f"{api_base}/v1/audio/speech",
+                json=payload,
+                headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
+            ) as resp,
+        ):
+            if resp.status_code != 200:
+                resp.read()
+                raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+            for chunk in resp.iter_bytes():
+                if not chunk:
+                    continue
+                raw = leftover + chunk
+                # s16le: never split a sample across chunk boundaries.
+                usable = len(raw) - (len(raw) % 2)
+                leftover = raw[usable:]
+                if usable == 0:
+                    continue
+                samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0
+                pending.append(samples)
+                pending_len += len(samples)
+                if pending_len < min_chunk_samples:
+                    continue
+                audio_chunk = np.concatenate(pending)
+                pending.clear()
+                pending_len = 0
+                if chunks_yielded < prebuffer_chunks:
+                    prebuffer.append(audio_chunk)
+                    chunks_yielded += 1
+                    if chunks_yielded == prebuffer_chunks:
+                        yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer))
+                        prebuffer.clear()
+                else:
+                    yield (PCM_SAMPLE_RATE, audio_chunk)
+        remaining = prebuffer + pending
+        if remaining:
+            yield (PCM_SAMPLE_RATE, np.concatenate(remaining))
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.") from None
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") from None
+
+
+def on_stream_change(stream: bool):
+    """Streaming only makes sense with raw PCM."""
+    if stream:
+        return gr.update(value="pcm", interactive=False)
+    return gr.update(value="wav", interactive=True)
+
+
+def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
+    with gr.Blocks(title="Audio8 TTS Preview Demo") as demo:
+        gr.Markdown("# Audio8 TTS Preview 0.6B - Text to Speech")
+        gr.Markdown(f"**Server:** `{api_base}` | **Model:** Audio8/Audio8-TTS-Preview-0.6b | **Output:** 44.1 kHz")
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                text_input = gr.Textbox(label="Text to Synthesize", placeholder="Enter text here...", lines=4)
+
+                with gr.Accordion("Voice Cloning (optional)", open=False):
+                    gr.Markdown(
+                        "Upload or link a short reference recording (5-30 s) and paste its **exact** transcript. "
+                        "A transcript that does not match the audio degrades both stability and speaker similarity."
+                    )
+                    ref_audio = gr.Audio(label="Reference Audio", type="numpy", sources=["upload", "microphone"])
+                    ref_audio_url = gr.Textbox(
+                        label="Reference Audio URL (alternative to upload)",
+                        placeholder="https://example.com/reference.wav",
+                        lines=1,
+                    )
+                    ref_text = gr.Textbox(
+                        label="Reference Audio Transcript (required for cloning)",
+                        placeholder="Exact transcript of the reference audio...",
+                        lines=2,
+                    )
+
+                with gr.Row():
+                    response_format = gr.Dropdown(
+                        choices=["wav", "mp3", "flac", "pcm"], value="wav", label="Audio Format", scale=1
+                    )
+                    stream_checkbox = gr.Checkbox(
+                        label="Stream output", value=False, info="Progressive PCM streaming", scale=1
+                    )
+
+                generate_btn = gr.Button("Generate Speech", variant="primary", size="lg")
+
+            with gr.Column(scale=2):
+                audio_output = gr.Audio(label="Generated Audio", interactive=False, streaming=True, autoplay=True)
+                gr.Markdown(
+                    "### About\n"
+                    "- **Audio8 TTS Preview** — 0.6B DualAR model (slow AR + fast AR codebook decoder)\n"
+                    "- **Languages**: Cantonese, Chinese, Dutch, English, French, German, "
+                    "Italian, Japanese, Korean, Polish, Spanish\n"
+                    "- **Voice cloning**: zero-shot from a reference clip + transcript\n"
+                    "- **44.1 kHz** output, ~21.5 codec frames per second"
+                )
+
+        stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format])
+        all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format]
+
+        def dispatch(stream_enabled, *args):
+            if stream_enabled:
+                yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds)
+            else:
+                yield generate_speech(api_base, *args)
+
+        generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output])
+        demo.queue()
+    return demo
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Gradio demo for Audio8 TTS Preview")
+    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})")
+    parser.add_argument("--host", default="0.0.0.0", help="Gradio host")
+    parser.add_argument("--port", type=int, default=7861, help="Gradio port")
+    parser.add_argument("--share", action="store_true", help="Create a public share link")
+    parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds of audio per yielded chunk")
+    args = parser.parse_args()
+
+    print(f"Connecting to vLLM server at: {args.api_base}")
+    build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds).launch(
+        server_name=args.host, server_port=args.port, share=args.share
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/text_to_speech/audio8_tts/run_gradio_demo.sh

@@ -0,0 +1,44 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+# Launch the Audio8 TTS Preview server + Gradio demo together.
+#
+# Usage:
+#   ./run_gradio_demo.sh
+#   CUDA_VISIBLE_DEVICES=0 PORT=8092 GRADIO_PORT=7861 ./run_gradio_demo.sh
+
+set -e
+
+MODEL="${MODEL:-Audio8/Audio8-TTS-Preview-0.6b}"
+PORT="${PORT:-8092}"
+GRADIO_PORT="${GRADIO_PORT:-7861}"
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+
+echo "Starting Audio8 TTS Preview server (port $PORT)..."
+vllm serve "$MODEL" \
+    --omni \
+    --host 0.0.0.0 \
+    --port "$PORT" &
+SERVER_PID=$!
+
+cleanup() {
+    echo "Stopping server (PID $SERVER_PID)..."
+    kill $SERVER_PID 2>/dev/null
+    wait $SERVER_PID 2>/dev/null
+}
+trap cleanup EXIT
+
+echo "Waiting for server to start..."
+for _ in $(seq 1 120); do
+    if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
+        echo "Server ready."
+        break
+    fi
+    sleep 2
+done
+
+echo "Starting Gradio demo (port $GRADIO_PORT)..."
+python "$SCRIPT_DIR/gradio_demo.py" \
+    --api-base "http://localhost:$PORT" \
+    --port "$GRADIO_PORT"

examples/online_serving/text_to_speech/audio8_tts/run_server.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+# Launch a vLLM-Omni server for Audio8 TTS Preview 0.6B.
+#
+# Usage:
+#   ./run_server.sh
+#   CUDA_VISIBLE_DEVICES=0 PORT=8092 ./run_server.sh
+
+set -e
+
+MODEL="${MODEL:-Audio8/Audio8-TTS-Preview-0.6b}"
+PORT="${PORT:-8092}"
+
+echo "Starting Audio8 TTS Preview server with model: $MODEL"
+
+# --trust-remote-code is deliberately NOT passed: vllm-omni registers its own
+# `arktts` config, and transformers would otherwise prefer the checkpoint's
+# remote code and bypass it.
+vllm serve "$MODEL" \
+    --omni \
+    --host 0.0.0.0 \
+    --port "$PORT"

examples/online_serving/text_to_speech/audio8_tts/speech_client.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Client for Audio8 TTS Preview via the /v1/audio/speech endpoint.
+
+Examples:
+    # Basic TTS
+    python speech_client.py --text "Welcome to Audio8 TTS."
+
+    # Zero-shot voice cloning (the transcript must match the reference audio)
+    python speech_client.py --text "Welcome to Audio8 TTS." \
+        --ref-audio ref.wav --ref-text "The exact transcript of the reference recording."
+
+    # Streaming PCM output
+    python speech_client.py --text "Welcome to Audio8 TTS." --stream --output output.pcm
+"""
+
+import argparse
+import base64
+import os
+
+import httpx
+
+DEFAULT_API_BASE = "http://localhost:8092"
+DEFAULT_API_KEY = "EMPTY"
+DEFAULT_MODEL = "Audio8/Audio8-TTS-Preview-0.6b"
+#: The codec runs at 44.1 kHz, so raw PCM must be played back at that rate.
+PCM_SAMPLE_RATE = 44100
+
+
+def encode_audio_to_base64(audio_path: str) -> str:
+    """Encode a local audio file as a base64 data URL."""
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+    ext = audio_path.lower().rsplit(".", 1)[-1]
+    mime_type = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"}.get(
+        ext, "audio/wav"
+    )
+    with open(audio_path, "rb") as handle:
+        audio_b64 = base64.b64encode(handle.read()).decode("utf-8")
+    return f"data:{mime_type};base64,{audio_b64}"
+
+
+def build_payload(args) -> dict:
+    payload = {
+        "model": args.model,
+        "input": args.text,
+        "response_format": args.response_format,
+    }
+    if args.ref_audio:
+        if args.ref_audio.startswith(("http://", "https://")):
+            payload["ref_audio"] = args.ref_audio
+        else:
+            payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
+    if args.ref_text:
+        payload["ref_text"] = args.ref_text
+    if args.max_new_tokens is not None:
+        payload["max_new_tokens"] = args.max_new_tokens
+    if args.stream:
+        payload["stream"] = True
+        payload["stream_format"] = "audio"
+        payload["response_format"] = "pcm"
+    return payload
+
+
+def run_tts(args) -> None:
+    payload = build_payload(args)
+    api_url = f"{args.api_base}/v1/audio/speech"
+    headers = {"Content-Type": "application/json", "Authorization": f"Bearer {args.api_key}"}
+
+    print(f"Model: {args.model}")
+    print(f"Text: {args.text}")
+    if args.ref_audio:
+        print(f"Voice cloning: ref_audio={args.ref_audio!r} ref_text={args.ref_text!r}")
+
+    if args.stream:
+        output_path = args.output or "output.pcm"
+        with (
+            httpx.Client(timeout=300.0) as client,
+            client.stream("POST", api_url, json=payload, headers=headers) as resp,
+        ):
+            if resp.status_code != 200:
+                print(f"Error {resp.status_code}: {resp.read().decode()}")
+                return
+            total_bytes = 0
+            with open(output_path, "wb") as handle:
+                for chunk in resp.iter_bytes():
+                    handle.write(chunk)
+                    total_bytes += len(chunk)
+        seconds = total_bytes / 2 / PCM_SAMPLE_RATE
+        print(f"Streamed {total_bytes} bytes (~{seconds:.2f}s of s16le @ {PCM_SAMPLE_RATE} Hz) to {output_path}")
+        print(f"Play with: ffplay -f s16le -ar {PCM_SAMPLE_RATE} -ac 1 {output_path}")
+        return
+
+    with httpx.Client(timeout=300.0) as client:
+        response = client.post(api_url, json=payload, headers=headers)
+    if response.status_code != 200:
+        print(f"Error {response.status_code}: {response.text}")
+        return
+    if response.headers.get("content-type", "").startswith("application/json"):
+        print(f"Error: {response.text}")
+        return
+
+    output_path = args.output or f"output.{args.response_format}"
+    with open(output_path, "wb") as handle:
+        handle.write(response.content)
+    print(f"Audio saved to: {output_path} ({len(response.content)} bytes)")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Audio8 TTS Preview client")
+    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
+    parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
+    parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name")
+    parser.add_argument("--text", required=True, help="Text to synthesize")
+    parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning (path or URL)")
+    parser.add_argument("--ref-text", default=None, help="Transcript of the reference audio")
+    parser.add_argument("--max-new-tokens", type=int, default=None, help="Cap the number of generated codec frames")
+    parser.add_argument("--stream", action="store_true", help="Stream PCM instead of returning a whole file")
+    parser.add_argument(
+        "--response-format",
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
+        help="Audio format (default: wav)",
+    )
+    parser.add_argument("--output", "-o", default=None, help="Output file path")
+    run_tts(parser.parse_args())
+
+
+if __name__ == "__main__":
+    main()

tests/e2e/offline_inference/test_audio8_tts.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""E2E offline tests for Audio8 TTS Preview (DualAR, 44.1 kHz codec).
+
+The prompt is pre-tokenized here (not built from ``mm_processor_kwargs``)
+because Audio8 TTS shares Fish Speech's protocol: text-only prompts are plain
+token ids, and voice-clone prompts are a placeholder of the exact final length
+whose real embeddings are built model-side from the encoded reference audio.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner, OmniRunnerHandler
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = os.environ.get("AUDIO8_TTS_MODEL_PATH", "Audio8/Audio8-TTS-Preview-0.6b")
+DEPLOY_CONFIG = get_deploy_config_path("audio8_tts.yaml")
+SAMPLE_RATE = 44100
+TEXT = "The weather is nice today, perfect for a walk in the park."
+
+# Four distinct inputs for the concurrency test: distinct text is required so a
+# leaked per-request codec buffer (which would make two requests decode to the
+# same audio) actually fails the pairwise-difference assertion.
+CONCURRENT_TEXTS = [
+    "The weather is nice today, perfect for a walk in the park.",
+    "She sold seashells by the seashore all through the summer.",
+    "Our train departs at a quarter past nine tomorrow morning.",
+    "He planted rows of tomatoes and basil behind the old house.",
+]
+
+# The checkpoint registers its own `arktts` config in vllm-omni, so remote code
+# must stay disabled (transformers would otherwise prefer the checkpoint's
+# auto_map and bypass the Qwen2-shaped view Qwen2Model needs).
+_OMNI_RUNNER_PARAM = (MODEL, DEPLOY_CONFIG, {"trust_remote_code": False})
+
+pytestmark = pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
+
+
+def _text_only_prompt(text: str = TEXT) -> dict:
+    from transformers import AutoTokenizer
+
+    from vllm_omni.model_executor.models.audio8_tts.prompt_utils import build_text_only_prompt_ids
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL)
+    prompt_ids, normalized_text = build_text_only_prompt_ids(tokenizer, text)
+    return {"prompt_token_ids": prompt_ids, "additional_information": {"text": [normalized_text]}}
+
+
+# At ``core_model`` (L1/L2) the harness patches every stage to
+# ``load_format: dummy``, so the AR runs on random weights: it never emits EOS
+# and always generates to ``max_tokens`` (1024 frames ~= 47.5 s). Duration
+# bounds are therefore only meaningful once real weights load, at
+# ``advanced_model`` / ``full_model``.
+_REAL_WEIGHT_RUN_LEVELS = frozenset({"advanced_model", "full_model"})
+
+
+def _duration_bounds(run_level: str) -> dict[str, float]:
+    if run_level not in _REAL_WEIGHT_RUN_LEVELS:
+        return {}
+    return {"min_duration_s": 0.5, "max_duration_s": 30.0}
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_text_to_audio_001(omni_runner: OmniRunner, run_level: str) -> None:
+    """Default deploy, single request, text-only synthesis.
+
+    Deploy Setting: audio8_tts.yaml
+    Input Modal: text
+    Output Modal: audio
+    """
+    prompt = _text_only_prompt()
+    request_config = {
+        "input": TEXT,
+        "prompt_token_ids": prompt["prompt_token_ids"],
+        "additional_information": prompt["additional_information"],
+        "response_format": "wav",
+        "expected_sample_rate": SAMPLE_RATE,
+        **_duration_bounds(run_level),
+    }
+    OmniRunnerHandler(omni_runner).send_tokenized_tts_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_text_to_audio_concurrent_002(omni_runner: OmniRunner, run_level: str) -> None:
+    """Four concurrent requests with distinct text: stage 0 runs ``max_num_seqs: 4``.
+
+    Per-request codec state is keyed by request id. The four inputs are
+    deliberately distinct and the decoded outputs are asserted pairwise
+    different, so leaked codec state (which would make two requests decode to
+    identical audio) fails here, and only here.
+    """
+    prompts = [_text_only_prompt(text) for text in CONCURRENT_TEXTS]
+    request_config = {
+        "input": CONCURRENT_TEXTS[0],
+        "prompt_token_ids": [p["prompt_token_ids"] for p in prompts],
+        "additional_information": [p["additional_information"] for p in prompts],
+        "response_format": "wav",
+        "expected_sample_rate": SAMPLE_RATE,
+        "assert_distinct_outputs": True,
+        **_duration_bounds(run_level),
+    }
+    OmniRunnerHandler(omni_runner).send_tokenized_tts_request(request_config)

tests/e2e/online_serving/test_audio8_tts.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""E2E online tests for Audio8 TTS Preview via /v1/audio/speech.
+
+Real model inference (no mocks): text-only synthesis, streaming, and zero-shot
+voice cloning. The reference audio is inlined as a base64 data URL because CI
+hosts cannot fetch external URLs.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = os.environ.get("AUDIO8_TTS_MODEL_PATH", "Audio8/Audio8-TTS-Preview-0.6b")
+DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0
+# The codec decodes at 44.1 kHz; raw PCM carries no header, so assertions that
+# depend on the rate (HNR pitch search) have to be told explicitly.
+SAMPLE_RATE = 44100
+MAX_CONCURRENT = 4
+
+# ~0.5 s of 44.1 kHz mono PCM_16 WAV.
+_MIN_AUDIO_BYTES = 40_000
+
+# Reused vendored asset (3.5 s, Chinese) plus its exact transcript: Audio8 TTS
+# requires the reference transcript to match the recording.
+REF_AUDIO_URL = get_asset_path("cosyvoice3/zero_shot_prompt.wav", as_data_url=True)
+REF_TEXT = "希望你以后能够做的比我还好呦。"
+
+
+def get_prompt(prompt_type: str = "en") -> str:
+    prompts = {
+        "en": "The weather is nice today, perfect for a walk in the park.",
+        "zh": "今天天气很好，非常适合去公园散步。",
+    }
+    return prompts.get(prompt_type, prompts["en"])
+
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("audio8_tts.yaml"),
+            # No --trust-remote-code: vllm-omni owns the `arktts` config.
+            server_args=["--disable-log-stats"],
+        ),
+        id="audio8_tts",
+    )
+]
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_001(omni_server, openai_client) -> None:
+    """Baseline smoke: default deploy, non-streaming WAV.
+
+    Deploy Setting: audio8_tts.yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=False
+    Datasets: few requests (max_num_seqs=4)
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": False,
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "response_format": "wav",
+        "min_audio_bytes": _MIN_AUDIO_BYTES,
+    }
+    openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_streaming_002(omni_server, openai_client) -> None:
+    """Streaming PCM: exercises the async_chunk path end to end.
+
+    Deploy Setting: audio8_tts.yaml (async_chunk: true)
+    Input Setting: stream=True, response_format=pcm
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": True,
+        "stream_format": "audio",
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "response_format": "pcm",
+        "min_audio_bytes": _MIN_AUDIO_BYTES,
+        # Raw PCM has no header: without this the harness assumes 24 kHz and the
+        # HNR pitch search lands on 147-735 Hz instead of 80-400 Hz, scoring
+        # clean 44.1 kHz speech ~1.9 dB too low.
+        "expected_sample_rate": SAMPLE_RATE,
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_voice_clone_003(omni_server, openai_client) -> None:
+    """Zero-shot voice cloning: the reference audio is encoded model-side and
+    its codec codes are spliced into the prompt.
+
+    Input Modal: text + ref_audio + ref_text
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt("zh"),
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+        "stream": False,
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "response_format": "wav",
+        "min_audio_bytes": _MIN_AUDIO_BYTES,
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_voice_clone_requires_ref_text_004(omni_server, openai_client) -> None:
+    """``ref_audio`` without ``ref_text`` must be rejected, not silently
+    synthesized with an untethered voice."""
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "ref_audio": REF_AUDIO_URL,
+        "stream": False,
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "response_format": "wav",
+        "status_code": (400, 422),
+        "err_message": "ref_text",
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_empty_input_is_rejected_005(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "input": "   ",
+        "stream": False,
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "response_format": "wav",
+        "status_code": (400, 422),
+        "err_message": "Input text cannot be empty",
+    }
+    openai_client.send_audio_speech_request(request_config)

tests/entrypoints/openai/test_audio8_tts_adapter.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Serving-side registration invariants for the Audio8 TTS adapter.
+
+The engine never names ``audio8_tts`` directly: ``/v1/audio/speech`` resolves the
+adapter by matching the deployed stage's ``model_stage`` key against the
+``stage_keys`` each adapter declares. That makes the pipeline definition and the
+adapter two halves of one contract, with nothing in the type system holding them
+together -- rename either side and the endpoint silently stops recognising the
+model (falling through to "unknown TTS model") instead of failing at startup.
+
+These tests pin the contract from both ends.
+"""
+
+import pytest
+
+from vllm_omni.entrypoints.openai.tts_adapters import (
+    all_tts_model_types,
+    all_tts_stage_keys,
+    detect_tts_model_type,
+    resolve_adapter,
+)
+from vllm_omni.entrypoints.openai.tts_adapters.audio8_tts import Audio8TTSAdapter
+from vllm_omni.model_executor.models.audio8_tts.pipeline import AUDIO8_TTS_PIPELINE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+SLOW_AR_STAGE_KEY = "audio8_tts_slow_ar"
+
+
+def _pipeline_stage_keys() -> list[str]:
+    return [stage.model_stage for stage in AUDIO8_TTS_PIPELINE.stages]
+
+
+def test_pipeline_stage_key_is_the_one_the_adapter_claims():
+    """The shipped pipeline must name the stage the adapter matches on."""
+    assert SLOW_AR_STAGE_KEY in _pipeline_stage_keys()
+    assert SLOW_AR_STAGE_KEY in Audio8TTSAdapter.stage_keys
+
+
+def test_detection_resolves_slow_ar_stage_to_the_audio8_adapter():
+    assert detect_tts_model_type(SLOW_AR_STAGE_KEY, "Audio8TTSSlowARForConditionalGeneration") == "audio8_tts"
+    assert resolve_adapter("audio8_tts") is Audio8TTSAdapter
+
+
+def test_adapter_is_registered_in_the_global_tables():
+    """Importing the package must register the adapter, not just define it."""
+    assert "audio8_tts" in all_tts_model_types()
+    assert SLOW_AR_STAGE_KEY in all_tts_stage_keys()
+
+
+def test_codec_decoder_stage_does_not_serve_speech():
+    """Only stage 0 answers /v1/audio/speech; the codec stage must not match."""
+    codec_stage_key = _pipeline_stage_keys()[1]
+    assert codec_stage_key != SLOW_AR_STAGE_KEY
+    assert detect_tts_model_type(codec_stage_key, "Audio8TTSCodecDecoder") != "audio8_tts"
+
+
+def test_apply_sampling_overrides_defaults_tts_local_seed_for_fast_ar():
+    """The deploy seed must reach the Fast AR through ``tts_local_seed``.
+
+    Only the Slow AR is seeded via ``sampling_metadata.generators``; the residual
+    codebooks read ``tts_local_seed`` and otherwise fall back to global RNG, which
+    makes an unseeded request's audio depend on batch composition. The adapter
+    defaults it from the stage-0 seed without clobbering an explicit value.
+    """
+    from types import SimpleNamespace
+
+    from vllm.sampling_params import SamplingParams
+
+    adapter = Audio8TTSAdapter.__new__(Audio8TTSAdapter)
+    request = SimpleNamespace(max_new_tokens=None)
+
+    # Deploy default seed propagates to tts_local_seed.
+    out = adapter.apply_sampling_overrides([SamplingParams(seed=42)], request)
+    assert out[0].extra_args["tts_local_seed"] == 42
+
+    # An explicit tts_local_seed is left intact (setdefault, not overwrite).
+    pinned = [SamplingParams(seed=42, extra_args={"tts_local_seed": 7})]
+    assert adapter.apply_sampling_overrides(pinned, request)[0].extra_args["tts_local_seed"] == 7
+
+    # No seed at all -> nothing injected (non-deterministic by request contract).
+    out_unseeded = adapter.apply_sampling_overrides([SamplingParams(seed=None)], request)
+    assert (out_unseeded[0].extra_args or {}).get("tts_local_seed") is None

tests/entrypoints/openai_api/test_tts_detection.py

@@ -132,6 +132,8 @@ _PIPELINE_STAGES = [
     "audex_thinker",
     "audex_tta_thinker",
     "audex_xcodec",
+    "audio8_tts_codec_decoder",
+    "audio8_tts_slow_ar",
     "audio_generation",
     "audio_tokenizer",
     "audio_vae",

tests/helpers/assertions.py

@@ -659,7 +659,11 @@ def _compute_pcm_hnr_db(pcm_samples: np.ndarray, sr: int = _PCM_SPEECH_SAMPLE_RA
     return float(np.mean(hnr_values)) if hnr_values else 0.0
 
 
-def _assert_pcm_int16_speech_hnr(audio_bytes: bytes, min_hnr_db: float = _MIN_PCM_SPEECH_HNR_DB) -> None:
+def _assert_pcm_int16_speech_hnr(
+    audio_bytes: bytes,
+    min_hnr_db: float = _MIN_PCM_SPEECH_HNR_DB,
+    sr: int = _PCM_SPEECH_SAMPLE_RATE_HZ,
+) -> None:
     """Validate harmonic-to-noise ratio on raw int16 PCM from /v1/audio/speech.
 
     min_hnr_db defaults to the global _MIN_PCM_SPEECH_HNR_DB (1.0 dB),
@@ -668,12 +672,16 @@ def _assert_pcm_int16_speech_hnr(audio_bytes: bytes, min_hnr_db: float = _MIN_PC
     intrinsically around -2 dB) can pass a lower per-test threshold via
     request_config["min_hnr_db"] to keep the catastrophic-failure check
     while not gating CI on a model-intrinsic property.
+
+    ``sr`` must be the stream's real sample rate: the autocorrelation lag window
+    is derived from it, so a wrong rate detunes the 80-400 Hz pitch search and
+    understates HNR for models that do not decode at 24 kHz.
     """
     assert audio_bytes is not None and len(audio_bytes) >= 2, "missing PCM bytes"
     assert len(audio_bytes) % 2 == 0, "PCM byte length must be aligned to int16"
     pcm_samples = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0
-    hnr = _compute_pcm_hnr_db(pcm_samples)
-    print(f"PCM speech HNR: {hnr:.2f} dB (threshold: {min_hnr_db} dB)")
+    hnr = _compute_pcm_hnr_db(pcm_samples, sr=sr)
+    print(f"PCM speech HNR: {hnr:.2f} dB (threshold: {min_hnr_db} dB, sr={sr})")
     assert hnr >= min_hnr_db, (
         f"Audio distortion detected: HNR={hnr:.2f} dB < {min_hnr_db} dB. "
         "Voice clone decoder may be losing ref_code speaker context on later chunks."
@@ -934,7 +942,15 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
     if run_level in {"advanced_model", "full_model"}:
         if req_fmt == "pcm" and response.audio_bytes:
             min_hnr_db = float(request_config.get("min_hnr_db", _MIN_PCM_SPEECH_HNR_DB))
-            _assert_pcm_int16_speech_hnr(response.audio_bytes, min_hnr_db=min_hnr_db)
+            # Raw PCM carries no header, so the HNR pitch search has to be told
+            # the rate. Defaulting to 24 kHz for a 44.1 kHz model shifts the
+            # search window to 147-735 Hz and misses the fundamental entirely,
+            # scoring clean speech ~1.9 dB lower than it is.
+            _assert_pcm_int16_speech_hnr(
+                response.audio_bytes,
+                min_hnr_db=min_hnr_db,
+                sr=int(request_config.get("expected_sample_rate") or _PCM_SPEECH_SAMPLE_RATE_HZ),
+            )
 
         transcript = _resolve_audio_transcript(response, request_config, run_level, speech_api=True)
         if transcript is not None:

tests/helpers/client.py

@@ -12,6 +12,7 @@ import copy
 import io
 import json
 import time
+from collections.abc import Mapping
 from dataclasses import asdict, dataclass
 from io import BytesIO
 from pathlib import Path
@@ -2123,6 +2124,123 @@ class OfflineOmniClient:
         assert_audio_speech_response(result, request_config, run_level="core_model")
         return result
 
+    def send_tokenized_tts_request(self, request_config: dict[str, Any]) -> OmniResponse:
+        """Offline TTS for models whose engine prompt is *pre-tokenized*.
+
+        Fish-Speech-style DualAR models (and any model whose prompt embeds codec
+        codes) cannot go through :meth:`send_audio_speech_request`: the prompt is
+        built by the caller as ``prompt_token_ids`` plus per-request
+        ``additional_information``, not from ``mm_processor_kwargs``.
+
+        request_config keys:
+          - ``prompt_token_ids``: ``list[int]`` or ``list[list[int]]`` (one entry
+            per request).
+          - ``additional_information``: dict, or list of dicts matching
+            ``prompt_token_ids``.
+          - ``expected_sample_rate``: asserted against the reported ``sr``.
+          - ``min_duration_s`` / ``max_duration_s``: bounds on decoded audio.
+          - plus every key ``assert_audio_speech_response`` understands
+            (``response_format``, ``min_audio_bytes``, ``input``, ...).
+        """
+        token_ids = request_config.get("prompt_token_ids")
+        if not token_ids:
+            raise ValueError("request_config must contain 'prompt_token_ids'")
+        batched_ids = token_ids if isinstance(token_ids[0], (list, tuple)) else [token_ids]
+
+        extra = request_config.get("additional_information")
+        if extra is None:
+            batched_extra: list[dict[str, Any] | None] = [None] * len(batched_ids)
+        elif isinstance(extra, dict):
+            batched_extra = [extra] * len(batched_ids)
+        else:
+            batched_extra = list(extra)
+        if len(batched_extra) != len(batched_ids):
+            raise ValueError("'additional_information' must match the number of prompts")
+
+        prompts: list[dict[str, Any]] = []
+        for ids, info in zip(batched_ids, batched_extra, strict=True):
+            prompt: dict[str, Any] = {"prompt_token_ids": list(ids)}
+            if info is not None:
+                prompt["additional_information"] = info
+            prompts.append(prompt)
+
+        start_time = time.perf_counter()
+        outputs = self.runner.omni.generate(prompts)
+        e2e_latency = time.perf_counter() - start_time
+
+        assert len(outputs) == len(prompts), f"Expected {len(prompts)} outputs, got {len(outputs)}"
+        expected_sr = request_config.get("expected_sample_rate")
+        min_duration = request_config.get("min_duration_s")
+        max_duration = request_config.get("max_duration_s")
+
+        result = OmniResponse()
+        wavs: list[torch.Tensor] = []
+        for index, stage_output in enumerate(outputs):
+            assert stage_output.outputs, f"Request {index} produced no output"
+            multimodal_output = stage_output.outputs[0].multimodal_output
+            assert isinstance(multimodal_output, Mapping), (
+                f"Request {index}: expected a mapping, got {type(multimodal_output)}"
+            )
+
+            # Never use ``a or b`` here: truthiness on a tensor raises.
+            audio = multimodal_output.get("model_outputs")
+            if audio is None:
+                audio = multimodal_output.get("audio")
+            assert audio is not None, f"Request {index}: no audio key in {sorted(multimodal_output)}"
+            if isinstance(audio, list):
+                chunks = [torch.as_tensor(chunk).float().cpu().reshape(-1) for chunk in audio if chunk is not None]
+                assert chunks, f"Request {index}: audio list contained no tensors"
+                audio = torch.cat(chunks, dim=0)
+            wav = torch.as_tensor(audio).float().cpu().reshape(-1)
+            assert wav.numel() > 0, f"Request {index}: decoded audio is empty"
+            wavs.append(wav)
+
+            sr_raw = multimodal_output.get("sr")
+            if isinstance(sr_raw, list):
+                sr_raw = sr_raw[-1] if sr_raw else None
+            assert sr_raw is not None, f"Request {index}: missing sample rate"
+            sample_rate = int(sr_raw.item() if hasattr(sr_raw, "item") else sr_raw)
+            if expected_sr is not None:
+                assert sample_rate == int(expected_sr), (
+                    f"Request {index}: expected sample rate {expected_sr}, got {sample_rate}"
+                )
+
+            duration_s = wav.numel() / sample_rate
+            if min_duration is not None:
+                assert duration_s >= float(min_duration), (
+                    f"Request {index}: audio too short ({duration_s:.2f}s < {min_duration}s)"
+                )
+            if max_duration is not None:
+                assert duration_s <= float(max_duration), (
+                    f"Request {index}: audio too long ({duration_s:.2f}s > {max_duration}s)"
+                )
+
+            wav_buf = io.BytesIO()
+            sf.write(wav_buf, wav.numpy(), samplerate=sample_rate, format="WAV", subtype="PCM_16")
+            result = OmniResponse(
+                success=True,
+                audio_bytes=wav_buf.getvalue(),
+                audio_format="audio/wav",
+                e2e_latency=e2e_latency,
+            )
+            assert_audio_speech_response(result, request_config, run_level="core_model")
+
+        if request_config.get("assert_distinct_outputs") and len(wavs) > 1:
+            # Distinct inputs must decode to distinct audio. Leaked per-request
+            # codec state shows up as two requests producing (near-)identical
+            # waveforms; different lengths already prove distinctness.
+            for i in range(len(wavs)):
+                for j in range(i + 1, len(wavs)):
+                    a, b = wavs[i], wavs[j]
+                    if a.numel() != b.numel():
+                        continue
+                    max_diff = (a - b).abs().max().item()
+                    assert max_diff > 1e-3, (
+                        f"Requests {i} and {j} produced near-identical audio "
+                        f"(max sample diff {max_diff:.2e}); possible codec-state crosstalk"
+                    )
+        return result
+
     def start_profile(self, profile_prefix: str | None = None, stages: list[int] | None = None) -> list[Any]:
         return self.runner.start_profile(profile_prefix=profile_prefix, stages=stages)
 

tests/model_executor/models/audio8_tts/test_audio8_tts_codec_decoder.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import functools
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+NUM_CODEBOOKS = 10
+#: Samples the fake codec returns per frame, so trims are easy to assert.
+SAMPLES_PER_FRAME = 100
+
+
+@functools.lru_cache(maxsize=1)
+def _decoder_cls():
+    from vllm_omni.model_executor.models.audio8_tts.audio8_tts_codec_decoder import Audio8TTSCodecDecoder
+
+    return Audio8TTSCodecDecoder
+
+
+class _FakeCodec:
+    """Returns ``SAMPLES_PER_FRAME`` ramp samples per frame."""
+
+    def __init__(self):
+        self.calls: list[tuple[int, ...]] = []
+
+    def decode(self, codes_bqf: torch.Tensor) -> torch.Tensor:
+        self.calls.append(tuple(codes_bqf.shape))
+        frames = int(codes_bqf.shape[-1])
+        total = frames * SAMPLES_PER_FRAME
+        return torch.arange(total, dtype=torch.float32).view(1, 1, total)
+
+
+def _make_decoder(codec: _FakeCodec, num_codebooks: int = NUM_CODEBOOKS, splits=None):
+    decoder = object.__new__(_decoder_cls())
+    torch.nn.Module.__init__(decoder)
+    decoder._codec = codec
+    decoder._codec_device = torch.device("cpu")
+    decoder._codec_dtype = torch.float32
+    decoder._num_codebooks = num_codebooks
+    decoder._sample_rate = 44100
+    decoder._frame_size = 2048
+    decoder._logged_stats = True
+    decoder._ensure_codec_loaded = lambda: None
+    if splits is not None:
+        decoder._split_request_ids = lambda ids, seq_token_counts=None: splits
+    return decoder
+
+
+def test_left_context_frames_are_trimmed_from_the_emitted_chunk():
+    """The decoder must emit *delta* audio: the left-context prefix exists only
+    so the causal codec has history and must never reach the client twice."""
+    codec = _FakeCodec()
+    decoder = _make_decoder(codec)
+    codes = torch.arange(NUM_CODEBOOKS * 4, dtype=torch.long).reshape(NUM_CODEBOOKS, 4)
+
+    out = decoder.forward(
+        input_ids=torch.tensor([0], dtype=torch.long),
+        runtime_additional_information=[{"codes": {"audio": codes}, "meta": {"left_context_size": 3}}],
+    )
+
+    audios = out.multimodal_outputs["model_outputs"]
+    assert len(audios) == 1
+    # 4 frames decoded, 3 of context trimmed => 1 frame of new audio.
+    assert audios[0].shape[0] == SAMPLES_PER_FRAME
+    assert codec.calls == [(1, NUM_CODEBOOKS, 4)]
+
+
+def test_tensor_codes_from_runtime_info_win_over_input_ids():
+    codec = _FakeCodec()
+    decoder = _make_decoder(codec, num_codebooks=2)
+    out = decoder.forward(
+        input_ids=torch.zeros(64, dtype=torch.long),
+        runtime_additional_information=[{"codes": {"audio": torch.tensor([[1, 2], [3, 4]], dtype=torch.long)}}],
+    )
+    assert codec.calls == [(1, 2, 2)]
+    assert out.multimodal_outputs["model_outputs"][0].shape[0] == 2 * SAMPLES_PER_FRAME
+
+
+def test_flat_input_ids_are_reshaped_codebook_major():
+    codec = _FakeCodec()
+    decoder = _make_decoder(codec, num_codebooks=2)
+    out = decoder.forward(input_ids=torch.arange(6, dtype=torch.long), runtime_additional_information=[{}])
+    assert codec.calls == [(1, 2, 3)]
+    assert out.multimodal_outputs["model_outputs"][0].shape[0] == 3 * SAMPLES_PER_FRAME
+
+
+def test_every_return_path_sets_model_outputs_and_sr():
+    """A branch that omits ``model_outputs`` makes the serving layer drop the
+    chunk silently, so assert the key exists on all the early-exit paths."""
+    codec = _FakeCodec()
+
+    empty = _make_decoder(codec).forward(input_ids=torch.zeros(0, dtype=torch.long))
+    assert empty.multimodal_outputs["model_outputs"][0].numel() == 0
+    assert empty.multimodal_outputs["sr"]
+
+    ragged = _make_decoder(codec, num_codebooks=4).forward(
+        input_ids=torch.arange(6, dtype=torch.long),  # 6 % 4 != 0
+        runtime_additional_information=[{}],
+    )
+    assert ragged.multimodal_outputs["model_outputs"][0].numel() == 0
+
+    codes = torch.zeros((NUM_CODEBOOKS, 2), dtype=torch.long)
+    no_new = _make_decoder(codec).forward(
+        input_ids=torch.tensor([0], dtype=torch.long),
+        runtime_additional_information=[{"codes": {"audio": codes}, "meta": {"left_context_size": 2}}],
+    )
+    assert no_new.multimodal_outputs["model_outputs"][0].numel() == 0
+    assert codec.calls == [], "a context-only chunk must not be decoded"
+
+
+def test_mixed_batch_keeps_per_request_alignment():
+    """An empty request must not shift the following request's audio into its
+    slot: the output list index is the request index."""
+    codec = _FakeCodec()
+    splits = [
+        torch.empty((0,), dtype=torch.long),
+        torch.arange(NUM_CODEBOOKS * 2, dtype=torch.long),
+    ]
+    decoder = _make_decoder(codec, splits=splits)
+    out = decoder.forward(
+        input_ids=torch.arange(NUM_CODEBOOKS * 2, dtype=torch.long),
+        runtime_additional_information=[{}, {"meta": {"left_context_size": 1}}],
+    )
+    audios = out.multimodal_outputs["model_outputs"]
+    assert len(audios) == 2
+    assert audios[0].numel() == 0
+    assert audios[1].shape[0] == SAMPLES_PER_FRAME
+
+
+def test_codes_with_wrong_codebook_count_fall_back_to_input_ids():
+    codec = _FakeCodec()
+    decoder = _make_decoder(codec, num_codebooks=2)
+    out = decoder.forward(
+        input_ids=torch.arange(4, dtype=torch.long),
+        runtime_additional_information=[{"codes": {"audio": torch.zeros((5, 3), dtype=torch.long)}}],
+    )
+    assert codec.calls == [(1, 2, 2)]
+    assert out.multimodal_outputs["model_outputs"][0].shape[0] == 2 * SAMPLES_PER_FRAME

tests/model_executor/models/audio8_tts/test_audio8_tts_prompt_utils.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+from vllm_omni.model_executor.models.audio8_tts.prompt_utils import (
+    build_text_only_prompt_ids,
+    build_voice_clone_prompt_ids,
+    estimate_voice_clone_prompt_len,
+    normalize_text,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _SegmentTokenizer:
+    """Tokenizer stub that maps each encoded string to one sentinel id.
+
+    Lets the tests assert the *segment sequence* the prompt builder emits, which
+    is what must stay byte-identical to the reference ``ArkttsProcessor``.
+    """
+
+    def __init__(self):
+        self.vocab: dict[str, int] = {}
+
+    def encode(self, text: str, add_special_tokens: bool = True):
+        assert add_special_tokens is False, "Audio8 prompts must not add extra special tokens"
+        return [self.vocab.setdefault(text, 1000 + len(self.vocab))]
+
+    def segment_of(self, token_id: int) -> str:
+        return next(text for text, tid in self.vocab.items() if tid == token_id)
+
+
+def _segments(tokenizer: _SegmentTokenizer, ids: list[int]) -> list[str]:
+    return [tokenizer.segment_of(i) for i in ids]
+
+
+def test_text_only_prompt_matches_reference_segment_order():
+    tokenizer = _SegmentTokenizer()
+    ids, normalized = build_text_only_prompt_ids(tokenizer, "  Hello   world  ")
+
+    assert normalized == "Hello world"
+    assert _segments(tokenizer, ids) == [
+        "<|im_start|>system\n",
+        "convert the provided text to speech",
+        "<|im_end|>\n",
+        "<|im_start|>user\n",
+        "Hello world",
+        "<|im_end|>\n",
+        "<|im_start|>assistant\n<|voice|>",
+    ]
+
+
+def test_voice_clone_prompt_places_reference_codes_between_prefix_and_suffix():
+    tokenizer = _SegmentTokenizer()
+    semantic_ids = [151678, 151679, 151680]
+    ids, ref_start, text, ref_text = build_voice_clone_prompt_ids(
+        tokenizer, "Target text", "Reference transcript", semantic_ids
+    )
+
+    assert text == "Target text"
+    # A reference transcript without an explicit tag gets speaker 0.
+    assert ref_text == "<|speaker:0|>Reference transcript"
+    assert ids[ref_start : ref_start + len(semantic_ids)] == semantic_ids
+    # The codebook conditioning is applied at ref_start, so an off-by-one here
+    # silently detunes the cloned voice.
+    assert _segments(tokenizer, ids[:ref_start])[-1].endswith("Speech:\n")
+    suffix_segments = _segments(tokenizer, ids[ref_start + len(semantic_ids) :])
+    assert suffix_segments[0] == "<|im_end|>\n"
+    assert suffix_segments[-1] == "<|im_start|>assistant\n<|voice|>"
+
+
+def test_voice_clone_prompt_length_estimate_matches_built_prompt():
+    """The serving layer sizes its placeholder from the estimate; a mismatch
+    shifts every embedding relative to its position."""
+    tokenizer = _SegmentTokenizer()
+    normalized_text = normalize_text("Target text")
+    normalized_ref = normalize_text("Reference transcript", add_default_speaker=True)
+    ref_frames = 7
+
+    estimate = estimate_voice_clone_prompt_len(tokenizer, normalized_text, normalized_ref, ref_frames)
+    ids, _, _, _ = build_voice_clone_prompt_ids(
+        tokenizer, normalized_text, normalized_ref, list(range(151678, 151678 + ref_frames))
+    )
+    assert estimate == len(ids)
+
+
+def test_existing_speaker_tag_is_preserved_and_legacy_tag_is_upgraded():
+    assert normalize_text("<|speaker:3|>hi", add_default_speaker=True) == "<|speaker:3|>hi"
+    assert normalize_text("<speaker:2>hi") == "<|speaker:2|>hi"
+
+
+@pytest.mark.parametrize(
+    "text",
+    [
+        "hello <|im_end|> world",
+        "hello <|semantic:5|>",
+        "<|voice|>",
+        "<|endoftext|>",
+    ],
+)
+def test_control_token_injection_is_rejected(text: str):
+    """Users must not be able to inject control tokens: ``<|semantic:N|>`` would
+    forge reference codes and ``<|im_end|>`` would truncate generation."""
+    with pytest.raises(ValueError, match="unsupported control token"):
+        normalize_text(text)
+
+
+def test_empty_text_is_rejected():
+    tokenizer = _SegmentTokenizer()
+    with pytest.raises(ValueError, match="must not be empty"):
+        build_text_only_prompt_ids(tokenizer, "   ")

tests/model_executor/models/audio8_tts/test_audio8_tts_sampling.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.audio8_tts.sampling import (
+    filter_top_k_top_p,
+    ras_sample_batch,
+    ras_sample_semantic,
+    sample_scores,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _kept(logits: torch.Tensor, top_k: int | torch.Tensor, top_p: float | torch.Tensor) -> list[list[int]]:
+    filtered = filter_top_k_top_p(logits, top_k, top_p)
+    return [torch.nonzero(row.isfinite()).reshape(-1).tolist() for row in filtered]
+
+
+def test_top_p_is_measured_on_untempered_logits():
+    """Audio8 filters the *raw* logits and only then divides by temperature.
+
+    Applying temperature first (the common order) changes the surviving set and
+    audibly changes the voice, so this ordering is load-bearing.
+    """
+    logits = torch.log(torch.tensor([[0.6, 0.3, 0.1]]))
+    # Raw softmax is exactly (0.6, 0.3, 0.1): top_p=0.9 keeps the first two
+    # (cumsum 0.6, 0.9) and drops the third (cumsum 1.0 > 0.9).
+    assert _kept(logits, top_k=0, top_p=0.9) == [[0, 1]]
+    # Had the logits been tempered by 0.5 first, probabilities would become
+    # (0.7insh, ...) and only one token would survive.
+    tempered = _kept(logits / 0.5, top_k=0, top_p=0.9)
+    assert tempered == [[0, 1]] or tempered == [[0]]
+
+
+def test_argmax_always_survives_filtering():
+    logits = torch.tensor([[10.0, -1.0, -2.0]])
+    # top_p=0.0 would otherwise remove every candidate and produce NaNs.
+    assert _kept(logits, top_k=1, top_p=0.0) == [[0]]
+
+
+def test_per_row_top_k_and_top_p_are_honoured():
+    logits = torch.log(torch.tensor([[0.5, 0.3, 0.15, 0.05], [0.5, 0.3, 0.15, 0.05]]))
+    kept = _kept(logits, top_k=torch.tensor([1, 3]), top_p=torch.tensor([1.0, 1.0]))
+    assert kept == [[0], [0, 1, 2]]
+
+
+def test_zero_temperature_row_falls_back_to_argmax():
+    logits = torch.tensor([[1.0, 5.0, 2.0], [1.0, 5.0, 2.0]])
+    out = sample_scores(
+        logits,
+        temperature=torch.tensor([0.0, 1.0]),
+        top_k=0,
+        top_p=1.0,
+        generator=torch.Generator().manual_seed(0),
+    )
+    assert int(out[0]) == 1
+
+
+def test_ras_substitutes_only_repeated_semantic_ids():
+    """RAS re-rolls a repeat from a flatter distribution; the EOS slot and
+    non-repeats must pass through untouched, or utterances get truncated."""
+    num_semantic = 2
+    # Row 0 will draw id 0 (a repeat) -> substituted.
+    # Row 1 will draw id 2 (the EOS slot) -> never substituted even if recent.
+    logits = torch.tensor(
+        [
+            [10.0, -20.0, -20.0],
+            [-20.0, -20.0, 10.0],
+        ]
+    )
+    recent = torch.tensor([[0, -1], [2, -1]])
+    out = ras_sample_semantic(
+        logits,
+        recent,
+        temperature=1e-6,  # deterministic: argmax within each draw
+        top_k=0,
+        top_p=1.0,
+        ras_temperature=1e-6,
+        ras_top_p=1.0,
+        num_semantic_ids=num_semantic,
+    )
+    # Both draws are argmax, so substitution is observable only through which
+    # branch was taken; assert the EOS row is preserved verbatim.
+    assert int(out[1]) == 2
+    assert int(out[0]) == 0
+
+
+def test_ras_without_history_returns_the_plain_draw():
+    logits = torch.tensor([[10.0, 0.0, -5.0]])
+    out = ras_sample_semantic(
+        logits,
+        None,
+        temperature=1e-6,
+        top_k=0,
+        top_p=1.0,
+        ras_temperature=1.0,
+        ras_top_p=0.9,
+        num_semantic_ids=2,
+    )
+    assert int(out[0]) == 0
+
+
+def test_sampling_is_reproducible_for_a_seeded_generator():
+    logits = torch.randn(4, 32, generator=torch.Generator().manual_seed(7))
+    first = sample_scores(logits, temperature=0.7, top_k=8, top_p=0.9, generator=torch.Generator().manual_seed(1234))
+    second = sample_scores(logits, temperature=0.7, top_k=8, top_p=0.9, generator=torch.Generator().manual_seed(1234))
+    assert torch.equal(first, second)
+
+
+def test_batch_routes_each_row_to_its_own_generator():
+    """Regression for the shared slot-0 generator bug.
+
+    A shared generator over ``[num_reqs, vocab]`` made every row draw from slot
+    0's RNG stream, so a request's audio depended on its batch neighbours and
+    per-request ``seed`` was ignored for slots > 0. Flat logits make the draw
+    purely noise-driven so the wrong generator changes the token.
+    """
+    # Flat logits: the sampled id is determined only by the generator's noise,
+    # so routing the generator to the wrong slot changes the result (coincidence ~1/256).
+    logits = torch.zeros(2, 256)
+    kw = dict(temperature=1.0, top_k=0, top_p=1.0, ras_temperature=1.0, ras_top_p=0.9, num_semantic_ids=200)
+
+    # (1) Each batched row equals sampling that row alone with its own generator.
+    batched = ras_sample_batch(
+        logits,
+        None,
+        generators={0: torch.Generator().manual_seed(11), 1: torch.Generator().manual_seed(22)},
+        num_reqs=2,
+        **kw,
+    )
+    row0 = ras_sample_semantic(logits[0:1], None, generator=torch.Generator().manual_seed(11), **kw)
+    row1 = ras_sample_semantic(logits[1:2], None, generator=torch.Generator().manual_seed(22), **kw)
+    assert int(batched[0]) == int(row0)
+    assert int(batched[1]) == int(row1)
+
+    # (2) Row 1 must not change when only slot 0's seed changes.
+    other = ras_sample_batch(
+        logits,
+        None,
+        generators={0: torch.Generator().manual_seed(9999), 1: torch.Generator().manual_seed(22)},
+        num_reqs=2,
+        **kw,
+    )
+    assert int(other[1]) == int(batched[1])
+
+
+def test_single_request_batch_keeps_the_fast_path():
+    """``num_reqs == 1`` must match a direct single-row draw (no per-row loop)."""
+    logits = torch.randn(1, 48, generator=torch.Generator().manual_seed(5))
+    kw = dict(temperature=0.8, top_k=16, top_p=0.9, ras_temperature=1.0, ras_top_p=0.9, num_semantic_ids=40)
+    batched = ras_sample_batch(logits, None, generators={0: torch.Generator().manual_seed(99)}, num_reqs=1, **kw)
+    direct = ras_sample_semantic(logits, None, generator=torch.Generator().manual_seed(99), **kw)
+    assert torch.equal(batched, direct)

tests/model_executor/stage_input_processors/test_audio8_tts_async_chunk.py

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.audio8_tts import (
+    extract_last_frame,
+    slow_ar_to_codec_decoder_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+NUM_CODEBOOKS = 3
+
+
+def _make_transfer_manager(*, chunk_frames=2, left_context=1, initial_chunk_frames=0):
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        put_req_chunk=defaultdict(int),
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "codec_chunk_frames": chunk_frames,
+                    "codec_left_context_frames": left_context,
+                    "initial_codec_chunk_frames": initial_chunk_frames,
+                }
+            }
+        ),
+    )
+
+
+def _make_request(req_id="req", *, finished=False):
+    return SimpleNamespace(
+        external_req_id=req_id,
+        additional_information=None,
+        is_finished=lambda: finished,
+    )
+
+
+def _frame(value: int) -> torch.Tensor:
+    return torch.full((NUM_CODEBOOKS,), value, dtype=torch.long)
+
+
+def _step(transfer_manager, request, frame_value: int | None, *, finished: bool = False):
+    multimodal_output = {} if frame_value is None else {"audio_codes": _frame(frame_value).unsqueeze(0)}
+    return slow_ar_to_codec_decoder_async_chunk(transfer_manager, multimodal_output, request, is_finished=finished)
+
+
+def test_codes_accumulate_across_steps_and_chunk_is_withheld_until_full():
+    """Codes must accumulate; emitting per-step would starve the codec of the
+    context it needs and resetting would truncate the utterance."""
+    transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1)
+    request = _make_request()
+
+    assert _step(transfer_manager, request, 1) is None
+    payload = _step(transfer_manager, request, 2)
+
+    assert payload is not None
+    assert len(transfer_manager.code_prompt_token_ids["req"]) == 2
+    # [num_codebooks, frames], codebook-major.
+    assert tuple(payload.codes.audio.shape) == (NUM_CODEBOOKS, 2)
+    assert payload.codes.audio[0].tolist() == [1, 2]
+    # Nothing precedes the first chunk, so there is no left context yet.
+    assert payload.meta.left_context_size == 0
+
+
+def test_second_chunk_carries_left_context_and_reports_its_size():
+    transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1)
+    request = _make_request()
+    for value in (1, 2, 3):
+        _step(transfer_manager, request, value)
+    payload = _step(transfer_manager, request, 4)
+
+    assert payload is not None
+    # 1 context frame + 2 new frames.
+    assert tuple(payload.codes.audio.shape) == (NUM_CODEBOOKS, 3)
+    assert payload.codes.audio[0].tolist() == [2, 3, 4]
+    assert payload.meta.left_context_size == 1
+
+
+def test_partial_tail_is_flushed_when_the_request_finishes():
+    transfer_manager = _make_transfer_manager(chunk_frames=4, left_context=2)
+    request = _make_request()
+    _step(transfer_manager, request, 1)
+    _step(transfer_manager, request, 2)
+    payload = _step(transfer_manager, request, 3, finished=True)
+
+    assert payload is not None
+    assert payload.codes.audio[0].tolist() == [1, 2, 3]
+    assert bool(payload.meta.finished) is True
+
+
+def test_finished_request_without_any_frame_emits_a_terminal_payload():
+    transfer_manager = _make_transfer_manager()
+    payload = slow_ar_to_codec_decoder_async_chunk(
+        transfer_manager, None, _make_request(finished=True), is_finished=True
+    )
+    assert payload is not None
+    assert payload.codes.audio.numel() == 0
+    assert bool(payload.meta.finished) is True
+
+
+def test_terminal_call_on_a_boundary_does_not_re_emit_the_last_chunk():
+    """A finished call carrying no new frame, arriving when the stream sits
+    exactly on a chunk boundary, must close the stream rather than re-ship the
+    last window. Before the guard, the final chunk_frames frames shipped twice."""
+    transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1)
+    request = _make_request()
+    for value in (1, 2, 3, 4):  # length ends exactly on a chunk boundary
+        _step(transfer_manager, request, value)
+    # Separate terminal call, no new frame: the boundary chunk is already sent.
+    payload = _step(transfer_manager, request, None, finished=True)
+    assert payload is not None
+    assert payload.codes.audio.numel() == 0
+    assert bool(payload.meta.finished) is True
+
+
+def test_final_frame_on_a_boundary_still_emits_when_it_arrives_finished():
+    """If the last frame both lands on a boundary and carries finished, its chunk
+    has not been shipped yet, so it must still be emitted (the guard must key on
+    'no new frame this call', not merely on being finished at a boundary)."""
+    transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1)
+    request = _make_request()
+    _step(transfer_manager, request, 1)
+    payload = _step(transfer_manager, request, 2, finished=True)
+    assert payload is not None
+    assert payload.codes.audio[0].tolist() == [1, 2]
+    assert bool(payload.meta.finished) is True
+
+
+def test_initial_chunk_shortens_time_to_first_audio():
+    transfer_manager = _make_transfer_manager(chunk_frames=4, left_context=2, initial_chunk_frames=1)
+    request = _make_request()
+    payload = _step(transfer_manager, request, 1)
+
+    assert payload is not None
+    assert payload.codes.audio[0].tolist() == [1]
+    assert payload.meta.left_context_size == 0
+
+
+def test_prefill_placeholder_frames_are_not_forwarded():
+    """Prefill emits an all-zero placeholder frame per position; forwarding it
+    would prepend silence and desynchronise the chunk boundaries."""
+    assert extract_last_frame({"audio_codes": torch.zeros((3, NUM_CODEBOOKS), dtype=torch.long)}) is None
+    assert extract_last_frame({}) is None
+    assert extract_last_frame({"audio_codes": torch.zeros((0, NUM_CODEBOOKS), dtype=torch.long)}) is None
+
+
+def test_explicit_validity_flag_overrides_the_all_zero_heuristic():
+    frame = torch.zeros((1, NUM_CODEBOOKS), dtype=torch.long)
+    out = extract_last_frame({"audio_codes": frame, "audio_code_valid": torch.tensor([True])})
+    assert out is not None and out.tolist() == [0] * NUM_CODEBOOKS
+
+
+def test_invalid_chunk_config_fails_loudly():
+    transfer_manager = _make_transfer_manager(chunk_frames=0)
+    request = _make_request()
+    with pytest.raises(ValueError, match="codec chunk config"):
+        _step(transfer_manager, request, 1)

tests/worker/test_gpu_ar_model_runner.py

@@ -1829,6 +1829,7 @@ class TestPreferModelSamplerNoneFallback:
             if _declares_prefer_model_sampler(p.read_text(encoding="utf-8", errors="ignore"))
         }
         expected = {
+            "audio8_tts",
             "cosyvoice3",
             "glm_tts",
             "higgs_audio_v2",

vllm_omni/config/pipeline_registry.py

@@ -39,6 +39,7 @@ from vllm_omni.model_executor.models.audex.pipeline import (
     AUDEX_TTA_PIPELINE,
     AUDEX_TTS_PIPELINE,
 )
+from vllm_omni.model_executor.models.audio8_tts.pipeline import AUDIO8_TTS_PIPELINE
 from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE
 from vllm_omni.model_executor.models.bagel.pipeline import (
     BAGEL_PIPELINE,
@@ -172,6 +173,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "voxtral_tts": VOXTRAL_TTS_PIPELINE,
     "glm_tts": GLM_TTS_PIPELINE,
     "fish_qwen3_omni": FISH_SPEECH_PIPELINE,
+    "arktts": AUDIO8_TTS_PIPELINE,
     "ming_flash_omni": MING_FLASH_OMNI_PIPELINE,
     "ming_flash_omni_tts": MING_FLASH_OMNI_TTS_PIPELINE,
     "ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE,

vllm_omni/deploy/audio8_tts.yaml

@@ -0,0 +1,80 @@
+# Audio8 TTS Preview 0.6B deploy: async-chunk streaming slow_ar -> codec_decoder.
+# Model: Audio8/Audio8-TTS-Preview-0.6b (model_type = "arktts").
+#
+# trust_remote_code is deliberately OFF: the checkpoint ships an `auto_map`, and
+# transformers prefers remote code over registered classes when it is enabled,
+# which would bypass vllm-omni's Audio8TTSConfig (the Qwen2-shaped view that
+# Qwen2Model needs). The tokenizer is a plain PreTrainedTokenizerFast and needs
+# no remote code either.
+async_chunk: true
+enable_chunked_prefill: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      # ~21.5 Hz codec; 25 frames ~= 1.16 s of audio.
+      codec_chunk_frames: 25
+      codec_left_context_frames: 25
+      initial_codec_chunk_frames: 4
+      # float32 keeps the Snake activations / weight-normed convs clean.
+      audio8_tts_codec_dtype: float32
+
+stages:
+  - stage_id: 0
+    # >= 4 so concurrent requests pipeline; a default of 1 starves the codec
+    # window and produces audible gaps under load.
+    #
+    # Measured on 1x H20 (20 prompts, mean audio 4.5 s): at client concurrency 8,
+    # raising this to 8 buys +27% throughput and 4.5x lower first-packet latency
+    # (1124 ms -> 247 ms) at the cost of a 0.22 s worst-case buffer deficit in
+    # 1 of 20 requests. 4 is the default because every measured point here had
+    # zero underrun; see examples/online_serving/text_to_speech/README.md.
+    max_num_seqs: 4
+    # 0.6B weights + a 2048-token KV cache need very little: sized for a small
+    # slice of the GPU so the codec stage and other tenants still fit.
+    gpu_memory_utilization: 0.2
+    enforce_eager: false
+    trust_remote_code: false
+    enable_prefix_caching: false
+    async_scheduling: false
+    # The checkpoint's max_seq_len is 2048 packed text+audio positions; with
+    # chunked prefill off, max_num_batched_tokens must be >= max_model_len.
+    max_model_len: 2048
+    max_num_batched_tokens: 2048
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      # generation_config.json of the released checkpoint.
+      temperature: 0.7
+      top_k: 50
+      top_p: 0.9
+      max_tokens: 1024
+      seed: 42
+      repetition_penalty: 1.0
+
+  - stage_id: 1
+    max_num_seqs: 1
+    # The codec is ~1.4 GB of weights plus per-chunk activations.
+    gpu_memory_utilization: 0.1
+    enforce_eager: true
+    trust_remote_code: false
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_model_len: 2048
+    max_num_batched_tokens: 2048
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0

vllm_omni/entrypoints/openai/serving_speech.py

@@ -712,6 +712,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             if self._tts_model_type in (
                 "cosyvoice3",
                 "fish_tts",
+                "audio8_tts",
                 "omnivoice",
                 "moss_tts_nano",
                 "glm_tts",
@@ -721,6 +722,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 label = {
                     "cosyvoice3": "CosyVoice3",
                     "fish_tts": "Fish Speech",
+                    "audio8_tts": "Audio8 TTS",
                     "omnivoice": "OmniVoice",
                     "moss_tts_nano": "MOSS-TTS-Nano",
                     "higgs_audio_v2": "Higgs-Audio V2",

vllm_omni/entrypoints/openai/tts_adapters/__init__.py

@@ -104,6 +104,7 @@ def tts_entry_stage_archs() -> frozenset[str]:
 from vllm_omni.entrypoints.openai.tts_adapters import (  # noqa: E402,F401
     audex,
     audex_tta,
+    audio8_tts,
     cosyvoice3,
     covo_audio,
     dots_tts,

vllm_omni/entrypoints/openai/tts_adapters/audio8_tts.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio8 TTS Preview serving adapter for ``/v1/audio/speech``."""
+
+import copy
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+from vllm.inputs import tokens_input
+from vllm.utils.async_utils import make_async
+
+from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+    ARTTSAdapter,
+    PreparedRequest,
+    apply_max_new_tokens,
+    conditioning_cache_salt,
+)
+from vllm_omni.model_executor.models.audio8_tts.codec_utils import (
+    estimate_reference_code_frames,
+)
+from vllm_omni.model_executor.models.audio8_tts.prompt_utils import (
+    build_text_only_prompt_ids,
+    estimate_voice_clone_prompt_len,
+    normalize_text,
+)
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+@register_tts_adapter
+class Audio8TTSAdapter(ARTTSAdapter):
+    """Audio8 TTS Preview 0.6B: text-only synthesis and zero-shot voice cloning."""
+
+    stage_keys = frozenset({"audio8_tts_slow_ar"})
+    name = "audio8_tts"
+
+    def __init__(self, ctx: Any) -> None:
+        super().__init__(ctx)
+        self._cached_tokenizer: Any = None
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        server = self.ctx.server
+        err = server._apply_uploaded_speaker(request)
+        if err:
+            return err
+        if not request.input or not request.input.strip():
+            return "Input text cannot be empty"
+
+        if request.ref_audio is not None:
+            fmt_err = server._validate_ref_audio_format(request.ref_audio)
+            if fmt_err:
+                return fmt_err
+            if not request.ref_text or not request.ref_text.strip():
+                return "Voice cloning requires 'ref_text' (transcript of the reference audio)"
+
+        if request.max_new_tokens is not None:
+            if request.max_new_tokens < self.max_new_tokens_min:
+                return f"max_new_tokens must be at least {self.max_new_tokens_min}"
+            if request.max_new_tokens > self.max_new_tokens_max:
+                return f"max_new_tokens cannot exceed {self.max_new_tokens_max}"
+
+        return None
+
+    def apply_sampling_overrides(
+        self,
+        sampling_params_list: list,
+        request: "OpenAICreateSpeechRequest",
+        prompt: dict[str, Any] | None = None,
+        request_id: str | None = None,
+    ) -> list:
+        sampling_params_list = apply_max_new_tokens(sampling_params_list, request)
+        # Seed the Fast AR (talker_mtp) residual sampling per request. The Slow AR
+        # is already seeded via sampling_metadata.generators, but the residual
+        # codebooks read tts_local_seed; without it they fall back to global RNG,
+        # so an unseeded request's audio would depend on batch composition. Default
+        # it from the deploy seed; an explicit request.seed still overrides it
+        # downstream. Mirrors how serving handles qwen3_tts.
+        if sampling_params_list:
+            default_seed = sampling_params_list[0].seed
+            if default_seed is not None:
+                sampling_params_list = copy.deepcopy(sampling_params_list)
+                stage0 = sampling_params_list[0]
+                if stage0.extra_args is None:
+                    stage0.extra_args = {}
+                stage0.extra_args.setdefault("tts_local_seed", int(default_seed))
+        return sampling_params_list
+
+    async def build(
+        self,
+        request: "OpenAICreateSpeechRequest",
+        sampling_params_list: list,
+        has_inline_ref_audio: bool,
+    ) -> PreparedRequest:
+        del sampling_params_list
+        server = self.ctx.server
+        ref_audio_data = None
+        if request.ref_audio is not None:
+            wav_list, sample_rate = await server._resolve_ref_audio(request.ref_audio)
+            ref_audio_data = (wav_list, sample_rate)
+        # Prompt building tokenizes and, for voice clone, allocates tensors; keep
+        # it off the event loop via the server's single-worker TTS executor.
+        build_prompt = make_async(self._build_prompt, executor=server._tts_executor)
+        prompt = await build_prompt(
+            request,
+            ref_audio_data=ref_audio_data,
+            has_inline_ref_audio=has_inline_ref_audio,
+        )
+        tts_params: dict = {}
+        prompt["cache_salt"] = conditioning_cache_salt(request, tts_params)
+        return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name)
+
+    def _tokenizer(self) -> Any:
+        """The Audio8 tokenizer, loaded once per process."""
+        if self._cached_tokenizer is None:
+            from transformers import AutoTokenizer
+
+            model_name = self.ctx.server.engine_client.model_config.model
+            self._cached_tokenizer = AutoTokenizer.from_pretrained(model_name)
+        return self._cached_tokenizer
+
+    def _estimate_prompt_len(self, text: str, ref_text: str, ref_audio: object) -> int:
+        """Exact clone-prompt length, without encoding the reference audio.
+
+        Deliberately not defensive: the placeholder this sizes must match the
+        length ``preprocess()`` actually builds, so a wrong guess corrupts the
+        request instead of failing it. Let errors surface to the caller.
+        """
+        if not isinstance(ref_audio, (list, tuple)) or len(ref_audio) != 2:
+            raise ValueError("Audio8 TTS reference audio must be a (samples, sample_rate) pair")
+        wav, sample_rate = ref_audio
+        ref_frames = estimate_reference_code_frames(len(wav), int(sample_rate))
+        return estimate_voice_clone_prompt_len(self._tokenizer(), text, ref_text, ref_frames)
+
+    def _build_prompt(
+        self,
+        request: "OpenAICreateSpeechRequest",
+        ref_audio_data: tuple[list[float], int] | None = None,
+        *,
+        has_inline_ref_audio: bool = False,
+    ) -> dict[str, Any]:
+        """Build the engine prompt for Audio8 TTS Preview.
+
+        Text-only prompts are tokenized here. Voice cloning cannot be: the
+        prompt embeds the reference audio's own codec codes, so the model-side
+        ``preprocess`` builds it and this path only reserves a placeholder of the
+        exact final length.
+        """
+        server = self.ctx.server
+        if ref_audio_data is None or not request.ref_text:
+            prompt_ids, normalized_text = build_text_only_prompt_ids(self._tokenizer(), request.input)
+            # Scalars are list-wrapped for the text-only path, matching the
+            # other TTS entrypoints' additional_information shape.
+            additional_information: dict[str, Any] = {"text": [normalized_text]}
+            if request.max_new_tokens is not None:
+                additional_information["max_new_tokens"] = [request.max_new_tokens]
+            prompt = tokens_input(prompt_token_ids=prompt_ids)
+            prompt["additional_information"] = additional_information
+            return prompt
+
+        wav_samples, sample_rate = ref_audio_data
+        normalized_text = normalize_text(request.input)
+        normalized_ref_text = normalize_text(request.ref_text, add_default_speaker=True)
+        placeholder_len = self._estimate_prompt_len(normalized_text, normalized_ref_text, ref_audio_data)
+
+        # Structured clone: scalars (not list-wrapped) because model-side
+        # preprocess() consumes these fields directly.
+        additional_information = {
+            "text": normalized_text,
+            "ref_text": normalized_ref_text,
+            "ref_audio_wav": torch.from_numpy(np.asarray(wav_samples, dtype=np.float32)),
+            "ref_audio_sr": int(sample_rate),
+            "audio8_structured_voice_clone": True,
+        }
+        if request.voice is not None:
+            voice_lower = request.voice.lower()
+            if voice_lower in server.uploaded_speakers and not has_inline_ref_audio:
+                # Lets the model cache this speaker's encoded codes.
+                additional_information["voice_name"] = voice_lower
+                additional_information["voice_created_at"] = server._voice_created_at(voice_lower)
+        if request.max_new_tokens is not None:
+            additional_information["max_new_tokens"] = request.max_new_tokens
+        prompt = tokens_input(prompt_token_ids=[1] * placeholder_len)
+        prompt["additional_information"] = additional_information
+        return prompt
+
+
+__all__ = ["Audio8TTSAdapter"]

vllm_omni/model_executor/models/audio8_tts/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio8 TTS Preview (DualAR, 44.1 kHz codec) model implementation."""

vllm_omni/model_executor/models/audio8_tts/audio8_tts_codec_decoder.py

@@ -0,0 +1,284 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio8 TTS Preview -- codec decoder (Stage 1).
+
+Consumes frames of ``num_codebooks`` codec codes and decodes them to a 44.1 kHz
+waveform. Streaming contract is delta: with ``async_chunk`` on, each chunk is
+``[left_context | new_frames]``; the whole window is decoded so the causal
+codec keeps continuous context, then only the new tail is returned.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.forward_context import get_forward_context, is_forward_context_available
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.audio8_tts.codec_utils import load_arktts_codec
+from vllm_omni.model_executor.models.audio8_tts.configuration_audio8_tts import (
+    ARKTTS_CODEC_FRAME_SIZE,
+    ARKTTS_CODEC_SAMPLE_RATE,
+)
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+logger = init_logger(__name__)
+
+_DTYPES = {
+    "fp16": torch.float16,
+    "float16": torch.float16,
+    "half": torch.float16,
+    "bf16": torch.bfloat16,
+    "bfloat16": torch.bfloat16,
+    "fp32": torch.float32,
+    "float32": torch.float32,
+    "": torch.float32,
+}
+
+
+def _connector_extra_config(vllm_config: VllmConfig) -> dict[str, Any]:
+    model_config = getattr(vllm_config, "model_config", None)
+    connector_cfg = getattr(model_config, "stage_connector_config", None)
+    if isinstance(connector_cfg, dict):
+        return connector_cfg.get("extra", connector_cfg)
+    extra = getattr(connector_cfg, "extra", None)
+    return extra if isinstance(extra, dict) else {}
+
+
+def _resolve_dtype(extra_cfg: dict[str, Any]) -> torch.dtype:
+    raw = str(extra_cfg.get("audio8_tts_codec_dtype", "float32")).strip().lower()
+    if raw not in _DTYPES:
+        raise ValueError(f"Invalid Audio8 TTS codec dtype: {raw!r}")
+    return _DTYPES[raw]
+
+
+class Audio8TTSCodecDecoder(nn.Module):
+    """Stage 1: codec codes -> waveform (runs under the generation runner)."""
+
+    input_modalities = "audio"
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        del prefix
+        self.vllm_config = vllm_config
+        self.model_path = vllm_config.model_config.model
+        config = vllm_config.model_config.hf_config
+
+        self.have_multimodal_outputs = True
+        self.has_preprocess = False
+        self.has_postprocess = False
+        self.enable_update_additional_information = True
+        self.requires_raw_input_tokens = True
+
+        self._num_codebooks = int(getattr(config, "num_codebooks", 10))
+        self._sample_rate = int(getattr(config, "codec_sample_rate", ARKTTS_CODEC_SAMPLE_RATE))
+        self._frame_size = int(getattr(config, "codec_frame_size", ARKTTS_CODEC_FRAME_SIZE))
+        self._codec_kwargs = {
+            "post_n_layer": int(getattr(config, "codec_post_n_layer", 8)),
+            "post_n_head": int(getattr(config, "codec_post_n_head", 16)),
+            "post_n_local_heads": int(getattr(config, "codec_post_n_local_heads", 8)),
+            "post_intermediate_size": int(getattr(config, "codec_post_intermediate_size", 1216)),
+        }
+        # float32 by default: the codec's Snake activations and weight-normed
+        # convs lose audible headroom in bf16.
+        self._codec_dtype = _resolve_dtype(_connector_extra_config(vllm_config))
+        self._codec: nn.Module | None = None
+        self._codec_device: torch.device | None = None
+        self._logged_stats = False
+
+    # -------------------- codec --------------------
+
+    def _ensure_codec_loaded(self) -> None:
+        if self._codec is not None:
+            return
+        device = self.vllm_config.device_config.device
+        self._codec = load_arktts_codec(
+            self.model_path,
+            device=device,
+            dtype=self._codec_dtype,
+            role="decode",
+            **self._codec_kwargs,
+        )
+        self._codec_device = torch.device(device)
+
+    # -------------------- vLLM hooks --------------------
+
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        if input_ids.numel() == 0:
+            return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
+        return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
+
+    def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None:
+        return None
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """No-op: the codec comes from ``codec.pth``, not the main checkpoint."""
+        del weights
+        return set()
+
+    # -------------------- request splitting --------------------
+
+    def _split_request_ids(
+        self,
+        ids: torch.Tensor,
+        seq_token_counts: list[int] | None = None,
+    ) -> list[torch.Tensor]:
+        if seq_token_counts is not None and len(seq_token_counts) > 1:
+            boundaries = [0]
+            for count in seq_token_counts:
+                boundaries.append(boundaries[-1] + count)
+            total = ids.numel()
+            return [ids[boundaries[i] : min(boundaries[i + 1], total)] for i in range(len(seq_token_counts))]
+        if is_forward_context_available():
+            slices = get_forward_context().ubatch_slices
+            if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices):
+                boundaries = [0]
+                for size in slices:
+                    boundaries.append(boundaries[-1] + size)
+                return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)]
+        return [ids]
+
+    def _codes_from_runtime_info(
+        self,
+        info: dict[str, Any] | None,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        """Prefer the tensor payload the connector carried over input_ids."""
+        if not isinstance(info, dict):
+            return None
+        codes = info.get("codes", {}).get("audio")
+        if not isinstance(codes, torch.Tensor) or codes.numel() == 0:
+            return None
+        codes = codes.to(device=device, dtype=torch.long, non_blocking=True).contiguous()
+        if codes.ndim != 2 or int(codes.shape[0]) != self._num_codebooks:
+            logger.warning(
+                "Audio8 TTS codec decoder expected codes of shape [%d, frames], got %s",
+                self._num_codebooks,
+                tuple(codes.shape),
+            )
+            return None
+        return codes
+
+    # -------------------- decode --------------------
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: Any = None,
+        inputs_embeds: torch.Tensor | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        """Decode one chunk per request.
+
+        ``input_ids`` is codebook-major flat codes
+        (``[cb0_f0 ... cb0_fN, cb1_f0 ...]``); tensor codes on the connector
+        win over ``input_ids``. Every return path must set ``model_outputs`` or
+        the serving layer silently drops a chunk of audio.
+        """
+        del positions, intermediate_tensors, inputs_embeds
+        sr_tensor = torch.tensor(self._sample_rate, dtype=torch.int32)
+        empty = torch.zeros((0,), dtype=torch.float32)
+
+        self._ensure_codec_loaded()
+        codec = self._codec
+        assert codec is not None
+
+        if input_ids is None or input_ids.numel() == 0:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]},
+            )
+
+        ids = input_ids.reshape(-1).to(dtype=torch.long)
+        per_request_ids = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
+        num_req = len(per_request_ids)
+
+        context_frames = [0] * num_req
+        if runtime_additional_information is not None:
+            for index, info in enumerate(runtime_additional_information[:num_req]):
+                meta = info.get("meta", {}) if isinstance(info, dict) else {}
+                context_frames[index] = int(meta.get("left_context_size", 0) or 0)
+
+        audios: list[torch.Tensor] = [empty] * num_req
+        for index, request_ids in enumerate(per_request_ids):
+            codes_qf = (
+                self._codes_from_runtime_info(runtime_additional_information[index], ids.device)
+                if runtime_additional_information is not None and index < len(runtime_additional_information)
+                else None
+            )
+            if codes_qf is None:
+                if request_ids.numel() < self._num_codebooks:
+                    continue
+                if int(request_ids.numel()) % self._num_codebooks != 0:
+                    logger.warning(
+                        "Audio8 TTS codec decoder got %d ids, not divisible by num_codebooks=%d; skipping request",
+                        int(request_ids.numel()),
+                        self._num_codebooks,
+                    )
+                    continue
+                codes_qf = request_ids.reshape(self._num_codebooks, -1)
+
+            total_frames = int(codes_qf.shape[1])
+            if total_frames <= 0:
+                continue
+            new_frames = total_frames - context_frames[index]
+            if new_frames <= 0:
+                logger.warning(
+                    "Audio8 TTS codec decoder chunk has no new frames (total=%d, context=%d)",
+                    total_frames,
+                    context_frames[index],
+                )
+                continue
+
+            if not self._logged_stats:
+                self._logged_stats = True
+                logger.info(
+                    "Audio8 TTS codec decoder: frames=%d codebooks=%d dtype=%s sr=%d",
+                    total_frames,
+                    self._num_codebooks,
+                    self._codec_dtype,
+                    self._sample_rate,
+                )
+
+            with torch.amp.autocast("cuda", enabled=False):
+                # Move codes to the codec's device explicitly: a mismatch here
+                # silently falls back to CPU or raises deep inside the decoder.
+                waveform = codec.decode(codes_qf.unsqueeze(0).to(device=self._codec_device))
+            waveform = waveform.reshape(-1).to(dtype=torch.float32)
+
+            # Trim the left-context prefix proportionally: the codec's padding /
+            # rounding means the decoded length is not exactly
+            # frames * frame_size.
+            if context_frames[index] > 0:
+                cut = int(round(context_frames[index] / total_frames * int(waveform.shape[0])))
+                cut = max(0, min(cut, int(waveform.shape[0])))
+                waveform = waveform[cut:]
+            if waveform.numel() > 0:
+                audios[index] = waveform.contiguous()
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"model_outputs": audios, "sr": [sr_tensor] * num_req},
+        )
+
+    def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
+        del kwargs
+        if isinstance(model_outputs, OmniOutput):
+            return model_outputs
+        if not (isinstance(model_outputs, tuple) and len(model_outputs) == 2):
+            raise TypeError(f"Audio8TTSCodecDecoder expected (audio_tensor, sr), got {type(model_outputs)}")
+        audio_tensor, sample_rate = model_outputs
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"model_outputs": audio_tensor, "sr": sample_rate},
+        )
+
+
+__all__ = ["Audio8TTSCodecDecoder"]

vllm_omni/model_executor/models/audio8_tts/audio8_tts_fast_ar.py

@@ -0,0 +1,346 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio8 TTS Fast AR: the residual-codebook predictor.
+
+Runs once per Slow AR step over ``num_codebooks + 1`` positions: position 0 is
+the projected Slow AR hidden state (primes the KV cache, logits discarded),
+position 1 is the sampled semantic code (not predicted, only fed in), positions
+2.. are the previously sampled residual codes; only codes 1..num_codebooks-1
+are sampled here.
+
+Not paged: at most 10 tokens per call, SDPA over a per-call pre-allocated KV
+cache. Matches the reference ``ArkttsModel._generate_codebooks``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.config import VllmConfig
+from vllm.config.vllm import set_current_vllm_config
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from .configuration_audio8_tts import Audio8TTSFastARConfig, Audio8TTSSlowARConfig
+from .sampling import sample_scores
+
+logger = init_logger(__name__)
+
+
+class _FastARAttention(nn.Module):
+    """GQA self-attention over the short codebook sequence.
+
+    ``forward_one`` decodes a single position against a caller-owned KV cache;
+    there is no batched-prefill path because the codebook loop is inherently
+    sequential.
+    """
+
+    def __init__(self, config: Audio8TTSFastARConfig, *, prefix: str = "") -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self._use_gqa = self.num_kv_heads != self.num_heads
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=self.hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.num_heads,
+            total_num_kv_heads=self.num_kv_heads,
+            bias=config.attention_qkv_bias,
+            prefix=f"{prefix}.qkv_proj",
+            disable_tp=True,
+        )
+        self.o_proj = RowParallelLinear(
+            input_size=self.num_heads * self.head_dim,
+            output_size=self.hidden_size,
+            bias=False,
+            prefix=f"{prefix}.o_proj",
+            disable_tp=True,
+        )
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=config.max_position_embeddings,
+            # Audio8 TTS uses interleaved (GPT-J) RoPE, not NeoX.
+            is_neox_style=False,
+            rope_parameters={"rope_theta": config.rope_theta, "rope_type": "default"},
+        )
+        if config.attention_qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        else:
+            self.q_norm = None
+            self.k_norm = None
+
+    def forward_one(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        cache_pos: int,
+    ) -> torch.Tensor:
+        bsz = int(hidden_states.shape[0])
+
+        qkv, _ = self.qkv_proj(hidden_states.reshape(bsz, -1))
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        if self.q_norm is not None:
+            q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape)
+            k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape)
+
+        q, k = self.rotary_emb(position_ids.reshape(-1), q, k)
+
+        q = q.view(bsz, self.num_heads, self.head_dim).unsqueeze(2)
+        k_cache[:bsz, :, cache_pos, :] = k.view(bsz, self.num_kv_heads, self.head_dim)
+        v_cache[:bsz, :, cache_pos, :] = v.view(bsz, self.num_kv_heads, self.head_dim)
+
+        attn_out = F.scaled_dot_product_attention(
+            q,
+            k_cache[:bsz, :, : cache_pos + 1, :],
+            v_cache[:bsz, :, : cache_pos + 1, :],
+            scale=self.scaling,
+            is_causal=False,
+            enable_gqa=self._use_gqa,
+        )
+        attn_out = attn_out.transpose(1, 2).reshape(bsz, -1)
+        output, _ = self.o_proj(attn_out)
+        return output
+
+
+class _FastARMLP(nn.Module):
+    """SiLU-gated MLP (``w1``/``w3`` gate-up, ``w2`` down)."""
+
+    def __init__(self, config: Audio8TTSFastARConfig, *, prefix: str = "") -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=config.hidden_size,
+            output_sizes=[config.intermediate_size] * 2,
+            bias=False,
+            prefix=f"{prefix}.gate_up_proj",
+            disable_tp=True,
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=config.intermediate_size,
+            output_size=config.hidden_size,
+            bias=False,
+            prefix=f"{prefix}.down_proj",
+            disable_tp=True,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(x)
+        gate, up = gate_up.chunk(2, dim=-1)
+        x, _ = self.down_proj(F.silu(gate) * up)
+        return x
+
+
+class _FastARDecoderLayer(nn.Module):
+    def __init__(self, config: Audio8TTSFastARConfig, *, prefix: str = "") -> None:
+        super().__init__()
+        self.self_attn = _FastARAttention(config, prefix=f"{prefix}.self_attn")
+        self.mlp = _FastARMLP(config, prefix=f"{prefix}.mlp")
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward_one(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        cache_pos: int,
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + self.self_attn.forward_one(
+            self.input_layernorm(hidden_states), position_ids, k_cache, v_cache, cache_pos
+        )
+        return hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
+
+
+class Audio8TTSFastAR(nn.Module):
+    """Residual-codebook predictor for Audio8 TTS.
+
+    ``forward`` returns all ``num_codebooks`` codes for one frame, with code 0
+    copied from the already-sampled semantic token.
+    """
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        config: Audio8TTSFastARConfig,
+        slow_ar_config: Audio8TTSSlowARConfig,
+        prefix: str = "fast_ar",
+    ) -> None:
+        super().__init__()
+        self._vllm_config = vllm_config
+        self.config = config
+        self.slow_ar_config = slow_ar_config
+
+        self.layers = nn.ModuleList(
+            [_FastARDecoderLayer(config, prefix=f"{prefix}.layers.{i}") for i in range(config.num_hidden_layers)]
+        )
+        self.fast_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.fast_output = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.fast_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        if slow_ar_config.hidden_size != config.hidden_size:
+            self.fast_project_in = nn.Linear(slow_ar_config.hidden_size, config.hidden_size, bias=True)
+        else:
+            # 0.6b: dim == fast_dim == 896, and the checkpoint has no
+            # fast_project_in tensor.
+            self.fast_project_in = nn.Identity()
+
+        self._num_codebooks = int(config.num_codebooks)
+        self._fast_dim = int(config.hidden_size)
+        self._k_cache: torch.Tensor | None = None
+        self._v_cache: torch.Tensor | None = None
+        self._pos_ids: torch.Tensor | None = None
+
+    def _ensure_buffers(self, bsz: int, device: torch.device, dtype: torch.dtype) -> None:
+        max_seq = self._num_codebooks + 1
+        if (
+            self._k_cache is not None
+            and self._pos_ids is not None
+            and self._k_cache.shape[1] >= bsz
+            and self._k_cache.device == device
+            and self._k_cache.dtype == dtype
+            and self._pos_ids.shape[0] >= bsz
+        ):
+            return
+        self._k_cache = torch.empty(
+            self.config.num_hidden_layers,
+            bsz,
+            self.config.num_key_value_heads,
+            max_seq,
+            self.config.head_dim,
+            dtype=dtype,
+            device=device,
+        )
+        self._v_cache = torch.empty_like(self._k_cache)
+        self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device).unsqueeze(0).expand(bsz, -1).contiguous()
+
+    def _forward_one(self, embed: torch.Tensor, position_ids: torch.Tensor, cache_pos: int) -> torch.Tensor:
+        assert self._k_cache is not None and self._v_cache is not None
+        hidden = embed
+        for layer_idx, layer in enumerate(self.layers):
+            hidden = layer.forward_one(
+                hidden,
+                position_ids,
+                self._k_cache[layer_idx],
+                self._v_cache[layer_idx],
+                cache_pos,
+            )
+        return hidden
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        slow_ar_hidden: torch.Tensor,
+        semantic_token_id: torch.Tensor,
+        *,
+        do_sample: bool = True,
+        temperature: float = 0.7,
+        top_k: int = 50,
+        top_p: float = 0.9,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor:
+        """Predict one frame of codec codes.
+
+        ``slow_ar_hidden`` must already be post-final-norm when
+        ``norm_fastlayer_input`` is set (it is for the released checkpoint) --
+        vLLM's Qwen2Model returns exactly that. Returns ``[B, num_codebooks]``
+        int64 codes; code 0 is the semantic code, codes 1.. come from this module.
+        """
+        bsz = int(slow_ar_hidden.shape[0])
+        num_cb = self._num_codebooks
+        device = slow_ar_hidden.device
+        dtype = slow_ar_hidden.dtype
+
+        semantic_begin = int(self.slow_ar_config.semantic_begin_id)
+        codebook_size = int(self.slow_ar_config.codebook_size)
+        # Non-semantic tokens (EOS) clamp to code 0; those frames are dropped
+        # downstream because the request finishes on that token.
+        semantic_code = (semantic_token_id.reshape(bsz) - semantic_begin).clamp(min=0, max=codebook_size - 1)
+
+        codes = torch.empty(bsz, num_cb, dtype=torch.long, device=device)
+        codes[:, 0] = semantic_code
+
+        self._ensure_buffers(bsz, device, dtype)
+        pos_ids = self._pos_ids
+        assert pos_ids is not None
+
+        # Position 0 primes the cache; its logits are unused.
+        projected = self.fast_project_in(slow_ar_hidden.reshape(bsz, -1))
+        self._forward_one(projected, pos_ids[:bsz, 0], 0)
+
+        embed = self.fast_embeddings(semantic_code)
+        for step in range(1, num_cb):
+            hidden = self._forward_one(embed, pos_ids[:bsz, step], step)
+            logits = self.fast_output(self.fast_norm(hidden)).float()
+            next_ids = sample_scores(
+                logits,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                do_sample=do_sample,
+                generator=generator,
+            )
+            codes[:, step] = next_ids
+            if step < num_cb - 1:
+                embed = self.fast_embeddings(next_ids)
+
+        return codes
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load already-remapped Fast AR weights (see the Slow AR remapper)."""
+        with set_current_vllm_config(self._vllm_config):
+            params_dict = dict(self.named_parameters(remove_duplicate=False))
+            loaded: set[str] = set()
+            stacked_params_mapping = [
+                ("qkv_proj", "q_proj", "q"),
+                ("qkv_proj", "k_proj", "k"),
+                ("qkv_proj", "v_proj", "v"),
+                ("gate_up_proj", "gate_proj", 0),
+                ("gate_up_proj", "up_proj", 1),
+            ]
+            for name, loaded_weight in weights:
+                handled = False
+                for param_name, weight_name, shard_id in stacked_params_mapping:
+                    if weight_name not in name:
+                        continue
+                    mapped = name.replace(weight_name, param_name)
+                    if mapped not in params_dict:
+                        continue
+                    param = params_dict[mapped]
+                    weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                    if weight_loader == default_weight_loader:
+                        weight_loader(param, loaded_weight)
+                    else:
+                        weight_loader(param, loaded_weight, shard_id)
+                    loaded.add(mapped)
+                    handled = True
+                    break
+                if handled or name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded.add(name)
+            return loaded
+
+
+__all__ = ["Audio8TTSFastAR"]

vllm_omni/model_executor/models/audio8_tts/audio8_tts_slow_ar.py

@@ -0,0 +1,755 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio8 TTS Preview -- Slow AR model (Stage 0).
+
+Text (+ optional reference voice) -> one semantic token per codec frame; the
+frame's remaining 9 codebooks are produced by the nested Fast AR.
+
+Backbone is vLLM's ``Qwen2Model`` (qkv bias, no q/k norm), with three
+adjustments: interleaved (GPT-J) RoPE, multi-codebook input embedding at
+semantic positions (summed, unlike Fish Speech's ``1/sqrt(Q+1)`` rescale), and
+logits masked to the semantic range + ``<|im_end|>`` then Repetition-Aware
+sampled.
+
+Streaming contract is delta: each decode step appends one frame of
+``num_codebooks`` codes to ``codes.audio``.
+"""
+
+from __future__ import annotations
+
+import copy
+from collections.abc import Iterable
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer
+from vllm.config import VllmConfig
+from vllm.distributed import get_pp_group
+from vllm.logger import init_logger
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.qwen2 import Qwen2Model
+from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix
+from vllm.sequence import IntermediateTensors
+from vllm.v1.outputs import SamplerOutput
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.utils.speaker_cache import get_speaker_cache
+
+from .audio8_tts_fast_ar import Audio8TTSFastAR
+from .codec_utils import encode_reference_audio_codes
+from .configuration_audio8_tts import (
+    Audio8TTSConfig,
+    Audio8TTSFastARConfig,
+    Audio8TTSSlowARConfig,
+)
+from .prompt_utils import build_voice_clone_prompt_ids
+from .sampling import SAMPLING_EPS, ras_sample_batch
+
+logger = init_logger(__name__)
+
+#: Defaults from ``generation_config.json`` of the released checkpoint.
+DEFAULT_TEMPERATURE = 0.7
+DEFAULT_TOP_K = 50
+DEFAULT_TOP_P = 0.9
+
+
+def _remap_audio8_tts_weights(
+    weights: Iterable[tuple[str, torch.Tensor]],
+    *,
+    q_size: int,
+    kv_size: int,
+    fast_q_size: int,
+    fast_kv_size: int,
+) -> Iterable[tuple[str, torch.Tensor]]:
+    """Rename Audio8 checkpoint tensors to vLLM / Qwen2 names.
+
+    The checkpoint is flat (no ``text_model.`` prefix). Slow AR tensors move
+    under ``model.``, Fast AR tensors under ``fast_ar.``, and ``wqkv`` splits
+    into ``{q,k,v}_proj``. See the body for the exact mapping table.
+    """
+
+    def rewrite_block(suffix: str) -> str:
+        for old, new in (
+            (".attention.wo.", ".self_attn.o_proj."),
+            (".attention.q_norm.", ".self_attn.q_norm."),
+            (".attention.k_norm.", ".self_attn.k_norm."),
+            (".attention_norm.", ".input_layernorm."),
+            (".feed_forward.w1.", ".mlp.gate_proj."),
+            (".feed_forward.w3.", ".mlp.up_proj."),
+            (".feed_forward.w2.", ".mlp.down_proj."),
+            (".ffn_norm.", ".post_attention_layernorm."),
+        ):
+            suffix = suffix.replace(old, new)
+        return suffix
+
+    for name, tensor in weights:
+        # RoPE tables are recomputed by vLLM.
+        if name.endswith(("freqs_cis", "fast_freqs_cis")) or "rotary_emb.inv_freq" in name:
+            continue
+
+        if name.startswith("fast_layers."):
+            suffix = name[len("fast_layers.") :]
+            if ".attention.wqkv." in suffix:
+                layer_prefix, _, param = suffix.partition(".attention.wqkv.")
+                yield f"fast_ar.layers.{layer_prefix}.self_attn.q_proj.{param}", tensor[:fast_q_size]
+                yield (
+                    f"fast_ar.layers.{layer_prefix}.self_attn.k_proj.{param}",
+                    tensor[fast_q_size : fast_q_size + fast_kv_size],
+                )
+                yield (
+                    f"fast_ar.layers.{layer_prefix}.self_attn.v_proj.{param}",
+                    tensor[fast_q_size + fast_kv_size :],
+                )
+                continue
+            yield f"fast_ar.layers.{rewrite_block(suffix)}", tensor
+            continue
+
+        if name.startswith(("fast_embeddings.", "fast_output.", "fast_norm.", "fast_project_in.")):
+            yield f"fast_ar.{name}", tensor
+            continue
+
+        if name.startswith("layers."):
+            suffix = name[len("layers.") :]
+            if ".attention.wqkv." in suffix:
+                layer_prefix, _, param = suffix.partition(".attention.wqkv.")
+                yield f"model.layers.{layer_prefix}.self_attn.q_proj.{param}", tensor[:q_size]
+                yield (
+                    f"model.layers.{layer_prefix}.self_attn.k_proj.{param}",
+                    tensor[q_size : q_size + kv_size],
+                )
+                yield f"model.layers.{layer_prefix}.self_attn.v_proj.{param}", tensor[q_size + kv_size :]
+                continue
+            yield f"model.layers.{rewrite_block(suffix)}", tensor
+            continue
+
+        if name == "embeddings.weight":
+            yield "model.embed_tokens.weight", tensor
+            continue
+        if name == "norm.weight":
+            yield "model.norm.weight", tensor
+            continue
+
+        yield name, tensor
+
+
+class Audio8TTSSlowARForConditionalGeneration(nn.Module):
+    """Stage 0: text -> semantic tokens + residual codec codes."""
+
+    prefer_model_sampler = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.model_path = vllm_config.model_config.model
+
+        config: Audio8TTSConfig = vllm_config.model_config.hf_config  # type: ignore[assignment]
+        self.config = config
+        self.text_config: Audio8TTSSlowARConfig = config.get_text_config()
+        self.fast_ar_config: Audio8TTSFastARConfig = config.fast_ar_config
+
+        self._semantic_begin_id = int(self.text_config.semantic_begin_id)
+        self._semantic_end_id = int(self.text_config.semantic_end_id)
+        self._num_semantic_ids = self._semantic_end_id - self._semantic_begin_id + 1
+        self._eos_token_id = int(self.text_config.eos_token_id)
+        self._pad_token_id = int(self.text_config.pad_token_id)
+        self._codebook_size = int(self.text_config.codebook_size)
+        self._num_codebooks = int(self.text_config.num_codebooks)
+        self._ras_window_size = int(config.ras_window_size)
+        self._ras_temperature = float(config.ras_temperature)
+        self._ras_top_p = float(config.ras_top_p)
+
+        self.have_multimodal_outputs = True
+        self.has_preprocess = True
+        self.has_postprocess = True
+        self.mtp_hidden_size = int(self.text_config.hidden_size)
+        self.talker_mtp_output_key = ("codes", "audio")
+        self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("hidden_states", "last")}
+        self.talker_mtp_graph_safe = True
+
+        self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+        self._fix_rope_style()
+
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(
+                self.text_config.vocab_size,
+                self.text_config.hidden_size,
+                quant_config=vllm_config.quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+        else:
+            self.lm_head = PPMissingLayer()
+        self.logits_processor = LogitsProcessor(self.text_config.vocab_size)
+        self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
+
+        # Summed multi-codebook input embedding table.
+        self.codebook_embeddings = nn.Embedding(
+            self._codebook_size * self._num_codebooks,
+            self.text_config.hidden_size,
+        )
+
+        # Fast AR shares the parent's VllmConfig but must not share the
+        # compilation context: see the Fish Speech note on why this is
+        # copy.copy and not dataclasses.replace (pydantic validators re-run and
+        # reject an already-rebound compilation backend).
+        fast_ar_compilation = copy.copy(vllm_config.compilation_config)
+        fast_ar_compilation.static_forward_context = {}
+        self._fast_ar_vllm_config = copy.copy(vllm_config)
+        self._fast_ar_vllm_config.compilation_config = fast_ar_compilation
+        from vllm.config.vllm import set_current_vllm_config
+
+        with set_current_vllm_config(self._fast_ar_vllm_config):
+            self.fast_ar = Audio8TTSFastAR(
+                vllm_config=self._fast_ar_vllm_config,
+                config=self.fast_ar_config,
+                slow_ar_config=self.text_config,
+                prefix="fast_ar",
+            )
+
+        # Constant logits mask: semantic codes plus <|im_end|>. Safe as a
+        # non-persistent buffer under vLLM (module built on the target device),
+        # unlike under ``PreTrainedModel.from_pretrained`` on transformers >= 5,
+        # which meta-inits and then empty_like's every non-persistent buffer.
+        vocab = int(self.text_config.vocab_size)
+        allowed = torch.zeros((vocab,), dtype=torch.bool)
+        allowed[self._semantic_begin_id : min(self._semantic_end_id + 1, vocab)] = True
+        if self._eos_token_id < vocab:
+            allowed[self._eos_token_id] = True
+        self.register_buffer("_semantic_allowed_mask", allowed, persistent=False)
+
+        self._speaker_cache = get_speaker_cache()
+        self._tokenizer = None
+
+    def _fix_rope_style(self) -> None:
+        """Rebuild RoPE as interleaved (GPT-J); vLLM's Qwen2 defaults to NeoX.
+
+        This also replaces the reference implementation's ``freqs_cis`` table
+        entirely: ``get_rope`` computes its cos/sin cache in the constructor, so
+        there is no checkpoint-absent buffer for a loader to leave uninitialised.
+        """
+        from vllm.model_executor.layers.rotary_embedding import get_rope
+
+        rope_parameters = dict(self.text_config.rope_parameters or {})
+        rope_parameters.setdefault("rope_type", "default")
+        for layer in self.model.layers:
+            attn = layer.self_attn
+            attn.rotary_emb = get_rope(
+                head_size=attn.head_dim,
+                max_position=self.text_config.max_position_embeddings,
+                is_neox_style=False,
+                rope_parameters=rope_parameters,
+            )
+        logger.info("Audio8 TTS: switched %d layers to interleaved RoPE", len(self.model.layers))
+
+    # -------------------- vLLM hooks --------------------
+
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **_: Any,
+    ) -> torch.Tensor | IntermediateTensors:
+        return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor | OmniOutput,
+        sampling_metadata: Any = None,
+    ) -> torch.Tensor | None:
+        if isinstance(hidden_states, OmniOutput):
+            hidden_states = hidden_states.text_hidden_states
+        if hidden_states is None:
+            return None
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        if logits is None:
+            return None
+        return logits.masked_fill(~self._semantic_allowed_mask, float("-inf"))
+
+    # -------------------- Repetition-Aware Sampling --------------------
+
+    def _compact_semantic_logits(self, logits: torch.Tensor) -> torch.Tensor:
+        """Slice the allowed vocabulary into ``[B, num_semantic_ids + 1]``.
+
+        Local ids ``0..num_semantic_ids-1`` are semantic codes; the last column
+        is ``<|im_end|>``. Sorting 4097 columns instead of the full 155776
+        vocabulary is the difference between a cheap and an expensive AR step,
+        and is exactly equivalent because everything else is already ``-inf``.
+        """
+        semantic = logits[:, self._semantic_begin_id : self._semantic_end_id + 1]
+        eos = logits[:, self._eos_token_id : self._eos_token_id + 1]
+        return torch.cat((semantic, eos), dim=-1)
+
+    def _local_to_token_id(self, local_ids: torch.Tensor) -> torch.Tensor:
+        is_eos = local_ids >= self._num_semantic_ids
+        return torch.where(is_eos, torch.full_like(local_ids, self._eos_token_id), local_ids + self._semantic_begin_id)
+
+    def _recent_local_ids(self, output_token_ids: list[list[int]], num_reqs: int, device: torch.device):
+        """Build the ``[B, W]`` RAS window in compact id space.
+
+        Reads the host-side decoded-token history the AR runner already
+        materialised, so this costs one small H2D copy and never syncs the GPU.
+        Returns ``None`` when no request has history yet.
+        """
+        window = self._ras_window_size
+        if window <= 0:
+            return None
+        rows: list[list[int]] = []
+        any_history = False
+        for req_idx in range(num_reqs):
+            history = output_token_ids[req_idx] if req_idx < len(output_token_ids) else []
+            recent = [int(token) for token in history[-window:]]
+            local = [
+                token - self._semantic_begin_id if self._semantic_begin_id <= token <= self._semantic_end_id else -1
+                for token in recent
+            ]
+            any_history = any_history or bool(local)
+            rows.append([-1] * (window - len(local)) + local)
+        if not any_history:
+            return None
+        return torch.tensor(rows, dtype=torch.long, device=device)
+
+    def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> SamplerOutput | None:
+        """Sample the semantic token with the reference filter order + RAS.
+
+        Returns ``None`` to fall back to vLLM's sampler when the request mix
+        needs features RAS does not model (logprobs, penalties, bad words).
+        """
+        if logits is None or logits.numel() == 0:
+            return None
+        if sampling_metadata.max_num_logprobs is not None:
+            return None
+        if not sampling_metadata.no_penalties:
+            return None
+        if sampling_metadata.bad_words_token_ids:
+            return None
+
+        logits = logits.to(torch.float32)
+        if sampling_metadata.allowed_token_ids_mask is not None:
+            num_reqs = int(logits.shape[0])
+            logits.masked_fill_(sampling_metadata.allowed_token_ids_mask[:num_reqs], float("-inf"))
+        for processor in sampling_metadata.logitsprocs.non_argmax_invariant:
+            logits = processor.apply(logits)
+
+        compact = self._compact_semantic_logits(logits)
+        num_reqs = int(compact.shape[0])
+        device = compact.device
+
+        temperature = sampling_metadata.temperature
+        if sampling_metadata.all_greedy or temperature is None:
+            local_ids = compact.argmax(dim=-1)
+        else:
+            temperature = temperature[:num_reqs]
+            top_p = DEFAULT_TOP_P if sampling_metadata.top_p is None else sampling_metadata.top_p[:num_reqs]
+            top_k = DEFAULT_TOP_K if sampling_metadata.top_k is None else sampling_metadata.top_k[:num_reqs]
+            recent = self._recent_local_ids(sampling_metadata.output_token_ids, num_reqs, device)
+            local_ids = ras_sample_batch(
+                compact,
+                recent,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                ras_temperature=self._ras_temperature,
+                ras_top_p=self._ras_top_p,
+                num_semantic_ids=self._num_semantic_ids,
+                generators=sampling_metadata.generators,
+                num_reqs=num_reqs,
+            )
+
+        token_ids = self._local_to_token_id(local_ids).to(dtype=torch.int32)
+        return SamplerOutput(sampled_token_ids=token_ids.unsqueeze(-1), logprobs_tensors=None)
+
+    # -------------------- Omni multimodal plumbing --------------------
+
+    def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
+        if isinstance(model_outputs, OmniOutput):
+            return model_outputs
+
+        hidden = model_outputs
+        info_dicts = kwargs.get("model_intermediate_buffer")
+        if info_dicts is None:
+            info_dicts = kwargs.get("runtime_additional_information") or []
+
+        audio_codes_list: list[torch.Tensor] = []
+        for info in info_dicts:
+            if not isinstance(info, dict):
+                continue
+            codes = info.get("codes", {}).get("audio")
+            if isinstance(codes, torch.Tensor):
+                audio_codes_list.append(codes)
+
+        if not audio_codes_list:
+            return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
+
+        audio_codes = torch.cat(audio_codes_list, dim=0)
+        hidden = hidden[: int(audio_codes.shape[0])]
+        return OmniOutput(text_hidden_states=hidden, multimodal_outputs={"audio_codes": audio_codes})
+
+    # -------------------- preprocess / postprocess --------------------
+
+    def preprocess(
+        self,
+        input_ids: torch.Tensor,
+        input_embeds: torch.Tensor | None,
+        **info_dict: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+        additional_information = info_dict.get("additional_information")
+        if isinstance(additional_information, dict):
+            merged: dict[str, Any] = {k: v for k, v in info_dict.items() if k != "additional_information"}
+            for key, value in additional_information.items():
+                merged.setdefault(key, value)
+            info_dict = merged
+
+        span_len = int(input_ids.shape[0])
+        if span_len <= 0:
+            embeds = input_embeds if input_embeds is not None else self.embed_input_ids(input_ids)
+            return input_ids, embeds, {}
+
+        if span_len > 1:
+            return self._preprocess_prefill(input_ids, info_dict, span_len)
+        return self._preprocess_decode(input_ids, info_dict)
+
+    def _prefill_chunk(
+        self,
+        prompt_embeds_buf: torch.Tensor,
+        start: int,
+        span_len: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Slice ``span_len`` rows out of the pinned CPU prompt-embed buffer."""
+        total = int(prompt_embeds_buf.shape[0])
+        begin = max(0, min(start, total))
+        end = max(0, min(start + span_len, total))
+        take = prompt_embeds_buf[begin:end]
+        if int(take.shape[0]) < span_len:
+            pad_rows = span_len - int(take.shape[0])
+            pad_embed = self.embed_input_ids(
+                torch.tensor([self._pad_token_id], device=device, dtype=torch.long)
+            ).reshape(1, -1)
+            take = torch.cat([take.to(device=device, dtype=torch.bfloat16), pad_embed.expand(pad_rows, -1)], dim=0)
+            return take.to(dtype=torch.bfloat16)
+        return take.to(device=device, dtype=torch.bfloat16, non_blocking=True)
+
+    def _preprocess_prefill(
+        self,
+        input_ids: torch.Tensor,
+        info_dict: dict[str, Any],
+        span_len: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+        device = input_ids.device
+        embed = info_dict.get("embed", {})
+        prompt_embeds_buf = embed.get("prefill")
+        is_first_chunk = not isinstance(prompt_embeds_buf, torch.Tensor) or prompt_embeds_buf.ndim != 2
+
+        if is_first_chunk:
+            if bool(info_dict.get("audio8_structured_voice_clone", False)):
+                prompt_embeds = self._build_voice_clone_prefill_embeds(info_dict)
+            else:
+                prompt_embeds = self.embed_input_ids(input_ids.reshape(1, -1).to(torch.long)).squeeze(0)
+            prompt_embeds_buf = prompt_embeds.detach().to("cpu", dtype=torch.bfloat16).contiguous()
+            if not prompt_embeds_buf.is_pinned():
+                prompt_embeds_buf = prompt_embeds_buf.pin_memory()
+            offset = 0
+        else:
+            offset = int(info_dict.get("meta", {}).get("prefill_offset", 0) or 0)
+
+        total_prompt_len = int(prompt_embeds_buf.shape[0])
+        chunk = self._prefill_chunk(prompt_embeds_buf, offset, span_len, device)
+        next_offset = min(offset + span_len, total_prompt_len)
+
+        # No codes are emitted during prefill; a zero frame per position keeps
+        # the codes buffer aligned with the hidden-state span.
+        zeros = torch.zeros((chunk.shape[0], self._num_codebooks), device=device, dtype=torch.long)
+        info_update = {
+            "embed": {"prefill": prompt_embeds_buf if next_offset < total_prompt_len else None},
+            "meta": {"prefill_offset": next_offset},
+            "codes": {"audio": zeros},
+        }
+        return torch.full_like(input_ids, self._pad_token_id), chunk, info_update
+
+    def _preprocess_decode(
+        self,
+        input_ids: torch.Tensor,
+        info_dict: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+        device = input_ids.device
+        token_embed = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)).to(
+            device=device, dtype=torch.bfloat16
+        )
+        last_hidden = info_dict.get("hidden_states", {}).get("last")
+        if not isinstance(last_hidden, torch.Tensor):
+            # First decode step right after prefill: no Fast AR input yet.
+            logger.warning("Audio8 TTS preprocess: hidden_states.last missing; emitting text-only embedding")
+            return input_ids, token_embed.reshape(1, -1), {}
+
+        # Codebook embeddings are added in talker_mtp, using this step's Fast AR
+        # output; adding them here would use the previous step's codes.
+        info_update = {
+            "mtp_inputs": (
+                last_hidden.to(device=device, dtype=torch.bfloat16).reshape(1, -1),
+                torch.zeros(1, self.text_config.hidden_size, device=device, dtype=torch.bfloat16),
+            ),
+        }
+        return input_ids, token_embed.reshape(1, -1), info_update
+
+    def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]:
+        if hidden_states.numel() == 0:
+            return {}
+        return {"hidden_states": {"last": hidden_states[-1, :].detach().contiguous()}}
+
+    # -------------------- voice cloning --------------------
+
+    def _get_tokenizer(self):
+        if self._tokenizer is None:
+            self._tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+        return self._tokenizer
+
+    def _codec_kwargs(self) -> dict[str, int]:
+        return {
+            "post_n_layer": int(self.config.codec_post_n_layer),
+            "post_n_head": int(self.config.codec_post_n_head),
+            "post_n_local_heads": int(self.config.codec_post_n_local_heads),
+            "post_intermediate_size": int(self.config.codec_post_intermediate_size),
+        }
+
+    def _build_voice_clone_prefill_embeds(self, info_dict: dict[str, Any]) -> torch.Tensor:
+        """Encode the reference audio, build the clone prompt, embed the codes."""
+        text = info_dict.get("text")
+        ref_text = info_dict.get("ref_text")
+        if not isinstance(text, str) or not isinstance(ref_text, str):
+            raise ValueError("Audio8 TTS voice cloning requires string 'text' and 'ref_text'")
+
+        device = self.codebook_embeddings.weight.device
+        cache_key = None
+        voice_name = info_dict.get("voice_name")
+        if isinstance(voice_name, str) and voice_name:
+            cache_key = self._speaker_cache.make_cache_key(
+                voice_name,
+                model_type="audio8_tts",
+                created_at=int(info_dict.get("voice_created_at") or 0),
+            )
+            cached = self._speaker_cache.get(cache_key)
+            if cached is not None:
+                logger.debug("Speaker cache HIT for Audio8 TTS speaker '%s'", voice_name)
+                return self._embed_voice_clone_prompt(
+                    text, ref_text, cached["ref_codes_fq"].to(device=device, dtype=torch.long)
+                )
+
+        ref_audio_sr = info_dict.get("ref_audio_sr")
+        if not isinstance(ref_audio_sr, int):
+            raise ValueError("Audio8 TTS voice cloning requires integer 'ref_audio_sr'")
+        ref_audio_wav = info_dict.get("ref_audio_wav")
+        if ref_audio_wav is None:
+            raise ValueError("Audio8 TTS voice cloning requires 'ref_audio_wav'")
+        if not isinstance(ref_audio_wav, torch.Tensor):
+            ref_audio_wav = torch.from_numpy(np.asarray(ref_audio_wav, dtype=np.float32))
+
+        ref_codes_fq = encode_reference_audio_codes(
+            self.model_path,
+            ref_audio_wav,
+            ref_audio_sr,
+            device=device,
+            **self._codec_kwargs(),
+        )
+        if cache_key is not None:
+            self._speaker_cache.put(cache_key, {"ref_codes_fq": ref_codes_fq.detach().cpu()})
+            logger.debug("Speaker cache STORE for Audio8 TTS speaker '%s'", voice_name)
+        return self._embed_voice_clone_prompt(text, ref_text, ref_codes_fq)
+
+    def _embed_voice_clone_prompt(
+        self,
+        text: str,
+        ref_text: str,
+        ref_codes_fq: torch.Tensor,
+    ) -> torch.Tensor:
+        """Embed the clone prompt, summing codebooks over the reference frames.
+
+        Args:
+            ref_codes_fq: ``[frames, num_codebooks]`` reference codec codes.
+        """
+        device = self.codebook_embeddings.weight.device
+        ref_codes_fq = ref_codes_fq.to(device=device, dtype=torch.long)
+        semantic_token_ids = (ref_codes_fq[:, 0] + self._semantic_begin_id).tolist()
+        prompt_ids, ref_start, _, _ = build_voice_clone_prompt_ids(
+            self._get_tokenizer(), text, ref_text, semantic_token_ids
+        )
+        prompt = torch.tensor(prompt_ids, dtype=torch.long, device=device)
+        embeds = self.embed_input_ids(prompt.unsqueeze(0)).squeeze(0).to(dtype=torch.bfloat16)
+
+        frames = int(ref_codes_fq.shape[0])
+        if frames <= 0:
+            return embeds
+        codebooks = min(int(ref_codes_fq.shape[1]), self._num_codebooks)
+        offsets = (torch.arange(codebooks, device=device, dtype=torch.long) * self._codebook_size).unsqueeze(0)
+        codes = ref_codes_fq[:, :codebooks].clamp(min=0, max=self._codebook_size - 1) + offsets
+        codebook_sum = self.codebook_embeddings(codes).sum(dim=1).to(dtype=embeds.dtype)
+
+        result = embeds.clone()
+        # Audio8 TTS sums the embeddings directly -- no 1/sqrt(Q+1) rescale.
+        result[ref_start : ref_start + frames] += codebook_sum
+        return result.to(dtype=torch.bfloat16)
+
+    # -------------------- GPU-side Fast AR fast path --------------------
+
+    @torch.inference_mode()
+    def talker_mtp(
+        self,
+        input_ids: torch.Tensor,
+        input_embeds: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
+        text_step: torch.Tensor,
+        seed: int | None = None,
+        **kwargs: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Run the Fast AR and fold its codes into this step's input embedding.
+
+        Returns:
+            ``(inputs_embeds, audio_codes)`` where ``audio_codes`` is
+            ``[B, num_codebooks]``.
+        """
+        del text_step
+        bsz = int(input_ids.shape[0])
+        device = input_embeds.device
+
+        token_ids = input_ids.reshape(bsz).to(dtype=torch.long, device=device)
+        past_hidden = last_talker_hidden.reshape(bsz, -1).to(dtype=torch.bfloat16, device=device)
+
+        temperature = kwargs.get("temperature")
+        temperature = DEFAULT_TEMPERATURE if temperature is None else float(temperature)
+        do_sample = kwargs.get("do_sample")
+        generator = kwargs.get("generator")
+        if generator is None and seed is not None:
+            generator = torch.Generator(device=device)
+            generator.manual_seed(int(seed))
+
+        _top_k = kwargs.get("top_k")
+        _top_p = kwargs.get("top_p")
+        audio_codes = self.fast_ar(
+            slow_ar_hidden=past_hidden,
+            semantic_token_id=token_ids,
+            do_sample=(temperature >= SAMPLING_EPS) if do_sample is None else bool(do_sample),
+            temperature=temperature,
+            top_k=int(DEFAULT_TOP_K if _top_k is None else _top_k),
+            top_p=float(DEFAULT_TOP_P if _top_p is None else _top_p),
+            generator=generator,
+        )
+
+        embeds = input_embeds.reshape(bsz, -1)
+        offsets = (torch.arange(self._num_codebooks, device=device, dtype=torch.long) * self._codebook_size).unsqueeze(
+            0
+        )
+        codebook_sum = self.codebook_embeddings(audio_codes + offsets).sum(dim=1).to(dtype=embeds.dtype)
+        is_semantic = (token_ids >= self._semantic_begin_id) & (token_ids <= self._semantic_end_id)
+        embeds = torch.where(is_semantic.unsqueeze(-1), embeds + codebook_sum, embeds)
+        return embeds, audio_codes.to(dtype=torch.long)
+
+    # -------------------- prompt length estimation --------------------
+
+    @staticmethod
+    def estimate_prompt_len_from_additional_information(
+        additional_information: dict[str, Any] | None,
+        **kwargs: Any,
+    ) -> int:
+        """Upper-bound the text-only prompt length for placeholder allocation."""
+        del kwargs
+        info = additional_information or {}
+        text = info.get("text", "")
+        if isinstance(text, list):
+            text = text[0] if text else ""
+        # ~1 token per character is a safe bound for CJK; +64 covers the
+        # chat-template scaffolding.
+        return max(2, len(str(text)) + 64)
+
+    # -------------------- weight loading --------------------
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        text_config = self.text_config
+        fast_config = self.fast_ar_config
+        remapped = _remap_audio8_tts_weights(
+            weights,
+            q_size=text_config.num_attention_heads * text_config.head_dim,
+            kv_size=text_config.num_key_value_heads * text_config.head_dim,
+            fast_q_size=fast_config.num_attention_heads * fast_config.head_dim,
+            fast_kv_size=fast_config.num_key_value_heads * fast_config.head_dim,
+        )
+
+        stacked_params_mapping = [
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded_params: set[str] = set()
+        unexpected: list[str] = []
+
+        for name, loaded_weight in remapped:
+            if name == "model.embed_tokens.weight" and self.text_config.tie_word_embeddings:
+                lm_key = "lm_head.weight"
+                if lm_key in params_dict:
+                    param = params_dict[lm_key]
+                    loader = getattr(param, "weight_loader", default_weight_loader)
+                    loader(param, loaded_weight)
+                    loaded_params.add(lm_key)
+
+            handled = False
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                mapped = name.replace(weight_name, param_name)
+                if mapped not in params_dict:
+                    continue
+                param = params_dict[mapped]
+                loader = getattr(param, "weight_loader", default_weight_loader)
+                if loader == default_weight_loader:
+                    loader(param, loaded_weight)
+                else:
+                    loader(param, loaded_weight, shard_id)
+                loaded_params.add(mapped)
+                handled = True
+                break
+            if handled:
+                continue
+
+            if name in params_dict:
+                param = params_dict[name]
+                loader = getattr(param, "weight_loader", default_weight_loader)
+                loader(param, loaded_weight)
+                loaded_params.add(name)
+            else:
+                unexpected.append(name)
+
+        if unexpected:
+            raise ValueError(
+                f"Audio8 TTS Slow AR received {len(unexpected)} unmapped checkpoint tensors, "
+                f"e.g. {sorted(unexpected)[:5]}. The weight remapper is out of sync with the checkpoint."
+            )
+        missing = sorted(set(params_dict) - loaded_params)
+        if missing:
+            raise ValueError(f"Audio8 TTS Slow AR is missing weights for {missing[:5]} ({len(missing)} total)")
+        logger.info("Loaded %d weights for Audio8TTSSlowARForConditionalGeneration", len(loaded_params))
+
+        # The reference precomputes RoPE in bf16; keeping fp32 cos/sin here
+        # shifts the logits enough to trigger early EOS (same failure mode as
+        # Fish Speech).
+        truncated = 0
+        for module in self.modules():
+            cache = getattr(module, "cos_sin_cache", None)
+            if isinstance(cache, torch.Tensor):
+                module.cos_sin_cache = cache.to(torch.bfloat16).to(cache.dtype)
+                truncated += 1
+        if truncated:
+            logger.info("Audio8 TTS: truncated %d RoPE cos_sin_cache buffers to bf16", truncated)
+
+        return loaded_params
+
+
+__all__ = ["Audio8TTSSlowARForConditionalGeneration"]

vllm_omni/model_executor/models/audio8_tts/codec.py

@@ -0,0 +1,600 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio8 TTS neural audio codec (44.1 kHz, 10 codebooks).
+
+Vendored from the reference ``modeling_arktts_codec.py`` shipped with
+``Audio8/Audio8-TTS-Preview-0.6b`` (Apache-2.0). The module tree and attribute
+names are reproduced verbatim so ``codec.pth`` loads under ``strict=True``.
+
+Inference-only deviations: no ``torch.jit.script`` on Snake (it fights
+``torch.compile`` / CUDA graphs), RoPE cached in a plain dict instead of a
+non-persistent buffer (the pattern that leaves the reference LM with
+uninitialised RoPE under ``from_pretrained`` on transformers >= 5), and
+:func:`build_arktts_codec` prunes the unused encoder or decoder half per stage.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.utils import weight_norm as legacy_weight_norm
+from torch.nn.utils.parametrizations import weight_norm
+
+# The codec's own transformer blocks always use base 10000, independent of the
+# 1e6 rope_base of the DualAR language model.
+_CODEC_ROPE_BASE = 10000.0
+
+
+def _rope_table(length: int, head_dim: int, base: float, device: torch.device) -> Tensor:
+    frequencies = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    phases = torch.outer(torch.arange(length, device=device), frequencies)
+    values = torch.polar(torch.ones_like(phases), phases)
+    return torch.stack((values.real, values.imag), dim=-1).to(torch.bfloat16)
+
+
+def _apply_rope(x: Tensor, values: Tensor) -> Tensor:
+    shaped = x.float().reshape(*x.shape[:-1], -1, 2)
+    values = values.view(1, shaped.shape[1], 1, shaped.shape[3], 2)
+    output = torch.stack(
+        (
+            shaped[..., 0] * values[..., 0] - shaped[..., 1] * values[..., 1],
+            shaped[..., 1] * values[..., 0] + shaped[..., 0] * values[..., 1],
+        ),
+        dim=-1,
+    )
+    return output.flatten(3).to(x.dtype)
+
+
+class ArkttsCodecRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.eps = float(eps)
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        output = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return output.to(x.dtype) * self.weight
+
+
+class ArkttsCodecLayerScale(nn.Module):
+    def __init__(self, dim: int, init_values: float = 1e-2) -> None:
+        super().__init__()
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x * self.gamma
+
+
+class ArkttsCodecTransformerConfig:
+    """Plain option holder for the codec's transformer blocks.
+
+    Deliberately not a dataclass / msgspec Struct: it mirrors a reference
+    dataclass whose only job is to carry constructor arguments, and it is
+    instantiated from Python literals inside this module only.
+    """
+
+    def __init__(
+        self,
+        n_layer: int,
+        n_head: int,
+        dim: int,
+        intermediate_size: int,
+        n_local_heads: int = -1,
+        head_dim: int = 64,
+        rope_base: float = _CODEC_ROPE_BASE,
+        norm_eps: float = 1e-5,
+        channels_first: bool = True,
+    ) -> None:
+        self.n_layer = int(n_layer)
+        self.n_head = int(n_head)
+        self.dim = int(dim)
+        self.intermediate_size = int(intermediate_size)
+        self.n_local_heads = int(n_head) if n_local_heads == -1 else int(n_local_heads)
+        self.head_dim = int(head_dim)
+        self.rope_base = float(rope_base)
+        self.norm_eps = float(norm_eps)
+        self.channels_first = bool(channels_first)
+
+
+class ArkttsCodecAttention(nn.Module):
+    def __init__(self, config: ArkttsCodecTransformerConfig) -> None:
+        super().__init__()
+        total = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+        self.wqkv = nn.Linear(config.dim, total, bias=False)
+        self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
+        self.n_head = config.n_head
+        self.n_local_heads = config.n_local_heads
+        self.head_dim = config.head_dim
+
+    def forward(self, x: Tensor, rope_values: Tensor, mask: Tensor) -> Tensor:
+        batch, length, _ = x.shape
+        query_size = self.n_head * self.head_dim
+        kv_size = self.n_local_heads * self.head_dim
+        query, key, value = self.wqkv(x).split((query_size, kv_size, kv_size), dim=-1)
+        query = query.view(batch, length, self.n_head, self.head_dim)
+        key = key.view(batch, length, self.n_local_heads, self.head_dim)
+        value = value.view(batch, length, self.n_local_heads, self.head_dim)
+        query = _apply_rope(query, rope_values).transpose(1, 2)
+        key = _apply_rope(key, rope_values).transpose(1, 2)
+        value = value.transpose(1, 2)
+        output = F.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=mask,
+            enable_gqa=self.n_local_heads != self.n_head,
+        )
+        output = output.transpose(1, 2).contiguous().view(batch, length, query_size)
+        return self.wo(output)
+
+
+class ArkttsCodecFeedForward(nn.Module):
+    def __init__(self, config: ArkttsCodecTransformerConfig) -> None:
+        super().__init__()
+        self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+        self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+        self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class ArkttsCodecTransformerBlock(nn.Module):
+    def __init__(self, config: ArkttsCodecTransformerConfig) -> None:
+        super().__init__()
+        self.attention = ArkttsCodecAttention(config)
+        self.feed_forward = ArkttsCodecFeedForward(config)
+        self.ffn_norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps)
+        self.attention_norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps)
+        self.attention_layer_scale = ArkttsCodecLayerScale(config.dim)
+        self.ffn_layer_scale = ArkttsCodecLayerScale(config.dim)
+
+    def forward(self, x: Tensor, rope_values: Tensor, mask: Tensor) -> Tensor:
+        hidden = x + self.attention_layer_scale(self.attention(self.attention_norm(x), rope_values, mask))
+        return hidden + self.ffn_layer_scale(self.feed_forward(self.ffn_norm(hidden)))
+
+
+class ArkttsCodecWindowTransformer(nn.Module):
+    """Causal transformer with a sliding attention window."""
+
+    def __init__(
+        self,
+        config: ArkttsCodecTransformerConfig,
+        input_dim: int,
+        window_size: int | None,
+    ) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList([ArkttsCodecTransformerBlock(config) for _ in range(config.n_layer)])
+        self.norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps)
+        self.window_size = window_size
+        self.channels_first = config.channels_first
+        self.input_proj = nn.Linear(input_dim, config.dim) if input_dim != config.dim else nn.Identity()
+        self.output_proj = nn.Linear(config.dim, input_dim) if input_dim != config.dim else nn.Identity()
+        self.look_ahead_conv = nn.Identity()
+        self.head_dim = config.head_dim
+        self.rope_base = config.rope_base
+        self._mask_cache: dict[tuple[int, str], Tensor] = {}
+        self._rope_cache: dict[tuple[int, str], Tensor] = {}
+
+    def _mask(self, length: int, device: torch.device) -> Tensor:
+        key = (length, str(device))
+        cached = self._mask_cache.get(key)
+        if cached is not None:
+            return cached
+        row = torch.arange(length, device=device)[:, None]
+        column = torch.arange(length, device=device)[None, :]
+        mask = column <= row
+        if self.window_size is not None:
+            mask &= column >= (row - self.window_size + 1).clamp_min(0)
+        mask = mask[None, None]
+        self._mask_cache[key] = mask
+        return mask
+
+    def _rope(self, length: int, device: torch.device) -> Tensor:
+        key = (length, str(device))
+        cached = self._rope_cache.get(key)
+        if cached is not None:
+            return cached
+        values = _rope_table(length, self.head_dim, self.rope_base, device)
+        self._rope_cache[key] = values
+        return values
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.channels_first:
+            x = x.transpose(1, 2)
+        x = self.look_ahead_conv(self.input_proj(x))
+        length = int(x.shape[1])
+        mask = self._mask(length, x.device)
+        rope_values = self._rope(length, x.device)
+        for layer in self.layers:
+            x = layer(x, rope_values, mask)
+        x = self.output_proj(self.norm(x))
+        return x.transpose(1, 2) if self.channels_first else x
+
+
+def _extra_padding(x: Tensor, kernel_size: int, stride: int, padding_total: int = 0) -> int:
+    length = x.shape[-1]
+    frames = (length - kernel_size + padding_total) / stride + 1
+    ideal = (math.ceil(frames) - 1) * stride + kernel_size - padding_total
+    return ideal - length
+
+
+class ArkttsCausalConv1d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        dilation: int = 1,
+        stride: int = 1,
+        groups: int = 1,
+    ) -> None:
+        super().__init__()
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, groups=groups)
+        self.stride = stride
+        self.kernel_size = (kernel_size - 1) * dilation + 1
+        self.padding = self.kernel_size - self.stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        right = _extra_padding(x, self.kernel_size, self.stride, self.padding)
+        return self.conv(F.pad(x, (self.padding, right))).contiguous()
+
+    def apply_weight_norm(self) -> ArkttsCausalConv1d:
+        self.conv = weight_norm(self.conv)
+        return self
+
+
+class ArkttsCausalConvTranspose1d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+    ) -> None:
+        super().__init__()
+        self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation)
+        self.stride = stride
+        self.kernel_size = kernel_size
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.conv(x)
+        crop = self.kernel_size - self.stride
+        return x[..., : x.shape[-1] - crop].contiguous() if crop else x.contiguous()
+
+    def apply_weight_norm(self) -> ArkttsCausalConvTranspose1d:
+        self.conv = weight_norm(self.conv)
+        return self
+
+
+def _causal_wn_conv(*args: Any, **kwargs: Any) -> ArkttsCausalConv1d:
+    return ArkttsCausalConv1d(*args, **kwargs).apply_weight_norm()
+
+
+def _causal_wn_transpose(*args: Any, **kwargs: Any) -> ArkttsCausalConvTranspose1d:
+    return ArkttsCausalConvTranspose1d(*args, **kwargs).apply_weight_norm()
+
+
+class ArkttsSnake1d(nn.Module):
+    def __init__(self, channels: int) -> None:
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x: Tensor) -> Tensor:
+        shape = x.shape
+        x = x.reshape(shape[0], shape[1], -1)
+        x = x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2)
+        return x.reshape(shape)
+
+
+class ArkttsResidualUnit(nn.Module):
+    def __init__(self, dim: int, dilation: int) -> None:
+        super().__init__()
+        self.block = nn.Sequential(
+            ArkttsSnake1d(dim),
+            _causal_wn_conv(dim, dim, kernel_size=7, dilation=dilation),
+            ArkttsSnake1d(dim),
+            _causal_wn_conv(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        output = self.block(x)
+        difference = x.shape[-1] - output.shape[-1]
+        if difference > 0:
+            x = x[..., :-difference]
+        return x + output
+
+
+class ArkttsEncoderBlock(nn.Module):
+    def __init__(self, dim: int, stride: int, transformer_layers: int) -> None:
+        super().__init__()
+        modules: list[nn.Module] = [
+            ArkttsResidualUnit(dim // 2, 1),
+            ArkttsResidualUnit(dim // 2, 3),
+            ArkttsResidualUnit(dim // 2, 9),
+            ArkttsSnake1d(dim // 2),
+            _causal_wn_conv(dim // 2, dim, kernel_size=2 * stride, stride=stride),
+        ]
+        if transformer_layers:
+            config = ArkttsCodecTransformerConfig(
+                n_layer=transformer_layers,
+                n_head=dim // 64,
+                dim=dim,
+                intermediate_size=dim * 3,
+            )
+            modules.append(ArkttsCodecWindowTransformer(config, dim, window_size=512))
+        else:
+            modules.append(nn.Identity())
+        self.block = nn.Sequential(*modules)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.block(x)
+
+
+class ArkttsEncoder(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        dim = 64
+        modules: list[nn.Module] = [_causal_wn_conv(1, dim, kernel_size=7)]
+        for stride, transformer_layers in zip((2, 4, 8, 8), (0, 0, 0, 4), strict=True):
+            dim *= 2
+            modules.append(ArkttsEncoderBlock(dim, stride, transformer_layers))
+        modules.extend((ArkttsSnake1d(dim), _causal_wn_conv(dim, 1024, kernel_size=3)))
+        self.block = nn.Sequential(*modules)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.block(x)
+
+
+class ArkttsDecoderBlock(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int, stride: int) -> None:
+        super().__init__()
+        self.block = nn.Sequential(
+            ArkttsSnake1d(input_dim),
+            _causal_wn_transpose(input_dim, output_dim, kernel_size=2 * stride, stride=stride),
+            ArkttsResidualUnit(output_dim, 1),
+            ArkttsResidualUnit(output_dim, 3),
+            ArkttsResidualUnit(output_dim, 9),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.block(x)
+
+
+class ArkttsDecoder(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        channels = 1536
+        modules: list[nn.Module] = [_causal_wn_conv(1024, channels, kernel_size=7)]
+        output_dim = channels
+        for index, stride in enumerate((8, 8, 4, 2)):
+            input_dim = channels // (2**index)
+            output_dim = channels // (2 ** (index + 1))
+            modules.append(ArkttsDecoderBlock(input_dim, output_dim, stride))
+        modules.extend(
+            (
+                ArkttsSnake1d(output_dim),
+                _causal_wn_conv(output_dim, 1, kernel_size=7),
+                nn.Tanh(),
+            )
+        )
+        self.model = nn.Sequential(*modules)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.model(x)
+
+
+class ArkttsVectorQuantizer(nn.Module):
+    def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int) -> None:
+        super().__init__()
+        self.codebook_size = int(codebook_size)
+        self.codebook_dim = int(codebook_dim)
+        self.in_proj = legacy_weight_norm(nn.Conv1d(input_dim, codebook_dim, kernel_size=1))
+        self.out_proj = legacy_weight_norm(nn.Conv1d(codebook_dim, input_dim, kernel_size=1))
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+    def decode_code(self, indices: Tensor) -> Tensor:
+        return F.embedding(indices, self.codebook.weight).transpose(1, 2)
+
+    def decode_latents(self, latents: Tensor) -> tuple[Tensor, Tensor]:
+        batch, _, length = latents.shape
+        flattened = F.normalize(latents.transpose(1, 2).reshape(batch * length, -1))
+        codebook = F.normalize(self.codebook.weight)
+        distances = (
+            flattened.pow(2).sum(1, keepdim=True)
+            - 2 * flattened @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = (-distances).argmax(1).view(batch, length)
+        return self.decode_code(indices), indices
+
+    def forward(self, z: Tensor) -> tuple[Tensor, Tensor]:
+        projected = self.in_proj(z)
+        quantized, indices = self.decode_latents(projected)
+        return self.out_proj(quantized), indices
+
+
+class ArkttsResidualQuantizer(nn.Module):
+    def __init__(self, input_dim: int, n_codebooks: int, codebook_size: int, codebook_dim: int) -> None:
+        super().__init__()
+        self.n_codebooks = int(n_codebooks)
+        self.codebook_size = int(codebook_size)
+        self.quantizers = nn.ModuleList(
+            [ArkttsVectorQuantizer(input_dim, codebook_size, codebook_dim) for _ in range(n_codebooks)]
+        )
+
+    def forward(self, z: Tensor) -> tuple[Tensor, Tensor]:
+        quantized_sum: Tensor | float = 0.0
+        residual = z
+        codes: list[Tensor] = []
+        for quantizer in self.quantizers:
+            quantized, indices = quantizer(residual)
+            quantized_sum = quantized_sum + quantized
+            residual = residual - quantized
+            codes.append(indices)
+        return quantized_sum, torch.stack(codes, dim=1)
+
+    def from_codes(self, codes: Tensor) -> Tensor:
+        output: Tensor | float = 0.0
+        for index in range(codes.shape[1]):
+            projected = self.quantizers[index].decode_code(codes[:, index])
+            output = output + self.quantizers[index].out_proj(projected)
+        return output
+
+
+class ArkttsConvNeXtBlock(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.dwconv = ArkttsCausalConv1d(dim, dim, kernel_size=7, groups=dim)
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = nn.Parameter(1e-6 * torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        residual = x
+        x = self.dwconv(x).permute(0, 2, 1)
+        x = self.pwconv2(self.act(self.pwconv1(self.norm(x))))
+        x = (self.gamma * x).permute(0, 2, 1)
+        return residual + x
+
+
+class ArkttsDownsampleQuantizer(nn.Module):
+    """1 semantic codebook (4096) + 9 residual codebooks (1024), 4x downsampled."""
+
+    def __init__(
+        self,
+        post_n_layer: int = 8,
+        post_n_head: int = 16,
+        post_n_local_heads: int = 8,
+        post_intermediate_size: int = 1216,
+    ) -> None:
+        super().__init__()
+        self.semantic_quantizer = ArkttsResidualQuantizer(1024, 1, 4096, 8)
+        self.quantizer = ArkttsResidualQuantizer(1024, 9, 1024, 8)
+        self.downsample = nn.Sequential(
+            nn.Sequential(ArkttsCausalConv1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)),
+            nn.Sequential(ArkttsCausalConv1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)),
+        )
+        self.upsample = nn.Sequential(
+            nn.Sequential(ArkttsCausalConvTranspose1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)),
+            nn.Sequential(ArkttsCausalConvTranspose1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)),
+        )
+        self.pre_module = ArkttsCodecWindowTransformer(
+            ArkttsCodecTransformerConfig(n_layer=8, n_head=16, dim=1024, intermediate_size=3072),
+            1024,
+            window_size=128,
+        )
+        self.post_module = ArkttsCodecWindowTransformer(
+            ArkttsCodecTransformerConfig(
+                n_layer=post_n_layer,
+                n_head=post_n_head,
+                n_local_heads=post_n_local_heads,
+                dim=1024,
+                intermediate_size=post_intermediate_size,
+            ),
+            1024,
+            window_size=128,
+        )
+        self.semantic_predictor_module = nn.Identity()
+
+    def forward(self, z: Tensor) -> tuple[Tensor, Tensor]:
+        original_length = z.shape[-1]
+        z = self.pre_module(self.downsample(z))
+        semantic, semantic_codes = self.semantic_quantizer(z)
+        residual, residual_codes = self.quantizer(z - semantic)
+        z = self.upsample(self.post_module(semantic + residual))
+        difference = original_length - z.shape[-1]
+        if difference > 0:
+            z = F.pad(z, (difference, 0))
+        elif difference < 0:
+            z = z[..., -difference:]
+        return z, torch.cat((semantic_codes, residual_codes), dim=1)
+
+    def decode(self, indices: Tensor) -> Tensor:
+        # Cloned because the clamps below are in-place and this is the caller's
+        # codes buffer; once per decode chunk, not per AR step.
+        indices = indices.clone()
+        indices[:, 0].clamp_(0, self.semantic_quantizer.codebook_size - 1)
+        indices[:, 1:].clamp_(0, self.quantizer.codebook_size - 1)
+        semantic = self.semantic_quantizer.from_codes(indices[:, :1])
+        residual = self.quantizer.from_codes(indices[:, 1:])
+        return self.upsample(self.post_module(semantic + residual))
+
+
+class ArkttsCodec(nn.Module):
+    """Audio8 TTS codec: waveform <-> 10 codebooks at ~21.5 frames/s."""
+
+    sample_rate = 44100
+    frame_length = 2048
+
+    def __init__(
+        self,
+        post_n_layer: int = 8,
+        post_n_head: int = 16,
+        post_n_local_heads: int = 8,
+        post_intermediate_size: int = 1216,
+    ) -> None:
+        super().__init__()
+        self.encoder = ArkttsEncoder()
+        self.quantizer = ArkttsDownsampleQuantizer(
+            post_n_layer=post_n_layer,
+            post_n_head=post_n_head,
+            post_n_local_heads=post_n_local_heads,
+            post_intermediate_size=post_intermediate_size,
+        )
+        self.decoder = ArkttsDecoder()
+
+    @torch.inference_mode()
+    def encode(self, audio: Tensor, audio_lengths: Tensor | None = None) -> tuple[Tensor, Tensor]:
+        """Encode ``[B, 1, samples]`` (or ``[B, samples]``) to ``[B, 10, frames]``."""
+        if audio.ndim == 2:
+            audio = audio[:, None]
+        if audio.ndim != 3 or audio.shape[1] != 1:
+            raise ValueError(f"audio must have shape [B, 1, samples], got {tuple(audio.shape)}")
+        original_length = int(audio.shape[-1])
+        right = math.ceil(original_length / self.frame_length) * self.frame_length - original_length
+        if right:
+            audio = F.pad(audio, (0, right))
+        if audio_lengths is None:
+            audio_lengths = torch.full((audio.shape[0],), original_length, device=audio.device, dtype=torch.long)
+        _, codes = self.quantizer(self.encoder(audio))
+        code_lengths = torch.ceil(audio_lengths.float() / self.frame_length).long()
+        max_codes = int(codes.shape[-1])
+        code_lengths = code_lengths.clamp_max(max_codes)
+        frame_index = torch.arange(max_codes, device=codes.device)
+        valid = frame_index[None, None, :] < code_lengths[:, None, None]
+        return torch.where(valid, codes, torch.full_like(codes, -1)), code_lengths
+
+    @torch.inference_mode()
+    def decode(self, codes: Tensor) -> Tensor:
+        """Decode ``[B, 10, frames]`` to ``[B, 1, samples]``."""
+        return self.decoder(self.quantizer.decode(codes.long()))
+
+
+def build_arktts_codec(
+    *,
+    post_n_layer: int = 8,
+    post_n_head: int = 16,
+    post_n_local_heads: int = 8,
+    post_intermediate_size: int = 1216,
+) -> ArkttsCodec:
+    """Construct the codec with uninitialised weights, on CPU, in eval mode."""
+    codec = ArkttsCodec(
+        post_n_layer=post_n_layer,
+        post_n_head=post_n_head,
+        post_n_local_heads=post_n_local_heads,
+        post_intermediate_size=post_intermediate_size,
+    )
+    codec.eval()
+    return codec
+
+
+__all__ = ["ArkttsCodec", "build_arktts_codec"]

vllm_omni/model_executor/models/audio8_tts/codec_utils.py

@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Loading / caching helpers for the Audio8 TTS ``codec.pth`` checkpoint.
+
+Both stages need the codec but for opposite halves: Stage 0 encodes reference
+audio for zero-shot voice cloning, Stage 1 decodes generated codes to waveform.
+The unused half is pruned before the module is moved to the device so a stage
+never pays GPU memory for weights it cannot reach.
+"""
+
+from __future__ import annotations
+
+import os
+from functools import lru_cache
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.utils.parametrize import remove_parametrizations
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.audio8_tts.codec import ArkttsCodec, build_arktts_codec
+from vllm_omni.model_executor.models.audio8_tts.configuration_audio8_tts import (
+    ARKTTS_CODEC_FRAME_SIZE,
+    ARKTTS_CODEC_SAMPLE_RATE,
+)
+
+logger = init_logger(__name__)
+
+CODEC_FILENAME = "codec.pth"
+
+_codec_cache: dict[tuple[str, str, str, str], ArkttsCodec] = {}
+
+
+def resolve_codec_path(model_path: str, filename: str = CODEC_FILENAME) -> str:
+    """Return a local path to ``filename``, downloading from the Hub if needed."""
+    local = os.path.join(model_path, filename)
+    if os.path.exists(local):
+        return local
+
+    from transformers.utils.hub import cached_file
+
+    cached = cached_file(model_path, filename)
+    if cached is None or not os.path.exists(cached):
+        raise FileNotFoundError(
+            f"{filename} not found for {model_path}; the Audio8 TTS checkpoint must ship its neural audio codec."
+        )
+    return cached
+
+
+def bake_weight_norm(module: nn.Module) -> int:
+    """Fold weight-norm parametrisations into plain weights (inference only)."""
+    baked = 0
+    for submodule in module.modules():
+        parametrizations = getattr(submodule, "parametrizations", None)
+        if not parametrizations:
+            continue
+        for name in list(parametrizations.keys()):
+            remove_parametrizations(submodule, name, leave_parametrized=True)
+            baked += 1
+    return baked
+
+
+def load_arktts_codec(
+    model_path: str,
+    *,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype = torch.float32,
+    role: str = "decode",
+    post_n_layer: int = 8,
+    post_n_head: int = 16,
+    post_n_local_heads: int = 8,
+    post_intermediate_size: int = 1216,
+) -> ArkttsCodec:
+    """Load ``codec.pth`` for ``role`` in {"encode", "decode", "both"}.
+
+    Instances are cached per (path, device, dtype, role) because both the Stage-0
+    voice-clone path and the Stage-1 decoder may ask for the same codec many
+    times per request.
+    """
+    if role not in {"encode", "decode", "both"}:
+        raise ValueError(f"role must be encode/decode/both, got {role!r}")
+    device = torch.device(device)
+    cache_key = (model_path, str(device), str(dtype), role)
+    cached_codec = _codec_cache.get(cache_key)
+    if cached_codec is not None:
+        return cached_codec
+
+    codec_path = resolve_codec_path(model_path)
+    codec = build_arktts_codec(
+        post_n_layer=post_n_layer,
+        post_n_head=post_n_head,
+        post_n_local_heads=post_n_local_heads,
+        post_intermediate_size=post_intermediate_size,
+    )
+
+    state = torch.load(codec_path, map_location="cpu", weights_only=True)
+    if "state_dict" in state:
+        state = state["state_dict"]
+    if any("generator." in key for key in state):
+        state = {key.replace("generator.", ""): value for key, value in state.items() if "generator." in key}
+    state = {key: value for key, value in state.items() if not key.endswith(("freqs_cis", "causal_mask"))}
+    codec.load_state_dict(state, strict=True)
+
+    baked = bake_weight_norm(codec)
+    if role == "encode":
+        codec.decoder = None
+    elif role == "decode":
+        codec.encoder = None
+        codec.quantizer.pre_module = None
+        codec.quantizer.downsample = None
+
+    codec = codec.to(device=device, dtype=dtype)
+    codec.eval()
+    _codec_cache[cache_key] = codec
+    logger.info(
+        "Loaded Audio8 TTS codec from %s (role=%s, device=%s, dtype=%s, baked_weight_norms=%d)",
+        codec_path,
+        role,
+        device,
+        dtype,
+        baked,
+    )
+    return codec
+
+
+@lru_cache(maxsize=8)
+def _resample_transform(source_sr: int, target_sr: int, device: torch.device, dtype: torch.dtype):
+    import torchaudio
+
+    return torchaudio.transforms.Resample(source_sr, target_sr).to(device=device, dtype=dtype)
+
+
+def prepare_reference_waveform(
+    wav_samples: list[float] | np.ndarray | torch.Tensor,
+    sample_rate: int,
+    *,
+    device: torch.device,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Down-mix to mono and resample reference audio to the codec's rate.
+
+    Returns a 1-D waveform tensor.
+    """
+    if isinstance(wav_samples, torch.Tensor):
+        wav = wav_samples.detach()
+    else:
+        wav = torch.as_tensor(np.asarray(wav_samples))
+    wav = wav.to(device=device, dtype=dtype)
+
+    if wav.ndim == 2:
+        # Accept [channels, samples] and [samples, channels].
+        if wav.shape[0] <= 8 and wav.shape[1] > wav.shape[0]:
+            wav = wav.mean(dim=0)
+        elif wav.shape[-1] <= 8 and wav.shape[0] > wav.shape[-1]:
+            wav = wav.mean(dim=-1)
+        else:
+            wav = wav.mean(dim=0)
+    elif wav.ndim > 2:
+        wav = wav.reshape(-1, wav.shape[-1]).mean(dim=0)
+    wav = wav.reshape(-1)
+
+    if wav.numel() == 0:
+        raise ValueError("Reference audio must not be empty")
+    if int(sample_rate) <= 0:
+        raise ValueError(f"Reference audio sample rate must be positive, got {sample_rate}")
+    if int(sample_rate) != ARKTTS_CODEC_SAMPLE_RATE:
+        transform = _resample_transform(int(sample_rate), ARKTTS_CODEC_SAMPLE_RATE, wav.device, wav.dtype)
+        wav = transform(wav.unsqueeze(0)).squeeze(0)
+    return wav.contiguous()
+
+
+@torch.no_grad()
+def encode_reference_audio_codes(
+    model_path: str,
+    wav_samples: list[float] | np.ndarray | torch.Tensor,
+    sample_rate: int,
+    *,
+    device: torch.device | str,
+    **codec_kwargs: int,
+) -> torch.Tensor:
+    """Encode reference audio into codec codes.
+
+    Returns:
+        ``[frames, num_codebooks]`` int64 codes on ``device``.
+    """
+    device = torch.device(device)
+    codec = load_arktts_codec(model_path, device=device, dtype=torch.float32, role="encode", **codec_kwargs)
+    wav = prepare_reference_waveform(wav_samples, sample_rate, device=device)
+    lengths = torch.tensor([wav.numel()], device=device, dtype=torch.long)
+    codes, code_lengths = codec.encode(wav.reshape(1, 1, -1), lengths)
+    frames = int(code_lengths[0].item())
+    codes_fq = codes[0, :, :frames].transpose(0, 1).to(dtype=torch.long).contiguous()
+    logger.info(
+        "Encoded Audio8 TTS reference audio: %d samples @ %d Hz -> frames=%d codebooks=%d",
+        int(wav.numel()),
+        int(sample_rate),
+        int(codes_fq.shape[0]),
+        int(codes_fq.shape[1]),
+    )
+    return codes_fq
+
+
+def estimate_reference_code_frames(num_samples: int, sample_rate: int) -> int:
+    """Frames the codec will emit for ``num_samples`` at ``sample_rate``.
+
+    Used by the serving layer to size the prompt placeholder without paying for
+    a codec encode on the API thread.
+    """
+    if num_samples <= 0 or sample_rate <= 0:
+        raise ValueError("Reference audio must have a positive length and sample rate")
+    resampled = max(1, -(-num_samples * ARKTTS_CODEC_SAMPLE_RATE // sample_rate))
+    return max(1, -(-resampled // ARKTTS_CODEC_FRAME_SIZE))
+
+
+__all__ = [
+    "CODEC_FILENAME",
+    "bake_weight_norm",
+    "encode_reference_audio_codes",
+    "estimate_reference_code_frames",
+    "load_arktts_codec",
+    "prepare_reference_waveform",
+    "resolve_codec_path",
+]

vllm_omni/model_executor/models/audio8_tts/configuration_audio8_tts.py

@@ -0,0 +1,267 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Configuration classes for Audio8 TTS Preview (``model_type = "arktts"``).
+
+Audio8's HuggingFace config is flat and uses DualAR field names (``dim``,
+``n_layer``, ``n_head``, ``fast_dim``, ...); this module re-exports them under
+Qwen2 attribute names so vLLM's ``Qwen2Model`` consumes the Slow AR half, and
+splits the Fast AR half into its own sub-config. The backbone is Qwen2 rather
+than Qwen3 (Fish Speech) because Audio8 uses ``qkv_bias=true`` with no q/k
+norm -- exactly the Qwen2 attention shape.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from transformers import PretrainedConfig
+
+#: Codec frame stride in waveform samples (encoder rates 2*4*8*8 = 512 times the
+#: quantizer's 2*2 downsample). 44100 / 2048 ~= 21.5 frames per second.
+ARKTTS_CODEC_FRAME_SIZE = 2048
+ARKTTS_CODEC_SAMPLE_RATE = 44100
+#: Only the first codebook is a 4096-entry semantic codebook; codebooks 1..N-1
+#: are 1024-entry residual codebooks (see ``ArkttsDownsampleQuantizer``).
+ARKTTS_SEMANTIC_CODEBOOK_SIZE = 4096
+ARKTTS_RESIDUAL_CODEBOOK_SIZE = 1024
+
+
+class Audio8TTSSlowARConfig(PretrainedConfig):
+    """Slow AR config exposed with Qwen2-compatible attribute names."""
+
+    model_type = "arktts_slow_ar"
+
+    def __init__(
+        self,
+        vocab_size: int = 155776,
+        dim: int = 896,
+        n_head: int = 14,
+        n_local_heads: int = 2,
+        head_dim: int = 64,
+        n_layer: int = 24,
+        intermediate_size: int = 4864,
+        max_seq_len: int = 2048,
+        rope_base: float = 1_000_000.0,
+        norm_eps: float = 1e-6,
+        attention_qkv_bias: bool = True,
+        attention_qk_norm: bool = False,
+        tie_word_embeddings: bool = True,
+        codebook_size: int = ARKTTS_SEMANTIC_CODEBOOK_SIZE,
+        num_codebooks: int = 10,
+        semantic_begin_id: int = 151678,
+        semantic_end_id: int = 155773,
+        eos_token_id: int = 151645,
+        pad_token_id: int = 151643,
+        **kwargs: Any,
+    ) -> None:
+        # Audio8 field names -> standard Transformers / Qwen2 names.
+        self.hidden_size = int(dim)
+        self.num_attention_heads = int(n_head)
+        self.num_key_value_heads = int(n_local_heads)
+        self.head_dim = int(head_dim)
+        self.num_hidden_layers = int(n_layer)
+        self.intermediate_size = int(intermediate_size)
+        self.max_position_embeddings = int(max_seq_len)
+        self.rms_norm_eps = float(norm_eps)
+        self.hidden_act = "silu"
+        # Qwen2Attention hardcodes qkv bias=True / o_proj bias=False, which is
+        # what Audio8 TTS ships. Keep the flags so a future checkpoint that
+        # flips them fails loudly in the model rather than loading silently.
+        self.attention_qkv_bias = bool(attention_qkv_bias)
+        self.attention_qk_norm = bool(attention_qk_norm)
+        # Set explicitly (rather than relying on vLLM's set_default_rope_theta)
+        # so that a checkpoint with a non-default rope_base is honoured.
+        self.rope_parameters = {"rope_type": "default", "rope_theta": float(rope_base)}
+
+        # Codec / codebook fields.
+        self.codebook_size = int(codebook_size)
+        self.num_codebooks = int(num_codebooks)
+        self.semantic_begin_id = int(semantic_begin_id)
+        self.semantic_end_id = int(semantic_end_id)
+
+        super().__init__(
+            vocab_size=int(vocab_size),
+            tie_word_embeddings=bool(tie_word_embeddings),
+            eos_token_id=int(eos_token_id),
+            pad_token_id=int(pad_token_id),
+            **kwargs,
+        )
+
+
+class Audio8TTSFastARConfig(PretrainedConfig):
+    """Fast AR config: the ``n_fast_layer`` residual-codebook predictor."""
+
+    model_type = "arktts_fast_ar"
+
+    def __init__(
+        self,
+        codebook_size: int = ARKTTS_SEMANTIC_CODEBOOK_SIZE,
+        num_codebooks: int = 10,
+        fast_dim: int = 896,
+        fast_n_head: int = 14,
+        fast_n_local_heads: int = 2,
+        fast_head_dim: int = 64,
+        n_fast_layer: int = 4,
+        fast_intermediate_size: int = 4864,
+        fast_attention_qkv_bias: bool = False,
+        fast_attention_qk_norm: bool = False,
+        rope_base: float = 1_000_000.0,
+        norm_eps: float = 1e-6,
+        **kwargs: Any,
+    ) -> None:
+        self.hidden_size = int(fast_dim)
+        self.num_attention_heads = int(fast_n_head)
+        self.num_key_value_heads = int(fast_n_local_heads)
+        self.head_dim = int(fast_head_dim)
+        self.num_hidden_layers = int(n_fast_layer)
+        self.intermediate_size = int(fast_intermediate_size)
+        # The Fast AR sequence is [slow hidden state, code_0, ..., code_{N-1}].
+        self.max_position_embeddings = int(num_codebooks)
+        self.rms_norm_eps = float(norm_eps)
+        self.hidden_act = "silu"
+        self.attention_qkv_bias = bool(fast_attention_qkv_bias)
+        self.attention_qk_norm = bool(fast_attention_qk_norm)
+        self.rope_theta = float(rope_base)
+        self.num_codebooks = int(num_codebooks)
+
+        super().__init__(vocab_size=int(codebook_size), **kwargs)
+
+
+class Audio8TTSConfig(PretrainedConfig):
+    """Top-level Audio8 TTS config (``model_type = "arktts"``).
+
+    Accepts the flat HF checkpoint fields and derives ``text_config`` (Slow AR)
+    and ``fast_ar_config`` (Fast AR).  ``get_text_config()`` returns the Slow AR
+    config, which is what ``Qwen2Model`` reads.
+    """
+
+    model_type = "arktts"
+    sub_configs = {
+        "text_config": Audio8TTSSlowARConfig,
+        "fast_ar_config": Audio8TTSFastARConfig,
+    }
+
+    def __init__(
+        self,
+        text_config: dict | Audio8TTSSlowARConfig | None = None,
+        fast_ar_config: dict | Audio8TTSFastARConfig | None = None,
+        vocab_size: int = 155776,
+        dim: int = 896,
+        n_head: int = 14,
+        n_local_heads: int = 2,
+        head_dim: int = 64,
+        n_layer: int = 24,
+        intermediate_size: int = 4864,
+        max_seq_len: int = 2048,
+        rope_base: float = 1_000_000.0,
+        norm_eps: float = 1e-6,
+        attention_qkv_bias: bool = True,
+        attention_qk_norm: bool = False,
+        tie_word_embeddings: bool = True,
+        codebook_size: int = ARKTTS_SEMANTIC_CODEBOOK_SIZE,
+        num_codebooks: int = 10,
+        semantic_begin_id: int = 151678,
+        semantic_end_id: int = 155773,
+        n_fast_layer: int = 4,
+        fast_dim: int = 896,
+        fast_n_head: int = 14,
+        fast_n_local_heads: int = 2,
+        fast_head_dim: int = 64,
+        fast_intermediate_size: int = 4864,
+        fast_attention_qkv_bias: bool = False,
+        fast_attention_qk_norm: bool = False,
+        norm_fastlayer_input: bool = True,
+        codec_filename: str = "codec.pth",
+        codec_sample_rate: int = ARKTTS_CODEC_SAMPLE_RATE,
+        codec_frame_size: int = ARKTTS_CODEC_FRAME_SIZE,
+        codec_post_n_layer: int = 8,
+        codec_post_n_head: int = 16,
+        codec_post_n_local_heads: int = 8,
+        codec_post_intermediate_size: int = 1216,
+        ras_window_size: int = 10,
+        ras_temperature: float = 1.0,
+        ras_top_p: float = 0.9,
+        eos_token_id: int = 151645,
+        pad_token_id: int = 151643,
+        **kwargs: Any,
+    ) -> None:
+        if isinstance(text_config, dict):
+            text_config = Audio8TTSSlowARConfig(**text_config)
+        if isinstance(fast_ar_config, dict):
+            fast_ar_config = Audio8TTSFastARConfig(**fast_ar_config)
+
+        self.text_config = text_config or Audio8TTSSlowARConfig(
+            vocab_size=vocab_size,
+            dim=dim,
+            n_head=n_head,
+            n_local_heads=n_local_heads,
+            head_dim=head_dim,
+            n_layer=n_layer,
+            intermediate_size=intermediate_size,
+            max_seq_len=max_seq_len,
+            rope_base=rope_base,
+            norm_eps=norm_eps,
+            attention_qkv_bias=attention_qkv_bias,
+            attention_qk_norm=attention_qk_norm,
+            tie_word_embeddings=tie_word_embeddings,
+            codebook_size=codebook_size,
+            num_codebooks=num_codebooks,
+            semantic_begin_id=semantic_begin_id,
+            semantic_end_id=semantic_end_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+        )
+        self.fast_ar_config = fast_ar_config or Audio8TTSFastARConfig(
+            codebook_size=codebook_size,
+            num_codebooks=num_codebooks,
+            fast_dim=fast_dim,
+            fast_n_head=fast_n_head,
+            fast_n_local_heads=fast_n_local_heads,
+            fast_head_dim=fast_head_dim,
+            n_fast_layer=n_fast_layer,
+            fast_intermediate_size=fast_intermediate_size,
+            fast_attention_qkv_bias=fast_attention_qkv_bias,
+            fast_attention_qk_norm=fast_attention_qk_norm,
+            rope_base=rope_base,
+            norm_eps=norm_eps,
+        )
+
+        # Fields the stages read directly off the top-level config.
+        self.codebook_size = int(codebook_size)
+        self.num_codebooks = int(num_codebooks)
+        self.semantic_begin_id = int(semantic_begin_id)
+        self.semantic_end_id = int(semantic_end_id)
+        self.norm_fastlayer_input = bool(norm_fastlayer_input)
+        self.codec_filename = str(codec_filename)
+        self.codec_sample_rate = int(codec_sample_rate)
+        self.codec_frame_size = int(codec_frame_size)
+        self.codec_post_n_layer = int(codec_post_n_layer)
+        self.codec_post_n_head = int(codec_post_n_head)
+        self.codec_post_n_local_heads = int(codec_post_n_local_heads)
+        self.codec_post_intermediate_size = int(codec_post_intermediate_size)
+        # Repetition-Aware Sampling (RAS): resample a repeated semantic token
+        # from a flatter distribution instead of masking it.
+        self.ras_window_size = int(ras_window_size)
+        self.ras_temperature = float(ras_temperature)
+        self.ras_top_p = float(ras_top_p)
+
+        super().__init__(
+            eos_token_id=int(eos_token_id),
+            pad_token_id=int(pad_token_id),
+            tie_word_embeddings=bool(tie_word_embeddings),
+            **kwargs,
+        )
+
+    def get_text_config(self, *args: Any, **kwargs: Any) -> Audio8TTSSlowARConfig:
+        return self.text_config
+
+
+__all__ = [
+    "ARKTTS_CODEC_FRAME_SIZE",
+    "ARKTTS_CODEC_SAMPLE_RATE",
+    "ARKTTS_RESIDUAL_CODEBOOK_SIZE",
+    "ARKTTS_SEMANTIC_CODEBOOK_SIZE",
+    "Audio8TTSConfig",
+    "Audio8TTSFastARConfig",
+    "Audio8TTSSlowARConfig",
+]

vllm_omni/model_executor/models/audio8_tts/pipeline.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio8 TTS Preview pipeline topology.
+
+Stage 0: ``audio8_tts_slow_ar``       -- text -> semantic tokens + codec codes.
+Stage 1: ``audio8_tts_codec_decoder`` -- codec codes -> 44.1 kHz waveform.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.audio8_tts"
+
+#: ``<|im_end|>``: the Slow AR ends the utterance with an end-of-turn token.
+AUDIO8_TTS_EOS_TOKEN_ID = 151645
+
+AUDIO8_TTS_PIPELINE = PipelineConfig(
+    model_type="arktts",
+    default_deploy_config_name="audio8_tts.yaml",
+    model_arch="Audio8TTSSlowARForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="audio8_tts_slow_ar",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.slow_ar_to_codec_decoder_async_chunk"),
+            sampling_constraints={
+                "detokenize": False,
+                "stop_token_ids": [AUDIO8_TTS_EOS_TOKEN_ID],
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="audio8_tts_codec_decoder",
+            model_arch="Audio8TTSCodecDecoder",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)

vllm_omni/model_executor/models/audio8_tts/prompt_utils.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio8 TTS prompt construction.
+
+Reproduces ``ArkttsProcessor._prompt_segments`` from the reference checkpoint,
+so tokenisation is byte-identical to HF inference.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import regex as re
+
+SYSTEM_PROMPT_TEXT_ONLY = "convert the provided text to speech"
+SYSTEM_PROMPT_CLONE_PREFIX = "convert the provided text to speech reference to the following:\n\nText:\n"
+SYSTEM_PROMPT_CLONE_SUFFIX = "\n\nSpeech:\n"
+
+_SPEAKER_TAG_PATTERN = re.compile(r"<\|speaker:\d+\|>")
+_LEGACY_SPEAKER_TAG_PATTERN = re.compile(r"<speaker:(\d+)>")
+_CONTROL_TOKEN_PATTERN = re.compile(r"<\|[^>]+\|>")
+
+
+def clean_text(text: str) -> str:
+    """Collapse all whitespace runs to single spaces (reference ``_clean_text``)."""
+    return " ".join(str(text).strip().split())
+
+
+def normalize_text(text: str, *, add_default_speaker: bool = False) -> str:
+    """Clean text, normalise speaker tags, reject other control tokens.
+
+    ``<|speaker:N|>`` is allowed; any other ``<|...|>`` token raises so callers
+    cannot inject ``<|semantic:...|>`` / ``<|im_end|>`` into a prompt.
+    """
+    normalized = clean_text(_LEGACY_SPEAKER_TAG_PATTERN.sub(r"<|speaker:\1|>", text))
+
+    disallowed = [
+        token for token in _CONTROL_TOKEN_PATTERN.findall(normalized) if not _SPEAKER_TAG_PATTERN.fullmatch(token)
+    ]
+    if disallowed:
+        tokens = ", ".join(sorted(set(disallowed)))
+        raise ValueError(f"Audio8 TTS input contains unsupported control token(s): {tokens}")
+
+    if add_default_speaker and not _SPEAKER_TAG_PATTERN.search(normalized):
+        normalized = f"<|speaker:0|>{normalized}"
+    return normalized
+
+
+def _encode(tokenizer: Any, text: str) -> list[int]:
+    return list(tokenizer.encode(text, add_special_tokens=False))
+
+
+def _encode_segments(tokenizer: Any, segments: list[str]) -> list[int]:
+    ids: list[int] = []
+    for segment in segments:
+        ids.extend(_encode(tokenizer, segment))
+    return ids
+
+
+def build_text_only_prompt_ids(tokenizer: Any, text: str) -> tuple[list[int], str]:
+    """Build the no-reference prompt.
+
+    Returns:
+        ``(prompt_token_ids, normalized_text)``.
+    """
+    normalized = normalize_text(text)
+    if not normalized:
+        raise ValueError("Audio8 TTS input text must not be empty")
+    ids = _encode_segments(
+        tokenizer,
+        [
+            "<|im_start|>system\n",
+            SYSTEM_PROMPT_TEXT_ONLY,
+            "<|im_end|>\n",
+            "<|im_start|>user\n",
+            normalized,
+            "<|im_end|>\n",
+            "<|im_start|>assistant\n<|voice|>",
+        ],
+    )
+    return ids, normalized
+
+
+def build_voice_clone_prompt_parts(
+    tokenizer: Any,
+    text: str,
+    ref_text: str,
+) -> tuple[list[int], list[int], str, str]:
+    """Build the clone prompt around the reference-code slot.
+
+    Caller splices ``[code + semantic_begin_id for code in ref_codes[:, 0]]``
+    between the returned prefix and suffix.
+    """
+    normalized = normalize_text(text)
+    if not normalized:
+        raise ValueError("Audio8 TTS input text must not be empty")
+    normalized_ref = normalize_text(ref_text, add_default_speaker=True)
+    if not normalized_ref:
+        raise ValueError("Audio8 TTS voice cloning requires a non-empty reference transcript")
+
+    prefix = _encode_segments(
+        tokenizer,
+        [
+            "<|im_start|>system\n",
+            SYSTEM_PROMPT_CLONE_PREFIX,
+            normalized_ref,
+            SYSTEM_PROMPT_CLONE_SUFFIX,
+        ],
+    )
+    suffix = _encode_segments(
+        tokenizer,
+        [
+            "<|im_end|>\n",
+            "<|im_start|>user\n",
+            normalized,
+            "<|im_end|>\n",
+            "<|im_start|>assistant\n<|voice|>",
+        ],
+    )
+    return prefix, suffix, normalized, normalized_ref
+
+
+def build_voice_clone_prompt_ids(
+    tokenizer: Any,
+    text: str,
+    ref_text: str,
+    semantic_token_ids: list[int],
+) -> tuple[list[int], int, str, str]:
+    """Build the full clone prompt.
+
+    ``ref_start_index`` is the position of the first reference frame, i.e.
+    where the residual codebook embeddings must be added.
+    """
+    prefix, suffix, normalized, normalized_ref = build_voice_clone_prompt_parts(tokenizer, text, ref_text)
+    return prefix + list(semantic_token_ids) + suffix, len(prefix), normalized, normalized_ref
+
+
+def estimate_voice_clone_prompt_len(
+    tokenizer: Any,
+    normalized_text: str,
+    normalized_ref_text: str,
+    ref_frames: int,
+) -> int:
+    """Exact clone prompt length, without encoding the reference audio."""
+    prefix, suffix, _, _ = build_voice_clone_prompt_parts(tokenizer, normalized_text, normalized_ref_text)
+    return len(prefix) + max(0, int(ref_frames)) + len(suffix)
+
+
+__all__ = [
+    "SYSTEM_PROMPT_CLONE_PREFIX",
+    "SYSTEM_PROMPT_CLONE_SUFFIX",
+    "SYSTEM_PROMPT_TEXT_ONLY",
+    "build_text_only_prompt_ids",
+    "build_voice_clone_prompt_ids",
+    "build_voice_clone_prompt_parts",
+    "clean_text",
+    "estimate_voice_clone_prompt_len",
+    "normalize_text",
+]

vllm_omni/model_executor/models/audio8_tts/sampling.py

@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio8 TTS sampling primitives, ported 1:1 from the reference checkpoint.
+
+Filter (top-k/top-p on **raw** logits) -> divide surviving logits by
+temperature -> Gumbel-max draw. This is the reverse of the usual filter order
+and is audibly different at ``temperature != 1``, which is why this cannot
+reuse ``models/common/nucleus_ras_sampling.py``. Per-row scalars are accepted
+so mixed batches run in one shot without a GPU->CPU sync.
+"""
+
+from __future__ import annotations
+
+import torch
+
+#: Below this temperature a request is treated as greedy, matching vLLM.
+SAMPLING_EPS = 1e-5
+
+
+def _as_column(
+    value: float | int | torch.Tensor,
+    batch_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Broadcast a scalar or ``[B]`` tensor to a ``[B, 1]`` tensor."""
+    if isinstance(value, torch.Tensor):
+        return value.to(device=device, dtype=dtype).reshape(-1, 1).expand(batch_size, 1)
+    return torch.full((batch_size, 1), float(value), device=device, dtype=dtype)
+
+
+def filter_top_k_top_p(
+    logits: torch.Tensor,
+    top_k: int | torch.Tensor,
+    top_p: float | torch.Tensor,
+) -> torch.Tensor:
+    """Mask candidates outside the top-k / top-p set of the raw logits.
+
+    ``top_k <= 0`` and ``top_p >= 1`` disable each filter. The highest-scoring
+    candidate is always kept, so no row can become all ``-inf``.
+    """
+    if logits.ndim != 2:
+        raise ValueError(f"logits must be 2-D [B, V], got {tuple(logits.shape)}")
+    batch_size, vocab = logits.shape
+    scalar_k = not isinstance(top_k, torch.Tensor)
+    scalar_p = not isinstance(top_p, torch.Tensor)
+    if scalar_k and scalar_p and (int(top_k) <= 0 or int(top_k) >= vocab) and float(top_p) >= 1.0:
+        return logits
+
+    sorted_scores, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+    cumulative = torch.cumsum(torch.softmax(sorted_scores, dim=-1), dim=-1)
+    positions = torch.arange(vocab, device=logits.device).unsqueeze(0)
+
+    keep_k = _as_column(top_k, batch_size, logits.device, torch.long)
+    # top_k <= 0 means "no cut"; represent it as the full vocabulary.
+    keep_k = torch.where(keep_k <= 0, torch.full_like(keep_k, vocab), keep_k)
+    threshold_p = _as_column(top_p, batch_size, logits.device, cumulative.dtype)
+
+    remove_sorted = (cumulative > threshold_p) | (positions >= keep_k)
+    # Always keep the argmax so at least one candidate survives.
+    remove_sorted[:, 0] = False
+    remove = torch.zeros_like(remove_sorted).scatter(-1, sorted_indices, remove_sorted)
+    return logits.masked_fill(remove, float("-inf"))
+
+
+def gumbel_argmax_sample(
+    scores: torch.Tensor,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    """Draw one index per row via Gumbel-max over ``softmax(scores)``."""
+    probabilities = torch.softmax(scores, dim=-1)
+    uniform = torch.rand(
+        probabilities.shape,
+        dtype=probabilities.dtype,
+        device=probabilities.device,
+        generator=generator,
+    )
+    # -log(u) is Exp(1); argmax(p / Exp(1)) samples from p.
+    noise = -torch.log(uniform.clamp_min(torch.finfo(probabilities.dtype).tiny))
+    return torch.argmax(probabilities / noise, dim=-1)
+
+
+def sample_scores(
+    logits: torch.Tensor,
+    *,
+    temperature: float | torch.Tensor,
+    top_k: int | torch.Tensor,
+    top_p: float | torch.Tensor,
+    do_sample: bool = True,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    """Filter -> temper -> draw, in the reference order.
+
+    Rows whose temperature is below :data:`SAMPLING_EPS` fall back to argmax.
+    """
+    filtered = filter_top_k_top_p(logits, top_k, top_p)
+    greedy = filtered.argmax(dim=-1)
+    if not do_sample:
+        return greedy
+    if not isinstance(temperature, torch.Tensor):
+        if float(temperature) < SAMPLING_EPS:
+            return greedy
+        return gumbel_argmax_sample(filtered / float(temperature), generator=generator)
+
+    temperature_col = _as_column(temperature, filtered.shape[0], filtered.device, filtered.dtype)
+    safe_temperature = torch.where(
+        temperature_col < SAMPLING_EPS,
+        torch.ones_like(temperature_col),
+        temperature_col,
+    )
+    sampled = gumbel_argmax_sample(filtered / safe_temperature, generator=generator)
+    return torch.where(temperature_col.reshape(-1) < SAMPLING_EPS, greedy, sampled)
+
+
+def ras_sample_semantic(
+    logits: torch.Tensor,
+    recent_ids: torch.Tensor | None,
+    *,
+    temperature: float | torch.Tensor,
+    top_k: int | torch.Tensor,
+    top_p: float | torch.Tensor,
+    ras_temperature: float,
+    ras_top_p: float,
+    do_sample: bool = True,
+    num_semantic_ids: int | None = None,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    """Repetition-Aware Sampling for the semantic (Slow AR) token.
+
+    Draws with the request's ``temperature`` / ``top_p``; if the drawn id is
+    already in ``recent_ids`` (pad with a value that cannot be sampled, e.g.
+    ``-1``), redraws from the flatter RAS distribution instead. The
+    ``num_semantic_ids`` slot and beyond (EOS) are never substituted, so a
+    repeat EOS still ends the utterance.
+
+    DualAR RAS *resamples* repeats rather than masking them — masking the
+    dominant token inflates EOS probability and truncates the utterance.
+    """
+    normal = sample_scores(
+        logits,
+        temperature=temperature,
+        top_k=top_k,
+        top_p=top_p,
+        do_sample=do_sample,
+        generator=generator,
+    )
+    if not do_sample or recent_ids is None or recent_ids.numel() == 0:
+        return normal
+
+    high = sample_scores(
+        logits,
+        temperature=ras_temperature,
+        top_k=top_k,
+        top_p=ras_top_p,
+        do_sample=True,
+        generator=generator,
+    )
+    repeated = (recent_ids == normal[:, None]).any(dim=1)
+    if num_semantic_ids is not None:
+        repeated &= normal < int(num_semantic_ids)
+    return torch.where(repeated, high, normal)
+
+
+def _row_value(value: object, index: int) -> object:
+    """Slice a per-row value for request ``index``; scalars pass through."""
+    return value[index : index + 1] if isinstance(value, torch.Tensor) else value
+
+
+def ras_sample_batch(
+    logits: torch.Tensor,
+    recent_ids: torch.Tensor | None,
+    *,
+    temperature: float | torch.Tensor,
+    top_k: int | torch.Tensor,
+    top_p: float | torch.Tensor,
+    ras_temperature: float,
+    ras_top_p: float,
+    num_semantic_ids: int | None,
+    generators: dict,
+    num_reqs: int,
+) -> torch.Tensor:
+    """Batched RAS that honors each request's own RNG generator.
+
+    vLLM keys ``generators`` by batch slot, so a single shared generator would
+    leak slot 0's RNG stream to every row and make a prompt's audio depend on
+    its batch neighbours (undercutting per-request ``seed``). The single-request
+    case keeps the batched fast path; larger batches loop so each row draws from
+    its own (possibly unseeded) generator.
+    """
+    common = dict(
+        ras_temperature=ras_temperature,
+        ras_top_p=ras_top_p,
+        num_semantic_ids=num_semantic_ids,
+    )
+    if num_reqs == 1:
+        return ras_sample_semantic(
+            logits,
+            recent_ids,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            generator=generators.get(0),
+            **common,
+        )
+    rows = [
+        ras_sample_semantic(
+            logits[i : i + 1],
+            None if recent_ids is None else recent_ids[i : i + 1],
+            temperature=_row_value(temperature, i),
+            top_k=_row_value(top_k, i),
+            top_p=_row_value(top_p, i),
+            generator=generators.get(i),
+            **common,
+        )
+        for i in range(num_reqs)
+    ]
+    return torch.cat(rows, dim=0)
+
+
+__all__ = [
+    "SAMPLING_EPS",
+    "filter_top_k_top_p",
+    "gumbel_argmax_sample",
+    "ras_sample_batch",
+    "ras_sample_semantic",
+    "sample_scores",
+]

vllm_omni/model_executor/models/registry.py

@@ -304,6 +304,17 @@ _OMNI_MODELS = {
         "gepard_talker",
         "GepardTalkerForConditionalGeneration",
     ),
+    ## audio8_tts (Audio8 TTS Preview 0.6B)
+    "Audio8TTSSlowARForConditionalGeneration": (
+        "audio8_tts",
+        "audio8_tts_slow_ar",
+        "Audio8TTSSlowARForConditionalGeneration",
+    ),
+    "Audio8TTSCodecDecoder": (
+        "audio8_tts",
+        "audio8_tts_codec_decoder",
+        "Audio8TTSCodecDecoder",
+    ),
     ## VoxCPM2
     "VoxCPM2TalkerForConditionalGeneration": (
         "voxcpm2",

vllm_omni/model_executor/stage_input_processors/audio8_tts.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio8 TTS Preview: Slow AR -> codec decoder stage input processor."""
+
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+
+from vllm_omni.data_entry_keys import (
+    CodesStruct,
+    MetaStruct,
+    OmniPayloadStruct,
+)
+
+#: Chunking defaults; ~21.5 codec frames per second, so 25 frames ~= 1.16 s.
+DEFAULT_CHUNK_FRAMES = 25
+DEFAULT_LEFT_CONTEXT_FRAMES = 25
+
+
+def _connector_extra(transfer_manager: Any) -> dict[str, Any]:
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    return raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+
+
+def _cfg_int(cfg: dict[str, Any], key: str, default: int) -> int:
+    value = cfg.get(key, default)
+    try:
+        return int(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Invalid Audio8 TTS integer config {key}={value!r}") from exc
+
+
+def _request_initial_chunk_frames(request: Any, default: int) -> int:
+    """Per-request override of ``initial_codec_chunk_frames``, if present."""
+    additional_information = getattr(request, "additional_information", None)
+    entries = getattr(additional_information, "entries", None)
+    if not entries or "initial_codec_chunk_frames" not in entries:
+        return default
+    entry = entries["initial_codec_chunk_frames"]
+    if entry.list_data is not None and len(entry.list_data) == 1:
+        return int(entry.list_data[0])
+    return default
+
+
+def extract_last_frame(multimodal_output: Mapping[str, Any]) -> torch.Tensor | None:
+    """Return the newest frame of codec codes, or ``None`` if this step had none.
+
+    Prefill steps emit an all-zero placeholder frame, which must not be
+    forwarded to the decoder.
+    """
+    audio_codes = multimodal_output.get("audio_codes")
+    if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
+        return None
+    if audio_codes.ndim == 1:
+        return audio_codes.to(device="cpu", dtype=torch.long).reshape(-1)
+    if audio_codes.ndim != 2:
+        raise ValueError(f"Invalid audio_codes shape for Audio8 TTS async_chunk: {tuple(audio_codes.shape)}")
+
+    frame = audio_codes[-1]
+    if frame.numel() == 0:
+        return None
+    valid = multimodal_output.get("audio_code_valid")
+    if isinstance(valid, torch.Tensor) and valid.numel() > 0:
+        is_valid = bool(valid.reshape(-1)[-1].item())
+    elif valid is not None:
+        is_valid = bool(valid)
+    else:
+        is_valid = bool(frame.any().item())
+    if not is_valid:
+        return None
+    return frame.to(device="cpu", dtype=torch.long).reshape(-1)
+
+
+def slow_ar_to_codec_decoder_async_chunk(
+    transfer_manager: Any,
+    multimodal_output: dict[str, Any] | None,
+    request: Any,
+    is_finished: bool = False,
+) -> OmniPayloadStruct | None:
+    """Emit ``[left_context | new_frames]`` chunks as the Slow AR produces codes.
+
+    Codes accumulate across AR steps and are shipped as a
+    ``[num_codebooks, frames]`` tensor. ``meta.left_context_size`` tells the
+    decoder how many leading frames to drop, keeping chunk boundaries
+    artefact-free. Returns ``None`` when the current chunk is not full yet.
+    """
+    request_id = request.external_req_id
+    finished = bool(is_finished or request.is_finished())
+    cfg = _connector_extra(transfer_manager)
+
+    appended_this_call = False
+    if isinstance(multimodal_output, Mapping):
+        frame = extract_last_frame(multimodal_output)
+        if frame is not None:
+            transfer_manager.code_prompt_token_ids[request_id].append(frame.detach())
+            appended_this_call = True
+    elif not finished:
+        return None
+
+    chunk_frames = _cfg_int(cfg, "codec_chunk_frames", DEFAULT_CHUNK_FRAMES)
+    left_context_frames = _cfg_int(cfg, "codec_left_context_frames", DEFAULT_LEFT_CONTEXT_FRAMES)
+    initial_chunk_frames = _request_initial_chunk_frames(request, _cfg_int(cfg, "initial_codec_chunk_frames", 0))
+    if chunk_frames <= 0 or left_context_frames < 0 or initial_chunk_frames < 0:
+        raise ValueError(
+            "Invalid Audio8 TTS codec chunk config: "
+            f"codec_chunk_frames={chunk_frames}, codec_left_context_frames={left_context_frames}, "
+            f"initial_codec_chunk_frames={initial_chunk_frames}"
+        )
+    initial_chunk_frames = min(initial_chunk_frames, chunk_frames)
+
+    frames = transfer_manager.code_prompt_token_ids[request_id]
+    length = len(frames)
+    if length <= 0:
+        if finished:
+            # Still emit a terminal payload so the decoder can close the stream.
+            return OmniPayloadStruct(
+                codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+                meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+            )
+        return None
+
+    # A smaller first chunk trades a little quality for time-to-first-audio.
+    if initial_chunk_frames > 0 and length <= chunk_frames:
+        already_sent = transfer_manager.put_req_chunk[request_id] * initial_chunk_frames
+        pending = length - already_sent
+        if pending <= 0:
+            return None
+        if pending < initial_chunk_frames and not finished:
+            return None
+        new_frames = min(pending, initial_chunk_frames)
+        left_context_size = max(0, length - new_frames)
+        window = frames[:length]
+    else:
+        initial_coverage = (
+            (chunk_frames // initial_chunk_frames) * initial_chunk_frames if initial_chunk_frames > 0 else 0
+        )
+        pending = (length - initial_coverage) % chunk_frames
+        if pending != 0 and not finished:
+            return None
+        if pending == 0 and finished and not appended_this_call:
+            # This boundary's chunk was already shipped when the stream first
+            # reached it; a terminal call carrying no new frame must close the
+            # stream, not re-send that window (otherwise the last chunk_frames
+            # frames ship twice).
+            return OmniPayloadStruct(
+                codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+                meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+            )
+        new_frames = pending if pending != 0 else chunk_frames
+        end_index = min(length, left_context_frames + new_frames)
+        left_context_size = max(0, end_index - new_frames)
+        window = frames[-end_index:]
+
+    # Codebook-major [num_codebooks, frames]; the tensor payload avoids
+    # expanding codec indices into Python ints across the connector boundary.
+    codes_qf = torch.stack(window, dim=0).transpose(0, 1).contiguous()
+    return OmniPayloadStruct(
+        codes=CodesStruct(audio=codes_qf),
+        meta=MetaStruct(
+            left_context_size=left_context_size,
+            finished=torch.tensor(finished, dtype=torch.bool),
+        ),
+    )
+
+
+__all__ = [
+    "DEFAULT_CHUNK_FRAMES",
+    "DEFAULT_LEFT_CONTEXT_FRAMES",
+    "extract_last_frame",
+    "slow_ar_to_codec_decoder_async_chunk",
+]

vllm_omni/transformers_utils/configs/__init__.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """
 Custom model configs that cannot be expressed via HuggingFace Transformers alone,
 following the same pattern as vllm.transformers_utils.configs.
@@ -10,6 +10,9 @@ from __future__ import annotations
 import importlib
 
 _CLASS_TO_MODULE: dict[str, str] = {
+    "Audio8TTSConfig": "vllm_omni.transformers_utils.configs.audio8_tts",
+    "Audio8TTSSlowARConfig": "vllm_omni.transformers_utils.configs.audio8_tts",
+    "Audio8TTSFastARConfig": "vllm_omni.transformers_utils.configs.audio8_tts",
     "HiggsAudioV3Config": "vllm_omni.transformers_utils.configs.higgs_audio_v3",
     "Mammothmoda2Config": "vllm_omni.transformers_utils.configs.mammoth_moda2",
     "Mammothmoda2Qwen2_5_VLConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2",
@@ -35,6 +38,9 @@ _CLASS_TO_MODULE: dict[str, str] = {
 }
 
 __all__ = [
+    "Audio8TTSConfig",
+    "Audio8TTSFastARConfig",
+    "Audio8TTSSlowARConfig",
     "HiggsAudioV3Config",
     "Mammothmoda2Config",
     "Mammothmoda2Qwen2_5_VLConfig",
@@ -75,6 +81,7 @@ def __dir__():
 
 # Eagerly import all config modules so their AutoConfig.register() side-effects
 # run as soon as `vllm_omni.transformers_utils.configs` is imported.
+from vllm_omni.transformers_utils.configs import audio8_tts as _audio8_tts  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech  # noqa: F401, E402

vllm_omni/transformers_utils/configs/audio8_tts.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio8 TTS config registration with transformers AutoConfig.
+
+Registers ``model_type = "arktts"`` (and the two sub-config types) so
+``AutoConfig.from_pretrained("Audio8/Audio8-TTS-Preview-0.6b")`` resolves to
+vllm-omni's Qwen2-shaped config instead of the checkpoint's remote code.
+
+This only wins when ``trust_remote_code`` is **off** -- transformers prefers a
+checkpoint's ``auto_map`` over registered classes otherwise -- which is why
+``deploy/audio8_tts.yaml`` sets ``trust_remote_code: false``.
+"""
+
+from transformers import AutoConfig
+
+from vllm_omni.model_executor.models.audio8_tts.configuration_audio8_tts import (
+    Audio8TTSConfig,
+    Audio8TTSFastARConfig,
+    Audio8TTSSlowARConfig,
+)
+
+AutoConfig.register("arktts", Audio8TTSConfig)
+AutoConfig.register("arktts_slow_ar", Audio8TTSSlowARConfig)
+AutoConfig.register("arktts_fast_ar", Audio8TTSFastARConfig)
+
+__all__ = [
+    "Audio8TTSConfig",
+    "Audio8TTSFastARConfig",
+    "Audio8TTSSlowARConfig",
+]