From 1adcc4c3a2f6f329c1e4e7ac3114f254f9b5f558 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 27 Sep 2023 22:21:18 -0400 Subject: [PATCH 1/5] Add a SamplerCustom Node. This node takes a list of sigmas and a sampler object as input. This lets people easily implement custom schedulers and samplers as nodes. More nodes will be added to it in the future. --- comfy/sample.py | 12 ++++ comfy_extras/nodes_custom_sampler.py | 98 ++++++++++++++++++++++++++++ nodes.py | 4 +- 3 files changed, 113 insertions(+), 1 deletion(-) create mode 100644 comfy_extras/nodes_custom_sampler.py diff --git a/comfy/sample.py b/comfy/sample.py index fe9f4118d..322272766 100644 --- a/comfy/sample.py +++ b/comfy/sample.py @@ -99,3 +99,15 @@ def sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative cleanup_additional_models(models) return samples + +def sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=None, callback=None, disable_pbar=False, seed=None): + real_model, positive_copy, negative_copy, noise_mask, models = prepare_sampling(model, noise.shape, positive, negative, noise_mask) + noise = noise.to(model.load_device) + latent_image = latent_image.to(model.load_device) + sigmas = sigmas.to(model.load_device) + + samples = comfy.samplers.sample(real_model, noise, positive_copy, negative_copy, cfg, model.load_device, sampler, sigmas, model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) + samples = samples.cpu() + cleanup_additional_models(models) + return samples + diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py new file mode 100644 index 000000000..062852629 --- /dev/null +++ b/comfy_extras/nodes_custom_sampler.py @@ -0,0 +1,98 @@ +import comfy.samplers +import comfy.sample +from comfy.k_diffusion import sampling as k_diffusion_sampling +import latent_preview + + +class KarrasScheduler: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), + "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), + "rho": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + } + } + RETURN_TYPES = ("SIGMAS",) + CATEGORY = "_for_testing/custom_sampling" + + FUNCTION = "get_sigmas" + + def get_sigmas(self, steps, sigma_max, sigma_min, rho): + sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho) + return (sigmas, ) + + +class KSamplerSelect: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"sampler_name": (comfy.samplers.KSAMPLER_NAMES, ), + } + } + RETURN_TYPES = ("SAMPLER",) + CATEGORY = "_for_testing/custom_sampling" + + FUNCTION = "get_sampler" + + def get_sampler(self, sampler_name): + sampler = comfy.samplers.ksampler(sampler_name)() + return (sampler, ) + +class SamplerCustom: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"model": ("MODEL",), + "add_noise": (["enable", "disable"], ), + "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}), + "positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "sampler": ("SAMPLER", ), + "sigmas": ("SIGMAS", ), + "latent_image": ("LATENT", ), + } + } + + RETURN_TYPES = ("LATENT","LATENT") + RETURN_NAMES = ("output", "denoised_output") + + FUNCTION = "sample" + + CATEGORY = "_for_testing/custom_sampling" + + def sample(self, model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image): + latent = latent_image + latent_image = latent["samples"] + if add_noise == "disable": + noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") + else: + batch_inds = latent["batch_index"] if "batch_index" in latent else None + noise = comfy.sample.prepare_noise(latent_image, noise_seed, batch_inds) + + noise_mask = None + if "noise_mask" in latent: + noise_mask = latent["noise_mask"] + + x0_output = {} + callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output) + + disable_pbar = False + samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed) + + out = latent.copy() + out["samples"] = samples + if "x0" in x0_output: + out_denoised = latent.copy() + out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu()) + else: + out_denoised = out + return (out, out_denoised) + +NODE_CLASS_MAPPINGS = { + "SamplerCustom": SamplerCustom, + "KarrasScheduler": KarrasScheduler, + "KSamplerSelect": KSamplerSelect, +} diff --git a/nodes.py b/nodes.py index a847db6fb..1232373be 100644 --- a/nodes.py +++ b/nodes.py @@ -1202,9 +1202,10 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, noise_mask = latent["noise_mask"] callback = latent_preview.prepare_callback(model, steps) + disable_pbar = False samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step, - force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, seed=seed) + force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) out = latent.copy() out["samples"] = samples return (out, ) @@ -1791,4 +1792,5 @@ def init_custom_nodes(): load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_clip_sdxl.py")) load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_canny.py")) load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_freelunch.py")) + load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_custom_sampler.py")) load_custom_nodes() From 1d7dfc07d5e76968c9137c17fca0f7ad77a7b9d8 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 27 Sep 2023 22:32:42 -0400 Subject: [PATCH 2/5] Make add_noise in SamplerCustom a boolean. --- comfy_extras/nodes_custom_sampler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 062852629..842a9de4f 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -2,7 +2,7 @@ import comfy.samplers import comfy.sample from comfy.k_diffusion import sampling as k_diffusion_sampling import latent_preview - +import torch class KarrasScheduler: @classmethod @@ -45,7 +45,7 @@ class SamplerCustom: def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), - "add_noise": (["enable", "disable"], ), + "add_noise": ("BOOLEAN", {"default": True}), "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}), "positive": ("CONDITIONING", ), @@ -66,7 +66,7 @@ class SamplerCustom: def sample(self, model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image): latent = latent_image latent_image = latent["samples"] - if add_noise == "disable": + if not add_noise: noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") else: batch_inds = latent["batch_index"] if "batch_index" in latent else None From d234ca558a7777b607a4f81aeb9e8703ef020977 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 28 Sep 2023 00:17:03 -0400 Subject: [PATCH 3/5] Add missing samplers to KSamplerSelect. --- comfy/samplers.py | 20 ++++++++++++-------- comfy_extras/nodes_custom_sampler.py | 4 ++-- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index 7668d7913..a7c240f40 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -711,6 +711,17 @@ def calculate_sigmas_scheduler(model, scheduler_name, steps): print("error invalid scheduler", self.scheduler) return sigmas +def sampler_class(name): + if name == "uni_pc": + sampler = UNIPC + elif name == "uni_pc_bh2": + sampler = UNIPCBH2 + elif name == "ddim": + sampler = DDIM + else: + sampler = ksampler(name) + return sampler + class KSampler: SCHEDULERS = SCHEDULER_NAMES SAMPLERS = SAMPLER_NAMES @@ -769,13 +780,6 @@ class KSampler: else: return torch.zeros_like(noise) - if self.sampler == "uni_pc": - sampler = UNIPC - elif self.sampler == "uni_pc_bh2": - sampler = UNIPCBH2 - elif self.sampler == "ddim": - sampler = DDIM - else: - sampler = ksampler(self.sampler) + sampler = sampler_class(self.sampler) return sample(self.model, noise, positive, negative, cfg, self.device, sampler(), sigmas, self.model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 842a9de4f..1c587dbd8 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -28,7 +28,7 @@ class KSamplerSelect: @classmethod def INPUT_TYPES(s): return {"required": - {"sampler_name": (comfy.samplers.KSAMPLER_NAMES, ), + {"sampler_name": (comfy.samplers.SAMPLER_NAMES, ), } } RETURN_TYPES = ("SAMPLER",) @@ -37,7 +37,7 @@ class KSamplerSelect: FUNCTION = "get_sampler" def get_sampler(self, sampler_name): - sampler = comfy.samplers.ksampler(sampler_name)() + sampler = comfy.samplers.sampler_class(sampler_name)() return (sampler, ) class SamplerCustom: From 2bf051fda87cfa94e5c99bbd88fc7f1434e9e1a2 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 28 Sep 2023 00:30:45 -0400 Subject: [PATCH 4/5] Add a basic node to generate sigmas from scheduler. --- comfy_extras/nodes_custom_sampler.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 1c587dbd8..aafde8f32 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -4,6 +4,26 @@ from comfy.k_diffusion import sampling as k_diffusion_sampling import latent_preview import torch + +class BasicScheduler: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"model": ("MODEL",), + "scheduler": (comfy.samplers.SCHEDULER_NAMES, ), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + } + } + RETURN_TYPES = ("SIGMAS",) + CATEGORY = "_for_testing/custom_sampling" + + FUNCTION = "get_sigmas" + + def get_sigmas(self, model, scheduler, steps): + sigmas = comfy.samplers.calculate_sigmas_scheduler(model.model, scheduler, steps).cpu() + return (sigmas, ) + + class KarrasScheduler: @classmethod def INPUT_TYPES(s): @@ -95,4 +115,5 @@ NODE_CLASS_MAPPINGS = { "SamplerCustom": SamplerCustom, "KarrasScheduler": KarrasScheduler, "KSamplerSelect": KSamplerSelect, + "BasicScheduler": BasicScheduler, } From 76e0f8fc8fe330b9568fab4b4a8049a62d141165 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 28 Sep 2023 00:40:09 -0400 Subject: [PATCH 5/5] Add function to split sigmas. --- comfy_extras/nodes_custom_sampler.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index aafde8f32..efe03ad24 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -43,6 +43,23 @@ class KarrasScheduler: sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho) return (sigmas, ) +class SplitSigmas: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"sigmas": ("SIGMAS", ), + "step": ("INT", {"default": 0, "min": 0, "max": 10000}), + } + } + RETURN_TYPES = ("SIGMAS","SIGMAS") + CATEGORY = "_for_testing/custom_sampling" + + FUNCTION = "get_sigmas" + + def get_sigmas(self, sigmas, step): + sigmas1 = sigmas[:step + 1] + sigmas2 = sigmas[step + 1:] + return (sigmas1, sigmas2) class KSamplerSelect: @classmethod @@ -116,4 +133,5 @@ NODE_CLASS_MAPPINGS = { "KarrasScheduler": KarrasScheduler, "KSamplerSelect": KSamplerSelect, "BasicScheduler": BasicScheduler, + "SplitSigmas": SplitSigmas, }