Add files via upload

2024-12-25 06:25:33 +01:00
parent d41ada5d30
commit c3402dc59b
1 changed files with 92 additions and 104 deletions
@@ -1,104 +1,92 @@
-import comfy.samplers
-import comfy.sample
-from comfy.k_diffusion import sampling as k_diffusion_sampling
-import latent_preview
-import torch
-import comfy.utils
-import node_helpers
-from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
-
-class Noise_EmptyNoise:
-    def __init__(self):
-        self.seed = 0
-
-    def generate_noise(self, input_latent):
-        latent_image = input_latent["samples"]
-        return torch.zeros(latent_image.shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
-
-class Noise_RandomNoise:
-    def __init__(self, seed):
-        self.seed = seed
-
-    def generate_noise(self, input_latent):
-        latent_image = input_latent["samples"]
-        batch_inds = input_latent["batch_index"] if "batch_index" in input_latent else None
-        return comfy.sample.prepare_noise(latent_image, self.seed, batch_inds)
-
-class DisableNoise:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {}}
-
-    RETURN_TYPES = ("NOISE",)
-    FUNCTION = "get_noise"
-    CATEGORY = "sampling/custom_sampling/noise"
-
-    def get_noise(self):
-        return (Noise_EmptyNoise(),)
-
-class FluxSettingsNode(ComfyNodeABC, DisableNoise):
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "model": ("MODEL", {"pos": (0, 50)}),  # Adjusted position for model
-                "conditioning": ("CONDITIONING", {"pos": (200, 50)}),  # Adjusted position for conditioning
-                "guidance": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step": 0.1}),
-                "sampler_name": (comfy.samplers.SAMPLER_NAMES, ),
-                "scheduler": (comfy.samplers.SCHEDULER_NAMES, ),
-                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-            }
-        }
-
-    RETURN_TYPES = ("CONDITIONING", "SAMPLER", "SIGMAS", "NOISE")
-
-    CATEGORY = "sampling/custom_sampling"
-
-    FUNCTION = "execute"
-
-    def apply_guidance(self, conditioning, guidance):
-        c = node_helpers.conditioning_set_values(conditioning, {"guidance": guidance})
-        return (c, )
-
-    def get_sampler(self, sampler_name):
-        sampler = comfy.samplers.sampler_object(sampler_name)
-        return (sampler, )
-
-    def get_sigmas(self, model, scheduler, steps, denoise):
-        total_steps = steps
-        if denoise < 1.0:
-            if denoise <= 0.0:
-                return (torch.FloatTensor([]),)
-            total_steps = int(steps / denoise)
-
-        sigmas = comfy.samplers.calculate_sigmas(
-            model.get_model_object("model_sampling"),
-            scheduler,
-            total_steps
-        ).cpu()
-
-        sigmas = sigmas[-(steps + 1):]
-        return (sigmas, )
-
-    def get_noise(self, noise_seed):
-        return (Noise_RandomNoise(noise_seed),)
-
-    def execute(self, model, conditioning, guidance, sampler_name, scheduler, steps, denoise, noise_seed):
-        c, = self.apply_guidance(conditioning, guidance)
-        sampler, = self.get_sampler(sampler_name)
-        sigmas, = self.get_sigmas(model, scheduler, steps, denoise)
-        noise, = self.get_noise(noise_seed)
-
-        return (c, sampler, sigmas, noise)
-
-NODE_CLASS_MAPPINGS = {
-    "FluxSettingsNode": FluxSettingsNode,
-    "DisableNoise": DisableNoise
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "FluxSettingsNode": "Flux Settings Node",
-    "DisableNoise": "Disable Noise",
-}
+import comfy.samplers
+import comfy.sample
+import torch
+import comfy.utils
+import node_helpers
+from comfy.comfy_types import ComfyNodeABC
+
+class Noise_EmptyNoise:
+    def __init__(self):
+        self.seed = 0
+
+    def generate_noise(self, input_latent):
+        latent_image = input_latent["samples"]
+        return torch.zeros(latent_image.shape, dtype=latent_image.dtype, layout=latent_image.layout, device=latent_image.device)
+
+class Noise_RandomNoise:
+    def __init__(self, seed):
+        self.seed = seed
+
+    def generate_noise(self, input_latent):
+        latent_image = input_latent["samples"]
+        batch_inds = input_latent.get("batch_index", None)
+        return comfy.sample.prepare_noise(latent_image, self.seed, batch_inds)
+
+class DisableNoise:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {}}
+
+    RETURN_TYPES = ("NOISE",)
+    FUNCTION = "get_noise"
+    CATEGORY = "sampling/custom_sampling/noise"
+
+    def get_noise(self):
+        return (Noise_EmptyNoise(),)
+
+class FluxSettingsNode(ComfyNodeABC):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL", {"pos": (0, 50)}),
+                "conditioning": ("CONDITIONING", {"pos": (200, 50)}),
+                "guidance": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step": 0.1}),
+                "sampler_name": (comfy.samplers.SAMPLER_NAMES, {"help": "Choose a sampling method"}),
+                "scheduler": (comfy.samplers.SCHEDULER_NAMES, {"help": "Choose a scheduler"}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING", "SAMPLER", "SIGMAS", "NOISE")
+    CATEGORY = "sampling/custom_sampling"
+    FUNCTION = "execute"
+
+    def apply_guidance(self, conditioning, guidance):
+        return (node_helpers.conditioning_set_values(conditioning, {"guidance": guidance}),)
+
+    def get_sampler(self, sampler_name):
+        if sampler_name not in comfy.samplers.SAMPLER_NAMES:
+            raise ValueError(f"Invalid sampler name: {sampler_name}")
+        return comfy.samplers.sampler_object(sampler_name),
+
+    def get_sigmas(self, model, scheduler, steps, denoise):
+        total_steps = int(steps / denoise) if denoise < 1.0 else steps
+        sigmas = comfy.samplers.calculate_sigmas(
+            model.get_model_object("model_sampling"),
+            scheduler,
+            total_steps
+        ).cpu()[-(steps + 1):]
+        return (sigmas,)
+
+    def get_noise(self, noise_seed):
+        return (Noise_RandomNoise(noise_seed),)
+
+    def execute(self, model, conditioning, guidance, sampler_name, scheduler, steps, denoise, noise_seed):
+        c = self.apply_guidance(conditioning, guidance)[0]
+        sampler = self.get_sampler(sampler_name)[0]
+        sigmas = self.get_sigmas(model, scheduler, steps, denoise)[0]
+        noise = self.get_noise(noise_seed)[0]
+        return c, sampler, sigmas, noise
+
+NODE_CLASS_MAPPINGS = {
+    "FluxSettingsNode": FluxSettingsNode,
+    "DisableNoise": DisableNoise
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "FluxSettingsNode": "Flux Settings Node",
+    "DisableNoise": "Disable Noise",
+}