Compare commits
14
Commits
dev
..
dev20240815
| Author | SHA1 | Date | |
|---|---|---|---|
 wailovet
|
ad50d8a1f1 | ||
|
喵哩个咪
|
9c60ef9046 | ||
|
喵哩个咪
|
18c0abe5b2 | ||
|
喵哩个咪
|
d518a2f31e | ||
|
喵哩个咪
|
6bd04f06e2 | ||
|
喵哩个咪
|
ebeb2c1bf3 | ||
|
喵哩个咪
|
b0dd687433 | ||
|
喵哩个咪
|
f8418440e7 | ||
|
喵哩个咪
|
1383e803db | ||
|
喵哩个咪
|
015267c8c0 | ||
|
wailovet
|
17c3256867 | ||
|
wailovet
|
996c94edde | ||
|
wailovet
|
8125869ef7 | ||
|
喵哩个咪
|
4a956a036d |
+2
-647
@@ -67,6 +67,7 @@ class MZ_KohyaSSDatasetConfig:
|
||||
"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",),
|
||||
@@ -76,17 +77,12 @@ class MZ_KohyaSSDatasetConfig:
|
||||
RETURN_TYPES = (f"STRING",)
|
||||
RETURN_NAMES = ("workspace_images_dir",)
|
||||
|
||||
# OUTPUT_NODE = True
|
||||
MZ_DESC = """
|
||||
如果训练类型是controlnet,必须传入预处理后的图片(conditioning_images)
|
||||
"""
|
||||
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME + "/kohya_ss"
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
return mz_train_tools_core.MZ_ImageSelecter_call(kwargs)
|
||||
return mz_train_tools_core.MZ_KohyaSSDatasetConfig_call(kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_ImagesCopyWorkspace"] = MZ_KohyaSSDatasetConfig
|
||||
@@ -97,51 +93,6 @@ NODE_CLASS_MAPPINGS["MZ_KohyaSSDatasetConfig"] = MZ_KohyaSSDatasetConfig
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSDatasetConfig"] = f"{AUTHOR_NAME} - KohyaSSDatasetConfig"
|
||||
|
||||
|
||||
class MZ_KohyaSSUseConfig_oldversion:
|
||||
train_config_template_dir = os.path.join(
|
||||
os.path.dirname(__file__), "configs", "kohya_ss_lora"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
train_config_templates = Utils.listdir(s.train_config_template_dir)
|
||||
|
||||
# 去掉json后缀
|
||||
train_config_templates = [os.path.splitext(x)[0]
|
||||
for x in train_config_templates]
|
||||
return {
|
||||
"required": {
|
||||
"workspace_config": ("MZ_TT_SS_WorkspaceConfig",),
|
||||
"workspace_images_dir": ("STRING", {"forceInput": True}),
|
||||
"train_config_template": (train_config_templates,),
|
||||
"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"}),
|
||||
},
|
||||
"optional": {
|
||||
"save_advanced_config": ("MZ_TT_SS_AdvConfig",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (f"MZ_TT_SS_TrainConfig",)
|
||||
RETURN_NAMES = ("train_config",)
|
||||
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME + "/kohya_ss" + "/v1"
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
kwargs["train_config_template_dir"] = self.train_config_template_dir
|
||||
return mz_train_tools_core.MZ_KohyaSSUseConfig_call(kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_KohyaSSUseConfig"] = MZ_KohyaSSUseConfig_oldversion
|
||||
NODE_DISPLAY_NAME_MAPPINGS[
|
||||
"MZ_KohyaSSUseConfig"] = f"{AUTHOR_NAME} - KohyaSSUseConfig(old version)"
|
||||
|
||||
|
||||
class MZ_KohyaSSAdvConfig:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -232,75 +183,6 @@ class AlwaysEqualProxy(str):
|
||||
return False
|
||||
|
||||
|
||||
class MZ_KohyaSSTrain_oldversion:
|
||||
|
||||
@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
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"train_config": ("MZ_TT_SS_TrainConfig",),
|
||||
"base_lora": (loras, {"default": "latest"}),
|
||||
"sample_generate": (["enable", "disable"], {"default": "enable"}),
|
||||
"sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
|
||||
},
|
||||
"optional": {
|
||||
"has_no_effect": (AlwaysEqualProxy("*"),),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
RETURN_NAMES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME + "/kohya_ss" + "/v1"
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
return mz_train_tools_core.MZ_KohyaSSTrain_call(kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_KohyaSSTrain"] = MZ_KohyaSSTrain_oldversion
|
||||
NODE_DISPLAY_NAME_MAPPINGS[
|
||||
"MZ_KohyaSSTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(old version)"
|
||||
|
||||
|
||||
class MZ_KohyaSSLoraTrain:
|
||||
train_config_template_dir = os.path.join(
|
||||
os.path.dirname(__file__), "configs", "kohya_ss_lora"
|
||||
@@ -367,7 +249,6 @@ class MZ_KohyaSSLoraTrain:
|
||||
return {
|
||||
"required": {
|
||||
"workspace_config": ("MZ_TT_SS_WorkspaceConfig",),
|
||||
"train_config_template": (train_config_templates,),
|
||||
"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}),
|
||||
@@ -393,25 +274,6 @@ class MZ_KohyaSSLoraTrain:
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
|
||||
# train_config = {
|
||||
# "workspace_config": kwargs["workspace_config"],
|
||||
# "train_config_template": kwargs["train_config_template"],
|
||||
# "ckpt_name": kwargs["ckpt_name"],
|
||||
# "max_train_steps": kwargs["max_train_steps"],
|
||||
# "max_train_epochs": kwargs["max_train_epochs"],
|
||||
# "save_every_n_epochs": kwargs["save_every_n_epochs"],
|
||||
# "learning_rate": kwargs["learning_rate"],
|
||||
# }
|
||||
|
||||
# train_config["train_config_template_dir"] = self.train_config_template_dir
|
||||
|
||||
# advanced_config = kwargs.get("advanced_config", None)
|
||||
|
||||
# if advanced_config is not None:
|
||||
# for k, v in advanced_config.items():
|
||||
# train_config[k] = v
|
||||
|
||||
# kwargs["train_config"] = train_config
|
||||
return mz_train_tools_core.MZ_KohyaSSTrain_call(kwargs)
|
||||
|
||||
|
||||
@@ -420,212 +282,6 @@ NODE_DISPLAY_NAME_MAPPINGS[
|
||||
"MZ_KohyaSSLoraTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(lora)"
|
||||
|
||||
|
||||
class MZ_KohyaSSControlnetTrain:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
models = [
|
||||
"latest",
|
||||
"empty",
|
||||
]
|
||||
|
||||
comfyui_full_m_path = []
|
||||
comfyui_basemodels = folder_paths.get_filename_list("controlnet")
|
||||
for b_model in comfyui_basemodels:
|
||||
m_path = folder_paths.get_full_path("controlnet", b_model)
|
||||
comfyui_full_m_path.append(m_path)
|
||||
|
||||
# 按创建时间排序
|
||||
comfyui_full_m_path = sorted(
|
||||
comfyui_full_m_path, key=lambda x: os.path.getctime(x), reverse=True)
|
||||
|
||||
models = models + comfyui_full_m_path
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"train_config": ("MZ_TT_SS_TrainConfig",),
|
||||
"base_controlnet": (models, {"default": "latest"}),
|
||||
"sample_generate": (["enable", "disable"], {"default": "enable"}),
|
||||
"sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
|
||||
},
|
||||
"optional": {
|
||||
"has_no_effect": (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_KohyaSSControlnetTrain"] = MZ_KohyaSSControlnetTrain
|
||||
NODE_DISPLAY_NAME_MAPPINGS[
|
||||
"MZ_KohyaSSControlnetTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(controlnet)"
|
||||
|
||||
|
||||
class MZ_KohyaSS_KohakuBlueleaf_HYHiDLoraTrain:
|
||||
KohyaSSLoraTrain = MZ_KohyaSSLoraTrain()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
models, folders, vae_models, unet_models, loras = HYDiT_paths()
|
||||
_kohya_input_types = s.KohyaSSLoraTrain.INPUT_TYPES()
|
||||
|
||||
# "workspace_config": kwargs["workspace_config"],
|
||||
# "train_config_template": kwargs["train_config_template"],
|
||||
# "ckpt_name": kwargs["ckpt_name"],
|
||||
# "max_train_steps": kwargs["max_train_steps"],
|
||||
# "max_train_epochs": kwargs["max_train_epochs"],
|
||||
# "save_every_n_epochs": kwargs["save_every_n_epochs"],
|
||||
# "learning_rate": kwargs["learning_rate"],
|
||||
|
||||
kohya_input_types = {"required": {}, "optional": {}}
|
||||
kohya_input_types["required"]["unet_path"] = (
|
||||
["auto"] + models + unet_models, {"default": "auto"})
|
||||
kohya_input_types["required"]["vae_ema_path"] = (
|
||||
["auto"] + folders + vae_models, {"default": "auto"})
|
||||
kohya_input_types["required"]["text_encoder_path"] = (
|
||||
["auto"] + folders, {"default": "auto"})
|
||||
kohya_input_types["required"]["tokenizer_path"] = (
|
||||
["auto"] + folders, {"default": "auto"})
|
||||
kohya_input_types["required"]["t5_encoder_path"] = (
|
||||
["none", "auto"] + folders, {"default": "none"})
|
||||
|
||||
for k, v in _kohya_input_types["required"].items():
|
||||
if k == "ckpt_name":
|
||||
continue
|
||||
if k == "ema_to_module":
|
||||
continue
|
||||
kohya_input_types["required"][k] = v
|
||||
|
||||
for k, v in _kohya_input_types["optional"].items():
|
||||
kohya_input_types["optional"][k] = v
|
||||
|
||||
return kohya_input_types
|
||||
|
||||
RETURN_TYPES = ()
|
||||
RETURN_NAMES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME + "/kohya_ss_HYDiT_kohakublueleaf"
|
||||
|
||||
def start(self, **kwargs):
|
||||
kwargs["hunyuan_models_config"] = {
|
||||
"unet_path": kwargs["unet_path"],
|
||||
"vae_ema_path": kwargs["vae_ema_path"],
|
||||
"text_encoder_path": kwargs["text_encoder_path"],
|
||||
"tokenizer_path": kwargs["tokenizer_path"],
|
||||
"t5_encoder_path": kwargs["t5_encoder_path"],
|
||||
}
|
||||
kwargs["ckpt_name"] = None
|
||||
return self.KohyaSSLoraTrain.start(**kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_KohyaSS_KohakuBlueleaf_HYHiDLoraTrain"] = MZ_KohyaSS_KohakuBlueleaf_HYHiDLoraTrain
|
||||
NODE_DISPLAY_NAME_MAPPINGS[
|
||||
"MZ_KohyaSS_KohakuBlueleaf_HYHiDLoraTrain"] = f"{AUTHOR_NAME} - KohyaSS_KohakuBlueleaf_HYHiDLoraTrain"
|
||||
|
||||
|
||||
class MZ_KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"lora_name": ("STRING", {"default": ""}),
|
||||
"branch": ("STRING", {"default": "0dc79edc01f2000de1dad5ad6d20d8b099bfafe2"}),
|
||||
"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_HYDiT_kohakublueleaf"
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
kwargs["branch_repoid"] = "KohakuBlueleaf/sd-scripts"
|
||||
kwargs["branch_local_name"] = "KohakuBlueleaf_kohya_ss_lora"
|
||||
return mz_train_tools_core.MZ_KohyaSSInitWorkspace_call(kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace"] = MZ_KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace
|
||||
NODE_DISPLAY_NAME_MAPPINGS[
|
||||
"MZ_KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace"] = f"{AUTHOR_NAME} - KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace"
|
||||
|
||||
|
||||
class MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
models, folders, vae_models, unet_models, _ = HYDiT_paths()
|
||||
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)
|
||||
return {
|
||||
"required": {
|
||||
"branch": ("STRING", {"default": "0dc79edc01f2000de1dad5ad6d20d8b099bfafe2"}),
|
||||
"source": ([
|
||||
"github",
|
||||
"githubfast",
|
||||
"521github",
|
||||
"kkgithub",
|
||||
], {"default": "github"}),
|
||||
"version": (["1.1", "1.2"], {"default": "1.2"}),
|
||||
"unet_path": (["auto"] + models + unet_models, {"default": "auto"}),
|
||||
"vae_ema_path": (["auto"] + folders + vae_models, {"default": "auto"}),
|
||||
"text_encoder_path": (["auto"] + folders, {"default": "auto"}),
|
||||
"tokenizer_path": (["auto"] + folders, {"default": "auto"}),
|
||||
"t5_encoder_path": (["none", "auto"] + folders, {"default": "none"}),
|
||||
"lora_path": (["none"] + comfyui_full_loras, {"default": "none"}),
|
||||
"seed": ("INT", {"default": 0}),
|
||||
"steps": ("INT", {"default": 20}),
|
||||
"cfg": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
|
||||
"scheduler": ([
|
||||
"euler_ancestral", "dpmpp_2m_sde"
|
||||
], {"default": "dpmpp_2m_sde"}),
|
||||
"prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
|
||||
"negative_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
|
||||
"width": ("INT", {"default": 1024, "max": 8192, "step": 16}),
|
||||
"height": ("INT", {"default": 1024, "max": 8192, "step": 16}),
|
||||
"keep_device": (["enable", "disable"], {"default": "enable"}),
|
||||
},
|
||||
"optional": {
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("images",)
|
||||
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME + "/kohya_ss_HYDiT_kohakublueleaf"
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
kwargs["branch_repoid"] = "KohakuBlueleaf/sd-scripts"
|
||||
kwargs["branch_local_name"] = "KohakuBlueleaf_kohya_ss_lora"
|
||||
return mz_train_tools_core.MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I_call(kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I"] = MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I
|
||||
NODE_DISPLAY_NAME_MAPPINGS[
|
||||
"MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I"] = f"{AUTHOR_NAME} - KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I"
|
||||
|
||||
|
||||
class MZ_LoadImagesFromDirectoryPath:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -668,304 +324,3 @@ class MZ_LoadImagesFromDirectoryPath:
|
||||
NODE_CLASS_MAPPINGS["MZ_LoadImagesFromDirectoryPath"] = MZ_LoadImagesFromDirectoryPath
|
||||
NODE_DISPLAY_NAME_MAPPINGS[
|
||||
"MZ_LoadImagesFromDirectoryPath"] = f"{AUTHOR_NAME} - LoadImagesFromDirectoryPath"
|
||||
|
||||
|
||||
from . import mz_train_tools_core_HYDiT
|
||||
|
||||
|
||||
class MZ_HYDiTInitWorkspace:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"train_name": ("STRING", {"default": ""}),
|
||||
"branch": ("STRING", {"default": "5657364143e44ac90f72aeb47b81bd505a95665d"}),
|
||||
"source": ([
|
||||
"github",
|
||||
"githubfast",
|
||||
"521github",
|
||||
"kkgithub",
|
||||
], {"default": "github"}),
|
||||
"seed": ("INT", {"default": 0}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MZ_TT_HYDiT_WorkspaceConfig",)
|
||||
RETURN_NAMES = ("workspace_config",)
|
||||
|
||||
FUNCTION = "start"
|
||||
|
||||
CATEGORY = CATEGORY_NAME + "/HYDiT_native"
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core_HYDiT)
|
||||
return mz_train_tools_core_HYDiT.MZ_HYDiTInitWorkspace_call(kwargs.copy())
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_HYDiTInitWorkspace"] = MZ_HYDiTInitWorkspace
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTInitWorkspace"] = f"{AUTHOR_NAME} - HYDiTInitWorkspace"
|
||||
|
||||
|
||||
class MZ_HYDiTDatasetConfig:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",),
|
||||
"resolution": ("INT", {"default": 1024}),
|
||||
"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}),
|
||||
},
|
||||
"optional": {
|
||||
"images": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("workspace_images_dir",)
|
||||
|
||||
FUNCTION = "start"
|
||||
|
||||
CATEGORY = CATEGORY_NAME + "/HYDiT_native"
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core_HYDiT)
|
||||
return mz_train_tools_core_HYDiT.MZ_HYDiTDatasetConfig_call(kwargs.copy())
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_HYDiTDatasetConfig"] = MZ_HYDiTDatasetConfig
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTDatasetConfig"] = f"{AUTHOR_NAME} - HYDiTDatasetConfig"
|
||||
|
||||
|
||||
class MZ_HYDiTAdvConfig:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"lr": ("STRING", {"default": "1e-5"}),
|
||||
"rope_real": (["enable", "disable"], {"default": "enable"}),
|
||||
# ['Wqkv', 'q_proj', 'kv_proj', 'out_proj']
|
||||
"target_modules_Wqkv": (["enable", "disable"], {"default": "enable"}),
|
||||
"target_modules_q_proj": (["enable", "disable"], {"default": "enable"}),
|
||||
"target_modules_kv_proj": (["enable", "disable"], {"default": "enable"}),
|
||||
"target_modules_out_proj": (["enable", "disable"], {"default": "enable"}),
|
||||
"warmup_min_lr": ("STRING", {"default": "1e-6"}),
|
||||
# parser.add_argument("--warmup-num-steps", type=float, default=0)
|
||||
# parser.add_argument("--weight-decay", type=float, default=0, help="weight-decay in optimizer")
|
||||
"weight_decay": ("FLOAT", {"default": 0}),
|
||||
"warmup_num_steps": ("FLOAT", {"default": 0}),
|
||||
# parser.add_argument("--uncond-p", type=float, default=0.2,
|
||||
# help="The probability of dropping training text used for CLIP feature extraction")
|
||||
# parser.add_argument("--uncond-p-t5", type=float, default=0.2,
|
||||
# help="The probability of dropping training text used for mT5 feature extraction")
|
||||
"uncond_p": ("FLOAT", {"default": 0.2}),
|
||||
"uncond_p_t5": ("FLOAT", {"default": 0.2}),
|
||||
|
||||
# parser.add_argument("--use-flash-attn", action="store_true", help="During training, "
|
||||
# "flash attention is used to accelerate training.")
|
||||
# parser.add_argument("--no-flash-attn", dest="use_flash_attn",
|
||||
# action="store_false", help="During training, flash attention is not used to accelerate training.")
|
||||
# parser.add_argument("--use-zero-stage", type=int, default=1, help="Use AngelPTM zero stage. Support 2 and 3")
|
||||
# parser.add_argument("--grad-accu-steps", type=int, default=1, help="Gradient accumulation steps.")
|
||||
"use_flash_attn": (["enable", "disable"], {"default": "disable"}),
|
||||
"use_zero_stage": ("INT", {"default": 2}),
|
||||
"grad_accu_steps": ("INT", {"default": 1}),
|
||||
# parser.add_argument("--extra-fp16", action="store_true", help="Use extra fp16 for vae and text_encoder.")
|
||||
"extra_fp16": (["enable", "disable"], {"default": "enable"}),
|
||||
# parser.add_argument("--qk-norm", action="store_true", help="Query Key normalization. See http://arxiv.org/abs/2302.05442 for details.")
|
||||
"qk_norm": (["enable", "disable"], {"default": "enable"}),
|
||||
# parser.add_argument("--norm", type=str, choices=["rms", "laryer"], default="layer", help="Normalization layer type")
|
||||
"norm": (["rms", "layer"], {"default": "layer"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MZ_TT_HYDiT_AdvConfig",)
|
||||
RETURN_NAMES = ("advanced_config",)
|
||||
|
||||
FUNCTION = "start"
|
||||
|
||||
CATEGORY = CATEGORY_NAME + "/HYDiT_native"
|
||||
|
||||
def start(self, **kwargs):
|
||||
return (kwargs.copy(),)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_HYDiTAdvConfig"] = MZ_HYDiTAdvConfig
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTAdvConfig"] = f"{AUTHOR_NAME} - HYDiTAdvConfig"
|
||||
|
||||
|
||||
def HYDiT_paths():
|
||||
hunyuan_models_path = os.path.join(
|
||||
Utils.get_comfyui_models_path(), "hunyuan")
|
||||
os.makedirs(hunyuan_models_path, exist_ok=True)
|
||||
|
||||
models = Utils.get_models_by_folder(hunyuan_models_path)
|
||||
|
||||
folders = Utils.get_folders_by_folder(hunyuan_models_path)
|
||||
|
||||
vae_models = Utils.get_models_by_folder(
|
||||
os.path.join(Utils.get_comfyui_models_path(), "vae"))
|
||||
unet_models = Utils.get_models_by_folder(
|
||||
os.path.join(Utils.get_comfyui_models_path(), "unet"))
|
||||
|
||||
workspaces_root = os.path.join(
|
||||
folder_paths.output_directory, "mz_train_workspaces")
|
||||
|
||||
loras = mz_train_tools_core_HYDiT.search_loras([
|
||||
workspaces_root,
|
||||
os.path.join(Utils.get_comfyui_models_path(), "loras"),
|
||||
])
|
||||
|
||||
return models, folders, vae_models, unet_models, loras
|
||||
|
||||
|
||||
class MZ_HYDiTTrain:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
models, folders, vae_models, unet_models, loras = HYDiT_paths()
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",),
|
||||
"unet_path": (["auto"] + models + unet_models, {"default": "auto"}),
|
||||
"ema_to_module": (["enable", "disable"], {"default": "enable"}),
|
||||
"vae_ema_path": (["auto"] + folders + vae_models, {"default": "auto"}),
|
||||
"text_encoder_path": (["auto"] + folders, {"default": "auto"}),
|
||||
"tokenizer_path": (["auto"] + folders, {"default": "auto"}),
|
||||
"t5_encoder_path": (["none", "auto"] + folders, {"default": "none"}),
|
||||
"resolution": ("INT", {"default": 1024, "step": 16}),
|
||||
"batch_size": ("INT", {"default": 1}),
|
||||
"epochs": ("INT", {"default": 50}),
|
||||
"ckpt_every": ("INT", {"default": 500}),
|
||||
"rank": ("INT", {"default": 8}),
|
||||
"base_lora": (["latest", "empty"] + loras, {"default": "latest"}),
|
||||
"sample_generate": (["enable", "disable"], {"default": "enable"}),
|
||||
"sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
|
||||
},
|
||||
"optional": {
|
||||
"advanced_config": ("MZ_TT_HYDiT_AdvConfig",),
|
||||
"workspace_images_dir": ("STRING", {"forceInput": True}),
|
||||
"has_no_effect": (AlwaysEqualProxy("*"),),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
RETURN_NAMES = ()
|
||||
|
||||
FUNCTION = "start"
|
||||
|
||||
MZ_DESC = """
|
||||
base_size The base resolution (n, n) from which to create multiple resolutions | Recommended values: 256/512/1024
|
||||
"""
|
||||
|
||||
OUTPUT_NODE = True
|
||||
|
||||
CATEGORY = CATEGORY_NAME + "/HYDiT_native"
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core_HYDiT)
|
||||
return mz_train_tools_core_HYDiT.MZ_HYDiTTrain_call(kwargs.copy())
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_HYDiTTrain"] = MZ_HYDiTTrain
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTTrain"] = f"{AUTHOR_NAME} - HYDiTTrain"
|
||||
|
||||
|
||||
class MZ_HYDiTSimpleT2I:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
hunyuan_models_path = os.path.join(
|
||||
Utils.get_comfyui_models_path(), "hunyuan")
|
||||
os.makedirs(hunyuan_models_path, exist_ok=True)
|
||||
|
||||
models = Utils.get_models_by_folder(hunyuan_models_path)
|
||||
folders = Utils.get_folders_by_folder(hunyuan_models_path)
|
||||
|
||||
vae_models = Utils.get_models_by_folder(
|
||||
os.path.join(Utils.get_comfyui_models_path(), "vae"))
|
||||
unet_models = Utils.get_models_by_folder(
|
||||
os.path.join(Utils.get_comfyui_models_path(), "unet"))
|
||||
|
||||
comfyui_full_loras = mz_train_tools_core_HYDiT.search_loras([
|
||||
os.path.join(Utils.get_comfyui_models_path(), "loras"),
|
||||
])
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"branch": ("STRING", {"default": "5657364143e44ac90f72aeb47b81bd505a95665d"}),
|
||||
"source": ([
|
||||
"github",
|
||||
"githubfast",
|
||||
"521github",
|
||||
"kkgithub",
|
||||
], {"default": "github"}),
|
||||
"unet_path": (["auto"] + models + unet_models, {"default": "auto"}),
|
||||
"vae_ema_path": (["auto"] + folders + vae_models, {"default": "auto"}),
|
||||
"text_encoder_path": (["auto"] + folders, {"default": "auto"}),
|
||||
"tokenizer_path": (["auto"] + folders, {"default": "auto"}),
|
||||
"t5_encoder_path": (["none", "auto"] + folders, {"default": "auto"}),
|
||||
"lora_path": (["none"] + comfyui_full_loras, {"default": "none"}),
|
||||
"seed": ("INT", {"default": 0}),
|
||||
"steps": ("INT", {"default": 20}),
|
||||
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
|
||||
"scheduler": ([
|
||||
"ddpm", "ddim", "dpmms", "uni_pc", "dpmpp_2m_karras"
|
||||
], {"default": "ddpm"}),
|
||||
"prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
|
||||
"negative_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
|
||||
"width": ("INT", {"default": 512, "max": 8192, "step": 16}),
|
||||
"height": ("INT", {"default": 512, "max": 8192, "step": 16}),
|
||||
"keep_device": (["enable", "disable"], {"default": "enable"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
|
||||
RETURN_NAMES = ("images",)
|
||||
|
||||
FUNCTION = "start"
|
||||
|
||||
CATEGORY = CATEGORY_NAME + "/HYDiT_native"
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core_HYDiT)
|
||||
return mz_train_tools_core_HYDiT.MZ_HYDiTSimpleT2I_call(kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_HYDiTSimpleT2I"] = MZ_HYDiTSimpleT2I
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTSimpleT2I"] = f"{AUTHOR_NAME} - HYDiTSimpleT2I"
|
||||
|
||||
|
||||
class MZ_TrainToolsDebug:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"object": (AlwaysEqualProxy("*"),),
|
||||
"indent": ("INT", {"default": 2}),
|
||||
"depth": ("INT", {"default": 5}),
|
||||
"width": ("INT", {"default": 80}),
|
||||
"compact": (["enable", "disable"], {"default": "enable"}),
|
||||
"sort_keys": (["enable", "disable"], {"default": "enable"}),
|
||||
"underscore_numbers": (["enable", "disable"], {"default": "enable"}),
|
||||
"index": ("STRING", {"default": ""}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("debug",)
|
||||
|
||||
FUNCTION = "start"
|
||||
|
||||
CATEGORY = CATEGORY_NAME
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
return mz_train_tools_core.MZ_TrainToolsDebug_call(kwargs.copy())
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_TrainToolsDebug"] = MZ_TrainToolsDebug
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_TrainToolsDebug"] = f"{AUTHOR_NAME} - TrainToolsDebug"
|
||||
|
||||
@@ -1,155 +0,0 @@
|
||||
import argparse
|
||||
|
||||
from index_kits.dataset.make_dataset_core import startup, make_multireso
|
||||
from index_kits.common import show_index_info
|
||||
from index_kits import __version__
|
||||
|
||||
|
||||
def common_args(parser):
|
||||
parser.add_argument('-t', '--target', type=str,
|
||||
required=True, help='Save path')
|
||||
|
||||
|
||||
def get_args():
|
||||
parser = argparse.ArgumentParser(description="""
|
||||
IndexKits is a tool to build and manage index files for large-scale datasets.
|
||||
It supports both base index and multi-resolution index.
|
||||
|
||||
Introduction
|
||||
------------
|
||||
This command line tool provides the following functionalities:
|
||||
1. Show index v2 information
|
||||
2. Build base index v2
|
||||
3. Build multi-resolution index v2
|
||||
|
||||
Examples
|
||||
--------
|
||||
1. Show index v2 information
|
||||
index_kits show /path/to/index.json
|
||||
|
||||
2. Build base index v2
|
||||
Default usage:
|
||||
index_kits base -c /path/to/config.yaml -t /path/to/index.json
|
||||
|
||||
Use multiple processes:
|
||||
index_kits base -c /path/to/config.yaml -t /path/to/index.json -w 40
|
||||
|
||||
3. Build multi-resolution index v2
|
||||
|
||||
Build with a configuration file:
|
||||
index_kits multireso -c /path/to/config.yaml -t /path/to/index_mb_gt512.json
|
||||
|
||||
Build by specifying arguments without a configuration file:
|
||||
index_kits multireso --src /path/to/index.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json
|
||||
|
||||
Build by specifying target-ratios:
|
||||
index_kits multireso --src /path/to/index.json --base-size 512 --target-ratios 1:1 4:3 3:4 16:9 9:16 --min-size 512 -t /path/to/index_mb_gt512.json
|
||||
|
||||
Build with multiple source index files.
|
||||
index_kits multireso --src /path/to/index1.json /path/to/index2.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json
|
||||
""", formatter_class=argparse.RawTextHelpFormatter)
|
||||
sub_parsers = parser.add_subparsers(dest='task', required=True)
|
||||
|
||||
# Show index message
|
||||
show_parser = sub_parsers.add_parser('show', description="""
|
||||
Show base/multireso index v2 information.
|
||||
|
||||
Example
|
||||
-------
|
||||
index_kits show /path/to/index.json
|
||||
""", formatter_class=argparse.RawTextHelpFormatter)
|
||||
show_parser.add_argument(
|
||||
'src', type=str, help='Path to a base/multireso index file.')
|
||||
show_parser.add_argument(
|
||||
'--arrow-files', action='store_true', help='Show arrow files only.')
|
||||
show_parser.add_argument('--depth', type=int, default=1,
|
||||
help='Arrow file depth. Default is 1, the level of last folder in the arrow file path. '
|
||||
'Set it to 0 to show the full path including `xxx/last_folder/*.arrow`.')
|
||||
|
||||
# Single resolution bucket
|
||||
base_parser = sub_parsers.add_parser('base', description="""
|
||||
Build base index v2.
|
||||
|
||||
Example
|
||||
-------
|
||||
index_kits base -c /path/to/config.yaml -t /path/to/index.json
|
||||
""", formatter_class=argparse.RawTextHelpFormatter)
|
||||
base_parser.add_argument('-c', '--config', type=str,
|
||||
required=True, help='Configuration file path')
|
||||
common_args(base_parser)
|
||||
base_parser.add_argument('-w', '--world-size', type=int, default=1)
|
||||
base_parser.add_argument('--work-dir', type=str,
|
||||
default='.', help='Work directory')
|
||||
base_parser.add_argument('--use-cache', action='store_true', help='Use cache to avoid reprocessing. '
|
||||
'Perform merge pkl results directly.')
|
||||
|
||||
# Multi-resolution bucket
|
||||
mo_parser = sub_parsers.add_parser('multireso', description="""
|
||||
Build multi-resolution index v2
|
||||
|
||||
Example
|
||||
-------
|
||||
Build with a configuration file:
|
||||
index_kits multireso -c /path/to/config.yaml -t /path/to/index_mb_gt512.json
|
||||
|
||||
Build by specifying arguments without a configuration file:
|
||||
index_kits multireso --src /path/to/index.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json
|
||||
|
||||
Build by specifying target-ratios:
|
||||
index_kits multireso --src /path/to/index.json --base-size 512 --target-ratios 1:1 4:3 3:4 16:9 9:16 --min-size 512 -t /path/to/index_mb_gt512.json
|
||||
|
||||
Build with multiple source index files.
|
||||
index_kits multireso --src /path/to/index1.json /path/to/index2.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json
|
||||
""", formatter_class=argparse.RawTextHelpFormatter)
|
||||
mo_parser.add_argument('-c', '--config', type=str, default=None,
|
||||
help='Configuration file path in a yaml format. Either --config or --src must be provided.')
|
||||
mo_parser.add_argument('-s', '--src', type=str, nargs='+', default=None,
|
||||
help='Source index files. Either --config or --src must be provided.')
|
||||
common_args(mo_parser)
|
||||
mo_parser.add_argument('--base-size', type=int, default=None,
|
||||
help="Base size. Typically set as 256/512/1024 according to image size you train model.")
|
||||
mo_parser.add_argument('--reso-step', type=int, default=None,
|
||||
help="Resolution step. Either reso_step or target_ratios must be provided.")
|
||||
mo_parser.add_argument('--target-ratios', type=str, nargs='+', default=None,
|
||||
help="Target ratios. Either reso_step or target_ratios must be provided.")
|
||||
mo_parser.add_argument('--md5-file', type=str, default=None,
|
||||
help='You can provide an md5 to height and width file to accelerate the process. '
|
||||
'It is a pickle file that contains a dict, which maps md5 to (height, width) tuple.')
|
||||
mo_parser.add_argument('--align', type=int, default=16,
|
||||
help="Used when --target-ratios is provided. Align size of source image height and width.")
|
||||
mo_parser.add_argument('--min-size', type=int, default=0,
|
||||
help="Minimum size. Images smaller than this size will be ignored.")
|
||||
|
||||
# Common
|
||||
parser.add_argument('-v', '--version', action='version',
|
||||
version=f'%(prog)s {__version__}')
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = get_args()
|
||||
if args.task == 'show':
|
||||
show_index_info(args.src,
|
||||
args.arrow_files,
|
||||
args.depth,
|
||||
)
|
||||
elif args.task == 'base':
|
||||
startup(args.config,
|
||||
args.target,
|
||||
args.world_size,
|
||||
args.work_dir,
|
||||
use_cache=args.use_cache,
|
||||
)
|
||||
elif args.task == 'multireso':
|
||||
make_multireso(args.target,
|
||||
args.config,
|
||||
args.src,
|
||||
args.base_size,
|
||||
args.reso_step,
|
||||
args.target_ratios,
|
||||
args.align,
|
||||
args.min_size,
|
||||
args.md5_file,
|
||||
)
|
||||
@@ -1,595 +0,0 @@
|
||||
import gc
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
from functools import partial
|
||||
from glob import glob
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import safetensors.torch
|
||||
|
||||
import deepspeed
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.distributed.optim import ZeroRedundancyOptimizer
|
||||
from torchvision.transforms import functional as TF
|
||||
from diffusers.models import AutoencoderKL
|
||||
from transformers import BertModel, BertTokenizer, logging as tf_logging
|
||||
|
||||
from hydit.config import get_args
|
||||
from hydit.lr_scheduler import WarmupLR
|
||||
from hydit.data_loader.arrow_load_stream import TextImageArrowStream
|
||||
from hydit.diffusion import create_diffusion
|
||||
from hydit.ds_config import deepspeed_config_from_args
|
||||
from hydit.modules.ema import EMA
|
||||
from hydit.modules.fp16_layers import Float16Module
|
||||
from hydit.modules.models import HUNYUAN_DIT_MODELS
|
||||
from hydit.modules.posemb_layers import init_image_posemb
|
||||
from hydit.utils.tools import create_logger, set_seeds, create_exp_folder, get_trainable_params
|
||||
from IndexKits.index_kits import ResolutionGroup
|
||||
from IndexKits.index_kits.sampler import DistributedSamplerWithStartIndex, BlockDistributedSampler
|
||||
from peft import LoraConfig, get_peft_model
|
||||
|
||||
|
||||
def deepspeed_initialize(args, logger, model, opt, deepspeed_config):
|
||||
logger.info(f"Initialize deepspeed...")
|
||||
logger.info(f" Using deepspeed optimizer")
|
||||
|
||||
def get_learning_rate_scheduler(warmup_min_lr, lr, warmup_num_steps, opt):
|
||||
return WarmupLR(opt, warmup_min_lr, lr, warmup_num_steps)
|
||||
|
||||
logger.info(
|
||||
f" Building scheduler with warmup_min_lr={args.warmup_min_lr}, warmup_num_steps={args.warmup_num_steps}")
|
||||
logger.info(
|
||||
f" deepspeed_config={deepspeed_config}")
|
||||
|
||||
model, opt, _, scheduler = deepspeed.initialize(
|
||||
model=model,
|
||||
model_parameters=get_trainable_params(model),
|
||||
config_params=deepspeed_config,
|
||||
args=args,
|
||||
lr_scheduler=partial(get_learning_rate_scheduler, args.warmup_min_lr,
|
||||
args.lr, args.warmup_num_steps) if args.warmup_num_steps > 0 else None,
|
||||
)
|
||||
return model, opt, scheduler
|
||||
|
||||
|
||||
def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoint_dir):
|
||||
def save_lora_weight(checkpoint_dir, client_state, tag=f"{train_steps:07d}.pt"):
|
||||
cur_ckpt_save_dir = f"{checkpoint_dir}/{tag}"
|
||||
if rank == 0:
|
||||
if args.use_fp16:
|
||||
model.module.module.save_pretrained(cur_ckpt_save_dir)
|
||||
else:
|
||||
model.module.save_pretrained(cur_ckpt_save_dir)
|
||||
|
||||
checkpoint_path = "[Not rank 0. Disabled output.]"
|
||||
|
||||
client_state = {
|
||||
"steps": train_steps,
|
||||
"epoch": epoch,
|
||||
"args": args
|
||||
}
|
||||
if ema is not None:
|
||||
client_state['ema'] = ema.state_dict()
|
||||
|
||||
dst_paths = []
|
||||
if train_steps % args.ckpt_every == 0:
|
||||
checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
|
||||
try:
|
||||
if args.training_parts == "lora":
|
||||
save_lora_weight(checkpoint_dir, client_state,
|
||||
tag=f"{train_steps:07d}.pt")
|
||||
else:
|
||||
model.save_checkpoint(
|
||||
checkpoint_dir, client_state=client_state, tag=f"{train_steps:07d}.pt")
|
||||
dst_paths.append(checkpoint_path)
|
||||
logger.info(f"Saved checkpoint to {checkpoint_path}")
|
||||
except:
|
||||
logger.error(f"Saved failed to {checkpoint_path}")
|
||||
|
||||
if train_steps % args.ckpt_latest_every == 0 or train_steps == args.max_training_steps:
|
||||
save_name = "latest.pt"
|
||||
checkpoint_path = f"{checkpoint_dir}/{save_name}"
|
||||
try:
|
||||
if args.training_parts == "lora":
|
||||
save_lora_weight(checkpoint_dir, client_state,
|
||||
tag=f"{save_name}")
|
||||
else:
|
||||
model.save_checkpoint(
|
||||
checkpoint_dir, client_state=client_state, tag=f"{save_name}")
|
||||
dst_paths.append(checkpoint_path)
|
||||
logger.info(f"Saved checkpoint to {checkpoint_path}")
|
||||
except:
|
||||
logger.error(f"Saved failed to {checkpoint_path}")
|
||||
|
||||
dist.barrier()
|
||||
if rank == 0 and len(dst_paths) > 0:
|
||||
# Delete optimizer states to avoid occupying too much disk space.
|
||||
for dst_path in dst_paths:
|
||||
for opt_state_path in glob(f"{dst_path}/zero_dp_rank_*_tp_rank_00_pp_rank_00_optim_states.pt"):
|
||||
os.remove(opt_state_path)
|
||||
|
||||
return checkpoint_path
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def prepare_model_inputs(args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img):
|
||||
try:
|
||||
from .hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar
|
||||
except:
|
||||
from hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar
|
||||
|
||||
image, text_embedding, text_embedding_mask, text_embedding_t5, text_embedding_mask_t5, kwargs = batch
|
||||
|
||||
# clip & mT5 text embedding
|
||||
text_embedding = text_embedding.to(device)
|
||||
text_embedding_mask = text_embedding_mask.to(device)
|
||||
encoder_hidden_states = text_encoder(
|
||||
text_embedding.to(device),
|
||||
attention_mask=text_embedding_mask.to(device),
|
||||
)[0]
|
||||
text_embedding_t5 = text_embedding_t5.to(device).squeeze(1)
|
||||
text_embedding_mask_t5 = text_embedding_mask_t5.to(device).squeeze(1)
|
||||
with torch.no_grad():
|
||||
output_t5 = text_encoder_t5(
|
||||
input_ids=text_embedding_t5,
|
||||
attention_mask=text_embedding_mask_t5 if T5_ENCODER['attention_mask'] else None,
|
||||
output_hidden_states=True
|
||||
)
|
||||
encoder_hidden_states_t5 = output_t5['hidden_states'][T5_ENCODER['layer_index']].detach(
|
||||
)
|
||||
|
||||
# additional condition
|
||||
image_meta_size = kwargs['image_meta_size'].to(device)
|
||||
style = kwargs['style'].to(device)
|
||||
|
||||
if args.extra_fp16:
|
||||
image = image.half()
|
||||
image_meta_size = image_meta_size.half() if image_meta_size is not None else None
|
||||
|
||||
# Map input images to latent space + normalize latents:
|
||||
image = image.to(device)
|
||||
vae_scaling_factor = vae.config.scaling_factor
|
||||
latents = vae.encode(image).latent_dist.sample().mul_(vae_scaling_factor)
|
||||
|
||||
# positional embedding
|
||||
_, _, height, width = image.shape
|
||||
reso = f"{height}x{width}"
|
||||
cos_cis_img, sin_cis_img = freqs_cis_img[reso]
|
||||
|
||||
# Model conditions
|
||||
model_kwargs = dict(
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
text_embedding_mask=text_embedding_mask,
|
||||
encoder_hidden_states_t5=encoder_hidden_states_t5,
|
||||
text_embedding_mask_t5=text_embedding_mask_t5,
|
||||
image_meta_size=image_meta_size,
|
||||
style=style,
|
||||
cos_cis_img=cos_cis_img,
|
||||
sin_cis_img=sin_cis_img,
|
||||
)
|
||||
|
||||
return latents, model_kwargs
|
||||
|
||||
|
||||
def Core(args):
|
||||
if args.training_parts == "lora":
|
||||
args.use_ema = False
|
||||
|
||||
assert torch.cuda.is_available(), "Training currently requires at least one GPU."
|
||||
|
||||
dist.init_process_group("nccl")
|
||||
world_size = dist.get_world_size()
|
||||
batch_size = args.batch_size
|
||||
grad_accu_steps = args.grad_accu_steps
|
||||
global_batch_size = world_size * batch_size * grad_accu_steps
|
||||
|
||||
rank = dist.get_rank()
|
||||
device = rank % torch.cuda.device_count()
|
||||
seed = args.global_seed * world_size + rank
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.cuda.set_device(device)
|
||||
print(f"Starting rank={rank}, seed={seed}, world_size={world_size}.")
|
||||
deepspeed_config = deepspeed_config_from_args(args, global_batch_size)
|
||||
|
||||
# Setup an experiment folder
|
||||
experiment_dir, checkpoint_dir, logger = create_exp_folder(args, rank)
|
||||
|
||||
# Log all the arguments
|
||||
logger.info(sys.argv)
|
||||
logger.info(str(args))
|
||||
# Save to a json file
|
||||
args_dict = vars(args)
|
||||
args_dict['world_size'] = world_size
|
||||
with open(f"{experiment_dir}/args.json", 'w') as f:
|
||||
json.dump(args_dict, f, indent=4)
|
||||
|
||||
# Disable the message "Some weights of the model checkpoint at ... were not used when initializing BertModel."
|
||||
# If needed, just comment the following line.
|
||||
tf_logging.set_verbosity_error()
|
||||
|
||||
# ===========================================================================
|
||||
# Building HYDIT
|
||||
# ===========================================================================
|
||||
|
||||
logger.info("Building HYDIT Model.")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Training sample base size, such as 256/512/1024. Notice that this size is
|
||||
# just a base size, not necessary the actual size of training samples. Actual
|
||||
# size of the training samples are correlated with `resolutions` when enabling
|
||||
# multi-resolution training.
|
||||
# ---------------------------------------------------------------------------
|
||||
image_size = args.image_size
|
||||
if len(image_size) == 1:
|
||||
image_size = [image_size[0], image_size[0]]
|
||||
if len(image_size) != 2:
|
||||
raise ValueError(f"Invalid image size: {args.image_size}")
|
||||
assert image_size[0] % 8 == 0 and image_size[1] % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder). " \
|
||||
f"got {image_size}"
|
||||
latent_size = [image_size[0] // 8, image_size[1] // 8]
|
||||
|
||||
# initialize model by deepspeed
|
||||
assert args.deepspeed, f"Must enable deepspeed in this script: train_deepspeed.py"
|
||||
with deepspeed.zero.Init(data_parallel_group=torch.distributed.group.WORLD,
|
||||
remote_device=None if args.remote_device == 'none' else args.remote_device,
|
||||
config_dict_or_path=deepspeed_config,
|
||||
mpu=None,
|
||||
enabled=args.zero_stage == 3):
|
||||
model = HUNYUAN_DIT_MODELS[args.model](args,
|
||||
input_size=latent_size,
|
||||
log_fn=logger.info,
|
||||
)
|
||||
# Multi-resolution / Single-resolution training.
|
||||
if args.multireso:
|
||||
resolutions = ResolutionGroup(image_size[0],
|
||||
align=16,
|
||||
step=args.reso_step,
|
||||
target_ratios=args.target_ratios).data
|
||||
else:
|
||||
resolutions = ResolutionGroup(image_size[0],
|
||||
align=16,
|
||||
target_ratios=['1:1']).data
|
||||
|
||||
freqs_cis_img = init_image_posemb(args.rope_img,
|
||||
resolutions=resolutions,
|
||||
patch_size=model.patch_size,
|
||||
hidden_size=model.hidden_size,
|
||||
num_heads=model.num_heads,
|
||||
log_fn=logger.info,
|
||||
rope_real=args.rope_real,
|
||||
)
|
||||
|
||||
# Create EMA model and convert to fp16 if needed.
|
||||
ema = None
|
||||
if args.use_ema:
|
||||
ema = EMA(args, model, device, logger)
|
||||
|
||||
# Setup FP16 main model:
|
||||
if args.use_fp16:
|
||||
model = Float16Module(model, args)
|
||||
logger.info(
|
||||
f" Using main model with data type {'fp16' if args.use_fp16 else 'fp32'}")
|
||||
|
||||
diffusion = create_diffusion(
|
||||
noise_schedule=args.noise_schedule,
|
||||
predict_type=args.predict_type,
|
||||
learn_sigma=args.learn_sigma,
|
||||
mse_loss_weight_type=args.mse_loss_weight_type,
|
||||
beta_start=args.beta_start,
|
||||
beta_end=args.beta_end,
|
||||
noise_offset=args.noise_offset,
|
||||
)
|
||||
|
||||
try:
|
||||
from .hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar, CustomizeEmbeds
|
||||
except:
|
||||
from hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar, CustomizeEmbeds
|
||||
|
||||
# Setup VAE
|
||||
logger.info(f" Loading vae from {VAE_EMA_PATH}")
|
||||
vae = AutoencoderKL.from_pretrained(VAE_EMA_PATH)
|
||||
# Setup BERT text encoder
|
||||
logger.info(f" Loading Bert text encoder from {TEXT_ENCODER}")
|
||||
text_encoder = BertModel.from_pretrained(
|
||||
TEXT_ENCODER, False, revision=None)
|
||||
# Setup BERT tokenizer:
|
||||
logger.info(f" Loading Bert tokenizer from {TOKENIZER}")
|
||||
tokenizer = BertTokenizer.from_pretrained(TOKENIZER)
|
||||
# Setup T5 text encoder
|
||||
from hydit.modules.text_encoder import MT5Embedder
|
||||
mt5_path = T5_ENCODER['MT5']
|
||||
if mt5_path is None:
|
||||
embedder_t5 = CustomizeEmbeds()
|
||||
else:
|
||||
embedder_t5 = MT5Embedder(
|
||||
mt5_path, torch_dtype=T5_ENCODER['torch_dtype'], max_length=args.text_len_t5)
|
||||
tokenizer_t5 = embedder_t5.tokenizer
|
||||
text_encoder_t5 = embedder_t5.model
|
||||
|
||||
if args.extra_fp16:
|
||||
logger.info(f" Using fp16 for extra modules: vae, text_encoder")
|
||||
vae = vae.half().to(device)
|
||||
text_encoder = text_encoder.half().to(device)
|
||||
text_encoder_t5 = text_encoder_t5.half().to(device)
|
||||
else:
|
||||
vae = vae.to(device)
|
||||
text_encoder = text_encoder.to(device)
|
||||
text_encoder_t5 = text_encoder_t5.to(device)
|
||||
|
||||
logger.info(
|
||||
f" Optimizer parameters: lr={args.lr}, weight_decay={args.weight_decay}")
|
||||
logger.info(" Using deepspeed optimizer")
|
||||
opt = None
|
||||
|
||||
# ===========================================================================
|
||||
# Building Dataset
|
||||
# ===========================================================================
|
||||
|
||||
logger.info(f"Building Streaming Dataset.")
|
||||
logger.info(f" Loading index file {args.index_file} (v2)")
|
||||
|
||||
dataset = TextImageArrowStream(args=args,
|
||||
resolution=image_size[0],
|
||||
random_flip=args.random_flip,
|
||||
log_fn=logger.info,
|
||||
index_file=args.index_file,
|
||||
multireso=args.multireso,
|
||||
batch_size=batch_size,
|
||||
world_size=world_size,
|
||||
random_shrink_size_cond=args.random_shrink_size_cond,
|
||||
merge_src_cond=args.merge_src_cond,
|
||||
uncond_p=args.uncond_p,
|
||||
text_ctx_len=args.text_len,
|
||||
tokenizer=tokenizer,
|
||||
uncond_p_t5=args.uncond_p_t5,
|
||||
text_ctx_len_t5=args.text_len_t5,
|
||||
tokenizer_t5=tokenizer_t5,
|
||||
)
|
||||
if args.multireso:
|
||||
sampler = BlockDistributedSampler(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed,
|
||||
shuffle=False, drop_last=True, batch_size=batch_size)
|
||||
else:
|
||||
sampler = DistributedSamplerWithStartIndex(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed,
|
||||
shuffle=False, drop_last=True)
|
||||
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, sampler=sampler,
|
||||
num_workers=args.num_workers, pin_memory=True, drop_last=True)
|
||||
logger.info(f" Dataset contains {len(dataset):,} images.")
|
||||
logger.info(f" Index file: {args.index_file}.")
|
||||
if args.multireso:
|
||||
logger.info(f' Using MultiResolutionBucketIndexV2 with step {dataset.index_manager.step} '
|
||||
f'and base size {dataset.index_manager.base_size}')
|
||||
logger.info(f'\n {dataset.index_manager.resolutions}')
|
||||
|
||||
# ===========================================================================
|
||||
# Loading parameter
|
||||
# ===========================================================================
|
||||
|
||||
logger.info(f"Loading parameter")
|
||||
start_epoch = 0
|
||||
start_epoch_step = 0
|
||||
train_steps = 0
|
||||
# Resume checkpoint if needed
|
||||
# if args.resume is not None or len(args.resume) > 0:
|
||||
if True:
|
||||
model, ema, start_epoch, start_epoch_step, train_steps = model_resume(
|
||||
args, model, ema, logger)
|
||||
|
||||
if args.training_parts == "lora":
|
||||
lora_ckpt = args.lora_ckpt
|
||||
if lora_ckpt is not None:
|
||||
lastest_checkpoint = lora_ckpt
|
||||
from peft.peft_model import PeftModel
|
||||
print(f"Loading lora model from {lastest_checkpoint}")
|
||||
if args.use_fp16:
|
||||
model.module = PeftModel.from_pretrained(
|
||||
model.module, lastest_checkpoint, is_trainable=True)
|
||||
else:
|
||||
model = PeftModel.from_pretrained(
|
||||
model, lastest_checkpoint, is_trainable=True)
|
||||
else:
|
||||
loraconfig = LoraConfig(
|
||||
r=args.rank,
|
||||
lora_alpha=args.rank,
|
||||
target_modules=args.target_modules
|
||||
)
|
||||
|
||||
if args.use_fp16:
|
||||
model.module = get_peft_model(model.module, loraconfig)
|
||||
else:
|
||||
model = get_peft_model(model, loraconfig)
|
||||
|
||||
logger.info(f" Training parts: {args.training_parts}")
|
||||
|
||||
model, opt, scheduler = deepspeed_initialize(
|
||||
args, logger, model, opt, deepspeed_config)
|
||||
|
||||
# ===========================================================================
|
||||
# Training
|
||||
# ===========================================================================
|
||||
|
||||
# print model structure
|
||||
# for name, param in model.named_parameters():
|
||||
# print(name, param.size())
|
||||
# raise Exception("stop")
|
||||
|
||||
model.train()
|
||||
if args.use_ema:
|
||||
ema.eval()
|
||||
|
||||
print(f" Worker {rank} ready.")
|
||||
dist.barrier()
|
||||
|
||||
iters_per_epoch = len(loader)
|
||||
logger.info(
|
||||
" ****************************** Running training ******************************")
|
||||
logger.info(f" Number GPUs: {world_size}")
|
||||
logger.info(f" Number training samples: {len(dataset):,}")
|
||||
logger.info(
|
||||
f" Number parameters: {sum(p.numel() for p in model.parameters()):,}")
|
||||
logger.info(
|
||||
f" Number trainable params: {sum(p.numel() for p in get_trainable_params(model)):,}")
|
||||
logger.info(
|
||||
" ------------------------------------------------------------------------------")
|
||||
logger.info(f" Iters per epoch: {iters_per_epoch:,}")
|
||||
logger.info(f" Batch size per device: {batch_size}")
|
||||
logger.info(
|
||||
f" Batch size all device: {batch_size * world_size * grad_accu_steps:,} (world_size * batch_size * grad_accu_steps)")
|
||||
logger.info(f" Gradient Accu steps: {args.grad_accu_steps}")
|
||||
logger.info(
|
||||
f" Total optimization steps: {args.epochs * iters_per_epoch // grad_accu_steps:,}")
|
||||
|
||||
logger.info(
|
||||
f" Training epochs: {start_epoch}/{args.epochs}")
|
||||
logger.info(
|
||||
f" Training epoch steps: {start_epoch_step:,}/{iters_per_epoch:,}")
|
||||
logger.info(
|
||||
f" Training total steps: {train_steps:,}/{min(args.max_training_steps, args.epochs * iters_per_epoch):,}")
|
||||
logger.info(
|
||||
" ------------------------------------------------------------------------------")
|
||||
logger.info(f" Noise schedule: {args.noise_schedule}")
|
||||
logger.info(
|
||||
f" Beta limits: ({args.beta_start}, {args.beta_end})")
|
||||
logger.info(f" Learn sigma: {args.learn_sigma}")
|
||||
logger.info(f" Prediction type: {args.predict_type}")
|
||||
logger.info(f" Noise offset: {args.noise_offset}")
|
||||
|
||||
logger.info(
|
||||
" ------------------------------------------------------------------------------")
|
||||
logger.info(
|
||||
f" Using EMA model: {args.use_ema} ({args.ema_dtype})")
|
||||
if args.use_ema:
|
||||
logger.info(
|
||||
f" Using EMA decay: {ema.max_value if args.use_ema else None}")
|
||||
logger.info(
|
||||
f" Using EMA warmup power: {ema.power if args.use_ema else None}")
|
||||
logger.info(f" Using main model fp16: {args.use_fp16}")
|
||||
logger.info(f" Using extra modules fp16: {args.extra_fp16}")
|
||||
logger.info(
|
||||
" ------------------------------------------------------------------------------")
|
||||
logger.info(f" Experiment directory: {experiment_dir}")
|
||||
logger.info(
|
||||
" *******************************************************************************")
|
||||
|
||||
if args.gc_interval > 0:
|
||||
gc.disable()
|
||||
gc.collect()
|
||||
|
||||
# Variables for monitoring/logging purposes:
|
||||
log_steps = 0
|
||||
running_loss = 0
|
||||
start_time = time.time()
|
||||
|
||||
if args.async_ema:
|
||||
ema_stream = torch.cuda.Stream()
|
||||
|
||||
easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5,
|
||||
target_height=768, target_width=1280, train_steps=0)
|
||||
pbar = PBar(args.epochs * len(loader))
|
||||
|
||||
# Training loop
|
||||
for epoch in range(start_epoch, args.epochs):
|
||||
logger.info(f" Start random shuffle with seed={seed}")
|
||||
# Makesure all processors use the same seed to shuffle dataset.
|
||||
dataset.shuffle(seed=args.global_seed + epoch, fast=True)
|
||||
logger.info(f" End of random shuffle")
|
||||
|
||||
# Move sampler to start_index
|
||||
if not args.multireso:
|
||||
start_index = start_epoch_step * world_size * batch_size
|
||||
if start_index != sampler.start_index:
|
||||
sampler.start_index = start_index
|
||||
# Reset start_epoch_step to zero, to ensure next epoch will start from the beginning.
|
||||
start_epoch_step = 0
|
||||
logger.info(f" Iters left this epoch: {len(loader):,}")
|
||||
|
||||
logger.info(f" Beginning epoch {epoch}...")
|
||||
step = 0
|
||||
for batch in loader:
|
||||
step += 1
|
||||
|
||||
latents, model_kwargs = prepare_model_inputs(
|
||||
args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img)
|
||||
|
||||
# training model by deepspeed while use fp16
|
||||
if args.use_fp16:
|
||||
if args.use_ema and args.async_ema:
|
||||
with torch.cuda.stream(ema_stream):
|
||||
ema.update(model.module.module, step=step)
|
||||
torch.cuda.current_stream().wait_stream(ema_stream)
|
||||
|
||||
loss_dict = diffusion.training_losses(
|
||||
model=model, x_start=latents, model_kwargs=model_kwargs)
|
||||
loss = loss_dict["loss"].mean()
|
||||
model.backward(loss)
|
||||
last_batch_iteration = (
|
||||
train_steps + 1) // (global_batch_size // (batch_size * world_size))
|
||||
model.step(
|
||||
lr_kwargs={'last_batch_iteration': last_batch_iteration})
|
||||
|
||||
if args.use_ema and not args.async_ema or (args.async_ema and step == len(loader) - 1):
|
||||
if args.use_fp16:
|
||||
ema.update(model.module.module, step=step)
|
||||
else:
|
||||
ema.update(model.module, step=step)
|
||||
|
||||
# ===========================================================================
|
||||
# Log loss values:
|
||||
# ===========================================================================
|
||||
running_loss += loss.item()
|
||||
log_steps += 1
|
||||
train_steps += 1
|
||||
if train_steps % args.log_every == 0:
|
||||
# Measure training speed:
|
||||
torch.cuda.synchronize()
|
||||
end_time = time.time()
|
||||
steps_per_sec = log_steps / (end_time - start_time)
|
||||
# Reduce loss history over all processes:
|
||||
avg_loss = torch.tensor(
|
||||
running_loss / log_steps, device=device)
|
||||
dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
|
||||
avg_loss = avg_loss.item() / world_size
|
||||
# get lr from deepspeed fused optimizer
|
||||
logger.info(f"(step={train_steps:07d}) " +
|
||||
(f"(update_step={train_steps // args.grad_accu_steps:07d}) " if args.grad_accu_steps > 1 else "") +
|
||||
f"Train Loss: {avg_loss:.4f}, "
|
||||
f"Lr: {opt.param_groups[0]['lr']:.6g}, "
|
||||
f"Steps/Sec: {steps_per_sec:.2f}, "
|
||||
f"Samples/Sec: {int(steps_per_sec * batch_size * world_size):d}")
|
||||
# Reset monitoring variables:
|
||||
running_loss = 0
|
||||
log_steps = 0
|
||||
start_time = time.time()
|
||||
|
||||
# collect gc:
|
||||
if args.gc_interval > 0 and (step % args.gc_interval == 0):
|
||||
gc.collect()
|
||||
|
||||
pbar.step(
|
||||
f"Epoch {epoch}, step {step}, loss {loss.item():.4f}", args.epochs * len(loader), train_steps)
|
||||
|
||||
if (train_steps % args.ckpt_every == 0 or train_steps % args.ckpt_latest_every == 0 # or train_steps == args.max_training_steps
|
||||
) and train_steps > 0:
|
||||
easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5,
|
||||
target_height=768, target_width=1280, train_steps=train_steps)
|
||||
save_checkpoint(args, rank, logger, model, ema,
|
||||
epoch, train_steps, checkpoint_dir)
|
||||
|
||||
if train_steps >= args.max_training_steps:
|
||||
logger.info(f"Breaking step loop at {train_steps}.")
|
||||
break
|
||||
|
||||
if train_steps >= args.max_training_steps:
|
||||
logger.info(f"Breaking epoch loop at {epoch}.")
|
||||
break
|
||||
|
||||
dist.destroy_process_group()
|
||||
@@ -1,175 +0,0 @@
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class SimpleNamespaceCNWarrper(SimpleNamespace):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.__dict__.update(kwargs)
|
||||
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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
conflict_handler='resolve',
|
||||
)
|
||||
parser.add_argument("--sys_path", type=str, default="")
|
||||
parser.add_argument("--train_config_file", type=str, default="")
|
||||
parser.add_argument("--mz_master_port", type=int, default=0)
|
||||
args = parser.parse_args()
|
||||
|
||||
master_port = args.mz_master_port
|
||||
|
||||
print(f"master_port = {master_port}")
|
||||
|
||||
try:
|
||||
from . import hook_HYDiT_utils
|
||||
except Exception as e:
|
||||
import hook_HYDiT_utils
|
||||
|
||||
hook_HYDiT_utils.set_master_port(master_port)
|
||||
|
||||
sys_path = args.sys_path
|
||||
if sys_path != "":
|
||||
sys.path.append(sys_path)
|
||||
|
||||
print("HYDi run hook")
|
||||
|
||||
try:
|
||||
from . import hook_HYDiT_main_train_deepspeed
|
||||
except Exception as e:
|
||||
import hook_HYDiT_main_train_deepspeed
|
||||
|
||||
import hydit.config
|
||||
|
||||
def _handle_conflict_error(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def parse_args(self, args=None, namespace=None):
|
||||
args, argv = self.parse_known_args(args, namespace)
|
||||
return args
|
||||
|
||||
train_config_file = args.train_config_file
|
||||
|
||||
if train_config_file == "":
|
||||
raise ValueError("train_config_file is empty")
|
||||
|
||||
train_config = {}
|
||||
with open(train_config_file, "r") as f:
|
||||
train_config = json.load(f)
|
||||
|
||||
argparse.ArgumentParser._handle_conflict_error = _handle_conflict_error
|
||||
argparse.ArgumentParser._handle_conflict_resolve = _handle_conflict_error
|
||||
argparse.ArgumentParser.parse_args = parse_args
|
||||
margs = hydit.config.get_args()
|
||||
margs.model = train_config.get("model", "DiT-g/2")
|
||||
|
||||
margs.task_flag = train_config.get("task_flag")
|
||||
|
||||
margs.resume_split = train_config.get("resume_split", True)
|
||||
|
||||
margs.ema_to_module = train_config.get("ema_to_module", True)
|
||||
|
||||
margs.deepspeed = False
|
||||
|
||||
margs.predict_type = train_config.get("predict_type", "v_prediction")
|
||||
|
||||
margs.training_parts = train_config.get("training_parts", "lora")
|
||||
|
||||
margs.batch_size = train_config.get("batch_size", 1)
|
||||
|
||||
margs.grad_accu_steps = train_config.get("grad_accu_steps", 1)
|
||||
|
||||
margs.global_seed = train_config.get("global_seed", 0)
|
||||
|
||||
margs.use_flash_attn = train_config.get("use_flash_attn", False)
|
||||
|
||||
margs.use_fp16 = train_config.get("use_fp16", True)
|
||||
|
||||
margs.qk_norm = train_config.get("qk_norm", True)
|
||||
|
||||
margs.ema_dtype = train_config.get("ema_dtype", "fp32")
|
||||
|
||||
margs.async_ema = False
|
||||
|
||||
margs.ckpt_latest_every = 0x7fffffff
|
||||
|
||||
margs.multireso = train_config.get("multireso", True)
|
||||
|
||||
margs.epochs = train_config.get("epochs", 50)
|
||||
|
||||
margs.target_ratios = train_config.get(
|
||||
"target_ratios", ['1:1', '3:4', '4:3', '16:9', '9:16'])
|
||||
|
||||
margs.rope_img = train_config.get("rope_img", "base1024")
|
||||
|
||||
margs.image_size = train_config.get("image_size", 1024)
|
||||
|
||||
margs.rope_real = train_config.get("rope_real", True)
|
||||
|
||||
margs.index_file = train_config.get("index_file", None)
|
||||
|
||||
margs.lr = train_config.get("lr", 1e-5)
|
||||
|
||||
margs.rank = train_config.get("rank", 8)
|
||||
|
||||
margs.noise_offset = train_config.get("noise_offset", 0.0)
|
||||
|
||||
margs.log_every = train_config.get("log_every", 99999999999999)
|
||||
|
||||
margs.use_zero_stage = train_config.get("use_zero_stage", 2)
|
||||
|
||||
margs.global_batch_size = train_config.get("global_batch_size", 1)
|
||||
|
||||
margs.deepspeed = True
|
||||
|
||||
margs.results_dir = train_config.get("results_dir")
|
||||
|
||||
margs.mse_loss_weight_type = train_config.get(
|
||||
"mse_loss_weight_type", "constant")
|
||||
|
||||
for k, v in train_config.items():
|
||||
if hasattr(margs, k):
|
||||
setattr(margs, k, v)
|
||||
|
||||
hook_HYDiT_utils.set_unet_path(
|
||||
train_config.get("unet_path"))
|
||||
|
||||
hook_HYDiT_utils.set_vae_ema_path(
|
||||
train_config.get("vae_ema_path"))
|
||||
|
||||
hook_HYDiT_utils.set_text_encoder_path(
|
||||
train_config.get("text_encoder_path"))
|
||||
|
||||
hook_HYDiT_utils.set_tokenizer_path(
|
||||
train_config.get("tokenizer_path"))
|
||||
|
||||
hook_HYDiT_utils.set_t5_encoder_path(
|
||||
train_config.get("t5_encoder_path"))
|
||||
|
||||
hook_HYDiT_utils.set_train_config(train_config)
|
||||
|
||||
try:
|
||||
# deepspeed/runtime/engine
|
||||
import deepspeed.runtime.engine
|
||||
deepspeed.runtime.engine.DeepSpeedEngine._do_sanity_check = lambda x: None
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
if type(margs.image_size) == int:
|
||||
margs.image_size = [margs.image_size, margs.image_size]
|
||||
hook_HYDiT_main_train_deepspeed.Core(margs)
|
||||
@@ -1,379 +0,0 @@
|
||||
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import torch
|
||||
|
||||
|
||||
UNET_PATH = "ckpts/t2i/model/pytorch_model_ema.pt"
|
||||
|
||||
|
||||
def set_unet_path(path):
|
||||
global UNET_PATH
|
||||
UNET_PATH = path
|
||||
|
||||
|
||||
VAE_EMA_PATH = "ckpts/t2i/sdxl-vae-fp16-fix"
|
||||
|
||||
|
||||
def set_vae_ema_path(path):
|
||||
global VAE_EMA_PATH
|
||||
VAE_EMA_PATH = path
|
||||
|
||||
|
||||
TOKENIZER = "ckpts/t2i/tokenizer"
|
||||
|
||||
|
||||
def set_tokenizer_path(path):
|
||||
global TOKENIZER
|
||||
TOKENIZER = path
|
||||
|
||||
|
||||
TEXT_ENCODER = 'ckpts/t2i/clip_text_encoder'
|
||||
|
||||
|
||||
def set_text_encoder_path(path):
|
||||
global TEXT_ENCODER
|
||||
TEXT_ENCODER = path
|
||||
|
||||
|
||||
T5_ENCODER = {
|
||||
'MT5': None,
|
||||
'attention_mask': True,
|
||||
'layer_index': -1,
|
||||
'attention_pool': True,
|
||||
'torch_dtype': torch.float16,
|
||||
'learnable_replace': True
|
||||
}
|
||||
|
||||
|
||||
def set_t5_encoder_path(path):
|
||||
global T5_ENCODER
|
||||
T5_ENCODER['MT5'] = path
|
||||
|
||||
|
||||
global TRAIN_CONFIG
|
||||
|
||||
|
||||
def set_train_config(train_config):
|
||||
global TRAIN_CONFIG
|
||||
TRAIN_CONFIG = train_config
|
||||
|
||||
|
||||
def easy_sample_images(
|
||||
args,
|
||||
vae=None,
|
||||
text_encoder=None,
|
||||
tokenizer=None,
|
||||
model=None,
|
||||
embedder_t5=None,
|
||||
target_height=768,
|
||||
target_width=1280,
|
||||
prompt="A photo of a girl with a hat on a sunny day",
|
||||
negative_prompt="",
|
||||
batch_size=1,
|
||||
guidance_scale=2.0,
|
||||
infer_steps=20,
|
||||
sampler='dpmpp_2m_karras',
|
||||
train_steps=0,
|
||||
seed=0,
|
||||
):
|
||||
from hydit.diffusion.pipeline import StableDiffusionPipeline
|
||||
from diffusers import schedulers
|
||||
from hydit.constants import SAMPLER_FACTORY
|
||||
from hydit.modules.posemb_layers import get_fill_resize_and_crop, get_2d_rotary_pos_embed
|
||||
from hydit.modules.models import HUNYUAN_DIT_CONFIG
|
||||
|
||||
import traceback
|
||||
with torch.cuda.amp.autocast():
|
||||
|
||||
workspace_dir = TRAIN_CONFIG.get("workspace_dir")
|
||||
sample_config_file = TRAIN_CONFIG.get("sample_config_file", None)
|
||||
if sample_config_file is None:
|
||||
print("sample_config_file is not set.")
|
||||
return
|
||||
try:
|
||||
sample_config = json.load(open(sample_config_file, "r"))
|
||||
except Exception as e:
|
||||
print(f"Failed to load sample_config_file: {sample_config_file}")
|
||||
return
|
||||
sample_images_dir = os.path.join(workspace_dir, "sample_images")
|
||||
os.makedirs(sample_images_dir, exist_ok=True)
|
||||
|
||||
sampler_factory = SAMPLER_FACTORY.copy()
|
||||
|
||||
sampler_factory["uni_pc"] = {
|
||||
'scheduler': 'UniPCMultistepScheduler',
|
||||
'name': 'UniPCMultistepScheduler',
|
||||
'kwargs': {
|
||||
'beta_schedule': 'scaled_linear',
|
||||
'beta_start': 0.00085,
|
||||
'beta_end': 0.03,
|
||||
'prediction_type': 'v_prediction',
|
||||
'trained_betas': None,
|
||||
'solver_order': 2,
|
||||
}
|
||||
}
|
||||
sampler_factory["dpmpp_2m_karras"] = {
|
||||
'scheduler': 'DPMSolverMultistepScheduler',
|
||||
'name': 'DPMSolverMultistepScheduler',
|
||||
'kwargs': {
|
||||
'beta_schedule': 'scaled_linear',
|
||||
'beta_start': 0.00085,
|
||||
'beta_end': 0.03,
|
||||
'prediction_type': 'v_prediction',
|
||||
'trained_betas': None,
|
||||
'solver_order': 2,
|
||||
'algorithm_type': 'dpmsolver++',
|
||||
"use_karras_sigmas": True,
|
||||
}
|
||||
}
|
||||
|
||||
# Load sampler from factory
|
||||
kwargs = sampler_factory[sampler]['kwargs']
|
||||
scheduler = sampler_factory[sampler]['scheduler']
|
||||
|
||||
# Build scheduler according to the sampler.
|
||||
scheduler_class = getattr(schedulers, scheduler)
|
||||
scheduler = scheduler_class(**kwargs)
|
||||
|
||||
# Set timesteps for inference steps.
|
||||
scheduler.set_timesteps(infer_steps, "cuda")
|
||||
|
||||
def calc_rope(height, width):
|
||||
model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"]
|
||||
patch_size = model_config['patch_size']
|
||||
head_size = model_config['hidden_size'] // model_config['num_heads']
|
||||
th = height // 8 // patch_size
|
||||
tw = width // 8 // patch_size
|
||||
base_size = 512 // 8 // patch_size
|
||||
start, stop = get_fill_resize_and_crop((th, tw), base_size)
|
||||
sub_args = [start, stop, (th, tw)]
|
||||
rope = get_2d_rotary_pos_embed(head_size, *sub_args)
|
||||
return rope
|
||||
|
||||
pipeline = StableDiffusionPipeline(vae=vae,
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=tokenizer,
|
||||
unet=model.module,
|
||||
scheduler=scheduler,
|
||||
feature_extractor=None,
|
||||
safety_checker=None,
|
||||
requires_safety_checker=False,
|
||||
embedder_t5=embedder_t5,
|
||||
)
|
||||
pipeline = pipeline.to("cuda")
|
||||
# attr _execution_device is not defined
|
||||
|
||||
style = torch.as_tensor([0, 0] * batch_size, device="cuda")
|
||||
|
||||
src_size_cond = (target_width, target_height)
|
||||
size_cond = list(src_size_cond) + [target_width, target_height, 0, 0]
|
||||
image_meta_size = torch.as_tensor(
|
||||
[size_cond] * 2 * batch_size, device="cuda",)
|
||||
|
||||
if type(sample_config) != list:
|
||||
sample_config = [sample_config]
|
||||
|
||||
for i, sample in enumerate(sample_config):
|
||||
prompt = sample.get("prompt", "")
|
||||
negative_prompt = sample.get("negative_prompt", "")
|
||||
guidance_scale = sample.get("cfg", guidance_scale)
|
||||
infer_steps = sample.get("steps", infer_steps)
|
||||
width = sample.get("width", target_width)
|
||||
height = sample.get("height", target_height)
|
||||
|
||||
freqs_cis_img = calc_rope(height, width)
|
||||
|
||||
try:
|
||||
generator = torch.Generator(device="cuda")
|
||||
generator.manual_seed(seed)
|
||||
samples = pipeline(
|
||||
height=height,
|
||||
width=width,
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
num_images_per_prompt=batch_size,
|
||||
guidance_scale=guidance_scale,
|
||||
num_inference_steps=infer_steps,
|
||||
style=style,
|
||||
return_dict=False,
|
||||
use_fp16=True,
|
||||
learn_sigma=args.learn_sigma,
|
||||
freqs_cis_img=freqs_cis_img,
|
||||
image_meta_size=image_meta_size,
|
||||
generator=generator,
|
||||
)[0]
|
||||
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"Failed to sample images: {e} ")
|
||||
# 打印堆栈信息
|
||||
traceback.print_exc()
|
||||
print(f"Failed to sample pipeline: {pipeline} ")
|
||||
|
||||
# print("samples:",type(samples),)
|
||||
# input("Press Enter to continue...")
|
||||
# print("samples:",samples,)
|
||||
|
||||
if type(samples) == list:
|
||||
pil_image = samples[0]
|
||||
else:
|
||||
pil_image = samples
|
||||
|
||||
sample_filename = f"{args.task_flag}_train_steps_{train_steps:07d}.png"
|
||||
sample_filename_path = os.path.join(
|
||||
sample_images_dir, sample_filename)
|
||||
pil_image.save(sample_filename_path)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def model_resume(args, model, ema, logger):
|
||||
"""
|
||||
Load pretrained weights.
|
||||
"""
|
||||
start_epoch = 0
|
||||
start_epoch_step = 0
|
||||
train_steps = 0
|
||||
resume_path = UNET_PATH
|
||||
|
||||
logger.info(f"Resume from checkpoint {resume_path}")
|
||||
|
||||
if args.resume_split:
|
||||
# Resume main model
|
||||
|
||||
resume_ckpt_module = torch.load(
|
||||
resume_path, map_location=lambda storage, loc: storage)
|
||||
model.load_state_dict(resume_ckpt_module, strict=False)
|
||||
|
||||
# Resume ema model
|
||||
if args.use_ema:
|
||||
if args.module_to_ema:
|
||||
if "resume_ckpt_module" in locals():
|
||||
logger.info(f" Resume ema model from main states.")
|
||||
ema.load_state_dict(resume_ckpt_module, strict=args.strict)
|
||||
else:
|
||||
logger.info(f" Resume ema model from module states.")
|
||||
resume_ckpt_module = torch.load(
|
||||
resume_path, map_location=lambda storage, loc: storage)
|
||||
ema.load_state_dict(resume_ckpt_module, strict=args.strict)
|
||||
else:
|
||||
if "resume_ckpt_ema" in locals():
|
||||
logger.info(f" Resume ema model from EMA states.")
|
||||
ema.load_state_dict(resume_ckpt_ema, strict=args.strict)
|
||||
else:
|
||||
logger.info(f" Resume ema model from EMA states.")
|
||||
resume_ckpt_ema = torch.load(resume_path,
|
||||
map_location=lambda storage, loc: storage)
|
||||
ema.load_state_dict(resume_ckpt_ema, strict=args.strict)
|
||||
else:
|
||||
raise ValueError(
|
||||
" “If `resume` is True, then either `resume_split` must be true.”")
|
||||
|
||||
return model, ema, start_epoch, start_epoch_step, train_steps
|
||||
|
||||
|
||||
import tqdm
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def set_master_port(port):
|
||||
global master_port
|
||||
master_port = port
|
||||
|
||||
|
||||
master_port = 0
|
||||
|
||||
|
||||
def LOG(log):
|
||||
if master_port == 0:
|
||||
raise Exception("master_port is 0")
|
||||
# 发送http
|
||||
try:
|
||||
resp = requests.request("post", f"http://127.0.0.1:{master_port}/log", data=json.dumps(log), headers={
|
||||
"Content-Type": "application/json"})
|
||||
except Exception as e:
|
||||
return
|
||||
if resp.status_code != 200:
|
||||
raise Exception(f"LOG failed: {resp.text}")
|
||||
|
||||
|
||||
# with tqdm(total=total_steps, initial=train_steps) as pbar:
|
||||
# pbar.update(1)
|
||||
# pbar.set_description(
|
||||
# f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {mean_loss / step:.4f}")
|
||||
class PBar:
|
||||
def __init__(self, total):
|
||||
self.pbar = tqdm.tqdm(total=total)
|
||||
|
||||
def step(self, desc, total_steps, train_steps):
|
||||
self.pbar.update(1)
|
||||
self.pbar.set_description(desc)
|
||||
|
||||
LOG({
|
||||
"type": "sample_images",
|
||||
"global_step": train_steps,
|
||||
"total_steps": total_steps,
|
||||
# "latent": noise_pred_latent_path,
|
||||
})
|
||||
|
||||
|
||||
from torch import nn
|
||||
|
||||
|
||||
class CustomizeEmbedsModel(nn.Module):
|
||||
dtype = torch.float16
|
||||
# x = torch.zeros(1, 1, 256, 2048)
|
||||
x = None
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__()
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
self.dtype = torch.float16
|
||||
return self
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
input_ids = kwargs.get("input_ids", None)
|
||||
if self.x is None:
|
||||
if input_ids is None:
|
||||
batch_size = 1
|
||||
else:
|
||||
batch_size = input_ids.shape[0]
|
||||
self.x = torch.zeros(1, batch_size, 256, 2048, dtype=self.dtype)
|
||||
|
||||
if kwargs.get("output_hidden_states", False):
|
||||
return {
|
||||
"hidden_states": self.x.to("cuda"),
|
||||
"input_ids": torch.zeros(1, 1),
|
||||
}
|
||||
return self.x
|
||||
|
||||
|
||||
class CustomizeTokenizer(dict):
|
||||
|
||||
added_tokens_encoder = []
|
||||
input_ids = torch.zeros(1, 256)
|
||||
attention_mask = torch.zeros(1, 256)
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
self['added_tokens_encoder'] = self.added_tokens_encoder
|
||||
self['input_ids'] = self.input_ids
|
||||
self['attention_mask'] = self.attention_mask
|
||||
|
||||
def tokenize(self, text):
|
||||
return text
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
|
||||
class CustomizeEmbeds():
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.tokenizer = CustomizeTokenizer()
|
||||
self.model = CustomizeEmbedsModel().to("cuda")
|
||||
self.max_length = 256
|
||||
@@ -1,326 +0,0 @@
|
||||
|
||||
import inspect
|
||||
import re
|
||||
from typing import Callable, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
from packaging import version
|
||||
from tqdm import tqdm
|
||||
from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
|
||||
|
||||
from diffusers import SchedulerMixin, StableDiffusionPipeline
|
||||
from diffusers.models import AutoencoderKL, UNet2DConditionModel
|
||||
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker
|
||||
from diffusers.utils import logging
|
||||
from PIL import Image
|
||||
from library import sdxl_model_util, sdxl_train_util, train_util
|
||||
|
||||
from library.sdxl_lpw_stable_diffusion import *
|
||||
|
||||
|
||||
from library.hunyuan_models import *
|
||||
|
||||
|
||||
def load_scheduler_sigmas(beta_start=0.00085, beta_end=0.018, num_train_timesteps=1000):
|
||||
betas = torch.linspace(beta_start**0.5, beta_end**0.5,
|
||||
num_train_timesteps, dtype=torch.float32) ** 2
|
||||
alphas = 1.0 - betas
|
||||
alphas_cumprod = torch.cumprod(alphas, dim=0)
|
||||
|
||||
sigmas = np.array(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5)
|
||||
sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32)
|
||||
sigmas = torch.from_numpy(sigmas)
|
||||
return alphas_cumprod, sigmas
|
||||
|
||||
|
||||
ATTN_MODE = "xformers"
|
||||
CLIP_TOKENS = 75 * 2 + 2
|
||||
DEVICE = "cuda"
|
||||
try:
|
||||
from k_diffusion.external import DiscreteVDDPMDenoiser
|
||||
from k_diffusion.sampling import sample_euler_ancestral, get_sigmas_exponential, sample_dpmpp_2m_sde
|
||||
except ImportError:
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
subprocess.check_call(
|
||||
[sys.executable, "-m", "pip", "install", "k-diffusion"])
|
||||
|
||||
from k_diffusion.external import DiscreteVDDPMDenoiser
|
||||
from k_diffusion.sampling import sample_euler_ancestral, get_sigmas_exponential, sample_dpmpp_2m_sde
|
||||
|
||||
|
||||
from library.hunyuan_utils import get_cond, calc_rope
|
||||
|
||||
|
||||
class HuanYuanDiffusionLongPromptWeightingPipeline:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae: AutoencoderKL,
|
||||
text_encoder,
|
||||
tokenizer,
|
||||
unet,
|
||||
scheduler,
|
||||
# clip_skip: int,
|
||||
safety_checker,
|
||||
feature_extractor,
|
||||
requires_safety_checker=False,
|
||||
clip_skip=0,
|
||||
):
|
||||
# clip skip is ignored currently
|
||||
# print("tokenizer: ", tokenizer)
|
||||
# print("text_encoder: ", text_encoder)
|
||||
self.unet = unet
|
||||
self.scheduler = scheduler
|
||||
self.safety_checker = safety_checker
|
||||
self.feature_extractor = feature_extractor
|
||||
self.requires_safety_checker = requires_safety_checker
|
||||
self.vae = vae
|
||||
self.vae_scale_factor = 2 ** (
|
||||
len(self.vae.config.block_out_channels) - 1)
|
||||
self.progress_bar = lambda x: tqdm(x, leave=False)
|
||||
|
||||
self.clip_skip = clip_skip
|
||||
self.tokenizers = tokenizer
|
||||
self.text_encoders = text_encoder
|
||||
|
||||
def to(self, device=None, dtype=None):
|
||||
if device is not None:
|
||||
self.device = device
|
||||
# self.vae.to(device=self.device)
|
||||
if dtype is not None:
|
||||
self.dtype = dtype
|
||||
|
||||
@property
|
||||
def _execution_device(self):
|
||||
r"""
|
||||
Returns the device on which the pipeline's models will be executed. After calling
|
||||
`pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
|
||||
hooks.
|
||||
"""
|
||||
if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
|
||||
return self.device
|
||||
for module in self.unet.modules():
|
||||
if (
|
||||
hasattr(module, "_hf_hook")
|
||||
and hasattr(module._hf_hook, "execution_device")
|
||||
and module._hf_hook.execution_device is not None
|
||||
):
|
||||
return torch.device(module._hf_hook.execution_device)
|
||||
return self.device
|
||||
|
||||
def check_inputs(self, prompt, height, width, strength, callback_steps):
|
||||
pass
|
||||
|
||||
def get_timesteps(self, num_inference_steps, strength, device, is_text2img):
|
||||
if is_text2img:
|
||||
return self.scheduler.timesteps.to(device), num_inference_steps
|
||||
else:
|
||||
# get the original timestep using init_timestep
|
||||
offset = self.scheduler.config.get("steps_offset", 0)
|
||||
init_timestep = int(num_inference_steps * strength) + offset
|
||||
init_timestep = min(init_timestep, num_inference_steps)
|
||||
|
||||
t_start = max(num_inference_steps - init_timestep + offset, 0)
|
||||
timesteps = self.scheduler.timesteps[t_start:].to(device)
|
||||
return timesteps, num_inference_steps - t_start
|
||||
|
||||
def run_safety_checker(self, image, device, dtype):
|
||||
return image, None
|
||||
|
||||
def decode_latents(self, latents):
|
||||
with torch.no_grad():
|
||||
latents = 1 / sdxl_model_util.VAE_SCALE_FACTOR * latents
|
||||
|
||||
image = self.vae.decode(latents.to(self.vae.dtype)).sample
|
||||
image = (image / 2 + 0.5).clamp(0, 1)
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
|
||||
image = image.cpu().permute(0, 2, 3, 1).float().numpy()
|
||||
return image
|
||||
|
||||
def prepare_extra_step_kwargs(self, generator, eta):
|
||||
return {}
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
prompt: Union[str, List[str]],
|
||||
negative_prompt: Optional[Union[str, List[str]]] = None,
|
||||
image: Union[torch.FloatTensor, PIL.Image.Image] = None,
|
||||
mask_image: Union[torch.FloatTensor, PIL.Image.Image] = None,
|
||||
height: int = 512,
|
||||
width: int = 512,
|
||||
num_inference_steps: int = 50,
|
||||
guidance_scale: float = 7.5,
|
||||
strength: float = 0.8,
|
||||
num_images_per_prompt: Optional[int] = 1,
|
||||
eta: float = 0.0,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
max_embeddings_multiples: Optional[int] = 3,
|
||||
output_type: Optional[str] = "pil",
|
||||
return_dict: bool = True,
|
||||
controlnet=None,
|
||||
controlnet_image=None,
|
||||
callback: Optional[Callable[[
|
||||
int, int, torch.FloatTensor], None]] = None,
|
||||
is_cancelled_callback: Optional[Callable[[], bool]] = None,
|
||||
callback_steps: int = 1,
|
||||
):
|
||||
|
||||
BETA_END = 0.018
|
||||
CFG_SCALE = guidance_scale
|
||||
STEPS = num_inference_steps
|
||||
patch_size = 2
|
||||
num_heads = 88
|
||||
|
||||
alphas, sigmas = load_scheduler_sigmas(beta_end=BETA_END)
|
||||
denoiser = self.unet
|
||||
clip_tokenizer = self.tokenizers[0]
|
||||
clip_encoder = self.text_encoders[0]
|
||||
|
||||
mt5_embedder = self.text_encoders[1]
|
||||
|
||||
vae = self.vae
|
||||
denoiser.eval()
|
||||
denoiser.disable_fp32_silu()
|
||||
denoiser.disable_fp32_layer_norm()
|
||||
denoiser.set_attn_mode(ATTN_MODE)
|
||||
vae.requires_grad_(False)
|
||||
mt5_embedder.to(torch.float16)
|
||||
|
||||
with torch.autocast("cuda"):
|
||||
clip_h, clip_m, mt5_h, mt5_m = get_cond(
|
||||
prompt,
|
||||
mt5_embedder,
|
||||
clip_tokenizer,
|
||||
clip_encoder,
|
||||
# Should be same as original implementation with max_length_clip=77
|
||||
# Support 75*n + 2
|
||||
max_length_clip=CLIP_TOKENS,
|
||||
)
|
||||
neg_clip_h, neg_clip_m, neg_mt5_h, neg_mt5_m = get_cond(
|
||||
negative_prompt,
|
||||
mt5_embedder,
|
||||
clip_tokenizer,
|
||||
clip_encoder,
|
||||
max_length_clip=CLIP_TOKENS,
|
||||
)
|
||||
clip_h = torch.concat([clip_h, neg_clip_h], dim=0)
|
||||
clip_m = torch.concat([clip_m, neg_clip_m], dim=0)
|
||||
mt5_h = torch.concat([mt5_h, neg_mt5_h], dim=0)
|
||||
mt5_m = torch.concat([mt5_m, neg_mt5_m], dim=0)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
style = torch.as_tensor([0] * 2, device=DEVICE)
|
||||
# src hw, dst hw, 0, 0
|
||||
size_cond = [height, width, height, width, 0, 0]
|
||||
image_meta_size = torch.as_tensor([size_cond] * 2, device=DEVICE)
|
||||
|
||||
freqs_cis_img = calc_rope(height, width, patch_size, num_heads)
|
||||
|
||||
denoiser_wrapper = DiscreteVDDPMDenoiser(
|
||||
# A quick patch for learn_sigma
|
||||
lambda *args, **kwargs: denoiser(* \
|
||||
args, **kwargs).chunk(2, dim=1)[0],
|
||||
alphas,
|
||||
False,
|
||||
).to(DEVICE)
|
||||
|
||||
def cfg_denoise_func(x, sigma):
|
||||
cond, uncond = denoiser_wrapper(
|
||||
x.repeat(2, 1, 1, 1),
|
||||
sigma.repeat(2),
|
||||
encoder_hidden_states=clip_h,
|
||||
text_embedding_mask=clip_m,
|
||||
encoder_hidden_states_t5=mt5_h,
|
||||
text_embedding_mask_t5=mt5_m,
|
||||
image_meta_size=image_meta_size,
|
||||
style=style,
|
||||
cos_cis_img=freqs_cis_img[0],
|
||||
sin_cis_img=freqs_cis_img[1],
|
||||
).chunk(2, dim=0)
|
||||
return uncond + (cond - uncond) * CFG_SCALE
|
||||
|
||||
sigmas = denoiser_wrapper.get_sigmas(STEPS).to(DEVICE)
|
||||
sigmas = get_sigmas_exponential(
|
||||
STEPS, denoiser_wrapper.sigma_min, denoiser_wrapper.sigma_max, DEVICE
|
||||
)
|
||||
x1 = torch.randn(1, 4, height // 8, width // 8,
|
||||
dtype=torch.float16, device=DEVICE)
|
||||
|
||||
with torch.autocast("cuda"):
|
||||
sample = sample_dpmpp_2m_sde(
|
||||
cfg_denoise_func,
|
||||
x1 * sigmas[0],
|
||||
sigmas,
|
||||
)
|
||||
torch.cuda.empty_cache()
|
||||
latents = sample
|
||||
return latents
|
||||
|
||||
def text2img(
|
||||
self,
|
||||
prompt: Union[str, List[str]],
|
||||
negative_prompt: Optional[Union[str, List[str]]] = None,
|
||||
height: int = 512,
|
||||
width: int = 512,
|
||||
num_inference_steps: int = 50,
|
||||
guidance_scale: float = 7.5,
|
||||
num_images_per_prompt: Optional[int] = 1,
|
||||
eta: float = 0.0,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
max_embeddings_multiples: Optional[int] = 3,
|
||||
output_type: Optional[str] = "pil",
|
||||
return_dict: bool = True,
|
||||
callback: Optional[Callable[[
|
||||
int, int, torch.FloatTensor], None]] = None,
|
||||
is_cancelled_callback: Optional[Callable[[], bool]] = None,
|
||||
callback_steps: int = 1,
|
||||
):
|
||||
|
||||
return self.__call__(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
height=height,
|
||||
width=width,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
eta=eta,
|
||||
generator=generator,
|
||||
latents=latents,
|
||||
max_embeddings_multiples=max_embeddings_multiples,
|
||||
output_type=output_type,
|
||||
return_dict=return_dict,
|
||||
callback=callback,
|
||||
is_cancelled_callback=is_cancelled_callback,
|
||||
callback_steps=callback_steps,
|
||||
)
|
||||
|
||||
def latents_to_image(self, latents):
|
||||
# 9. Post-processing
|
||||
image = self.decode_latents(latents.to(self.vae.dtype))
|
||||
image = self.numpy_to_pil(image)
|
||||
return image
|
||||
|
||||
# copy from pil_utils.py
|
||||
def numpy_to_pil(self, images: np.ndarray) -> Image.Image:
|
||||
"""
|
||||
Convert a numpy image or a batch of images to a PIL image.
|
||||
"""
|
||||
if images.ndim == 3:
|
||||
images = images[None, ...]
|
||||
images = (images * 255).round().astype("uint8")
|
||||
if images.shape[-1] == 1:
|
||||
# special case for grayscale (single channel) images
|
||||
pil_images = [Image.fromarray(
|
||||
image.squeeze(), mode="L") for image in images]
|
||||
else:
|
||||
pil_images = [Image.fromarray(image) for image in images]
|
||||
|
||||
return pil_images
|
||||
+7
-183
@@ -27,15 +27,12 @@ def config2args(train_parser: argparse.ArgumentParser, config):
|
||||
return args
|
||||
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
import numpy as np
|
||||
import tempfile
|
||||
import safetensors.torch
|
||||
|
||||
import hook_kohya_ss_utils
|
||||
|
||||
import sys
|
||||
sys.path.append(os.path.dirname(__file__))
|
||||
try:
|
||||
import hook_kohya_ss_utils
|
||||
except:
|
||||
from . import hook_kohya_ss_utils
|
||||
|
||||
other_config = {}
|
||||
original_save_model = None
|
||||
@@ -200,182 +197,9 @@ class SimpleNamespaceCNWarrper(SimpleNamespace):
|
||||
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
|
||||
library.train_util.sample_images = utils_sample_images
|
||||
|
||||
global sample_images_pipe_class
|
||||
sample_images_pipe_class = library.train_util.StableDiffusionLongPromptWeightingPipeline
|
||||
|
||||
train_args = config2args(train_controlnet.setup_parser(), train_config)
|
||||
|
||||
LOG({
|
||||
"type": "start_train",
|
||||
})
|
||||
|
||||
train_controlnet.train(train_args)
|
||||
|
||||
|
||||
def run_lora_hunyuan1_2():
|
||||
hook_kohya_ss_utils.hook_kohya_ss()
|
||||
|
||||
# 覆盖sample_images生成函数,包括进度条和生成图片功能
|
||||
import hunyuan_train_network
|
||||
|
||||
# 还未实现
|
||||
hunyuan_train_network.HunYuanNetworkTrainer.sample_images = sample_images
|
||||
|
||||
# def empty_sample_images(*args, **kwargs):
|
||||
# pass
|
||||
# sample_images_pipe_class = empty_sample_images
|
||||
|
||||
import hook_kohya_ss_hunyuan_pipe
|
||||
global sample_images_pipe_class
|
||||
sample_images_pipe_class = hook_kohya_ss_hunyuan_pipe.HuanYuanDiffusionLongPromptWeightingPipeline
|
||||
# 配置对应的pipeline
|
||||
import library.hunyuan_utils
|
||||
|
||||
print(json.dumps(other_config, indent=4))
|
||||
hunyuan_models_config = other_config.get("hunyuan_models_config", None)
|
||||
|
||||
from transformers import (
|
||||
AutoTokenizer,
|
||||
T5Tokenizer,
|
||||
BertModel,
|
||||
BertTokenizer,
|
||||
)
|
||||
from diffusers import AutoencoderKL, LMSDiscreteScheduler
|
||||
|
||||
def hunyuan_load_tokenizers():
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
hunyuan_models_config["tokenizer_path"],
|
||||
local_files_only=True,
|
||||
)
|
||||
tokenizer.eos_token_id = tokenizer.sep_token_id
|
||||
t5_encoder_path = hunyuan_models_config.get("t5_encoder_path", None)
|
||||
if t5_encoder_path == "none":
|
||||
t5_encoder_path = None
|
||||
tokenizer2 = None
|
||||
if t5_encoder_path is not None:
|
||||
tokenizer2 = T5Tokenizer.from_pretrained(
|
||||
t5_encoder_path,
|
||||
local_files_only=True,
|
||||
)
|
||||
|
||||
return [tokenizer, tokenizer2]
|
||||
|
||||
library.hunyuan_utils.load_tokenizers = hunyuan_load_tokenizers
|
||||
|
||||
def hunyuan_load_model(model_path: str, dtype=torch.float16, device="cuda", use_extra_cond=False, dit_path=None):
|
||||
|
||||
dit_path = hunyuan_models_config.get("unet_path", None)
|
||||
|
||||
import library.hunyuan_models
|
||||
# from hunyuan_models import MT5Embedder, HunYuanDiT, BertModel, DiT_g_2
|
||||
MT5Embedder = library.hunyuan_models.MT5Embedder
|
||||
HunYuanDiT = library.hunyuan_models.HunYuanDiT
|
||||
BertModel = library.hunyuan_models.BertModel
|
||||
DiT_g_2 = library.hunyuan_models.DiT_g_2
|
||||
|
||||
denoiser, patch_size, head_dim = DiT_g_2(
|
||||
input_size=(128, 128), use_extra_cond=use_extra_cond)
|
||||
if dit_path is not None:
|
||||
state_dict = torch.load(dit_path)
|
||||
if 'state_dict' in state_dict:
|
||||
state_dict = state_dict['state_dict']
|
||||
else:
|
||||
state_dict = torch.load(os.path.join(
|
||||
model_path, "denoiser/pytorch_model_module.pt"))
|
||||
denoiser.load_state_dict(state_dict)
|
||||
denoiser.to(device).to(dtype)
|
||||
|
||||
clip_tokenizer = AutoTokenizer.from_pretrained(
|
||||
hunyuan_models_config["tokenizer_path"],
|
||||
local_files_only=True,
|
||||
)
|
||||
clip_tokenizer.eos_token_id = 2
|
||||
clip_encoder = (
|
||||
BertModel.from_pretrained(
|
||||
hunyuan_models_config["text_encoder_path"],
|
||||
local_files_only=True,
|
||||
).to(device).to(dtype)
|
||||
)
|
||||
|
||||
t5_encoder_path = hunyuan_models_config.get("t5_encoder_path", None)
|
||||
if t5_encoder_path == "none":
|
||||
t5_encoder_path = None
|
||||
mt5_embedder = None
|
||||
if t5_encoder_path is not None:
|
||||
mt5_embedder = (
|
||||
MT5Embedder(
|
||||
model_dir=hunyuan_models_config["t5_encoder_path"],
|
||||
torch_dtype=dtype,
|
||||
max_length=256)
|
||||
.to(device)
|
||||
.to(dtype)
|
||||
)
|
||||
else:
|
||||
|
||||
batch_size = train_args.train_batch_size
|
||||
import library.config_util
|
||||
user_config = library.config_util.load_user_config(
|
||||
train_args.dataset_config)
|
||||
datasets = user_config.get("datasets", [])
|
||||
if len(datasets) > 0:
|
||||
batch_size = datasets[0].get("batch_size", batch_size)
|
||||
mt5_embedder = (
|
||||
hook_kohya_ss_utils.CustomizeMT5Embedder(
|
||||
batch_size=batch_size,
|
||||
)
|
||||
.to(device)
|
||||
.to(dtype)
|
||||
)
|
||||
|
||||
vae = (
|
||||
AutoencoderKL.from_pretrained(
|
||||
hunyuan_models_config["vae_ema_path"],
|
||||
local_files_only=True,
|
||||
)
|
||||
.to(device)
|
||||
.to(dtype)
|
||||
)
|
||||
vae.requires_grad_(False)
|
||||
return (
|
||||
denoiser,
|
||||
patch_size,
|
||||
head_dim,
|
||||
clip_tokenizer,
|
||||
clip_encoder,
|
||||
mt5_embedder,
|
||||
vae,
|
||||
)
|
||||
|
||||
library.hunyuan_utils.load_model = hunyuan_load_model
|
||||
|
||||
trainer = hunyuan_train_network.HunYuanNetworkTrainer()
|
||||
train_args = config2args(
|
||||
hunyuan_train_network.setup_parser(), train_config)
|
||||
print(f"train_args = {train_args}")
|
||||
|
||||
LOG({
|
||||
"type": "start_train",
|
||||
})
|
||||
trainer.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,
|
||||
"run_lora_hunyuan1_2": run_lora_hunyuan1_2,
|
||||
}
|
||||
|
||||
|
||||
@@ -445,4 +269,4 @@ if __name__ == "__main__":
|
||||
except Exception as e:
|
||||
print(f"Exception: {e}")
|
||||
if sys.platform == "win32":
|
||||
input("Press Enter to continue...")
|
||||
input("Press Enter to continue...")
|
||||
|
||||
+102
-387
@@ -16,65 +16,11 @@ import folder_paths
|
||||
import nodes
|
||||
|
||||
|
||||
git_accelerate_urls = {
|
||||
"githubfast": "githubfast.com",
|
||||
"521github": "521github.com",
|
||||
"kkgithub": "kkgithub.com",
|
||||
}
|
||||
|
||||
|
||||
def MZ_KohyaSSCloneRepo_call(args={}):
|
||||
mz_dir = Utils.get_minus_zone_models_path()
|
||||
|
||||
branch_repoid = args.get("branch_repoid", "kohya-ss/sd-scripts")
|
||||
branch_local_name = args.get("branch_local_name", "kohya_ss_lora")
|
||||
|
||||
git_url = f"https://github.com/{branch_repoid}"
|
||||
source = args.get("source", "github")
|
||||
kohya_ss_lora_dir = os.path.join(mz_dir, "train_tools", branch_local_name)
|
||||
if git_accelerate_urls.get(source, None) is not None:
|
||||
git_url = f"https://{git_accelerate_urls[source]}/{branch_repoid}"
|
||||
try:
|
||||
if not os.path.exists(kohya_ss_lora_dir) or not os.path.exists(os.path.join(kohya_ss_lora_dir, ".git")):
|
||||
subprocess.run(
|
||||
["git", "clone", "--depth", "1", git_url, kohya_ss_lora_dir], check=True)
|
||||
|
||||
# 切换远程分支 git remote set-branches origin 'main'
|
||||
branch = args.get("branch", "main")
|
||||
|
||||
# 查看本地分支是否一致
|
||||
short_result = subprocess.run(
|
||||
["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=kohya_ss_lora_dir, stdout=subprocess.PIPE, check=True)
|
||||
result = subprocess.run(
|
||||
["git", "rev-parse", "HEAD"], cwd=kohya_ss_lora_dir, stdout=subprocess.PIPE, check=True)
|
||||
|
||||
short_current_branch = short_result.stdout.decode().strip()
|
||||
long_current_branch = result.stdout.decode().strip()
|
||||
print(
|
||||
f"当前分支(current branch): {long_current_branch}({short_current_branch})")
|
||||
print(f"目标分支(target branch): {branch}")
|
||||
|
||||
if branch != long_current_branch and branch != short_current_branch:
|
||||
subprocess.run(
|
||||
["git", "remote", "set-branches", "origin", branch], cwd=kohya_ss_lora_dir, check=True)
|
||||
subprocess.run(
|
||||
["git", "fetch", "--depth", "1", "origin", branch], cwd=kohya_ss_lora_dir, check=True)
|
||||
|
||||
# 恢复所有文件
|
||||
subprocess.run(
|
||||
["git", "checkout", "."], cwd=kohya_ss_lora_dir, check=True)
|
||||
|
||||
subprocess.run(
|
||||
["git", "checkout", branch], cwd=kohya_ss_lora_dir, check=True)
|
||||
|
||||
except Exception as e:
|
||||
raise Exception(f"克隆kohya-ss/sd-scripts或者切换分支时出现异常,详细信息请查看控制台...")
|
||||
|
||||
# 初始化工具仓库和工作区
|
||||
|
||||
|
||||
def MZ_KohyaSSInitWorkspace_call(args={}):
|
||||
MZ_KohyaSSCloneRepo_call(args)
|
||||
Utils.clone_repo(args)
|
||||
|
||||
workspace_name = args.get("lora_name", None)
|
||||
workspace_name = workspace_name.strip()
|
||||
@@ -97,7 +43,7 @@ def MZ_KohyaSSInitWorkspace_call(args={}):
|
||||
)
|
||||
|
||||
|
||||
def MZ_ImageSelecter_call(args={}):
|
||||
def MZ_KohyaSSDatasetConfig_call(args={}):
|
||||
images = args.get("images")
|
||||
pil_images = Utils.tensors2pil_list(images)
|
||||
|
||||
@@ -175,7 +121,9 @@ def MZ_ImageSelecter_call(args={}):
|
||||
|
||||
dataset_config_extension = args.get("dataset_config_extension")
|
||||
generate_dataset_config(
|
||||
os.path.join(workspace_dir, "dataset" + dataset_config_extension),
|
||||
workspace_dir=workspace_dir,
|
||||
output_path=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"),
|
||||
@@ -183,6 +131,7 @@ def MZ_ImageSelecter_call(args={}):
|
||||
conditioning_data_dir=conditioning_images_dir,
|
||||
caption_extension=caption_extension,
|
||||
num_repeats=args.get("num_repeats"),
|
||||
is_reg=args.get("auto_reg") == "enable",
|
||||
)
|
||||
return (
|
||||
train_images_dir,
|
||||
@@ -343,7 +292,7 @@ def check_install():
|
||||
import logging
|
||||
|
||||
|
||||
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, ):
|
||||
def generate_dataset_config(workspace_dir, output_path, enable_bucket=True, resolution=512, batch_size=1, image_dir=None, conditioning_data_dir=None, caption_extension=".caption", num_repeats=10, is_reg=False):
|
||||
|
||||
config = {
|
||||
'general': {
|
||||
@@ -367,6 +316,13 @@ def generate_dataset_config(output_path, enable_bucket=True, resolution=512, bat
|
||||
if conditioning_data_dir is not None:
|
||||
config["datasets"][0]["subsets"][0]["conditioning_data_dir"] = conditioning_data_dir
|
||||
|
||||
if is_reg:
|
||||
config["datasets"][0]["subsets"].append({
|
||||
"is_reg": True,
|
||||
"image_dir": os.path.join(workspace_dir, "reg_images"),
|
||||
"num_repeats": 1,
|
||||
})
|
||||
|
||||
if output_path.endswith(".toml"):
|
||||
check_install()
|
||||
import toml
|
||||
@@ -479,12 +435,8 @@ def run_hook_kohya_ss_run_file(workspace_dir, output_name, kohya_ss_tool_dir, tr
|
||||
|
||||
from .mz_train_tools_utils import HSubprocess
|
||||
|
||||
screen_name = None
|
||||
if use_screen:
|
||||
screen_name = "mz_train_tools_core"
|
||||
|
||||
process_instance = HSubprocess(
|
||||
cmd_list, screen_name=screen_name)
|
||||
cmd_list)
|
||||
process_instance.wait()
|
||||
|
||||
stop_server()
|
||||
@@ -495,6 +447,24 @@ def run_hook_kohya_ss_run_file(workspace_dir, output_name, kohya_ss_tool_dir, tr
|
||||
raise Exception(f"训练失败!!! 具体报错信息请查看控制台...")
|
||||
|
||||
|
||||
def get_models_type(model_path):
|
||||
|
||||
from comfy.model_detection import model_config_from_unet, unet_prefix_from_state_dict, detect_unet_config
|
||||
from comfy.supported_models import SDXL, SD15
|
||||
import safetensors.torch
|
||||
sd = safetensors.torch.load_file(model_path, device="cpu")
|
||||
unet_prefix = unet_prefix_from_state_dict(sd)
|
||||
|
||||
model_type = model_config_from_unet(sd, unet_prefix)
|
||||
|
||||
if isinstance(model_type, SD15):
|
||||
return "lora_sd1_5"
|
||||
if isinstance(model_type, SDXL):
|
||||
return "lora_sdxl"
|
||||
|
||||
raise Exception(f"不支持的模型类型: {model_type}")
|
||||
|
||||
|
||||
def generate_kohya_ss_config(args):
|
||||
args = args.copy()
|
||||
workspace_config = args.get("workspace_config", {}).copy()
|
||||
@@ -512,7 +482,11 @@ def generate_kohya_ss_config(args):
|
||||
|
||||
workspace_config_file = os.path.join(workspace_dir, "config.json")
|
||||
|
||||
train_config_template = args.get("train_config_template", None)
|
||||
ckpt_name = args.get("ckpt_name", "")
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
|
||||
train_config_template = get_models_type(ckpt_path)
|
||||
|
||||
train_config_template_dir = os.path.join(
|
||||
os.path.dirname(__file__), "configs", "kohya_ss_lora")
|
||||
train_config_template_file = os.path.join(
|
||||
@@ -525,7 +499,6 @@ def generate_kohya_ss_config(args):
|
||||
config["metadata"]["train_type"] = train_config_template
|
||||
ckpt_name = args.get("ckpt_name", "")
|
||||
if ckpt_name != "" and ckpt_name is not None:
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
config["train_config"]["pretrained_model_name_or_path"] = ckpt_path
|
||||
|
||||
# output_dir
|
||||
@@ -582,6 +555,41 @@ def generate_kohya_ss_config(args):
|
||||
return config
|
||||
|
||||
|
||||
def generate_regularization_folder(ckpt_name, workspace_dir, caption_extension,
|
||||
width=512, height=512):
|
||||
train_images_dir = os.path.join(workspace_dir, "train_images")
|
||||
saved_images_path = os.listdir(train_images_dir)
|
||||
|
||||
reg_dir = os.path.join(workspace_dir, "reg_images")
|
||||
|
||||
caption_files = []
|
||||
caption_value = []
|
||||
for i, filename in enumerate(saved_images_path):
|
||||
if filename.lower().endswith(caption_extension):
|
||||
caption_files.append(filename)
|
||||
with open(os.path.join(train_images_dir, filename), "r", encoding="utf-8") as f:
|
||||
caption_value.append(f.read())
|
||||
|
||||
if len(caption_files) == 0:
|
||||
return
|
||||
os.makedirs(reg_dir, exist_ok=True)
|
||||
for filename, caption in zip(caption_files, caption_value):
|
||||
if caption == "":
|
||||
continue
|
||||
|
||||
with open(os.path.join(reg_dir, filename), "w", encoding="utf-8") as f:
|
||||
f.write(caption)
|
||||
|
||||
image_filename = filename.replace(caption_extension, ".webp")
|
||||
Utils.comfyui_fast_generate_picture(
|
||||
ckpt_name=ckpt_name,
|
||||
prompt=caption,
|
||||
output=os.path.join(reg_dir, image_filename),
|
||||
width=width,
|
||||
height=height,
|
||||
)
|
||||
|
||||
|
||||
def MZ_KohyaSSTrain_call(args={}):
|
||||
args = args.copy()
|
||||
workspace_config = args.get("workspace_config").copy()
|
||||
@@ -678,8 +686,36 @@ def MZ_KohyaSSTrain_call(args={}):
|
||||
with open(workspace_config_file, "w", encoding="utf-8") as f:
|
||||
json.dump(config, f, indent=4, ensure_ascii=False)
|
||||
|
||||
dataset_config_file = config["train_config"]["dataset_config"]
|
||||
with open(dataset_config_file, "r", encoding="utf-8") as f:
|
||||
dataset_config = json.load(f)
|
||||
|
||||
if dataset_config is None:
|
||||
raise Exception("dataset_config is None")
|
||||
|
||||
subsets_list = dataset_config.get("datasets")[0].get("subsets")
|
||||
is_reg = True if len(subsets_list) == 2 and subsets_list[1].get(
|
||||
"is_reg") else False
|
||||
if is_reg:
|
||||
caption_extension = subsets_list[0].get("caption_extension")
|
||||
ckpt_name = args.get("ckpt_name")
|
||||
resolution = dataset_config.get("datasets")[0].get("resolution")
|
||||
if f"{resolution}".find("x") != -1:
|
||||
resolutionsp = resolution.split("x")
|
||||
width = int(resolutionsp[0])
|
||||
height = int(resolutionsp[1])
|
||||
elif f"{resolution}".find("*") != -1:
|
||||
resolutionsp = resolution.split("*")
|
||||
width = int(resolutionsp[0])
|
||||
height = int(resolutionsp[1])
|
||||
else:
|
||||
width = height = int(resolution)
|
||||
|
||||
generate_regularization_folder(
|
||||
ckpt_name, workspace_dir, caption_extension, width=width, height=height)
|
||||
|
||||
# raise Exception(
|
||||
# f"config: {json.dumps(config, indent=4, ensure_ascii=False)}")
|
||||
# f"config: {json.dumps(dataset_config, indent=4, ensure_ascii=False)}")
|
||||
|
||||
output_name = config["train_config"].get("output_name")
|
||||
|
||||
@@ -704,21 +740,6 @@ def MZ_KohyaSSTrain_call(args={}):
|
||||
|
||||
run_hook_kohya_ss_run_file(
|
||||
workspace_dir, output_name, kohya_ss_tool_dir, "run_controlnet_sd1_5")
|
||||
elif train_type == "lora_hunyuan1_2" or train_type == "lora_hunyuan1_1":
|
||||
hunyuan_models_config = args.get(
|
||||
"hunyuan_models_config", {})
|
||||
from .mz_train_tools_core_HYDiT import check_model_auto_download
|
||||
|
||||
hunyuan_models_config["version"] = config.get(
|
||||
"metadata").get("version")
|
||||
config["other_config"]["hunyuan_models_config"] = check_model_auto_download(
|
||||
hunyuan_models_config)
|
||||
|
||||
with open(workspace_config_file, "w", encoding="utf-8") as f:
|
||||
json.dump(config, f, indent=4, ensure_ascii=False)
|
||||
|
||||
run_hook_kohya_ss_run_file(
|
||||
workspace_dir, output_name, kohya_ss_tool_dir, "run_lora_hunyuan1_2")
|
||||
else:
|
||||
raise Exception(
|
||||
f"暂时不支持的训练类型: {train_type}")
|
||||
@@ -726,309 +747,3 @@ def MZ_KohyaSSTrain_call(args={}):
|
||||
return (
|
||||
"训练完成",
|
||||
)
|
||||
|
||||
|
||||
def MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I_call(args={}):
|
||||
args = args.copy()
|
||||
MZ_KohyaSSCloneRepo_call(args)
|
||||
from .mz_train_tools_core_HYDiT import check_model_auto_download
|
||||
args = check_model_auto_download(args)
|
||||
import numpy as np
|
||||
import torch
|
||||
seed = args.get("seed", 0)
|
||||
torch.manual_seed(seed)
|
||||
from packaging import version
|
||||
from transformers import AutoTokenizer, BertModel
|
||||
from diffusers.models import AutoencoderKL
|
||||
try:
|
||||
from k_diffusion.external import DiscreteVDDPMDenoiser
|
||||
from k_diffusion.sampling import sample_euler_ancestral, get_sigmas_exponential, sample_dpmpp_2m_sde
|
||||
except ImportError:
|
||||
subprocess.check_call(
|
||||
[sys.executable, "-m", "pip", "install", "k-diffusion"])
|
||||
from k_diffusion.external import DiscreteVDDPMDenoiser
|
||||
from k_diffusion.sampling import sample_euler_ancestral, get_sigmas_exponential, sample_dpmpp_2m_sde
|
||||
|
||||
branch_local_name = args.get("branch_local_name")
|
||||
kohya_ss_tool_dir = os.path.join(
|
||||
Utils.get_minus_zone_models_path(), "train_tools", branch_local_name)
|
||||
if kohya_ss_tool_dir not in sys.path:
|
||||
sys.path.append(kohya_ss_tool_dir)
|
||||
from library.hunyuan_models import DiT_g_2, MT5Embedder
|
||||
from library.hunyuan_utils import get_cond, calc_rope
|
||||
from networks.lora import create_network_from_weights
|
||||
|
||||
def load_scheduler_sigmas(beta_start=0.00085, beta_end=0.018, num_train_timesteps=1000):
|
||||
betas = torch.linspace(beta_start**0.5, beta_end **
|
||||
0.5, num_train_timesteps, dtype=torch.float32) ** 2
|
||||
alphas = 1.0 - betas
|
||||
alphas_cumprod = torch.cumprod(alphas, dim=0)
|
||||
|
||||
sigmas = np.array(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5)
|
||||
sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32)
|
||||
sigmas = torch.from_numpy(sigmas)
|
||||
return alphas_cumprod, sigmas
|
||||
|
||||
version = args.get("version")
|
||||
BETA_END = None
|
||||
USE_EXTRA_COND = None
|
||||
if version == "1.1":
|
||||
BETA_END = 0.03
|
||||
USE_EXTRA_COND = True
|
||||
else:
|
||||
BETA_END = 0.018
|
||||
USE_EXTRA_COND = False
|
||||
|
||||
ATTN_MODE = "xformers"
|
||||
CLIP_TOKENS = 75 * 2 + 2
|
||||
dtype = DTYPE = torch.float16
|
||||
device = DEVICE = "cuda"
|
||||
|
||||
image = None
|
||||
|
||||
with torch.inference_mode(True), torch.no_grad():
|
||||
alphas, sigmas = load_scheduler_sigmas(beta_end=BETA_END)
|
||||
|
||||
tokenizer_path = args.get("tokenizer_path")
|
||||
clip_tokenizer = AutoTokenizer.from_pretrained(
|
||||
tokenizer_path, local_files_only=True)
|
||||
clip_tokenizer.eos_token_id = 2
|
||||
|
||||
text_encoder_path = args.get("text_encoder_path")
|
||||
|
||||
clip_encoder = Utils.model_cache_get(
|
||||
model_type="HYDiT_clip_encoder", model_path=text_encoder_path,)
|
||||
if clip_encoder is None:
|
||||
clip_encoder = (
|
||||
BertModel.from_pretrained(
|
||||
text_encoder_path, local_files_only=True).to(device).to(dtype)
|
||||
)
|
||||
Utils.model_cache_set(
|
||||
model_type="HYDiT_clip_encoder", model_path=text_encoder_path, model=clip_encoder)
|
||||
|
||||
t5_encoder_path = args.get("t5_encoder_path")
|
||||
if t5_encoder_path != "none" and os.path.exists(t5_encoder_path):
|
||||
mt5_embedder = Utils.model_cache_get(
|
||||
model_type="HYDiT_mt5_embedder", model_path=t5_encoder_path,)
|
||||
if mt5_embedder is None:
|
||||
mt5_embedder = (
|
||||
MT5Embedder(t5_encoder_path, torch_dtype=dtype,
|
||||
max_length=256,).to(device).to(dtype)
|
||||
)
|
||||
Utils.model_cache_set(
|
||||
model_type="HYDiT_mt5_embedder", model_path=t5_encoder_path, model=mt5_embedder)
|
||||
else:
|
||||
from .mz_train_tools_utils import CustomizeMT5Embedder
|
||||
mt5_embedder = (
|
||||
CustomizeMT5Embedder(
|
||||
batch_size=1,
|
||||
)
|
||||
.to(device)
|
||||
.to(dtype)
|
||||
)
|
||||
|
||||
vae_ema_path = args.get("vae_ema_path")
|
||||
vae = Utils.model_cache_get(
|
||||
model_type="HYDiT_vae", model_path=vae_ema_path,)
|
||||
if vae is None:
|
||||
vae = (
|
||||
AutoencoderKL.from_pretrained(
|
||||
vae_ema_path, local_files_only=True)
|
||||
.to(device)
|
||||
.to(dtype)
|
||||
)
|
||||
Utils.model_cache_set(
|
||||
model_type="HYDiT_vae", model_path=vae_ema_path, model=vae)
|
||||
|
||||
unet_path = args.get("unet_path")
|
||||
lora_path = args.get("lora_path")
|
||||
|
||||
denoiser_args = Utils.model_cache_get(
|
||||
model_type="HYDiT_unet_merge_lora", model_path=f"{unet_path}_{lora_path}")
|
||||
|
||||
if denoiser_args is None:
|
||||
denoiser, patch_size, head_dim = DiT_g_2(
|
||||
input_size=(128, 128), use_extra_cond=USE_EXTRA_COND)
|
||||
state_dict = torch.load(unet_path)
|
||||
denoiser.load_state_dict(state_dict)
|
||||
denoiser.to(device).to(dtype)
|
||||
denoiser.eval()
|
||||
denoiser.disable_fp32_silu()
|
||||
denoiser.disable_fp32_layer_norm()
|
||||
denoiser.set_attn_mode(ATTN_MODE)
|
||||
|
||||
if lora_path is not None and lora_path != "none":
|
||||
if not os.path.exists(lora_path):
|
||||
raise Exception(f"lora_path: {lora_path} 不存在")
|
||||
lora_net, state_dict = create_network_from_weights(
|
||||
multiplier=1.0,
|
||||
file=lora_path,
|
||||
vae=vae,
|
||||
text_encoder=[clip_encoder, mt5_embedder],
|
||||
unet=denoiser,
|
||||
)
|
||||
lora_net.apply_to(
|
||||
text_encoder=[clip_encoder, mt5_embedder],
|
||||
unet=denoiser,
|
||||
)
|
||||
lora_net.load_state_dict(state_dict)
|
||||
lora_net = lora_net.to(DEVICE, dtype=DTYPE)
|
||||
|
||||
Utils.model_cache_set(
|
||||
model_type="HYDiT_unet_merge_lora", model_path=f"{unet_path}_{lora_path}", model=(denoiser, patch_size, head_dim))
|
||||
else:
|
||||
denoiser, patch_size, head_dim = denoiser_args
|
||||
|
||||
vae.requires_grad_(False)
|
||||
mt5_embedder.to(torch.float16)
|
||||
prompt = args.get("prompt")
|
||||
negative_prompt = args.get("negative_prompt")
|
||||
with torch.autocast("cuda"):
|
||||
clip_h, clip_m, mt5_h, mt5_m = get_cond(
|
||||
prompt,
|
||||
mt5_embedder,
|
||||
clip_tokenizer,
|
||||
clip_encoder,
|
||||
# Should be same as original implementation with max_length_clip=77
|
||||
# Support 75*n + 2
|
||||
max_length_clip=CLIP_TOKENS,
|
||||
)
|
||||
neg_clip_h, neg_clip_m, neg_mt5_h, neg_mt5_m = get_cond(
|
||||
negative_prompt,
|
||||
mt5_embedder,
|
||||
clip_tokenizer,
|
||||
clip_encoder,
|
||||
max_length_clip=CLIP_TOKENS,
|
||||
)
|
||||
clip_h = torch.concat([clip_h, neg_clip_h], dim=0)
|
||||
clip_m = torch.concat([clip_m, neg_clip_m], dim=0)
|
||||
mt5_h = torch.concat([mt5_h, neg_mt5_h], dim=0)
|
||||
mt5_m = torch.concat([mt5_m, neg_mt5_m], dim=0)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
style = torch.as_tensor([0] * 2, device=DEVICE)
|
||||
W = args.get("width")
|
||||
H = args.get("height")
|
||||
|
||||
size_cond = [H, W, H, W, 0, 0]
|
||||
image_meta_size = torch.as_tensor([size_cond] * 2, device=DEVICE)
|
||||
freqs_cis_img = calc_rope(H, W, patch_size, head_dim)
|
||||
|
||||
denoiser_wrapper = DiscreteVDDPMDenoiser(
|
||||
# A quick patch for learn_sigma
|
||||
lambda *args, **kwargs: denoiser(* \
|
||||
args, **kwargs).chunk(2, dim=1)[0],
|
||||
alphas,
|
||||
False,
|
||||
).to(DEVICE)
|
||||
|
||||
CFG_SCALE = cfg = args.get("cfg", 5.0)
|
||||
STEPS = steps = args.get("steps", 25)
|
||||
|
||||
def cfg_denoise_func(x, sigma):
|
||||
cond, uncond = denoiser_wrapper(
|
||||
x.repeat(2, 1, 1, 1),
|
||||
sigma.repeat(2),
|
||||
encoder_hidden_states=clip_h,
|
||||
text_embedding_mask=clip_m,
|
||||
encoder_hidden_states_t5=mt5_h,
|
||||
text_embedding_mask_t5=mt5_m,
|
||||
image_meta_size=image_meta_size,
|
||||
style=style,
|
||||
cos_cis_img=freqs_cis_img[0],
|
||||
sin_cis_img=freqs_cis_img[1],
|
||||
).chunk(2, dim=0)
|
||||
return uncond + (cond - uncond) * CFG_SCALE
|
||||
|
||||
sigmas = denoiser_wrapper.get_sigmas(STEPS).to(DEVICE)
|
||||
sigmas = get_sigmas_exponential(
|
||||
STEPS, denoiser_wrapper.sigma_min, denoiser_wrapper.sigma_max, DEVICE
|
||||
)
|
||||
x1 = torch.randn(1, 4, H // 8, W // 8,
|
||||
dtype=torch.float16, device=DEVICE)
|
||||
|
||||
pbar = Utils.progress_bar(STEPS, "sdxl")
|
||||
preview = pbar.get_previewer()
|
||||
|
||||
def generate_callback(args):
|
||||
try:
|
||||
i = args.get("i")
|
||||
latents = args.get("denoised")
|
||||
# 判断是否存在decode_latent_to_preview_image
|
||||
if hasattr(preview, "decode_latent_to_preview_image"):
|
||||
pil_img = preview.decode_latent_to_preview_image(
|
||||
None,
|
||||
latents,
|
||||
)[1]
|
||||
else:
|
||||
pil_img = None
|
||||
pbar.update(i, STEPS, pil_img)
|
||||
except Exception as e:
|
||||
print("generate_callback error:", e)
|
||||
raise e
|
||||
|
||||
with torch.autocast("cuda"):
|
||||
scheduler = args.get("scheduler")
|
||||
if scheduler == "euler_ancestral":
|
||||
sample = sample_euler_ancestral(
|
||||
cfg_denoise_func,
|
||||
x1 * sigmas[0],
|
||||
sigmas,
|
||||
callback=generate_callback,
|
||||
)
|
||||
else:
|
||||
sample = sample_dpmpp_2m_sde(
|
||||
cfg_denoise_func,
|
||||
x1 * sigmas[0],
|
||||
sigmas,
|
||||
callback=generate_callback,
|
||||
)
|
||||
torch.cuda.empty_cache()
|
||||
with torch.no_grad():
|
||||
latent = sample / 0.13025
|
||||
image = vae.decode(latent).sample
|
||||
image = (image / 2 + 0.5).clamp(0, 1)
|
||||
image = image.permute(0, 2, 3, 1)
|
||||
|
||||
keep_device = args.get("keep_device", "enable")
|
||||
if keep_device == "disable":
|
||||
Utils.model_cache_clean(
|
||||
model_type="HYDiT_clip_encoder")
|
||||
Utils.model_cache_clean(
|
||||
model_type="HYDiT_mt5_embedder")
|
||||
Utils.model_cache_clean(
|
||||
model_type="HYDiT_vae")
|
||||
Utils.model_cache_clean(
|
||||
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,)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+95
-15
@@ -29,6 +29,90 @@ CACHE_POOL = {}
|
||||
|
||||
|
||||
class Utils:
|
||||
def comfyui_fast_generate_picture(ckpt_name, prompt, output, width=512, height=512):
|
||||
if os.path.exists(output):
|
||||
return
|
||||
import nodes
|
||||
import node_helpers
|
||||
model, clip, vae = nodes.CheckpointLoaderSimple().load_checkpoint(
|
||||
ckpt_name=ckpt_name)
|
||||
seed = 66
|
||||
steps = 20
|
||||
cfg = 7
|
||||
sampler_name = "euler_ancestral"
|
||||
scheduler = "normal"
|
||||
|
||||
# EmptyLatentImage
|
||||
latent_image = nodes.EmptyLatentImage().generate(width, height)[0]
|
||||
|
||||
# CLIPTextEncode
|
||||
positive = nodes.CLIPTextEncode().encode(clip, prompt)[0]
|
||||
negative = nodes.CLIPTextEncode().encode(clip, "")[0]
|
||||
|
||||
# KSampler
|
||||
latent = nodes.KSampler().sample(
|
||||
model=model,
|
||||
cfg=cfg, seed=seed, steps=steps, sampler_name=sampler_name, scheduler=scheduler, latent_image=latent_image, positive=positive, negative=negative)[0]
|
||||
|
||||
# VAEDecode
|
||||
image = nodes.VAEDecode().decode(vae, latent)[0]
|
||||
|
||||
# SaveImage
|
||||
Utils.tensor2pil(image).save(output)
|
||||
|
||||
def clone_repo(args={}):
|
||||
git_accelerate_urls = {
|
||||
"githubfast": "githubfast.com",
|
||||
"521github": "521github.com",
|
||||
"kkgithub": "kkgithub.com",
|
||||
}
|
||||
mz_dir = Utils.get_minus_zone_models_path()
|
||||
|
||||
branch_repoid = args.get("branch_repoid", "kohya-ss/sd-scripts")
|
||||
branch_local_name = args.get("branch_local_name", "kohya_ss_lora")
|
||||
|
||||
git_url = f"https://github.com/{branch_repoid}"
|
||||
source = args.get("source", "github")
|
||||
kohya_ss_lora_dir = os.path.join(
|
||||
mz_dir, "train_tools", branch_local_name)
|
||||
if git_accelerate_urls.get(source, None) is not None:
|
||||
git_url = f"https://{git_accelerate_urls[source]}/{branch_repoid}"
|
||||
try:
|
||||
if not os.path.exists(kohya_ss_lora_dir) or not os.path.exists(os.path.join(kohya_ss_lora_dir, ".git")):
|
||||
subprocess.run(
|
||||
["git", "clone", "--depth", "1", git_url, kohya_ss_lora_dir], check=True)
|
||||
|
||||
# 切换远程分支 git remote set-branches origin 'main'
|
||||
branch = args.get("branch", "main")
|
||||
|
||||
# 查看本地分支是否一致
|
||||
short_result = subprocess.run(
|
||||
["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=kohya_ss_lora_dir, stdout=subprocess.PIPE, check=True)
|
||||
result = subprocess.run(
|
||||
["git", "rev-parse", "HEAD"], cwd=kohya_ss_lora_dir, stdout=subprocess.PIPE, check=True)
|
||||
|
||||
short_current_branch = short_result.stdout.decode().strip()
|
||||
long_current_branch = result.stdout.decode().strip()
|
||||
print(
|
||||
f"当前分支(current branch): {long_current_branch}({short_current_branch})")
|
||||
print(f"目标分支(target branch): {branch}")
|
||||
|
||||
if branch != long_current_branch and branch != short_current_branch:
|
||||
subprocess.run(
|
||||
["git", "remote", "set-branches", "origin", branch], cwd=kohya_ss_lora_dir, check=True)
|
||||
subprocess.run(
|
||||
["git", "fetch", "--depth", "1", "origin", branch], cwd=kohya_ss_lora_dir, check=True)
|
||||
|
||||
# 恢复所有文件
|
||||
subprocess.run(
|
||||
["git", "checkout", "."], cwd=kohya_ss_lora_dir, check=True)
|
||||
|
||||
subprocess.run(
|
||||
["git", "checkout", branch], cwd=kohya_ss_lora_dir, check=True)
|
||||
|
||||
except Exception as e:
|
||||
raise Exception(f"克隆kohya-ss/sd-scripts或者切换分支时出现异常,详细信息请查看控制台...")
|
||||
|
||||
def Md5(str):
|
||||
return hashlib.md5(str.encode('utf-8')).hexdigest()
|
||||
|
||||
@@ -1143,33 +1227,20 @@ class HSubprocess:
|
||||
screen_name = None
|
||||
_mswindows = False
|
||||
|
||||
def __init__(self, args, screen_name=None):
|
||||
def __init__(self, args):
|
||||
self._mswindows = (sys.platform == "win32")
|
||||
if self._mswindows:
|
||||
self.args = ["start", "/w"]
|
||||
self.args.extend(args)
|
||||
return
|
||||
|
||||
self.screen_name = screen_name
|
||||
if screen_name is not None:
|
||||
try:
|
||||
subprocess.check_call(["screen", "-v"])
|
||||
except Exception as e:
|
||||
raise Exception("Please install screen first.")
|
||||
|
||||
screen_cmd = ["screen", "-R", screen_name, "-m", ]
|
||||
screen_cmd.extend(args)
|
||||
|
||||
self.args = screen_cmd
|
||||
else:
|
||||
self.args = args
|
||||
self.args = args
|
||||
|
||||
def stop(self):
|
||||
if self._mswindows:
|
||||
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:
|
||||
@@ -1193,6 +1264,15 @@ class HSubprocess:
|
||||
self.process_instance_pid = None
|
||||
|
||||
def wait(self):
|
||||
if self._mswindows:
|
||||
self.wait_win()
|
||||
else:
|
||||
self.wait_linux()
|
||||
|
||||
def wait_linux(self):
|
||||
subprocess.run(self.args, check=True)
|
||||
|
||||
def wait_win(self):
|
||||
with subprocess.Popen(
|
||||
self.args,
|
||||
stdin=subprocess.PIPE,
|
||||
|
||||
Reference in New Issue
Block a user