diff --git a/default_style.py b/default_style.py new file mode 100644 index 0000000..424f8eb --- /dev/null +++ b/default_style.py @@ -0,0 +1,29 @@ +import os +from tools.log import logger + +import numpy as np +import json + + +def set_style_config(json_path, output_path): + with open(json_path, "r") as f: + json_dict = json.load(f) + json_dict["data"]["num_styles"] = 1 + json_dict["data"]["style2id"] = {"Neutral": 0} + with open(output_path, "w") as f: + json.dump(json_dict, f, indent=2) + logger.info(f"Update style config (only Neutral style) to {output_path}") + + +def save_mean_vector(wav_dir, output_path): + embs = [] + for file in os.listdir(wav_dir): + if file.endswith(".npy"): + xvec = np.load(os.path.join(wav_dir, file)) + embs.append(np.expand_dims(xvec, axis=0)) + + x = np.concatenate(embs, axis=0) # (N, 256) + mean = np.mean(x, axis=0) # (256,) + only_mean = np.stack([mean]) # (1, 256) + np.save(output_path, only_mean) + logger.info(f"Saved mean style vector to {output_path}") diff --git a/make_default_style.py b/make_default_style.py deleted file mode 100644 index ca708f6..0000000 --- a/make_default_style.py +++ /dev/null @@ -1,17 +0,0 @@ -import os -import numpy as np -import argparse - -parser = argparse.ArgumentParser() -parser.add_argument("--wav_dir", type=str, default="data/wav") - -embs = [] -names = [] -for file in os.listdir(wav_dir): - if file.endswith(".npy"): - xvec = np.load(os.path.join(wav_dir, file)) - embs.append(np.expand_dims(xvec, axis=0)) - names.append(file) - -x = np.concatenate(embs, axis=0) -x = np.squeeze(x) diff --git a/train_ms.py b/train_ms.py index 92177bb..cb49a82 100644 --- a/train_ms.py +++ b/train_ms.py @@ -4,6 +4,7 @@ import gc import os import platform import shutil +import sys import torch import torch.distributed as dist @@ -29,6 +30,7 @@ from models import DurationDiscriminator, MultiPeriodDiscriminator, SynthesizerT from text.symbols import symbols from tools.log import logger from tools.stdout_wrapper import get_stdout +import default_style torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = ( @@ -41,6 +43,9 @@ torch.backends.cuda.enable_mem_efficient_sdp( True ) # Not available if torch version is lower than 2.0 torch.backends.cuda.enable_math_sdp(True) + +IS_COLAB = "google.colab" in sys.modules + global_step = 0 @@ -106,9 +111,39 @@ def run(): data = f.read() with open(config.train_ms_config.config_path, "w", encoding="utf-8") as f: f.write(data) + + """ + Path constants are a bit complicated... + TODO: Refactor or rename these? + (Both `config.yml` and `config.json` are used, which is confusing I think.) + + args.model: For saving all info needed for training. + default: `Data/{model_name}`. + hps.model_dir = model_dir: For saving checkpoints (for resuming training). + default: `Data/{model_name}/models`. + + config.out_dir: Root directory of model assets needed for inference. + default: `model_assets`. + + out_dir: For saving resulting models (for inference). + default: `model_assets/{model_name}`, which is used for inference. + """ + if IS_COLAB: + config.out_dir = "/content/drive/MyDrive/Style-Bert-VITS2/model_assets" + logger.info( + "Colab detected, so use mounted Google Drive as directory for saving resulting models:" + ) + logger.info(config.out_dir) + os.makedirs(config.out_dir, exist_ok=True) out_dir = os.path.join(config.out_dir, config.model_name) os.makedirs(out_dir, exist_ok=True) - shutil.copy(args.config, os.path.join(out_dir, "config.json")) + + # Save default style to out_dir + default_style.set_style_config(args.config, os.path.join(out_dir, "config.json")) + default_style.save_mean_vector( + os.path.join(args.model, "wavs"), + os.path.join(out_dir, "style_vectors.npy"), + ) torch.manual_seed(hps.train.seed) torch.cuda.set_device(local_rank) diff --git a/webui_train.py b/webui_train.py index fb894cd..db72bdd 100644 --- a/webui_train.py +++ b/webui_train.py @@ -10,16 +10,17 @@ import yaml from tools.log import logger from tools.subprocess_utils import run_script_with_log, second_elem_of - -is_colab = "google.colab" in sys.modules +IS_COLAB = "google.colab" in sys.modules def get_path(model_name): assert model_name != "", "モデル名は空にできません" - if is_colab: + if IS_COLAB: + logger.info("Colab detected, so use mounted Google Drive as dataset path:") dataset_path = os.path.join( "/content/drive/MyDrive/Style-Bert-VITS2/Data", model_name ) + logger.info(dataset_path) else: dataset_path = os.path.join("Data", model_name) lbl_path = os.path.join(dataset_path, "esd.list")