diff --git a/style_bert_vits2/models/infer.py b/style_bert_vits2/models/infer.py index 93bf8c2..d3ec963 100644 --- a/style_bert_vits2/models/infer.py +++ b/style_bert_vits2/models/infer.py @@ -80,9 +80,9 @@ def get_net_g(model_path: str, version: str, device: str, hps: HyperParameters): net_g.state_dict() _ = net_g.eval() if model_path.endswith(".pth") or model_path.endswith(".pt"): - _ = utils.load_checkpoint(model_path, net_g, None, skip_optimizer=True) + _ = utils.checkpoints.load_checkpoint(model_path, net_g, None, skip_optimizer=True) elif model_path.endswith(".safetensors"): - _ = utils.load_safetensors(model_path, net_g, True) + _ = utils.safetensors.load_safetensors(model_path, net_g, True) else: raise ValueError(f"Unknown model format: {model_path}") return net_g diff --git a/style_bert_vits2/models/utils.py b/style_bert_vits2/models/utils.py deleted file mode 100644 index 37ea269..0000000 --- a/style_bert_vits2/models/utils.py +++ /dev/null @@ -1,357 +0,0 @@ -import argparse -import glob -import json -import logging -import os -import re -import subprocess - -import numpy as np -import torch -from safetensors import safe_open -from safetensors.torch import save_file -from scipy.io.wavfile import read - -from style_bert_vits2.logging import logger - - -MATPLOTLIB_FLAG = False - - -def load_checkpoint( - checkpoint_path, model, optimizer=None, skip_optimizer=False, for_infer=False -): - assert os.path.isfile(checkpoint_path) - checkpoint_dict = torch.load(checkpoint_path, map_location="cpu") - iteration = checkpoint_dict["iteration"] - learning_rate = checkpoint_dict["learning_rate"] - logger.info( - f"Loading model and optimizer at iteration {iteration} from {checkpoint_path}" - ) - if ( - optimizer is not None - and not skip_optimizer - and checkpoint_dict["optimizer"] is not None - ): - optimizer.load_state_dict(checkpoint_dict["optimizer"]) - elif optimizer is None and not skip_optimizer: - # else: Disable this line if Infer and resume checkpoint,then enable the line upper - new_opt_dict = optimizer.state_dict() - new_opt_dict_params = new_opt_dict["param_groups"][0]["params"] - new_opt_dict["param_groups"] = checkpoint_dict["optimizer"]["param_groups"] - new_opt_dict["param_groups"][0]["params"] = new_opt_dict_params - optimizer.load_state_dict(new_opt_dict) - - saved_state_dict = checkpoint_dict["model"] - if hasattr(model, "module"): - state_dict = model.module.state_dict() - else: - state_dict = model.state_dict() - - new_state_dict = {} - for k, v in state_dict.items(): - try: - # assert "emb_g" not in k - new_state_dict[k] = saved_state_dict[k] - assert saved_state_dict[k].shape == v.shape, ( - saved_state_dict[k].shape, - v.shape, - ) - except: - # For upgrading from the old version - if "ja_bert_proj" in k: - v = torch.zeros_like(v) - logger.warning( - f"Seems you are using the old version of the model, the {k} is automatically set to zero for backward compatibility" - ) - elif "enc_q" in k and for_infer: - continue - else: - logger.error(f"{k} is not in the checkpoint {checkpoint_path}") - - new_state_dict[k] = v - - if hasattr(model, "module"): - model.module.load_state_dict(new_state_dict, strict=False) - else: - model.load_state_dict(new_state_dict, strict=False) - - logger.info("Loaded '{}' (iteration {})".format(checkpoint_path, iteration)) - - return model, optimizer, learning_rate, iteration - - -def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path): - logger.info( - "Saving model and optimizer state at iteration {} to {}".format( - iteration, checkpoint_path - ) - ) - if hasattr(model, "module"): - state_dict = model.module.state_dict() - else: - state_dict = model.state_dict() - torch.save( - { - "model": state_dict, - "iteration": iteration, - "optimizer": optimizer.state_dict(), - "learning_rate": learning_rate, - }, - checkpoint_path, - ) - - -def clean_checkpoints(path_to_models="logs/44k/", n_ckpts_to_keep=2, sort_by_time=True): - """Freeing up space by deleting saved ckpts - - Arguments: - path_to_models -- Path to the model directory - n_ckpts_to_keep -- Number of ckpts to keep, excluding G_0.pth and D_0.pth - sort_by_time -- True -> chronologically delete ckpts - False -> lexicographically delete ckpts - """ - import re - - ckpts_files = [ - f - for f in os.listdir(path_to_models) - if os.path.isfile(os.path.join(path_to_models, f)) - ] - - def name_key(_f): - return int(re.compile("._(\\d+)\\.pth").match(_f).group(1)) - - def time_key(_f): - return os.path.getmtime(os.path.join(path_to_models, _f)) - - sort_key = time_key if sort_by_time else name_key - - def x_sorted(_x): - return sorted( - [f for f in ckpts_files if f.startswith(_x) and not f.endswith("_0.pth")], - key=sort_key, - ) - - to_del = [ - os.path.join(path_to_models, fn) - for fn in ( - x_sorted("G_")[:-n_ckpts_to_keep] - + x_sorted("D_")[:-n_ckpts_to_keep] - + x_sorted("WD_")[:-n_ckpts_to_keep] - + x_sorted("DUR_")[:-n_ckpts_to_keep] - ) - ] - - def del_info(fn): - return logger.info(f"Free up space by deleting ckpt {fn}") - - def del_routine(x): - return [os.remove(x), del_info(x)] - - [del_routine(fn) for fn in to_del] - - -def load_safetensors(checkpoint_path, model, for_infer=False): - """ - Load safetensors model. - """ - - tensors = {} - iteration = None - with safe_open(checkpoint_path, framework="pt", device="cpu") as f: - for key in f.keys(): - if key == "iteration": - iteration = f.get_tensor(key).item() - tensors[key] = f.get_tensor(key) - if hasattr(model, "module"): - result = model.module.load_state_dict(tensors, strict=False) - else: - result = model.load_state_dict(tensors, strict=False) - for key in result.missing_keys: - if key.startswith("enc_q") and for_infer: - continue - logger.warning(f"Missing key: {key}") - for key in result.unexpected_keys: - if key == "iteration": - continue - logger.warning(f"Unexpected key: {key}") - if iteration is None: - logger.info(f"Loaded '{checkpoint_path}'") - else: - logger.info(f"Loaded '{checkpoint_path}' (iteration {iteration})") - return model, iteration - - -def save_safetensors(model, iteration, checkpoint_path, is_half=False, for_infer=False): - """ - Save model with safetensors. - """ - if hasattr(model, "module"): - state_dict = model.module.state_dict() - else: - state_dict = model.state_dict() - keys = [] - for k in state_dict: - if "enc_q" in k and for_infer: - continue # noqa: E701 - keys.append(k) - - new_dict = ( - {k: state_dict[k].half() for k in keys} - if is_half - else {k: state_dict[k] for k in keys} - ) - new_dict["iteration"] = torch.LongTensor([iteration]) - logger.info(f"Saved safetensors to {checkpoint_path}") - save_file(new_dict, checkpoint_path) - - -def summarize( - writer, - global_step, - scalars={}, - histograms={}, - images={}, - audios={}, - audio_sampling_rate=22050, -): - for k, v in scalars.items(): - writer.add_scalar(k, v, global_step) - for k, v in histograms.items(): - writer.add_histogram(k, v, global_step) - for k, v in images.items(): - writer.add_image(k, v, global_step, dataformats="HWC") - for k, v in audios.items(): - writer.add_audio(k, v, global_step, audio_sampling_rate) - - -def is_resuming(dir_path): - # JP-ExtraバージョンではDURがなくWDがあったり変わるため、Gのみで判断する - g_list = glob.glob(os.path.join(dir_path, "G_*.pth")) - # d_list = glob.glob(os.path.join(dir_path, "D_*.pth")) - # dur_list = glob.glob(os.path.join(dir_path, "DUR_*.pth")) - return len(g_list) > 0 - - -def latest_checkpoint_path(dir_path, regex="G_*.pth"): - f_list = glob.glob(os.path.join(dir_path, regex)) - f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f)))) - try: - x = f_list[-1] - except IndexError: - raise ValueError(f"No checkpoint found in {dir_path} with regex {regex}") - return x - - -def plot_spectrogram_to_numpy(spectrogram): - global MATPLOTLIB_FLAG - if not MATPLOTLIB_FLAG: - import matplotlib - - matplotlib.use("Agg") - MATPLOTLIB_FLAG = True - mpl_logger = logging.getLogger("matplotlib") - mpl_logger.setLevel(logging.WARNING) - import matplotlib.pylab as plt - import numpy as np - - fig, ax = plt.subplots(figsize=(10, 2)) - im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none") - plt.colorbar(im, ax=ax) - plt.xlabel("Frames") - plt.ylabel("Channels") - plt.tight_layout() - - fig.canvas.draw() - data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="") - data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) - plt.close() - return data - - -def plot_alignment_to_numpy(alignment, info=None): - global MATPLOTLIB_FLAG - if not MATPLOTLIB_FLAG: - import matplotlib - - matplotlib.use("Agg") - MATPLOTLIB_FLAG = True - mpl_logger = logging.getLogger("matplotlib") - mpl_logger.setLevel(logging.WARNING) - import matplotlib.pylab as plt - import numpy as np - - fig, ax = plt.subplots(figsize=(6, 4)) - im = ax.imshow( - alignment.transpose(), aspect="auto", origin="lower", interpolation="none" - ) - fig.colorbar(im, ax=ax) - xlabel = "Decoder timestep" - if info is not None: - xlabel += "\n\n" + info - plt.xlabel(xlabel) - plt.ylabel("Encoder timestep") - plt.tight_layout() - - fig.canvas.draw() - data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="") - data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) - plt.close() - return data - - -def load_wav_to_torch(full_path): - sampling_rate, data = read(full_path) - return torch.FloatTensor(data.astype(np.float32)), sampling_rate - - -def load_filepaths_and_text(filename, split="|"): - with open(filename, encoding="utf-8") as f: - filepaths_and_text = [line.strip().split(split) for line in f] - return filepaths_and_text - - -def get_logger(model_dir, filename="train.log"): - global logger - logger = logging.getLogger(os.path.basename(model_dir)) - logger.setLevel(logging.DEBUG) - - formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s") - if not os.path.exists(model_dir): - os.makedirs(model_dir) - h = logging.FileHandler(os.path.join(model_dir, filename)) - h.setLevel(logging.DEBUG) - h.setFormatter(formatter) - logger.addHandler(h) - return logger - - -def get_steps(model_path): - matches = re.findall(r"\d+", model_path) - return matches[-1] if matches else None - - -def check_git_hash(model_dir): - source_dir = os.path.dirname(os.path.realpath(__file__)) - if not os.path.exists(os.path.join(source_dir, ".git")): - logger.warning( - "{} is not a git repository, therefore hash value comparison will be ignored.".format( - source_dir - ) - ) - return - - cur_hash = subprocess.getoutput("git rev-parse HEAD") - - path = os.path.join(model_dir, "githash") - if os.path.exists(path): - saved_hash = open(path).read() - if saved_hash != cur_hash: - logger.warning( - "git hash values are different. {}(saved) != {}(current)".format( - saved_hash[:8], cur_hash[:8] - ) - ) - else: - open(path, "w").write(cur_hash) diff --git a/style_bert_vits2/models/utils/__init__.py b/style_bert_vits2/models/utils/__init__.py new file mode 100644 index 0000000..f488e09 --- /dev/null +++ b/style_bert_vits2/models/utils/__init__.py @@ -0,0 +1,156 @@ +import glob +import logging +import os +import re +import subprocess + +import numpy as np +import torch +from scipy.io.wavfile import read + +from style_bert_vits2.logging import logger +from style_bert_vits2.models.utils import checkpoints # type: ignore +from style_bert_vits2.models.utils import safetensors # type: ignore + + +MATPLOTLIB_FLAG = False + + +def summarize( + writer, + global_step, + scalars={}, + histograms={}, + images={}, + audios={}, + audio_sampling_rate=22050, +): + for k, v in scalars.items(): + writer.add_scalar(k, v, global_step) + for k, v in histograms.items(): + writer.add_histogram(k, v, global_step) + for k, v in images.items(): + writer.add_image(k, v, global_step, dataformats="HWC") + for k, v in audios.items(): + writer.add_audio(k, v, global_step, audio_sampling_rate) + + +def is_resuming(dir_path): + # JP-ExtraバージョンではDURがなくWDがあったり変わるため、Gのみで判断する + g_list = glob.glob(os.path.join(dir_path, "G_*.pth")) + # d_list = glob.glob(os.path.join(dir_path, "D_*.pth")) + # dur_list = glob.glob(os.path.join(dir_path, "DUR_*.pth")) + return len(g_list) > 0 + + +def plot_spectrogram_to_numpy(spectrogram): + global MATPLOTLIB_FLAG + if not MATPLOTLIB_FLAG: + import matplotlib + + matplotlib.use("Agg") + MATPLOTLIB_FLAG = True + mpl_logger = logging.getLogger("matplotlib") + mpl_logger.setLevel(logging.WARNING) + import matplotlib.pylab as plt + import numpy as np + + fig, ax = plt.subplots(figsize=(10, 2)) + im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none") + plt.colorbar(im, ax=ax) + plt.xlabel("Frames") + plt.ylabel("Channels") + plt.tight_layout() + + fig.canvas.draw() + data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="") + data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) + plt.close() + return data + + +def plot_alignment_to_numpy(alignment, info=None): + global MATPLOTLIB_FLAG + if not MATPLOTLIB_FLAG: + import matplotlib + + matplotlib.use("Agg") + MATPLOTLIB_FLAG = True + mpl_logger = logging.getLogger("matplotlib") + mpl_logger.setLevel(logging.WARNING) + import matplotlib.pylab as plt + import numpy as np + + fig, ax = plt.subplots(figsize=(6, 4)) + im = ax.imshow( + alignment.transpose(), aspect="auto", origin="lower", interpolation="none" + ) + fig.colorbar(im, ax=ax) + xlabel = "Decoder timestep" + if info is not None: + xlabel += "\n\n" + info + plt.xlabel(xlabel) + plt.ylabel("Encoder timestep") + plt.tight_layout() + + fig.canvas.draw() + data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="") + data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) + plt.close() + return data + + +def load_wav_to_torch(full_path): + sampling_rate, data = read(full_path) + return torch.FloatTensor(data.astype(np.float32)), sampling_rate + + +def load_filepaths_and_text(filename, split="|"): + with open(filename, encoding="utf-8") as f: + filepaths_and_text = [line.strip().split(split) for line in f] + return filepaths_and_text + + +def get_logger(model_dir, filename="train.log"): + global logger + logger = logging.getLogger(os.path.basename(model_dir)) + logger.setLevel(logging.DEBUG) + + formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s") + if not os.path.exists(model_dir): + os.makedirs(model_dir) + h = logging.FileHandler(os.path.join(model_dir, filename)) + h.setLevel(logging.DEBUG) + h.setFormatter(formatter) + logger.addHandler(h) + return logger + + +def get_steps(model_path): + matches = re.findall(r"\d+", model_path) + return matches[-1] if matches else None + + +def check_git_hash(model_dir): + source_dir = os.path.dirname(os.path.realpath(__file__)) + if not os.path.exists(os.path.join(source_dir, ".git")): + logger.warning( + "{} is not a git repository, therefore hash value comparison will be ignored.".format( + source_dir + ) + ) + return + + cur_hash = subprocess.getoutput("git rev-parse HEAD") + + path = os.path.join(model_dir, "githash") + if os.path.exists(path): + saved_hash = open(path).read() + if saved_hash != cur_hash: + logger.warning( + "git hash values are different. {}(saved) != {}(current)".format( + saved_hash[:8], cur_hash[:8] + ) + ) + else: + open(path, "w").write(cur_hash) diff --git a/style_bert_vits2/models/utils/checkpoints.py b/style_bert_vits2/models/utils/checkpoints.py new file mode 100644 index 0000000..f26f8fc --- /dev/null +++ b/style_bert_vits2/models/utils/checkpoints.py @@ -0,0 +1,194 @@ +import glob +import os +import re +from pathlib import Path +from typing import Any, Optional, Union + +import torch + +from style_bert_vits2.logging import logger + + +def load_checkpoint( + checkpoint_path: Union[str, Path], + model: torch.nn.Module, + optimizer: Optional[torch.optim.Optimizer] = None, + skip_optimizer: bool = False, + for_infer: bool = False +) -> tuple[torch.nn.Module, Optional[torch.optim.Optimizer], float, int]: + """ + 指定されたパスからチェックポイントを読み込み、モデルとオプティマイザーを更新する。 + + Args: + checkpoint_path (Union[str, Path]): チェックポイントファイルのパス + model (torch.nn.Module): 更新するモデル + optimizer (Optional[torch.optim.Optimizer]): 更新するオプティマイザー。None の場合は更新しない + skip_optimizer (bool): オプティマイザーの更新をスキップするかどうかのフラグ + for_infer (bool): 推論用に読み込むかどうかのフラグ + + Returns: + tuple[torch.nn.Module, Optional[torch.optim.Optimizer], float, int]: 更新されたモデルとオプティマイザー、学習率、イテレーション番号 + """ + + assert os.path.isfile(checkpoint_path) + checkpoint_dict = torch.load(checkpoint_path, map_location="cpu") + iteration = checkpoint_dict["iteration"] + learning_rate = checkpoint_dict["learning_rate"] + logger.info( + f"Loading model and optimizer at iteration {iteration} from {checkpoint_path}" + ) + if ( + optimizer is not None + and not skip_optimizer + and checkpoint_dict["optimizer"] is not None + ): + optimizer.load_state_dict(checkpoint_dict["optimizer"]) + elif optimizer is None and not skip_optimizer: + # else: Disable this line if Infer and resume checkpoint,then enable the line upper + new_opt_dict = optimizer.state_dict() # type: ignore + new_opt_dict_params = new_opt_dict["param_groups"][0]["params"] + new_opt_dict["param_groups"] = checkpoint_dict["optimizer"]["param_groups"] + new_opt_dict["param_groups"][0]["params"] = new_opt_dict_params + optimizer.load_state_dict(new_opt_dict) # type: ignore + + saved_state_dict = checkpoint_dict["model"] + if hasattr(model, "module"): + state_dict = model.module.state_dict() + else: + state_dict = model.state_dict() + + new_state_dict = {} + for k, v in state_dict.items(): + try: + # assert "emb_g" not in k + new_state_dict[k] = saved_state_dict[k] + assert saved_state_dict[k].shape == v.shape, ( + saved_state_dict[k].shape, + v.shape, + ) + except: + # For upgrading from the old version + if "ja_bert_proj" in k: + v = torch.zeros_like(v) + logger.warning( + f"Seems you are using the old version of the model, the {k} is automatically set to zero for backward compatibility" + ) + elif "enc_q" in k and for_infer: + continue + else: + logger.error(f"{k} is not in the checkpoint {checkpoint_path}") + + new_state_dict[k] = v + + if hasattr(model, "module"): + model.module.load_state_dict(new_state_dict, strict=False) + else: + model.load_state_dict(new_state_dict, strict=False) + + logger.info("Loaded '{}' (iteration {})".format(checkpoint_path, iteration)) + + return model, optimizer, learning_rate, iteration + + +def save_checkpoint( + model: torch.nn.Module, + optimizer: Union[torch.optim.Optimizer, torch.optim.AdamW], + learning_rate: float, + iteration: int, + checkpoint_path: Union[str, Path], +) -> None: + """ + モデルとオプティマイザーの状態を指定されたパスに保存する。 + + Args: + model (torch.nn.Module): 保存するモデル + optimizer (Union[torch.optim.Optimizer, torch.optim.AdamW]): 保存するオプティマイザー + learning_rate (float): 学習率 + iteration (int): イテレーション数 + checkpoint_path (Union[str, Path]): 保存先のパス + """ + logger.info(f"Saving model and optimizer state at iteration {iteration} to {checkpoint_path}") + if hasattr(model, "module"): + state_dict = model.module.state_dict() + else: + state_dict = model.state_dict() + torch.save( + { + "model": state_dict, + "iteration": iteration, + "optimizer": optimizer.state_dict(), + "learning_rate": learning_rate, + }, + checkpoint_path, + ) + + +def clean_checkpoints(model_dir_path: Union[str, Path] = "logs/44k/", n_ckpts_to_keep: int = 2, sort_by_time: bool = True) -> None: + """ + 指定されたディレクトリから古いチェックポイントを削除して空き容量を確保する + + Args: + model_dir_path (Union[str, Path]): モデルが保存されているディレクトリのパス + n_ckpts_to_keep (int): 保持するチェックポイントの数（G_0.pth と D_0.pth を除く） + sort_by_time (bool): True の場合、時間順に削除。False の場合、名前順に削除 + """ + + ckpts_files = [ + f + for f in os.listdir(model_dir_path) + if os.path.isfile(os.path.join(model_dir_path, f)) + ] + + def name_key(_f: str) -> int: + return int(re.compile("._(\\d+)\\.pth").match(_f).group(1)) # type: ignore + + def time_key(_f: str) -> float: + return os.path.getmtime(os.path.join(model_dir_path, _f)) + + sort_key = time_key if sort_by_time else name_key + + def x_sorted(_x: str) -> list[str]: + return sorted( + [f for f in ckpts_files if f.startswith(_x) and not f.endswith("_0.pth")], + key=sort_key, + ) + + to_del = [ + os.path.join(model_dir_path, fn) + for fn in ( + x_sorted("G_")[:-n_ckpts_to_keep] + + x_sorted("D_")[:-n_ckpts_to_keep] + + x_sorted("WD_")[:-n_ckpts_to_keep] + + x_sorted("DUR_")[:-n_ckpts_to_keep] + ) + ] + + def del_info(fn: str) -> None: + return logger.info(f"Free up space by deleting ckpt {fn}") + + def del_routine(x: str) -> list[Any]: + return [os.remove(x), del_info(x)] + + [del_routine(fn) for fn in to_del] + + +def get_latest_checkpoint_path(model_dir_path: Union[str, Path], regex: str = "G_*.pth") -> str: + """ + 指定されたディレクトリから最新のチェックポイントのパスを取得する + + Args: + model_dir_path (Union[str, Path]): モデルが保存されているディレクトリのパス + regex (str): チェックポイントのファイル名の正規表現 + + Returns: + str: 最新のチェックポイントのパス + """ + + f_list = glob.glob(os.path.join(str(model_dir_path), regex)) + f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f)))) + try: + x = f_list[-1] + except IndexError: + raise ValueError(f"No checkpoint found in {model_dir_path} with regex {regex}") + + return x diff --git a/style_bert_vits2/models/utils/safetensors.py b/style_bert_vits2/models/utils/safetensors.py new file mode 100644 index 0000000..8917c77 --- /dev/null +++ b/style_bert_vits2/models/utils/safetensors.py @@ -0,0 +1,91 @@ +from pathlib import Path +from typing import Any, Optional, Union + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from style_bert_vits2.logging import logger + + +def load_safetensors( + checkpoint_path: Union[str, Path], + model: torch.nn.Module, + for_infer: bool = False, +) -> tuple[torch.nn.Module, Optional[int]]: + """ + 指定されたパスから safetensors モデルを読み込み、モデルとイテレーションを返す。 + + Args: + checkpoint_path (Union[str, Path]): モデルのチェックポイントファイルのパス + model (torch.nn.Module): 読み込む対象のモデル + for_infer (bool): 推論用に読み込むかどうかのフラグ + + Returns: + tuple[torch.nn.Module, Optional[int]]: 読み込まれたモデルとイテレーション番号（存在する場合） + """ + + tensors: dict[str, Any] = {} + iteration: Optional[int] = None + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as f: # type: ignore + for key in f.keys(): + if key == "iteration": + iteration = f.get_tensor(key).item() + tensors[key] = f.get_tensor(key) + if hasattr(model, "module"): + result = model.module.load_state_dict(tensors, strict=False) + else: + result = model.load_state_dict(tensors, strict=False) + for key in result.missing_keys: + if key.startswith("enc_q") and for_infer: + continue + logger.warning(f"Missing key: {key}") + for key in result.unexpected_keys: + if key == "iteration": + continue + logger.warning(f"Unexpected key: {key}") + if iteration is None: + logger.info(f"Loaded '{checkpoint_path}'") + else: + logger.info(f"Loaded '{checkpoint_path}' (iteration {iteration})") + + return model, iteration + + +def save_safetensors( + model: torch.nn.Module, + iteration: int, + checkpoint_path: Union[str, Path], + is_half: bool = False, + for_infer: bool = False, +) -> None: + """ + モデルを safetensors 形式で保存する。 + + Args: + model (torch.nn.Module): 保存するモデル + iteration (int): イテレーション番号 + checkpoint_path (Union[str, Path]): 保存先のパス + is_half (bool): モデルを半精度で保存するかどうかのフラグ + for_infer (bool): 推論用に保存するかどうかのフラグ + """ + + if hasattr(model, "module"): + state_dict = model.module.state_dict() + else: + state_dict = model.state_dict() + keys = [] + for k in state_dict: + if "enc_q" in k and for_infer: + continue # noqa: E701 + keys.append(k) + + new_dict = ( + {k: state_dict[k].half() for k in keys} + if is_half + else {k: state_dict[k] for k in keys} + ) + new_dict["iteration"] = torch.LongTensor([iteration]) + logger.info(f"Saved safetensors to {checkpoint_path}") + + save_file(new_dict, checkpoint_path) diff --git a/train_ms.py b/train_ms.py index b2cb02b..3d50c03 100644 --- a/train_ms.py +++ b/train_ms.py @@ -248,10 +248,7 @@ def run(): drop_last=False, collate_fn=collate_fn, ) - if ( - "use_noise_scaled_mas" in hps.model.keys() - and hps.model.use_noise_scaled_mas is True - ): + if hps.model.use_noise_scaled_mas is True: logger.info("Using noise scaled MAS for VITS2") mas_noise_scale_initial = 0.01 noise_scale_delta = 2e-6 @@ -259,10 +256,7 @@ def run(): logger.info("Using normal MAS for VITS1") mas_noise_scale_initial = 0.0 noise_scale_delta = 0.0 - if ( - "use_duration_discriminator" in hps.model.keys() - and hps.model.use_duration_discriminator is True - ): + if hps.model.use_duration_discriminator is True: logger.info("Using duration discriminator for VITS2") net_dur_disc = DurationDiscriminator( hps.model.hidden_channels, @@ -271,10 +265,7 @@ def run(): 0.1, gin_channels=hps.model.gin_channels if hps.data.n_speakers != 0 else 0, ).cuda(local_rank) - if ( - "use_spk_conditioned_encoder" in hps.model.keys() - and hps.model.use_spk_conditioned_encoder is True - ): + if hps.model.use_spk_conditioned_encoder is True: if hps.data.n_speakers == 0: raise ValueError( "n_speakers must be > 0 when using spk conditioned encoder to train multi-speaker model" @@ -370,31 +361,25 @@ def run(): if utils.is_resuming(model_dir): if net_dur_disc is not None: - _, _, dur_resume_lr, epoch_str = utils.load_checkpoint( - utils.latest_checkpoint_path(model_dir, "DUR_*.pth"), + _, _, dur_resume_lr, epoch_str = utils.checkpoints.load_checkpoint( + utils.checkpoints.get_latest_checkpoint_path(model_dir, "DUR_*.pth"), net_dur_disc, optim_dur_disc, - skip_optimizer=( - hps.train.skip_optimizer if "skip_optimizer" in hps.train else True - ), + skip_optimizer=hps.train.skip_optimizer, ) if not optim_dur_disc.param_groups[0].get("initial_lr"): optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr - _, optim_g, g_resume_lr, epoch_str = utils.load_checkpoint( - utils.latest_checkpoint_path(model_dir, "G_*.pth"), + _, optim_g, g_resume_lr, epoch_str = utils.checkpoints.load_checkpoint( + utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"), net_g, optim_g, - skip_optimizer=( - hps.train.skip_optimizer if "skip_optimizer" in hps.train else True - ), + skip_optimizer=hps.train.skip_optimizer, ) - _, optim_d, d_resume_lr, epoch_str = utils.load_checkpoint( - utils.latest_checkpoint_path(model_dir, "D_*.pth"), + _, optim_d, d_resume_lr, epoch_str = utils.checkpoints.load_checkpoint( + utils.checkpoints.get_latest_checkpoint_path(model_dir, "D_*.pth"), net_d, optim_d, - skip_optimizer=( - hps.train.skip_optimizer if "skip_optimizer" in hps.train else True - ), + skip_optimizer=hps.train.skip_optimizer, ) if not optim_g.param_groups[0].get("initial_lr"): optim_g.param_groups[0]["initial_lr"] = g_resume_lr @@ -404,21 +389,21 @@ def run(): epoch_str = max(epoch_str, 1) # global_step = (epoch_str - 1) * len(train_loader) global_step = int( - utils.get_steps(utils.latest_checkpoint_path(model_dir, "G_*.pth")) + utils.get_steps(utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth")) ) logger.info( f"******************Found the model. Current epoch is {epoch_str}, gloabl step is {global_step}*********************" ) else: try: - _ = utils.load_safetensors( + _ = utils.safetensors.load_safetensors( os.path.join(model_dir, "G_0.safetensors"), net_g ) - _ = utils.load_safetensors( + _ = utils.safetensors.load_safetensors( os.path.join(model_dir, "D_0.safetensors"), net_d ) if net_dur_disc is not None: - _ = utils.load_safetensors( + _ = utils.safetensors.load_safetensors( os.path.join(model_dir, "DUR_0.safetensors"), net_dur_disc ) logger.info("Loaded the pretrained models.") @@ -511,14 +496,16 @@ def run(): if epoch == hps.train.epochs: # Save the final models - utils.save_checkpoint( + assert optim_g is not None + utils.checkpoints.save_checkpoint( net_g, optim_g, hps.train.learning_rate, epoch, os.path.join(model_dir, "G_{}.pth".format(global_step)), ) - utils.save_checkpoint( + assert optim_d is not None + utils.checkpoints.save_checkpoint( net_d, optim_d, hps.train.learning_rate, @@ -526,14 +513,15 @@ def run(): os.path.join(model_dir, "D_{}.pth".format(global_step)), ) if net_dur_disc is not None: - utils.save_checkpoint( + assert optim_dur_disc is not None + utils.checkpoints.save_checkpoint( net_dur_disc, optim_dur_disc, hps.train.learning_rate, epoch, os.path.join(model_dir, "DUR_{}.pth".format(global_step)), ) - utils.save_safetensors( + utils.safetensors.save_safetensors( net_g, epoch, os.path.join( @@ -804,14 +792,15 @@ def train_and_evaluate( ): if not hps.speedup: evaluate(hps, net_g, eval_loader, writer_eval) - utils.save_checkpoint( + assert hps.model_dir is not None + utils.checkpoints.save_checkpoint( net_g, optim_g, hps.train.learning_rate, epoch, os.path.join(hps.model_dir, "G_{}.pth".format(global_step)), ) - utils.save_checkpoint( + utils.checkpoints.save_checkpoint( net_d, optim_d, hps.train.learning_rate, @@ -819,7 +808,7 @@ def train_and_evaluate( os.path.join(hps.model_dir, "D_{}.pth".format(global_step)), ) if net_dur_disc is not None: - utils.save_checkpoint( + utils.checkpoints.save_checkpoint( net_dur_disc, optim_dur_disc, hps.train.learning_rate, @@ -828,13 +817,13 @@ def train_and_evaluate( ) keep_ckpts = config.train_ms_config.keep_ckpts if keep_ckpts > 0: - utils.clean_checkpoints( - path_to_models=hps.model_dir, + utils.checkpoints.clean_checkpoints( + model_dir_path=hps.model_dir, n_ckpts_to_keep=keep_ckpts, sort_by_time=True, ) # Save safetensors (for inference) to `model_assets/{model_name}` - utils.save_safetensors( + utils.safetensors.save_safetensors( net_g, epoch, os.path.join( diff --git a/train_ms_jp_extra.py b/train_ms_jp_extra.py index a464d29..5ed68a1 100644 --- a/train_ms_jp_extra.py +++ b/train_ms_jp_extra.py @@ -247,10 +247,7 @@ def run(): drop_last=False, collate_fn=collate_fn, ) - if ( - "use_noise_scaled_mas" in hps.model.keys() - and hps.model.use_noise_scaled_mas is True - ): + if hps.model.use_noise_scaled_mas is True: logger.info("Using noise scaled MAS for VITS2") mas_noise_scale_initial = 0.01 noise_scale_delta = 2e-6 @@ -258,10 +255,7 @@ def run(): logger.info("Using normal MAS for VITS1") mas_noise_scale_initial = 0.0 noise_scale_delta = 0.0 - if ( - "use_duration_discriminator" in hps.model.keys() - and hps.model.use_duration_discriminator is True - ): + if hps.model.use_duration_discriminator is True: logger.info("Using duration discriminator for VITS2") net_dur_disc = DurationDiscriminator( hps.model.hidden_channels, @@ -272,19 +266,13 @@ def run(): ).cuda(local_rank) else: net_dur_disc = None - if ( - "use_wavlm_discriminator" in hps.model.keys() - and hps.model.use_wavlm_discriminator is True - ): + if hps.model.use_wavlm_discriminator is True: net_wd = WavLMDiscriminator( hps.model.slm.hidden, hps.model.slm.nlayers, hps.model.slm.initial_channel ).cuda(local_rank) else: net_wd = None - if ( - "use_spk_conditioned_encoder" in hps.model.keys() - and hps.model.use_spk_conditioned_encoder is True - ): + if hps.model.use_spk_conditioned_encoder is True: if hps.data.n_speakers == 0: raise ValueError( "n_speakers must be > 0 when using spk conditioned encoder to train multi-speaker model" @@ -394,15 +382,11 @@ def run(): if utils.is_resuming(model_dir): if net_dur_disc is not None: try: - _, _, dur_resume_lr, epoch_str = utils.load_checkpoint( - utils.latest_checkpoint_path(model_dir, "DUR_*.pth"), + _, _, dur_resume_lr, epoch_str = utils.checkpoints.load_checkpoint( + utils.checkpoints.get_latest_checkpoint_path(model_dir, "DUR_*.pth"), net_dur_disc, optim_dur_disc, - skip_optimizer=( - hps.train.skip_optimizer - if "skip_optimizer" in hps.train - else True - ), + skip_optimizer=hps.train.skip_optimizer, ) if not optim_dur_disc.param_groups[0].get("initial_lr"): optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr @@ -412,15 +396,11 @@ def run(): print("Initialize dur_disc") if net_wd is not None: try: - _, optim_wd, wd_resume_lr, epoch_str = utils.load_checkpoint( - utils.latest_checkpoint_path(model_dir, "WD_*.pth"), + _, optim_wd, wd_resume_lr, epoch_str = utils.checkpoints.load_checkpoint( + utils.checkpoints.get_latest_checkpoint_path(model_dir, "WD_*.pth"), net_wd, optim_wd, - skip_optimizer=( - hps.train.skip_optimizer - if "skip_optimizer" in hps.train - else True - ), + skip_optimizer=hps.train.skip_optimizer, ) if not optim_wd.param_groups[0].get("initial_lr"): optim_wd.param_groups[0]["initial_lr"] = wd_resume_lr @@ -430,21 +410,17 @@ def run(): logger.info("Initialize wavlm") try: - _, optim_g, g_resume_lr, epoch_str = utils.load_checkpoint( - utils.latest_checkpoint_path(model_dir, "G_*.pth"), + _, optim_g, g_resume_lr, epoch_str = utils.checkpoints.load_checkpoint( + utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"), net_g, optim_g, - skip_optimizer=( - hps.train.skip_optimizer if "skip_optimizer" in hps.train else True - ), + skip_optimizer=hps.train.skip_optimizer, ) - _, optim_d, d_resume_lr, epoch_str = utils.load_checkpoint( - utils.latest_checkpoint_path(model_dir, "D_*.pth"), + _, optim_d, d_resume_lr, epoch_str = utils.checkpoints.load_checkpoint( + utils.checkpoints.get_latest_checkpoint_path(model_dir, "D_*.pth"), net_d, optim_d, - skip_optimizer=( - hps.train.skip_optimizer if "skip_optimizer" in hps.train else True - ), + skip_optimizer=hps.train.skip_optimizer, ) if not optim_g.param_groups[0].get("initial_lr"): optim_g.param_groups[0]["initial_lr"] = g_resume_lr @@ -454,7 +430,7 @@ def run(): epoch_str = max(epoch_str, 1) # global_step = (epoch_str - 1) * len(train_loader) global_step = int( - utils.get_steps(utils.latest_checkpoint_path(model_dir, "G_*.pth")) + utils.get_steps(utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth")) ) logger.info( f"******************Found the model. Current epoch is {epoch_str}, gloabl step is {global_step}*********************" @@ -468,18 +444,18 @@ def run(): global_step = 0 else: try: - _ = utils.load_safetensors( + _ = utils.safetensors.load_safetensors( os.path.join(model_dir, "G_0.safetensors"), net_g ) - _ = utils.load_safetensors( + _ = utils.safetensors.load_safetensors( os.path.join(model_dir, "D_0.safetensors"), net_d ) if net_dur_disc is not None: - _ = utils.load_safetensors( + _ = utils.safetensors.load_safetensors( os.path.join(model_dir, "DUR_0.safetensors"), net_dur_disc ) if net_wd is not None: - _ = utils.load_safetensors( + _ = utils.safetensors.load_safetensors( os.path.join(model_dir, "WD_0.safetensors"), net_wd ) logger.info("Loaded the pretrained models.") @@ -586,14 +562,16 @@ def run(): scheduler_wd.step() if epoch == hps.train.epochs: # Save the final models - utils.save_checkpoint( + assert optim_g is not None + utils.checkpoints.save_checkpoint( net_g, optim_g, hps.train.learning_rate, epoch, os.path.join(model_dir, "G_{}.pth".format(global_step)), ) - utils.save_checkpoint( + assert optim_d is not None + utils.checkpoints.save_checkpoint( net_d, optim_d, hps.train.learning_rate, @@ -601,7 +579,8 @@ def run(): os.path.join(model_dir, "D_{}.pth".format(global_step)), ) if net_dur_disc is not None: - utils.save_checkpoint( + assert optim_dur_disc is not None + utils.checkpoints.save_checkpoint( net_dur_disc, optim_dur_disc, hps.train.learning_rate, @@ -609,14 +588,15 @@ def run(): os.path.join(model_dir, "DUR_{}.pth".format(global_step)), ) if net_wd is not None: - utils.save_checkpoint( + assert optim_wd is not None + utils.checkpoints.save_checkpoint( net_wd, optim_wd, hps.train.learning_rate, epoch, os.path.join(model_dir, "WD_{}.pth".format(global_step)), ) - utils.save_safetensors( + utils.safetensors.save_safetensors( net_g, epoch, os.path.join( @@ -949,14 +929,14 @@ def train_and_evaluate( ): if not hps.speedup: evaluate(hps, net_g, eval_loader, writer_eval) - utils.save_checkpoint( + utils.checkpoints.save_checkpoint( net_g, optim_g, hps.train.learning_rate, epoch, os.path.join(hps.model_dir, "G_{}.pth".format(global_step)), ) - utils.save_checkpoint( + utils.checkpoints.save_checkpoint( net_d, optim_d, hps.train.learning_rate, @@ -964,7 +944,7 @@ def train_and_evaluate( os.path.join(hps.model_dir, "D_{}.pth".format(global_step)), ) if net_dur_disc is not None: - utils.save_checkpoint( + utils.checkpoints.save_checkpoint( net_dur_disc, optim_dur_disc, hps.train.learning_rate, @@ -972,7 +952,7 @@ def train_and_evaluate( os.path.join(hps.model_dir, "DUR_{}.pth".format(global_step)), ) if net_wd is not None: - utils.save_checkpoint( + utils.checkpoints.save_checkpoint( net_wd, optim_wd, hps.train.learning_rate, @@ -981,13 +961,13 @@ def train_and_evaluate( ) keep_ckpts = config.train_ms_config.keep_ckpts if keep_ckpts > 0: - utils.clean_checkpoints( - path_to_models=hps.model_dir, + utils.checkpoints.clean_checkpoints( + model_dir_path=hps.model_dir, n_ckpts_to_keep=keep_ckpts, sort_by_time=True, ) # Save safetensors (for inference) to `model_assets/{model_name}` - utils.save_safetensors( + utils.safetensors.save_safetensors( net_g, epoch, os.path.join(