diff --git a/docs/.nav.yml b/docs/.nav.yml index 4aa8fee43..b9f39ebfa 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -13,9 +13,13 @@ nav: - examples/README.md - Offline Inference: - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md + - Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md + - Helios Video Generation: user_guide/examples/offline_inference/helios.md + - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md + - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md @@ -25,6 +29,8 @@ nav: - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md - Online Serving: - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md + - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md + - Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md - Image-To-Image: user_guide/examples/online_serving/image_to_image.md - Image-To-Video: user_guide/examples/online_serving/image_to_video.md diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 52e75989d..4c709425e 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -50,6 +50,7 @@ th { |`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | |`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | |`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | +|`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ## List of Supported Models for NPU diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index c346b5b97..46b9154bb 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -1,6 +1,9 @@ # Speech API -vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation using Qwen3-TTS models. +vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. Supported TTS models include: + +- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz. +- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz. Each server instance runs a single model (specified at startup via `vllm serve --omni`). @@ -9,13 +12,22 @@ Each server instance runs a single model (specified at startup via `vllm serve < ### Start the Server ```bash -# CustomVoice model (predefined speakers) +# Qwen3-TTS: CustomVoice model (predefined speakers) vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ --omni \ --port 8091 \ --trust-remote-code \ --enforce-eager + +# Fish Speech S2 Pro +vllm-omni serve fishaudio/s2-pro \ + --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager \ + --gpu-memory-utilization 0.9 ``` ### Generate Speech @@ -317,6 +329,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \ ## Supported Models +### Qwen3-TTS + | Model | Task Type | Description | |-------|-----------|-------------| | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control | @@ -325,6 +339,14 @@ curl -X POST http://localhost:8091/v1/audio/speech \ | `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant | | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning | +### Fish Speech S2 Pro + +| Model | Description | +|-------|-------------| +| `fishaudio/s2-pro` | 4B dual-AR TTS with DAC codec (44.1 kHz). Supports text-to-speech and voice cloning. | + +Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech online serving example](../user_guide/examples/online_serving/fish_speech.md) for details. + ## Error Responses ### 400 Bad Request diff --git a/docs/user_guide/examples/offline_inference/fish_speech.md b/docs/user_guide/examples/offline_inference/fish_speech.md new file mode 100644 index 000000000..dd9066afd --- /dev/null +++ b/docs/user_guide/examples/offline_inference/fish_speech.md @@ -0,0 +1,97 @@ +# Fish Speech S2 Pro + +Source . + + +This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio. + +## Model Overview + +[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses: + +- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone) +- **Fast AR**: Predicts residual codebook tokens from semantic tokens +- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform + +The model produces high-quality speech with voice cloning capabilities. + +## Prerequisites + +Install the `fish-speech` package for the DAC codec: + +```bash +pip install fish-speech +``` + +## 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. + +## Quick Start + +### Text-Only Synthesis + +```bash +python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system." +``` + +Generated audio files are saved to `output_audio/` by default. + +### Voice Cloning + +Provide a reference audio file and its transcript to clone a voice: + +```bash +python end2end.py \ + --text "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Transcript of the reference audio." +``` + +## Streaming Mode + +Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config): + +```bash +python end2end.py --text "Hello, this is a test." --streaming +``` + +Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes. + +## Configuration + +| Argument | Default | Description | +|----------|---------|-------------| +| `--text` | `"Hello, this is a test..."` | Text to synthesize | +| `--ref-audio` | None | Path to reference audio for voice cloning | +| `--ref-text` | None | Transcript of the reference audio | +| `--model` | `fishaudio/s2-pro` | HuggingFace model path | +| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML | +| `--output-dir` | `output_audio` | Output directory for WAV files | +| `--streaming` | False | Enable streaming via AsyncOmni | +| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) | + +## Notes + +- Output audio is 44.1 kHz mono WAV. +- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory. +- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message. + +## Example materials + +??? abstract "diag_codebook.py" + ``````py + --8<-- "examples/offline_inference/fish_speech/diag_codebook.py" + `````` +??? abstract "diag_reference_compare.py" + ``````py + --8<-- "examples/offline_inference/fish_speech/diag_reference_compare.py" + `````` +??? abstract "end2end.py" + ``````py + --8<-- "examples/offline_inference/fish_speech/end2end.py" + `````` +??? abstract "standalone_test.py" + ``````py + --8<-- "examples/offline_inference/fish_speech/standalone_test.py" + `````` diff --git a/docs/user_guide/examples/online_serving/fish_speech.md b/docs/user_guide/examples/online_serving/fish_speech.md new file mode 100644 index 000000000..2c76b1bca --- /dev/null +++ b/docs/user_guide/examples/online_serving/fish_speech.md @@ -0,0 +1,194 @@ +# Fish Speech S2 Pro + +Source . + + +## Model + +| Model | Description | +|-------|-------------| +| `fishaudio/s2-pro` | Fish Speech S2 Pro -- 4B dual-AR TTS model with DAC codec (44.1 kHz) | + +## Gradio Demo + +An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support. + +```bash +# Option 1: Launch server + Gradio together +./run_gradio_demo.sh + +# Option 2: If server is already running +python gradio_demo.py --api-base http://localhost:8091 +``` + +Then open http://localhost:7860 in your browser. + +Features: + +- Text-to-speech synthesis +- Voice cloning from uploaded audio or URL +- Streaming mode (progressive PCM playback) + +## Launch the Server + +```bash +vllm-omni serve fishaudio/s2-pro \ + --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager \ + --gpu-memory-utilization 0.9 +``` + +Or use the convenience script: + +```bash +./run_server.sh +``` + +## Send TTS Request + +### Using curl + +```bash +# Basic TTS +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "default", + "response_format": "wav" + }' --output output.wav +``` + +### Voice Cloning + +Provide a reference audio (URL or base64 data URL) and its transcript: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is a cloned voice.", + "voice": "default", + "ref_audio": "https://example.com/reference.wav", + "ref_text": "Transcript of the reference audio." + }' --output cloned.wav +``` + +### Using Python + +```python +import httpx + +# Basic TTS +response = httpx.post( + "http://localhost:8091/v1/audio/speech", + json={ + "input": "Hello, how are you?", + "voice": "default", + "response_format": "wav", + }, + timeout=300.0, +) + +with open("output.wav", "wb") as f: + f.write(response.content) +``` + +### Using the CLI Client + +```bash +cd examples/online_serving/fish_speech + +# Basic TTS +python speech_client.py --text "Hello, how are you?" + +# Voice cloning +python speech_client.py \ + --text "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Transcript of the reference audio." + +# Streaming PCM output +python speech_client.py --text "Hello world" --stream --output output.pcm +``` + +The CLI client supports: + +- `--api-base`: API base URL (default: `http://localhost:8091`) +- `--model` (or `-m`): Model name (default: `fishaudio/s2-pro`) +- `--text`: Text to synthesize (required) +- `--ref-audio`: Reference audio for voice cloning (local path or URL) +- `--ref-text`: Transcript of the reference audio +- `--stream`: Enable streaming (PCM output) +- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav) +- `--output` (or `-o`): Output file path + +## Streaming + +Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "default", + "stream": true, + "response_format": "pcm" + }' --no-buffer | play -t raw -r 44100 -e signed -b 16 -c 1 - +``` + +**Note:** Fish Speech outputs at 44.1 kHz (unlike Qwen3-TTS which outputs at 24 kHz). + +## API Parameters + +Fish Speech uses the same `/v1/audio/speech` endpoint as Qwen3-TTS. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full parameter documentation. + +Key parameters for Fish Speech: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `input` | string | **required** | Text to synthesize | +| `voice` | string | "default" | Voice name (use "default" for Fish Speech) | +| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus | +| `ref_audio` | string | null | Reference audio URL or base64 data URL for voice cloning | +| `ref_text` | string | null | Transcript of reference audio (required for voice cloning) | +| `max_new_tokens` | int | 4096 | Maximum tokens to generate | +| `stream` | bool | false | Stream raw PCM chunks | + +## Prerequisites + +Install the `fish-speech` package for the DAC codec: + +```bash +pip install fish-speech +``` + +## Troubleshooting + +1. **No audio output**: Make sure the `fish-speech` package is installed for the DAC decoder +2. **Connection refused**: Ensure the server is running on the correct port +3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings +4. **Out of memory**: Reduce `--gpu-memory-utilization` or use a GPU with more VRAM + +## Example materials + +??? abstract "gradio_demo.py" + ``````py + --8<-- "examples/online_serving/fish_speech/gradio_demo.py" + `````` +??? abstract "run_gradio_demo.sh" + ``````sh + --8<-- "examples/online_serving/fish_speech/run_gradio_demo.sh" + `````` +??? abstract "run_server.sh" + ``````sh + --8<-- "examples/online_serving/fish_speech/run_server.sh" + `````` +??? abstract "speech_client.py" + ``````py + --8<-- "examples/online_serving/fish_speech/speech_client.py" + `````` diff --git a/examples/offline_inference/fish_speech/README.md b/examples/offline_inference/fish_speech/README.md new file mode 100644 index 000000000..9af853096 --- /dev/null +++ b/examples/offline_inference/fish_speech/README.md @@ -0,0 +1,75 @@ +# Fish Speech S2 Pro + +This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio. + +## Model Overview + +[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses: + +- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone) +- **Fast AR**: Predicts residual codebook tokens from semantic tokens +- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform + +The model produces high-quality speech with voice cloning capabilities. + +## Prerequisites + +Install the `fish-speech` package for the DAC codec: + +```bash +pip install fish-speech +``` + +## 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. + +## Quick Start + +### Text-Only Synthesis + +```bash +python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system." +``` + +Generated audio files are saved to `output_audio/` by default. + +### Voice Cloning + +Provide a reference audio file and its transcript to clone a voice: + +```bash +python end2end.py \ + --text "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Transcript of the reference audio." +``` + +## Streaming Mode + +Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config): + +```bash +python end2end.py --text "Hello, this is a test." --streaming +``` + +Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes. + +## Configuration + +| Argument | Default | Description | +|----------|---------|-------------| +| `--text` | `"Hello, this is a test..."` | Text to synthesize | +| `--ref-audio` | None | Path to reference audio for voice cloning | +| `--ref-text` | None | Transcript of the reference audio | +| `--model` | `fishaudio/s2-pro` | HuggingFace model path | +| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML | +| `--output-dir` | `output_audio` | Output directory for WAV files | +| `--streaming` | False | Enable streaming via AsyncOmni | +| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) | + +## Notes + +- Output audio is 44.1 kHz mono WAV. +- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory. +- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message. diff --git a/examples/offline_inference/fish_speech/end2end.py b/examples/offline_inference/fish_speech/end2end.py new file mode 100644 index 000000000..9e95f7e5c --- /dev/null +++ b/examples/offline_inference/fish_speech/end2end.py @@ -0,0 +1,252 @@ +"""Offline inference demo for Fish Speech S2 Pro via vLLM Omni. + +Generates speech from text using the fishaudio/s2-pro model. +Supports both text-only synthesis and voice cloning with reference audio. + +Usage: + # Text-only synthesis + python end2end.py --text "Hello, this is a test." + + # Voice cloning with reference audio + python end2end.py --text "Hello, this is a test." --ref-audio ref.wav --ref-text "Reference text." + + # Streaming mode + python end2end.py --text "Hello, this is a test." --streaming +""" + +import asyncio +import logging +import os +import time + +import soundfile as sf +import torch + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import AsyncOmni, Omni + +logger = logging.getLogger(__name__) + +DEFAULT_MODEL = "fishaudio/s2-pro" +DEFAULT_STAGE_CONFIG = os.path.join( + os.path.dirname(__file__), + "..", + "..", + "..", + "vllm_omni", + "model_executor", + "stage_configs", + "fish_speech_s2_pro.yaml", +) + + +def build_prompt( + text: str, + ref_audio_path: str | None = None, + ref_text: str | None = None, + model_name: str = DEFAULT_MODEL, +) -> dict: + """Build a prompt for Fish Speech S2 Pro. + + Uses Qwen3 chat template format with <|voice|> modality marker: + <|im_start|>user\n<|speaker:0|>text<|im_end|>\n<|im_start|>assistant\n<|voice|> + + The model then generates semantic tokens autoregressively. + """ + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + + # Build user message with speaker tag prefix. + user_text = f"<|speaker:0|>{text}" + + # Use chat template for user turn + assistant generation prompt. + messages = [{"role": "user", "content": user_text}] + prompt_ids = tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + ) + + # Append <|voice|> token (151673) to signal voice generation mode. + voice_token_id = tokenizer.encode("<|voice|>", add_special_tokens=False) + prompt_ids = prompt_ids + voice_token_id + + additional_information = { + "text": [text], + "max_new_tokens": [4096], + } + + if ref_audio_path: + additional_information["ref_audio"] = [ref_audio_path] + if ref_text: + additional_information["ref_text"] = [ref_text] + + return { + "prompt_token_ids": prompt_ids, + "additional_information": additional_information, + } + + +def _save_wav(output_dir: str, request_id: str, mm: dict) -> None: + """Concatenate audio chunks and write to a wav file.""" + audio_data = mm["audio"] + sr_raw = mm["sr"] + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val) + audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data + out_wav = os.path.join(output_dir, f"output_{request_id}.wav") + sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV") + logger.info("Request %s: saved audio to %s (sr=%d)", request_id, out_wav, sr) + + +def main(args): + """Run offline inference with Omni.""" + output_dir = args.output_dir + os.makedirs(output_dir, exist_ok=True) + + stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG + model_name = args.model or DEFAULT_MODEL + + inputs = [ + build_prompt( + text=args.text, + ref_audio_path=args.ref_audio, + ref_text=args.ref_text, + model_name=model_name, + ) + ] + + omni = Omni( + model=model_name, + stage_configs_path=stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + + t_start = time.perf_counter() + for stage_outputs in omni.generate(inputs): + for output in stage_outputs.request_output: + _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output) + t_end = time.perf_counter() + logger.info("Total inference time: %.1f ms", (t_end - t_start) * 1000) + + +async def main_streaming(args): + """Run offline inference with AsyncOmni for streaming audio chunks.""" + output_dir = args.output_dir + os.makedirs(output_dir, exist_ok=True) + + stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG + model_name = args.model or DEFAULT_MODEL + + prompt = build_prompt( + text=args.text, + ref_audio_path=args.ref_audio, + ref_text=args.ref_text, + model_name=model_name, + ) + + omni = AsyncOmni( + model=model_name, + stage_configs_path=stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + + request_id = "0" + t_start = time.perf_counter() + t_prev = t_start + chunk_idx = 0 + + async for stage_output in omni.generate(prompt, request_id=request_id): + mm = stage_output.request_output.outputs[0].multimodal_output + if not stage_output.finished: + t_now = time.perf_counter() + audio = mm.get("audio") + n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1) + dt_ms = (t_now - t_prev) * 1000 + ttfa_ms = (t_now - t_start) * 1000 + if chunk_idx == 0: + logger.info("Request %s: chunk %d samples=%d TTFA=%.1fms", request_id, chunk_idx, n, ttfa_ms) + else: + logger.info("Request %s: chunk %d samples=%d inter_chunk=%.1fms", request_id, chunk_idx, n, dt_ms) + t_prev = t_now + chunk_idx += 1 + else: + t_end = time.perf_counter() + total_ms = (t_end - t_start) * 1000 + logger.info("Request %s: done total=%.1fms chunks=%d", request_id, total_ms, chunk_idx) + _save_wav(output_dir, request_id, mm) + + +def parse_args(): + parser = FlexibleArgumentParser( + description="Fish Speech S2 Pro offline inference with vLLM Omni", + ) + parser.add_argument( + "--text", + type=str, + default="Hello, this is a test of the Fish Speech text to speech system.", + help="Text to synthesize.", + ) + parser.add_argument( + "--ref-audio", + type=str, + default=None, + help="Path to reference audio for voice cloning.", + ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Reference text corresponding to the reference audio.", + ) + parser.add_argument( + "--model", + type=str, + default=None, + help=f"HuggingFace model path (default: {DEFAULT_MODEL}).", + ) + parser.add_argument( + "--stage-configs-path", + type=str, + default=None, + help="Path to stage configs YAML.", + ) + parser.add_argument( + "--output-dir", + default="output_audio", + help="Output directory for generated wav files.", + ) + parser.add_argument( + "--log-stats", + action="store_true", + default=False, + help="Enable writing detailed statistics.", + ) + parser.add_argument( + "--stage-init-timeout", + type=int, + default=600, + help="Timeout for initializing stages in seconds.", + ) + parser.add_argument( + "--streaming", + action="store_true", + default=False, + help="Stream audio chunks as they arrive via AsyncOmni.", + ) + return parser.parse_args() + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + args = parse_args() + if args.streaming: + asyncio.run(main_streaming(args)) + else: + main(args) diff --git a/examples/online_serving/fish_speech/README.md b/examples/online_serving/fish_speech/README.md new file mode 100644 index 000000000..ae968d3ba --- /dev/null +++ b/examples/online_serving/fish_speech/README.md @@ -0,0 +1,172 @@ +# Fish Speech S2 Pro + +## Model + +| Model | Description | +|-------|-------------| +| `fishaudio/s2-pro` | Fish Speech S2 Pro -- 4B dual-AR TTS model with DAC codec (44.1 kHz) | + +## Gradio Demo + +An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support. + +```bash +# Option 1: Launch server + Gradio together +./run_gradio_demo.sh + +# Option 2: If server is already running +python gradio_demo.py --api-base http://localhost:8091 +``` + +Then open http://localhost:7860 in your browser. + +Features: + +- Text-to-speech synthesis +- Voice cloning from uploaded audio or URL +- Streaming mode (progressive PCM playback) + +## Launch the Server + +```bash +vllm-omni serve fishaudio/s2-pro \ + --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager \ + --gpu-memory-utilization 0.9 +``` + +Or use the convenience script: + +```bash +./run_server.sh +``` + +## Send TTS Request + +### Using curl + +```bash +# Basic TTS +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "default", + "response_format": "wav" + }' --output output.wav +``` + +### Voice Cloning + +Provide a reference audio (URL or base64 data URL) and its transcript: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is a cloned voice.", + "voice": "default", + "ref_audio": "https://example.com/reference.wav", + "ref_text": "Transcript of the reference audio." + }' --output cloned.wav +``` + +### Using Python + +```python +import httpx + +# Basic TTS +response = httpx.post( + "http://localhost:8091/v1/audio/speech", + json={ + "input": "Hello, how are you?", + "voice": "default", + "response_format": "wav", + }, + timeout=300.0, +) + +with open("output.wav", "wb") as f: + f.write(response.content) +``` + +### Using the CLI Client + +```bash +cd examples/online_serving/fish_speech + +# Basic TTS +python speech_client.py --text "Hello, how are you?" + +# Voice cloning +python speech_client.py \ + --text "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Transcript of the reference audio." + +# Streaming PCM output +python speech_client.py --text "Hello world" --stream --output output.pcm +``` + +The CLI client supports: + +- `--api-base`: API base URL (default: `http://localhost:8091`) +- `--model` (or `-m`): Model name (default: `fishaudio/s2-pro`) +- `--text`: Text to synthesize (required) +- `--ref-audio`: Reference audio for voice cloning (local path or URL) +- `--ref-text`: Transcript of the reference audio +- `--stream`: Enable streaming (PCM output) +- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav) +- `--output` (or `-o`): Output file path + +## Streaming + +Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "default", + "stream": true, + "response_format": "pcm" + }' --no-buffer | play -t raw -r 44100 -e signed -b 16 -c 1 - +``` + +**Note:** Fish Speech outputs at 44.1 kHz (unlike Qwen3-TTS which outputs at 24 kHz). + +## API Parameters + +Fish Speech uses the same `/v1/audio/speech` endpoint as Qwen3-TTS. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full parameter documentation. + +Key parameters for Fish Speech: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `input` | string | **required** | Text to synthesize | +| `voice` | string | "default" | Voice name (use "default" for Fish Speech) | +| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus | +| `ref_audio` | string | null | Reference audio URL or base64 data URL for voice cloning | +| `ref_text` | string | null | Transcript of reference audio (required for voice cloning) | +| `max_new_tokens` | int | 4096 | Maximum tokens to generate | +| `stream` | bool | false | Stream raw PCM chunks | + +## Prerequisites + +Install the `fish-speech` package for the DAC codec: + +```bash +pip install fish-speech +``` + +## Troubleshooting + +1. **No audio output**: Make sure the `fish-speech` package is installed for the DAC decoder +2. **Connection refused**: Ensure the server is running on the correct port +3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings +4. **Out of memory**: Reduce `--gpu-memory-utilization` or use a GPU with more VRAM diff --git a/examples/online_serving/fish_speech/gradio_demo.py b/examples/online_serving/fish_speech/gradio_demo.py new file mode 100644 index 000000000..a8e51b274 --- /dev/null +++ b/examples/online_serving/fish_speech/gradio_demo.py @@ -0,0 +1,272 @@ +"""Gradio demo for Fish Speech S2 Pro online serving via /v1/audio/speech API. + +Supports: + - Text-to-speech synthesis + - Voice cloning from reference audio (upload or URL) + - Streaming (progressive PCM chunks) and non-streaming modes + +Usage: + # Start the server first (see run_server.sh), then: + python gradio_demo.py --api-base http://localhost:8091 + + # Or use run_gradio_demo.sh to start both server and demo together. +""" + +import argparse +import base64 +import io + +import gradio as gr +import httpx +import numpy as np +import soundfile as sf + +PCM_SAMPLE_RATE = 44100 +DEFAULT_API_BASE = "http://localhost:8091" + + +def encode_audio_to_base64(audio_data: tuple) -> str: + """Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL.""" + sample_rate, audio_np = audio_data + if audio_np.dtype != np.int16: + if audio_np.dtype in (np.float32, np.float64): + audio_np = np.clip(audio_np, -1.0, 1.0) + audio_np = (audio_np * 32767).astype(np.int16) + else: + audio_np = audio_np.astype(np.int16) + buf = io.BytesIO() + sf.write(buf, audio_np, sample_rate, format="WAV") + wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8") + return f"data:audio/wav;base64,{wav_b64}" + + +def build_payload( + text: str, + ref_audio: tuple | None, + ref_audio_url: str, + ref_text: str, + response_format: str = "wav", + stream: bool = False, +) -> dict: + """Build the /v1/audio/speech request payload.""" + if not text or not text.strip(): + raise gr.Error("Please enter text to synthesize.") + + payload: dict = { + "input": text.strip(), + "voice": "default", + "response_format": "pcm" if stream else response_format, + "stream": stream, + } + + # Voice cloning: ref_audio takes priority over URL. + ref_url_stripped = ref_audio_url.strip() if ref_audio_url else "" + if ref_audio is not None: + payload["ref_audio"] = encode_audio_to_base64(ref_audio) + elif ref_url_stripped: + payload["ref_audio"] = ref_url_stripped + + if "ref_audio" in payload: + if not ref_text or not ref_text.strip(): + raise gr.Error("Voice cloning requires a transcript of the reference audio.") + payload["ref_text"] = ref_text.strip() + + return payload + + +def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format): + """Non-streaming: call /v1/audio/speech and return full audio.""" + payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=False) + + try: + with httpx.Client(timeout=300.0) as client: + resp = client.post( + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) + except httpx.TimeoutException: + raise gr.Error("Request timed out. The server may be busy.") + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") + + if resp.status_code != 200: + raise gr.Error(f"Server error ({resp.status_code}): {resp.text}") + + content_type = resp.headers.get("content-type", "") + if "application/json" in content_type: + try: + raise gr.Error(f"Server error: {resp.json()}") + except ValueError: + pass + + try: + if response_format == "pcm": + audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0 + return (PCM_SAMPLE_RATE, audio_np) + audio_np, sample_rate = sf.read(io.BytesIO(resp.content)) + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + return (sample_rate, audio_np.astype(np.float32)) + except Exception as e: + raise gr.Error(f"Failed to decode audio: {e}") + + +def generate_speech_stream(api_base, text, ref_audio, ref_audio_url, ref_text, response_format, chunk_seconds=0.5): + """Streaming: yield PCM chunks for progressive playback.""" + payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=True) + + min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds) + prebuffer_chunks = 2 + pending: list[np.ndarray] = [] + pending_len = 0 + chunks_yielded = 0 + prebuffer: list[np.ndarray] = [] + + try: + leftover = b"" + with httpx.Client(timeout=300.0) as client: + with client.stream( + "POST", + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) as resp: + if resp.status_code != 200: + resp.read() + raise gr.Error(f"Server error ({resp.status_code}): {resp.text}") + for chunk in resp.iter_bytes(): + if not chunk: + continue + raw = leftover + chunk + usable = len(raw) - (len(raw) % 2) + leftover = raw[usable:] + if usable == 0: + continue + float_samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0 + pending.append(float_samples) + pending_len += len(float_samples) + if pending_len >= min_chunk_samples: + audio_chunk = np.concatenate(pending) + pending.clear() + pending_len = 0 + if chunks_yielded < prebuffer_chunks: + prebuffer.append(audio_chunk) + chunks_yielded += 1 + if chunks_yielded == prebuffer_chunks: + yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer)) + prebuffer.clear() + else: + yield (PCM_SAMPLE_RATE, audio_chunk) + remaining = prebuffer + pending + if remaining: + yield (PCM_SAMPLE_RATE, np.concatenate(remaining)) + except httpx.TimeoutException: + raise gr.Error("Request timed out. The server may be busy.") + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") + + +def on_stream_change(stream: bool): + """Lock format to PCM when streaming.""" + if stream: + return gr.update(value="pcm", interactive=False) + return gr.update(value="wav", interactive=True) + + +def build_interface(api_base: str, stream_chunk_seconds: float = 0.5): + """Build the Gradio interface.""" + with gr.Blocks(title="Fish Speech S2 Pro Demo") as demo: + gr.Markdown("# Fish Speech S2 Pro - Text to Speech") + gr.Markdown(f"**Server:** `{api_base}` | **Model:** fishaudio/s2-pro | **Output:** 44.1kHz") + + with gr.Row(): + with gr.Column(scale=3): + text_input = gr.Textbox( + label="Text to Synthesize", + placeholder="Enter text here...", + lines=4, + ) + + with gr.Accordion("Voice Cloning (optional)", open=False): + gr.Markdown( + "Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice." + ) + ref_audio = gr.Audio( + label="Reference Audio", + type="numpy", + sources=["upload", "microphone"], + ) + ref_audio_url = gr.Textbox( + label="Reference Audio URL (alternative to upload)", + placeholder="https://example.com/reference.wav", + lines=1, + ) + ref_text = gr.Textbox( + label="Reference Audio Transcript (required for cloning)", + placeholder="Exact transcript of the reference audio...", + lines=2, + ) + + with gr.Row(): + response_format = gr.Dropdown( + choices=["wav", "mp3", "flac", "pcm"], + value="wav", + label="Audio Format", + scale=1, + ) + stream_checkbox = gr.Checkbox( + label="Stream output", + value=False, + info="Progressive PCM streaming", + scale=1, + ) + + generate_btn = gr.Button("Generate Speech", variant="primary", size="lg") + + with gr.Column(scale=2): + audio_output = gr.Audio( + label="Generated Audio", + interactive=False, + streaming=True, + autoplay=True, + ) + gr.Markdown( + "### About\n" + "- **Fish Speech S2 Pro** by FishAudio: 4B dual-AR model\n" + "- **Voice cloning**: upload 10-30s reference + transcript\n" + "- **Streaming**: real-time PCM output\n" + "- **44.1kHz** output via DAC codec" + ) + + stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format]) + + all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format] + + def dispatch(stream_enabled, *args): + if stream_enabled: + yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds) + else: + yield generate_speech(api_base, *args) + + generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output]) + demo.queue() + return demo + + +def main(): + parser = argparse.ArgumentParser(description="Gradio demo for Fish Speech S2 Pro") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})") + parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)") + parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)") + parser.add_argument("--share", action="store_true", help="Share publicly") + parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds per streaming chunk") + args = parser.parse_args() + + print(f"Connecting to vLLM server at: {args.api_base}") + demo = build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds) + demo.launch(server_name=args.host, server_port=args.port, share=args.share) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/fish_speech/run_gradio_demo.sh b/examples/online_serving/fish_speech/run_gradio_demo.sh new file mode 100755 index 000000000..98a696644 --- /dev/null +++ b/examples/online_serving/fish_speech/run_gradio_demo.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# Launch Fish Speech S2 Pro server + Gradio demo together. +# +# Usage: +# ./run_gradio_demo.sh +# CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh + +set -e + +MODEL="${MODEL:-fishaudio/s2-pro}" +PORT="${PORT:-8091}" +GRADIO_PORT="${GRADIO_PORT:-7860}" +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)" + +echo "Starting Fish Speech S2 Pro server (port $PORT)..." +FLASHINFER_DISABLE_VERSION_CHECK=1 \ +vllm-omni serve "$MODEL" \ + --stage-configs-path "$REPO_ROOT/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --gpu-memory-utilization 0.9 \ + --trust-remote-code \ + --enforce-eager \ + --omni & +SERVER_PID=$! + +cleanup() { + echo "Stopping server (PID $SERVER_PID)..." + kill $SERVER_PID 2>/dev/null + wait $SERVER_PID 2>/dev/null +} +trap cleanup EXIT + +# Wait for server to be ready. +echo "Waiting for server to start..." +for i in $(seq 1 120); do + if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then + echo "Server ready." + break + fi + sleep 2 +done + +echo "Starting Gradio demo (port $GRADIO_PORT)..." +python "$SCRIPT_DIR/gradio_demo.py" \ + --api-base "http://localhost:$PORT" \ + --port "$GRADIO_PORT" diff --git a/examples/online_serving/fish_speech/run_server.sh b/examples/online_serving/fish_speech/run_server.sh new file mode 100755 index 000000000..59c09c7fe --- /dev/null +++ b/examples/online_serving/fish_speech/run_server.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Launch vLLM-Omni server for Fish Speech S2 Pro +# +# Usage: +# ./run_server.sh +# CUDA_VISIBLE_DEVICES=0 ./run_server.sh + +set -e + +MODEL="${MODEL:-fishaudio/s2-pro}" +PORT="${PORT:-8091}" + +echo "Starting Fish Speech S2 Pro server with model: $MODEL" + +FLASHINFER_DISABLE_VERSION_CHECK=1 \ +vllm-omni serve "$MODEL" \ + --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \ + --host 0.0.0.0 \ + --port "$PORT" \ + --gpu-memory-utilization 0.9 \ + --trust-remote-code \ + --enforce-eager \ + --omni diff --git a/examples/online_serving/fish_speech/speech_client.py b/examples/online_serving/fish_speech/speech_client.py new file mode 100644 index 000000000..e3ace37e5 --- /dev/null +++ b/examples/online_serving/fish_speech/speech_client.py @@ -0,0 +1,129 @@ +"""Client for Fish Speech S2 Pro via /v1/audio/speech endpoint. + +Examples: + # Basic TTS + python speech_client.py --text "Hello, how are you?" + + # Voice cloning + python speech_client.py --text "Hello, how are you?" \ + --ref-audio ref.wav --ref-text "This is the reference transcript." + + # Streaming PCM output + python speech_client.py --text "Hello world" --stream --output output.pcm +""" + +import argparse +import base64 +import os + +import httpx + +DEFAULT_API_BASE = "http://localhost:8091" +DEFAULT_API_KEY = "EMPTY" + + +def encode_audio_to_base64(audio_path: str) -> str: + """Encode a local audio file to base64 data URL.""" + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + ext = audio_path.lower().rsplit(".", 1)[-1] + mime_map = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"} + mime_type = mime_map.get(ext, "audio/wav") + with open(audio_path, "rb") as f: + audio_b64 = base64.b64encode(f.read()).decode("utf-8") + return f"data:{mime_type};base64,{audio_b64}" + + +def run_tts(args) -> None: + """Generate speech via /v1/audio/speech API.""" + payload = { + "model": args.model, + "input": args.text, + "voice": "default", + "response_format": args.response_format, + } + + # Voice cloning parameters. + if args.ref_audio: + if args.ref_audio.startswith(("http://", "https://")): + payload["ref_audio"] = args.ref_audio + else: + payload["ref_audio"] = encode_audio_to_base64(args.ref_audio) + if args.ref_text: + payload["ref_text"] = args.ref_text + + if args.stream: + payload["stream"] = True + payload["response_format"] = "pcm" + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + if args.ref_audio: + print(f"Voice cloning: ref_audio={args.ref_audio}, ref_text={args.ref_text}") + print("Generating audio...") + + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + } + + if args.stream: + output_path = args.output or "output.pcm" + with httpx.Client(timeout=300.0) as client: + with client.stream("POST", api_url, json=payload, headers=headers) as resp: + if resp.status_code != 200: + print(f"Error: {resp.status_code}") + print(resp.read().decode()) + return + total_bytes = 0 + with open(output_path, "wb") as f: + for chunk in resp.iter_bytes(): + f.write(chunk) + total_bytes += len(chunk) + print(f"Streamed {total_bytes} bytes to: {output_path}") + else: + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + + try: + text = response.content.decode("utf-8") + if text.startswith('{"error"'): + print(f"Error: {text}") + return + except UnicodeDecodeError: + pass + + output_path = args.output or "output.wav" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="Fish Speech S2 Pro TTS client") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default="fishaudio/s2-pro", help="Model name") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning (path or URL)") + parser.add_argument("--ref-text", default=None, help="Transcript of reference audio") + parser.add_argument("--stream", action="store_true", help="Enable streaming (PCM output)") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "mp3", "flac", "pcm", "aac", "opus"], + help="Audio format (default: wav)", + ) + parser.add_argument("--output", "-o", default=None, help="Output file path") + args = parser.parse_args() + run_tts(args) + + +if __name__ == "__main__": + main() diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 331adda8e..3b5d67542 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -29,8 +29,8 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) -# TTS Configuration (currently supports Qwen3-TTS) -_TTS_MODEL_STAGES: set[str] = {"qwen3_tts"} +# TTS Configuration +_TTS_MODEL_STAGES: set[str] = {"qwen3_tts", "fish_speech_slow_ar"} _TTS_LANGUAGES: set[str] = { "Auto", "Chinese", @@ -140,6 +140,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Find and cache the TTS stage (if any) during initialization self._tts_stage = self._find_tts_stage() self._is_tts = self._tts_stage is not None + self._is_fish_speech = ( + self._tts_stage is not None and getattr(self._tts_stage, "model_stage", None) == "fish_speech_slow_ar" + ) + self._fish_speech_tokenizer = None # Cache TTS configuration values (computed once, reused per request) self._max_instructions_length = self._compute_max_instructions_length() @@ -780,6 +784,79 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return params + def _build_fish_speech_prompt( + self, + request: OpenAICreateSpeechRequest, + ref_audio_data: tuple[list[float], int] | None = None, + ) -> dict[str, Any]: + """Build prompt for Fish Speech S2 Pro. + + Without voice cloning: + <|im_start|>user\\n<|speaker:0|>{text}<|im_end|>\\n<|im_start|>assistant\\n<|voice|> + + With voice cloning (ref_audio + ref_text): + <|im_start|>system\\n<|speaker:0|>{ref_text}<|audio_start|>{semantic_tokens}<|audio_end|><|im_end|> + <|im_start|>user\\n<|speaker:0|>{text}<|im_end|>\\n<|im_start|>assistant\\n<|voice|> + """ + from transformers import AutoTokenizer + + if self._fish_speech_tokenizer is None: + model_name = self.engine_client.model_config.model + self._fish_speech_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + + tokenizer = self._fish_speech_tokenizer + model_name = self.engine_client.model_config.model + + if ref_audio_data is not None and request.ref_text: + # Voice cloning: encode reference audio and build system message. + from vllm_omni.model_executor.models.fish_speech.dac_encoder import ( + encode_reference_audio, + ) + + wav_samples, sr = ref_audio_data + semantic_token_ids = encode_reference_audio(model_name, wav_samples, sr) + + # Build system message with ref text + audio tokens. + audio_start_id = tokenizer.encode("<|audio_start|>", add_special_tokens=False) + audio_end_id = tokenizer.encode("<|audio_end|>", add_special_tokens=False) + + # System content: <|speaker:0|>{ref_text}<|audio_start|>{codes}<|audio_end|> + prefix_text = f"<|speaker:0|>{request.ref_text}" + prefix_ids = tokenizer.encode(prefix_text, add_special_tokens=False) + system_content_ids = prefix_ids + audio_start_id + semantic_token_ids + audio_end_id + + # Manually build system turn: <|im_start|>system\n{content}<|im_end|>\n + im_start = tokenizer.encode("<|im_start|>", add_special_tokens=False) + im_end = tokenizer.encode("<|im_end|>", add_special_tokens=False) + system_tag = tokenizer.encode("system\n", add_special_tokens=False) + newline = tokenizer.encode("\n", add_special_tokens=False) + system_ids = im_start + system_tag + system_content_ids + im_end + newline + + # User turn via chat template. + user_text = f"<|speaker:0|>{request.input}" + user_messages = [{"role": "user", "content": user_text}] + user_ids = tokenizer.apply_chat_template(user_messages, tokenize=True, add_generation_prompt=True) + prompt_ids = system_ids + user_ids + else: + # No voice cloning: simple user message. + user_text = f"<|speaker:0|>{request.input}" + messages = [{"role": "user", "content": user_text}] + prompt_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True) + + # Append <|voice|> token to signal voice generation. + voice_token_id = tokenizer.encode("<|voice|>", add_special_tokens=False) + prompt_ids = prompt_ids + voice_token_id + + additional_information: dict[str, Any] = { + "text": [request.input], + "max_new_tokens": [request.max_new_tokens or 4096], + } + + return { + "prompt_token_ids": prompt_ids, + "additional_information": additional_information, + } + async def _prepare_speech_generation( self, request: OpenAICreateSpeechRequest, @@ -787,7 +864,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if self.engine_client.errored: raise self.engine_client.dead_error - if self._is_tts: + if self._is_fish_speech: + if not request.input or not request.input.strip(): + raise ValueError("Input text cannot be empty") + ref_audio_data = None + if request.ref_audio is not None: + if not request.ref_text or not request.ref_text.strip(): + raise ValueError("Voice cloning requires 'ref_text' (transcript of the reference audio)") + wav_list, sr = await self._resolve_ref_audio(request.ref_audio) + ref_audio_data = (wav_list, sr) + prompt = self._build_fish_speech_prompt(request, ref_audio_data=ref_audio_data) + tts_params = {} + elif self._is_tts: validation_error = self._validate_tts_request(request) if validation_error: raise ValueError(validation_error) @@ -804,14 +892,29 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): prompt = {"prompt": request.input} request_id = f"speech-{random_uuid()}" + model_type = ( + "fish_speech" + if self._is_fish_speech + else tts_params.get("task_type", ["unknown"])[0] + if self._is_tts + else "generic" + ) logger.info( - "TTS speech request %s: text=%r, task_type=%s", + "TTS speech request %s: text=%r, model=%s", request_id, request.input[:50] + "..." if len(request.input) > 50 else request.input, - tts_params.get("task_type", ["unknown"])[0], + model_type, ) sampling_params_list = self.engine_client.default_sampling_params_list + + # Override Stage-0 max_tokens if caller specified max_new_tokens (Fish Speech). + if self._is_fish_speech and request.max_new_tokens is not None and sampling_params_list: + import copy + + sampling_params_list = copy.deepcopy(sampling_params_list) + sampling_params_list[0].max_tokens = request.max_new_tokens + generator = self.engine_client.generate( prompt=prompt, request_id=request_id, diff --git a/vllm_omni/model_executor/models/fish_speech/__init__.py b/vllm_omni/model_executor/models/fish_speech/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/models/fish_speech/configuration_fish_speech.py b/vllm_omni/model_executor/models/fish_speech/configuration_fish_speech.py new file mode 100644 index 000000000..9ed4ed013 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/configuration_fish_speech.py @@ -0,0 +1,154 @@ +"""Configuration classes for Fish Speech S2 Pro (fish_qwen3_omni). + +The HuggingFace config uses field names from the original Fish Speech codebase +(dim, n_head, n_layer, etc.). This module re-exports them under the standard +Transformers/Qwen3 attribute names so that vLLM's ``Qwen3Model`` can consume +the Slow AR config transparently. +""" + +from __future__ import annotations + +from transformers import PretrainedConfig + + +class FishSpeechSlowARConfig(PretrainedConfig): + """Slow AR (text_model) config -- Qwen3-based transformer. + + Maps Fish Speech field names to Qwen3-compatible attribute names so + ``vllm.model_executor.models.qwen3.Qwen3Model`` works out of the box. + """ + + model_type = "fish_qwen3" + + def __init__( + self, + vocab_size: int = 155776, + dim: int = 2560, + n_head: int = 32, + n_local_heads: int = 8, + head_dim: int = 128, + n_layer: int = 36, + intermediate_size: int = 9728, + attention_qk_norm: bool = True, + rope_base: float = 1_000_000.0, + max_seq_len: int = 32768, + tie_word_embeddings: bool = True, + codebook_size: int = 4096, + num_codebooks: int = 10, + semantic_begin_id: int = 0, + semantic_end_id: int = 0, + scale_codebook_embeddings: bool = False, + rms_norm_eps: float = 1e-6, + **kwargs, + ): + # Fish Speech field names → standard Transformers/Qwen3 names. + self.hidden_size = dim + self.num_attention_heads = n_head + self.num_key_value_heads = n_local_heads + self.head_dim = head_dim + self.num_hidden_layers = n_layer + self.intermediate_size = intermediate_size + self.max_position_embeddings = max_seq_len + self.rope_theta = rope_base + self.rms_norm_eps = rms_norm_eps + self.hidden_act = "silu" + self.attention_bias = False + self.attention_qk_norm = attention_qk_norm + + # Fish Speech codec / codebook fields. + self.codebook_size = codebook_size + self.num_codebooks = num_codebooks + self.semantic_begin_id = semantic_begin_id + self.semantic_end_id = semantic_end_id + self.scale_codebook_embeddings = scale_codebook_embeddings + + super().__init__( + vocab_size=vocab_size, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +class FishSpeechFastARConfig(PretrainedConfig): + """Fast AR (audio_decoder) config -- 4-layer residual codebook predictor.""" + + model_type = "fish_qwen3_audio_decoder" + + def __init__( + self, + vocab_size: int = 4096, + num_codebooks: int = 10, + dim: int = 2560, + n_head: int = 32, + n_local_heads: int = 8, + head_dim: int = 128, + n_layer: int = 4, + intermediate_size: int = 9728, + max_seq_len: int = 11, + text_dim: int = 2560, + audio_hidden_dim: int = 5120, + attention_qk_norm: bool = False, + rope_base: float = 1_000_000.0, + rms_norm_eps: float = 1e-6, + **kwargs, + ): + self.hidden_size = dim + self.num_attention_heads = n_head + self.num_key_value_heads = n_local_heads + self.head_dim = head_dim + self.num_hidden_layers = n_layer + self.intermediate_size = intermediate_size + self.max_position_embeddings = max_seq_len + self.rope_theta = rope_base + self.rms_norm_eps = rms_norm_eps + self.hidden_act = "silu" + self.attention_bias = False + self.attention_qk_norm = attention_qk_norm + self.num_codebooks = num_codebooks + self.text_dim = text_dim + self.audio_hidden_dim = audio_hidden_dim + + super().__init__(vocab_size=vocab_size, **kwargs) + + +class FishSpeechConfig(PretrainedConfig): + """Top-level config for Fish Speech S2 Pro (fish_qwen3_omni). + + Wraps ``text_config`` (Slow AR) and ``audio_decoder_config`` (Fast AR). + """ + + model_type = "fish_qwen3_omni" + sub_configs = { + "text_config": FishSpeechSlowARConfig, + "audio_decoder_config": FishSpeechFastARConfig, + } + + def __init__( + self, + text_config: dict | FishSpeechSlowARConfig | None = None, + audio_decoder_config: dict | FishSpeechFastARConfig | None = None, + semantic_start_token_id: int = 151678, + semantic_end_token_id: int = 155773, + audio_pad_token_id: int = 151677, + **kwargs, + ): + if isinstance(text_config, dict): + text_config = FishSpeechSlowARConfig(**text_config) + if isinstance(audio_decoder_config, dict): + audio_decoder_config = FishSpeechFastARConfig(**audio_decoder_config) + + self.text_config = text_config or FishSpeechSlowARConfig() + self.audio_decoder_config = audio_decoder_config or FishSpeechFastARConfig() + + self.semantic_start_token_id = semantic_start_token_id + self.semantic_end_token_id = semantic_end_token_id + self.audio_pad_token_id = audio_pad_token_id + + # Propagate semantic IDs into text_config for convenience. + self.text_config.semantic_begin_id = semantic_start_token_id + self.text_config.semantic_end_id = semantic_end_token_id + + super().__init__(**kwargs) + + def get_text_config(self, **kwargs) -> FishSpeechSlowARConfig: + return self.text_config diff --git a/vllm_omni/model_executor/models/fish_speech/dac_encoder.py b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py new file mode 100644 index 000000000..084a863ad --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py @@ -0,0 +1,111 @@ +"""DAC codec encoder for Fish Speech S2 Pro voice cloning. + +Encodes reference audio into VQ codes for use as prompt conditioning. +Runs on CPU in the API server process -- loaded lazily on first use. +""" + +from __future__ import annotations + +import os + +import numpy as np +import torch +import torch.nn as nn +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.fish_speech.dac_utils import ( + DAC_SAMPLE_RATE, + build_dac_codec, +) + +logger = init_logger(__name__) + +_codec_cache: dict[str, nn.Module] = {} + + +def _load_dac_codec(model_path: str) -> nn.Module: + """Load the DAC codec model from codec.pth (cached per model_path).""" + if model_path in _codec_cache: + return _codec_cache[model_path] + + codec_path = os.path.join(model_path, "codec.pth") + if not os.path.exists(codec_path): + from transformers.utils.hub import cached_file + + cached = cached_file(model_path, "codec.pth") + if cached is not None: + codec_path = cached + + if not os.path.exists(codec_path): + raise FileNotFoundError( + f"codec.pth not found for {model_path}. Required for voice cloning with Fish Speech S2 Pro." + ) + + codec = build_dac_codec() + + state_dict = torch.load(codec_path, map_location="cpu", weights_only=True) + if "generator" in state_dict: + state_dict = state_dict["generator"] + codec.load_state_dict(state_dict, strict=False) + codec.eval() + + _codec_cache[model_path] = codec + logger.info("Loaded DAC codec encoder from %s (CPU)", codec_path) + return codec + + +def _resample(wav: np.ndarray, sr: int, target_sr: int) -> np.ndarray: + """Resample audio using torchaudio's polyphase resampling.""" + if sr == target_sr: + return wav + import torchaudio + + wav_t = torch.from_numpy(wav).unsqueeze(0).float() + wav_t = torchaudio.functional.resample(wav_t, sr, target_sr) + return wav_t.squeeze(0).numpy() + + +@torch.no_grad() +def encode_reference_audio( + model_path: str, + wav_samples: list[float] | np.ndarray, + sample_rate: int, +) -> list[int]: + """Encode reference audio into semantic token IDs for prompt conditioning. + + Args: + model_path: HuggingFace model path (for locating codec.pth). + wav_samples: Audio waveform samples (mono, float). + sample_rate: Sample rate of the input audio. + + Returns: + List of semantic token IDs (151678 + code_value for each frame). + """ + codec = _load_dac_codec(model_path) + + wav = np.asarray(wav_samples, dtype=np.float32) + if wav.ndim > 1: + wav = np.mean(wav, axis=-1) + + # Resample to DAC sample rate (44100). + wav = _resample(wav, sample_rate, DAC_SAMPLE_RATE) + + # Encode: [1, 1, T] -> codes [1, num_codebooks, num_frames] + wav_tensor = torch.from_numpy(wav).unsqueeze(0).unsqueeze(0).float() + feature_lengths = torch.tensor([wav_tensor.shape[-1]]) + codes, feature_lengths_out = codec.encode(wav_tensor, feature_lengths) + + # Extract semantic codebook (index 0) - shape [num_frames]. + semantic_codes = codes[0, 0, :].tolist() + + # Convert to semantic token IDs: <|semantic:{i}|> = 151678 + i + SEMANTIC_TOKEN_OFFSET = 151678 + semantic_token_ids = [SEMANTIC_TOKEN_OFFSET + int(c) for c in semantic_codes] + + logger.info( + "Encoded reference audio: %d samples @ %dHz -> %d semantic tokens", + len(wav_samples), + sample_rate, + len(semantic_token_ids), + ) + return semantic_token_ids diff --git a/vllm_omni/model_executor/models/fish_speech/dac_utils.py b/vllm_omni/model_executor/models/fish_speech/dac_utils.py new file mode 100644 index 000000000..f1f2590a3 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/dac_utils.py @@ -0,0 +1,92 @@ +"""Shared DAC codec construction for Fish Speech S2 Pro. + +Used by both the encoder (voice cloning, CPU) and the decoder (Stage 1, GPU). +""" + +from __future__ import annotations + +import torch.nn as nn + +# Default DAC codec config matching s2-pro's modded_dac_vq.yaml. +DAC_SAMPLE_RATE = 44100 +DAC_HOP_LENGTH = 2048 # 512 (decoder upsample) * 4 (quantizer upsample) +DAC_NUM_CODEBOOKS = 10 # 1 semantic + 9 residual + + +def build_dac_codec() -> nn.Module: + """Construct a DAC codec model (uninitialized weights). + + Returns the model on CPU in eval mode -- caller is responsible for + loading weights and moving to the target device. + """ + from fish_speech.models.dac.modded_dac import ( + DAC, + ModelArgs, + WindowLimitedTransformer, + ) + from fish_speech.models.dac.rvq import DownsampleResidualVectorQuantize + + base_transformer_kwargs = dict( + block_size=16384, + n_local_heads=-1, + head_dim=64, + rope_base=10000, + norm_eps=1e-5, + dropout_rate=0.0, + attn_dropout_rate=0.0, + channels_first=True, + ) + + def _make_transformer_config(**kw): + return ModelArgs(**{**base_transformer_kwargs, **kw}) + + quantizer_transformer_config = ModelArgs( + block_size=4096, + n_layer=8, + n_head=16, + dim=1024, + intermediate_size=3072, + n_local_heads=-1, + head_dim=64, + rope_base=10000, + norm_eps=1e-5, + dropout_rate=0.0, + attn_dropout_rate=0.0, + channels_first=True, + ) + post_module = WindowLimitedTransformer( + causal=True, + window_size=128, + input_dim=1024, + config=quantizer_transformer_config, + ) + pre_module = WindowLimitedTransformer( + causal=True, + window_size=128, + input_dim=1024, + config=quantizer_transformer_config, + ) + quantizer = DownsampleResidualVectorQuantize( + input_dim=1024, + n_codebooks=9, + codebook_size=1024, + codebook_dim=8, + quantizer_dropout=0.0, + downsample_factor=[2, 2], + post_module=post_module, + pre_module=pre_module, + semantic_codebook_size=4096, + ) + codec = DAC( + encoder_dim=64, + encoder_rates=[2, 4, 8, 8], + decoder_dim=1536, + decoder_rates=[8, 8, 4, 2], + quantizer=quantizer, + sample_rate=DAC_SAMPLE_RATE, + causal=True, + encoder_transformer_layers=[0, 0, 0, 4], + decoder_transformer_layers=[4, 0, 0, 0], + transformer_general_config=_make_transformer_config, + ) + return codec diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py new file mode 100644 index 000000000..a889f66fd --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py @@ -0,0 +1,272 @@ +"""Fish Speech S2 Pro -- DAC Decoder (Stage 1). + +Loads the DAC codec from ``codec.pth`` and decodes codebook indices +[num_codebooks, T] → audio waveform at 44.1 kHz. + +Analogous to ``Qwen3TTSCode2Wav`` in qwen3_tts. + +Requires the ``fish-speech`` package for the DAC model architecture. +Install with: ``pip install fish-speech`` +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.fish_speech.dac_utils import ( + DAC_HOP_LENGTH, + DAC_NUM_CODEBOOKS, + DAC_SAMPLE_RATE, + build_dac_codec, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + + +class FishSpeechDACDecoder(nn.Module): + """Stage-1 DAC decoder for Fish Speech S2 Pro (GenerationModelRunner). + + Consumes frame-aligned codec tokens from input_ids and decodes waveform + via the DAC codec decoder. + """ + + input_modalities = "audio" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + + self._codec: nn.Module | None = None + self._num_codebooks: int = DAC_NUM_CODEBOOKS + self._output_sample_rate: int = DAC_SAMPLE_RATE + self._hop_length: int = DAC_HOP_LENGTH + self._logged_codec_stats = False + + def _ensure_codec_loaded(self) -> None: + if self._codec is not None: + return + + codec_path = os.path.join(self.model_path, "codec.pth") + if not os.path.exists(codec_path): + # Try HuggingFace cache. + try: + from transformers.utils.hub import cached_file + + cached = cached_file(self.model_path, "codec.pth") + if cached is not None: + codec_path = cached + except Exception: + pass + + if not os.path.exists(codec_path): + raise FileNotFoundError( + f"codec.pth not found at {codec_path}. Make sure the Fish Speech S2 Pro model includes codec.pth." + ) + + codec = build_dac_codec() + + # Load weights. + state_dict = torch.load(codec_path, map_location="cpu", weights_only=True) + # Some checkpoints wrap under "generator" key. + if "generator" in state_dict: + state_dict = state_dict["generator"] + codec.load_state_dict(state_dict, strict=False) + + device = self.vllm_config.device_config.device + codec = codec.to(device=device, dtype=torch.float32) + codec.eval() + self._codec = codec + + logger.info( + "Fish Speech DAC codec loaded from %s (device=%s, sample_rate=%d)", + codec_path, + device, + self._output_sample_rate, + ) + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + if input_ids.numel() == 0: + return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) + return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None: + return None + + def _split_request_ids( + self, + ids: torch.Tensor, + seq_token_counts: list[int] | None = None, + ) -> list[torch.Tensor]: + if seq_token_counts is not None and len(seq_token_counts) > 1: + boundaries = [0] + for count in seq_token_counts: + boundaries.append(boundaries[-1] + count) + n = ids.numel() + return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))] + if is_forward_context_available(): + slices = get_forward_context().ubatch_slices + if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices): + boundaries = [0] + for s in slices: + boundaries.append(boundaries[-1] + s) + return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)] + return [ids] + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + """Decode codec codes into audio waveform. + + input_ids layout per request: flat codes [num_codebooks * num_frames]. + Codes are codebook-major: [cb0_f0, cb0_f1, ..., cb0_fN, cb1_f0, ...]. + """ + self._ensure_codec_loaded() + assert self._codec is not None + + q = self._num_codebooks + sr_val = self._output_sample_rate + sr_tensor = torch.tensor(sr_val, dtype=torch.int32) + empty = torch.zeros((0,), dtype=torch.float32) + + if input_ids is None or input_ids.numel() == 0: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]}, + ) + + ids = input_ids.reshape(-1).to(dtype=torch.long) + request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts")) + + parsed: list[tuple[int, int]] = [] + valid_codes_qf: list[torch.Tensor] = [] + valid_indices: list[int] = [] + left_context_size = [0] * len(request_ids_list) + if runtime_additional_information is not None: + for i, info in enumerate(runtime_additional_information): + if i >= len(left_context_size): + break + if "left_context_size" in info: + left_context_size[i] = info["left_context_size"] + + for i, req_ids in enumerate(request_ids_list): + if req_ids.numel() < 1: + parsed.append((0, 0)) + continue + ctx_frames = left_context_size[i] + flat = req_ids + n = flat.numel() + if n == 0 or n % q != 0: + if n > 0: + logger.warning( + "DAC decoder input_ids length %d not divisible by num_codebooks %d; returning empty audio.", + n, + q, + ) + parsed.append((0, 0)) + continue + frames = n // q + codes_qf = flat.reshape(q, frames) + parsed.append((ctx_frames, frames)) + valid_codes_qf.append(codes_qf) + valid_indices.append(i) + + num_req = len(request_ids_list) + if not valid_codes_qf: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": [empty] * num_req, + "sr": [sr_tensor] * num_req, + }, + ) + + if not self._logged_codec_stats: + self._logged_codec_stats = True + try: + c = valid_codes_qf[0] + logger.info( + "DAC decoder: frames=%d q=%d uniq=%d range=[%d,%d] batch=%d", + c.shape[1], + q, + int(torch.unique(c).numel()), + int(c.min().item()), + int(c.max().item()), + len(valid_codes_qf), + ) + except Exception: + pass + + # Decode each request individually. + wav_tensors: list[torch.Tensor] = [] + for codes_qf in valid_codes_qf: + codes_bqf = codes_qf.unsqueeze(0) # [1, num_codebooks, num_frames] + num_frames = codes_qf.shape[1] + feature_lengths = torch.tensor([num_frames], device=codes_bqf.device) + with torch.cuda.amp.autocast(enabled=False): + wav, audio_lengths = self._codec.decode(codes_bqf, feature_lengths) + # wav shape: [1, 1, wav_len] + wav_tensors.append(wav.squeeze(0).squeeze(0)) # [wav_len] + + audios: list[torch.Tensor] = [empty] * num_req + srs = [sr_tensor] * num_req + + for j, idx in enumerate(valid_indices): + ctx_frames, actual_frames = parsed[idx] + wav = wav_tensors[j] + # Trim context frames (left overlap for streaming). + if ctx_frames > 0: + cut = ctx_frames * self._hop_length + if cut < wav.shape[0]: + wav = wav[cut:] + else: + logger.warning( + "Context trim %d >= decoded length %d; returning empty audio.", + cut, + wav.shape[0], + ) + continue + if wav.shape[0] > 0: + audios[idx] = wav.to(dtype=torch.float32).reshape(-1) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audios, "sr": srs}, + ) + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + if not (isinstance(model_outputs, tuple) and len(model_outputs) == 2): + raise TypeError(f"FishSpeechDACDecoder expected (audio_tensor, sr), got {type(model_outputs)}") + audio_tensor, sr = model_outputs + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audio_tensor, "sr": sr}, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # DAC codec weights are loaded lazily from codec.pth, not from the main checkpoint. + return set() diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py new file mode 100644 index 000000000..f73914c28 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py @@ -0,0 +1,470 @@ +"""Fast AR (audio_decoder) for Fish Speech S2 Pro. + +4-layer transformer that predicts residual codebook codes (1..num_codebooks-1) +autoregressively after each Slow AR step. Analogous to Qwen3 TTS's +CodePredictor but with its own embedding table, RoPE, and output head. + +Uses re-prefill (no KV cache): each AR step forwards the full growing +sequence through the 4-layer transformer. This trades ~O(T^2) attention +FLOPs (negligible for T=10, 4 layers) for zero KV cache management. + +Optimisations: + - torch.compile on model forward (kernel fusion) + - Pre-allocated projection / embedding buffer + - Pre-allocated position_ids + - Inline sampling +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.config import VllmConfig +from vllm.config.vllm import set_current_vllm_config +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.utils import is_pp_missing_parameter + +from .configuration_fish_speech import FishSpeechFastARConfig, FishSpeechSlowARConfig + +logger = init_logger(__name__) + + +# =================================================================== +# Standalone Fast AR Layers (no vLLM paged attention) +# =================================================================== + + +class _FastARAttention(nn.Module): + """Multi-head attention using F.scaled_dot_product_attention (SDPA). + + Supports fused QKV, RoPE, optional q/k normalization, and native GQA. + Input: [B, seq_len, hidden_size]. + """ + + def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self._use_gqa = self.num_kv_heads != self.num_heads + + self.qkv_proj = QKVParallelLinear( + hidden_size=self.hidden_size, + head_size=self.head_dim, + total_num_heads=self.num_heads, + total_num_kv_heads=self.num_kv_heads, + bias=False, + prefix=f"{prefix}.qkv_proj", + disable_tp=True, + ) + self.o_proj = RowParallelLinear( + input_size=self.num_heads * self.head_dim, + output_size=self.hidden_size, + bias=False, + prefix=f"{prefix}.o_proj", + disable_tp=True, + ) + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + is_neox_style=False, # Fish Speech uses interleaved (GPT-J) RoPE + rope_parameters={"rope_theta": config.rope_theta, "rope_type": "default"}, + ) + # Fish Speech Fast AR typically has attention_qk_norm=false. + if config.attention_qk_norm: + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + else: + self.q_norm = None + self.k_norm = None + + def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + bsz, seq_len, _ = hidden_states.shape + + qkv, _ = self.qkv_proj(hidden_states.reshape(bsz * seq_len, -1)) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + if self.q_norm is not None: + q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape) + k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape) + + q, k = self.rotary_emb(position_ids, q, k) + + q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + + attn_out = F.scaled_dot_product_attention( + q, + k, + v, + scale=self.scaling, + is_causal=True, + enable_gqa=self._use_gqa, + ) + attn_out = attn_out.transpose(1, 2).reshape(bsz * seq_len, -1) + output, _ = self.o_proj(attn_out) + return output.view(bsz, seq_len, -1) + + +class _FastARMLP(nn.Module): + """SiLU-gated MLP, matching Qwen3/LLaMA MLP structure.""" + + def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[config.intermediate_size] * 2, + bias=False, + prefix=f"{prefix}.gate_up_proj", + disable_tp=True, + ) + self.down_proj = RowParallelLinear( + input_size=config.intermediate_size, + output_size=config.hidden_size, + bias=False, + prefix=f"{prefix}.down_proj", + disable_tp=True, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + gate, up = gate_up.chunk(2, dim=-1) + x = F.silu(gate) * up + x, _ = self.down_proj(x) + return x + + +class _FastARDecoderLayer(nn.Module): + """Transformer decoder layer for Fast AR (SDPA, no KV cache).""" + + def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.self_attn = _FastARAttention(config, prefix=f"{prefix}.self_attn") + self.mlp = _FastARMLP(config, prefix=f"{prefix}.mlp") + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_ids) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +# =================================================================== +# Fast AR Transformer Model +# =================================================================== + + +class FishSpeechFastARModel(nn.Module): + """4-layer transformer for residual codebook prediction (re-prefill).""" + + def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.config = config + self.layers = nn.ModuleList( + [_FastARDecoderLayer(config, prefix=f"{prefix}.layers.{i}") for i in range(config.num_hidden_layers)] + ) + # NOTE: final norm is handled by FishSpeechFastAR.fast_norm (one norm weight + # in checkpoint: audio_decoder.norm.weight → fast_ar.fast_norm.weight). + + def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + hidden_states = inputs_embeds + for layer in self.layers: + hidden_states = layer(hidden_states, position_ids) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + if mapped.endswith(".bias") and mapped not in params_dict: + continue + if is_pp_missing_parameter(mapped, self): + continue + if mapped.endswith("scale"): + mapped = maybe_remap_kv_scale_name(mapped, params_dict) + if mapped is None: + continue + param = params_dict.get(mapped) + if param is None: + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if weight_loader == default_weight_loader: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(mapped) + break + else: + mapped = maybe_remap_kv_scale_name(name, params_dict) + if mapped is None: + continue + if name.endswith(".bias") and mapped not in params_dict: + continue + if is_pp_missing_parameter(mapped, self): + continue + param = params_dict.get(mapped) + if param is None: + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(mapped) + return loaded_params + + +# =================================================================== +# Fast AR Wrapper (optimised re-prefill + torch.compile) +# =================================================================== + + +class FishSpeechFastAR(nn.Module): + """vLLM-native Fast AR for Fish Speech (residual codebooks). + + Re-prefill approach: each AR step forwards the full growing sequence + through the 4-layer transformer. No KV cache needed. + + Optimisations over baseline: + 1. torch.compile on model forward -- kernel fusion. + 2. Pre-allocated embedding buffer [B, max_seq, H]. + 3. Pre-allocated position_ids. + 4. Inline sampling. + """ + + def __init__( + self, + *, + vllm_config: VllmConfig, + config: FishSpeechFastARConfig, + slow_ar_config: FishSpeechSlowARConfig, + prefix: str = "fast_ar", + ) -> None: + super().__init__() + self._vllm_config = vllm_config + self.config = config + self.slow_ar_config = slow_ar_config + + self.model = FishSpeechFastARModel(config, prefix=f"{prefix}.model") + + # Codebook embeddings for Fast AR (separate from Slow AR's codebook_embeddings). + self.fast_embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + + # Output projection for residual codes. + self.fast_output = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Layer norm for Fast AR output. + self.fast_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # Projection from Slow AR hidden dim → Fast AR hidden dim. + # s2-pro has same dim for both (2560) and no fast_project_in in checkpoint, + # so use identity when dims match. + if slow_ar_config.hidden_size != config.hidden_size: + self.fast_project_in = nn.Linear(slow_ar_config.hidden_size, config.hidden_size, bias=True) + else: + self.fast_project_in = nn.Identity() + + self._num_codebooks = config.num_codebooks + self._fast_dim = config.hidden_size + + # Pre-allocated buffers (lazily initialised on first forward). + self._embed_buf: torch.Tensor | None = None + self._pos_ids: torch.Tensor | None = None + self._compiled_model_fwd: object | None = None + + def _ensure_buffers(self, bsz: int, device: torch.device, dtype: torch.dtype) -> None: + max_seq = self._num_codebooks + 1 # hidden_state + num_codebooks codes + if ( + self._embed_buf is not None + and self._embed_buf.shape[0] >= bsz + and self._embed_buf.device == device + and self._embed_buf.dtype == dtype + ): + return + self._embed_buf = torch.zeros(bsz, max_seq, self._fast_dim, dtype=dtype, device=device) + self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device) + + def _setup_compile(self) -> None: + if self._compiled_model_fwd is not None: + return + # TODO: Enable torch.compile for performance. Eager for now to avoid + # potential graph-break issues during initial bring-up. + self._compiled_model_fwd = self.model.forward + + @torch.inference_mode() + def forward( + self, + slow_ar_hidden: torch.Tensor, + semantic_token_id: torch.Tensor, + *, + do_sample: bool = True, + temperature: float = 0.8, + top_k: int = 30, + top_p: float = 0.9, + ) -> torch.Tensor: + """Predict residual codebook codes 0..num_codebooks-1 autoregressively. + + Args: + slow_ar_hidden: [B, hidden_size] last hidden state from Slow AR. + semantic_token_id: [B] or [B, 1] sampled semantic token IDs (in vocab space). + + Returns: + all_codes: [B, num_codebooks] codes for codebooks 0..num_codebooks-1. + Code 0 is the semantic code (token - semantic_begin_id). + """ + bsz = int(slow_ar_hidden.shape[0]) + num_cb = self._num_codebooks + device = slow_ar_hidden.device + dtype = slow_ar_hidden.dtype + + semantic_begin = self.slow_ar_config.semantic_begin_id + semantic_end = self.slow_ar_config.semantic_end_id + codebook_size = semantic_end - semantic_begin + 1 # 4096 + # Convert vocab-space semantic token to codebook index. + # Clamp to valid range: im_end or other non-semantic tokens map to 0 (pad). + semantic_code = (semantic_token_id.reshape(bsz) - semantic_begin).clamp(min=0, max=codebook_size - 1) + + all_codes = torch.empty(bsz, num_cb, dtype=torch.long, device=device) + all_codes[:, 0] = semantic_code + + self._ensure_buffers(bsz, device, dtype) + self._setup_compile() + + embed_buf = self._embed_buf + pos_ids = self._pos_ids + model_fwd = self._compiled_model_fwd + + # Position 0: projected Slow AR hidden state. + projected = self.fast_project_in(slow_ar_hidden.reshape(bsz, -1)) + embed_buf[:bsz, 0, :] = projected + + # Position 1: embedding of semantic code. + code_embed = self.fast_embeddings(semantic_code.clamp(min=0)) + embed_buf[:bsz, 1, :] = code_embed + + use_sampling = do_sample and temperature > 0 + inv_temperature = 1.0 / max(temperature, 1e-6) if use_sampling else 0.0 + + # Residual codebook size (1024) vs semantic codebook size (4096). + # The fast_output head has codebook_size (4096) outputs, but residual + # codebooks only have 1024 entries. Truncate logits for steps > 0. + residual_codebook_size = 1024 + + for step in range(1, num_cb): + seq_len = step + 1 + step_input = embed_buf[:bsz, :seq_len, :] + step_pos_ids = pos_ids[:seq_len] if bsz == 1 else pos_ids[:seq_len].repeat(bsz) + + hidden_out = model_fwd(step_input, step_pos_ids) + logits = self.fast_output(self.fast_norm(hidden_out[:, -1, :])) + + # Residual codebooks (step >= 1) only have 1024 entries. + if step >= 1: + logits = logits[:, :residual_codebook_size] + + if use_sampling: + scaled = logits * inv_temperature + if top_k > 0: + topk_vals, _ = scaled.topk(min(top_k, scaled.shape[-1]), dim=-1) + scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) + if top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(scaled, descending=True) + cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) + sorted_indices_to_remove = cumulative_probs - F.softmax(sorted_logits, dim=-1) >= top_p + sorted_logits[sorted_indices_to_remove] = float("-inf") + scaled = sorted_logits.scatter(1, sorted_indices, sorted_logits) + probs = F.softmax(scaled, dim=-1) + next_ids = torch.multinomial(probs, num_samples=1) + else: + next_ids = logits.argmax(dim=-1, keepdim=True) + + all_codes[:, step] = next_ids.reshape(bsz) + + if step < num_cb - 1: + new_embed = self.fast_embeddings(next_ids.reshape(bsz)) + embed_buf[:bsz, step + 1, :] = new_embed + + return all_codes + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights -- Fast AR weights have already been remapped by the parent.""" + with set_current_vllm_config(self._vllm_config): + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded: set[str] = set() + + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + handled = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + if mapped not in params_dict: + continue + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if weight_loader == default_weight_loader: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + loaded.add(mapped) + handled = True + break + + if handled: + continue + + if name in params_dict: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded.add(name) + + return loaded diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py new file mode 100644 index 000000000..9c41d87f9 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py @@ -0,0 +1,669 @@ +"""Fish Speech S2 Pro -- Slow AR model (Stage 0). + +Uses vLLM's ``Qwen3Model`` as the transformer backbone. Adds: + - Multi-codebook input embedding (text + summed codebook embeddings at + semantic-token positions). + - Semantic logit masking. + - Nested Fast AR for residual codebook prediction (``talker_mtp``). + - ``preprocess`` / ``postprocess`` hooks for vLLM-omni's AR scheduler. + +Analogous to ``Qwen3TTSTalkerForConditionalGeneration`` in qwen3_tts. +""" + +from __future__ import annotations + +import dataclasses +import math +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from transformers import AutoTokenizer +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +from vllm.logger import init_logger +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .configuration_fish_speech import FishSpeechConfig, FishSpeechFastARConfig, FishSpeechSlowARConfig +from .fish_speech_fast_ar import FishSpeechFastAR + +logger = init_logger(__name__) + + +def _remap_fish_speech_weights( + weights: Iterable[tuple[str, torch.Tensor]], + n_head: int, + n_local_heads: int, + head_dim: int, + fast_n_head: int, + fast_n_local_heads: int, + fast_head_dim: int, +) -> Iterable[tuple[str, torch.Tensor]]: + """Transform Fish Speech HF weight names/values to Qwen3-compatible format. + + Key transformations: + - ``wqkv`` → split into ``q_proj``, ``k_proj``, ``v_proj`` + - ``wo`` → ``o_proj`` + - ``w1`` → ``gate_proj``, ``w3`` → ``up_proj``, ``w2`` → ``down_proj`` + - ``attention_norm`` → ``input_layernorm`` + - ``ffn_norm`` → ``post_attention_layernorm`` + - ``text_model.model.embeddings`` → ``model.embed_tokens`` + - ``audio_decoder.*`` → ``fast_ar.*`` with similar transforms + """ + q_size_text = n_head * head_dim + kv_size_text = n_local_heads * head_dim + q_size_fast = fast_n_head * fast_head_dim + kv_size_fast = fast_n_local_heads * fast_head_dim + + for name, tensor in weights: + # --- Text model (Slow AR) --- + if name.startswith("text_model.model."): + suffix = name[len("text_model.model.") :] + + # Embeddings + if suffix == "embeddings.weight": + yield "model.embed_tokens.weight", tensor + continue + + # Norm + if suffix == "norm.weight": + yield "model.norm.weight", tensor + continue + + # Layer weights + if suffix.startswith("layers."): + # layers.{N}.attention.wqkv.weight → split into q/k/v + if ".attention.wqkv.weight" in suffix: + layer_prefix = suffix.split(".attention.wqkv.weight")[0] + q = tensor[:q_size_text, :] + k = tensor[q_size_text : q_size_text + kv_size_text, :] + v = tensor[q_size_text + kv_size_text :, :] + yield f"model.{layer_prefix}.self_attn.q_proj.weight", q + yield f"model.{layer_prefix}.self_attn.k_proj.weight", k + yield f"model.{layer_prefix}.self_attn.v_proj.weight", v + continue + + new_suffix = suffix + new_suffix = new_suffix.replace(".attention.wo.", ".self_attn.o_proj.") + new_suffix = new_suffix.replace(".attention.q_norm.", ".self_attn.q_norm.") + new_suffix = new_suffix.replace(".attention.k_norm.", ".self_attn.k_norm.") + new_suffix = new_suffix.replace(".attention_norm.", ".input_layernorm.") + new_suffix = new_suffix.replace(".feed_forward.w1.", ".mlp.gate_proj.") + new_suffix = new_suffix.replace(".feed_forward.w3.", ".mlp.up_proj.") + new_suffix = new_suffix.replace(".feed_forward.w2.", ".mlp.down_proj.") + new_suffix = new_suffix.replace(".ffn_norm.", ".post_attention_layernorm.") + yield f"model.{new_suffix}", tensor + continue + + # Fallback for any other text_model.model.* weights + yield f"model.{suffix}", tensor + continue + + # --- Audio decoder (Fast AR) --- + if name.startswith("audio_decoder."): + suffix = name[len("audio_decoder.") :] + + # Codebook embeddings (belongs to the main model, not Fast AR). + if suffix == "codebook_embeddings.weight": + yield "codebook_embeddings.weight", tensor + continue + + # Fast AR embeddings, output, norm. + if suffix == "embeddings.weight": + yield "fast_ar.fast_embeddings.weight", tensor + continue + if suffix == "output.weight": + yield "fast_ar.fast_output.weight", tensor + continue + if suffix == "norm.weight": + yield "fast_ar.fast_norm.weight", tensor + continue + + # Fast AR projection in. + if suffix.startswith("fast_project_in."): + yield f"fast_ar.fast_project_in.{suffix[len('fast_project_in.') :]}", tensor + continue + + # Fast AR layer weights. + if suffix.startswith("layers."): + if ".attention.wqkv.weight" in suffix: + layer_prefix = suffix.split(".attention.wqkv.weight")[0] + q = tensor[:q_size_fast, :] + k = tensor[q_size_fast : q_size_fast + kv_size_fast, :] + v = tensor[q_size_fast + kv_size_fast :, :] + yield f"fast_ar.model.{layer_prefix}.self_attn.q_proj.weight", q + yield f"fast_ar.model.{layer_prefix}.self_attn.k_proj.weight", k + yield f"fast_ar.model.{layer_prefix}.self_attn.v_proj.weight", v + continue + + new_suffix = suffix + new_suffix = new_suffix.replace(".attention.wo.", ".self_attn.o_proj.") + new_suffix = new_suffix.replace(".attention.q_norm.", ".self_attn.q_norm.") + new_suffix = new_suffix.replace(".attention.k_norm.", ".self_attn.k_norm.") + new_suffix = new_suffix.replace(".attention_norm.", ".input_layernorm.") + new_suffix = new_suffix.replace(".feed_forward.w1.", ".mlp.gate_proj.") + new_suffix = new_suffix.replace(".feed_forward.w3.", ".mlp.up_proj.") + new_suffix = new_suffix.replace(".feed_forward.w2.", ".mlp.down_proj.") + new_suffix = new_suffix.replace(".ffn_norm.", ".post_attention_layernorm.") + yield f"fast_ar.model.{new_suffix}", tensor + continue + + yield f"fast_ar.{suffix}", tensor + continue + + # Pass through any other weights. + yield name, tensor + + +class FishSpeechSlowARForConditionalGeneration(nn.Module): + """vLLM-AR Slow AR model for Fish Speech S2 Pro. + + Stage 0: text → semantic tokens (+ residual codebook codes via Fast AR). + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + + config: FishSpeechConfig = vllm_config.model_config.hf_config # type: ignore[assignment] + self.config = config + self.text_config: FishSpeechSlowARConfig = config.text_config + self.fast_ar_config: FishSpeechFastARConfig = config.audio_decoder_config + + self._semantic_begin_id = int(config.semantic_start_token_id) + self._semantic_end_id = int(config.semantic_end_token_id) + self._audio_pad_token_id = int(config.audio_pad_token_id) + self._codebook_size = int(self.text_config.codebook_size) + self._num_codebooks = int(self.text_config.num_codebooks) + + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = True + self.mtp_hidden_size = int(self.text_config.hidden_size) + self.talker_mtp_output_key = "audio_codes" + + # Qwen3 transformer backbone. + self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + + # Fish Speech uses interleaved (GPT-J) RoPE, not NeoX style. + # vLLM's Qwen3Attention defaults to NeoX (is_neox_style=True). + # Replace with interleaved RoPE to match training. + self._fix_rope_style() + + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.text_config.vocab_size, + self.text_config.hidden_size, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(self.text_config.vocab_size) + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + + # Multi-codebook embedding table: codebook_size * num_codebooks entries. + self.codebook_embeddings = nn.Embedding( + self._codebook_size * self._num_codebooks, + self.text_config.hidden_size, + ) + + # Fast AR (residual codebook predictor). + predictor_compilation = dataclasses.replace(vllm_config.compilation_config) + predictor_compilation.static_forward_context = {} + self._fast_ar_vllm_config = dataclasses.replace(vllm_config, compilation_config=predictor_compilation) + from vllm.config.vllm import set_current_vllm_config as _set_cfg + + with _set_cfg(self._fast_ar_vllm_config): + self.fast_ar = FishSpeechFastAR( + vllm_config=self._fast_ar_vllm_config, + config=self.fast_ar_config, + slow_ar_config=self.text_config, + prefix="fast_ar", + ) + + # Constant logit mask: allow only semantic tokens + im_end. + vocab = int(self.text_config.vocab_size) + semantic_mask = torch.zeros((vocab,), dtype=torch.bool) + lo = self._semantic_begin_id + hi = min(self._semantic_end_id + 1, vocab) + if hi > lo: + semantic_mask[lo:hi] = True + # Also allow <|im_end|> (token 151645 in Qwen3 tokeniser). + im_end_id = 151645 + if im_end_id < vocab: + semantic_mask[im_end_id] = True + self.register_buffer("_semantic_allowed_mask", semantic_mask, persistent=False) + + # Tokeniser (lazy). + self._tokenizer = None + + def _fix_rope_style(self) -> None: + """Replace NeoX-style RoPE with interleaved (GPT-J) style. + + Fish Speech was trained with interleaved RoPE (complex-number pairs), + but vLLM's Qwen3Attention defaults to NeoX style. We rebuild the + rotary embedding with ``is_neox_style=False`` for each attention layer. + """ + from vllm.model_executor.layers.rotary_embedding import get_rope + + for layer in self.model.layers: + attn = layer.self_attn + # Extract parameters from the existing RoPE to rebuild it. + head_dim = attn.head_dim + max_position = self.text_config.max_position_embeddings + rope_params = getattr(self.text_config, "rope_scaling", None) or {} + rope_params.setdefault("rope_theta", getattr(self.text_config, "rope_theta", 1000000.0)) + attn.rotary_emb = get_rope( + head_size=head_dim, + max_position=max_position, + is_neox_style=False, + rope_parameters=rope_params, + ) + logger.info("Fixed RoPE style to interleaved (GPT-J) for %d layers", len(self.model.layers)) + + # -------------------- vLLM required hooks -------------------- + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor | IntermediateTensors: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: Any = None, + ) -> torch.Tensor | None: + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is None: + return None + + # Mask to semantic tokens + im_end only. + logits = logits.masked_fill(~self._semantic_allowed_mask, float("-inf")) + return logits + + # -------------------- Omni multimodal output plumbing -------------------- + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + + hidden = model_outputs + info_dicts = kwargs.get("model_intermediate_buffer") + if info_dicts is None: + info_dicts = kwargs.get("runtime_additional_information") or [] + + audio_codes_list: list[torch.Tensor] = [] + for info in info_dicts: + if not isinstance(info, dict): + continue + ac = info.get("audio_codes") + if isinstance(ac, torch.Tensor): + audio_codes_list.append(ac) + + if not audio_codes_list: + logger.debug("make_omni_output: no audio_codes found in info_dicts (len=%d)", len(info_dicts)) + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + + audio_codes = torch.cat(audio_codes_list, dim=0) + span_len = int(audio_codes.shape[0]) + hidden = hidden[:span_len] + mm: dict[str, torch.Tensor] = {"audio_codes": audio_codes} + return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm) + + # -------------------- preprocess / postprocess -------------------- + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + additional_information = info_dict.get("additional_information") + if isinstance(additional_information, dict): + merged: dict[str, Any] = {k: v for k, v in info_dict.items() if k != "additional_information"} + for k, v in additional_information.items(): + merged.setdefault(k, v) + info_dict = merged + + span_len = int(input_ids.shape[0]) + if span_len <= 0: + return input_ids, input_embeds if input_embeds is not None else self.embed_input_ids(input_ids), {} + + if span_len > 1: + # --- Prefill --- + prompt_embeds_cpu = info_dict.get("slow_ar_prompt_embeds") + is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2 + dev = input_ids.device + + if is_first_prefill: + prompt_embeds = self._build_prefill_embeds(input_ids, info_dict) + prompt_embeds_cpu = prompt_embeds.detach().to("cpu").contiguous() + + info_update: dict[str, Any] = { + "slow_ar_prompt_embeds": prompt_embeds_cpu, + "prefill_offset": 0, + } + + take = prompt_embeds_cpu[:span_len] + if int(take.shape[0]) < span_len: + pad_n = span_len - int(take.shape[0]) + pad_embed = self.embed_input_ids( + torch.tensor([self._audio_pad_token_id], device=dev, dtype=torch.long) + ).reshape(1, -1) + take = torch.cat([take, pad_embed.detach().cpu().expand(pad_n, -1)], dim=0) + prompt_embeds = take.to(device=dev, dtype=torch.bfloat16) + info_update["prefill_offset"] = span_len + + zeros = torch.zeros( + (prompt_embeds.shape[0], self._num_codebooks), + device=dev, + dtype=torch.long, + ) + info_update["audio_codes"] = zeros + + input_ids_out = input_ids.clone() + input_ids_out[:] = self._audio_pad_token_id + return input_ids_out, prompt_embeds, info_update + + else: + # Subsequent prefill chunk. + offset = int(info_dict.get("prefill_offset", 0) or 0) + s = max(0, min(offset, int(prompt_embeds_cpu.shape[0]))) + e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0]))) + take = prompt_embeds_cpu[s:e] + if int(take.shape[0]) < span_len: + pad_n = span_len - int(take.shape[0]) + pad_embed = self.embed_input_ids( + torch.tensor([self._audio_pad_token_id], device=dev, dtype=torch.long) + ).reshape(1, -1) + take = torch.cat([take, pad_embed.detach().cpu().expand(pad_n, -1)], dim=0) + prompt_embeds = take.to(device=dev, dtype=torch.bfloat16) + + zeros = torch.zeros((prompt_embeds.shape[0], self._num_codebooks), device=dev, dtype=torch.long) + return ( + input_ids.clone().fill_(self._audio_pad_token_id), + prompt_embeds, + { + "prefill_offset": offset + span_len, + "audio_codes": zeros, + }, + ) + + # --- Decode: span_len == 1 --- + dev = input_ids.device + + last_hidden_cpu = info_dict.get("last_slow_ar_hidden") + if not isinstance(last_hidden_cpu, torch.Tensor): + # First decode step after prefill -- just embed the token directly. + logger.warning( + "preprocess decode: last_slow_ar_hidden not found (keys=%s), " + "returning plain embed (mtp_inputs will NOT be set)", + list(info_dict.keys()), + ) + embeds = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)).reshape(1, -1) + return input_ids, embeds.to(dtype=torch.bfloat16), {} + + token_embed = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)).to( + device=dev, dtype=torch.bfloat16 + ) # [1, 1, H] + + # Codebook embeddings are added by talker_mtp (using the CURRENT step's + # FastAR hidden state) instead of here (which would use the PREVIOUS step's + # codes, one step too old). Preprocess always returns plain text embed. + inputs_embeds_out = token_embed.reshape(1, -1) + + info_update = { + "mtp_inputs": ( + last_hidden_cpu.to(device=dev, dtype=torch.bfloat16).reshape(1, -1), + torch.zeros(1, self.text_config.hidden_size, device=dev, dtype=torch.bfloat16), + ), + } + return input_ids, inputs_embeds_out, info_update + + def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: + if hidden_states.numel() == 0: + logger.debug("postprocess: empty hidden_states") + return {} + last = hidden_states[-1, :].detach().to("cpu").contiguous() + logger.debug("postprocess: saved last_slow_ar_hidden shape=%s", tuple(last.shape)) + return {"last_slow_ar_hidden": last} + + # -------------------- prompt construction -------------------- + + def _get_tokenizer(self): + if self._tokenizer is None: + self._tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True) + return self._tokenizer + + def _build_prefill_embeds( + self, + input_ids: torch.Tensor, + info_dict: dict[str, Any], + ) -> torch.Tensor: + """Build prefill embeddings, adding codebook embeddings at semantic positions. + + For text-only prefill (no reference audio), this is just embed_input_ids. + For voice cloning, reference codes are embedded with codebook offsets. + """ + dev = input_ids.device + # Basic text embeddings. + base_embeds = self.embed_input_ids(input_ids.reshape(1, -1).to(torch.long)) # [1, T, H] + + # Check for reference codebook codes (for voice cloning). + ref_codes = info_dict.get("ref_codes") + if not isinstance(ref_codes, torch.Tensor) or ref_codes.numel() == 0: + return base_embeds.squeeze(0).to(dtype=torch.bfloat16) + + # ref_codes: [T_ref, num_codebooks] -- codebook codes for reference audio positions. + ref_codes = ref_codes.to(device=dev, dtype=torch.long) + ref_positions = info_dict.get("ref_positions") + if not isinstance(ref_positions, torch.Tensor): + return base_embeds.squeeze(0).to(dtype=torch.bfloat16) + + ref_positions = ref_positions.to(device=dev, dtype=torch.long).reshape(-1) + seq_len = int(input_ids.shape[0]) + codebook_sum = torch.zeros_like(base_embeds) # [1, T, H] + + for pos_idx in range(int(ref_positions.shape[0])): + pos = int(ref_positions[pos_idx].item()) + if pos < 0 or pos >= seq_len: + continue + for cb_idx in range(min(int(ref_codes.shape[1]), self._num_codebooks)): + code = ref_codes[pos_idx, cb_idx].clamp(min=0) + code_with_offset = code + cb_idx * self._codebook_size + emb = self.codebook_embeddings(code_with_offset.unsqueeze(0)) + codebook_sum[0, pos, :] += emb.squeeze(0).to(dtype=base_embeds.dtype) + + result = base_embeds + codebook_sum + return result.squeeze(0).to(dtype=torch.bfloat16) + + # -------------------- GPU-side MTP fast-path -------------------- + + @torch.inference_mode() + def talker_mtp( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor, + last_talker_hidden: torch.Tensor, + text_step: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """GPU fast-path: run Fast AR to predict residual codebook codes. + + Returns (inputs_embeds, audio_codes). + + The embedding is: text_embed(token) + sum(codebook_embed(code_i + i * codebook_size)) + where codes come from FastAR(last_talker_hidden). + + This matches the reference Fish Speech inference flow: + - At step t, the Slow AR embedding includes codes from FastAR(hidden_{t-1}) + - last_talker_hidden IS hidden_{t-1} (from postprocess of the previous step) + - Preprocess provides the plain text embed; we add codebook embeddings here + """ + bsz = int(input_ids.shape[0]) + dev = input_embeds.device + + input_ids = input_ids.reshape(bsz, 1).to(dtype=torch.long, device=dev) + past_hidden = last_talker_hidden.reshape(bsz, -1).to(dtype=torch.bfloat16, device=dev) + + # Run Fast AR to predict all num_codebooks codes. + audio_codes = self.fast_ar( + slow_ar_hidden=past_hidden, + semantic_token_id=input_ids.reshape(bsz), + do_sample=True, + temperature=0.8, + top_k=30, + top_p=0.9, + ) # [B, num_codebooks] + + # Add codebook embeddings to the input embedding (from preprocess). + # This ensures the Slow AR sees codes from FastAR(hidden_{t-1}). + inputs_embeds_out = input_embeds.reshape(bsz, -1).clone() + + for b in range(bsz): + token_id = int(input_ids[b, 0].item()) + is_semantic = self._semantic_begin_id <= token_id <= self._semantic_end_id + if is_semantic: + codes = audio_codes[b] # [num_codebooks] + codebook_sum = torch.zeros(self.text_config.hidden_size, device=dev, dtype=torch.bfloat16) + for i in range(self._num_codebooks): + code_with_offset = codes[i].clamp(min=0) + i * self._codebook_size + emb = self.codebook_embeddings(code_with_offset.unsqueeze(0)) + codebook_sum += emb.squeeze(0).to(dtype=torch.bfloat16) + + # Normalize by sqrt(num_codebooks + 1) as in the reference model + # (scale_codebook_embeddings=True for fish_qwen3_omni). + inputs_embeds_out[b] = (inputs_embeds_out[b] + codebook_sum) / math.sqrt(self._num_codebooks + 1) + + return inputs_embeds_out, audio_codes.to(dtype=torch.long) + + # -------------------- Prompt length estimation -------------------- + + @staticmethod + def estimate_prompt_len_from_additional_information( + additional_information: dict[str, Any] | None, + **kwargs: Any, + ) -> int: + """Estimate prompt length for placeholder allocation.""" + info = additional_information or {} + text = info.get("text", [""])[0] if isinstance(info.get("text"), list) else info.get("text", "") + # Conservative estimate: tokenize text length + overhead. + return max(2, len(str(text)) // 2 + 64) + + # -------------------- Weight loading -------------------- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights with Fish Speech → Qwen3 format transformation. + + Transforms weight names (wqkv → q/k/v split, w1/w2/w3 → gate/up/down, + etc.) and routes to the correct sub-modules. + """ + n_head = self.text_config.num_attention_heads + n_local_heads = self.text_config.num_key_value_heads + head_dim = self.text_config.head_dim + fast_n_head = self.fast_ar_config.num_attention_heads + fast_n_local_heads = self.fast_ar_config.num_key_value_heads + fast_head_dim = self.fast_ar_config.head_dim + + remapped = _remap_fish_speech_weights( + weights, + n_head, + n_local_heads, + head_dim, + fast_n_head, + fast_n_local_heads, + fast_head_dim, + ) + + # Qwen3Model uses stacked_params_mapping for q/k/v → qkv_proj + # and gate/up → gate_up_proj. Feed the remapped weights through + # the standard Qwen3 loading path. + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + + for name, loaded_weight in remapped: + if "rotary_emb.inv_freq" in name: + continue + + # Handle tied embeddings → lm_head. + if name == "model.embed_tokens.weight" and self.text_config.tie_word_embeddings: + # Also load into lm_head if present. + lm_key = "lm_head.weight" + if lm_key in params_dict: + p = params_dict[lm_key] + wl = getattr(p, "weight_loader", default_weight_loader) + wl(p, loaded_weight) + loaded_params.add(lm_key) + + # Try stacked params mapping (q/k/v → qkv_proj, gate/up → gate_up_proj). + handled = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + if mapped not in params_dict: + continue + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if weight_loader == default_weight_loader: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(mapped) + handled = True + break + + if handled: + continue + + # Direct parameter mapping. + if name in params_dict: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + logger.info("Loaded %d weights for FishSpeechSlowARForConditionalGeneration", len(loaded_params)) + + # Truncate RoPE cos/sin caches to bf16 precision to match training. + # Without this, f32 RoPE values cause logit divergence and premature EOS. + truncated = 0 + for module in self.modules(): + if hasattr(module, "cos_sin_cache") and isinstance(module.cos_sin_cache, torch.Tensor): + cache = module.cos_sin_cache + module.cos_sin_cache = cache.to(torch.bfloat16).to(cache.dtype) + truncated += 1 + if truncated: + logger.info("Truncated %d RoPE cos_sin_cache buffers to bf16 precision", truncated) + + return loaded_params diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 9bf2b4fdb..f708b6ff9 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -129,6 +129,17 @@ _OMNI_MODELS = { "hunyuan_image3", "HunyuanImage3ForConditionalGeneration", ), + ## fish_speech (Fish Speech S2 Pro) + "FishSpeechSlowARForConditionalGeneration": ( + "fish_speech", + "fish_speech_slow_ar", + "FishSpeechSlowARForConditionalGeneration", + ), + "FishSpeechDACDecoder": ( + "fish_speech", + "fish_speech_dac_decoder", + "FishSpeechDACDecoder", + ), } diff --git a/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml b/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml new file mode 100644 index 000000000..7f5ba1202 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml @@ -0,0 +1,97 @@ +async_chunk: true +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 4 + engine_args: + model_stage: fish_speech_slow_ar + model_arch: FishSpeechSlowARForConditionalGeneration + hf_overrides: + architectures: [FishSpeechSlowARForConditionalGeneration] + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + enforce_eager: false + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: latent + gpu_memory_utilization: 0.4 + distributed_executor_backend: "mp" + max_num_batched_tokens: 2048 + max_model_len: 32768 + custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.fish_speech.slow_ar_to_dac_decoder_async_chunk + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.8 + top_k: 30 + top_p: 0.9 + max_tokens: 200 + seed: 42 + detokenize: false + repetition_penalty: 1.0 + # <|im_end|> token -- stop when model emits end-of-turn. + stop_token_ids: [151645] + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: dac_decoder + model_arch: FishSpeechDACDecoder + hf_overrides: + architectures: [FishSpeechDACDecoder] + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio + gpu_memory_utilization: 0.2 + distributed_executor_backend: "mp" + max_num_batched_tokens: 8192 + max_model_len: 32768 + engine_input_source: [0] + final_output: true + final_output_type: audio + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + detokenize: true + repetition_penalty: 1.0 + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + + connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + # Chunk sizes for streaming -- ~21 Hz codec. + # 25 frames ≈ 1.16s of audio at 21.5 Hz. + codec_chunk_frames: 25 + codec_left_context_frames: 25 + initial_codec_chunk_frames: 0 + + edges: + - from: 0 + to: 1 + window_size: -1 diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py new file mode 100644 index 000000000..5be70bc18 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/fish_speech.py @@ -0,0 +1,151 @@ +"""Stage input processor for Fish Speech S2 Pro: Slow AR → DAC Decoder.""" + +from typing import Any + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_NUM_CODEBOOKS = 10 # 1 semantic + 9 residual + + +def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None: + """Extract the last frame of audio codes from the pooling output.""" + audio_codes = pooling_output.get("audio_codes") + if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: + return None + if audio_codes.ndim == 2: + frame = audio_codes[-1] + if frame.numel() == 0 or not bool(frame.any().item()): + return None + return frame.to(torch.long).reshape(-1) + if audio_codes.ndim == 1: + return audio_codes.to(torch.long).reshape(-1) + raise ValueError(f"Invalid audio_codes shape for Fish Speech async_chunk: {tuple(audio_codes.shape)}") + + +def slow_ar_to_dac_decoder( + stage_list: list[Any], + engine_input_source: list[int], + prompt: Any = None, + requires_multimodal_data: bool = False, +) -> list[Any]: + """Non-async processor: wait for Slow AR to finish, then pass all codes to DAC decoder.""" + from vllm_omni.inputs.data import OmniTokensPrompt + from vllm_omni.model_executor.stage_input_processors.qwen3_omni import _validate_stage_inputs + + slow_ar_outputs = _validate_stage_inputs(stage_list, engine_input_source) + dac_inputs: list[OmniTokensPrompt] = [] + + for output in slow_ar_outputs: + out = output.outputs[0] + # audio_codes shape: [num_frames, num_codebooks] + audio_codes = out.multimodal_output["audio_codes"].to(torch.long) + # Filter zero-padded frames. + valid_mask = audio_codes.any(dim=1) + audio_codes = audio_codes[valid_mask] + # Codebook-major flat: [num_codebooks * num_frames] + codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist() + dac_inputs.append( + OmniTokensPrompt( + prompt_token_ids=codec_codes, + multi_modal_data=None, + mm_processor_kwargs=None, + additional_information=None, + ) + ) + return dac_inputs + + +def slow_ar_to_dac_decoder_async_chunk( + transfer_manager: Any, + pooling_output: dict[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> dict[str, Any] | None: + """Async streaming processor: emit code chunks as they are produced. + + Accumulates per-step codes and emits fixed-size chunks with left context + overlap for smooth audio transitions, analogous to + ``talker2code2wav_async_chunk`` in Qwen3 TTS. + """ + request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + + if isinstance(pooling_output, dict): + frame = _extract_last_frame(pooling_output) + if frame is not None: + codec_codes = frame.cpu().tolist() + transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) + elif not finished: + return None + + connector = getattr(transfer_manager, "connector", None) + raw_cfg = getattr(connector, "config", {}) or {} + cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} + chunk_size = int(cfg.get("codec_chunk_frames", 25)) + left_context_size_config = int(cfg.get("codec_left_context_frames", 25)) + initial_chunk_size = int(cfg.get("initial_codec_chunk_frames", 0)) + + # Per-request override. + additional_information = getattr(request, "additional_information", None) + if ( + additional_information is not None + and hasattr(additional_information, "entries") + and "initial_codec_chunk_frames" in additional_information.entries + ): + entry = additional_information.entries["initial_codec_chunk_frames"] + if entry.list_data is not None and len(entry.list_data) == 1: + initial_chunk_size = int(entry.list_data[0]) + + if chunk_size <= 0 or left_context_size_config < 0 or initial_chunk_size < 0: + raise ValueError( + f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, " + f"codec_left_context_frames={left_context_size_config}, " + f"initial_codec_chunk_frames={initial_chunk_size}" + ) + if initial_chunk_size > chunk_size: + initial_chunk_size = chunk_size + + length = len(transfer_manager.code_prompt_token_ids[request_id]) + + if length <= 0: + if finished: + return { + "code_predictor_codes": [], + "finished": torch.tensor(True, dtype=torch.bool), + } + return None + + in_initial_phase = initial_chunk_size > 0 and length <= chunk_size + + if in_initial_phase: + already_sent = transfer_manager.put_req_chunk[request_id] * initial_chunk_size + pending = length - already_sent + if pending <= 0: + return None + if pending < initial_chunk_size and not finished: + return None + context_length = min(pending, initial_chunk_size) + left_context_size = max(0, length - context_length) + window_frames = transfer_manager.code_prompt_token_ids[request_id][:length] + else: + initial_coverage = (chunk_size // initial_chunk_size) * initial_chunk_size if initial_chunk_size > 0 else 0 + adjusted = length - initial_coverage + chunk_length = adjusted % chunk_size + if chunk_length != 0 and not finished: + return None + context_length = chunk_length if chunk_length != 0 else chunk_size + end_index = min(length, left_context_size_config + context_length) + left_context_size = max(0, int(end_index - context_length)) + window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] + + # Pack into codebook-major flat codes. + code_predictor_codes = torch.tensor(window_frames).transpose(0, 1).reshape(-1).tolist() + + return { + "code_predictor_codes": code_predictor_codes, + "left_context_size": left_context_size, + "finished": torch.tensor(finished, dtype=torch.bool), + } diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 136a680a3..59b23f914 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -14,6 +14,9 @@ _CLASS_TO_MODULE: dict[str, str] = { "Mammothmoda2Qwen2_5_VLConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", "Mammothmoda2Qwen2_5_VLTextConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", "Mammothmoda2Qwen2_5_VLVisionConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", + "FishSpeechConfig": "vllm_omni.transformers_utils.configs.fish_speech", + "FishSpeechSlowARConfig": "vllm_omni.transformers_utils.configs.fish_speech", + "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech", } __all__ = [ @@ -21,6 +24,9 @@ __all__ = [ "Mammothmoda2Qwen2_5_VLConfig", "Mammothmoda2Qwen2_5_VLTextConfig", "Mammothmoda2Qwen2_5_VLVisionConfig", + "FishSpeechConfig", + "FishSpeechSlowARConfig", + "FishSpeechFastARConfig", ] @@ -39,4 +45,5 @@ def __dir__(): # Eagerly import all config modules so their AutoConfig.register() side-effects # run as soon as `vllm_omni.transformers_utils.configs` is imported. +from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/fish_speech.py b/vllm_omni/transformers_utils/configs/fish_speech.py new file mode 100644 index 000000000..44b5a67eb --- /dev/null +++ b/vllm_omni/transformers_utils/configs/fish_speech.py @@ -0,0 +1,24 @@ +"""Fish Speech S2 Pro config registration with transformers AutoConfig. + +Registers FishSpeechConfig (model_type="fish_qwen3_omni") and sub-configs +so that ``AutoConfig.from_pretrained("fishaudio/s2-pro")`` returns the +correct config class. +""" + +from transformers import AutoConfig + +from vllm_omni.model_executor.models.fish_speech.configuration_fish_speech import ( + FishSpeechConfig, + FishSpeechFastARConfig, + FishSpeechSlowARConfig, +) + +AutoConfig.register("fish_qwen3_omni", FishSpeechConfig) +AutoConfig.register("fish_qwen3", FishSpeechSlowARConfig) +AutoConfig.register("fish_qwen3_audio_decoder", FishSpeechFastARConfig) + +__all__ = [ + "FishSpeechConfig", + "FishSpeechSlowARConfig", + "FishSpeechFastARConfig", +]