From f376f59e383099ded5c88726d210c3b682c12efa Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Wed, 27 Mar 2024 06:54:07 -0600
Subject: [PATCH] Add KRestartSamplerCustomNoise when possible

---
 README.md    |  7 ++++
 __init__.py  |  3 ++
 changelog.md |  4 ++-
 py/nodes.py  | 94 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 107 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 75e73f7..2dc7628 100644
--- a/README.md
+++ b/README.md
@@ -98,6 +98,13 @@ From a usage perspective, using positive alpha will tend to create a colorful ef
 
 Noise from the `SonarCustomNoise` node and `SonarPowerNoise` can be freely mixed.
 
+### `KRestartSamplerCustomNoise`
+
+If you have a recent enough version of [ComfyUI_restart_sampling](https://github.com/ssitu/ComfyUI_restart_sampling/)
+installed, you'll also get the `KRestartSamplerCustomNoise` node which is exactly the same as `KRestartSamplerCustom`
+except for adding an optional custom noise input.
+See the restart sampling repo for more information: https://github.com/ssitu/ComfyUI_restart_sampling
+
 ## Sonar Sampler Parameters
 
 Very abbreviated section. The init type can make a big difference. If you use `RANDOM` you can get away with setting `direction` to high values (like up to `2.25` or so) and absurdly low values (like `-30.0`). It's also possible to set `momentum` and `momentum_hist` to negative values, although whether it's a good idea...
diff --git a/__init__.py b/__init__.py
index 49c3410..2912708 100644
--- a/__init__.py
+++ b/__init__.py
@@ -15,4 +15,7 @@ NODE_CLASS_MAPPINGS = {
 
 NODE_DISPLAY_NAME_MAPPINGS = {}
 
+if hasattr(nodes, "KRestartSamplerCustomNoise"):
+    NODE_CLASS_MAPPINGS["KRestartSamplerCustomNoise"] = nodes.KRestartSamplerCustomNoise
+
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
diff --git a/changelog.md b/changelog.md
index 5e19a75..1f8708a 100644
--- a/changelog.md
+++ b/changelog.md
@@ -2,13 +2,15 @@
 
 Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top.
 
-## 20240325
+## 20240327
 
 * Fixed issue when using Sonar samplers in normal sampling nodes/via stuff like `KSamplerSelect`.
 * Add `pyramid` (non-high-res) noise type.
 * Allow selecting `brownian` noise in custom noise nodes (but it won't work with `NoisyLatentLike`).
 * Use `brownian` as the default noise type for `SamplerSonarDPMPP`.
 * Make overriding the selected noise type in Sonar samplers a warning instead of a hard error.
+* Improve noise scaling (may change seeds).
+* Add `KRestartSamplerCustomNoise` if the user has a recent enough version of ComfyUI_restart_sampling installed.
 
 ## 20240320
 
diff --git a/py/nodes.py b/py/nodes.py
index c76b0bc..a06cda8 100644
--- a/py/nodes.py
+++ b/py/nodes.py
@@ -583,3 +583,97 @@ class SamplerNodeConfigOverride:
             extra_args=extra_args,
             **kwargs,
         )
+
+
+try:
+    import custom_nodes.ComfyUI_restart_sampling as rs
+
+    if not hasattr(rs.restart_sampling, "DEFAULT_SEGMENTS"):
+        # Dumb test but this should only exist in restart sampling versions that
+        # support plugging in custom noise.
+        raise NotImplementedError  # noqa: TRY301
+
+    class KRestartSamplerCustomNoise:
+        @classmethod
+        def INPUT_TYPES(cls):
+            return {
+                "required": {
+                    "model": ("MODEL",),
+                    "add_noise": (["enable", "disable"],),
+                    "noise_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": ("SAMPLER",),
+                    "scheduler": (tuple(rs.restart_sampling.SCHEDULER_MAPPING.keys()),),
+                    "positive": ("CONDITIONING",),
+                    "negative": ("CONDITIONING",),
+                    "latent_image": ("LATENT",),
+                    "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
+                    "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
+                    "return_with_leftover_noise": (["disable", "enable"],),
+                    "segments": (
+                        "STRING",
+                        {
+                            "default": rs.restart_sampling.DEFAULT_SEGMENTS,
+                            "multiline": False,
+                        },
+                    ),
+                    "restart_scheduler": (rs.nodes.get_supported_restart_schedulers(),),
+                    "chunked_mode": ("BOOLEAN", {"default": True}),
+                },
+                "optional": {
+                    "custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
+                },
+            }
+
+        RETURN_TYPES = ("LATENT", "LATENT")
+        RETURN_NAMES = ("output", "denoised_output")
+        FUNCTION = "sample"
+        CATEGORY = "sampling"
+
+        def sample(
+            self,
+            model,
+            add_noise,
+            noise_seed,
+            steps,
+            cfg,
+            sampler,
+            scheduler,
+            positive,
+            negative,
+            latent_image,
+            start_at_step,
+            end_at_step,
+            return_with_leftover_noise,
+            segments,
+            restart_scheduler,
+            chunked_mode=False,
+            custom_noise_opt=None,
+        ):
+            return rs.restart_sampling.restart_sampling(
+                model,
+                noise_seed,
+                steps,
+                cfg,
+                sampler,
+                scheduler,
+                positive,
+                negative,
+                latent_image,
+                segments,
+                restart_scheduler,
+                disable_noise=add_noise == "disable",
+                step_range=(start_at_step, end_at_step),
+                force_full_denoise=return_with_leftover_noise != "enable",
+                output_only=False,
+                chunked_mode=chunked_mode,
+                custom_noise=custom_noise_opt.make_noise_sampler
+                if custom_noise_opt
+                else None,
+            )
+except (ImportError, NotImplementedError):
+    pass