From e3b2ef9b35f8ad1216ed4496e16ccdc492a2ecff Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Sun, 30 Jun 2024 22:34:27 +0800
Subject: [PATCH] 'update'

---
 __init__.py                  |  2 +-
 mz_train_tools_core_HYDiT.py | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/__init__.py b/__init__.py
index 77e78ef..7397c2a 100644
--- a/__init__.py
+++ b/__init__.py
@@ -647,7 +647,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"
+                    "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/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py
index 10c5df3..2a3f4d4 100644
--- a/mz_train_tools_core_HYDiT.py
+++ b/mz_train_tools_core_HYDiT.py
@@ -872,6 +872,20 @@ def MZ_HYDiTSimpleT2I_call(args={}):
             # 'algorithm_type': 'dpmsolver++',
         }
     }
+    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,
+        }
+    }
 
     width = args.get("width")
     height = args.get("height")