Feat: save default style, and colab train support (maybe)

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}")

make_default_style.py

@@ -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)

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)

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")