diff --git a/__init__.py b/__init__.py index 502d71d..0bd942f 100644 --- a/__init__.py +++ b/__init__.py @@ -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": { diff --git a/configs/kohya_ss_lora/lora_hunyuan1_1.json b/configs/kohya_ss_lora/lora_hunyuan1_1.json new file mode 100644 index 0000000..188e079 --- /dev/null +++ b/configs/kohya_ss_lora/lora_hunyuan1_1.json @@ -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 + } +} diff --git a/configs/kohya_ss_lora/lora_hunyuan1_2.json b/configs/kohya_ss_lora/lora_hunyuan1_2.json new file mode 100644 index 0000000..5832a26 --- /dev/null +++ b/configs/kohya_ss_lora/lora_hunyuan1_2.json @@ -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" + } +} diff --git a/configs/kohya_ss_lora/lora_sdxl.json b/configs/kohya_ss_lora/lora_sdxl.json index 3218019..3c5634a 100644 --- a/configs/kohya_ss_lora/lora_sdxl.json +++ b/configs/kohya_ss_lora/lora_sdxl.json @@ -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", diff --git a/hook_kohya_ss_hunyuan_pipe.py b/hook_kohya_ss_hunyuan_pipe.py new file mode 100644 index 0000000..1614a25 --- /dev/null +++ b/hook_kohya_ss_hunyuan_pipe.py @@ -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 diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py index 73b821d..a564069 100644 --- a/hook_kohya_ss_run.py +++ b/hook_kohya_ss_run.py @@ -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, } diff --git a/hook_kohya_ss_utils.py b/hook_kohya_ss_utils.py index e1f6031..0f65a97 100644 --- a/hook_kohya_ss_utils.py +++ b/hook_kohya_ss_utils.py @@ -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 diff --git a/kohya_ss_lora b/kohya_ss_lora new file mode 160000 index 0000000..0dc79ed --- /dev/null +++ b/kohya_ss_lora @@ -0,0 +1 @@ +Subproject commit 0dc79edc01f2000de1dad5ad6d20d8b099bfafe2 diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index db1a3e0..c5cc94a 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -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}") diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py index 90359a8..f3427fe 100644 --- a/mz_train_tools_core_HYDiT.py +++ b/mz_train_tools_core_HYDiT.py @@ -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):