Refactor: Use PathConfig and pathlib instead of paths.yml loading

This commit is contained in:
litagin02
2024-05-25 10:25:22 +09:00
parent 2771fbd209
commit 8976a3ed2f
22 changed files with 266 additions and 253 deletions

View File

@@ -1,14 +1,13 @@
import json
import os
from pathlib import Path
import gradio as gr
import numpy as np
import torch
import yaml
from safetensors import safe_open
from safetensors.torch import save_file
from config import get_path_config
from style_bert_vits2.constants import DEFAULT_STYLE, GRADIO_THEME
from style_bert_vits2.logging import logger
from style_bert_vits2.tts_model import TTSModel, TTSModelHolder
@@ -20,15 +19,17 @@ speech_style_keys = ["enc_p"]
tempo_keys = ["sdp", "dp"]
device = "cuda" if torch.cuda.is_available() else "cpu"
# Get path settings
with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f:
path_config: dict[str, str] = yaml.safe_load(f.read())
# dataset_root = path_config["dataset_root"]
assets_root = path_config["assets_root"]
path_config = get_path_config()
assets_root = path_config.assets_root
def merge_style(model_name_a, model_name_b, weight, output_name, style_triple_list):
def merge_style(
model_name_a: str,
model_name_b: str,
weight: float,
output_name: str,
style_triple_list: list[tuple[str, str, str]],
) -> tuple[Path, list[str]]:
"""
style_triple_list: list[(model_aでのスタイル名, model_bでのスタイル名, 出力するスタイル名)]
"""
@@ -41,18 +42,14 @@ def merge_style(model_name_a, model_name_b, weight, output_name, style_triple_li
raise ValueError(f"No element with {DEFAULT_STYLE} output style name found.")
style_vectors_a = np.load(
os.path.join(assets_root, model_name_a, "style_vectors.npy")
assets_root / model_name_a / "style_vectors.npy"
) # (style_num_a, 256)
style_vectors_b = np.load(
os.path.join(assets_root, model_name_b, "style_vectors.npy")
assets_root / model_name_b / "style_vectors.npy"
) # (style_num_b, 256)
with open(
os.path.join(assets_root, model_name_a, "config.json"), "r", encoding="utf-8"
) as f:
with open(assets_root / model_name_a / "config.json", "r", encoding="utf-8") as f:
config_a = json.load(f)
with open(
os.path.join(assets_root, model_name_b, "config.json"), "r", encoding="utf-8"
) as f:
with open(assets_root / model_name_b / "config.json", "r", encoding="utf-8") as f:
config_b = json.load(f)
style2id_a = config_a["data"]["style2id"]
style2id_b = config_b["data"]["style2id"]
@@ -73,21 +70,19 @@ def merge_style(model_name_a, model_name_b, weight, output_name, style_triple_li
new_style2id[style_out] = len(new_style_vecs) - 1
new_style_vecs = np.array(new_style_vecs)
output_style_path = os.path.join(assets_root, output_name, "style_vectors.npy")
output_style_path = assets_root / output_name / "style_vectors.npy"
np.save(output_style_path, new_style_vecs)
new_config = config_a.copy()
new_config["data"]["num_styles"] = len(new_style2id)
new_config["data"]["style2id"] = new_style2id
new_config["model_name"] = output_name
with open(
os.path.join(assets_root, output_name, "config.json"), "w", encoding="utf-8"
) as f:
with open(assets_root / output_name / "config.json", "w", encoding="utf-8") as f:
json.dump(new_config, f, indent=2, ensure_ascii=False)
# recipe.jsonを読み込んで、style_triple_listを追記
info_path = os.path.join(assets_root, output_name, "recipe.json")
if os.path.exists(info_path):
info_path = assets_root / output_name / "recipe.json"
if info_path.exists():
with open(info_path, "r", encoding="utf-8") as f:
info = json.load(f)
else:
@@ -99,11 +94,13 @@ def merge_style(model_name_a, model_name_b, weight, output_name, style_triple_li
return output_style_path, list(new_style2id.keys())
def lerp_tensors(t, v0, v1):
def lerp_tensors(t: float, v0: torch.Tensor, v1: torch.Tensor):
return v0 * (1 - t) + v1 * t
def slerp_tensors(t, v0, v1, dot_thres=0.998):
def slerp_tensors(
t: float, v0: torch.Tensor, v1: torch.Tensor, dot_thres: float = 0.998
):
device = v0.device
v0c = v0.cpu().numpy()
v1c = v1.cpu().numpy()
@@ -123,23 +120,23 @@ def slerp_tensors(t, v0, v1, dot_thres=0.998):
def merge_models(
model_path_a,
model_path_b,
voice_weight,
voice_pitch_weight,
speech_style_weight,
tempo_weight,
output_name,
use_slerp_instead_of_lerp,
model_path_a: str,
model_path_b: str,
voice_weight: float,
voice_pitch_weight: float,
speech_style_weight: float,
tempo_weight: float,
output_name: str,
use_slerp_instead_of_lerp: bool,
):
"""model Aを起点に、model Bの各要素を重み付けしてマージする。
safetensors形式を前提とする。"""
model_a_weight = {}
model_a_weight: dict[str, torch.Tensor] = {}
with safe_open(model_path_a, framework="pt", device="cpu") as f:
for k in f.keys():
model_a_weight[k] = f.get_tensor(k)
model_b_weight = {}
model_b_weight: dict[str, torch.Tensor] = {}
with safe_open(model_path_b, framework="pt", device="cpu") as f:
for k in f.keys():
model_b_weight[k] = f.get_tensor(k)
@@ -161,10 +158,8 @@ def merge_models(
slerp_tensors if use_slerp_instead_of_lerp else lerp_tensors
)(weight, model_a_weight[key], model_b_weight[key])
merged_model_path = os.path.join(
assets_root, output_name, f"{output_name}.safetensors"
)
os.makedirs(os.path.dirname(merged_model_path), exist_ok=True)
merged_model_path = assets_root / output_name / f"{output_name}.safetensors"
merged_model_path.parent.mkdir(parents=True, exist_ok=True)
save_file(merged_model_weight, merged_model_path)
info = {
@@ -175,24 +170,22 @@ def merge_models(
"speech_style_weight": speech_style_weight,
"tempo_weight": tempo_weight,
}
with open(
os.path.join(assets_root, output_name, "recipe.json"), "w", encoding="utf-8"
) as f:
with open(assets_root / output_name / "recipe.json", "w", encoding="utf-8") as f:
json.dump(info, f, indent=2, ensure_ascii=False)
return merged_model_path
def merge_models_gr(
model_name_a,
model_path_a,
model_name_b,
model_path_b,
output_name,
voice_weight,
voice_pitch_weight,
speech_style_weight,
tempo_weight,
use_slerp_instead_of_lerp,
model_name_a: str,
model_path_a: str,
model_name_b: str,
model_path_b: str,
output_name: str,
voice_weight: float,
voice_pitch_weight: float,
speech_style_weight: float,
tempo_weight: float,
use_slerp_instead_of_lerp: bool,
):
if output_name == "":
return "Error: 新しいモデル名を入力してください。"
@@ -210,10 +203,10 @@ def merge_models_gr(
def merge_style_gr(
model_name_a,
model_name_b,
weight,
output_name,
model_name_a: str,
model_name_b: str,
weight: float,
output_name: str,
style_triple_list_str: str,
):
if output_name == "":
@@ -245,12 +238,14 @@ def merge_style_gr(
)
def simple_tts(model_name, text, style=DEFAULT_STYLE, style_weight=1.0):
model_path = os.path.join(assets_root, model_name, f"{model_name}.safetensors")
config_path = os.path.join(assets_root, model_name, "config.json")
style_vec_path = os.path.join(assets_root, model_name, "style_vectors.npy")
def simple_tts(
model_name: str, text: str, style: str = DEFAULT_STYLE, style_weight: float = 1.0
):
model_path = assets_root / model_name / f"{model_name}.safetensors"
config_path = assets_root / model_name / "config.json"
style_vec_path = assets_root / model_name / "style_vectors.npy"
model = TTSModel(Path(model_path), Path(config_path), Path(style_vec_path), device)
model = TTSModel(model_path, config_path, style_vec_path, device)
return model.infer(text, style=style, style_weight=style_weight)
@@ -259,13 +254,13 @@ def update_two_model_names_dropdown(model_holder: TTSModelHolder):
return new_names, new_files, new_names, new_files
def load_styles_gr(model_name_a, model_name_b):
config_path_a = os.path.join(assets_root, model_name_a, "config.json")
def load_styles_gr(model_name_a: str, model_name_b: str):
config_path_a = assets_root / model_name_a / "config.json"
with open(config_path_a, "r", encoding="utf-8") as f:
config_a = json.load(f)
styles_a = list(config_a["data"]["style2id"].keys())
config_path_b = os.path.join(assets_root, model_name_b, "config.json")
config_path_b = assets_root / model_name_b / "config.json"
with open(config_path_b, "r", encoding="utf-8") as f:
config_b = json.load(f)
styles_b = list(config_b["data"]["style2id"].keys())
@@ -336,7 +331,10 @@ def create_merge_app(model_holder: TTSModelHolder) -> gr.Blocks:
)
return app
initial_id = 0
initial_model_files = model_holder.model_files_dict[model_names[initial_id]]
# initial_model_files = model_holder.model_files_dict[model_names[initial_id]]
initial_model_files = [
str(f) for f in model_holder.model_files_dict[model_names[initial_id]]
]
with gr.Blocks(theme=GRADIO_THEME) as app:
gr.Markdown(

View File

@@ -1,27 +1,23 @@
import json
import os
import shutil
from pathlib import Path
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import yaml
from scipy.spatial.distance import pdist, squareform
from sklearn.cluster import DBSCAN, AgglomerativeClustering, KMeans
from sklearn.manifold import TSNE
from umap import UMAP
from config import config
from config import get_path_config
from style_bert_vits2.constants import DEFAULT_STYLE, GRADIO_THEME
from style_bert_vits2.logging import logger
# Get path settings
with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f:
path_config: dict[str, str] = yaml.safe_load(f.read())
dataset_root = Path(path_config["dataset_root"])
# assets_root = path_config["assets_root"]
path_config = get_path_config()
dataset_root = path_config.dataset_root
assets_root = path_config.assets_root
MAX_CLUSTER_NUM = 10
MAX_AUDIO_NUM = 10
@@ -39,11 +35,7 @@ centroids = []
def load(model_name: str, reduction_method: str):
global wav_files, x, x_reduced, mean
# wavs_dir = os.path.join(dataset_root, model_name, "wavs")
wavs_dir = dataset_root / model_name / "wavs"
# style_vector_files = [
# os.path.join(wavs_dir, f) for f in os.listdir(wavs_dir) if f.endswith(".npy")
# ]
style_vector_files = [f for f in wavs_dir.rglob("*.npy") if f.is_file()]
# foo.wav.npy -> foo.wav
wav_files = [f.with_suffix("") for f in style_vector_files]
@@ -196,21 +188,21 @@ def do_clustering_gradio(n_clusters=4, method="KMeans"):
] * MAX_AUDIO_NUM
def save_style_vectors_from_clustering(model_name, style_names_str: str):
def save_style_vectors_from_clustering(model_name: str, style_names_str: str):
"""centerとcentroidsを保存する"""
result_dir = os.path.join(config.assets_root, model_name)
os.makedirs(result_dir, exist_ok=True)
result_dir = assets_root / model_name
result_dir.mkdir(parents=True, exist_ok=True)
style_vectors = np.stack([mean] + centroids)
style_vector_path = os.path.join(result_dir, "style_vectors.npy")
if os.path.exists(style_vector_path):
style_vector_path = result_dir / "style_vectors.npy"
if style_vector_path.exists():
logger.info(f"Backup {style_vector_path} to {style_vector_path}.bak")
shutil.copy(style_vector_path, f"{style_vector_path}.bak")
np.save(style_vector_path, style_vectors)
logger.success(f"Saved style vectors to {style_vector_path}")
# config.jsonの更新
config_path = os.path.join(result_dir, "config.json")
if not os.path.exists(config_path):
config_path = result_dir / "config.json"
if not config_path.exists():
return f"{config_path}が存在しません。"
style_names = [name.strip() for name in style_names_str.split(",")]
style_name_list = [DEFAULT_STYLE] + style_names
@@ -233,7 +225,7 @@ def save_style_vectors_from_clustering(model_name, style_names_str: str):
def save_style_vectors_from_files(
model_name, audio_files_str: str, style_names_str: str
model_name: str, audio_files_str: str, style_names_str: str
):
"""音声ファイルからスタイルベクトルを作成して保存する"""
global mean
@@ -241,8 +233,8 @@ def save_style_vectors_from_files(
return "Error: スタイルベクトルを読み込んでください。"
mean = np.mean(x, axis=0)
result_dir = os.path.join(config.assets_root, model_name)
os.makedirs(result_dir, exist_ok=True)
result_dir = assets_root / model_name
result_dir.mkdir(parents=True, exist_ok=True)
audio_files = [name.strip() for name in audio_files_str.split(",")]
style_names = [name.strip() for name in style_names_str.split(",")]
if len(audio_files) != len(style_names):
@@ -252,23 +244,23 @@ def save_style_vectors_from_files(
return "スタイル名が重複しています。"
style_vectors = [mean]
wavs_dir = os.path.join(dataset_root, model_name, "wavs")
wavs_dir = dataset_root / model_name / "wavs"
for audio_file in audio_files:
path = os.path.join(wavs_dir, audio_file)
if not os.path.exists(path):
path = wavs_dir / audio_file
if not path.exists():
return f"{path}が存在しません。"
style_vectors.append(np.load(f"{path}.npy"))
style_vectors = np.stack(style_vectors)
assert len(style_name_list) == len(style_vectors)
style_vector_path = os.path.join(result_dir, "style_vectors.npy")
if os.path.exists(style_vector_path):
style_vector_path = result_dir / "style_vectors.npy"
if style_vector_path.exists():
logger.info(f"Backup {style_vector_path} to {style_vector_path}.bak")
shutil.copy(style_vector_path, f"{style_vector_path}.bak")
np.save(style_vector_path, style_vectors)
# config.jsonの更新
config_path = os.path.join(result_dir, "config.json")
if not os.path.exists(config_path):
config_path = result_dir / "config.json"
if not config_path.exists():
return f"{config_path}が存在しません。"
logger.info(f"Backup {config_path} to {config_path}.bak")
shutil.copy(config_path, f"{config_path}.bak")

View File

@@ -1,5 +1,4 @@
import json
import os
import shutil
import socket
import subprocess
@@ -12,7 +11,8 @@ from pathlib import Path
import gradio as gr
import yaml
from dataclasses import dataclass
from config import get_path_config
from style_bert_vits2.logging import logger
from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
from style_bert_vits2.utils.subprocess import run_script_with_log, second_elem_of
@@ -21,20 +21,27 @@ from style_bert_vits2.utils.subprocess import run_script_with_log, second_elem_o
logger_handler = None
tensorboard_executed = False
# Get path settings
with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f:
path_config: dict[str, str] = yaml.safe_load(f.read())
dataset_root = Path(path_config["dataset_root"])
path_config = get_path_config()
dataset_root = path_config.dataset_root
def get_path(model_name: str) -> tuple[Path, Path, Path, Path, Path]:
@dataclass
class PathsForPreprocess:
dataset_path: Path
esd_path: Path
train_path: Path
val_path: Path
config_path: Path
def get_path(model_name: str) -> PathsForPreprocess:
assert model_name != "", "モデル名は空にできません"
dataset_path = dataset_root / model_name
lbl_path = dataset_path / "esd.list"
esd_path = dataset_path / "esd.list"
train_path = dataset_path / "train.list"
val_path = dataset_path / "val.list"
config_path = dataset_path / "config.json"
return dataset_path, lbl_path, train_path, val_path, config_path
return PathsForPreprocess(dataset_path, esd_path, train_path, val_path, config_path)
def initialize(
@@ -51,14 +58,14 @@ def initialize(
log_interval: int,
):
global logger_handler
dataset_path, _, train_path, val_path, config_path = get_path(model_name)
paths = get_path(model_name)
# 前処理のログをファイルに保存する
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
file_name = f"preprocess_{timestamp}.log"
if logger_handler is not None:
logger.remove(logger_handler)
logger_handler = logger.add(os.path.join(dataset_path, file_name))
logger_handler = logger.add(paths.dataset_path / file_name)
logger.info(
f"Step 1: start initialization...\nmodel_name: {model_name}, batch_size: {batch_size}, epochs: {epochs}, save_every_steps: {save_every_steps}, freeze_ZH_bert: {freeze_ZH_bert}, freeze_JP_bert: {freeze_JP_bert}, freeze_EN_bert: {freeze_EN_bert}, freeze_style: {freeze_style}, freeze_decoder: {freeze_decoder}, use_jp_extra: {use_jp_extra}"
@@ -71,8 +78,8 @@ def initialize(
with open(default_config_path, "r", encoding="utf-8") as f:
config = json.load(f)
config["model_name"] = model_name
config["data"]["training_files"] = str(train_path)
config["data"]["validation_files"] = str(val_path)
config["data"]["training_files"] = str(paths.train_path)
config["data"]["validation_files"] = str(paths.val_path)
config["train"]["batch_size"] = batch_size
config["train"]["epochs"] = epochs
config["train"]["eval_interval"] = save_every_steps
@@ -89,14 +96,14 @@ def initialize(
# 今はデフォルトであるが、以前は非JP-Extra版になくバグの原因になるので念のため
config["data"]["use_jp_extra"] = use_jp_extra
model_path = dataset_path / "models"
model_path = paths.dataset_path / "models"
if model_path.exists():
logger.warning(
f"Step 1: {model_path} already exists, so copy it to backup to {model_path}_backup"
)
shutil.copytree(
src=model_path,
dst=dataset_path / "models_backup",
dst=paths.dataset_path / "models_backup",
dirs_exist_ok=True,
)
shutil.rmtree(model_path)
@@ -110,14 +117,14 @@ def initialize(
logger.error(f"Step 1: {pretrained_dir} folder not found.")
return False, f"Step 1, Error: {pretrained_dir}フォルダが見つかりません。"
with open(config_path, "w", encoding="utf-8") as f:
with open(paths.config_path, "w", encoding="utf-8") as f:
json.dump(config, f, indent=2, ensure_ascii=False)
if not Path("config.yml").exists():
shutil.copy(src="default_config.yml", dst="config.yml")
with open("config.yml", "r", encoding="utf-8") as f:
yml_data = yaml.safe_load(f)
yml_data["model_name"] = model_name
yml_data["dataset_path"] = str(dataset_path)
yml_data["dataset_path"] = str(paths.dataset_path)
with open("config.yml", "w", encoding="utf-8") as f:
yaml.dump(yml_data, f, allow_unicode=True)
logger.success("Step 1: initialization finished.")
@@ -126,7 +133,7 @@ def initialize(
def resample(model_name: str, normalize: bool, trim: bool, num_processes: int):
logger.info("Step 2: start resampling...")
dataset_path, _, _, _, _ = get_path(model_name)
dataset_path = get_path(model_name).dataset_path
input_dir = dataset_path / "raw"
output_dir = dataset_path / "wavs"
cmd = [
@@ -159,21 +166,24 @@ def preprocess_text(
model_name: str, use_jp_extra: bool, val_per_lang: int, yomi_error: str
):
logger.info("Step 3: start preprocessing text...")
_, lbl_path, train_path, val_path, config_path = get_path(model_name)
if not lbl_path.exists():
logger.error(f"Step 3: {lbl_path} not found.")
return False, f"Step 3, Error: 書き起こしファイル {lbl_path} が見つかりません。"
paths = get_path(model_name)
if not paths.esd_path.exists():
logger.error(f"Step 3: {paths.esd_path} not found.")
return (
False,
f"Step 3, Error: 書き起こしファイル {paths.esd_path} が見つかりません。",
)
cmd = [
"preprocess_text.py",
"--config-path",
str(config_path),
str(paths.config_path),
"--transcription-path",
str(lbl_path),
str(paths.esd_path),
"--train-path",
str(train_path),
str(paths.train_path),
"--val-path",
str(val_path),
str(paths.val_path),
"--val-per-lang",
str(val_per_lang),
"--yomi_error",
@@ -201,7 +211,7 @@ def preprocess_text(
def bert_gen(model_name: str):
logger.info("Step 4: start bert_gen...")
_, _, _, _, config_path = get_path(model_name)
config_path = get_path(model_name).config_path
success, message = run_script_with_log(
["bert_gen.py", "--config", str(config_path)]
)
@@ -220,7 +230,7 @@ def bert_gen(model_name: str):
def style_gen(model_name: str, num_processes: int):
logger.info("Step 5: start style_gen...")
_, _, _, _, config_path = get_path(model_name)
config_path = get_path(model_name).config_path
success, message = run_script_with_log(
[
"style_gen.py",
@@ -319,17 +329,23 @@ def train(
use_jp_extra: bool = True,
speedup: bool = False,
):
dataset_path, _, _, _, config_path = get_path(model_name)
paths = get_path(model_name)
# 学習再開の場合を考えて念のためconfig.ymlの名前等を更新
with open("config.yml", "r", encoding="utf-8") as f:
yml_data = yaml.safe_load(f)
yml_data["model_name"] = model_name
yml_data["dataset_path"] = str(dataset_path)
yml_data["dataset_path"] = str(paths.dataset_path)
with open("config.yml", "w", encoding="utf-8") as f:
yaml.dump(yml_data, f, allow_unicode=True)
train_py = "train_ms.py" if not use_jp_extra else "train_ms_jp_extra.py"
cmd = [train_py, "--config", str(config_path), "--model", str(dataset_path)]
cmd = [
train_py,
"--config",
str(paths.config_path),
"--model",
str(paths.dataset_path),
]
if skip_style:
cmd.append("--skip_default_style")
if speedup: