diff --git a/__init__.py b/__init__.py index ce8790b..f7f0d11 100644 --- a/__init__.py +++ b/__init__.py @@ -59,12 +59,14 @@ class MZ_KohyaSSDatasetConfig: "enable_bucket": (["enable", "disable"], {"default": "enable"}), "resolution": ("INT", {"default": 1024}), "num_repeats": ("INT", {"default": 1}), + "caption_extension": ([".caption", ".txt"], {"default": ".caption"}), "batch_size": ("INT", {"default": 1}), "force_clear": (["enable", "disable"], {"default": "disable"}), "force_clear_only_images": (["enable", "disable"], {"default": "disable"}), "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",), @@ -345,9 +347,23 @@ class MZ_KohyaSSLoraTrain: train_config_templates = Utils.listdir(s.train_config_template_dir) + priority = [ + "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) return { "required": { "workspace_config": ("MZ_TT_SS_WorkspaceConfig",), @@ -935,6 +951,7 @@ class MZ_TrainToolsDebug: "compact": (["enable", "disable"], {"default": "enable"}), "sort_keys": (["enable", "disable"], {"default": "enable"}), "underscore_numbers": (["enable", "disable"], {"default": "enable"}), + "index": ("STRING", {"default": ""}), }, } @@ -946,19 +963,8 @@ class MZ_TrainToolsDebug: CATEGORY = CATEGORY_NAME def start(self, **kwargs): - - from pprint import pprint, pformat - object = kwargs["object"] - indent = kwargs["indent"] - depth = kwargs["depth"] - width = kwargs["width"] - compact = kwargs["compact"] == "enable" - sort_keys = kwargs["sort_keys"] == "enable" - underscore_numbers = kwargs["underscore_numbers"] == "enable" - - debug = pformat(object, indent=indent, depth=depth, width=width, - compact=compact, sort_dicts=sort_keys, underscore_numbers=underscore_numbers) - return (debug,) + importlib.reload(mz_train_tools_core) + return mz_train_tools_core.MZ_TrainToolsDebug_call(kwargs.copy()) NODE_CLASS_MAPPINGS["MZ_TrainToolsDebug"] = MZ_TrainToolsDebug diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py index d52f123..b6364e7 100644 --- a/hook_kohya_ss_run.py +++ b/hook_kohya_ss_run.py @@ -400,49 +400,54 @@ def LOG(log): if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--sys_path", type=str, default="") - parser.add_argument("--config", type=str, default="") - parser.add_argument("--train_func", type=str, default="") - parser.add_argument("--master_port", type=int, default=0) - args = parser.parse_args() + try: + parser = argparse.ArgumentParser() + parser.add_argument("--sys_path", type=str, default="") + parser.add_argument("--config", type=str, default="") + parser.add_argument("--train_func", type=str, default="") + parser.add_argument("--master_port", type=int, default=0) + args = parser.parse_args() - master_port = args.master_port + master_port = args.master_port - print(f"master_port = {master_port}") + print(f"master_port = {master_port}") - sys_path = args.sys_path - if sys_path != "": - sys.path.append(sys_path) + sys_path = args.sys_path + if sys_path != "": + sys.path.append(sys_path) - config_file = args.config - if config_file == "": - raise Exception("train_config is empty") + config_file = args.config + if config_file == "": + raise Exception("train_config is empty") - global_config = {} - with open(config_file, "r") as f: - _global_config = f.read() - global_config = json.loads(_global_config) + global_config = {} + with open(config_file, "r") as f: + _global_config = f.read() + global_config = json.loads(_global_config) - train_config = global_config.get("train_config") - print(f"""=======================train_config======================= -{json.dumps(train_config, indent=4, ensure_ascii=False)} - """) + train_config = global_config.get("train_config") + print(f"""=======================train_config======================= + {json.dumps(train_config, indent=4, ensure_ascii=False)} + """) - other_config = global_config.get("other_config", {}) - print(f"""=======================other_config======================= -{json.dumps(other_config, indent=4, ensure_ascii=False)} - """) + other_config = global_config.get("other_config", {}) + print(f"""=======================other_config======================= + {json.dumps(other_config, indent=4, ensure_ascii=False)} + """) - train_func = args.train_func - if train_func == "": - raise Exception("train_func is empty") + train_func = args.train_func + if train_func == "": + raise Exception("train_func is empty") - print(f"train_func = {train_func}") + print(f"train_func = {train_func}") - time.sleep(2) - LOG({ - "type": "Read configuration completed!", - }) + time.sleep(2) + LOG({ + "type": "Read configuration completed!", + }) - func_map[train_func]() + func_map[train_func]() + except Exception as e: + print(f"Exception: {e}") + if sys.platform == "win32": + input("Press Enter to continue...") \ No newline at end of file diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 4009ecd..fe4e468 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -162,19 +162,26 @@ def MZ_ImageSelecter_call(args={}): with open(os.path.join(train_images_dir, caption_filename), "w", encoding="utf-8") as f: f.write(same_caption) - dataset_config_path = os.path.join(workspace_dir, "dataset.toml") - if conditioning_images is None: conditioning_images_dir = None - generate_toml_config( - dataset_config_path, + 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" + dataset_config_extension), enable_bucket=args.get("enable_bucket") == "enable", resolution=args.get("resolution"), batch_size=args.get("batch_size"), image_dir=train_images_dir, conditioning_data_dir=conditioning_images_dir, - caption_extension=".caption", + caption_extension=caption_extension, num_repeats=args.get("num_repeats"), ) return ( @@ -227,7 +234,11 @@ def MZ_KohyaSSUseConfig_call(args={}): config["train_config"]["output_name"] = f"{workspace_name}_{train_config_template}_{datetime}" dataset_config_path = os.path.join( - workspace_dir, "dataset.toml") + workspace_dir, "dataset.json") + if not os.path.exists(dataset_config_path): + dataset_config_path = os.path.join( + workspace_dir, "dataset.toml") + config["train_config"]["dataset_config"] = dataset_config_path config["train_config"]["max_train_steps"] = str( @@ -332,9 +343,8 @@ def check_install(): import logging -def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_size=1, image_dir=None, conditioning_data_dir=None, caption_extension=".caption", num_repeats=10, ): - check_install() - import toml +def generate_dataset_config(output_path, enable_bucket=True, resolution=512, batch_size=1, image_dir=None, conditioning_data_dir=None, caption_extension=".caption", num_repeats=10, ): + config = { 'general': { 'enable_bucket': enable_bucket, @@ -346,7 +356,6 @@ def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_ 'subsets': [ { 'image_dir': image_dir, - 'conditioning_data_dir': conditioning_data_dir, 'caption_extension': caption_extension, 'num_repeats': num_repeats, }, @@ -355,8 +364,19 @@ def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_ ], } - with open(output_path, "w", encoding="utf-8") as f: - toml.dump(config, f) + if conditioning_data_dir is not None: + config["datasets"][0]["subsets"][0]["conditioning_data_dir"] = conditioning_data_dir + + if output_path.endswith(".toml"): + check_install() + import toml + with open(output_path, "w", encoding="utf-8") as f: + toml.dump(config, f) + elif output_path.endswith(".json"): + with open(output_path, "w", encoding="utf-8") as f: + json.dump(config, f, indent=4, ensure_ascii=False) + else: + raise Exception(f"不支持的文件格式: {output_path}") from PIL import Image @@ -462,7 +482,7 @@ def run_hook_kohya_ss_run_file(workspace_dir, output_name, kohya_ss_tool_dir, tr screen_name = None if use_screen: screen_name = "mz_train_tools_core" - + process_instance = HSubprocess( cmd_list, screen_name=screen_name) process_instance.wait() @@ -517,7 +537,11 @@ def generate_kohya_ss_config(args): config["train_config"]["output_name"] = f"{workspace_name}_{train_config_template}_{datetime}" dataset_config_path = os.path.join( - workspace_dir, "dataset.toml") + workspace_dir, "dataset.json") + if not os.path.exists(dataset_config_path): + dataset_config_path = os.path.join( + workspace_dir, "dataset.toml") + config["train_config"]["dataset_config"] = dataset_config_path config["train_config"]["max_train_steps"] = str( @@ -978,3 +1002,33 @@ def MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I_call(args={}): model_type="HYDiT_unet_merge_lora") return (image,) + + +def MZ_TrainToolsDebug_call(kwargs): + from pprint import pprint, pformat + object = kwargs["object"] + indent = kwargs["indent"] + depth = kwargs["depth"] + width = kwargs["width"] + compact = kwargs["compact"] == "enable" + sort_keys = kwargs["sort_keys"] == "enable" + underscore_numbers = kwargs["underscore_numbers"] == "enable" + + index = kwargs["index"] + # 用点号分割 + index = index.split(".") + for i in index: + if i == "": + continue + if isinstance(object, (list, tuple)): + object = object[int(i)] + elif isinstance(object, dict): + object = object.get(i) + elif hasattr(object, i): + object = getattr(object, i) + else: + object = object[i] + + debug = pformat(object, indent=indent, depth=depth, width=width, + compact=compact, sort_dicts=sort_keys, underscore_numbers=underscore_numbers) + return (debug,) diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index a08b549..15aad24 100644 --- a/mz_train_tools_utils.py +++ b/mz_train_tools_utils.py @@ -1166,9 +1166,11 @@ class HSubprocess: def stop(self): if self._mswindows: - if self.process_instance_pid is not None: - os.system( - f'taskkill /F /FI "WINDOWTITLE eq hook_kohya_ss_run" /T') + print('taskkill /F /FI "WINDOWTITLE eq hook_kohya_ss_run" /T') + os.system( + f'taskkill /F /FI "WINDOWTITLE eq hook_kohya_ss_run" /T') + + if self.process_instance is not None: if self.screen_name is not None: subprocess.run(