Add KRestartSamplerCustomNoise when possible

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...

__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"]

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
 

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