[New Model] Add support for tencent/Covo-Audio-Chat (#2293)

Signed-off-by: Dnoob <dxpouo@gmail.com>
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/models/supported_models.md

@@ -68,6 +68,7 @@ th {
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
 | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
 | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
+| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
 |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
 | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.

examples/offline_inference/covo_audio/README.md

@@ -0,0 +1,76 @@
+# Covo-Audio-Chat (Offline Inference)
+
+## Setup
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+
+> **Note**
+> Covo-Audio code2wav requires `torchdiffeq`. Install it with: `pip install torchdiffeq`
+
+## Run examples
+
+Get into the example folder:
+```bash
+cd examples/offline_inference/covo_audio
+```
+
+### Audio input chat
+
+Using the default audio asset:
+```bash
+python end2end.py
+```
+
+Using a custom audio file:
+```bash
+python end2end.py --audio-path /path/to/audio.wav
+```
+
+Using a local model:
+```bash
+python end2end.py -m /path/to/Covo-Audio-Chat --output-dir ./my_output
+```
+
+### Command-line Arguments
+
+| Argument | Short | Default | Description |
+|----------|-------|---------|-------------|
+| `--model-name` | `-m` | `tencent/Covo-Audio-Chat` | Model path or HuggingFace model ID |
+| `--text` | `-t` | `请回答这段音频里的问题。` | Text prompt / question for the audio |
+| `--audio-path` | `-a` | default audio asset | Path to local audio file |
+| `--sampling-rate` | | `16000` | Sampling rate for audio loading (Hz) |
+| `--output-dir` | | `./output_audio` | Output directory for generated files |
+| `--num-prompts` | | `1` | Number of prompts to generate |
+| `--stage-configs-path` | | (auto) | Path to stage configs YAML file |
+| `--log-stats` | | `false` | Enable detailed statistics logging |
+| `--stage-init-timeout` | | `300` | Stage initialization timeout (seconds) |
+| `--batch-timeout` | | `5` | Batching timeout (seconds) |
+| `--init-timeout` | | `300` | Overall initialization timeout (seconds) |
+| `--shm-threshold-bytes` | | `65536` | Shared memory threshold (bytes) |
+
+## Pipeline
+
+Covo-Audio-Chat uses a 2-stage pipeline:
+
+- **Stage 0 (fused_thinker_talker):** The 7B LLM generates interleaved text and audio tokens in a single autoregressive pass.
+- **Stage 1 (code2wav):** A BigVGAN-based vocoder converts the extracted audio codes into a 24kHz WAV waveform.
+
+## Output
+
+The script generates two files per request in the output directory:
+
+- `{request_id}.txt` -- prompt and generated text
+- `{request_id}.wav` -- generated audio (24kHz WAV)
+
+## FAQ
+
+If you encounter `ModuleNotFoundError: No module named 'librosa'`, install it with:
+```bash
+pip install librosa
+```
+
+## Environment
+
+- GPU: 1x A100 (80 GiB)
+- Stage 0 (7B LLM): ~16 GiB VRAM
+- Stage 1 (BigVGAN vocoder): ~2 GiB VRAM

examples/offline_inference/covo_audio/end2end.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+This example shows how to use vLLM-Omni for running offline inference
+with the correct prompt format on Covo-Audio-Chat.
+
+Usage:
+    python end2end.py --audio-path /path/to/audio.wav
+"""
+
+import os
+
+import soundfile as sf
+from vllm.assets.audio import AudioAsset
+from vllm.multimodal.media.audio import load_audio
+from vllm.sampling_params import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.model_executor.models.covo_audio.prompt_utils import (
+    COVO_AUDIO_INPUT_PREFIX,
+    build_covo_audio_chat_prompt,
+)
+
+SEED = 42
+
+
+def get_audio_query(
+    question: str | None = None,
+    audio_path: str | None = None,
+    sampling_rate: int = 16000,
+) -> dict:
+    if question is None:
+        question = "请回答这段音频里的问题。"
+    user_content = COVO_AUDIO_INPUT_PREFIX + question
+    prompt = build_covo_audio_chat_prompt(user_content)
+
+    if audio_path is None:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+    else:
+        import numpy as np
+
+        audio_signal, sr = load_audio(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+
+    return {
+        "prompt": prompt,
+        "multi_modal_data": {"audio": audio_data},
+        "modalities": ["audio"],
+    }
+
+
+def main(args):
+    query_result = get_audio_query(
+        question=args.text,
+        audio_path=args.audio_path,
+        sampling_rate=args.sampling_rate,
+    )
+
+    omni = Omni(
+        model=args.model_name,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+    )
+
+    # Stage 0: fused_thinker_talker
+    # stop_token_ids=[151645] (<|im_end|>) and ignore_eos=True are required
+    # so the model generates interleaved text+audio tokens before stopping.
+    thinker_sampling_params = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=2048,
+        seed=SEED,
+        detokenize=True,
+        repetition_penalty=1.05,
+        stop_token_ids=[151645],
+        ignore_eos=True,
+    )
+    # Stage 1: code2wav (audio codes, not real token IDs — skip detokenize)
+    code2wav_sampling_params = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=2048,
+        seed=SEED,
+        detokenize=False,
+        repetition_penalty=1.1,
+    )
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        code2wav_sampling_params,
+    ]
+
+    prompts = [query_result for _ in range(args.num_prompts)]
+
+    omni_outputs = omni.generate(prompts, sampling_params_list)
+
+    output_dir = args.output_dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    for stage_outputs in omni_outputs:
+        output = stage_outputs.request_output
+        if stage_outputs.final_output_type == "text":
+            request_id = output.request_id
+            text_output = output.outputs[0].text
+            prompt_text = output.prompt
+            out_txt = os.path.join(output_dir, f"{request_id}.txt")
+            lines = [
+                "Prompt:\n",
+                str(prompt_text) + "\n",
+                "vllm_text_output:\n",
+                str(text_output).strip() + "\n",
+            ]
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.writelines(lines)
+            except Exception as e:
+                print(f"[Warn] Failed writing text file {out_txt}: {e}")
+            print(f"Request ID: {request_id}, Text saved to {out_txt}")
+        elif stage_outputs.final_output_type == "audio":
+            request_id = output.request_id
+            audio_tensor = output.outputs[0].multimodal_output.get("audio")
+            if audio_tensor is None:
+                continue
+            output_wav = os.path.join(output_dir, f"{request_id}.wav")
+            audio_numpy = audio_tensor.float().detach().cpu().numpy()
+            if audio_numpy.ndim > 1:
+                audio_numpy = audio_numpy.flatten()
+            sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+            print(f"Request ID: {request_id}, Audio saved to {output_wav}")
+
+    omni.close()
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Offline inference demo for Covo-Audio-Chat")
+    parser.add_argument(
+        "--model-name",
+        "-m",
+        type=str,
+        default="tencent/Covo-Audio-Chat",
+        help="Model path or HuggingFace model ID.",
+    )
+    parser.add_argument(
+        "--text",
+        "-t",
+        type=str,
+        default=None,
+        help="Text prompt / question for the audio.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file. Uses default asset if not provided.",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Sampling rate for audio loading (default: 16000).",
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to stage configs YAML file.",
+    )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics.",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing a single stage in seconds.",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds.",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing stages in seconds.",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="./output_audio",
+        help="Output directory for generated files.",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of prompts to generate.",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/online_serving/covo_audio/README.md

@@ -0,0 +1,93 @@
+# Covo-Audio-Chat
+
+## Installation
+
+Please refer to [README.md](../../../README.md)
+
+> **Note**
+> Covo-Audio code2wav requires `torchdiffeq`. Install it with: `pip install torchdiffeq`
+
+## Run examples (Covo-Audio-Chat)
+
+### Launch the Server
+
+```bash
+CUDA_VISIBLE_DEVICES=0 vllm serve tencent/Covo-Audio-Chat --omni \
+    --trust-remote-code --port 18091
+```
+
+> The default deploy config is located at `vllm_omni/deploy/covo_audio.yaml` and is loaded automatically by the model registry — no `--deploy-config` flag needed for default use.
+
+### Send Request
+
+Get into the example folder:
+```bash
+cd examples/online_serving/covo_audio
+```
+
+#### Audio input chat
+
+Using the default audio asset:
+```bash
+python openai_chat_completion_client.py
+```
+
+Using a custom audio file:
+```bash
+python openai_chat_completion_client.py --audio-path /path/to/audio.wav
+```
+
+#### Streaming mode
+
+```bash
+python openai_chat_completion_client.py --audio-path /path/to/audio.wav --stream
+```
+
+### Command-line Arguments
+
+The Python client supports the following arguments:
+
+- `--prompt` (or `-p`): Text prompt (default: `"请回答这段音频里的问题。"`)
+- `--audio-path`: Path to audio file. Supports local file paths or HTTP/HTTPS URLs. Common audio formats: WAV, MP3, OGG, FLAC, M4A. Uses the default audio asset if omitted
+- `--output-audio-path` (or `-o`): Output path for generated audio (default: `./audio_0.wav`)
+- `--model`: Model name for the API. Auto-detected from server if omitted
+- `--stream`: Enable streaming mode
+
+## Pipeline
+
+Covo-Audio-Chat uses a 2-stage pipeline:
+
+- **Stage 0 (fused_thinker_talker):** The 7B LLM generates interleaved text and audio tokens in a single autoregressive pass.
+- **Stage 1 (code2wav):** A BigVGAN-based vocoder converts the extracted audio codes into a 24kHz WAV waveform.
+
+## Environment
+
+- GPU: 1x H100 (80 GiB)
+- Stage 0 (7B LLM): ~16 GiB VRAM
+- Stage 1 (BigVGAN vocoder): ~2 GiB VRAM
+
+## Important: System Prompt
+
+Covo-Audio-Chat requires a specific system prompt to enable interleaved text+audio generation. Without it, Stage 0 produces text-only tokens, and the output audio will be **silent**.
+
+The system prompt must include the instruction:
+
+```
+请用文本和音频进行对话，交替生成5个文本token和15个音频token，音频部分使用发音人：default_female
+```
+
+The full system prompt is defined in `vllm_omni/model_executor/models/covo_audio/prompt_utils.py` (`COVO_AUDIO_SYSTEM_PROMPT`).
+
+## FAQ
+
+If you encounter `ModuleNotFoundError: No module named 'librosa'`, install it with:
+```bash
+pip install librosa
+```
+
+### Audio output is silent (0.1s, all zeros)
+
+This means Stage 0 did not generate any audio tokens. Most likely causes:
+
+1. **Missing system prompt** -- see the "Important: System Prompt" section above. The model needs the interleaved generation instruction in the system message.
+2. **Missing `ignore_eos: true` / `stop_token_ids: [151645]`** in Stage 0 sampling params -- without these, the model stops at `<|endoftext|>` before generating audio tokens.

examples/online_serving/covo_audio/openai_chat_completion_client.py

@@ -0,0 +1,216 @@
+"""
+Client example for Covo-Audio-Chat with vllm-omni.
+
+Usage
+-----
+# Start the server first:
+#   CUDA_VISIBLE_DEVICES=0 vllm serve tencent/Covo-Audio-Chat --omni --trust-remote-code
+
+# Audio input chat (uses default audio asset if --audio-path not provided):
+python openai_chat_completion_client.py
+python openai_chat_completion_client.py --audio-path /path/to/audio.wav
+
+# Streaming mode:
+python openai_chat_completion_client.py --stream
+"""
+
+import base64
+import os
+import sys
+import time
+
+from openai import OpenAI
+from vllm.assets.audio import AudioAsset
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.model_executor.models.covo_audio.prompt_utils import (
+    COVO_AUDIO_SYSTEM_PROMPT,
+)
+
+openai_api_key = "EMPTY"
+openai_api_base = "http://localhost:18091/v1"
+
+client = OpenAI(api_key=openai_api_key, base_url=openai_api_base)
+
+SEED = 42
+
+
+def encode_base64_content_from_file(file_path: str) -> str:
+    """Encode a local file to base64 format."""
+    with open(file_path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+
+def get_audio_url(audio_path: str | None) -> str:
+    """Convert an audio path to a data URL for the API."""
+    if not audio_path:
+        return AudioAsset("mary_had_lamb").url
+
+    if audio_path.startswith(("data:", "http://", "https://")):
+        return audio_path
+
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+    ext = audio_path.rsplit(".", 1)[-1].lower() if "." in audio_path else "wav"
+    mime_map = {
+        "wav": "audio/wav",
+        "mp3": "audio/mpeg",
+        "mpeg": "audio/mpeg",
+        "ogg": "audio/ogg",
+        "flac": "audio/flac",
+        "m4a": "audio/mp4",
+    }
+    mime = mime_map.get(ext, "audio/wav")
+    b64 = encode_base64_content_from_file(audio_path)
+    return f"data:{mime};base64,{b64}"
+
+
+def build_audio_messages(audio_path: str | None, prompt: str) -> list[dict]:
+    """Build messages for an audio chat request."""
+    audio_url = get_audio_url(audio_path)
+    return [
+        {"role": "system", "content": COVO_AUDIO_SYSTEM_PROMPT},
+        {
+            "role": "user",
+            "content": [
+                {"type": "audio_url", "audio_url": {"url": audio_url}},
+                {"type": "text", "text": prompt},
+            ],
+        },
+    ]
+
+
+def get_model_name(override: str | None = None) -> str:
+    """Return the model name to use in API calls."""
+    if override:
+        return override
+    # Auto-detect from the server
+    models = client.models.list()
+    return models.data[0].id
+
+
+def run(args) -> None:
+    model_name = get_model_name(args.model)
+    thinker_params = {
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "top_k": -1,
+        "max_tokens": 2048,
+        "seed": SEED,
+        "detokenize": True,
+        "repetition_penalty": 1.05,
+        "stop_token_ids": [151645],
+        "ignore_eos": True,
+    }
+    code2wav_params = {
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "top_k": -1,
+        "max_tokens": 2048,
+        "seed": SEED,
+        "detokenize": False,
+        "repetition_penalty": 1.05,
+    }
+
+    extra_body = {
+        "sampling_params_list": [thinker_params, code2wav_params],
+    }
+
+    messages = build_audio_messages(args.audio_path, args.prompt)
+
+    start = time.perf_counter()
+
+    if not args.stream:
+        completion = client.chat.completions.create(
+            messages=messages,
+            model=model_name,
+            extra_body=extra_body,
+            stream=False,
+        )
+        elapsed = time.perf_counter() - start
+
+        for choice in completion.choices:
+            if choice.message.audio:
+                audio_bytes = base64.b64decode(choice.message.audio.data)
+                out_path = args.output_audio_path
+                os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
+                with open(out_path, "wb") as f:
+                    f.write(audio_bytes)
+                print(f"Audio saved to {out_path} ({len(audio_bytes)} bytes)")
+            elif choice.message.content:
+                print("Chat completion output from text:", choice.message.content)
+
+        print(f"Elapsed: {elapsed:.2f}s", file=sys.stderr)
+
+    else:
+        stream = client.chat.completions.create(
+            messages=messages,
+            model=model_name,
+            extra_body=extra_body,
+            stream=True,
+        )
+        audio_idx = 0
+        for chunk in stream:
+            modality = getattr(chunk, "modality", None)
+            elapsed = time.perf_counter() - start
+            for choice in chunk.choices:
+                content = getattr(choice.delta, "content", None) if hasattr(choice, "delta") else None
+                if not content:
+                    continue
+                if modality == "audio":
+                    audio_bytes = base64.b64decode(content)
+                    out_dir = os.path.dirname(args.output_audio_path) or "."
+                    out_path = os.path.join(out_dir, f"audio_{audio_idx}.wav")
+                    os.makedirs(out_dir, exist_ok=True)
+                    with open(out_path, "wb") as f:
+                        f.write(audio_bytes)
+                    print(f"[{elapsed:.2f}s] Audio chunk saved to {out_path}", file=sys.stderr)
+                    audio_idx += 1
+                elif modality == "text":
+                    print(content, end="", flush=True)
+
+        print()  # trailing newline
+        elapsed = time.perf_counter() - start
+        print(f"Elapsed: {elapsed:.2f}s", file=sys.stderr)
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Covo-Audio-Chat client example")
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        type=str,
+        default="请回答这段音频里的问题。",
+        help="Text prompt.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        type=str,
+        default=None,
+        help="Path to audio file. Uses default audio asset if omitted.",
+    )
+    parser.add_argument(
+        "--output-audio-path",
+        "-o",
+        type=str,
+        default="./audio_0.wav",
+        help="Output path for generated audio.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="Model name for the API. Auto-detected from server if omitted.",
+    )
+    parser.add_argument(
+        "--stream",
+        action="store_true",
+        help="Enable streaming mode.",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run(args)

recipes/README.md

@@ -29,6 +29,8 @@ recipes/
   Qwen-Image on `1x A100 80GB`, including optional step-wise continuous batching replay
 - [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md): online serving recipe for
   multimodal chat on `1x A100 80GB`
+- [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md): online
+  serving recipe for audio chat on `1x A100 80GB`
 - [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md): text-to-speech serving recipe
   for Qwen3-TTS (CustomVoice / VoiceDesign / Base) on `1x H100/A100 80GB`
 - [`LTX/LTX-2.md`](./LTX/LTX-2.md): text-to-video and image-to-video serving

recipes/Tencent/Covo-Audio-Chat.md

@@ -0,0 +1,92 @@
+# Covo-Audio-Chat for audio chat on 1x A100 80GB
+
+## Summary
+
+- Vendor: Tencent
+- Model: `tencent/Covo-Audio-Chat`
+- Task: Audio chat — audio (and optional text) input, interleaved text + audio output
+- Mode: Online serving with the OpenAI-compatible API (offline inference also supported)
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for serving
+`tencent/Covo-Audio-Chat` with vLLM-Omni on a single 80 GB-class GPU and
+validate the deployment with the bundled client example.
+
+## References
+
+- Related example under `examples/`:
+  [`examples/online_serving/covo_audio/README.md`](../../examples/online_serving/covo_audio/README.md)
+- Offline end-to-end script:
+  [`examples/offline_inference/covo_audio/end2end.py`](../../examples/offline_inference/covo_audio/end2end.py)
+- Deploy config:
+  [`vllm_omni/deploy/covo_audio.yaml`](../../vllm_omni/deploy/covo_audio.yaml)
+- Pipeline topology:
+  [`vllm_omni/model_executor/models/covo_audio/pipeline.py`](../../vllm_omni/model_executor/models/covo_audio/pipeline.py)
+
+## Hardware Support
+
+This recipe currently documents one tested reference configuration for
+CUDA GPU serving. The fused thinker+talker stage and the BigVGAN code2wav
+stage both run on the same physical device.
+
+## GPU
+
+### 1x A100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: NVIDIA CUDA 12.x with an A100 80 GB GPU
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+- Extra Python deps: `torchdiffeq`, `librosa`
+
+#### Command
+
+Install the two runtime-only deps that are not pulled in by vLLM-Omni itself:
+
+```bash
+pip install torchdiffeq librosa
+```
+
+Start the server from the repository root:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 vllm serve tencent/Covo-Audio-Chat --omni \
+    --trust-remote-code --port 18091
+```
+
+The deploy config at `vllm_omni/deploy/covo_audio.yaml` is loaded
+automatically by the model registry.
+
+#### Verification
+
+With the server ready, send a test request via the bundled client:
+
+```bash
+cd examples/online_serving/covo_audio
+python openai_chat_completion_client.py
+```
+
+Expected: the client prints the generated text and writes an output WAV
+file to the current directory.
+
+#### Notes
+
+- Memory usage: Stage 0 (7B LLM) ~16 GiB + stage 1 (BigVGAN vocoder)
+  ~2 GiB, both resident on the same device.
+- Key flags: `--omni` and `--trust-remote-code` are both required.
+- System prompt: the bundled `openai_chat_completion_client.py` and
+  `end2end.py` already include `COVO_AUDIO_SYSTEM_PROMPT` (defined in
+  `vllm_omni/model_executor/models/covo_audio/prompt_utils.py`) in the
+  system message by default. When integrating your own client, make sure
+  to include it — without it, stage 0 emits text-only tokens and the
+  output audio is silent.
+- Stage 1 runs with `enforce_eager: true` and `dtype: float32` — see the
+  header of `deploy/covo_audio.yaml` for the specific reasons.
+- Known limitations: this starter recipe covers the single-GPU online-serving
+  path only; multi-GPU layouts and platform-specific overrides are not
+  documented here.

tests/e2e/offline_inference/test_covo_audio_expansion.py

@@ -0,0 +1,47 @@
+"""
+E2E tests for Covo-Audio-Chat model with audio input and audio/text output.
+"""
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.models.covo_audio.prompt_utils import (
+    COVO_AUDIO_INPUT_PREFIX,
+    build_covo_audio_chat_prompt,
+)
+
+models = ["tencent/Covo-Audio-Chat"]
+stage_config = get_deploy_config_path("covo_audio.yaml")
+test_params = [(model, stage_config) for model in models]
+
+
+def get_question():
+    return "请回答这段音频里的问题。"
+
+
+@pytest.mark.full_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_audio_to_audio(omni_runner, omni_runner_handler) -> None:
+    """
+    Test audio input → text + audio output.
+    Deploy Setting: default yaml
+    Input Modal: text + audio
+    Output Modal: audio
+    """
+    audio = generate_synthetic_audio(2, 1, 16000)["np_array"]
+    if len(audio.shape) == 2:
+        audio = audio.squeeze()
+
+    user_content = COVO_AUDIO_INPUT_PREFIX + get_question()
+    prompt = build_covo_audio_chat_prompt(user_content)
+
+    request_config = {
+        "prompts": prompt,
+        "audios": (audio, 16000),
+        "modalities": ["audio"],
+    }
+    omni_runner_handler.send_omni_request(request_config)

tests/e2e/online_serving/test_covo_audio_expansion.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for Covo-Audio-Chat model.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio
+from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.models.covo_audio.prompt_utils import (
+    COVO_AUDIO_SYSTEM_PROMPT,
+)
+
+model = "tencent/Covo-Audio-Chat"
+stage_config_path = get_deploy_config_path("covo_audio.yaml")
+
+test_params = [
+    pytest.param(
+        OmniServerParams(
+            model=model,
+            stage_config_path=stage_config_path,
+            server_args=["--trust-remote-code"],
+        ),
+        id="default",
+    ),
+]
+
+
+@pytest.mark.full_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_audio_to_audio_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text + audio
+    Output Modal: audio
+    Input Setting: stream=False
+    Datasets: single request
+    """
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(2, 1, sample_rate=16000)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt={"role": "system", "content": COVO_AUDIO_SYSTEM_PROMPT},
+        audio_data_url=audio_data_url,
+        content_text="请回答这段音频里的问题。",
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "modalities": ["audio"],
+        "stream": False,
+    }
+
+    openai_client.send_omni_request(request_config)

vllm_omni/config/pipeline_registry.py

@@ -49,6 +49,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.qwen3_tts.pipeline",
         "QWEN3_TTS_PIPELINE",
     ),
+    "covo_audio": (
+        "vllm_omni.model_executor.models.covo_audio.pipeline",
+        "COVO_AUDIO_PIPELINE",
+    ),
     "bagel": (
         "vllm_omni.model_executor.models.bagel.pipeline",
         "BAGEL_PIPELINE",

vllm_omni/deploy/covo_audio.yaml

@@ -0,0 +1,46 @@
+# Covo-Audio-Chat deploy: fused_thinker_talker → code2wav on a single GPU.
+# Verified on 1x A100-80G.
+#
+# Fields omitted from a stage fall back to StageDeployConfig defaults (see
+# vllm_omni/config/stage_config.py). Non-default choices in this file:
+#   * async_chunk: false — pipeline defines no async-chunk processor funcs.
+#   * Stage 1 max_num_batched_tokens=2048, max_model_len=2048 — audio code
+#     sequences are short; the 32k default is unnecessary here.
+#   * Stage 1 dtype: float32 — BigVGAN vocoder weights are fp32; the default
+#     bf16 triggers mat1/mat2 dtype mismatch during warmup.
+#   * Stage 1 enforce_eager: true — the code2wav forward pass performs a
+#     GPU→CPU sync (int(code.flatten()[0]) in covo_audio.py) which is
+#     illegal during CUDA graph capture.
+async_chunk: false
+trust_remote_code: true
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.5
+    mm_processor_cache_gb: 0
+    devices: "0"
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      repetition_penalty: 1.1
+
+  - stage_id: 1
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.2
+    enforce_eager: true
+    dtype: float32
+    max_num_batched_tokens: 2048
+    max_model_len: 2048
+    mm_processor_cache_gb: 0
+    devices: "0"
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      repetition_penalty: 1.1

vllm_omni/entrypoints/openai/serving_speech.py

@@ -64,6 +64,7 @@ _QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"}
 _FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"}
 _COSYVOICE3_TTS_MODEL_STAGES = {"cosyvoice3_talker"}
 _OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"}
+_COVO_AUDIO_MODEL_STAGES = {"fused_thinker_talker"}
 _VOXCPM_TTS_MODEL_STAGES = {"latent_generator", "vae"}
 _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"}
 _MING_TTS_MODEL_STAGES = {"ming_tts"}
@@ -74,6 +75,7 @@ _TTS_MODEL_STAGES: set[str] = (
     | _FISH_TTS_MODEL_STAGES
     | _COSYVOICE3_TTS_MODEL_STAGES
     | _OMNIVOICE_TTS_MODEL_STAGES
+    | _COVO_AUDIO_MODEL_STAGES
     | _VOXCPM_TTS_MODEL_STAGES
     | _VOXCPM2_TTS_MODEL_STAGES
     | _MING_TTS_MODEL_STAGES
@@ -325,6 +327,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             and getattr(getattr(self._tts_stage, "engine_args", None), "model_stage", None) == "fish_speech_slow_ar"
         )
         self._fish_speech_tokenizer = None
+        self._covo_audio_tokenizer = None
 
         self._is_cosyvoice3 = (
             self._tts_stage is not None
@@ -490,6 +493,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return "cosyvoice3"
         if model_stage in _OMNIVOICE_TTS_MODEL_STAGES:
             return "omnivoice"
+        if model_stage in _COVO_AUDIO_MODEL_STAGES:
+            model_arch = getattr(self._tts_stage.engine_args, "model_arch", None)
+            if model_arch and "CovoAudio" in model_arch:
+                return "covo_audio"
         if model_stage in (_VOXCPM_TTS_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES):
             has_vae_stage = any(
                 getattr(getattr(stage, "engine_args", None), "model_stage", None) == "vae"
@@ -1874,6 +1881,41 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             "mm_processor_kwargs": mm_kwargs,
         }
 
+    # ---- Covo-Audio helpers ----
+
+    def _build_covo_audio_prompt(
+        self,
+        request: OpenAICreateSpeechRequest,
+    ) -> dict[str, Any]:
+        """Build a chat-style prompt for Covo-Audio-Chat.
+
+        Covo-Audio requires a specific system prompt that instructs the model
+        to interleave text and audio tokens in its output.  We render the
+        messages through the chat template and pass prompt_token_ids so that
+        the engine does not need to re-tokenize.
+        """
+        from transformers import AutoTokenizer
+
+        from vllm_omni.model_executor.models.covo_audio.prompt_utils import (
+            build_covo_audio_prompt_token_ids,
+        )
+
+        if self._covo_audio_tokenizer is None:
+            model_name = self.engine_client.model_config.model
+            try:
+                self._covo_audio_tokenizer = AutoTokenizer.from_pretrained(
+                    model_name,
+                    trust_remote_code=True,
+                )
+            except Exception as exc:
+                raise RuntimeError(f"Failed to load Covo-Audio tokenizer from '{model_name}': {exc}") from exc
+
+        prompt_ids = build_covo_audio_prompt_token_ids(
+            self._covo_audio_tokenizer,
+            request.input,
+        )
+        return {"prompt_token_ids": prompt_ids}
+
     def _apply_cosyvoice3_dynamic_tokens(
         self,
         sampling_params_list: list,
@@ -2014,6 +2056,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 prompt["lang"] = request.language
             if request.instructions:
                 prompt["instruct"] = request.instructions
+        elif self._tts_model_type == "covo_audio":
+            prompt = self._build_covo_audio_prompt(request)
+            tts_params = {}
         elif self._tts_model_type == "voxcpm2":
             # voxcpm2 doesn't use `_apply_uploaded_speaker` because the prompt builder needs the
             # raw waveform tuple for prefill-length accounting, not a base64 data URL.
@@ -2100,6 +2145,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request_id = request_id or f"speech-{random_uuid()}"
         if self._is_fish_speech:
             model_type = "fish_speech"
+        elif self._tts_model_type == "covo_audio":
+            model_type = "covo_audio"
         elif self._tts_model_type == "voxtral_tts":
             model_type = "voxtral_tts"
         elif self._tts_model_type == "cosyvoice3":

vllm_omni/model_executor/models/covo_audio/config_covo_audio.py

@@ -0,0 +1,72 @@
+# Copyright 2026 Tencent.
+from transformers.configuration_utils import PretrainedConfig
+
+# audio_token_index: token IDs >= this value are audio codes.
+# Matches config.audio_token_index and len(tokenizer) = 151671.
+# Note: the LLM total vocab_size is 168055 (includes audio tokens), which
+# is different from this boundary value.
+COVO_AUDIO_TOKEN_INDEX = 151671
+
+
+class CovoAudioCode2WavConfig(PretrainedConfig):
+    model_type = "covo_audio_code2wav"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.wav_input_sr = 24000
+        self.n_speakers = 1
+        self.n_styles = 1
+
+        self.token2latent = {
+            "upsample_factor": 4,
+            "token_vocab_size": 16384,
+            "token_input_dim": 512,
+            "z_dim": 64,
+            "spkr_embed_dim": 1024,
+            "spkr_mask_ratio": 0.0,
+            "bert_mask_rate0": 0.7,
+            "bert_mask_rate1": 1.0,
+            "random_maskrate": 0.3,
+            "cfg_dropout": 0.2,
+            "sigma": 1e-05,
+            "transformer": {
+                "num_layers": 12,
+                "hidden_size": 1024,
+                "ffn_hidden_size": 4096,
+                "num_heads": 16,
+                "modulation": True,
+                "alibi_bias": False,
+                "rotary_bias": True,
+                "qk_norm": False,
+                "max_position_embeddings": 4096,
+                "attn_dropout": 0.0,
+                "dropout": 0.1,
+            },
+        }
+
+        self.wavegan = {
+            "type": "BigVGANFlowVAE",
+            "upsample_rates": [5, 3, 2, 2, 2, 2],
+            "upsample_kernel_sizes": [10, 6, 4, 4, 4, 4],
+            "upsample_initial_channel": 1536,
+            "resblock": "1",
+            "resblock_kernel_sizes": [3, 7, 11],
+            "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+            "downsample_rates": [2, 2, 2, 2, 3, 5],
+            "downsample_channels": [12, 24, 48, 96, 192, 384, 768],
+            "activation": "snakebeta",
+            "snake_logscale": True,
+            "latent_dim": 64,
+            "use_flow": True,
+            "use_vae": True,
+            "kl_weight": 5,
+            "causal": True,
+            "flow_hidden_channels": 256,
+        }
+
+        self.inference = {
+            "s_steps": 10,
+            "cfg_alpha": 1.0,
+            "dynamic_the": False,
+            "rescale_logits": False,
+        }

vllm_omni/model_executor/models/covo_audio/covo_audio.py

@@ -0,0 +1,426 @@
+# Copyright 2026 Tencent.
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any
+
+import numpy as np
+import torch
+from torch import nn
+from transformers import BatchFeature, WhisperFeatureExtractor
+from vllm.config import VllmConfig
+from vllm.config.multimodal import BaseDummyOptions
+from vllm.inputs import MultiModalDataDict
+from vllm.model_executor.models import SupportsPP
+from vllm.model_executor.models.interfaces import SupportsMultiModal
+from vllm.model_executor.models.utils import init_vllm_registered_model
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalFieldConfig
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    ProcessorInputs,
+    PromptReplacement,
+    PromptUpdateDetails,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.v1.sample.metadata import SamplingMetadata
+
+from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
+
+# Max tokens per audio: Whisper 30s → 3000 mel frames → encoder 2x + adapter 8x = 16x → 188
+MAX_AUDIO_TOKENS = 188
+
+
+def _calc_audio_num_tokens(num_samples: int, sample_rate: int) -> int:
+    """Calculate the number of audio tokens for a given audio.
+
+    Matches the original Covo-Audio calc_seq_len: 4 rounds of stride-2
+    downsampling (Whisper encoder 2x + AudioAdapter 3-layer 8x = 16x total).
+    """
+    centiseconds = num_samples * 100 // sample_rate
+    seq_len = centiseconds
+    for _ in range(4):
+        seq_len = (seq_len + 1) // 2
+    return seq_len
+
+
+class CovoAudioProcessingInfo(BaseProcessingInfo):
+    def get_hf_config(self):
+        return self.ctx.get_hf_config()
+
+    def get_hf_processor(self, **kwargs):
+        return None
+
+    def get_feature_extractor(self, **kwargs):
+        return WhisperFeatureExtractor(
+            feature_size=128,
+            sampling_rate=16000,
+        )
+
+    def get_data_parser(self) -> MultiModalDataParser:
+        return MultiModalDataParser(
+            target_sr=16000,
+        )
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": None}
+
+    def get_mm_max_tokens_per_item(self, seq_len, mm_counts):
+        return {"audio": MAX_AUDIO_TOKENS}
+
+
+class CovoAudioDummyInputsBuilder(BaseDummyInputsBuilder[CovoAudioProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        num_audios = mm_counts.get("audio", 0)
+        return (" <|begofcAUDIO|><|cAUDIO|><|endofcAUDIO|>" * num_audios).strip()
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+    ) -> MultiModalDataDict:
+        num_audios = mm_counts.get("audio", 0)
+        if num_audios == 0:
+            return {}
+        # 30s at 16kHz = Whisper's max input length, ensures profile_run
+        # allocates memory for the worst-case MAX_AUDIO_TOKENS (188) tokens.
+        dummy_audio = np.zeros((16000 * 30,), dtype=np.float32)
+        return {"audio": [(dummy_audio, 16000)] * num_audios}
+
+    def get_dummy_processor_inputs(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+    ) -> ProcessorInputs:
+        dummy_text = self.get_dummy_text(mm_counts)
+        dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
+        dummy_mm_items = self.info.parse_mm_data(dummy_mm_data)
+        return ProcessorInputs(
+            prompt=dummy_text,
+            mm_data_items=dummy_mm_items,
+        )
+
+
+class CovoAudioMultiModalProcessor(BaseMultiModalProcessor[CovoAudioProcessingInfo]):
+    def _hf_processor_applies_updates(self, prompt_text, mm_items, hf_processor_mm_kwargs, tokenization_kwargs) -> bool:
+        return False
+
+    def _call_hf_processor(
+        self, prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, Any], tok_kwargs: Mapping[str, object]
+    ) -> BatchFeature:
+        tokenizer = self.info.get_tokenizer()
+        feature_extractor = self.info.get_feature_extractor()
+
+        if isinstance(prompt, str):
+            prompt_ids = tokenizer.encode(prompt)
+        else:
+            prompt_ids = prompt
+
+        audios = mm_data.get("audios", []) or mm_data.get("audio", [])
+        if not audios:
+            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
+
+        raw_audios = []
+        num_tokens_list = []
+        for item in audios:
+            if isinstance(item, tuple):
+                audio_array, _sr = item
+            else:
+                audio_array = item
+            raw_audios.append(audio_array)
+            # MultiModalDataParser resamples to target_sr=16000, so always
+            # compute token count at 16 kHz regardless of the original sr.
+            num_tokens_list.append(_calc_audio_num_tokens(len(audio_array), 16000))
+
+        features = feature_extractor(
+            raw_audios,
+            sampling_rate=16000,
+            return_tensors="pt",
+        )
+
+        return BatchFeature(
+            data={
+                "input_ids": torch.tensor([prompt_ids], dtype=torch.int64),
+                "audio_features": features["input_features"],
+                "audio_num_tokens": torch.tensor(num_tokens_list, dtype=torch.int64),
+            },
+            tensor_type=None,
+        )
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        return {
+            "audio_features": MultiModalFieldConfig.batched("audio"),
+            "audio_num_tokens": MultiModalFieldConfig.batched("audio"),
+        }
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs,
+    ) -> Sequence[PromptReplacement]:
+        tokenizer = self.info.get_tokenizer()
+        vocab = tokenizer.get_vocab()
+        audio_token_id = vocab["<|cAUDIO|>"]
+
+        out_mm_data = out_mm_kwargs.get_data()
+
+        def get_replacement(item_idx: int):
+            num_tokens_tensor = out_mm_data.get("audio_num_tokens")
+            if num_tokens_tensor is not None:
+                num_tokens = int(num_tokens_tensor[item_idx])
+            else:
+                num_tokens = MAX_AUDIO_TOKENS
+            return PromptUpdateDetails.select_token_id(
+                [audio_token_id] * num_tokens,
+                embed_token_id=audio_token_id,
+            )
+
+        return [
+            PromptReplacement("audio", "<|cAUDIO|>", get_replacement),
+        ]
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    CovoAudioMultiModalProcessor,
+    info=CovoAudioProcessingInfo,
+    dummy_inputs=CovoAudioDummyInputsBuilder,
+)
+class CovoAudioForConditionalGeneration(
+    nn.Module,
+    SupportsPP,
+    SupportsMultiModal,
+    CustomProcessMixin,
+):
+    requires_raw_input_tokens: bool = True
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        if modality.startswith("audio"):
+            return "<|begofcAUDIO|><|cAUDIO|><|endofcAUDIO|>"
+        return None
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.has_preprocess = False
+        self.have_multimodal_outputs = True
+        self.vllm_config = vllm_config
+
+        self.model_stage = vllm_config.model_config.model_stage
+        if self.model_stage == "fused_thinker_talker":
+            self.has_preprocess = True
+            self.set_custom_preprocess(self.fused_thinker_talker_preprocess)
+            self.fused_thinker_talker = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                hf_config=vllm_config.model_config.hf_config,
+                architectures=["CovoAudioLLMModel"],
+            )
+            self.code2wav = None
+            self.model = self.fused_thinker_talker
+
+        elif self.model_stage == "code2wav":
+            self.code2wav = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                architectures=["CovoAudioCode2WavModel"],
+            )
+            self.fused_thinker_talker = None
+            self.model = self.code2wav
+        else:
+            raise ValueError("Invalid model stage")
+
+        self.make_empty_intermediate_tensors = (
+            self.fused_thinker_talker.make_empty_intermediate_tensors
+            if self.model_stage == "fused_thinker_talker"
+            else lambda *args, **kwargs: None
+        )
+
+    def fused_thinker_talker_preprocess(
+        self,
+        input_ids: torch.Tensor | None,
+        input_embeds: torch.Tensor | None,
+        **info_dict: dict,
+    ):
+        if input_embeds is None and input_ids is not None:
+            input_embeds = self.fused_thinker_talker.embed_input_ids(input_ids)
+        return input_ids, input_embeds, info_dict
+
+    def embed_multimodal(self, **kwargs: object):
+        if self.model_stage == "fused_thinker_talker":
+            return self.fused_thinker_talker.embed_multimodal(**kwargs)
+        return None
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        is_multimodal: bool = False,
+    ) -> torch.Tensor:
+        if self.model_stage == "code2wav":
+            hidden_size = self.vllm_config.model_config.get_hidden_size()
+            return torch.zeros(
+                input_ids.numel(),
+                hidden_size,
+                dtype=self.vllm_config.model_config.dtype,
+                device=input_ids.device,
+            )
+        return self.model.embed_input_ids(input_ids, multimodal_embeddings, is_multimodal=is_multimodal)
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        sampling_metadata: SamplingMetadata | None = None,
+    ) -> torch.Tensor | None:
+        if isinstance(hidden_states, OmniOutput):
+            hidden_states = hidden_states.text_hidden_states
+        return self.model.compute_logits(hidden_states)
+
+    @staticmethod
+    def _module_device(module: nn.Module) -> torch.device:
+        try:
+            return next(module.parameters()).device
+        except StopIteration:
+            # No parameters; fall back to buffers or cpu
+            for _, buf in module.named_buffers(recurse=True):
+                return buf.device
+            return torch.device("cpu")
+
+    def generate_codes(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        **kwargs: object,
+    ):
+        llm_dev = self._module_device(self.fused_thinker_talker)
+
+        # Fallback: dummy input_ids matching embedding length
+        if input_ids is None:
+            input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=llm_dev).unsqueeze(
+                0
+            )  # (1, seq_len)
+
+        # Thinker (ensure inputs on thinker's device)
+        if input_ids is not None and input_ids.device != llm_dev:
+            input_ids = input_ids.to(llm_dev)
+        if positions is not None and positions.device != llm_dev:
+            positions = positions.to(llm_dev)
+        if inputs_embeds is not None and inputs_embeds.device != llm_dev:
+            inputs_embeds = inputs_embeds.to(llm_dev)
+
+        # Run llm
+        llm_output = self.fused_thinker_talker(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+        # CovoAudio interleaves audio tokens directly in the LLM output
+        # token stream; the stage input processor extracts them from
+        # token_ids, so there is no separate codec output here.
+        return None, llm_output
+
+    def generate_audio(self, code: torch.Tensor):
+        code2wav_dev = self._module_device(self.code2wav)
+
+        if isinstance(code, torch.Tensor):
+            code_tensor = code.to(dtype=torch.long, device=code2wav_dev)
+        else:
+            code_tensor = torch.as_tensor(code, dtype=torch.long, device=code2wav_dev)
+
+        if code_tensor.ndim == 2 and code_tensor.shape[0] == 1:
+            code_tensor = code_tensor.squeeze(0)
+
+        with torch.inference_mode():
+            audio_tensor = self.code2wav(input_ids=code_tensor)
+
+        return audio_tensor
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights for all components of the omni model."""
+        loaded_weights: set[str] = set()
+
+        # code2wav weights are loaded in its __init__ via torch.load;
+        # skip iterating safetensors entirely for this stage.
+        if self.code2wav:
+            c2w_loaded = self.code2wav.load_weights(iter([]))
+            c2w_loaded = add_prefix_to_loaded_weights(c2w_loaded, "code2wav")
+            loaded_weights.update(c2w_loaded)
+            return loaded_weights
+
+        llm_weights = ((k, v) for k, v in weights if k.startswith(("llm", "encoder", "audio_adapter")))
+        if self.fused_thinker_talker:
+            llm_loaded = self.fused_thinker_talker.load_weights(llm_weights)
+            llm_loaded = add_prefix_to_loaded_weights(llm_loaded, "fused_thinker_talker")
+            loaded_weights.update(llm_loaded)
+
+        return loaded_weights
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        sampling_metadata: SamplingMetadata | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors | OmniOutput:
+        """Route forward to the active stage.
+
+        - fused_thinker_talker: run Whisper + LLM → text hidden states
+          (audio tokens are interleaved in the output token stream).
+        - code2wav: run BigVGAN vocoder on audio codes → waveform.
+        """
+        if self.model_stage == "fused_thinker_talker":
+            _, text_hidden_states = self.generate_codes(
+                input_ids=input_ids,
+                positions=positions,
+                inputs_embeds=inputs_embeds,
+                intermediate_tensors=intermediate_tensors,
+                **kwargs,
+            )
+
+            return OmniOutput(
+                text_hidden_states=text_hidden_states.reshape(-1, text_hidden_states.shape[-1]),
+                multimodal_outputs={},
+            )
+
+        if self.model_stage == "code2wav":
+            code = (
+                input_ids
+                if input_ids is not None
+                else torch.zeros(
+                    inputs_embeds.shape[0],
+                    dtype=torch.long,
+                    device=inputs_embeds.device,
+                )
+            )
+
+            # Text-only response: stage processor sends a single sentinel [-1]
+            # token.  Skip BigVGAN inference and return a short silence so the
+            # serving layer produces a valid WAV instead of noise.
+            if code.numel() == 1 and int(code.flatten()[0]) == -1:
+                # 0.1s of silence at 24 kHz
+                return OmniOutput(
+                    text_hidden_states=None,
+                    multimodal_outputs={"model_outputs": torch.zeros(1, 2400, device=code.device)},
+                )
+
+            audio_tensor = self.generate_audio(code)
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={
+                    "model_outputs": audio_tensor.reshape(1, -1) if audio_tensor is not None else audio_tensor
+                },
+            )

vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py

@@ -0,0 +1,126 @@
+# Copyright 2026 Tencent.
+import os
+from collections.abc import Iterable
+
+import numpy as np
+import torch
+from torch import nn
+from vllm.config import VllmConfig
+from vllm.model_executor.models import SupportsPP
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
+
+from .config_covo_audio import CovoAudioCode2WavConfig
+from .token2wav import JsonHParams, Token2WavDecoder
+
+
+class CovoAudioCode2WavForConditionalGeneration(nn.Module, SupportsPP):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        self.vllm_config = vllm_config
+        self.config = vllm_config.model_config.hf_config
+        model_name = vllm_config.model_config.model
+        if os.path.isdir(model_name):
+            model_path = model_name
+        else:
+            from huggingface_hub import snapshot_download
+
+            model_path = snapshot_download(model_name)
+        token2wav_path = os.path.join(model_path, "token2wav")
+
+        code2wav_config = CovoAudioCode2WavConfig()
+        model_config = JsonHParams(
+            **{
+                "token2latent": code2wav_config.token2latent,
+                "wavegan": code2wav_config.wavegan,
+                "wav_input_sr": code2wav_config.wav_input_sr,
+                "global_mean_var": os.path.join(token2wav_path, "global_mean_var.npy"),
+            }
+        )
+
+        self.decoder = Token2WavDecoder(model_config)
+        ckpt = torch.load(
+            os.path.join(token2wav_path, "model.pt"),
+            map_location="cpu",
+            weights_only=True,
+        )
+        self.decoder.load_state_dict(ckpt)
+        self.decoder.eval()
+        self.infer_config = code2wav_config.inference
+        self._decoder_moved = False
+        self._sampler = Sampler()
+
+        # Load bundled speaker prompt files for voice timbre control.
+        prompt_dir = os.path.join(os.path.dirname(__file__), "speaker_prompt")
+        self._prompt_token = torch.from_numpy(np.load(os.path.join(prompt_dir, "prompt_token.npy"))).long()
+        self._prompt_latent = torch.from_numpy(np.load(os.path.join(prompt_dir, "prompt_latent.npy"))).float()
+        self._spkr_embed = torch.from_numpy(np.load(os.path.join(prompt_dir, "speaker_embed.npy"))).float()
+
+    def _ensure_on_device(self, device: torch.device) -> None:
+        """Move decoder and prompt tensors to *device* once."""
+        if self._decoder_moved:
+            return
+        self.decoder = self.decoder.to(device)
+        self._prompt_token = self._prompt_token.to(device)
+        self._prompt_latent = self._prompt_latent.to(device)
+        self._spkr_embed = self._spkr_embed.to(device)
+        self._decoder_moved = True
+
+    def forward(self, input_ids: torch.Tensor, **kwargs):
+        self._ensure_on_device(input_ids.device)
+        data = {
+            "target_token": input_ids.unsqueeze(0),
+            "sample_rate": 24000,
+            "prompt_token": self._prompt_token,
+            "prompt_latent": self._prompt_latent,
+            "spkr_embed": self._spkr_embed,
+        }
+        audio = self.decoder.inference(data, **self.infer_config)
+        return audio
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        is_multimodal: bool = False,
+    ) -> torch.Tensor:
+        hidden_size = self.vllm_config.model_config.get_hidden_size()
+        return torch.zeros(
+            input_ids.numel(),
+            hidden_size,
+            dtype=self.vllm_config.model_config.dtype,
+            device=input_ids.device,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        if hidden_states.numel() == 0:
+            return None
+        batch_size = hidden_states.shape[0] if hidden_states.ndim > 0 else 1
+        vocab_size = self.config.llm_config.vocab_size
+        return torch.zeros(
+            (batch_size, vocab_size),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> SamplerOutput | None:
+        if logits is None or logits.numel() == 0:
+            return None
+        return self._sampler(logits=logits, sampling_metadata=sampling_metadata)
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        # Weights are loaded in __init__ via torch.load
+        # Return all decoder state_dict keys so vLLM knows they are initialized
+        return {f"decoder.{k}" for k in self.decoder.state_dict()}

vllm_omni/model_executor/models/covo_audio/covo_audio_llm.py

@@ -0,0 +1,152 @@
+# Copyright 2026 Tencent.
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+from transformers.models.whisper.modeling_whisper import WhisperEncoder
+from vllm.config import VllmConfig
+from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
+from vllm.sequence import IntermediateTensors
+
+
+class DownsampleLayer(nn.Module):
+    """
+    Downsample layer with 1D convolution and linear layers.
+    """
+
+    def __init__(self, input_dim, output_dim, hidden_dim=2048):
+        super().__init__()
+        self.conv1d = nn.Conv1d(in_channels=input_dim, out_channels=input_dim, kernel_size=3, stride=2, padding=1)
+        self.linear1 = nn.Linear(input_dim, hidden_dim)
+        self.relu1 = nn.ReLU()
+        self.linear2 = nn.Linear(hidden_dim, output_dim)
+        self.relu2 = nn.ReLU()
+
+    def forward(self, x):
+        # x: (B, T, C)
+        x = x.transpose(1, 2)  # -> (B, C, T)
+        x = self.conv1d(x)  # -> (B, C, T // 2)
+        x = x.transpose(1, 2)  # -> (B, T // 2, C)
+        x = self.relu1(x)
+        x = self.linear1(x)  # -> (B, T // 2, hidden_dim)
+        x = self.relu2(x)
+        x = self.linear2(x)  # -> (B, T // 2, output_dim)
+        return x
+
+
+class AudioAdapter(nn.Module):
+    """
+    Audio adapter with downsample layers.
+    """
+
+    def __init__(self, input_dim, output_dim, downsample=8):
+        """
+        Args:
+            input_dim (int): input feature dimension (number of channels)
+            output_dim (int): output feature dimension
+            downsample (int): total downsampling factor, must be a power of 2
+        """
+        super().__init__()
+        assert downsample >= 2 and (downsample & (downsample - 1)) == 0, "downsample must be a power of 2"
+
+        num_layers = (
+            downsample.bit_length() - 1
+        )  # calculate how many downsampling steps are needed to reach the target factor
+
+        layers = []
+        in_dim = input_dim
+        for i in range(num_layers):
+            is_last = i == num_layers - 1
+            out_dim = output_dim if is_last else input_dim
+            layers.append(DownsampleLayer(in_dim, out_dim))
+            in_dim = out_dim
+
+        self.downsample_layers = nn.ModuleList(layers)
+
+    def forward(self, x):
+        """
+        Args:
+            x (Tensor): shape (B, T, C),C=input_dim
+        Returns:
+            Tensor: shape (B, T // downsample, output_dim)
+        """
+        for layer in self.downsample_layers:
+            x = layer(x)
+        return x
+
+
+class CovoAudioLLMForConditionalGeneration(nn.Module, SupportsPP, SupportsMultiModal):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        self.llm = init_vllm_registered_model(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "llm"),
+            architectures=["Qwen2ForCausalLM"],
+        )
+        self.encoder = WhisperEncoder(config.encoder_config)
+        self.audio_adapter = AudioAdapter(
+            config.whisper_feats_dim, config.llm_config.hidden_size, config.adapter_downsample
+        )
+        self.make_empty_intermediate_tensors = self.llm.make_empty_intermediate_tensors
+
+    def get_language_model(self) -> torch.nn.Module:
+        return self.llm
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        is_multimodal=None,
+    ) -> torch.Tensor:
+        return super().embed_input_ids(
+            input_ids,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+        )
+
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
+        audio_features = kwargs["audio_features"]
+        audio_num_tokens = kwargs.get("audio_num_tokens")
+        # Match encoder dtype (bfloat16) since audio_features from WhisperFeatureExtractor are float32
+        encoder_dtype = next(self.encoder.parameters()).dtype
+        audio_features = audio_features.to(dtype=encoder_dtype)
+        feats = self.encoder(audio_features).last_hidden_state
+        features = self.audio_adapter(feats)
+        # Truncate each audio to its actual length (Whisper pads all to 30s)
+        # Return a list of per-item tensors since lengths may differ.
+        batch_size = features.shape[0]
+        result = []
+        for i in range(batch_size):
+            if audio_num_tokens is not None:
+                n = int(audio_num_tokens[i])
+                result.append(features[i, :n, :])
+            else:
+                result.append(features[i])
+        return result
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+
+        hidden_state = self.llm(input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds)
+        return hidden_state
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        return self.llm.compute_logits(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/model_executor/models/covo_audio/pipeline.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Covo-Audio pipeline topology (frozen).
+
+Stage 0: fused_thinker_talker — multimodal understanding + interleaved text/audio
+         token generation (AR).
+Stage 1: code2wav              — audio codes → waveform (generation).
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.covo_audio"
+
+COVO_AUDIO_PIPELINE = PipelineConfig(
+    model_type="covo_audio",
+    model_arch="CovoAudioForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="fused_thinker_talker",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="latent",
+            sampling_constraints={
+                "detokenize": True,
+                "stop_token_ids": [151645],
+                "ignore_eos": True,
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="code2wav",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            custom_process_input_func=f"{_PROC}.llm2code2wav",
+            sampling_constraints={"detokenize": False},
+        ),
+    ),
+)

vllm_omni/model_executor/models/covo_audio/prompt_utils.py

@@ -0,0 +1,56 @@
+"""Shared Covo-Audio prompt constants and construction helpers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+COVO_AUDIO_INPUT_PREFIX = "<|begofcAUDIO|><|cAUDIO|><|endofcAUDIO|>"
+
+COVO_AUDIO_SYSTEM_PROMPT = (
+    '你是"小腾"，英文名是"Covo"，由腾讯开发的AI助手。\n'
+    "1、请使用简洁、口语化的语言和用户聊天，"
+    "你的态度积极、耐心，像一位值得信赖的朋友。\n"
+    "2、不要使用列表或编号，避免输出网址、表情符号和复杂的公式。\n"
+    "3、不评价竞争对手，不发表主观政治观点，"
+    "针对色情类、政治类、恐怖类、歧视类、暴力类的用户问题，"
+    "你要妥善应对潜在的安全风险，并给出幽默，情绪安抚以及安全的劝导。\n"
+    "请用文本和音频进行对话，交替生成5个文本token和15个音频token，"
+    "音频部分使用发音人：default_female"
+)
+
+
+def build_covo_audio_chat_prompt(
+    user_content: str,
+    system_prompt: str = COVO_AUDIO_SYSTEM_PROMPT,
+) -> str:
+    """Build a raw chat-template prompt string for Covo-Audio-Chat.
+
+    Used by offline inference and tests where the prompt is passed as a
+    plain string (tokenization happens downstream).
+    """
+    return (
+        f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+        f"<|im_start|>user\n{user_content}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+
+
+def build_covo_audio_prompt_token_ids(
+    tokenizer: Any,
+    user_text: str,
+    system_prompt: str = COVO_AUDIO_SYSTEM_PROMPT,
+) -> list[int]:
+    """Build tokenized prompt IDs for Covo-Audio-Chat.
+
+    Used by the OpenAI-compatible serving layer where prompt_token_ids
+    are passed directly to the engine.
+    """
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_text},
+    ]
+    return tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+    )

vllm_omni/model_executor/models/covo_audio/token2wav.py

@@ -0,0 +1,996 @@
+# Copyright 2026 Tencent.
+# token2wav: audio token codes -> waveform (inference only)
+# Pipeline: Token -> Latent (flow matching) -> Waveform (BigVGAN)
+
+import math
+from collections import OrderedDict, namedtuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch import pow, sin
+from torch.nn import Parameter
+from torch.nn.utils import weight_norm
+
+# --------------- Utilities ---------------
+
+
+class JsonHParams:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, dict):
+                v = JsonHParams(**v)
+            if isinstance(v, str) and v.lower() in ["non", "none", "nil", "null"]:
+                v = None
+            self[k] = v
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def items(self):
+        return self.__dict__.items()
+
+    def values(self):
+        return self.__dict__.values()
+
+    def __len__(self):
+        return len(self.__dict__)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+    def __contains__(self, key):
+        return key in self.__dict__
+
+    def __repr__(self):
+        return self.__dict__.__repr__()
+
+    def get(self, key, default=None):
+        return self.__dict__.get(key, default)
+
+
+def exists(val):
+    return val is not None
+
+
+def eval_decorator(fn):
+    def inner(model, *args, **kwargs):
+        was_training = model.training
+        model.eval()
+        out = fn(model, *args, **kwargs)
+        model.train(was_training)
+        return out
+
+    return inner
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+@torch.autocast(enabled=False, device_type="cuda")
+def apply_rotary_pos_emb(pos, t):
+    return t * pos.cos() + rotate_half(t) * pos.sin()
+
+
+if "sinc" in dir(torch):
+    sinc = torch.sinc
+else:
+
+    def sinc(x: torch.Tensor):
+        return torch.where(
+            x == 0, torch.tensor(1.0, device=x.device, dtype=x.dtype), torch.sin(math.pi * x) / math.pi / x
+        )
+
+
+def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+
+    delta_f = 4 * half_width
+    A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if A > 50.0:
+        beta = 0.1102 * (A - 8.7)
+    elif A >= 21.0:
+        beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+    if even:
+        time = torch.arange(-half_size, half_size) + 0.5
+    else:
+        time = torch.arange(kernel_size) - half_size
+    if cutoff == 0:
+        filter_ = torch.zeros_like(time)
+    else:
+        filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+        filter_ /= filter_.sum()
+
+    return filter_.view(1, 1, kernel_size)
+
+
+# --------------- Basic Layers ---------------
+
+
+class Dropout(nn.Module):
+    def __init__(self, p: float = 0.5, inplace: bool = False, force_drop: bool = False, **kwargs):
+        super().__init__()
+        if p < 0.0 or p > 1.0:
+            raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
+        self.p = p
+        self.inplace = inplace
+        self.force_drop = force_drop
+
+    def forward(self, x, **kwargs):
+        return F.dropout(x, p=self.p, training=True if self.force_drop else self.training, inplace=self.inplace)
+
+
+class EmbeddingTable(nn.Embedding):
+    def __init__(self, num_embeddings: int, embedding_dim: int, pad_id=-1, **kwargs):
+        super().__init__(num_embeddings, embedding_dim, **kwargs)
+        nn.init.normal_(self.weight, 0.0, embedding_dim**-0.5)
+        self.pad_id = pad_id
+        self.output_dim = embedding_dim
+
+    def forward(self, x):
+        if self.pad_id is not None:
+            mask = x == self.pad_id
+            x = x.masked_fill(mask, 0)
+        outputs = super().forward(x)
+        if self.pad_id is not None:
+            outputs = outputs.masked_fill(mask.unsqueeze(-1), 0.0)
+        return outputs
+
+
+class Linear(nn.Linear):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        bias: bool = True,
+        w_init_gain: str = "linear",
+        activation=None,
+        **kwargs,
+    ):
+        super().__init__(in_channels, out_channels, bias=bias)
+        self.activation = activation if activation is not None else nn.Identity()
+        self.output_dim = out_channels
+        if w_init_gain is not None:
+            if isinstance(w_init_gain, str):
+                gain = nn.init.calculate_gain(w_init_gain)
+            else:
+                gain = w_init_gain
+            nn.init.xavier_uniform_(self.weight, gain=gain)
+        if bias:
+            nn.init.constant_(self.bias, 0.0)
+
+    def forward(self, x, **kwargs):
+        return self.activation(super().forward(x))
+
+
+class Conv1d(nn.Conv1d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 1,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        padding_mode: str = "zeros",
+        bias: bool = True,
+        padding=None,
+        causal: bool = False,
+        bn: bool = False,
+        activation=None,
+        w_init_gain=None,
+        input_transpose: bool = False,
+        **kwargs,
+    ):
+        self.causal = causal
+        if padding is None:
+            if causal:
+                padding = 0
+                self.left_padding = dilation * (kernel_size - 1)
+            else:
+                padding = get_padding(kernel_size, dilation)
+
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            padding_mode=padding_mode,
+            bias=bias,
+        )
+
+        self.in_channels = in_channels
+        self.transpose = input_transpose
+        self.bn = nn.BatchNorm1d(out_channels) if bn else nn.Identity()
+        self.activation = activation if activation is not None else nn.Identity()
+        if w_init_gain is not None:
+            nn.init.xavier_uniform_(self.weight, gain=nn.init.calculate_gain(w_init_gain))
+
+    def forward(self, x):
+        if self.transpose or x.size(1) != self.in_channels:
+            assert x.size(2) == self.in_channels
+            x = x.transpose(1, 2)
+            self.transpose = True
+
+        if self.causal:
+            x = F.pad(x.unsqueeze(2), (self.left_padding, 0, 0, 0)).squeeze(2)
+
+        outputs = self.activation(self.bn(super().forward(x)))
+        return outputs.transpose(1, 2) if self.transpose else outputs
+
+
+class ConvTranspose1d(nn.ConvTranspose1d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        output_padding: int = 0,
+        groups: int = 1,
+        bias: bool = True,
+        dilation: int = 1,
+        padding=None,
+        padding_mode: str = "zeros",
+        causal: bool = False,
+        input_transpose: bool = False,
+        **kwargs,
+    ):
+        if padding is None:
+            padding = 0 if causal else (kernel_size - stride) // 2
+        if causal:
+            assert padding == 0, "padding is not allowed in causal ConvTranspose1d."
+            assert kernel_size == 2 * stride, "kernel_size must be equal to 2*stride in Causal ConvTranspose1d."
+
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            output_padding=output_padding,
+            groups=groups,
+            bias=bias,
+            dilation=dilation,
+            padding_mode=padding_mode,
+        )
+
+        self.causal = causal
+        self.stride = stride
+        self.transpose = input_transpose
+
+    def forward(self, x):
+        if self.transpose or x.size(1) != self.in_channels:
+            assert x.size(2) == self.in_channels
+            x = x.transpose(1, 2)
+            self.transpose = True
+
+        x = super().forward(x)
+        if self.causal:
+            x = x[:, :, : -self.stride]
+        return x.transpose(1, 2) if self.transpose else x
+
+
+class ConvPositionEmbed(nn.Module):
+    def __init__(self, hidden_size: int, kernel_size: int = 31, groups: int = 16):
+        super().__init__()
+        self.conv = Conv1d(
+            hidden_size, hidden_size, kernel_size, groups=groups, input_transpose=True, activation=nn.GELU()
+        )
+
+    def forward(self, x, mask=None):
+        if mask is not None:
+            mask = mask[..., None]
+            x = x.masked_fill(~mask, 0.0)
+        x = self.conv(x)
+        if mask is not None:
+            x = x.masked_fill(~mask, 0.0)
+        return x
+
+
+# --------------- Anti-aliasing (BigVGAN) ---------------
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(
+        self,
+        cutoff=0.5,
+        half_width=0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ):
+        super().__init__()
+        if cutoff < -0.0:
+            raise ValueError("Minimum cutoff must be larger than zero.")
+        if cutoff > 0.5:
+            raise ValueError("A cutoff above 0.5 does not make sense.")
+        self.kernel_size = kernel_size
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.padding = padding
+        self.padding_mode = padding_mode
+        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+        self.register_buffer("filter", filter)
+
+    def forward(self, x):
+        _, C, _ = x.shape
+        if self.padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+        return out
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        self.stride = ratio
+        self.pad = self.kernel_size // ratio - 1
+        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
+        self.register_buffer("filter", filter)
+
+    def forward(self, x):
+        _, C, _ = x.shape
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+        x = x[..., self.pad_left : -self.pad_right]
+        return x
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=self.kernel_size
+        )
+
+    def forward(self, x):
+        return self.lowpass(x)
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self, activation, up_ratio: int = 2, down_ratio: int = 2, up_kernel_size: int = 12, down_kernel_size: int = 12
+    ):
+        super().__init__()
+        self.up_ratio = up_ratio
+        self.down_ratio = down_ratio
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    def forward(self, x):
+        x = self.upsample(x)
+        x = self.act(x)
+        x = self.downsample(x)
+        return x
+
+
+# --------------- Activations ---------------
+
+
+class SnakeBeta(nn.Module):
+    def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
+        super().__init__()
+        self.in_features = in_features
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+            self.beta = Parameter(torch.zeros(in_features) * alpha)
+        else:
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+            self.beta = Parameter(torch.ones(in_features) * alpha)
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+        return x
+
+
+# --------------- Attention ---------------
+
+AttentionConfig = namedtuple("AttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"])
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, theta=50000):
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+    @property
+    def device(self):
+        return self.inv_freq.device
+
+    @torch.autocast(enabled=False, device_type="cuda")
+    def forward(self, t):
+        if not torch.is_tensor(t):
+            t = torch.arange(t, device=self.device)
+        t = t.type_as(self.inv_freq)
+        freqs = torch.einsum("i , j -> i j", t, self.inv_freq)
+        freqs = torch.cat((freqs, freqs), dim=-1)
+        return freqs
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.0,
+        dropout: float = 0.0,
+        max_position_embeddings: int = 4096,
+        norm_layer: nn.Module = nn.LayerNorm,
+        rotary_bias: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        assert hidden_size % num_heads == 0, "hidden_size should be divisible by num_heads"
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.scale = self.head_dim**-0.5
+        self.max_position_embeddings = max_position_embeddings
+        self.rotary_bias = rotary_bias
+
+        self.q_proj = Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.k_proj = Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.v_proj = Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+
+        self.attn_drop = Dropout(attn_drop)
+        self.o_proj = Linear(hidden_size, hidden_size)
+        self.o_dropout = Dropout(dropout)
+
+        self.cpu_config = AttentionConfig(True, True, True)
+        device_properties = torch.cuda.get_device_properties(torch.device("cuda"))
+        if device_properties.major == 8 and device_properties.minor == 0:
+            self.cuda_config = AttentionConfig(True, True, True)
+        else:
+            self.cuda_config = AttentionConfig(False, True, True)
+
+        if self.rotary_bias:
+            self.rotary = RotaryEmbedding(self.head_dim)
+
+    def forward(self, q, k=None, v=None, mask=None):
+        k = k or q
+        v = v or q
+        B, L, C = q.shape
+        B, S, C = v.shape
+        if mask is not None:
+            if mask.ndim == 2:
+                assert L == S
+                mask = rearrange(mask, "b j -> b 1 1 j")
+                mask = mask.expand(-1, self.num_heads, L, -1)
+            elif mask.ndim == 3:
+                assert mask.size(1) == L and mask.size(2) == S
+                mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
+
+        q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)
+        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
+        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)
+        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        config = self.cuda_config if q.is_cuda else self.cpu_config
+        attn_bias = torch.zeros(B, self.num_heads, L, S, dtype=q.dtype, device=q.device)
+
+        if self.rotary_bias:
+            if L == S:
+                rotary_emb = self.rotary(L)
+                q, k = map(lambda x: apply_rotary_pos_emb(rotary_emb, x), (q, k))
+            else:
+                q_rotary_emb = self.rotary(L)
+                k_rotary_emb = self.rotary(S)
+                q = apply_rotary_pos_emb(q_rotary_emb, q)
+                k = apply_rotary_pos_emb(k_rotary_emb, k)
+
+        if mask is not None:
+            attn_bias.masked_fill_(mask.logical_not(), float("-inf"))
+
+        with torch.backends.cuda.sdp_kernel(**config._asdict()):
+            out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=attn_bias, dropout_p=self.attn_drop.p if self.training else 0.0
+            )
+
+        out = rearrange(out, "b h n d -> b n (h d)")
+        out = self.o_dropout(self.o_proj(out))
+        return out
+
+
+# --------------- DiT Modules ---------------
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
+            device=t.device
+        )
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, output_size):
+        super().__init__()
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = Linear(hidden_size, output_size, bias=True)
+
+    def forward(self, x, c, mask=None):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+# --------------- Transformer ---------------
+
+
+class Mlp(nn.Module):
+    def __init__(self, hidden_size, ffn_hidden_size=4096, act_layer=nn.GELU, dropout=0.0, **kwargs):
+        super().__init__()
+        self.fc1 = Linear(hidden_size, ffn_hidden_size)
+        self.act = act_layer()
+        self.fc2 = Linear(ffn_hidden_size, hidden_size)
+        self.drop = Dropout(dropout)
+
+    def forward(self, x, **kwargs):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        attention: nn.Module,
+        ffn: nn.Module,
+        hidden_size: int = 1024,
+        modulation: bool = False,
+        eps: float = 1e-6,
+        **kwargs,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps)
+        self.attn = attention
+        self.ffn = ffn
+        self.modulation = modulation
+        if modulation:
+            self.modulation_layer = nn.Sequential(nn.SiLU(), Linear(hidden_size, 6 * hidden_size, bias=True))
+            nn.init.constant_(self.modulation_layer[-1].weight, 0.0)
+            nn.init.constant_(self.modulation_layer[-1].bias, 0.0)
+
+    def forward(self, x, condition=None, mask=None):
+        if condition is None:
+            assert not self.modulation, "Without global condition, must set modulation to False"
+        else:
+            assert self.modulation, "With global condition, must set modulation to True"
+            shift_attn, scale_attn, gate_attn, shift_ffn, scale_ffn, gate_ffn = self.modulation_layer(condition).chunk(
+                6, dim=1
+            )
+
+        if condition is not None:
+            x = x + gate_attn.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_attn, scale_attn), mask=mask)
+        else:
+            x = x + self.attn(self.norm1(x), mask=mask)
+
+        if condition is not None:
+            x = x + gate_ffn.unsqueeze(1) * self.ffn(modulate(self.norm2(x), shift_ffn, scale_ffn))
+        else:
+            x = x + self.ffn(self.norm2(x), mask=mask)
+        return x
+
+
+# --------------- Flow Matching: Token -> Latent ---------------
+
+
+class Token2latentFlowMatching(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.model_dim = config.transformer.hidden_size
+        self.token_input_dim = config.get("token_input_dim", self.model_dim)
+        self.target_dim = config.z_dim
+
+        self.spkr_embed_dim = config.get("spkr_embed_dim", 512)
+        self.cond_proj = Linear(self.model_dim + self.spkr_embed_dim, self.model_dim)
+        self.token_pad_id = -1
+        if config.upsample_factor > 1:
+            self.token_proj = nn.Sequential(
+                Linear(self.token_input_dim, self.model_dim, bias=True),
+                ConvTranspose1d(
+                    self.model_dim,
+                    self.model_dim,
+                    stride=config.upsample_factor,
+                    kernel_size=config.upsample_factor * 2,
+                ),
+            )
+        else:
+            self.token_proj = Linear(self.token_input_dim, self.model_dim)
+
+        self.transformer_input_proj = Linear(self.model_dim + self.target_dim * 2, self.model_dim)
+        self.time_embedder = TimestepEmbedder(self.model_dim)
+        self.conv_embed = ConvPositionEmbed(hidden_size=self.model_dim, kernel_size=31, groups=16)
+        self.blocks = nn.ModuleList()
+        for _ in range(config.transformer.num_layers):
+            attn_block = MultiHeadAttention(**config.transformer)
+            ffn_block = Mlp(act_layer=lambda: nn.GELU(approximate="tanh"), **config.transformer)
+            self.blocks.append(TransformerBlock(attn_block, ffn_block, **config.transformer))
+        self.output_layer = FinalLayer(config.transformer.hidden_size, self.target_dim)
+
+    def cond_mask_spkr_embed(self, x, spkr_embed):
+        b, device = x.size(0), x.device
+        if not exists(spkr_embed):
+            spkr_embed = torch.zeros(b, self.spkr_embed_dim, device=device, dtype=x.dtype)
+        return spkr_embed
+
+    def vectorfield_forward(self, inputs, times, self_attn_mask, g_cond=None):
+        t = self.time_embedder(times)
+        if g_cond is None:
+            c = t
+        else:
+            cond_inp = torch.cat([t, g_cond], dim=-1)
+            c = self.cond_proj(cond_inp)
+
+        ut = self.transformer_input_proj(inputs)
+        ut = self.conv_embed(ut, mask=self_attn_mask) + ut
+
+        for block in self.blocks:
+            ut = block(ut, c, mask=self_attn_mask)
+        ut = self.output_layer(ut, c)
+        return ut
+
+    @eval_decorator
+    @torch.no_grad()
+    def inference(
+        self,
+        *,
+        token: torch.Tensor,
+        prefix_target: torch.Tensor | None = None,
+        spkr_embed: torch.Tensor | None = None,
+        s_steps: int | None = 10,
+        cfg_alpha: float | None = 2.0,
+        rescale_logits: bool = False,
+        **kwargs,
+    ):
+        b, device = token.size(0), token.device
+        assert b == 1, "Only support batch_size == 1 when inference"
+
+        token = self.token_proj(token)
+        ge = self.cond_mask_spkr_embed(token, spkr_embed)
+        tgt_lens = token.size(1)
+        latent_cond = torch.zeros(b, tgt_lens, self.target_dim).to(device)
+        if prefix_target is not None:
+            latent_cond[:, : prefix_target.size(1), :] = prefix_target
+        sample, trajectory = self.sample(
+            tokens=token, audio=latent_cond, steps=s_steps, alpha=cfg_alpha, g_cond=ge, rescale_logits=rescale_logits
+        )
+        return sample
+
+    def sample(self, tokens, audio, steps, alpha=None, g_cond=None, rescale_logits=False):
+        noise = torch.randn([audio.size(0), audio.shape[1], self.target_dim], device=audio.device)
+        times = torch.linspace(0, 1, steps, device=audio.device)
+
+        def solver(t, z):
+            if alpha is None:
+                output = torch.cat([tokens, audio, z], dim=-1)
+                return self.vectorfield_forward(inputs=output, times=t.unsqueeze(0), self_attn_mask=None, g_cond=g_cond)
+            tokens_empty = torch.zeros(*audio.shape[:2], self.model_dim, device=tokens.device, dtype=tokens.dtype)
+            audio_empty = audio
+            tokens_t = torch.cat([tokens_empty, tokens], dim=0)
+            audio_t = torch.cat([audio_empty, audio], dim=0)
+            audio_noizy_t = torch.cat([z, z], dim=0)
+            t_t = torch.stack([t, t], dim=0)
+            c = g_cond
+            if g_cond is not None:
+                c = torch.cat([g_cond, g_cond], dim=0)
+            output = torch.cat([tokens_t, audio_t, audio_noizy_t], dim=-1)
+            predicted_mix = self.vectorfield_forward(inputs=output, times=t_t, self_attn_mask=None, g_cond=c)
+            predicted_conditioned = predicted_mix[1].unsqueeze(0)
+            predicted_unconditioned = predicted_mix[0].unsqueeze(0)
+
+            prediction = predicted_conditioned + (predicted_conditioned - predicted_unconditioned) * alpha
+            prediction_rescaled = predicted_conditioned.std() * (prediction / prediction.std())
+            if rescale_logits:
+                return prediction_rescaled
+            else:
+                return prediction
+
+        try:
+            from torchdiffeq import odeint
+        except ImportError:
+            raise ImportError("Covo-Audio code2wav requires `torchdiffeq`. Install it with: pip install torchdiffeq")
+        trajectory = odeint(solver, noise, times, atol=1e-5, rtol=1e-5, method="midpoint")
+        return trajectory[-1], trajectory
+
+
+class Token2latentFlowMatchingWithEmbed(Token2latentFlowMatching):
+    def __init__(self, config, **kwargs):
+        super().__init__(config)
+        self.vocab_size = config.token_vocab_size
+        self.token_embedding = EmbeddingTable(
+            num_embeddings=self.vocab_size, embedding_dim=self.token_input_dim, pad_id=self.token_pad_id
+        )
+
+    def inference(
+        self, *, token, prefix_target=None, spkr_embed=None, s_steps=10, cfg_alpha=2, rescale_logits=False, **kwargs
+    ):
+        token = self.token_embedding(token)
+        return super().inference(
+            token=token,
+            prefix_target=prefix_target,
+            spkr_embed=spkr_embed,
+            s_steps=s_steps,
+            cfg_alpha=cfg_alpha,
+            rescale_logits=rescale_logits,
+            **kwargs,
+        )
+
+
+# --------------- Vocoder: Latent -> Waveform (BigVGAN) ---------------
+
+
+class AMPBlock1(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5), causal=True):
+        super().__init__()
+        self.h = h
+
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], causal=causal)),
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], causal=causal)),
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], causal=causal)),
+            ]
+        )
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)),
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)),
+                weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)),
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+        self.num_layers = len(self.convs1) + len(self.convs2)
+
+        self.activations = nn.ModuleList(
+            [
+                Activation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale))
+                for _ in range(self.num_layers)
+            ]
+        )
+
+    def forward(self, x):
+        acts1, acts2 = self.activations[::2], self.activations[1::2]
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+            xt = a1(x)
+            xt = c1(xt)
+            xt = a2(xt)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+
+class BigVGANFlowVAE(nn.Module):
+    def __init__(self, h):
+        super().__init__()
+        self.config = h
+        self.h = h
+        causal = h.causal
+        self.hop_size = np.prod(h.downsample_rates)
+
+        self.num_kernels = len(h.resblock_kernel_sizes)
+        self.num_upsamples = len(h.upsample_rates)
+
+        self.conv_pre = weight_norm(Conv1d(h.latent_dim, h.upsample_initial_channel, 7, 1, causal=False))
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        weight_norm(
+                            ConvTranspose1d(
+                                h.upsample_initial_channel // (2**i),
+                                h.upsample_initial_channel // (2 ** (i + 1)),
+                                k,
+                                u,
+                                causal=causal,
+                            )
+                        )
+                    ]
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = h.upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
+                self.resblocks.append(AMPBlock1(h, ch, k, d, causal=causal))
+
+        activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale)
+        self.activation_post = Activation1d(activation=activation_post)
+
+        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, causal=causal))
+
+        for i in range(len(self.ups)):
+            self.ups[i].apply(init_weights)
+        self.conv_post.apply(init_weights)
+
+    def inference_from_latents(self, x, do_sample=True, noise_scale=1.0):
+        if self.h.use_vae and do_sample:
+            assert x.size(1) == self.h.latent_dim * 2, f"Input must be like [B, D, H], got {x.shape}"
+            m_q, logs_q = torch.split(x, self.h.latent_dim, dim=1)
+            x = m_q + torch.randn_like(m_q) * torch.exp(logs_q) * noise_scale
+        else:
+            assert x.size(1) == self.h.latent_dim, f"Input must be like [B, D, H], got {x.shape}"
+
+        x = self.conv_pre(x)
+
+        for i in range(self.num_upsamples):
+            for i_up in range(len(self.ups[i])):
+                x = self.ups[i][i_up](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+
+        x = self.activation_post(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+        return x
+
+
+# --------------- Token2WavDecoder ---------------
+
+
+class Token2WavDecoder(nn.Module):
+    def __init__(self, config, **kwargs):
+        super().__init__()
+        self.config = config
+
+        self.wavegan = BigVGANFlowVAE(config.wavegan)
+        self.wavegan_hop_size = np.prod(config.wavegan.downsample_rates)
+        self.global_mean_var = getattr(config, "global_mean_var", None)
+        if self.global_mean_var is not None:
+            mean_var_data = torch.from_numpy(np.load(self.global_mean_var)).float().squeeze()
+            global_mean, global_var = mean_var_data.chunk(2, 0)
+            self.register_buffer("global_mean", global_mean)
+            self.register_buffer("global_var", global_var)
+
+        self.token2latent = Token2latentFlowMatchingWithEmbed(config.token2latent)
+
+        self.upsample_factor = self.token2latent.config.get("upsample_factor", 1)
+        self.wav_input_sr = config.get("wav_input_sr", 24000)
+
+        self.trainable_module = ["wavegan", "token2latent"]
+
+    def state_dict(self):
+        param_dict = OrderedDict()
+        for name in self.trainable_module:
+            state = self.get_submodule(name).state_dict(prefix=f"{name}.")
+            param_dict.update(state)
+        return param_dict
+
+    def load_state_dict(self, param_dict):
+        for name in self.trainable_module:
+            module_state = OrderedDict()
+            name_len = len(name)
+            for k, v in param_dict.items():
+                if k.startswith(f"{name}."):
+                    new_k = k[name_len + 1 :]
+                    module_state[new_k] = v
+            self.get_submodule(name).load_state_dict(module_state, strict=False)
+
+    @torch.no_grad()
+    def preprocess_infer_data(self, data):
+        zero_spkr = data.get("zero_spkr", False)
+
+        token = data["target_token"]
+        prefix_target = None
+        res_prefix_len = 0
+        prompt_token = data.get("prompt_token", None)
+        prompt_latent = data.get("prompt_latent", None)
+        if exists(prompt_token):
+            token = torch.cat([prompt_token, token], dim=1)
+        if exists(prompt_latent):
+            prefix_target = prompt_latent
+            res_prefix_len = prompt_latent.shape[1]
+
+        spkr_embed = None
+        if not zero_spkr:
+            spkr_embed = data.get("spkr_embed", None)
+        res = {"token": token, "prefix_target": prefix_target, "spkr_embed": spkr_embed}
+        return res, res_prefix_len
+
+    @eval_decorator
+    @torch.no_grad()
+    def inference(self, data, **kwargs):
+        infer_data, prefix_len = self.preprocess_infer_data(data)
+        prefix_target = infer_data["prefix_target"]
+        res_latents = self.token2latent.inference(**infer_data, **kwargs)
+        if self.global_mean_var is not None:
+            res_latents = self.global_mean + res_latents * torch.sqrt(self.global_var)
+            if exists(prefix_target):
+                prefix_target = self.global_mean + prefix_target * torch.sqrt(self.global_var)
+        if exists(prefix_target):
+            res_latents[:, : prefix_target.shape[1]] = prefix_target
+        res_latents = res_latents[:, prefix_len:]
+        audio = self.wavegan.inference_from_latents(res_latents.transpose(1, 2), do_sample=False)
+        return audio

vllm_omni/model_executor/models/registry.py

@@ -174,6 +174,32 @@ _OMNI_MODELS = {
         "VoxtralTTSAudioGenerationForConditionalGeneration",
     ),
     "VoxtralTTSAudioTokenizer": ("voxtral_tts", "voxtral_tts_audio_tokenizer", "VoxtralTTSAudioTokenizer"),
+    ## covo_audio
+    "CovoAudioForCausalLM": (
+        "covo_audio",
+        "covo_audio",
+        "CovoAudioForConditionalGeneration",
+    ),
+    "CovoAudioForConditionalGeneration": (
+        "covo_audio",
+        "covo_audio",
+        "CovoAudioForConditionalGeneration",
+    ),
+    "CovoAudioModel": (
+        "covo_audio",
+        "covo_audio",
+        "CovoAudioForConditionalGeneration",
+    ),
+    "CovoAudioLLMModel": (
+        "covo_audio",
+        "covo_audio_llm",
+        "CovoAudioLLMForConditionalGeneration",
+    ),
+    "CovoAudioCode2WavModel": (
+        "covo_audio",
+        "covo_audio_code2wav",
+        "CovoAudioCode2WavForConditionalGeneration",
+    ),
     ## MOSS-TTS-Nano
     "MossTTSNanoForCausalLM": (
         "moss_tts_nano",

vllm_omni/model_executor/stage_input_processors/covo_audio.py

@@ -0,0 +1,33 @@
+# Copyright 2026 Tencent.
+from typing import Any
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+from vllm_omni.model_executor.models.covo_audio.config_covo_audio import COVO_AUDIO_TOKEN_INDEX
+
+
+def llm2code2wav(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    talker_outputs = source_outputs
+    code2wav_inputs = []
+
+    for i, talker_output in enumerate(talker_outputs):
+        output = talker_output.outputs[0]
+        token_ids = output.token_ids
+
+        audio_codes = [t - COVO_AUDIO_TOKEN_INDEX for t in token_ids if t >= COVO_AUDIO_TOKEN_INDEX]
+
+        if not audio_codes:
+            audio_codes = [-1]
+
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=audio_codes,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+
+    return code2wav_inputs