diff --git a/examples/offline_inference/moss_tts_nano/end2end.py b/examples/offline_inference/moss_tts_nano/end2end.py index 9eb0be671..adb8cadbb 100644 --- a/examples/offline_inference/moss_tts_nano/end2end.py +++ b/examples/offline_inference/moss_tts_nano/end2end.py @@ -40,9 +40,7 @@ from vllm_omni import Omni # noqa: E402 MODEL = "OpenMOSS-Team/MOSS-TTS-Nano" _REPO_ROOT = Path(__file__).resolve().parents[4] -DEFAULT_STAGE_CONFIG = str( - _REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "moss_tts_nano.yaml" -) +DEFAULT_STAGE_CONFIG = str(_REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "moss_tts_nano.yaml") BATCH_SAMPLES = [ {"text": "Hello, this is a test of MOSS-TTS-Nano.", "voice": "Ava", "label": "en_hello"}, @@ -118,10 +116,7 @@ def main(args) -> None: if args.batch: print(f"Running batch synthesis ({len(BATCH_SAMPLES)} samples)...") - inputs = [ - build_request(s["text"], voice=s["voice"], seed=args.seed) - for s in BATCH_SAMPLES - ] + inputs = [build_request(s["text"], voice=s["voice"], seed=args.seed) for s in BATCH_SAMPLES] params_list = [sampling_params] * len(inputs) else: print(f"Synthesizing: {args.text!r}") diff --git a/examples/online_serving/moss_tts_nano/gradio_demo.py b/examples/online_serving/moss_tts_nano/gradio_demo.py index 6eeb8d4c0..f120f838c 100644 --- a/examples/online_serving/moss_tts_nano/gradio_demo.py +++ b/examples/online_serving/moss_tts_nano/gradio_demo.py @@ -27,9 +27,7 @@ import logging try: import gradio as gr except ImportError: - raise ImportError( - "gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'" - ) from None + raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None import httpx import numpy as np @@ -602,8 +600,8 @@ def create_app(api_base: str): if stream: return ( gr.update(value="pcm", interactive=False), - gr.update(visible=True), # player - gr.update(visible=False), # audio + gr.update(visible=True), # player + gr.update(visible=False), # audio ) return ( gr.update(value="wav", interactive=True), @@ -619,9 +617,9 @@ def create_app(api_base: str): def on_reset(): return ( - "", # text - None, # audio_output - "", # hidden_payload + "", # text + None, # audio_output + "", # hidden_payload PLAYER_HTML, # reset player ) @@ -672,9 +670,7 @@ def main(): parser = argparse.ArgumentParser( description="Gradio demo for MOSS-TTS-Nano with gapless AudioWorklet streaming.", ) - parser.add_argument( - "--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})" - ) + 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 via Gradio tunnel") @@ -686,6 +682,7 @@ def main(): app = create_app(args.api_base) import uvicorn + uvicorn.run(app, host=args.host, port=args.port) diff --git a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py index 2ef79c91c..5b9401caf 100644 --- a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py +++ b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py @@ -19,7 +19,6 @@ so that vLLM initialises distributed state before any CUDA allocations occur. from __future__ import annotations -import logging import tempfile import threading from collections.abc import Iterable @@ -40,8 +39,10 @@ def _patch_torchaudio_load() -> None: """Patch torchaudio.load to use soundfile if torchcodec is unavailable.""" try: import torchaudio + torchaudio # noqa import torchcodec # noqa: F401 + return except Exception: pass @@ -67,6 +68,7 @@ def _patch_torchaudio_load() -> None: try: import torchaudio + torchaudio.load = _soundfile_load torchaudio.save = _soundfile_save logger.info("Patched torchaudio.load/save to use soundfile (torchcodec unavailable)") @@ -151,12 +153,16 @@ class MossTTSNanoForGeneration(nn.Module): logger.info("Loading MOSS-TTS-Nano LM from %s (dtype=%s)", self.model_path, tts_dtype) from transformers import AutoModelForCausalLM + lm = AutoModelForCausalLM.from_pretrained( - self.model_path, trust_remote_code=True, torch_dtype=tts_dtype, + self.model_path, + trust_remote_code=True, + torch_dtype=tts_dtype, ) if device.type == "cuda": try: import flash_attn # noqa: F401 + lm._set_attention_implementation("flash_attention_2") logger.info("MOSS-TTS-Nano using flash_attention_2") except ImportError: @@ -168,13 +174,17 @@ class MossTTSNanoForGeneration(nn.Module): logger.info("MOSS-TTS-Nano LM loaded on %s", device) codec_path: str = getattr( - self.config, "audio_tokenizer_pretrained_name_or_path", + self.config, + "audio_tokenizer_pretrained_name_or_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano", ) logger.info("Loading MOSS-Audio-Tokenizer-Nano from %s", codec_path) from transformers import AutoModel + audio_tokenizer = AutoModel.from_pretrained( - codec_path, trust_remote_code=True, torch_dtype=torch.float32, + codec_path, + trust_remote_code=True, + torch_dtype=torch.float32, ) audio_tokenizer.to(device=device) audio_tokenizer.eval() @@ -208,7 +218,6 @@ class MossTTSNanoForGeneration(nn.Module): yield torch.zeros((sr,), dtype=torch.float32), True return - voice: str = str(_pick(info, "voice", _DEFAULT_VOICE)) mode: str = str(_pick(info, "mode", _DEFAULT_MODE)) prompt_audio_path: str | None = _pick(info, "prompt_audio_path", None) if prompt_audio_path is not None: @@ -402,7 +411,9 @@ class MossTTSNanoForGeneration(nn.Module): vocab_size = int(getattr(self.config, "vocab_size", 32000)) num_rows = int(hidden_states.shape[0]) logits = torch.zeros( - (num_rows, vocab_size), dtype=torch.float32, device=hidden_states.device, + (num_rows, vocab_size), + dtype=torch.float32, + device=hidden_states.device, ) eos_id = 2 if vocab_size > 2 else 0 safe_id = 1 if vocab_size > 1 and 1 != eos_id else 0