import numpy as np
import gradio as gr
import torch
import os
import warnings
from gradio.processing_utils import convert_to_16_bit_wav
from typing import Dict, List, Optional, Union
import utils
from infer import get_net_g, infer
from models import SynthesizerTrn
from models_jp_extra import SynthesizerTrn as SynthesizerTrnJPExtra
from .log import logger
from .constants import (
DEFAULT_ASSIST_TEXT_WEIGHT,
DEFAULT_LENGTH,
DEFAULT_LINE_SPLIT,
DEFAULT_NOISE,
DEFAULT_NOISEW,
DEFAULT_SDP_RATIO,
DEFAULT_SPLIT_INTERVAL,
DEFAULT_STYLE,
DEFAULT_STYLE_WEIGHT,
)
class Model:
def __init__(
self, model_path: str, config_path: str, style_vec_path: str, device: str
):
self.model_path: str = model_path
self.config_path: str = config_path
self.device: str = device
self.style_vec_path: str = style_vec_path
self.hps: utils.HParams = utils.get_hparams_from_file(self.config_path)
self.spk2id: Dict[str, int] = self.hps.data.spk2id
self.id2spk: Dict[int, str] = {v: k for k, v in self.spk2id.items()}
self.num_styles: int = self.hps.data.num_styles
if hasattr(self.hps.data, "style2id"):
self.style2id: Dict[str, int] = self.hps.data.style2id
else:
self.style2id: Dict[str, int] = {str(i): i for i in range(self.num_styles)}
if len(self.style2id) != self.num_styles:
raise ValueError(
f"Number of styles ({self.num_styles}) does not match the number of style2id ({len(self.style2id)})"
)
self.style_vectors: np.ndarray = np.load(self.style_vec_path)
if self.style_vectors.shape[0] != self.num_styles:
raise ValueError(
f"The number of styles ({self.num_styles}) does not match the number of style vectors ({self.style_vectors.shape[0]})"
)
self.net_g: Union[SynthesizerTrn, SynthesizerTrnJPExtra, None] = None
def load_net_g(self):
self.net_g = get_net_g(
model_path=self.model_path,
version=self.hps.version,
device=self.device,
hps=self.hps,
)
def get_style_vector(self, style_id: int, weight: float = 1.0) -> np.ndarray:
mean = self.style_vectors[0]
style_vec = self.style_vectors[style_id]
style_vec = mean + (style_vec - mean) * weight
return style_vec
def get_style_vector_from_audio(
self, audio_path: str, weight: float = 1.0
) -> np.ndarray:
from style_gen import get_style_vector
xvec = get_style_vector(audio_path)
mean = self.style_vectors[0]
xvec = mean + (xvec - mean) * weight
return xvec
def infer(
self,
text: str,
language: str = "JP",
sid: int = 0,
reference_audio_path: Optional[str] = None,
sdp_ratio: float = DEFAULT_SDP_RATIO,
noise: float = DEFAULT_NOISE,
noisew: float = DEFAULT_NOISEW,
length: float = DEFAULT_LENGTH,
line_split: bool = DEFAULT_LINE_SPLIT,
split_interval: float = DEFAULT_SPLIT_INTERVAL,
assist_text: Optional[str] = None,
assist_text_weight: float = DEFAULT_ASSIST_TEXT_WEIGHT,
use_assist_text: bool = False,
style: str = DEFAULT_STYLE,
style_weight: float = DEFAULT_STYLE_WEIGHT,
given_tone: Optional[list[int]] = None,
) -> tuple[int, np.ndarray]:
logger.info(f"Start generating audio data from text:\n{text}")
if language != "JP" and self.hps.version.endswith("JP-Extra"):
raise ValueError(
"The model is trained with JP-Extra, but the language is not JP"
)
if reference_audio_path == "":
reference_audio_path = None
if assist_text == "" or not use_assist_text:
assist_text = None
if self.net_g is None:
self.load_net_g()
if reference_audio_path is None:
style_id = self.style2id[style]
style_vector = self.get_style_vector(style_id, style_weight)
else:
style_vector = self.get_style_vector_from_audio(
reference_audio_path, style_weight
)
if not line_split:
with torch.no_grad():
audio = infer(
text=text,
sdp_ratio=sdp_ratio,
noise_scale=noise,
noise_scale_w=noisew,
length_scale=length,
sid=sid,
language=language,
hps=self.hps,
net_g=self.net_g,
device=self.device,
assist_text=assist_text,
assist_text_weight=assist_text_weight,
style_vec=style_vector,
given_tone=given_tone,
)
else:
texts = text.split("\n")
texts = [t for t in texts if t != ""]
audios = []
with torch.no_grad():
for i, t in enumerate(texts):
audios.append(
infer(
text=t,
sdp_ratio=sdp_ratio,
noise_scale=noise,
noise_scale_w=noisew,
length_scale=length,
sid=sid,
language=language,
hps=self.hps,
net_g=self.net_g,
device=self.device,
assist_text=assist_text,
assist_text_weight=assist_text_weight,
style_vec=style_vector,
)
)
if i != len(texts) - 1:
audios.append(np.zeros(int(44100 * split_interval)))
audio = np.concatenate(audios)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
audio = convert_to_16_bit_wav(audio)
logger.info("Audio data generated successfully")
return (self.hps.data.sampling_rate, audio)
class ModelHolder:
def __init__(self, root_dir: str, device: str):
self.root_dir: str = root_dir
self.device: str = device
self.model_files_dict: Dict[str, List[str]] = {}
self.current_model: Optional[Model] = None
self.model_names: List[str] = []
self.models: List[Model] = []
self.refresh()
def refresh(self):
self.model_files_dict = {}
self.model_names = []
self.current_model = None
model_dirs = [
d
for d in os.listdir(self.root_dir)
if os.path.isdir(os.path.join(self.root_dir, d))
]
for model_name in model_dirs:
model_dir = os.path.join(self.root_dir, model_name)
model_files = [
os.path.join(model_dir, f)
for f in os.listdir(model_dir)
if f.endswith(".pth") or f.endswith(".pt") or f.endswith(".safetensors")
]
if len(model_files) == 0:
logger.warning(
f"No model files found in {self.root_dir}/{model_name}, so skip it"
)
continue
self.model_files_dict[model_name] = model_files
self.model_names.append(model_name)
def load_model_gr(
self, model_name: str, model_path: str
) -> tuple[gr.Dropdown, gr.Button, gr.Dropdown]:
if model_name not in self.model_files_dict:
raise ValueError(f"Model `{model_name}` is not found")
if model_path not in self.model_files_dict[model_name]:
raise ValueError(f"Model file `{model_path}` is not found")
if (
self.current_model is not None
and self.current_model.model_path == model_path
):
# Already loaded
speakers = list(self.current_model.spk2id.keys())
styles = list(self.current_model.style2id.keys())
return (
gr.Dropdown(choices=styles, value=styles[0]),
gr.Button(interactive=True, value="音声合成"),
gr.Dropdown(choices=speakers, value=speakers[0]),
)
self.current_model = Model(
model_path=model_path,
config_path=os.path.join(self.root_dir, model_name, "config.json"),
style_vec_path=os.path.join(self.root_dir, model_name, "style_vectors.npy"),
device=self.device,
)
speakers = list(self.current_model.spk2id.keys())
styles = list(self.current_model.style2id.keys())
return (
gr.Dropdown(choices=styles, value=styles[0]),
gr.Button(interactive=True, value="音声合成"),
gr.Dropdown(choices=speakers, value=speakers[0]),
)
def update_model_files_gr(self, model_name: str) -> gr.Dropdown:
model_files = self.model_files_dict[model_name]
return gr.Dropdown(choices=model_files, value=model_files[0])
def update_model_names_gr(self) -> tuple[gr.Dropdown, gr.Dropdown, gr.Button]:
self.refresh()
initial_model_name = self.model_names[0]
initial_model_files = self.model_files_dict[initial_model_name]
return (
gr.Dropdown(choices=self.model_names, value=initial_model_name),
gr.Dropdown(choices=initial_model_files, value=initial_model_files[0]),
gr.Button(interactive=False), # For tts_button
)