150 строки
4.9 KiB
Python
150 строки
4.9 KiB
Python
"""Minimal FastAPI server exposing MOSS-TTS generation and voice cloning. https://github.com/OpenMOSS/MOSS-TTS
|
|
|
|
Other checkpoints usable via MOSSTTS_MODEL (same build_user_message/generate/decode
|
|
API unless noted):
|
|
- OpenMOSS-Team/MOSS-TTS-v1.5 8B, 24kHz, the flagship model
|
|
- OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 4B, 48kHz stereo (default here)
|
|
- OpenMOSS-Team/MOSS-TTSD-v1.0 8B, same arch as v1.5, earlier release
|
|
- OpenMOSS-Team/MOSS-TTS-Nano 0.1B, CPU-friendly (as few as 4 cores),
|
|
streaming-oriented; API not verified
|
|
- OpenMOSS-Team/MOSS-TTS-Realtime 1.7B, multi-turn/streaming voice-agent
|
|
model with a different generation API
|
|
(not build_user_message/generate/decode);
|
|
NOT compatible with this server as-is
|
|
"""
|
|
import io
|
|
import os
|
|
import tempfile
|
|
|
|
import soundfile as sf
|
|
import torch
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel
|
|
from transformers import AutoModel, AutoProcessor
|
|
|
|
MODEL_NAME = os.environ.get("MOSSTTS_MODEL", "OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5")
|
|
ATTN_IMPLEMENTATION = os.environ.get("MOSSTTS_ATTN_IMPL", "sdpa")
|
|
MAX_NEW_TOKENS = int(os.environ.get("MOSSTTS_MAX_NEW_TOKENS", "4096"))
|
|
|
|
_forced_device = os.environ.get("MOSSTTS_DEVICE", "").strip()
|
|
if _forced_device:
|
|
DEVICE = _forced_device
|
|
elif torch.cuda.is_available(): # ROCm builds expose the GPU via the CUDA API too
|
|
DEVICE = "cuda"
|
|
else:
|
|
DEVICE = "cpu"
|
|
|
|
# bfloat16 is not well supported on CPU; use full precision there.
|
|
DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
|
|
|
|
if DEVICE == "cuda":
|
|
# Match the backends enabled in the upstream quickstart; avoids a broken
|
|
# cuDNN SDP kernel on some GPU combos (same class of issue as chatterbox).
|
|
torch.backends.cuda.enable_cudnn_sdp(False)
|
|
torch.backends.cuda.enable_flash_sdp(True)
|
|
torch.backends.cuda.enable_mem_efficient_sdp(True)
|
|
torch.backends.cuda.enable_math_sdp(True)
|
|
|
|
app = FastAPI(title="MOSS-TTS server")
|
|
|
|
model = None
|
|
processor = None
|
|
|
|
|
|
@app.on_event("startup")
|
|
def load_model():
|
|
global model, processor
|
|
print(f"Loading MOSS-TTS processor/model {MODEL_NAME} on {DEVICE}...")
|
|
|
|
processor = AutoProcessor.from_pretrained(
|
|
MODEL_NAME,
|
|
trust_remote_code=True,
|
|
)
|
|
processor.audio_tokenizer = processor.audio_tokenizer.to(DEVICE)
|
|
|
|
model = AutoModel.from_pretrained(
|
|
MODEL_NAME,
|
|
trust_remote_code=True,
|
|
attn_implementation=ATTN_IMPLEMENTATION,
|
|
dtype=DTYPE,
|
|
).to(DEVICE)
|
|
model.eval()
|
|
|
|
print("MOSS-TTS model loaded.")
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok", "model": MODEL_NAME, "device": DEVICE}
|
|
|
|
|
|
def _run_generation(text: str, language: str, reference: list[str] | None) -> tuple:
|
|
if model is None or processor is None:
|
|
raise HTTPException(status_code=503, detail="Model not loaded yet")
|
|
|
|
message_kwargs = {"text": text}
|
|
if language:
|
|
message_kwargs["language"] = language
|
|
if reference:
|
|
message_kwargs["reference"] = reference
|
|
|
|
conversation = [processor.build_user_message(**message_kwargs)]
|
|
batch = processor([conversation], mode="generation")
|
|
input_ids = batch["input_ids"].to(DEVICE)
|
|
attention_mask = batch["attention_mask"].to(DEVICE)
|
|
|
|
with torch.no_grad():
|
|
outputs = model.generate(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
max_new_tokens=MAX_NEW_TOKENS,
|
|
)
|
|
|
|
decoded = processor.decode(outputs)[0]
|
|
audio = decoded.audio_codes_list[0]
|
|
if audio.ndim > 1:
|
|
audio = audio.squeeze(0)
|
|
|
|
return audio.detach().cpu().to(torch.float32).numpy(), processor.model_config.sampling_rate
|
|
|
|
|
|
class SpeechRequest(BaseModel):
|
|
text: str
|
|
language: str = ""
|
|
|
|
|
|
@app.post("/speech")
|
|
def speech(req: SpeechRequest):
|
|
audio, sr = _run_generation(req.text, req.language, reference=None)
|
|
|
|
buf = io.BytesIO()
|
|
sf.write(buf, audio, sr, format="WAV")
|
|
buf.seek(0)
|
|
return StreamingResponse(buf, media_type="audio/wav")
|
|
|
|
|
|
@app.post("/speech/clone")
|
|
def speech_clone(
|
|
text: str = Form(...),
|
|
language: str = Form(""),
|
|
ref_audio: UploadFile = File(...),
|
|
):
|
|
"""Clone a voice from ref_audio and synthesize text with it."""
|
|
ref_bytes = ref_audio.file.read()
|
|
|
|
suffix = os.path.splitext(ref_audio.filename or "")[1] or ".wav"
|
|
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
|
|
tmp.write(ref_bytes)
|
|
tmp_path = tmp.name
|
|
|
|
try:
|
|
audio, sr = _run_generation(text, language, reference=[tmp_path])
|
|
finally:
|
|
os.remove(tmp_path)
|
|
|
|
buf = io.BytesIO()
|
|
sf.write(buf, audio, sr, format="WAV")
|
|
buf.seek(0)
|
|
return StreamingResponse(buf, media_type="audio/wav")
|