diff --git a/__init__.py b/__init__.py index 56bd805..852d390 100644 --- a/__init__.py +++ b/__init__.py @@ -560,22 +560,30 @@ class MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I: 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"}), + "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": "ddpm"}), + ], {"default": "dpmpp_2m_sde"}), "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}), + "width": ("INT", {"default": 1024, "max": 8192, "step": 16}), + "height": ("INT", {"default": 1024, "max": 8192, "step": 16}), "keep_device": (["enable", "disable"], {"default": "enable"}), }, "optional": { @@ -590,9 +598,15 @@ class MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I: 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): diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index c986018..89b9463 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -22,10 +22,8 @@ git_accelerate_urls = { "kkgithub": "kkgithub.com", } -# 初始化工具仓库和工作区 - -def MZ_KohyaSSInitWorkspace_call(args={}): +def MZ_KohyaSSCloneRepo_call(args={}): mz_dir = Utils.get_minus_zone_models_path() branch_repoid = args.get("branch_repoid", "kohya-ss/sd-scripts") @@ -72,6 +70,12 @@ def MZ_KohyaSSInitWorkspace_call(args={}): except Exception as e: raise Exception(f"克隆kohya-ss/sd-scripts或者切换分支时出现异常,详细信息请查看控制台...") +# 初始化工具仓库和工作区 + + +def MZ_KohyaSSInitWorkspace_call(args={}): + MZ_KohyaSSCloneRepo_call(args) + workspace_name = args.get("lora_name", None) workspace_name = workspace_name.strip() @@ -592,26 +596,272 @@ def MZ_KohyaSSTrain_call(args={}): def MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I_call(args={}): - # "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": "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"}), - 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 - from .hook_kohya_ss_hunyuan_pipe import run_kohya_ss_hunyuan_pipe + 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") + # src hw, dst hw, 0, 0 + 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") + pil_img = preview.decode_latent_to_preview_image( + None, + latents, + )[1] + 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,) diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index 04c099d..da336c8 100644 --- a/mz_train_tools_utils.py +++ b/mz_train_tools_utils.py @@ -169,6 +169,31 @@ class Utils: CACHE_POOL[key] = value return True + def model_cache_clean(model_type): + global CACHE_POOL + for key in list(CACHE_POOL.keys()): + if key.startswith(f"model_cache_{model_type}"): + del CACHE_POOL[key] + torch.cuda.empty_cache() + return True + + def model_cache_get(model_type, model_path): + resp = Utils.cache_get(f"model_cache_{model_type}") + if resp is None: + return None + cache_model_path = resp.get("model_path") + if cache_model_path != model_path: + Utils.model_cache_clean(f"model_cache_{model_type}") + return None + return resp.get("model") + + def model_cache_set(model_type, model_path, model): + key = f"model_cache_{model_type}" + if Utils.cache_get(key) is not None: + Utils.model_cache_clean(key) + Utils.cache_set(key, {"model_path": model_path, "model": model}) + return True + def get_minus_zone_models_path(): models_path = os.path.join( folder_paths.models_dir, "minus_zone_models") @@ -930,7 +955,7 @@ class Utils: def get_models_by_folder(dir_path): models = [] for root, dirs, files in os.walk(dir_path): - # 排除隐藏文件夹 + # 排除隐藏文件夹 dirs[:] = [d for d in dirs if not d.startswith('.')] for file in files: if file.endswith(".pth") or file.endswith(".pt") or file.endswith(".pkl") or file.endswith(".onnx") or file.endswith(".safetensors"): @@ -940,9 +965,172 @@ class Utils: def get_folders_by_folder(dir_path): folders = [] for root, dirs, files in os.walk(dir_path): - - # 排除隐藏文件夹 + + # 排除隐藏文件夹 dirs[:] = [d for d in dirs if not d.startswith('.')] for dir in dirs: folders.append(os.path.join(root, dir)) return folders + + +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)