From 2441d098b898a5106a75d3591070260a929e5e7f Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 25 Feb 2024 01:04:05 +0200
Subject: [PATCH] Update nodes.py

---
 nodes.py | 38 +++++++++++++++++++++++++++++++++++---
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/nodes.py b/nodes.py
index a24e1d0..8a35ce9 100644
--- a/nodes.py
+++ b/nodes.py
@@ -3659,7 +3659,37 @@ class ImageNormalize_Neg1_To_1:
         images = images * 2.0 - 1.0
 
         return (images,)    
-    
+
+class X0_passlatent(comfy.model_sampling.EPS):
+    def calculate_denoised(self, sigma, model_output, model_input):
+        return model_output
+    def calculate_input(self, sigma, noise):
+        return noise
+class SingleStepSampling:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              }}
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
+
+    CATEGORY = "KJNodes"
+
+    def patch(self, model):
+        m = model.clone()
+
+        sampling_base = comfy.model_sampling.ModelSamplingDiscrete
+        sampling_type = X0_passlatent
+
+        class ModelSamplingAdvanced(sampling_base, sampling_type):
+            pass
+
+        model_sampling = ModelSamplingAdvanced(model.model.model_config)
+
+        m.add_object_patch("model_sampling", model_sampling)
+        return (m, )
+
 NODE_CLASS_MAPPINGS = {
     "INTConstant": INTConstant,
     "FloatConstant": FloatConstant,
@@ -3727,7 +3757,8 @@ NODE_CLASS_MAPPINGS = {
     "ImageUpscaleWithModelBatched": ImageUpscaleWithModelBatched,
     "ScaleBatchPromptSchedule": ScaleBatchPromptSchedule,
     "EffnetEncode": EffnetEncode,
-    "ImageNormalize_Neg1_To_1": ImageNormalize_Neg1_To_1
+    "ImageNormalize_Neg1_To_1": ImageNormalize_Neg1_To_1,
+    "SingleStepSampling": SingleStepSampling,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "INTConstant": "INT Constant",
@@ -3795,5 +3826,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ImageUpscaleWithModelBatched": "ImageUpscaleWithModelBatched",
     "ScaleBatchPromptSchedule": "ScaleBatchPromptSchedule",
     "EffnetEncode": "EffnetEncode",
-    "ImageNormalize_Neg1_To_1": "ImageNormalize_Neg1_To_1"
+    "ImageNormalize_Neg1_To_1": "ImageNormalize_Neg1_To_1",
+    "SingleStepSampling": "SingleStepSampling",
 }
\ No newline at end of file