'update'
This commit is contained in:
wailovet
+233
-57
@@ -95,7 +95,7 @@ NODE_CLASS_MAPPINGS["MZ_KohyaSSDatasetConfig"] = MZ_KohyaSSDatasetConfig
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSDatasetConfig"] = f"{AUTHOR_NAME} - KohyaSSDatasetConfig"
|
||||
|
||||
|
||||
class MZ_KohyaSSUseConfig:
|
||||
class MZ_KohyaSSUseConfig_oldversion:
|
||||
train_config_template_dir = os.path.join(
|
||||
os.path.dirname(__file__), "configs", "kohya_ss_lora"
|
||||
)
|
||||
@@ -127,7 +127,7 @@ class MZ_KohyaSSUseConfig:
|
||||
RETURN_NAMES = ("train_config",)
|
||||
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME
|
||||
CATEGORY = CATEGORY_NAME + "/v1"
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
@@ -135,8 +135,9 @@ class MZ_KohyaSSUseConfig:
|
||||
return mz_train_tools_core.MZ_KohyaSSUseConfig_call(kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_KohyaSSUseConfig"] = MZ_KohyaSSUseConfig
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSUseConfig"] = f"{AUTHOR_NAME} - KohyaSSUseConfig"
|
||||
NODE_CLASS_MAPPINGS["MZ_KohyaSSUseConfig"] = MZ_KohyaSSUseConfig_oldversion
|
||||
NODE_DISPLAY_NAME_MAPPINGS[
|
||||
"MZ_KohyaSSUseConfig"] = f"{AUTHOR_NAME} - KohyaSSUseConfig(old version)"
|
||||
|
||||
|
||||
class MZ_KohyaSSAdvConfig:
|
||||
@@ -148,6 +149,8 @@ class MZ_KohyaSSAdvConfig:
|
||||
"sdpa": (["enable", "disable"], {"default": "disable"}),
|
||||
"fp8_base": (["enable", "disable"], {"default": "disable"}),
|
||||
"mixed_precision": (["no", "fp16", "bf16"], {"default": "fp16"}),
|
||||
"gradient_accumulation_steps": ("INT", {"default": 1}),
|
||||
"gradient_checkpointing": (["enable", "disable"], {"default": "disable"}),
|
||||
"cache_latents": (["enable", "disable"], {"default": "enable"}),
|
||||
"cache_latents_to_disk": (["enable", "disable"], {"default": "enable"}),
|
||||
"network_dim": ("INT", {"default": 16}),
|
||||
@@ -226,7 +229,7 @@ class AlwaysEqualProxy(str):
|
||||
return False
|
||||
|
||||
|
||||
class MZ_KohyaSSLoraTrain:
|
||||
class MZ_KohyaSSTrain_oldversion:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -283,17 +286,118 @@ class MZ_KohyaSSLoraTrain:
|
||||
RETURN_NAMES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME
|
||||
CATEGORY = CATEGORY_NAME + "/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_KohyaSSLoraTrain
|
||||
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"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
loras = [
|
||||
"latest",
|
||||
"empty",
|
||||
]
|
||||
|
||||
workspaces_dir = os.path.join(
|
||||
folder_paths.output_directory, "mz_train_workspaces")
|
||||
|
||||
# 使用walk查询所有的workspace中的所有lora模型,lora存放在每个workspace的output目录下
|
||||
workspaces_loras = []
|
||||
for root, dirs, files in os.walk(workspaces_dir):
|
||||
|
||||
# 排除隐藏文件夹
|
||||
dirs[:] = [d for d in dirs if not d.startswith(".")]
|
||||
if root.endswith("output"):
|
||||
for file in files:
|
||||
if file.endswith(".safetensors"):
|
||||
workspaces_loras.append(
|
||||
os.path.join(root, file)
|
||||
)
|
||||
|
||||
# 按创建时间排序
|
||||
workspaces_loras = sorted(
|
||||
workspaces_loras, key=lambda x: os.path.getctime(x), reverse=True)
|
||||
|
||||
comfyui_full_loras = []
|
||||
comfyui_loras = folder_paths.get_filename_list("loras")
|
||||
for lora in comfyui_loras:
|
||||
lora_path = folder_paths.get_full_path("loras", lora)
|
||||
comfyui_full_loras.append(lora_path)
|
||||
|
||||
# 按创建时间排序
|
||||
comfyui_full_loras = sorted(
|
||||
comfyui_full_loras, key=lambda x: os.path.getctime(x), reverse=True)
|
||||
|
||||
loras = loras + workspaces_loras + comfyui_full_loras
|
||||
|
||||
train_config_templates = Utils.listdir(s.train_config_template_dir)
|
||||
|
||||
# 去掉json后缀
|
||||
train_config_templates = [os.path.splitext(x)[0]
|
||||
for x in train_config_templates]
|
||||
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}),
|
||||
"save_every_n_epochs": ("INT", {"default": 10}),
|
||||
"learning_rate": ("STRING", {"default": "1e-5"}),
|
||||
|
||||
"base_lora": (loras, {"default": "latest"}),
|
||||
"sample_generate": (["enable", "disable"], {"default": "enable"}),
|
||||
"sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
|
||||
},
|
||||
"optional": {
|
||||
"advanced_config": ("MZ_TT_SS_AdvConfig",),
|
||||
"caption_completed_flag": (AlwaysEqualProxy("*"),),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
RETURN_NAMES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME
|
||||
|
||||
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)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_KohyaSSLoraTrain"] = MZ_KohyaSSLoraTrain
|
||||
NODE_DISPLAY_NAME_MAPPINGS[
|
||||
"MZ_KohyaSSLoraTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(lora)"
|
||||
@@ -348,6 +452,102 @@ 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()
|
||||
|
||||
# "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"}),
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
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_LoadImagesFromDirectoryPath:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -461,36 +661,6 @@ NODE_CLASS_MAPPINGS["MZ_HYDiTDatasetConfig"] = MZ_HYDiTDatasetConfig
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTDatasetConfig"] = f"{AUTHOR_NAME} - HYDiTDatasetConfig"
|
||||
|
||||
|
||||
# class MZ_HYDiTUseConfig:
|
||||
# @classmethod
|
||||
# def INPUT_TYPES(s):
|
||||
|
||||
# return {
|
||||
# "required": {
|
||||
# "workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",),
|
||||
# "workspace_images_dir": ("STRING", {"forceInput": True}),
|
||||
# },
|
||||
# "optional": {
|
||||
# }
|
||||
# }
|
||||
|
||||
# RETURN_TYPES = ("MZ_TT_HYDiT_TrainConfig",)
|
||||
# RETURN_NAMES = ("train_config",)
|
||||
|
||||
# FUNCTION = "start"
|
||||
|
||||
# CATEGORY = CATEGORY_NAME
|
||||
|
||||
# def start(self, **kwargs):
|
||||
# importlib.reload(mz_train_tools_core_HYDiT)
|
||||
# # kwargs["train_config_template_dir"] = self.train_config_template_dir
|
||||
# return mz_train_tools_core_HYDiT.MZ_HYDiTUseConfig_call(kwargs)
|
||||
|
||||
|
||||
# NODE_CLASS_MAPPINGS["MZ_HYDiTUseConfig"] = MZ_HYDiTUseConfig
|
||||
# NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTUseConfig"] = f"{AUTHOR_NAME} - HYDiTUseConfig"
|
||||
|
||||
|
||||
class MZ_HYDiTAdvConfig:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -548,30 +718,36 @@ 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):
|
||||
|
||||
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"),
|
||||
])
|
||||
models, folders, vae_models, unet_models, loras = HYDiT_paths()
|
||||
|
||||
return {
|
||||
"required": {
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
{
|
||||
"metadata": {
|
||||
"train_type": "lora_hunyuan1_1",
|
||||
"version": "1.1"
|
||||
},
|
||||
"train_config": {
|
||||
"max_train_steps": "4500",
|
||||
"xformers": true,
|
||||
"sdpa": false,
|
||||
"fp8_base": false,
|
||||
"mixed_precision": "fp16",
|
||||
"cache_latents": true,
|
||||
"cache_latents_to_disk": true,
|
||||
"network_dim": "16",
|
||||
"network_alpha": "8",
|
||||
"network_module": "networks.lora",
|
||||
"network_train_unet_only": true,
|
||||
"learning_rate": "1e-5",
|
||||
"lr_scheduler": "cosine_with_restarts",
|
||||
"optimizer_type": "AdamW",
|
||||
"save_every_n_epochs": "20",
|
||||
"shuffle_caption": true,
|
||||
"lr_warmup_steps": "0",
|
||||
"save_precision": "fp16",
|
||||
"lr_scheduler_num_cycles": "1",
|
||||
"persistent_data_loader_workers": true,
|
||||
"no_metadata": true,
|
||||
"noise_offset": "0.1",
|
||||
"output_dir": "",
|
||||
"output_name": "",
|
||||
"no_half_vae": true,
|
||||
"lowram": false,
|
||||
"max_token_length": "225",
|
||||
"use_extra_cond": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
{
|
||||
"metadata": {
|
||||
"train_type": "lora_hunyuan1_2",
|
||||
"version": "1.2"
|
||||
},
|
||||
"train_config": {
|
||||
"max_train_steps": "4500",
|
||||
"xformers": true,
|
||||
"sdpa": false,
|
||||
"fp8_base": false,
|
||||
"mixed_precision": "fp16",
|
||||
"cache_latents": true,
|
||||
"cache_latents_to_disk": true,
|
||||
"network_dim": "4",
|
||||
"network_alpha": "4",
|
||||
"network_module": "networks.lora",
|
||||
"network_train_unet_only": true,
|
||||
"learning_rate": "1e-5",
|
||||
"lr_scheduler": "cosine_with_restarts",
|
||||
"optimizer_type": "AdamW",
|
||||
"save_every_n_epochs": "20",
|
||||
"shuffle_caption": false,
|
||||
"lr_warmup_steps": "0",
|
||||
"save_precision": "fp16",
|
||||
"lr_scheduler_num_cycles": "1",
|
||||
"persistent_data_loader_workers": true,
|
||||
"no_metadata": true,
|
||||
"noise_offset": "0.00",
|
||||
"output_dir": "",
|
||||
"output_name": "",
|
||||
"no_half_vae": true,
|
||||
"lowram": false,
|
||||
"max_token_length": "150"
|
||||
}
|
||||
}
|
||||
@@ -19,7 +19,7 @@
|
||||
"lr_scheduler": "cosine_with_restarts",
|
||||
"optimizer_type": "AdamW",
|
||||
"save_every_n_epochs": "20",
|
||||
"shuffle_caption": true,
|
||||
"shuffle_caption": false,
|
||||
"lr_warmup_steps": "0",
|
||||
"save_precision": "fp16",
|
||||
"lr_scheduler_num_cycles": "1",
|
||||
|
||||
@@ -0,0 +1,326 @@
|
||||
|
||||
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
|
||||
+177
-24
@@ -68,19 +68,21 @@ def sample_images(self, *args, **kwargs):
|
||||
|
||||
prompt_dict_list = other_config.get("prompt_dict_list", [])
|
||||
if len(prompt_dict_list) == 0:
|
||||
seed = other_config.get("seed", 0)
|
||||
prompt_dict = {
|
||||
"controlnet_image": other_config.get("controlnet_image", None),
|
||||
"prompt": other_config.get("sample_prompt", ""),
|
||||
"seed": seed,
|
||||
"negative_prompt": "",
|
||||
"enum": 0,
|
||||
"sample_sampler": "euler_a",
|
||||
"sample_steps": 20,
|
||||
"scale": 5.0,
|
||||
}
|
||||
sample_prompt = other_config.get("sample_prompt", None)
|
||||
if sample_prompt is not None:
|
||||
seed = other_config.get("seed", 0)
|
||||
prompt_dict = {
|
||||
"controlnet_image": other_config.get("controlnet_image", None),
|
||||
"prompt": other_config.get("sample_prompt", ""),
|
||||
"seed": seed,
|
||||
"negative_prompt": "",
|
||||
"enum": 0,
|
||||
"sample_sampler": "euler_a",
|
||||
"sample_steps": 20,
|
||||
"scale": 5.0,
|
||||
}
|
||||
|
||||
prompt_dict_list.append(prompt_dict)
|
||||
prompt_dict_list.append(prompt_dict)
|
||||
else:
|
||||
for i, prompt_dict in enumerate(prompt_dict_list):
|
||||
if prompt_dict.get("controlnet_image", None) is None:
|
||||
@@ -91,18 +93,20 @@ def sample_images(self, *args, **kwargs):
|
||||
prompt_dict["negative_prompt"] = ""
|
||||
if prompt_dict.get("enum", None) is None:
|
||||
prompt_dict["enum"] = i
|
||||
hook_kohya_ss_utils.generate_image(
|
||||
pipe_class=sample_images_pipe_class,
|
||||
cmd_args=cmd_args,
|
||||
accelerator=accelerator,
|
||||
epoch=epoch,
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=tokenizer,
|
||||
unet=unet,
|
||||
vae=vae,
|
||||
prompt_dict_list=prompt_dict_list,
|
||||
controlnet=controlnet,
|
||||
)
|
||||
|
||||
if prompt_dict_list is not None and len(prompt_dict_list) > 0:
|
||||
hook_kohya_ss_utils.generate_image(
|
||||
pipe_class=sample_images_pipe_class,
|
||||
cmd_args=cmd_args,
|
||||
accelerator=accelerator,
|
||||
epoch=epoch,
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=tokenizer,
|
||||
unet=unet,
|
||||
vae=vae,
|
||||
prompt_dict_list=prompt_dict_list,
|
||||
controlnet=controlnet,
|
||||
)
|
||||
|
||||
LOG({
|
||||
"type": "sample_images",
|
||||
@@ -201,10 +205,159 @@ def run_controlnet_sd1_5():
|
||||
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_config = json.loads(train_config_json)
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
|
||||
+168
-1
@@ -181,6 +181,170 @@ class TokenizersWrapper:
|
||||
f"TokenizersWrapper: {self.typed} is not supported")
|
||||
|
||||
|
||||
from transformers import AutoTokenizer, MT5EncoderModel
|
||||
from torch import nn
|
||||
|
||||
|
||||
class CustomizeEmbedsModel(nn.Module):
|
||||
dtype = torch.float16
|
||||
shared = None
|
||||
# x = torch.zeros(1, 1, 256, 2048)
|
||||
x = None
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__()
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
# print("CustomizeEmbedsModel forward: args:", args)
|
||||
# print("CustomizeEmbedsModel forward: kwargs:", kwargs)
|
||||
input_ids = kwargs.get("input_ids", None)
|
||||
# if self.x is None:
|
||||
if True:
|
||||
if input_ids is None:
|
||||
batch_size = 1
|
||||
else:
|
||||
batch_size = input_ids.shape[0]
|
||||
|
||||
attention_mask = kwargs.get("attention_mask")
|
||||
attention_mask_dim = attention_mask.shape[1]
|
||||
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 = None
|
||||
attention_mask = None
|
||||
batch_size = 1
|
||||
|
||||
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
|
||||
self.batch_size = kwargs.get("batch_size", 1)
|
||||
|
||||
def tokenize(self, text):
|
||||
return text
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
# print("CustomizeTokenizer args:", args)
|
||||
# print("CustomizeTokenizer kwargs:", kwargs)
|
||||
|
||||
value = args[0]
|
||||
if isinstance(value, str):
|
||||
batch_size = 1
|
||||
else:
|
||||
batch_size = value.shape[0]
|
||||
|
||||
# print(f"CustomizeTokenizer batch_size: {batch_size}")
|
||||
# if self.input_ids is not None:
|
||||
# return self
|
||||
|
||||
self.input_ids = torch.zeros(batch_size, 256)
|
||||
self.attention_mask = torch.zeros(batch_size, 256)
|
||||
self['input_ids'] = self.input_ids
|
||||
self['attention_mask'] = self.attention_mask
|
||||
|
||||
# print("CustomizeTokenizer input_ids:", self.input_ids.shape)
|
||||
# print("CustomizeTokenizer attention_mask:", self.attention_mask.shape)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class CustomizeEmbeds():
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.tokenizer = CustomizeTokenizer()
|
||||
self.model = CustomizeEmbedsModel().to("cuda")
|
||||
self.max_length = 256
|
||||
|
||||
|
||||
class CustomizeMT5Embedder(nn.Module):
|
||||
device = torch.device("cuda")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_dir="t5-v1_1-xxl",
|
||||
model_kwargs=None,
|
||||
torch_dtype=None,
|
||||
use_tokenizer_only=False,
|
||||
max_length=128,
|
||||
batch_size=1,
|
||||
):
|
||||
super().__init__()
|
||||
self.torch_dtype = torch_dtype or torch.bfloat16
|
||||
self.max_length = max_length
|
||||
self.tokenizer = CustomizeTokenizer(
|
||||
batch_size=batch_size
|
||||
)
|
||||
self.model = CustomizeEmbedsModel().to("cuda")
|
||||
|
||||
def gradient_checkpointing_enable(self):
|
||||
pass
|
||||
|
||||
def gradient_checkpointing_disable(self):
|
||||
pass
|
||||
|
||||
def get_tokens_and_mask(self, texts):
|
||||
text_tokens_and_mask = self.tokenizer(
|
||||
texts,
|
||||
max_length=self.max_length,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_attention_mask=True,
|
||||
add_special_tokens=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
tokens = text_tokens_and_mask["input_ids"][0]
|
||||
mask = text_tokens_and_mask["attention_mask"][0]
|
||||
return tokens, mask
|
||||
|
||||
def get_text_embeddings(self, texts, attention_mask=True, layer_index=-1):
|
||||
text_tokens_and_mask = self.tokenizer(
|
||||
texts,
|
||||
max_length=self.max_length,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_attention_mask=True,
|
||||
add_special_tokens=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
outputs = self.model(
|
||||
input_ids=text_tokens_and_mask["input_ids"],
|
||||
attention_mask=(
|
||||
text_tokens_and_mask["attention_mask"]
|
||||
if attention_mask
|
||||
else None
|
||||
),
|
||||
output_hidden_states=True,
|
||||
)
|
||||
text_encoder_embs = outputs["hidden_states"][layer_index].detach()
|
||||
|
||||
return text_encoder_embs, text_tokens_and_mask["attention_mask"].to(self.device)
|
||||
|
||||
def get_input_ids(self, caption):
|
||||
return self.tokenizer(
|
||||
caption,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
max_length=self.max_length,
|
||||
return_tensors="pt",
|
||||
).input_ids
|
||||
|
||||
def get_hidden_states(self, input_ids, layer_index=-1):
|
||||
return self.get_text_embeddings(input_ids, layer_index=layer_index)
|
||||
|
||||
|
||||
def load_tokenizers(*args, **kwargs):
|
||||
return TokenizersWrapper("clip_large")
|
||||
|
||||
@@ -350,7 +514,10 @@ def _load_target_model(args: argparse.Namespace, weight_dtype, device="cpu", une
|
||||
return text_encoder, vae, unet, load_stable_diffusion_format
|
||||
|
||||
|
||||
def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encoder, unet, epoch, prompt_dict_list):
|
||||
def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encoder, unet, epoch, prompt_dict_list, **kwargs):
|
||||
if pipe_class is None:
|
||||
print("pipe_class is None")
|
||||
return
|
||||
import library.train_util
|
||||
|
||||
# for multi gpu distributed inference. this is a singleton, so it's safe to use it here
|
||||
|
||||
Submodule
+1
Submodule kohya_ss_lora added at 0dc79edc01
+55
-44
@@ -27,11 +27,15 @@ git_accelerate_urls = {
|
||||
|
||||
def MZ_KohyaSSInitWorkspace_call(args={}):
|
||||
mz_dir = Utils.get_minus_zone_models_path()
|
||||
git_url = "https://github.com/kohya-ss/sd-scripts"
|
||||
|
||||
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", "kohya_ss_lora")
|
||||
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]}/kohya-ss/sd-scripts"
|
||||
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(
|
||||
@@ -52,7 +56,7 @@ def MZ_KohyaSSInitWorkspace_call(args={}):
|
||||
f"当前分支(current branch): {long_current_branch}({short_current_branch})")
|
||||
print(f"目标分支(target branch): {branch}")
|
||||
|
||||
if branch != result.stdout.decode() and branch != short_result.stdout.decode():
|
||||
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(
|
||||
@@ -186,26 +190,26 @@ def MZ_KohyaSSUseConfig_call(args={}):
|
||||
raise Exception(f"工作区不存在: {workspace_dir}")
|
||||
|
||||
workspace_config_file = os.path.join(workspace_dir, "config.json")
|
||||
|
||||
train_config_template = args.get("train_config_template", None)
|
||||
if not os.path.exists(workspace_config_file):
|
||||
train_config_template_dir = args.get("train_config_template_dir", None)
|
||||
train_config_template_dir = args.get("train_config_template_dir", None)
|
||||
train_config_template_file = os.path.join(
|
||||
train_config_template_dir, train_config_template + ".json")
|
||||
|
||||
train_config_template_file = os.path.join(
|
||||
train_config_template_dir, train_config_template + ".json")
|
||||
|
||||
shutil.copy(train_config_template_file, workspace_config_file)
|
||||
# if not os.path.exists(workspace_config_file):
|
||||
# train_config_template_dir = args.get("train_config_template_dir", None)
|
||||
# train_config_template_file = os.path.join(
|
||||
# train_config_template_dir, train_config_template + ".json")
|
||||
# shutil.copy(train_config_template_file, workspace_config_file)
|
||||
|
||||
config = None
|
||||
with open(workspace_config_file, "r", encoding="utf-8") as f:
|
||||
with open(train_config_template_file, "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
config["metadata"]["train_type"] = train_config_template
|
||||
ckpt_name = args.get("ckpt_name", "")
|
||||
if ckpt_name == "":
|
||||
raise Exception("未选择模型文件(ckpt_name is required)")
|
||||
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
|
||||
config["train_config"]["pretrained_model_name_or_path"] = ckpt_path
|
||||
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
|
||||
output_dir = os.path.join(workspace_dir, "output")
|
||||
@@ -234,6 +238,8 @@ def MZ_KohyaSSUseConfig_call(args={}):
|
||||
args.get("learning_rate"))
|
||||
|
||||
advanced_config = args.get("save_advanced_config", {}).copy()
|
||||
if len(advanced_config) == 0:
|
||||
advanced_config = args.get("advanced_config", {}).copy()
|
||||
|
||||
for k in advanced_config:
|
||||
if type(advanced_config[k]) == str and advanced_config[k] == "":
|
||||
@@ -338,15 +344,6 @@ def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_
|
||||
toml.dump(config, f)
|
||||
|
||||
|
||||
# def latent2image(previews, latent_file):
|
||||
# samples = nodes.LoadLatent().load(latent_file)[0]
|
||||
# with torch.no_grad():
|
||||
|
||||
# samples = samples["samples"].to("cuda")
|
||||
# preview_bytes = previews.decode_latent_to_preview_image(
|
||||
# "", samples)
|
||||
# return preview_bytes[1]
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
@@ -371,6 +368,8 @@ def get_sample_images(train_config):
|
||||
pil_image = Image.open(os.path.join(sample_images_dir, image_file))
|
||||
pil_images.append([pil_image])
|
||||
pre_render_texts_x.append(image_file)
|
||||
if pil_images is None or len(pil_images) == 0:
|
||||
return Image.new("RGB", (512, 512), (255, 255, 255))
|
||||
result = Utils.xy_image(
|
||||
pre_render_images=pil_images,
|
||||
pre_render_texts_x=pre_render_texts_x,
|
||||
@@ -395,6 +394,8 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
|
||||
taesd_type = "sd1_5"
|
||||
if trainer_func.find("sdxl") != -1:
|
||||
taesd_type = "sdxl"
|
||||
if trainer_func.find("hunyuan1_1") != -1:
|
||||
taesd_type = "sdxl"
|
||||
|
||||
pb = Utils.progress_bar(train_config.get("max_train_steps"), taesd_type)
|
||||
|
||||
@@ -445,20 +446,19 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
|
||||
|
||||
|
||||
def MZ_KohyaSSTrain_call(args={}):
|
||||
args = args.copy()
|
||||
workspace_config = args.get("workspace_config", {})
|
||||
|
||||
# pb = Utils.progress_bar(1)
|
||||
# previewer = pb.get_previewer()
|
||||
|
||||
# img = latent2image(previewer, "/tmp/sample_images_4.latent")
|
||||
|
||||
# pb.update(1, 1, img)
|
||||
# raise Exception(f"MZ_KohyaSSTrain_call: {args}")
|
||||
train_config = args.get("train_config", {})
|
||||
workspace_config = train_config.get("workspace_config", {})
|
||||
MZ_KohyaSSUseConfig_call(train_config)
|
||||
|
||||
workspace_name = workspace_config.get("workspace_name", None)
|
||||
workspace_dir = os.path.join(
|
||||
folder_paths.output_directory, "mz_train_workspaces", workspace_name)
|
||||
|
||||
if not os.path.exists(workspace_dir):
|
||||
raise Exception(f"工作区不存在: {workspace_dir}")
|
||||
|
||||
workspace_config_file = os.path.join(workspace_dir, "config.json")
|
||||
|
||||
if not os.path.exists(workspace_config_file):
|
||||
@@ -471,8 +471,10 @@ def MZ_KohyaSSTrain_call(args={}):
|
||||
if config is None:
|
||||
raise Exception(f"读取配置文件失败: {workspace_config_file}")
|
||||
|
||||
branch_local_name = workspace_config.get(
|
||||
"branch_local_name", "kohya_ss_lora")
|
||||
kohya_ss_tool_dir = os.path.join(
|
||||
Utils.get_minus_zone_models_path(), "train_tools", "kohya_ss_lora")
|
||||
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)
|
||||
@@ -512,16 +514,14 @@ def MZ_KohyaSSTrain_call(args={}):
|
||||
if "network_dropout" in train_config:
|
||||
del train_config["network_dropout"]
|
||||
|
||||
|
||||
|
||||
|
||||
base_controlnet = args.get("base_controlnet", "empty")
|
||||
if base_controlnet == "empty":
|
||||
pass
|
||||
elif base_controlnet == "latest":
|
||||
workspace_controlnet_dir = os.path.join(workspace_dir, "output")
|
||||
if os.path.exists(workspace_controlnet_dir):
|
||||
workspace_controlnet_files = Utils.listdir(workspace_controlnet_dir)
|
||||
workspace_controlnet_files = Utils.listdir(
|
||||
workspace_controlnet_dir)
|
||||
workspace_controlnet_files = list(
|
||||
filter(lambda x: x.endswith(".safetensors"), workspace_controlnet_files))
|
||||
workspace_controlnet_files = list(
|
||||
@@ -538,12 +538,10 @@ def MZ_KohyaSSTrain_call(args={}):
|
||||
pass
|
||||
|
||||
if base_controlnet != "empty" and os.path.exists(base_controlnet):
|
||||
train_config["controlnet_model_name_or_path"] = base_controlnet
|
||||
|
||||
train_config["controlnet_model_name_or_path"] = base_controlnet
|
||||
|
||||
train_type = config.get("metadata").get("train_type")
|
||||
|
||||
|
||||
sample_generate = args.get("sample_generate", "enable")
|
||||
sample_prompt = args.get("sample_prompt", "")
|
||||
if sample_generate == "enable":
|
||||
@@ -565,12 +563,25 @@ def MZ_KohyaSSTrain_call(args={}):
|
||||
workspace_dir, "conditioning_images")
|
||||
conditioning_images_onec = ""
|
||||
if os.path.exists(conditioning_images_dir):
|
||||
conditioning_images_onec = Utils.listdir(conditioning_images_dir)[0]
|
||||
conditioning_images_onec = Utils.listdir(
|
||||
conditioning_images_dir)[0]
|
||||
other_config["controlnet_image"] = os.path.join(
|
||||
conditioning_images_dir, conditioning_images_onec)
|
||||
|
||||
|
||||
run_hook_kohya_ss_run_file(
|
||||
kohya_ss_tool_dir, train_config, "run_controlnet_sd1_5", other_config)
|
||||
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")
|
||||
other_config["hunyuan_models_config"] = check_model_auto_download(
|
||||
hunyuan_models_config)
|
||||
|
||||
run_hook_kohya_ss_run_file(
|
||||
kohya_ss_tool_dir, train_config, "run_lora_hunyuan1_2", other_config)
|
||||
else:
|
||||
raise Exception(
|
||||
f"暂时不支持的训练类型: {train_type}")
|
||||
|
||||
@@ -182,6 +182,9 @@ def MZ_HYDiTDatasetConfig_call(args={}):
|
||||
|
||||
|
||||
HYDiT_MODEL = {
|
||||
"HunyuanDiT-v1.2/t2i/model/pytorch_model_module.pt": {
|
||||
"url": "https://www.modelscope.cn/models/wailovet/hy1.2/resolve/master/t2i/model/pytorch_model_module.pt",
|
||||
},
|
||||
"HunyuanDiT/t2i/model/pytorch_model_ema.pt": {
|
||||
"url": "https://www.modelscope.cn/models/wailovet/hy1.1/resolve/master/t2i/model/pytorch_model_ema.pt",
|
||||
},
|
||||
@@ -238,11 +241,21 @@ def check_model_auto_download(args):
|
||||
hunyuan_base_path = os.path.join(
|
||||
Utils.get_comfyui_models_path(), "hunyuan")
|
||||
if unet_path == "auto":
|
||||
ema_to_module = args.get("ema_to_module", "enable") == "enable"
|
||||
if ema_to_module:
|
||||
download_file = "HunyuanDiT/t2i/model/pytorch_model_ema.pt"
|
||||
version = args.get("version", None)
|
||||
if version is None:
|
||||
ema_to_module = args.get("ema_to_module", "enable") == "enable"
|
||||
if ema_to_module:
|
||||
download_file = "HunyuanDiT/t2i/model/pytorch_model_ema.pt"
|
||||
else:
|
||||
download_file = "HunyuanDiT/t2i/model/pytorch_model_module.pt"
|
||||
else:
|
||||
download_file = "HunyuanDiT/t2i/model/pytorch_model_module.pt"
|
||||
if version == "1.2":
|
||||
download_file = "HunyuanDiT-v1.2/t2i/model/pytorch_model_module.pt"
|
||||
elif version == "1.1":
|
||||
download_file = "HunyuanDiT/t2i/model/pytorch_model_module.pt"
|
||||
else:
|
||||
raise Exception(f"未知的版本号: {version}")
|
||||
|
||||
download_fullpath = os.path.join(
|
||||
hunyuan_base_path, download_file)
|
||||
if not os.path.exists(download_fullpath):
|
||||
|
||||
Reference in New Issue
Block a user