diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 8edafec1e..f30f34758 100644 --- a/docs/models/supported_models.md +++ b/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. diff --git a/examples/offline_inference/covo_audio/README.md b/examples/offline_inference/covo_audio/README.md new file mode 100644 index 000000000..d13783d39 --- /dev/null +++ b/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 diff --git a/examples/offline_inference/covo_audio/end2end.py b/examples/offline_inference/covo_audio/end2end.py new file mode 100644 index 000000000..ae7e5c6bb --- /dev/null +++ b/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) diff --git a/examples/online_serving/covo_audio/README.md b/examples/online_serving/covo_audio/README.md new file mode 100644 index 000000000..6894507ef --- /dev/null +++ b/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. diff --git a/examples/online_serving/covo_audio/openai_chat_completion_client.py b/examples/online_serving/covo_audio/openai_chat_completion_client.py new file mode 100644 index 000000000..d3aa042d5 --- /dev/null +++ b/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) diff --git a/recipes/README.md b/recipes/README.md index de07d22da..bccc4b222 100644 --- a/recipes/README.md +++ b/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 diff --git a/recipes/Tencent/Covo-Audio-Chat.md b/recipes/Tencent/Covo-Audio-Chat.md new file mode 100644 index 000000000..4f2eab635 --- /dev/null +++ b/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. diff --git a/tests/e2e/offline_inference/test_covo_audio_expansion.py b/tests/e2e/offline_inference/test_covo_audio_expansion.py new file mode 100644 index 000000000..3c6e5efb7 --- /dev/null +++ b/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) diff --git a/tests/e2e/online_serving/test_covo_audio_expansion.py b/tests/e2e/online_serving/test_covo_audio_expansion.py new file mode 100644 index 000000000..a1baca554 --- /dev/null +++ b/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) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 3d44d1bff..1067316a2 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/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", diff --git a/vllm_omni/deploy/covo_audio.yaml b/vllm_omni/deploy/covo_audio.yaml new file mode 100644 index 000000000..613a91b4a --- /dev/null +++ b/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 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 05a52b537..256b9a56b 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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": diff --git a/vllm_omni/model_executor/models/covo_audio/__init__.py b/vllm_omni/model_executor/models/covo_audio/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py b/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py new file mode 100644 index 000000000..8153c2f38 --- /dev/null +++ b/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, + } diff --git a/vllm_omni/model_executor/models/covo_audio/covo_audio.py b/vllm_omni/model_executor/models/covo_audio/covo_audio.py new file mode 100644 index 000000000..42f80bb74 --- /dev/null +++ b/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 + }, + ) diff --git a/vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py b/vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py new file mode 100644 index 000000000..4b0b2fe4b --- /dev/null +++ b/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()} diff --git a/vllm_omni/model_executor/models/covo_audio/covo_audio_llm.py b/vllm_omni/model_executor/models/covo_audio/covo_audio_llm.py new file mode 100644 index 000000000..09a9d6b66 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/covo_audio/pipeline.py b/vllm_omni/model_executor/models/covo_audio/pipeline.py new file mode 100644 index 000000000..5b1a31d6e --- /dev/null +++ b/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}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/covo_audio/prompt_utils.py b/vllm_omni/model_executor/models/covo_audio/prompt_utils.py new file mode 100644 index 000000000..99338ae8e --- /dev/null +++ b/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, + ) diff --git a/vllm_omni/model_executor/models/covo_audio/speaker_prompt/prompt_latent.npy b/vllm_omni/model_executor/models/covo_audio/speaker_prompt/prompt_latent.npy new file mode 100644 index 000000000..33b9ee069 Binary files /dev/null and b/vllm_omni/model_executor/models/covo_audio/speaker_prompt/prompt_latent.npy differ diff --git a/vllm_omni/model_executor/models/covo_audio/speaker_prompt/prompt_token.npy b/vllm_omni/model_executor/models/covo_audio/speaker_prompt/prompt_token.npy new file mode 100644 index 000000000..7610be3a8 Binary files /dev/null and b/vllm_omni/model_executor/models/covo_audio/speaker_prompt/prompt_token.npy differ diff --git a/vllm_omni/model_executor/models/covo_audio/speaker_prompt/speaker_embed.npy b/vllm_omni/model_executor/models/covo_audio/speaker_prompt/speaker_embed.npy new file mode 100644 index 000000000..67c315755 Binary files /dev/null and b/vllm_omni/model_executor/models/covo_audio/speaker_prompt/speaker_embed.npy differ diff --git a/vllm_omni/model_executor/models/covo_audio/token2wav.py b/vllm_omni/model_executor/models/covo_audio/token2wav.py new file mode 100644 index 000000000..f6bed125c --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 5c2ef65d0..cc900d471 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/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", diff --git a/vllm_omni/model_executor/stage_input_processors/covo_audio.py b/vllm_omni/model_executor/stage_input_processors/covo_audio.py new file mode 100644 index 000000000..a0a964bdd --- /dev/null +++ b/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