From e1ce75e0f187a1396a6e4f575fd06ea62bae6588 Mon Sep 17 00:00:00 2001 From: wailovet Date: Tue, 9 Jul 2024 20:41:05 +0800 Subject: [PATCH] 'update' --- __init__.py | 11 ++++++----- mz_train_tools_core.py | 11 +++++++++-- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/__init__.py b/__init__.py index 436a246..f7f0d11 100644 --- a/__init__.py +++ b/__init__.py @@ -66,6 +66,7 @@ class MZ_KohyaSSDatasetConfig: "same_caption_generate": (["enable", "disable"], {"default": "disable"}), "same_caption": ("STRING", {"default": "", "dynamicPrompts": True, "multiline": True}), "image_format": (["png", "jpg", "webp"], {"default": "webp"}), + "dataset_config_extension": ([".toml", ".json"], {"default": ".json"}), }, "optional": { "conditioning_images": ("IMAGE",), @@ -346,23 +347,23 @@ class MZ_KohyaSSLoraTrain: train_config_templates = Utils.listdir(s.train_config_template_dir) - priority = [ - "lora", + "lora", "1_2" "1_1" ] # 去掉json后缀 train_config_templates = [os.path.splitext(x)[0] for x in train_config_templates] - + def priority_sort(x): for p in priority: if x.find(p) != -1: return priority.index(p) return 999 - - train_config_templates = sorted(train_config_templates, key=priority_sort) + + train_config_templates = sorted( + train_config_templates, key=priority_sort) return { "required": { "workspace_config": ("MZ_TT_SS_WorkspaceConfig",), diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 007691c..fe4e468 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -167,8 +167,15 @@ def MZ_ImageSelecter_call(args={}): caption_extension = args.get("caption_extension", ".caption") + if os.path.exists(os.path.join(workspace_dir, "dataset.json")): + os.remove(os.path.join(workspace_dir, "dataset.json")) + + if os.path.exists(os.path.join(workspace_dir, "dataset.toml")): + os.remove(os.path.join(workspace_dir, "dataset.toml")) + + dataset_config_extension = args.get("dataset_config_extension") generate_dataset_config( - os.path.join(workspace_dir, "dataset.json"), + os.path.join(workspace_dir, "dataset" + dataset_config_extension), enable_bucket=args.get("enable_bucket") == "enable", resolution=args.get("resolution"), batch_size=args.get("batch_size"), @@ -348,7 +355,7 @@ def generate_dataset_config(output_path, enable_bucket=True, resolution=512, bat 'batch_size': batch_size, 'subsets': [ { - 'image_dir': image_dir, + 'image_dir': image_dir, 'caption_extension': caption_extension, 'num_repeats': num_repeats, },