239 lines
7.3 KiB
Python
239 lines
7.3 KiB
Python
import argparse
|
|
import hashlib
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
|
|
import numpy as np
|
|
import uvicorn
|
|
from fastapi import FastAPI, File, Form, UploadFile
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
logging.getLogger("matplotlib").setLevel(logging.WARNING)
|
|
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.append(os.path.join(ROOT_DIR, "../../.."))
|
|
sys.path.append(os.path.join(ROOT_DIR, "../../../third_party/Matcha-TTS"))
|
|
|
|
from cosyvoice.cli.cosyvoice import AutoModel
|
|
|
|
app = FastAPI(title="CosyVoice 3 streaming API")
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
cosyvoice = None
|
|
model_name = ""
|
|
fp16_enabled = False
|
|
speaker_cache_lock = threading.Lock()
|
|
cancel_events = {}
|
|
cancel_events_lock = threading.Lock()
|
|
|
|
|
|
def persist_upload(upload: UploadFile) -> str:
|
|
suffix = os.path.splitext(upload.filename or "")[1] or ".wav"
|
|
with tempfile.NamedTemporaryFile(prefix="cosyvoice-prompt-", suffix=suffix, delete=False) as target:
|
|
upload.file.seek(0)
|
|
shutil.copyfileobj(upload.file, target)
|
|
return target.name
|
|
|
|
|
|
def speaker_cache_id(prompt_text: str, prompt_wav: str) -> str:
|
|
digest = hashlib.sha256(prompt_text.encode("utf-8"))
|
|
with open(prompt_wav, "rb") as source:
|
|
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return "voice-" + digest.hexdigest()[:24]
|
|
|
|
|
|
def ensure_cached_speaker(prompt_text: str, prompt_wav: str) -> str:
|
|
cache_id = speaker_cache_id(prompt_text, prompt_wav)
|
|
with speaker_cache_lock:
|
|
if cache_id not in cosyvoice.frontend.spk2info:
|
|
cosyvoice.add_zero_shot_spk(prompt_text, prompt_wav, cache_id)
|
|
logging.info("cached zero-shot speaker %s", cache_id)
|
|
return cache_id
|
|
|
|
|
|
def register_cancel_event(request_id: str):
|
|
request_id = request_id.strip()
|
|
if not request_id:
|
|
return None
|
|
|
|
event = threading.Event()
|
|
with cancel_events_lock:
|
|
previous = cancel_events.get(request_id)
|
|
if previous is not None:
|
|
previous.set()
|
|
cancel_events[request_id] = event
|
|
return event
|
|
|
|
|
|
def release_cancel_event(request_id: str, event):
|
|
if not request_id or event is None:
|
|
return
|
|
with cancel_events_lock:
|
|
if cancel_events.get(request_id) is event:
|
|
cancel_events.pop(request_id, None)
|
|
|
|
|
|
def pcm_stream(model_output, cleanup_path: str = "", request_id: str = "", cancel_event=None):
|
|
iterator = iter(model_output)
|
|
try:
|
|
while True:
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
break
|
|
try:
|
|
item = next(iterator)
|
|
except StopIteration:
|
|
break
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
break
|
|
audio = item["tts_speech"].detach().cpu().numpy()
|
|
yield (audio * (2**15)).astype(np.int16).tobytes()
|
|
finally:
|
|
close = getattr(iterator, "close", None)
|
|
if callable(close):
|
|
close()
|
|
release_cancel_event(request_id, cancel_event)
|
|
if cleanup_path:
|
|
try:
|
|
os.remove(cleanup_path)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
|
|
def response(model_output, cleanup_path: str = "", request_id: str = ""):
|
|
request_id = request_id.strip()[:128]
|
|
cancel_event = register_cancel_event(request_id)
|
|
return StreamingResponse(
|
|
pcm_stream(model_output, cleanup_path, request_id, cancel_event),
|
|
media_type="application/octet-stream",
|
|
headers={
|
|
"X-Sample-Rate": str(cosyvoice.sample_rate),
|
|
"X-Audio-Format": "pcm_s16le",
|
|
"X-Accel-Buffering": "no",
|
|
"Cache-Control": "no-store, no-transform",
|
|
},
|
|
)
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {
|
|
"status": "ok",
|
|
"model": model_name,
|
|
"sample_rate": cosyvoice.sample_rate,
|
|
"streaming": True,
|
|
"fp16": fp16_enabled,
|
|
"cached_speakers": len(cosyvoice.frontend.spk2info),
|
|
"active_streams": len(cancel_events),
|
|
}
|
|
|
|
|
|
@app.post("/cancel/{request_id}")
|
|
def cancel(request_id: str):
|
|
with cancel_events_lock:
|
|
event = cancel_events.get(request_id)
|
|
if event is not None:
|
|
event.set()
|
|
return {"cancelled": event is not None, "request_id": request_id}
|
|
|
|
|
|
@app.get("/inference_sft")
|
|
@app.post("/inference_sft")
|
|
def inference_sft(tts_text: str = Form(), spk_id: str = Form(), request_id: str = Form("")):
|
|
return response(cosyvoice.inference_sft(tts_text, spk_id, stream=True), request_id=request_id)
|
|
|
|
|
|
@app.get("/inference_zero_shot")
|
|
@app.post("/inference_zero_shot")
|
|
def inference_zero_shot(
|
|
tts_text: str = Form(),
|
|
prompt_text: str = Form(),
|
|
prompt_wav: UploadFile = File(),
|
|
request_id: str = Form(""),
|
|
):
|
|
prompt_path = persist_upload(prompt_wav)
|
|
cache_id = ensure_cached_speaker(prompt_text, prompt_path)
|
|
return response(
|
|
cosyvoice.inference_zero_shot(
|
|
tts_text,
|
|
prompt_text,
|
|
prompt_path,
|
|
zero_shot_spk_id=cache_id,
|
|
stream=True,
|
|
),
|
|
prompt_path,
|
|
request_id,
|
|
)
|
|
|
|
|
|
@app.get("/inference_cross_lingual")
|
|
@app.post("/inference_cross_lingual")
|
|
def inference_cross_lingual(
|
|
tts_text: str = Form(),
|
|
prompt_wav: UploadFile = File(),
|
|
request_id: str = Form(""),
|
|
):
|
|
prompt_path = persist_upload(prompt_wav)
|
|
return response(
|
|
cosyvoice.inference_cross_lingual(tts_text, prompt_path, stream=True),
|
|
prompt_path,
|
|
request_id,
|
|
)
|
|
|
|
|
|
@app.get("/inference_instruct")
|
|
@app.post("/inference_instruct")
|
|
def inference_instruct(
|
|
tts_text: str = Form(),
|
|
spk_id: str = Form(),
|
|
instruct_text: str = Form(),
|
|
request_id: str = Form(""),
|
|
):
|
|
return response(
|
|
cosyvoice.inference_instruct(tts_text, spk_id, instruct_text, stream=True),
|
|
request_id=request_id,
|
|
)
|
|
|
|
|
|
@app.get("/inference_instruct2")
|
|
@app.post("/inference_instruct2")
|
|
def inference_instruct2(
|
|
tts_text: str = Form(),
|
|
instruct_text: str = Form(),
|
|
prompt_wav: UploadFile = File(),
|
|
request_id: str = Form(""),
|
|
):
|
|
prompt_path = persist_upload(prompt_wav)
|
|
return response(
|
|
cosyvoice.inference_instruct2(tts_text, instruct_text, prompt_path, stream=True),
|
|
prompt_path,
|
|
request_id,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--port", type=int, default=50000)
|
|
parser.add_argument("--model_dir", type=str, default="FunAudioLLM/Fun-CosyVoice3-0.5B-2512")
|
|
parser.add_argument("--fp16", action="store_true", help="Run the PyTorch model in FP16 on CUDA")
|
|
args = parser.parse_args()
|
|
model_name = args.model_dir
|
|
fp16_enabled = bool(args.fp16)
|
|
cosyvoice = AutoModel(model_dir=args.model_dir, fp16=fp16_enabled)
|
|
|
|
default_prompt_wav = os.path.join(ROOT_DIR, "../../../asset/zero_shot_prompt.wav")
|
|
default_prompt_text = "You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。"
|
|
if os.path.isfile(default_prompt_wav):
|
|
ensure_cached_speaker(default_prompt_text, default_prompt_wav)
|
|
|
|
uvicorn.run(app, host="0.0.0.0", port=args.port)
|