diff --git a/examples/openai_server.py b/examples/openai_server.py index 7845b2f..d10065f 100644 --- a/examples/openai_server.py +++ b/examples/openai_server.py @@ -72,6 +72,24 @@ last_voices_mtime: float = 0.0 default_voice: Optional[str] = None SAMPLE_RATE = 24000 # updated once the model loads _model_lock = threading.Lock() # prevent concurrent GPU inference +aligner_model = None + +def _get_aligner(): + global aligner_model + if aligner_model is None: + try: + from qwen_asr import Qwen3ForcedAligner + import torch + except ImportError: + raise HTTPException(status_code=500, detail="qwen-asr is not installed. Run: pip install qwen-asr") + logger.info("Loading Qwen3-ForcedAligner-0.6B...") + aligner_model = Qwen3ForcedAligner.from_pretrained( + "Qwen/Qwen3-ForcedAligner-0.6B", + dtype=torch.bfloat16, + device_map="cuda" + ) + logger.info("Aligner loaded.") + return aligner_model def _voice_seed(voice_name: str) -> int: @@ -99,8 +117,8 @@ class SpeechRequest(BaseModel): model: str = "tts-1" input: str voice: str = "alloy" - response_format: str = "wav" # wav | pcm | mp3 - speed: float = 1.0 # accepted but not yet applied + response_format: str = "wav" # wav | pcm | mp3 | zip + speed: float = 1.0 # scales audio tempo # --------------------------------------------------------------------------- @@ -243,7 +261,7 @@ def _load_voice_clone_prompt(voice_cfg: dict, voice_name: str, tts_model): return None -async def _stream_chunks(voice_cfg: dict, text: str, voice_name: str) -> AsyncGenerator[bytes, None]: +async def _stream_chunks(voice_cfg: dict, text: str, voice_name: str, speed: float) -> AsyncGenerator[bytes, None]: """ Run generate_voice_clone_streaming in a background thread and yield raw PCM bytes for each chunk as they arrive. @@ -252,6 +270,31 @@ async def _stream_chunks(voice_cfg: dict, text: str, voice_name: str) -> AsyncGe _DONE = object() def producer(): + process = None + if speed != 1.0: + import subprocess + cmd = [ + "ffmpeg", "-y", "-loglevel", "error", + "-f", "s16le", "-ar", str(SAMPLE_RATE), "-ac", "1", "-i", "pipe:0", + "-filter:a", f"atempo={speed}", + "-f", "s16le", "-ar", str(SAMPLE_RATE), "-ac", "1", "pipe:1" + ] + process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE) + + def ffmpeg_reader(): + try: + while True: + out = process.stdout.read(4096) + if not out: + break + q.put(out) + except Exception as e: + q.put(e) + finally: + q.put(_DONE) + + threading.Thread(target=ffmpeg_reader, daemon=True).start() + try: with _model_lock: _seed_rng(voice_cfg.get("seed", _voice_seed(voice_name))) @@ -268,11 +311,22 @@ async def _stream_chunks(voice_cfg: dict, text: str, voice_name: str) -> AsyncGe top_k=voice_cfg.get("top_k", 50), top_p=voice_cfg.get("top_p", 0.9), ): - q.put(chunk) + raw = _to_pcm16(chunk) + if process: + process.stdin.write(raw) + process.stdin.flush() + else: + q.put(raw) except Exception as exc: q.put(exc) finally: - q.put(_DONE) + if process: + try: + process.stdin.close() + except Exception: + pass + else: + q.put(_DONE) thread = threading.Thread(target=producer, daemon=True) thread.start() @@ -284,7 +338,7 @@ async def _stream_chunks(voice_cfg: dict, text: str, voice_name: str) -> AsyncGe break if isinstance(item, Exception): raise item - yield _to_pcm16(item) + yield item # --------------------------------------------------------------------------- @@ -301,7 +355,8 @@ async def health(): async def create_speech(req: SpeechRequest): if tts_model is None: raise HTTPException(status_code=503, detail="Model not loaded") - if not req.input.strip(): + req.input = req.input.strip() + if not req.input: raise HTTPException(status_code=400, detail="'input' text is empty") voice_cfg = resolve_voice(req.voice) @@ -311,16 +366,17 @@ async def create_speech(req: SpeechRequest): "wav": "audio/wav", "pcm": "audio/pcm", "mp3": "audio/mpeg", + "zip": "application/zip", } if fmt not in _CONTENT_TYPES: raise HTTPException( status_code=400, - detail=f"response_format {fmt!r} not supported. Use: wav, pcm, mp3", + detail=f"response_format {fmt!r} not supported. Use: wav, pcm, mp3, zip", ) content_type = _CONTENT_TYPES[fmt] - # --- MP3: generate all audio, then encode (non-streaming) --- - if fmt == "mp3": + # --- MP3 / ZIP: generate all audio, then encode (non-streaming) --- + if fmt in ("mp3", "zip"): loop = asyncio.get_event_loop() def _generate(): @@ -341,13 +397,46 @@ async def create_speech(req: SpeechRequest): audio_arrays, sr = await loop.run_in_executor(None, _generate) audio = audio_arrays[0] if audio_arrays else np.zeros(1, dtype=np.float32) + + if req.speed != 1.0: + import subprocess + cmd = [ + "ffmpeg", "-y", "-loglevel", "error", + "-f", "f32le", "-ar", str(sr), "-ac", "1", "-i", "pipe:0", + "-filter:a", f"atempo={req.speed}", + "-f", "f32le", "-ar", str(sr), "-ac", "1", "pipe:1" + ] + process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE) + process.stdin.write(audio.tobytes()) + process.stdin.close() + out = process.stdout.read() + audio = np.frombuffer(out, dtype=np.float32) + + if fmt == "zip": + def _align(): + aligner = _get_aligner() + res = aligner.align(audio=(audio, sr), text=req.input, language=voice_cfg.get("language", "Auto")) + import dataclasses + return [dataclasses.asdict(x) for x in res] + + align_data = await loop.run_in_executor(None, _align) + + import zipfile + mp3_bytes = _to_mp3_bytes(audio, sr) + zip_buf = io.BytesIO() + with zipfile.ZipFile(zip_buf, "w", zipfile.ZIP_DEFLATED) as zf: + zf.writestr("audio.mp3", mp3_bytes) + zf.writestr("timer.json", json.dumps(align_data, ensure_ascii=False)) + + return Response(content=zip_buf.getvalue(), media_type=content_type) + return Response(content=_to_mp3_bytes(audio, sr), media_type=content_type) # --- WAV / PCM: stream chunks as they are generated --- async def audio_stream(): if fmt == "wav": yield _wav_header(SAMPLE_RATE) # stream with unknown data length - async for raw_chunk in _stream_chunks(voice_cfg, req.input, req.voice): + async for raw_chunk in _stream_chunks(voice_cfg, req.input, req.voice, req.speed): yield raw_chunk return StreamingResponse(audio_stream(), media_type=content_type)