From ebdf607f0df540a25c18ff43e390cb8bedb47518 Mon Sep 17 00:00:00 2001
From: xieyongliang <yongliangxie@hotmail.com>
Date: Mon, 27 Mar 2023 23:37:16 +0800
Subject: [PATCH] fix issues with 1.x SD models training

---
 webui.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/webui.py b/webui.py
index e117d4ede84..216071e1742 100644
--- a/webui.py
+++ b/webui.py
@@ -750,10 +750,11 @@ def train():
                 os.makedirs(os.path.dirname("/opt/ml/model/ControlNet/"), exist_ok=True)
                 train_steps=int(db_config.revision)
                 model_file_basename = f'{db_model_name}_{train_steps}_lora' if db_config.use_lora else f'{db_model_name}_{train_steps}'
-                f1=os.path.join(sd_models_dir, db_model_name, f'{model_file_basename}.yaml')
-                if os.path.exists(f1):
-                    shutil.copy(f1,"/opt/ml/model/Stable-diffusion/")
-                if db_save_safetensors:
+                if db_config.v2:
+                    f1=os.path.join(sd_models_dir, db_model_name, f'{model_file_basename}.yaml')
+                    if os.path.exists(f1):
+                        shutil.copy(f1,"/opt/ml/model/Stable-diffusion/")
+                if db_config.save_safetensors:
                     f2=os.path.join(sd_models_dir, db_model_name, f'{model_file_basename}.safetensors')
                     if os.path.exists(f2):
                         shutil.copy(f2,"/opt/ml/model/Stable-diffusion/")