nodes/MinusZoneAI-ComfyUI-TrainTools-MZ
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #15 from MinusZoneAI/dev

Browse Source 
Fix some problems
This commit is contained in:
喵哩个咪
2024-07-09 20:42:11 +08:00
committed by GitHub
parent 18c0abe5b2 e1ce75e0f1
commit 9c60ef9046
4 changed files with 132 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+19 -13
View File
@@ -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
hook_kohya_ss_run.py
+40 -35
View File
@@ -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...")
mz_train_tools_core.py
+68 -14
View File
@@ -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,)
mz_train_tools_utils.py
+5 -3
View File
@@ -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(