From 682398b22f52b5882832f4616126d20e861ba795 Mon Sep 17 00:00:00 2001 From: coding-crying Date: Sat, 27 Dec 2025 05:12:37 -0800 Subject: [PATCH] feat: add OpenAI-compatible streaming server with WebSocket support Adds an OpenAI-compatible TTS server with both HTTP and WebSocket endpoints: ## Endpoints - `POST /v1/audio/speech` - HTTP streaming (OpenAI-compatible) - `WS /v1/audio/speech/stream` - Bidirectional WebSocket streaming ## WebSocket Protocol 1. Connect to `ws://host:port/v1/audio/speech/stream` 2. Send config: `{"voice": "speaker.wav", "speed": 1.0}` 3. Send text chunks: `{"text": "Hello"}` 4. Send end signal: `{"event": "end"}` 5. Receive binary PCM audio (int16, 24kHz, mono) 6. Server closes connection when complete ## Changes - `openai_server.py`: FastAPI server with HTTP + WebSocket TTS endpoints - `run_openai_server.sh`: Launch script with venv activation - `cosyvoice/llm/llm.py`: Add `inference_bistream` to CosyVoice3LM for streaming support ## Usage ```bash bash run_openai_server.sh # Server runs on http://0.0.0.0:50000 ``` Tested with CosyVoice3-0.5B model for real-time voice agent applications. --- cosyvoice/llm/llm.py | 100 ++++++++++++++ openai_server.py | 318 +++++++++++++++++++++++++++++++++++++++++++ run_openai_server.sh | 18 +++ 3 files changed, 436 insertions(+) create mode 100644 openai_server.py create mode 100755 run_openai_server.sh diff --git a/cosyvoice/llm/llm.py b/cosyvoice/llm/llm.py index eacde5b3a..122266edc 100644 --- a/cosyvoice/llm/llm.py +++ b/cosyvoice/llm/llm.py @@ -737,3 +737,103 @@ def inference( # 5. step by step decode for token in self.inference_wrapper(lm_input, sampling, min_len, max_len, uuid): yield token + + @torch.inference_mode() + def inference_bistream( + self, + text: Generator, + prompt_text: torch.Tensor, + prompt_text_len: torch.Tensor, + prompt_speech_token: torch.Tensor, + prompt_speech_token_len: torch.Tensor, + embedding: torch.Tensor, + sampling: int = 25, + max_token_text_ratio: float = 20, + min_token_text_ratio: float = 2, + ) -> Generator[torch.Tensor, None, None]: + + device = prompt_text.device + # 1. prepare input + sos_emb = self.speech_embedding.weight[self.sos].reshape(1, 1, -1) + if prompt_speech_token_len != 0: + prompt_speech_token_emb = self.speech_embedding(prompt_speech_token) + else: + prompt_speech_token_emb = torch.zeros(1, 0, self.llm_input_size, dtype=prompt_text.dtype).to(device) + lm_input = torch.concat([sos_emb], dim=1) + + # 2. iterate text + out_tokens = [] + cache = None + # NOTE init prompt_text as text_cache as it is basically impossible prompt_speech_token/prompt_text < 15/5 + text_cache = self.llm.model.model.embed_tokens(prompt_text) + next_fill_index = (int(prompt_speech_token.shape[1] / self.mix_ratio[1]) + 1) * self.mix_ratio[1] - prompt_speech_token.shape[1] + for this_text in text: + text_cache = torch.concat([text_cache, self.llm.model.model.embed_tokens(this_text)], dim=1) + # prompt_speech_token_emb not empty, try append to lm_input + while prompt_speech_token_emb.size(1) != 0: + if text_cache.size(1) >= self.mix_ratio[0]: + lm_input_text, lm_input_speech = text_cache[:, :self.mix_ratio[0]], prompt_speech_token_emb[:, :self.mix_ratio[1]] + logging.info('append {} text token {} speech token'.format(lm_input_text.size(1), lm_input_speech.size(1))) + lm_input = torch.concat([lm_input, lm_input_text, lm_input_speech], dim=1) + text_cache, prompt_speech_token_emb = text_cache[:, self.mix_ratio[0]:], prompt_speech_token_emb[:, self.mix_ratio[1]:] + else: + logging.info('not enough text token to decode, wait for more') + break + # no prompt_speech_token_emb remain, can decode some speech token + if prompt_speech_token_emb.size(1) == 0: + if (len(out_tokens) != 0 and out_tokens[-1] == self.fill_token) or (len(out_tokens) == 0 and lm_input.size(1) == 1): + logging.info('get fill token, need to append more text token') + if text_cache.size(1) >= self.mix_ratio[0]: + lm_input_text = text_cache[:, :self.mix_ratio[0]] + logging.info('append {} text token'.format(lm_input_text.size(1))) + if len(out_tokens) != 0 and out_tokens[-1] == self.fill_token: + lm_input = lm_input_text + else: + lm_input = torch.concat([lm_input, lm_input_text], dim=1) + text_cache = text_cache[:, self.mix_ratio[0]:] + else: + logging.info('not enough text token to decode, wait for more') + continue + while True: + seq_len = lm_input.shape[1] if cache is None else lm_input.shape[1] + cache[0][0].size(2) + y_pred, cache = self.llm.forward_one_step(lm_input, + masks=torch.tril(torch.ones((1, seq_len, seq_len), device=lm_input.device)).to(torch.bool), + cache=cache) + logp = self.llm_decoder(y_pred[:, -1]).log_softmax(dim=-1) + if next_fill_index != -1 and len(out_tokens) == next_fill_index: + top_ids = self.fill_token + next_fill_index += (self.mix_ratio[1] + 1) + else: + top_ids = self.sampling_ids(logp.squeeze(dim=0), out_tokens, sampling, ignore_eos=True) + if top_ids == self.fill_token: + next_fill_index = len(out_tokens) + self.mix_ratio[1] + 1 + logging.info('fill_token index {} next fill_token index {}'.format(len(out_tokens), next_fill_index)) + out_tokens.append(top_ids) + if top_ids >= self.speech_token_size: + if top_ids == self.fill_token: + break + else: + raise ValueError('should not get token {}'.format(top_ids)) + yield top_ids + lm_input = self.speech_embedding.weight[top_ids].reshape(1, 1, -1) + + # 3. final decode + task_id_emb = self.speech_embedding.weight[self.task_id].reshape(1, 1, -1) + lm_input = torch.concat([lm_input, text_cache, task_id_emb], dim=1) + logging.info('no more text token, decode until met eos') + while True: + seq_len = lm_input.shape[1] if cache is None else lm_input.shape[1] + cache[0][0].size(2) + y_pred, cache = self.llm.forward_one_step(lm_input, + masks=torch.tril(torch.ones((1, seq_len, seq_len), device=lm_input.device)).to(torch.bool), + cache=cache) + logp = self.llm_decoder(y_pred[:, -1]).log_softmax(dim=-1) + top_ids = self.sampling_ids(logp.squeeze(dim=0), out_tokens, sampling, ignore_eos=False) + out_tokens.append(top_ids) + if top_ids >= self.speech_token_size: + if top_ids == self.eos_token: + break + else: + raise ValueError('should not get token {}'.format(top_ids)) + # in stream mode, yield token one by one + yield top_ids + lm_input = self.speech_embedding.weight[top_ids].reshape(1, 1, -1) diff --git a/openai_server.py b/openai_server.py new file mode 100644 index 000000000..75afbb38e --- /dev/null +++ b/openai_server.py @@ -0,0 +1,318 @@ +import os +import sys +import argparse +import logging +import asyncio +import threading +import subprocess +import numpy as np +import json +from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect +from fastapi.responses import StreamingResponse, JSONResponse +from fastapi.middleware.cors import CORSMiddleware +from pydantic import BaseModel +from typing import Optional, Literal, Generator + +# Adjust path to include CosyVoice modules +ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) +sys.path.append(ROOT_DIR) +sys.path.append(os.path.join(ROOT_DIR, 'third_party', 'Matcha-TTS')) + +try: + from cosyvoice.cli.cosyvoice import AutoModel +except ImportError as e: + print(f"Error importing CosyVoice: {e}") + sys.exit(1) + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("openai_server") + +app = FastAPI(title="CosyVoice OpenAI API") + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +# Global configuration +cosyvoice = None +MODEL_DIR = "pretrained_models/Fun-CosyVoice3-0.5B" +DEFAULT_PROMPT_WAV = "asset/zero_shot_prompt.wav" + +DEFAULT_PROMPT_TEXT = "希望你以后能够做的比我还好呦。" + +# CUSTOM VOICE MAP +VOICE_MAP = { + "russian": { + "text": "Всем привет, дорогие друзья! Сейчас 6.20 и мы с вами успели. Сегодня мы с вами встречаем восход солнца.", + "wav": "asset/russian_prompt.wav" + }, + "english": { + "text": "And then later on, fully acquiring that company. So keeping management in line, interest in line with the asset that's coming into the family is a reason why sometimes we don't buy the whole thing.", + "wav": "asset/cross_lingual_prompt.wav" + } +} + + +class SpeechRequest(BaseModel): + model: Optional[str] = "tts-1" + input: str + voice: Optional[str] = "中文女" + response_format: Optional[Literal['mp3', 'opus', 'aac', 'flac', 'wav', 'pcm']] = "mp3" + speed: Optional[float] = 1.0 + +def get_ffmpeg_cmd(format: str, sample_rate: int): + cmd = ['ffmpeg', '-f', 's16le', '-ar', str(sample_rate), '-ac', '1', '-i', 'pipe:0'] + if format == 'mp3': + cmd.extend(['-f', 'mp3', 'pipe:1']) + elif format == 'opus': + cmd.extend(['-f', 'opus', '-c:a', 'libopus', 'pipe:1']) + elif format == 'aac': + cmd.extend(['-f', 'adts', 'pipe:1']) + elif format == 'flac': + cmd.extend(['-f', 'flac', 'pipe:1']) + elif format == 'wav': + cmd.extend(['-f', 'wav', 'pipe:1']) + else: + raise ValueError(f"Unsupported format via ffmpeg: {format}") + return cmd + +@app.on_event("startup") +async def startup_event(): + global cosyvoice + logger.info(f"Loading model from {MODEL_DIR}...") + try: + if not os.path.exists(MODEL_DIR): + logger.warning(f"Model directory {MODEL_DIR} not found. Please run download_models.py first.") + else: + cosyvoice = AutoModel(model_dir=MODEL_DIR) + logger.info(f"Model loaded. Sample rate: {cosyvoice.sample_rate}") + available_spks = cosyvoice.list_available_spks() + logger.info(f"Available speakers: {available_spks}") + except Exception as e: + logger.error(f"Failed to load model: {e}") + +@app.get("/v1/models") +async def list_models(): + return JSONResponse(content={ + "object": "list", + "data": [ + { + "id": "cosyvoice-tts", + "object": "model", + "created": 1234567890, + "owned_by": "cosyvoice", + } + ] + }) + +@app.post("/v1/audio/speech") +async def text_to_speech(req: SpeechRequest): + if not cosyvoice: + raise HTTPException(status_code=500, detail="Model not loaded or invalid model directory.") + + text = req.input + spk_id = req.voice + speed = req.speed if req.speed else 1.0 + + # Custom Voice Logic + prompt_text = DEFAULT_PROMPT_TEXT + prompt_wav = DEFAULT_PROMPT_WAV + + # Check if requested voice matches our map (case-insensitive) + for key, val in VOICE_MAP.items(): + if key.lower() in spk_id.lower(): + prompt_text = val["text"] + prompt_wav = val["wav"] + logger.info(f"Using custom prompt for voice '{spk_id}': {val['wav']}") + break + format = req.response_format + + available_spks = cosyvoice.list_available_spks() + use_zero_shot = spk_id not in available_spks + + logger.info(f"TTS Request: text='{text[:20]}...', voice='{spk_id}', format={format}, speed={speed}, zero_shot={use_zero_shot}") + + async def audio_generator(): + loop = asyncio.get_running_loop() + queue = asyncio.Queue() + sentinel = object() + ffmpeg_proc = None + + if format != 'pcm': + try: + cmd = get_ffmpeg_cmd(format, cosyvoice.sample_rate) + ffmpeg_proc = subprocess.Popen( + cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL + ) + except Exception as e: + logger.error(f"ffmpeg error: {e}") + return + + def producer_thread(): + try: + if use_zero_shot: + generator = cosyvoice.inference_zero_shot(text, prompt_text, prompt_wav, stream=True, speed=speed) + else: + generator = cosyvoice.inference_sft(text, spk_id, stream=True, speed=speed) + + for i in generator: + raw_data = (i['tts_speech'].numpy() * (2 ** 15)).astype(np.int16).tobytes() + if ffmpeg_proc: + try: + ffmpeg_proc.stdin.write(raw_data) + ffmpeg_proc.stdin.flush() + except (BrokenPipeError, OSError): + break + else: + loop.call_soon_threadsafe(queue.put_nowait, raw_data) + + if ffmpeg_proc: + ffmpeg_proc.stdin.close() + else: + loop.call_soon_threadsafe(queue.put_nowait, sentinel) + except Exception as e: + logger.error(f"Producer error: {e}") + if ffmpeg_proc: + ffmpeg_proc.terminate() + loop.call_soon_threadsafe(queue.put_nowait, sentinel) + + def reader_thread(): + if not ffmpeg_proc: return + try: + while True: + chunk = ffmpeg_proc.stdout.read(4096) + if not chunk: break + loop.call_soon_threadsafe(queue.put_nowait, chunk) + ffmpeg_proc.wait() + except Exception as e: + logger.error(f"Reader error: {e}") + finally: + loop.call_soon_threadsafe(queue.put_nowait, sentinel) + + threading.Thread(target=producer_thread, daemon=True).start() + if ffmpeg_proc: + threading.Thread(target=reader_thread, daemon=True).start() + + while True: + chunk = await queue.get() + if chunk is sentinel: break + yield chunk + + media_type = f"audio/{format}" + if format == 'pcm': media_type = "application/octet-stream" + elif format == 'mp3': media_type = "audio/mpeg" + elif format == 'wav': media_type = "audio/wav" + + return StreamingResponse(audio_generator(), media_type=media_type) + +@app.websocket("/v1/audio/speech/stream") +async def websocket_endpoint(websocket: WebSocket): + await websocket.accept() + logger.info("WebSocket connection established") + + try: + # 1. Receive config + config_msg = await websocket.receive_text() + config = json.loads(config_msg) + spk_id = config.get("voice", "中文女") + speed = config.get("speed", 1.0) + language = config.get("language") + # format = config.get("response_format", "pcm") # For simplicity, WS yields PCM for now + + available_spks = cosyvoice.list_available_spks() + use_zero_shot = spk_id not in available_spks + + # 2. Setup text generator for bi-streaming + text_queue = asyncio.Queue() + + def sync_text_generator(): + # This generator will be consumed by CosyVoice in a background thread + first_chunk = True + while True: + # We need to get items from the async queue in a sync way + future = asyncio.run_coroutine_threadsafe(text_queue.get(), loop) + val = future.result() + if val is None: + break + + # Prepend language tag if provided in config + if first_chunk and language: + val = f"<|{language}|>" + val + first_chunk = False + elif first_chunk: + first_chunk = False + + yield val + + loop = asyncio.get_running_loop() + + # 3. Start inference in a separate thread + inference_finished_event = asyncio.Event() + + def run_inference(): + try: + # Use the sync generator which bridges to the async queue + if use_zero_shot: + # Note: inference_zero_shot with generator is only supported if the model is CosyVoice2/3 + # and it calls inference_bistream internally. + gen = cosyvoice.inference_zero_shot(sync_text_generator(), DEFAULT_PROMPT_TEXT, DEFAULT_PROMPT_WAV, stream=True, speed=speed) + else: + gen = cosyvoice.inference_sft(sync_text_generator(), spk_id, stream=True, speed=speed) + + for output in gen: + audio_data = (output['tts_speech'].numpy() * (2 ** 15)).astype(np.int16).tobytes() + asyncio.run_coroutine_threadsafe(websocket.send_bytes(audio_data), loop) + except Exception as e: + logger.error(f"WS Inference error: {e}") + import traceback + logger.error(traceback.format_exc()) + finally: + loop.call_soon_threadsafe(inference_finished_event.set) + + inference_thread = threading.Thread(target=run_inference, daemon=True) + inference_thread.start() + + # 4. Receive text tokens + while True: + data = await websocket.receive_text() + msg = json.loads(data) + text_chunk = msg.get("text") + if text_chunk: + await text_queue.put(text_chunk) + if msg.get("event") == "flush": + # CosyVoice doesn't have an explicit flush per chunk, + # but we can handle it if we want to restart the generator + pass + if msg.get("event") == "end": + await text_queue.put(None) + break + + # Wait for inference to finish sending all audio + await inference_finished_event.wait() + + # Explicitly close WebSocket after all audio sent + await websocket.close() + logger.info("WebSocket closed after audio complete") + + except WebSocketDisconnect: + logger.info("WebSocket disconnected") + except Exception as e: + logger.error(f"WebSocket error: {e}") + finally: + # Ensure thread stops + await text_queue.put(None) + +if __name__ == "__main__": + import uvicorn + parser = argparse.ArgumentParser() + parser.add_argument("--port", type=int, default=50000) + parser.add_argument("--host", type=str, default="0.0.0.0") + parser.add_argument("--model_dir", type=str, default="pretrained_models/Fun-CosyVoice3-0.5B") + args = parser.parse_args() + + MODEL_DIR = args.model_dir + uvicorn.run(app, host=args.host, port=args.port) \ No newline at end of file diff --git a/run_openai_server.sh b/run_openai_server.sh new file mode 100755 index 000000000..67355f476 --- /dev/null +++ b/run_openai_server.sh @@ -0,0 +1,18 @@ +#!/bin/bash +MODEL_DIR="pretrained_models/Fun-CosyVoice3-0.5B" + +if [ ! -d "$MODEL_DIR" ]; then + echo "Model directory $MODEL_DIR not found. Downloading models..." + source venv/bin/activate + python download_models.py +else + echo "Model directory found." + source venv/bin/activate +fi + +echo "Starting OpenAI-compatible server..." +# Check if uvicorn is installed, if not try pip install +python -c "import uvicorn" 2>/dev/null || pip install uvicorn fastapi + +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.13/site-packages/nvidia/cudnn/lib +python openai_server.py --port 50000 --model_dir "$MODEL_DIR"