'update'

__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):

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,)

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)