'update'

__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",),

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,
                     },