From ff38b93af415ba4c38729d3b9ea4f46969e21486 Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Sat, 6 Jul 2024 00:33:34 +0800
Subject: [PATCH] 'update'

---
 __init__.py                                | 44 ++++++++++++++++++++++
 configs/kohya_ss_lora/lora_hunyuan1_1.json |  1 +
 configs/kohya_ss_lora/lora_hunyuan1_2.json |  4 +-
 hook_kohya_ss_hunyuan_pipe.py              |  2 +-
 mz_train_tools_core.py                     | 26 +++++++++++++
 5 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/__init__.py b/__init__.py
index c8a67b9..56bd805 100644
--- a/__init__.py
+++ b/__init__.py
@@ -549,6 +549,50 @@ NODE_DISPLAY_NAME_MAPPINGS[
     "MZ_KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace"] = f"{AUTHOR_NAME} - KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace"
 
 
+class MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I:
+    @classmethod
+    def INPUT_TYPES(s):
+        models, folders, vae_models, unet_models, _ = HYDiT_paths()
+        comfyui_full_loras = []
+        comfyui_loras = folder_paths.get_filename_list("loras")
+        for lora in comfyui_loras:
+            lora_path = folder_paths.get_full_path("loras", lora)
+            comfyui_full_loras.append(lora_path)
+        return {
+            "required": {
+                "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"}),
+            },
+            "optional": {
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("images",)
+
+    FUNCTION = "start"
+    CATEGORY = CATEGORY_NAME + "/kohya_ss_HYDiT_kohakublueleaf"
+
+    def start(self, **kwargs):
+        importlib.reload(mz_train_tools_core)
+        return mz_train_tools_core.MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I_call(kwargs)
+
+
 class MZ_LoadImagesFromDirectoryPath:
     @classmethod
     def INPUT_TYPES(s):
diff --git a/configs/kohya_ss_lora/lora_hunyuan1_1.json b/configs/kohya_ss_lora/lora_hunyuan1_1.json
index b7ca453..036ce8e 100644
--- a/configs/kohya_ss_lora/lora_hunyuan1_1.json
+++ b/configs/kohya_ss_lora/lora_hunyuan1_1.json
@@ -29,6 +29,7 @@
     "output_dir": "",
     "output_name": "",
     "no_half_vae": true,
+    "v_parameterization": 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
index 5832a26..8a0bfac 100644
--- a/configs/kohya_ss_lora/lora_hunyuan1_2.json
+++ b/configs/kohya_ss_lora/lora_hunyuan1_2.json
@@ -29,7 +29,9 @@
     "output_dir": "",
     "output_name": "",
     "no_half_vae": true,
+    "v_parameterization": true,
     "lowram": false,
-    "max_token_length": "150"
+    "max_token_length": "150",
+    "use_extra_cond": false
   }
 }
diff --git a/hook_kohya_ss_hunyuan_pipe.py b/hook_kohya_ss_hunyuan_pipe.py
index 1614a25..8d5868e 100644
--- a/hook_kohya_ss_hunyuan_pipe.py
+++ b/hook_kohya_ss_hunyuan_pipe.py
@@ -258,7 +258,7 @@ class HuanYuanDiffusionLongPromptWeightingPipeline:
                 x1 * sigmas[0],
                 sigmas,
             )
-            # torch.cuda.empty_cache()
+            torch.cuda.empty_cache()
         latents = sample
         return latents
 
diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py
index c5cc94a..c986018 100644
--- a/mz_train_tools_core.py
+++ b/mz_train_tools_core.py
@@ -589,3 +589,29 @@ def MZ_KohyaSSTrain_call(args={}):
     return (
         "训练完成",
     )
+
+
+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()
+    from .mz_train_tools_core_HYDiT import check_model_auto_download
+    args = check_model_auto_download(args)
+
+    from .hook_kohya_ss_hunyuan_pipe import run_kohya_ss_hunyuan_pipe