fix input_train_data_set_files value

liasece · Mar 25, 2023 · 4fa1280 · 4fa1280
1 parent 0391a26
commit 4fa1280
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -196,7 +196,7 @@ English (TODO) [google translate](https://github-com.translate.goog/liasece/sd-w
 
 其中：
 
-n: 训练数据集图片数量 "Train number of repetitions" （-1 时表示训练图片数为（上传图片数\*4））。大多数情况下你不需要修改默认值，你可以改成你上传的图片数量的 k 倍
+n: 训练数据集图片数量 "Train number of repetitions" （-1 时表示训练图片数为 min(512 / 上传图片数, 4) ）。大多数情况下你不需要修改默认值，你可以改成你上传的图片数量的 k 倍
 
 i: 你的显卡性能
 

diff --git a/liasece_sd_webui_train_tools/dateset_ui.py b/liasece_sd_webui_train_tools/dateset_ui.py
@@ -91,7 +91,10 @@ def on_ui_update_dataset_click(id: str, project: str, version: str, input_train_
         # copy to dataset path
         shutil.copyfile(f.name, file_path)
     if train_num_repetitions < 0:
-        train_num_repetitions = len(input_train_data_set_files)*4
+        if len(input_train_data_set_files) > 0:
+            train_num_repetitions = int(512 / len(input_train_data_set_files))
+        if train_num_repetitions <= 4:
+            train_num_repetitions = 4
     processed_output_path = os.path.join(processed_path, str(train_num_repetitions)+"_"+project)
     os.makedirs(processed_output_path, exist_ok=True)
     os.makedirs(origin_preload_data_path, exist_ok=True)