diff --git a/app.py b/app.py index 09ced91..5f270c8 100644 --- a/app.py +++ b/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) diff --git a/bert_gen.py b/bert_gen.py index d72ce4e..e41484d 100644 --- a/bert_gen.py +++ b/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语句作为占位符 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/preprocess_text.py b/preprocess_text.py index 38683c8..d905f19 100644 --- a/preprocess_text.py +++ b/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( diff --git a/resample.py b/resample.py index 28a2991..4ece450 100644 --- a/resample.py +++ b/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 diff --git a/slice.py b/slice.py index c4faf0a..3cd2399 100644 --- a/slice.py +++ b/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, diff --git a/style_gen.py b/style_gen.py index 7d7f65e..6fdf9e3 100644 --- a/style_gen.py +++ b/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) + np.save(f"{wav_path}.npy", style_vec) # `test.wav` -> `test.wav.npy` + 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, ) ) diff --git a/tools/log.py b/tools/log.py index 85526cb..51dca5f 100644 --- a/tools/log.py +++ b/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 = ( "{time:MM-DD HH:mm:ss} |{level:^8}| {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) diff --git a/tools/stdout_wrapper.py b/tools/stdout_wrapper.py new file mode 100644 index 0000000..23c6e76 --- /dev/null +++ b/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 diff --git a/tools/subprocess_utils.py b/tools/subprocess_utils.py index dea2e0d..ad15575 100644 --- a/tools/subprocess_utils.py +++ b/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, ) diff --git a/train_ms.py b/train_ms.py index 7d170ad..4fe4857 100644 --- a/train_ms.py +++ b/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 diff --git a/transcribe.py b/transcribe.py index 4f2b955..7d2dd1a 100644 --- a/transcribe.py +++ b/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") diff --git a/webui_train.py b/webui_train.py index 7df212a..acc49fd 100644 --- a/webui_train.py +++ b/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.")