Files
jarvis-models/src/blackbox/melotts.py
2024-05-23 15:51:35 +08:00

86 lines
2.8 KiB
Python

import io
import time
import requests
from fastapi import Request, Response, status
from fastapi.responses import JSONResponse
from injector import inject
from injector import singleton
from ..log.logging_time import logging_time
from ..configuration import MeloConf
from .blackbox import Blackbox
import soundfile
from melo.api import TTS
import logging
logger = logging.getLogger(__name__)
@singleton
class MeloTTS(Blackbox):
mode: str
url: str
speed: int
device: str
language: str
speaker: str
@logging_time(logger=logger)
def model_init(self, melo_config: MeloConf) -> None:
self.speed = melo_config.speed
self.device = melo_config.device
self.language = melo_config.language
self.speaker = melo_config.speaker
self.device = melo_config.device
self.url = ''
self.mode = melo_config.mode
self.melotts = None
self.speaker_ids = None
if self.mode == 'local':
self.melotts = TTS(language=self.language, device=self.device)
self.speaker_ids = self.melotts.hps.data.spk2id
else:
self.url = melo_config.url
logging.info('#### Initializing MeloTTS Service in ' + self.device + ' mode...')
@inject
def __init__(self, melo_config: MeloConf) -> None:
self.model_init(melo_config)
def __call__(self, *args, **kwargs):
return self.processing(*args, **kwargs)
def valid(self, *args, **kwargs) -> bool:
text = args[0]
return isinstance(text, str)
@logging_time(logger=logger)
def processing(self, *args, **kwargs) -> io.BytesIO | bytes:
text = args[0]
current_time = time.time()
if self.mode == 'local':
audio = self.melotts.tts_to_file(text, self.speaker_ids[self.speaker], speed=self.speed)
f = io.BytesIO()
soundfile.write(f, audio, 44100, format='wav')
f.seek(0)
print("#### MeloTTS Service consume - local : ", (time.time() - current_time))
return f.read()
else:
message = {
"text": text
}
response = requests.post(self.url, json=message)
print("#### MeloTTS Service consume - docker : ", (time.time()-current_time))
return response.content
async def fast_api_handler(self, request: Request) -> Response:
try:
data = await request.json()
except:
return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST)
text = data.get("text")
if text is None:
return JSONResponse(content={"error": "text is required"}, status_code=status.HTTP_400_BAD_REQUEST)
return Response(content=self.processing(text), media_type="audio/wav", headers={"Content-Disposition": "attachment; filename=audio.wav"})