From 1dc4e45480beb0417ab2bfcd45e89995635f99d9 Mon Sep 17 00:00:00 2001
From: litagin02 <litagin02@gmail.com>
Date: Fri, 12 Jan 2024 09:05:28 +0900
Subject: [PATCH] Improve: when initialize, if models dir exists, backup and
 delete

---
 webui_train.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/webui_train.py b/webui_train.py
index ac461bb..7b631a7 100644
--- a/webui_train.py
+++ b/webui_train.py
@@ -60,13 +60,18 @@ def initialize(model_name, batch_size, epochs, save_every_steps, bf16_run):
     config["train"]["eval_interval"] = save_every_steps
 
     model_path = os.path.join(dataset_path, "models")
+    if os.path.exists(model_path):
+        logger.warning(f"Step 1: {model_path} already exists, so copy it to backup.")
+        shutil.copytree(
+            src=model_path,
+            dst=os.path.join(dataset_path, "models_backup"),
+        )
+        shutil.rmtree(model_path)
     try:
         shutil.copytree(
             src="pretrained",
             dst=model_path,
         )
-    except FileExistsError:
-        logger.warning(f"Step 1: {model_path} already exists.")
     except FileNotFoundError:
         logger.error("Step 1: `pretrained` folder not found.")
         return False, "Step 1, Error: pretrainedフォルダが見つかりません。"