From ac255212638a665f82cef609ccd4fe21c2ce3ec2 Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Wed, 3 Jul 2024 00:17:34 +0800
Subject: [PATCH] 'update'

---
 __init__.py                  |  2 +-
 hook_HYDiT_utils.py          | 41 ++++++++++++++++++++++++++++++------
 mz_train_tools_core_HYDiT.py |  2 +-
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/__init__.py b/__init__.py
index e7863f8..5986001 100644
--- a/__init__.py
+++ b/__init__.py
@@ -660,7 +660,7 @@ class MZ_HYDiTSimpleT2I:
                 "steps": ("INT", {"default": 20}),
                 "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
                 "scheduler": ([
-                    "ddpm", "ddim", "dpmms", "uni_pc", "DPMPP_2M_Karras"
+                    "ddpm", "ddim", "dpmms", "uni_pc", "dpmpp_2m_karras"
                 ], {"default": "ddpm"}),
                 "prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
                 "negative_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
diff --git a/hook_HYDiT_utils.py b/hook_HYDiT_utils.py
index 8d40308..84a70c5 100644
--- a/hook_HYDiT_utils.py
+++ b/hook_HYDiT_utils.py
@@ -100,9 +100,38 @@ def easy_sample_images(
         sample_images_dir = os.path.join(workspace_dir, "sample_images")
         os.makedirs(sample_images_dir, exist_ok=True)
 
+        sampler_factory = SAMPLER_FACTORY.copy()
+
+        sampler_factory["uni_pc"] = {
+            'scheduler': 'UniPCMultistepScheduler',
+            'name': 'UniPCMultistepScheduler',
+            'kwargs': {
+                'beta_schedule': 'scaled_linear',
+                'beta_start': 0.00085,
+                'beta_end': 0.03,
+                'prediction_type': 'v_prediction',
+                'trained_betas': None,
+                'solver_order': 2,
+            }
+        }
+        sampler_factory["dpmpp_2m_karras"] = {
+            'scheduler': 'DPMSolverMultistepScheduler',
+            'name': 'DPMSolverMultistepScheduler',
+            'kwargs': {
+                'beta_schedule': 'scaled_linear',
+                'beta_start': 0.00085,
+                'beta_end': 0.03,
+                'prediction_type': 'v_prediction',
+                'trained_betas': None,
+                'solver_order': 2,
+                'algorithm_type': 'dpmsolver++',
+                "use_karras_sigmas": True,
+            }
+        }
+
         # Load sampler from factory
-        kwargs = SAMPLER_FACTORY[sampler]['kwargs']
-        scheduler = SAMPLER_FACTORY[sampler]['scheduler']
+        kwargs = sampler_factory[sampler]['kwargs']
+        scheduler = sampler_factory[sampler]['scheduler']
 
         # Build scheduler according to the sampler.
         scheduler_class = getattr(schedulers, scheduler)
@@ -293,6 +322,7 @@ class CustomizeEmbedsModel(nn.Module):
     dtype = torch.float16
     # x = torch.zeros(1, 1, 256, 2048)
     x = None
+
     def __init__(self, *args, **kwargs):
         super().__init__()
 
@@ -300,13 +330,12 @@ class CustomizeEmbedsModel(nn.Module):
         self.dtype = torch.float16
         return self
 
-    def forward(self, *args, **kwargs): 
+    def forward(self, *args, **kwargs):
         input_ids = kwargs.get("input_ids", None)
         if self.x is None:
             batch_size = input_ids.shape[0]
             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"),
@@ -326,10 +355,10 @@ class CustomizeTokenizer(dict):
         self['input_ids'] = self.input_ids
         self['attention_mask'] = self.attention_mask
 
-    def tokenize(self, text):  
+    def tokenize(self, text):
         return text
 
-    def __call__(self, *args, **kwargs): 
+    def __call__(self, *args, **kwargs):
         return self
 
 
diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py
index 18e4390..d1c4623 100644
--- a/mz_train_tools_core_HYDiT.py
+++ b/mz_train_tools_core_HYDiT.py
@@ -887,7 +887,7 @@ def MZ_HYDiTSimpleT2I_call(args={}):
             # 'algorithm_type': 'dpmsolver++',
         }
     }
-    SAMPLER_FACTORY["DPMPP_2M_Karras"] = {
+    SAMPLER_FACTORY["dpmpp_2m_karras"] = {
         'scheduler': 'DPMSolverMultistepScheduler',
         'name': 'DPMSolverMultistepScheduler',
         'kwargs': {