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

'update'

Browse Source
This commit is contained in:
wailovet
2024-06-25 02:30:26 +08:00
parent 2bf060a47f
commit 698f04b9fc
4 changed files with 107 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+11 -2
View File
@@ -51,7 +51,6 @@ NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSInitWorkspace"] = f"{AUTHOR_NAME} - KohyaS
class MZ_ImagesCopyWorkspace:
@classmethod
def INPUT_TYPES(s):
return {
@@ -62,12 +61,17 @@ class MZ_ImagesCopyWorkspace:
"resolution": ("INT", {"default": 512}),
"num_repeats": ("INT", {"default": 1}),
"batch_size": ("INT", {"default": 1}),
# "class_name": ("STRING", {"default": "girl", "dynamicPrompts": True}),
"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}),
"description": ("STRING", {"default": """
如果训练类型是controlnet,必须传入预处理后的图片(conditioning_images)
""", "multiline": True}),
},
"optional": {
"conditioning_images": ("IMAGE",),
}
}
RETURN_TYPES = (f"STRING",)
@@ -323,3 +327,8 @@ class MZ_LoadImagesFromDirectoryPath:
NODE_CLASS_MAPPINGS["MZ_LoadImagesFromDirectoryPath"] = MZ_LoadImagesFromDirectoryPath
NODE_DISPLAY_NAME_MAPPINGS[
"MZ_LoadImagesFromDirectoryPath"] = f"{AUTHOR_NAME} - LoadImagesFromDirectoryPath"
# 别名
NODE_CLASS_MAPPINGS["MZ_DatasetConfig"] = MZ_ImagesCopyWorkspace
NODE_DISPLAY_NAME_MAPPINGS["MZ_DatasetConfig"] = f"{AUTHOR_NAME} - DatasetConfig"
configs/kohya_ss_lora/controlnet_sd1_5.json
+28
View File
@@ -0,0 +1,28 @@
{
"metadata": {
"train_type": "controlnet_sd1_5"
},
"train_config": {
"pretrained_model_name_or_path": "",
"max_train_steps": "4500",
"xformers": true,
"sdpa": false,
"fp8_base": false,
"mixed_precision": "fp16",
"cache_latents": true,
"cache_latents_to_disk": true,
"learning_rate": "1e-5",
"lr_scheduler": "cosine_with_restarts",
"optimizer_type": "AdamW",
"save_every_n_epochs": "20",
"shuffle_caption": true,
"lr_warmup_steps": "0",
"save_precision": "fp16",
"lr_scheduler_num_cycles": "1",
"persistent_data_loader_workers": true,
"noise_offset": "0.1",
"output_dir": "",
"output_name": "",
"lowram": false
}
}
hook_kohya_ss_run.py
+42
View File
@@ -161,9 +161,51 @@ def run_lora_sdxl():
trainer.train(train_args)
from types import SimpleNamespace
class SimpleNamespaceCNWarrper(SimpleNamespace):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.__dict__.update(kwargs) # or self.__dict__ = kwargs
self.__dict__["mid_block_type"] = "UNetMidBlock2DCrossAttn"
self.__iter__ = lambda: iter(kwargs.keys())
# is not iterable
def __iter__(self):
return iter(self.__dict__.keys())
# object has no attribute 'num_attention_heads'
def __getattr__(self, name):
return self.__dict__.get(name, None)
def run_controlnet_sd1_5():
import types
types.SimpleNamespace = SimpleNamespaceCNWarrper
hook_kohya_ss_utils.hook_kohya_ss()
# 覆盖sample_images生成函数,包括进度条和生成图片功能
import train_controlnet
# 配置对应的pipeline
import library.train_util
global sample_images_pipe_class
sample_images_pipe_class = library.train_util.StableDiffusionLongPromptWeightingPipeline
train_config = json.loads(train_config_json)
train_args = config2args(train_controlnet.setup_parser(), train_config)
LOG({
"type": "start_train",
})
train_controlnet.train(train_args)
func_map = {
"run_lora_sd1_5": run_lora_sd1_5,
"run_lora_sdxl": run_lora_sdxl,
"run_controlnet_sd1_5": run_controlnet_sd1_5,
}
mz_train_tools_core.py
+26 -6
View File
@@ -93,6 +93,12 @@ def MZ_ImageSelecter_call(args={}):
images = args.get("images")
pil_images = Utils.tensors2pil_list(images)
conditioning_images = args.get("conditioning_images", None)
conditioning_pil_images = None
if conditioning_images is not None:
conditioning_pil_images = Utils.tensors2pil_list(
conditioning_images)
resolution = args.get("resolution", 512)
workspace_config = args.get("workspace_config", {})
@@ -106,6 +112,9 @@ def MZ_ImageSelecter_call(args={}):
train_images_dir = os.path.join(workspace_dir, "train_images")
os.makedirs(train_images_dir, exist_ok=True)
conditioning_images_dir = os.path.join(
workspace_dir, "conditioning_images")
force_clear = args.get("force_clear") == "enable"
force_clear_only_images = args.get("force_clear_only_images") == "enable"
if force_clear:
@@ -120,10 +129,16 @@ def MZ_ImageSelecter_call(args={}):
saved_images_path = []
for i, pil_image in enumerate(pil_images):
pil_image = Utils.resize_max(pil_image, resolution, resolution)
width, height = pil_image.size
filename = hashlib.md5(pil_image.tobytes()).hexdigest() + ".png"
pil_image.save(os.path.join(train_images_dir, filename))
saved_images_path.append(filename)
if conditioning_pil_images is not None:
os.makedirs(conditioning_images_dir, exist_ok=True)
conditioning_pil_images[i].resize((width, height)).save(
os.path.join(conditioning_images_dir, filename))
same_caption_generate = args.get("same_caption_generate") == "enable"
if same_caption_generate:
same_caption = args.get("same_caption").strip()
@@ -136,12 +151,17 @@ def MZ_ImageSelecter_call(args={}):
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,
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",
num_repeats=args.get("num_repeats"),
)
@@ -290,7 +310,7 @@ def check_install():
import logging
def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_size=1, image_dir=None, caption_extension=".caption", num_repeats=10, ):
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
config = {
@@ -304,6 +324,7 @@ 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,
},
@@ -374,11 +395,8 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
if trainer_func.find("sdxl") != -1:
taesd_type = "sdxl"
pb = Utils.progress_bar(train_config.get("max_train_steps"), taesd_type)
import traceback
import comfy.model_management
@@ -401,7 +419,6 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
max_side = max(xy_img.width, xy_img.height)
# print(f"global_step: {global_step}, max_train_steps: {max_train_steps}")
total_steps = resp.get("total_steps")
pb.update(
int(global_step), int(total_steps), ("JPEG", xy_img, max_side))
@@ -517,6 +534,9 @@ def MZ_KohyaSSTrain_call(args={}):
elif train_type == "lora_sdxl":
run_hook_kohya_ss_run_file(
kohya_ss_tool_dir, train_config, "run_lora_sdxl", other_config)
elif train_type == "controlnet_sd1_5":
run_hook_kohya_ss_run_file(
kohya_ss_tool_dir, train_config, "run_controlnet_sd1_5", other_config)
else:
raise Exception(
f"暂时不支持的训练类型: {train_type}")