Merge branch 'colab' into dev

app.py

@@ -360,6 +360,9 @@ if __name__ == "__main__":
     parser.add_argument(
         "--dir", "-d", type=str, help="Model directory", default=config.out_dir
     )
+    parser.add_argument(
+        "--share", action="store_true", help="Share this app publicly", default=False
+    )
     args = parser.parse_args()
     model_dir = args.dir
 
@@ -518,4 +521,4 @@ if __name__ == "__main__":
             outputs=[style, ref_audio_path],
         )
 
-    app.launch(inbrowser=True)
+    app.launch(inbrowser=True, share=args.share)

bert_gen.py

@@ -1,5 +1,4 @@
 import argparse
-import sys
 from multiprocessing import Pool
 
 import torch
@@ -10,6 +9,7 @@ import commons
 import utils
 from config import config
 from text import cleaned_text_to_sequence, get_bert
+from tools.stdout_wrapper import SAFE_STDOUT
 
 
 def process_line(x):
@@ -76,7 +76,7 @@ if __name__ == "__main__":
             for _ in tqdm(
                 pool.imap_unordered(process_line, zip(lines, add_blank)),
                 total=len(lines),
-                file=sys.stdout,
+                file=SAFE_STDOUT,
             ):
                 # 这里是缩进的代码块，表示循环体
                 pass  # 使用pass语句作为占位符

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

preprocess_text.py

@@ -1,6 +1,5 @@
 import json
 import os
-import sys
 from collections import defaultdict
 from random import shuffle
 from typing import Optional
@@ -10,6 +9,7 @@ from tqdm import tqdm
 
 from config import config
 from text.cleaner import clean_text
+from tools.stdout_wrapper import SAFE_STDOUT
 
 preprocess_text_config = config.preprocess_text_config
 
@@ -52,7 +52,7 @@ def preprocess(
                 lines = trans_file.readlines()
                 # print(lines, ' ', len(lines))
                 if len(lines) != 0:
-                    for line in tqdm(lines, file=sys.stdout):
+                    for line in tqdm(lines, file=SAFE_STDOUT):
                         try:
                             utt, spk, language, text = line.strip().split("|")
                             norm_text, phones, tones, word2ph = clean_text(

resample.py

@@ -1,6 +1,5 @@
 import argparse
 import os
-import sys
 from multiprocessing import Pool, cpu_count
 
 import librosa
@@ -10,6 +9,7 @@ from tqdm import tqdm
 
 from config import config
 from tools.log import logger
+from tools.stdout_wrapper import SAFE_STDOUT
 
 
 def normalize_audio(data, sr):
@@ -97,7 +97,7 @@ if __name__ == "__main__":
 
     pool = Pool(processes=processes)
     for _ in tqdm(
-        pool.imap_unordered(process, tasks), file=sys.stdout, total=len(tasks)
+        pool.imap_unordered(process, tasks), file=SAFE_STDOUT, total=len(tasks)
     ):
         pass
 

slice.py

@@ -1,12 +1,13 @@
 import argparse
 import os
 import shutil
-import sys
 
 import soundfile as sf
 import torch
 from tqdm import tqdm
 
+from tools.stdout_wrapper import SAFE_STDOUT
+
 vad_model, utils = torch.hub.load(
     repo_or_dir="snakers4/silero-vad",
     model="silero_vad",
@@ -106,7 +107,7 @@ if __name__ == "__main__":
         shutil.rmtree(output_dir)
 
     total_sec = 0
-    for wav_file in tqdm(wav_files, file=sys.stdout):
+    for wav_file in tqdm(wav_files, file=SAFE_STDOUT):
         time_sec = split_wav(
             wav_file,
             output_dir,

style_gen.py

@@ -1,6 +1,5 @@
 import argparse
 import concurrent.futures
-import sys
 import warnings
 
 import numpy as np
@@ -9,6 +8,7 @@ from tqdm import tqdm
 
 import utils
 from config import config
+from tools.stdout_wrapper import SAFE_STDOUT
 
 warnings.filterwarnings("ignore", category=UserWarning)
 from pyannote.audio import Inference, Model
@@ -25,8 +25,13 @@ def extract_style_vector(wav_path):
 
 def save_style_vector(wav_path):
     style_vec = extract_style_vector(wav_path)
-    # `test.wav` -> `test.wav.npy`
+    np.save(f"{wav_path}.npy", style_vec)  # `test.wav` -> `test.wav.npy`
-    np.save(f"{wav_path}.npy", style_vec)
+    return style_vec
+
+
+def save_average_style_vector(style_vectors, filename="style_vectors.npy"):
+    average_vector = np.mean(style_vectors, axis=0)
+    np.save(filename, average_vector)
 
 
 if __name__ == "__main__":
@@ -59,7 +64,7 @@ if __name__ == "__main__":
             tqdm(
                 executor.map(save_style_vector, wavnames),
                 total=len(wavnames),
-                file=sys.stdout,
+                file=SAFE_STDOUT,
             )
         )
 

tools/log.py

@@ -2,8 +2,8 @@
 logger封装
 """
 from loguru import logger
-import sys
 
+from .stdout_wrapper import SAFE_STDOUT
 
 # 移除所有默认的处理器
 logger.remove()
@@ -13,4 +13,4 @@ log_format = (
     "<g>{time:MM-DD HH:mm:ss}</g> |<lvl>{level:^8}</lvl>| {file}:{line} | {message}"
 )
 
-logger.add(sys.stdout, format=log_format, backtrace=True, diagnose=True)
+logger.add(SAFE_STDOUT, format=log_format, backtrace=True, diagnose=True)

tools/stdout_wrapper.py

@@ -0,0 +1,34 @@
+import sys
+import tempfile
+
+
+class StdoutWrapper:
+    def __init__(self):
+        self.temp_file = tempfile.NamedTemporaryFile(mode="w+", delete=False)
+        self.original_stdout = sys.stdout
+
+    def write(self, message: str):
+        self.temp_file.write(message)
+        self.temp_file.flush()
+        print(message, end="", file=self.original_stdout)
+
+    def flush(self):
+        self.temp_file.flush()
+
+    def read(self):
+        self.temp_file.seek(0)
+        return self.temp_file.read()
+
+    def close(self):
+        self.temp_file.close()
+
+    def fileno(self):
+        return self.temp_file.fileno()
+
+
+try:
+    import google.colab
+
+    SAFE_STDOUT = StdoutWrapper()
+except ImportError:
+    SAFE_STDOUT = sys.stdout

tools/subprocess_utils.py

@@ -2,6 +2,7 @@ import subprocess
 import sys
 
 from .log import logger
+from .stdout_wrapper import SAFE_STDOUT
 
 python = sys.executable
 
@@ -10,7 +11,7 @@ def run_script_with_log(cmd: list[str]) -> tuple[bool, str]:
     logger.info(f"Running: {' '.join(cmd)}")
     result = subprocess.run(
         [python] + cmd,
-        stdout=sys.stdout,
+        stdout=SAFE_STDOUT,  # type: ignore
         stderr=subprocess.PIPE,
         text=True,
     )

train_ms.py

@@ -17,6 +17,7 @@ from tqdm import tqdm
 
 # logging.getLogger("numba").setLevel(logging.WARNING)
 import commons
+import default_style
 import utils
 from config import config
 from data_utils import (
@@ -29,6 +30,7 @@ from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
 from models import DurationDiscriminator, MultiPeriodDiscriminator, SynthesizerTrn
 from text.symbols import symbols
 from tools.log import logger
+from tools.stdout_wrapper import SAFE_STDOUT
 
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = (
@@ -41,6 +43,14 @@ 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)
+
+try:
+    import google.colab
+
+    IS_COLAB = True
+except ImportError:
+    IS_COLAB = False
+
 global_step = 0
 
 
@@ -106,9 +116,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)
@@ -427,7 +467,7 @@ def train_and_evaluate(
         ja_bert,
         en_bert,
         style_vec,
-    ) in enumerate(tqdm(train_loader, file=sys.stdout)):
+    ) in enumerate(tqdm(train_loader, file=SAFE_STDOUT)):
         if net_g.module.use_noise_scaled_mas:
             current_mas_noise_scale = (
                 net_g.module.mas_noise_scale_initial

transcribe.py

@@ -5,6 +5,8 @@ import sys
 from faster_whisper import WhisperModel
 from tqdm import tqdm
 
+from tools.stdout_wrapper import SAFE_STDOUT
+
 
 def transcribe(wav_path, initial_prompt=None):
     segments, _ = model.transcribe(
@@ -45,7 +47,7 @@ if __name__ == "__main__":
         os.rename(output_file, output_file + ".bak")
 
     with open(output_file, "w", encoding="utf-8") as f:
-        for wav_file in tqdm(wav_files, file=sys.stdout):
+        for wav_file in tqdm(wav_files, file=SAFE_STDOUT):
             file_name = os.path.basename(wav_file)
             text = transcribe(wav_file, initial_prompt=initial_prompt)
             f.write(f"{file_name}|{speaker_name}|JP|{text}\n")

webui_train.py

@@ -10,10 +10,24 @@ import yaml
 from tools.log import logger
 from tools.subprocess_utils import run_script_with_log, second_elem_of
 
+try:
+    import google.colab
+
+    IS_COLAB = True
+except ImportError:
+    IS_COLAB = False
+
 
 def get_path(model_name):
     assert model_name != "", "モデル名は空にできません"
-    dataset_path = os.path.join("Data", model_name)
+    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")
     train_path = os.path.join(dataset_path, "train.list")
     val_path = os.path.join(dataset_path, "val.list")
@@ -39,9 +53,12 @@ def initialize(model_name, batch_size, epochs, save_every_steps, bf16_run):
 
     model_path = os.path.join(dataset_path, "models")
     try:
-        shutil.copytree(src="pretrained", dst=model_path)
+        shutil.copytree(
+            src="pretrained",
+            dst=model_path,
+        )
     except FileExistsError:
-        logger.error(f"Step 1: {model_path} already exists.")
+        logger.warning(f"Step 1: {model_path} already exists.")
         return False, f"Step1, Error: モデルフォルダ {model_path} が既に存在します。問題なければ削除してください。"
     except FileNotFoundError:
         logger.error("Step 1: `pretrained` folder not found.")