228 lines
7.0 KiB
Python
228 lines
7.0 KiB
Python
import asyncio
|
|
import logging
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import torch
|
|
from wyoming.audio import AudioChunk, AudioStart, AudioStop
|
|
from wyoming.error import Error
|
|
from wyoming.event import Event
|
|
from wyoming.info import Describe, Info
|
|
from wyoming.server import AsyncEventHandler
|
|
from wyoming.tts import Synthesize
|
|
|
|
from style_bert_vits2.nlp import bert_models
|
|
from style_bert_vits2.constants import Languages
|
|
from style_bert_vits2.tts_model import TTSModel
|
|
|
|
_LOGGER = logging.getLogger(__name__)
|
|
|
|
_VOICE_LOCK = asyncio.Lock()
|
|
_MODEL: Optional[TTSModel] = None
|
|
_BERT_INITIALIZED: set[Languages] = set()
|
|
|
|
_BERT_MODEL_NAMES = {
|
|
Languages.JP: "ku-nlp/deberta-v2-large-japanese-char-wwm",
|
|
Languages.EN: "microsoft/deberta-v3-large",
|
|
Languages.ZH: "hfl/chinese-roberta-wwm-ext-large",
|
|
}
|
|
|
|
_HIRAGANA_KATAKANA = re.compile(r"[\u3040-\u309F\u30A0-\u30FF]")
|
|
_CJK = re.compile(r"[\u4E00-\u9FFF]")
|
|
|
|
|
|
def _optimize_gpu():
|
|
if not torch.cuda.is_available():
|
|
return
|
|
|
|
torch.cuda.empty_cache = lambda: None
|
|
_LOGGER.info("Disabled torch.cuda.empty_cache() to prevent re-allocation on every inference")
|
|
|
|
if torch.version.hip is not None and hasattr(torch.backends, 'miopen'):
|
|
torch.backends.miopen.benchmark = True
|
|
_LOGGER.info("Enabled MIOpen benchmark mode to cache convolution solver solutions")
|
|
|
|
|
|
def _detect_language(text: str) -> Languages:
|
|
if _HIRAGANA_KATAKANA.search(text):
|
|
return Languages.JP
|
|
if _CJK.search(text):
|
|
return Languages.ZH
|
|
return Languages.EN
|
|
|
|
|
|
def _load_bert_for_language(language: Languages, device: str, half: bool = False) -> None:
|
|
if language in _BERT_INITIALIZED:
|
|
return
|
|
|
|
model_name = _BERT_MODEL_NAMES[language]
|
|
_LOGGER.info("Loading BERT model for %s (%s)", language.name, model_name)
|
|
bert_models.load_model(language, model_name)
|
|
bert_models.load_tokenizer(language, model_name)
|
|
bert = bert_models.__loaded_models.get(language)
|
|
if bert is not None:
|
|
bert.eval()
|
|
bert.to(device).float()
|
|
_LOGGER.info("BERT model for %s moved to %s", language.name, device)
|
|
|
|
_BERT_INITIALIZED.add(language)
|
|
|
|
|
|
def _find_model_files(model_dir: Path):
|
|
model_dir = model_dir.resolve()
|
|
safetensors = list(model_dir.glob("*.safetensors"))
|
|
config = model_dir / "config.json"
|
|
style = model_dir / "style_vectors.npy"
|
|
|
|
if safetensors and config.exists():
|
|
return safetensors[0], config, style if style.exists() else None
|
|
|
|
for subdir in sorted(model_dir.iterdir()):
|
|
if not subdir.is_dir():
|
|
continue
|
|
safetensors = list(subdir.glob("*.safetensors"))
|
|
config = subdir / "config.json"
|
|
style = subdir / "style_vectors.npy"
|
|
if safetensors and config.exists():
|
|
return safetensors[0], config, style if style.exists() else None
|
|
|
|
raise FileNotFoundError(
|
|
f"No .safetensors files found in {model_dir} or its subdirectories"
|
|
)
|
|
|
|
|
|
def _load_model(model_dir: Path, device: str, half: bool = False) -> TTSModel:
|
|
model_path, config_path, style_path = _find_model_files(model_dir)
|
|
|
|
_LOGGER.info("Creating TTSModel (model=%s, config=%s, device=%s)",
|
|
model_path.name, config_path.name, device)
|
|
|
|
model = TTSModel(
|
|
model_path=model_path,
|
|
config_path=config_path,
|
|
style_vec_path=style_path,
|
|
device=device,
|
|
)
|
|
|
|
_LOGGER.info("Loading model weights...")
|
|
model.load()
|
|
|
|
net_g = getattr(model, "_TTSModel__net_g", None)
|
|
if net_g is not None:
|
|
net_g = net_g.float()
|
|
setattr(model, "_TTSModel__net_g", net_g)
|
|
|
|
_LOGGER.info("Model loaded successfully")
|
|
return model
|
|
|
|
|
|
class GLaDOSEventHandler(AsyncEventHandler):
|
|
def __init__(
|
|
self,
|
|
wyoming_info: Info,
|
|
model_dir: Path,
|
|
device: str,
|
|
*args,
|
|
half: bool = False,
|
|
preload: bool = False,
|
|
**kwargs,
|
|
) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
self.wyoming_info_event = wyoming_info.event()
|
|
self.model_dir = model_dir
|
|
self.device = device
|
|
self.half = half
|
|
|
|
_optimize_gpu()
|
|
|
|
if preload:
|
|
_LOGGER.info("Pre-loading model at startup...")
|
|
global _MODEL
|
|
_MODEL = _load_model(model_dir, device, half)
|
|
|
|
async def handle_event(self, event: Event) -> bool:
|
|
if Describe.is_type(event.type):
|
|
await self.write_event(self.wyoming_info_event)
|
|
return True
|
|
|
|
if not Synthesize.is_type(event.type):
|
|
return True
|
|
|
|
synthesize = Synthesize.from_event(event)
|
|
return await self._handle_synthesize(synthesize)
|
|
|
|
async def _handle_synthesize(self, synthesize: Synthesize) -> bool:
|
|
global _MODEL
|
|
|
|
text = synthesize.text.strip()
|
|
if not text:
|
|
return True
|
|
|
|
language = _detect_language(text)
|
|
|
|
speaker_id = 0
|
|
style = "Neutral"
|
|
if synthesize.voice is not None and synthesize.voice.speaker:
|
|
try:
|
|
speaker_id = int(synthesize.voice.speaker)
|
|
except ValueError:
|
|
pass
|
|
|
|
_LOGGER.info("Synthesizing: text='%s' language=%s speaker=%s style=%s",
|
|
text[:80], language.name, speaker_id, style)
|
|
|
|
try:
|
|
async with _VOICE_LOCK:
|
|
if _MODEL is None:
|
|
_LOGGER.info("Loading GLaDOS model from %s on %s",
|
|
self.model_dir, self.device)
|
|
_MODEL = _load_model(self.model_dir, self.device, self.half)
|
|
|
|
_load_bert_for_language(language, self.device, self.half)
|
|
|
|
sr, audio = await asyncio.to_thread(
|
|
_MODEL.infer,
|
|
text=text,
|
|
language=language,
|
|
speaker_id=speaker_id,
|
|
style=style,
|
|
)
|
|
|
|
audio_int16 = np.round(audio).astype(np.int16)
|
|
raw_bytes = audio_int16.tobytes()
|
|
|
|
rate = sr
|
|
width = 2
|
|
channels = 1
|
|
|
|
await self.write_event(
|
|
AudioStart(rate=rate, width=width, channels=channels).event()
|
|
)
|
|
|
|
samples_per_chunk = 1024
|
|
bytes_per_sample = width * channels
|
|
bytes_per_chunk = bytes_per_sample * samples_per_chunk
|
|
|
|
for i in range(0, len(raw_bytes), bytes_per_chunk):
|
|
chunk = raw_bytes[i:i + bytes_per_chunk]
|
|
await self.write_event(
|
|
AudioChunk(
|
|
audio=chunk,
|
|
rate=rate,
|
|
width=width,
|
|
channels=channels,
|
|
).event()
|
|
)
|
|
|
|
await self.write_event(AudioStop().event())
|
|
return True
|
|
|
|
except Exception as err:
|
|
_LOGGER.exception("Synthesis failed")
|
|
await self.write_event(
|
|
Error(text=str(err), code=err.__class__.__name__).event()
|
|
)
|
|
return True
|