diff --git a/train_ms.py b/train_ms.py index cb49a82..a917569 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 ( @@ -30,7 +31,6 @@ 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 = ( @@ -44,7 +44,12 @@ torch.backends.cuda.enable_mem_efficient_sdp( ) # Not available if torch version is lower than 2.0 torch.backends.cuda.enable_math_sdp(True) -IS_COLAB = "google.colab" in sys.modules +try: + import google.colab + + IS_COLAB = True +except ImportError: + IS_COLAB = False global_step = 0 diff --git a/webui_train.py b/webui_train.py index db72bdd..c44f464 100644 --- a/webui_train.py +++ b/webui_train.py @@ -10,7 +10,12 @@ 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 +try: + import google.colab + + IS_COLAB = True +except ImportError: + IS_COLAB = False def get_path(model_name): @@ -48,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.")