From 1c3ed8e7bb71c1dd415b20f617cbfa72d911d585 Mon Sep 17 00:00:00 2001
From: ssit <ssit5568@gmail.com>
Date: Sun, 9 Jul 2023 23:30:08 -0400
Subject: [PATCH] Added node for more control

---
 nodes.py | 33 +++++++++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index af41467..fe1f2ee 100644
--- a/nodes.py
+++ b/nodes.py
@@ -15,7 +15,7 @@ def get_supported_restart_schedulers():
     return list(SCHEDULER_MAPPING.keys())
 
 
-class KRestartSampler:
+class KRestartSamplerSimple:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -31,7 +31,6 @@ class KRestartSampler:
                 "latent_image": ("LATENT", ),
                 "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                 "segments": ("STRING", {"default": "[3,2,0.06,0.30],[3,1,0.30,0.59]", "multiline": False}),
-                # "restart_scheduler": (get_supported_restart_schedulers(), {"default": "karras"}),
             }
         }
 
@@ -43,10 +42,40 @@ class KRestartSampler:
         return restart_sampling(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise, segments, scheduler)
 
 
+class KRestartSampler:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL", ),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
+                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
+                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
+                "positive": ("CONDITIONING", ),
+                "negative": ("CONDITIONING", ),
+                "latent_image": ("LATENT", ),
+                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "segments": ("STRING", {"default": "[3,2,0.06,0.30],[3,1,0.30,0.59]", "multiline": False}),
+                "restart_scheduler": (get_supported_restart_schedulers(), ),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "sample"
+    CATEGORY = "sampling"
+
+    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise, segments, restart_scheduler):
+        return restart_sampling(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise, segments, restart_scheduler)
+
+
 NODE_CLASS_MAPPINGS = {
+    "KRestartSamplerSimple": KRestartSamplerSimple,
     "KRestartSampler": KRestartSampler,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
+    "KRestartSamplerSimple": "KSampler With Restarts (Simple)",
     "KRestartSampler": "KSampler With Restarts",
 }