327 lines
12 KiB
Python
327 lines
12 KiB
Python
import inspect
|
|
import json
|
|
import os
|
|
import folder_paths
|
|
import importlib
|
|
from .mz_train_tools_utils import Utils
|
|
from . import mz_train_tools_core
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
}
|
|
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
}
|
|
|
|
AUTHOR_NAME = "MinusZone"
|
|
CATEGORY_NAME = f"{AUTHOR_NAME} - TrainTools"
|
|
|
|
|
|
class MZ_KohyaSSInitWorkspace:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"lora_name": ("STRING", {"default": ""}),
|
|
"branch": ("STRING", {"default": "71e2c91330a9d866ec05cdd10584bbb962896a99"}),
|
|
"source": ([
|
|
"github",
|
|
"githubfast",
|
|
"521github",
|
|
"kkgithub",
|
|
], {"default": "github"}),
|
|
"seed": ("INT", {"default": 0}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MZ_TT_SS_WorkspaceConfig",)
|
|
RETURN_NAMES = ("workspace_config",)
|
|
|
|
FUNCTION = "start"
|
|
CATEGORY = CATEGORY_NAME + "/kohya_ss"
|
|
|
|
def start(self, **kwargs):
|
|
importlib.reload(mz_train_tools_core)
|
|
return mz_train_tools_core.MZ_KohyaSSInitWorkspace_call(kwargs)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS["MZ_KohyaSSInitWorkspace"] = MZ_KohyaSSInitWorkspace
|
|
NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSInitWorkspace"] = f"{AUTHOR_NAME} - KohyaSSInitWorkspace"
|
|
|
|
|
|
class MZ_KohyaSSDatasetConfig:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"workspace_config": ("MZ_TT_SS_WorkspaceConfig",),
|
|
"images": ("IMAGE",),
|
|
"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"}),
|
|
"auto_reg": (["enable", "disable"], {"default": "enable"}),
|
|
},
|
|
"optional": {
|
|
"conditioning_images": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (f"STRING",)
|
|
RETURN_NAMES = ("workspace_images_dir",)
|
|
|
|
FUNCTION = "start"
|
|
CATEGORY = CATEGORY_NAME + "/kohya_ss"
|
|
|
|
def start(self, **kwargs):
|
|
importlib.reload(mz_train_tools_core)
|
|
return mz_train_tools_core.MZ_KohyaSSDatasetConfig_call(kwargs)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS["MZ_ImagesCopyWorkspace"] = MZ_KohyaSSDatasetConfig
|
|
NODE_DISPLAY_NAME_MAPPINGS["MZ_ImagesCopyWorkspace"] = f"{AUTHOR_NAME} - ImagesCopyWorkspace"
|
|
|
|
# 别名
|
|
NODE_CLASS_MAPPINGS["MZ_KohyaSSDatasetConfig"] = MZ_KohyaSSDatasetConfig
|
|
NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSDatasetConfig"] = f"{AUTHOR_NAME} - KohyaSSDatasetConfig"
|
|
|
|
|
|
class MZ_KohyaSSAdvConfig:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"xformers": (["enable", "disable"], {"default": "enable"}),
|
|
"sdpa": (["enable", "disable"], {"default": "disable"}),
|
|
"fp8_base": (["enable", "disable"], {"default": "disable"}),
|
|
"mixed_precision": (["no", "fp16", "bf16"], {"default": "fp16"}),
|
|
"gradient_accumulation_steps": ("INT", {"default": 1}),
|
|
"gradient_checkpointing": (["enable", "disable"], {"default": "disable"}),
|
|
"cache_latents": (["enable", "disable"], {"default": "enable"}),
|
|
"cache_latents_to_disk": (["enable", "disable"], {"default": "enable"}),
|
|
"network_dim": ("INT", {"default": 16}),
|
|
"network_alpha": ("INT", {"default": 8}),
|
|
"network_module": ([
|
|
"networks.lora",
|
|
"networks.dylora",
|
|
"networks.oft",
|
|
], {"default": "networks.lora"}),
|
|
"network_train_unet_only": (["enable", "disable"], {"default": "enable"}),
|
|
# linear, cosine, cosine_with_restarts, polynomial, constant (default), constant_with_warmup, adafactor
|
|
"lr_scheduler": ([
|
|
"linear",
|
|
"cosine",
|
|
"cosine_with_restarts",
|
|
"polynomial",
|
|
"constant",
|
|
"constant_with_warmup",
|
|
"adafactor",
|
|
], {"default": "cosine"}),
|
|
"lr_scheduler_num_cycles": ("INT", {"default": 1}),
|
|
# AdamW (default), AdamW8bit, PagedAdamW, PagedAdamW8bit, PagedAdamW32bit, Lion8bit, PagedLion8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation(DAdaptAdamPreprint), DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptAdanIP, DAdaptLion, DAdaptSGD, AdaFactor
|
|
"optimizer_type": ([
|
|
"AdamW",
|
|
"AdamW8bit",
|
|
"PagedAdamW",
|
|
"PagedAdamW8bit",
|
|
"PagedAdamW32bit",
|
|
"Lion8bit",
|
|
"PagedLion8bit",
|
|
"Lion",
|
|
"SGDNesterov",
|
|
"SGDNesterov8bit",
|
|
"DAdaptation",
|
|
"DAdaptAdaGrad",
|
|
"DAdaptAdam",
|
|
"DAdaptAdan",
|
|
"DAdaptAdanIP",
|
|
"DAdaptLion",
|
|
"DAdaptSGD",
|
|
"AdaFactor",
|
|
], {"default": "AdamW"}),
|
|
"lr_warmup_steps": ("INT", {"default": 0}),
|
|
"unet_lr": ("STRING", {"default": ""}),
|
|
"text_encoder_lr": ("STRING", {"default": ""}),
|
|
"shuffle_caption": (["enable", "disable"], {"default": "disable"}),
|
|
"save_precision": (["float", "fp16", "bf16"], {"default": "fp16"}),
|
|
"persistent_data_loader_workers": (["enable", "disable"], {"default": "enable"}),
|
|
"no_metadata": (["enable", "disable"], {"default": "enable"}),
|
|
"noise_offset": ("FLOAT", {"default": 0.1}),
|
|
"no_half_vae": (["enable", "disable"], {"default": "enable"}),
|
|
"lowram": (["enable", "disable"], {"default": "disable"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MZ_TT_SS_AdvConfig",)
|
|
RETURN_NAMES = ("advanced_config",)
|
|
|
|
FUNCTION = "start"
|
|
|
|
CATEGORY = CATEGORY_NAME + "/kohya_ss"
|
|
|
|
def start(self, **kwargs):
|
|
importlib.reload(mz_train_tools_core)
|
|
return mz_train_tools_core.MZ_KohyaSSAdvConfig_call(kwargs)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS["MZ_KohyaSSAdvConfig"] = MZ_KohyaSSAdvConfig
|
|
NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSAdvConfig"] = f"{AUTHOR_NAME} - KohyaSSAdvConfig"
|
|
|
|
|
|
class AlwaysEqualProxy(str):
|
|
def __eq__(self, _):
|
|
return True
|
|
|
|
def __ne__(self, _):
|
|
return False
|
|
|
|
|
|
class MZ_KohyaSSLoraTrain:
|
|
train_config_template_dir = os.path.join(
|
|
os.path.dirname(__file__), "configs", "kohya_ss_lora"
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
loras = [
|
|
"latest",
|
|
"empty",
|
|
]
|
|
|
|
workspaces_dir = os.path.join(
|
|
folder_paths.output_directory, "mz_train_workspaces")
|
|
|
|
# 使用walk查询所有的workspace中的所有lora模型,lora存放在每个workspace的output目录下
|
|
workspaces_loras = []
|
|
for root, dirs, files in os.walk(workspaces_dir):
|
|
|
|
# 排除隐藏文件夹
|
|
dirs[:] = [d for d in dirs if not d.startswith(".")]
|
|
if root.endswith("output"):
|
|
for file in files:
|
|
if file.endswith(".safetensors"):
|
|
workspaces_loras.append(
|
|
os.path.join(root, file)
|
|
)
|
|
|
|
# 按创建时间排序
|
|
workspaces_loras = sorted(
|
|
workspaces_loras, key=lambda x: os.path.getctime(x), reverse=True)
|
|
|
|
comfyui_full_loras = []
|
|
comfyui_loras = folder_paths.get_filename_list("loras")
|
|
for lora in comfyui_loras:
|
|
lora_path = folder_paths.get_full_path("loras", lora)
|
|
comfyui_full_loras.append(lora_path)
|
|
|
|
# 按创建时间排序
|
|
comfyui_full_loras = sorted(
|
|
comfyui_full_loras, key=lambda x: os.path.getctime(x), reverse=True)
|
|
|
|
loras = loras + workspaces_loras + comfyui_full_loras
|
|
|
|
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",),
|
|
"ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
|
|
"max_train_steps": ("INT", {"default": 0, "min": 0, "max": 0x7fffffff}),
|
|
"max_train_epochs": ("INT", {"default": 100, "min": 0, "max": 0x7fffffff}),
|
|
"save_every_n_epochs": ("INT", {"default": 10}),
|
|
"learning_rate": ("STRING", {"default": "1e-5"}),
|
|
|
|
"base_lora": (loras, {"default": "latest"}),
|
|
"sample_generate": (["enable", "disable"], {"default": "enable"}),
|
|
"sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
|
|
},
|
|
"optional": {
|
|
"advanced_config": ("MZ_TT_SS_AdvConfig",),
|
|
"caption_completed_flag": (AlwaysEqualProxy("*"),),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
RETURN_NAMES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "start"
|
|
CATEGORY = CATEGORY_NAME + "/kohya_ss"
|
|
|
|
def start(self, **kwargs):
|
|
importlib.reload(mz_train_tools_core)
|
|
|
|
return mz_train_tools_core.MZ_KohyaSSTrain_call(kwargs)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS["MZ_KohyaSSLoraTrain"] = MZ_KohyaSSLoraTrain
|
|
NODE_DISPLAY_NAME_MAPPINGS[
|
|
"MZ_KohyaSSLoraTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(lora)"
|
|
|
|
|
|
class MZ_LoadImagesFromDirectoryPath:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"directory": ("STRING", {"default": "X://path/to/images"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
|
|
FUNCTION = "start"
|
|
|
|
CATEGORY = CATEGORY_NAME + "/tools"
|
|
|
|
def start(self, **kwargs):
|
|
from PIL import Image
|
|
images = []
|
|
image_dir = kwargs["directory"]
|
|
if not os.path.exists(image_dir):
|
|
return (images,)
|
|
images = Utils.listdir(image_dir)
|
|
|
|
images = [x for x in images if x.lower().endswith(
|
|
".png") or x.lower().endswith(".jpg") or file.lower().endswith(".webp")]
|
|
images = [os.path.join(image_dir, x) for x in images]
|
|
|
|
pil_images = []
|
|
for image in images:
|
|
pil_images.append(Image.open(image))
|
|
|
|
tensor_images = []
|
|
for pil_image in pil_images:
|
|
tensor_images.append(Utils.pil2tensor(pil_image))
|
|
|
|
return (Utils.list_tensor2tensor(tensor_images),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS["MZ_LoadImagesFromDirectoryPath"] = MZ_LoadImagesFromDirectoryPath
|
|
NODE_DISPLAY_NAME_MAPPINGS[
|
|
"MZ_LoadImagesFromDirectoryPath"] = f"{AUTHOR_NAME} - LoadImagesFromDirectoryPath"
|