313 lines
11 KiB
Python
313 lines
11 KiB
Python
import os
|
||
import io
|
||
|
||
from fastapi import FastAPI, Response, Body, HTTPException
|
||
from fastapi.responses import StreamingResponse, JSONResponse
|
||
from contextlib import asynccontextmanager
|
||
import uvicorn
|
||
import xml.etree.ElementTree as ET
|
||
|
||
from kokoro import KPipeline, KModel
|
||
# import soundfile as sf
|
||
import wave
|
||
import numpy as np
|
||
from scipy.signal import resample
|
||
|
||
import torch
|
||
from torch import Tensor
|
||
from torch.nn import functional as F
|
||
from typing import Optional, List
|
||
import re
|
||
from dataclasses import dataclass
|
||
|
||
import kokoro.istftnet as _ist
|
||
|
||
def _inverse_no_complex(self, magnitude, phase):
|
||
"""
|
||
解决 MLU 设备上不支持复数计算
|
||
"""
|
||
device = magnitude.device
|
||
dtype = magnitude.dtype
|
||
|
||
win_dev = torch.hann_window(self.win_length, device=device, dtype=dtype)
|
||
|
||
real = magnitude * torch.cos(phase)
|
||
imag = magnitude * torch.sin(phase)
|
||
spec_ri = torch.stack([real, imag], dim=-1).contiguous() # (..., 2)
|
||
|
||
real_cpu = real.to("cpu")
|
||
imag_cpu = imag.to("cpu")
|
||
spec_complex_cpu = torch.complex(real_cpu, imag_cpu) # (..,) 复数张量
|
||
win_cpu = torch.hann_window(self.win_length, device="cpu", dtype=dtype)
|
||
|
||
wav_cpu = torch.istft(
|
||
spec_complex_cpu,
|
||
n_fft=self.filter_length,
|
||
hop_length=self.hop_length,
|
||
win_length=self.win_length,
|
||
window=win_cpu,
|
||
center=True,
|
||
normalized=False,
|
||
onesided=True,
|
||
)
|
||
return wav_cpu.to(device).unsqueeze(-2)
|
||
|
||
def _transform_no_complex(self, input_data):
|
||
"""
|
||
纯实数 STFT：return_complex=False，随后手动求幅度与相位
|
||
"""
|
||
z = torch.stft(
|
||
input_data,
|
||
n_fft=self.filter_length,
|
||
hop_length=self.hop_length,
|
||
win_length=self.win_length,
|
||
window=self.window.to(input_data.device, dtype=input_data.dtype),
|
||
return_complex=False,
|
||
center=True,
|
||
normalized=False,
|
||
)
|
||
real = z[..., 0]
|
||
imag = z[..., 1]
|
||
magnitude = torch.sqrt(real * real + imag * imag)
|
||
phase = torch.atan2(imag, real)
|
||
return magnitude, phase
|
||
|
||
# 替换 Kokoro 的 STFT.inverse 实现
|
||
_ist.TorchSTFT.inverse = _inverse_no_complex
|
||
_ist.TorchSTFT.transform = _transform_no_complex
|
||
|
||
|
||
repo_id = 'hexgrad/Kokoro-82M-v1.1-zh'
|
||
MODEL_SR = 24000
|
||
TARGET_SR = 16000
|
||
# How much silence to insert between paragraphs: 5000 is about 0.2 seconds
|
||
N_ZEROS = 20
|
||
model = None
|
||
en_empty_pipeline = None
|
||
en_voice = os.getenv('EN_VOICE', 'af_maple.pt')
|
||
zh_voice = os.getenv('ZH_VOICE', 'zf_046.pt')
|
||
model_dir = os.getenv('MODEL_DIR', '/model/hexgrad')
|
||
model_name = os.getenv('MODEL_NAME','kokoro-v1_1-zh.pth')
|
||
|
||
# model_1_1_dir = os.path.join(model_dir, 'Kokoro-82M-v1.1-zh')
|
||
# model_1_0_dir = os.path.join(model_dir, 'Kokoro-82M')
|
||
# repo_id_1_0 = 'hexgrad/Kokoro-82M'
|
||
|
||
@dataclass
|
||
class LanguagePipeline:
|
||
pipeline: KPipeline
|
||
voice_pt: str
|
||
|
||
pipeline_dict: dict[str, LanguagePipeline] = {}
|
||
|
||
def en_callable(text):
|
||
if text == 'Kokoro':
|
||
return 'kˈOkəɹO'
|
||
elif text == 'Sol':
|
||
return 'sˈOl'
|
||
return next(en_empty_pipeline(text)).phonemes
|
||
|
||
# HACK: Mitigate rushing caused by lack of training data beyond ~100 tokens
|
||
# Simple piecewise linear fn that decreases speed as len_ps increases
|
||
def speed_callable(len_ps):
|
||
speed = 0.8
|
||
if len_ps <= 83:
|
||
speed = 1
|
||
elif len_ps < 183:
|
||
speed = 1 - (len_ps - 83) / 500
|
||
return speed
|
||
|
||
|
||
# from https://huggingface.co/spaces/coqui/voice-chat-with-mistral/blob/main/app.py
|
||
def wave_header_chunk(frame_input=b"", channels=1, sample_width=2, sample_rate=32000):
|
||
# This will create a wave header then append the frame input
|
||
# It should be first on a streaming wav file
|
||
# Other frames better should not have it (else you will hear some artifacts each chunk start)
|
||
wav_buf = io.BytesIO()
|
||
with wave.open(wav_buf, "wb") as vfout:
|
||
vfout.setnchannels(channels)
|
||
vfout.setsampwidth(sample_width)
|
||
vfout.setframerate(sample_rate)
|
||
vfout.writeframes(frame_input)
|
||
|
||
wav_buf.seek(0)
|
||
return wav_buf.read()
|
||
|
||
|
||
def resample_audio(data: np.ndarray, original_rate: int, target_rate: int):
|
||
ori_dtype = data.dtype
|
||
# data = normalize_audio(data)
|
||
number_of_samples = int(len(data) * float(target_rate) / original_rate)
|
||
resampled_data = resample(data, number_of_samples)
|
||
# resampled_data = normalize_audio(resampled_data)
|
||
return resampled_data.astype(ori_dtype)
|
||
|
||
def audio_postprocess(data: np.ndarray, original_rate: int, target_rate: int):
|
||
audio = resample_audio(data, original_rate, target_rate)
|
||
if audio.dtype == np.float32:
|
||
audio = np.int16(audio * 32767)
|
||
audio = np.concatenate([audio, np.zeros(N_ZEROS, dtype=np.int16)])
|
||
return audio
|
||
|
||
def init():
|
||
global model, en_empty_pipeline
|
||
global model_1_0
|
||
global pipeline_dict
|
||
# device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
device = 'mlu'
|
||
model = KModel(repo_id=repo_id, model=os.path.join(model_dir, model_name), config=os.path.join(model_dir, 'config.json')).to(device).eval()
|
||
en_empty_pipeline = KPipeline(lang_code='a', repo_id=repo_id, model=False)
|
||
en_pipeline = KPipeline(lang_code='a', repo_id=repo_id, model=model)
|
||
zh_pipeline = KPipeline(lang_code='z', repo_id=repo_id, model=model, en_callable=en_callable)
|
||
en_voice_pt = os.path.join(model_dir, 'voices', en_voice)
|
||
zh_voice_pt = os.path.join(model_dir, 'voices', zh_voice)
|
||
pipeline_dict['zh'] = LanguagePipeline(pipeline=zh_pipeline, voice_pt=zh_voice_pt)
|
||
pipeline_dict['en'] = LanguagePipeline(pipeline=en_pipeline, voice_pt=en_voice_pt)
|
||
|
||
# v1.0 model for other languages
|
||
# model_1_0 = KModel(repo_id=repo_id_1_0, model=os.path.join(model_1_0_dir, 'kokoro-v1_0.pth'), config=os.path.join(model_1_0_dir, 'config.json')).to(device).eval()
|
||
# # es
|
||
# es_pipeline = KPipeline(lang_code='e', repo_id=repo_id_1_0, model=model_1_0)
|
||
# es_voice_pt = os.path.join(model_1_0_dir, 'voices', 'ef_dora.pt')
|
||
# pipeline_dict['es'] = LanguagePipeline(pipeline=es_pipeline, voice_pt=es_voice_pt)
|
||
# # fr
|
||
# fr_pipeline = KPipeline(lang_code='f', repo_id=repo_id_1_0, model=model_1_0)
|
||
# fr_voice_pt = os.path.join(model_1_0_dir, 'voices', 'ff_siwis.pt')
|
||
# pipeline_dict['fr'] = LanguagePipeline(pipeline=fr_pipeline, voice_pt=fr_voice_pt)
|
||
# # hi
|
||
# hi_pipeline = KPipeline(lang_code='h', repo_id=repo_id_1_0, model=model_1_0)
|
||
# hi_voice_pt = os.path.join(model_1_0_dir, 'voices', 'hf_alpha.pt')
|
||
# pipeline_dict['hi'] = LanguagePipeline(pipeline=hi_pipeline, voice_pt=hi_voice_pt)
|
||
# # it
|
||
# it_pipeline = KPipeline(lang_code='i', repo_id=repo_id_1_0, model=model_1_0)
|
||
# it_voice_pt = os.path.join(model_1_0_dir, 'voices', 'if_sara.pt')
|
||
# pipeline_dict['it'] = LanguagePipeline(pipeline=it_pipeline, voice_pt=it_voice_pt)
|
||
# # ja
|
||
# ja_pipeline = KPipeline(lang_code='j', repo_id=repo_id_1_0, model=model_1_0)
|
||
# ja_voice_pt = os.path.join(model_1_0_dir, 'voices', 'jf_alpha.pt')
|
||
# pipeline_dict['ja'] = LanguagePipeline(pipeline=ja_pipeline, voice_pt=ja_voice_pt)
|
||
# # pt
|
||
# pt_pipeline = KPipeline(lang_code='p', repo_id=repo_id_1_0, model=model_1_0)
|
||
# pt_voice_pt = os.path.join(model_1_0_dir, 'voices', 'pf_dora.pt')
|
||
# pipeline_dict['pt'] = LanguagePipeline(pipeline=pt_pipeline, voice_pt=pt_voice_pt)
|
||
|
||
warmup()
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
init()
|
||
yield
|
||
pass
|
||
|
||
app = FastAPI(lifespan=lifespan)
|
||
|
||
def warmup():
|
||
zh_pipeline = pipeline_dict['zh'].pipeline
|
||
voice = pipeline_dict['zh'].voice_pt
|
||
generator = zh_pipeline(text="语音合成测试TTS。", voice=voice, speed=speed_callable)
|
||
for _ in generator:
|
||
pass
|
||
|
||
xml_namespace = "{http://www.w3.org/XML/1998/namespace}"
|
||
|
||
symbols = ',.!?;:()[]{}<>，。！？；：【】《》……"“”_—'
|
||
def contains_words(text):
|
||
return any(char not in symbols for char in text)
|
||
|
||
def cut_sentences(text) -> list[str]:
|
||
text = text.strip()
|
||
splits = re.split(r"([.;?!、。？！；])", text)
|
||
sentences = []
|
||
for i in range(0, len(splits), 2):
|
||
if i + 1 < len(splits):
|
||
s = splits[i] + splits[i + 1]
|
||
else:
|
||
s = splits[i]
|
||
s = s.strip()
|
||
if s:
|
||
sentences.append(s)
|
||
return sentences
|
||
|
||
LANGUAGE_ALIASES = {
|
||
'z': 'zh',
|
||
'a': 'en',
|
||
'e': 'es',
|
||
'f': 'fr',
|
||
'h': 'hi',
|
||
'i': 'it',
|
||
'j': 'ja',
|
||
'p': 'pt',
|
||
}
|
||
|
||
@app.post("/")
|
||
@app.post("/tts")
|
||
def predict(ssml: str = Body(...), include_header: bool = False):
|
||
try:
|
||
root = ET.fromstring(ssml)
|
||
voice_element = root.find(".//voice")
|
||
if voice_element is not None:
|
||
transcription = voice_element.text.strip()
|
||
language = voice_element.get(f'{xml_namespace}lang', "zh").strip()
|
||
# voice_name = voice_element.get("name", "zh-f-soft-1").strip()
|
||
else:
|
||
return JSONResponse(status_code=400, content={"message": "Invalid SSML format: <voice> element not found."})
|
||
except ET.ParseError as e:
|
||
return JSONResponse(status_code=400, content={"message": "Invalid SSML format", "Exception": str(e)})
|
||
|
||
if not contains_words(transcription):
|
||
audio = np.zeros(N_ZEROS, dtype=np.int16).tobytes()
|
||
if include_header:
|
||
audio_header = wave_header_chunk(sample_rate=TARGET_SR)
|
||
audio = audio_header + audio
|
||
return Response(audio, media_type='audio/wav')
|
||
|
||
if language not in pipeline_dict:
|
||
if language in LANGUAGE_ALIASES:
|
||
language = LANGUAGE_ALIASES[language]
|
||
else:
|
||
return JSONResponse(status_code=400, content={"message": f"Language '{language}' not supported."})
|
||
|
||
|
||
def streaming_generator():
|
||
texts = cut_sentences(transcription)
|
||
has_yield = False
|
||
for text in texts:
|
||
if text.strip() and contains_words(text):
|
||
pipeline = pipeline_dict[language].pipeline
|
||
voice = pipeline_dict[language].voice_pt
|
||
if language == 'zh':
|
||
generator = pipeline(text=text, voice=voice, speed=speed_callable)
|
||
else:
|
||
generator = pipeline(text=text, voice=voice)
|
||
|
||
for (_, _, audio) in generator:
|
||
if include_header and not has_yield:
|
||
has_yield = True
|
||
yield wave_header_chunk(sample_rate=TARGET_SR)
|
||
yield audio_postprocess(audio.numpy(), MODEL_SR, TARGET_SR).tobytes()
|
||
|
||
return StreamingResponse(streaming_generator(), media_type='audio/wav')
|
||
|
||
|
||
@app.get("/health")
|
||
@app.get("/ready")
|
||
async def ready():
|
||
return JSONResponse(status_code=200, content={"status": "ok"})
|
||
|
||
@app.get("/health_check")
|
||
async def health_check():
|
||
try:
|
||
a = torch.ones(10, 20, dtype=torch.float32, device='cuda')
|
||
b = torch.ones(20, 10, dtype=torch.float32, device='cuda')
|
||
c = torch.matmul(a, b)
|
||
if c.sum() == 10 * 20 * 10:
|
||
return {"status": "ok"}
|
||
else:
|
||
raise HTTPException(status_code=503)
|
||
except Exception as e:
|
||
print(f'health_check failed')
|
||
raise HTTPException(status_code=503)
|
||
|
||
if __name__ == "__main__":
|
||
uvicorn.run(app, host="0.0.0.0", port=80)
|