From fa43dbe8a7b9dd7c9949e778e1d2ec224f86792f Mon Sep 17 00:00:00 2001 From: Sn0w123 Date: Sat, 25 Oct 2025 23:05:18 +0200 Subject: [PATCH] Fix #15 --- src/simple_ksampler.py | 226 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 198 insertions(+), 28 deletions(-) diff --git a/src/simple_ksampler.py b/src/simple_ksampler.py index cef0cbf..90b7af0 100644 --- a/src/simple_ksampler.py +++ b/src/simple_ksampler.py @@ -1,18 +1,198 @@ -from nodes import VAEDecode, EmptyLatentImage, CLIPTextEncode -from comfy_extras.nodes_custom_sampler import ( - SamplerCustom, - BasicScheduler, - PolyexponentialScheduler, - VPScheduler, - SplitSigmasDenoise, -) +import torch +import latent_preview import comfy.samplers - +from nodes import VAEDecode, EmptyLatentImage, CLIPTextEncode +from comfy.k_diffusion import sampling as k_diffusion_sampling from comfy_extras.nodes_align_your_steps import AlignYourStepsScheduler from ..sn0w import Logger, Utility from .custom_schedulers.custom_schedulers import CustomSchedulers +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 _SplitSigmasDenoise: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "sigmas": ("SIGMAS",), + "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("SIGMAS", "SIGMAS") + RETURN_NAMES = ("high_sigmas", "low_sigmas") + CATEGORY = "sampling/custom_sampling/sigmas" + + FUNCTION = "get_sigmas" + + def get_sigmas(self, sigmas, denoise): + steps = max(sigmas.shape[-1] - 1, 0) + total_steps = round(steps * denoise) + sigmas1 = sigmas[:-(total_steps)] + sigmas2 = sigmas[-(total_steps + 1) :] + return (sigmas1, sigmas2) + + +class _VPScheduler: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "beta_d": ("FLOAT", {"default": 19.9, "min": 0.0, "max": 5000.0, "step": 0.01, "round": False}), # TODO: fix default values + "beta_min": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 5000.0, "step": 0.01, "round": False}), + "eps_s": ("FLOAT", {"default": 0.001, "min": 0.0, "max": 1.0, "step": 0.0001, "round": False}), + } + } + + RETURN_TYPES = ("SIGMAS",) + CATEGORY = "sampling/custom_sampling/schedulers" + + FUNCTION = "get_sigmas" + + def get_sigmas(self, steps, beta_d, beta_min, eps_s): + sigmas = k_diffusion_sampling.get_sigmas_vp(n=steps, beta_d=beta_d, beta_min=beta_min, eps_s=eps_s) + return (sigmas,) + + +class _PolyexponentialScheduler: + @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": 5000.0, "step": 0.01, "round": False}), + "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 5000.0, "step": 0.01, "round": False}), + "rho": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "round": False}), + } + } + + RETURN_TYPES = ("SIGMAS",) + CATEGORY = "sampling/custom_sampling/schedulers" + + FUNCTION = "get_sigmas" + + def get_sigmas(self, steps, sigma_max, sigma_min, rho): + sigmas = k_diffusion_sampling.get_sigmas_polyexponential(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho) + return (sigmas,) + + +class _BasicScheduler: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "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}), + } + } + + RETURN_TYPES = ("SIGMAS",) + CATEGORY = "sampling/custom_sampling/schedulers" + + FUNCTION = "get_sigmas" + + 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,) + + +class _SamplerCustom: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "add_noise": ("BOOLEAN", {"default": True}), + "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "control_after_generate": True}), + "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "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 = "sampling/custom_sampling" + + def sample(self, model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image): + latent = latent_image + latent_image = latent["samples"] + latent = latent.copy() + latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image) + latent["samples"] = latent_image + + if not add_noise: + noise = Noise_EmptyNoise().generate_noise(latent) + else: + noise = Noise_RandomNoise(noise_seed).generate_noise(latent) + + 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 = not comfy.utils.PROGRESS_BAR_ENABLED + 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) + + class SimpleSamplerCustom: logger = Logger() @@ -27,12 +207,8 @@ class SimpleSamplerCustom: def build_scheduler_list(cls): custom_scheduler_names = list(cls.custom_schedulers.get_scheduler_settings().keys()) cls.logger.log(custom_scheduler_names, "DEBUG") - cls.scheduler_list = ( - comfy.samplers.KSampler.SCHEDULERS + ["align_your_steps", "polyexponential", "vp"] + custom_scheduler_names - ) - cls.scheduler_settings = { - name: settings for name, settings in cls.custom_schedulers.get_scheduler_settings().items() - } + cls.scheduler_list = comfy.samplers.KSampler.SCHEDULERS + ["align_your_steps", "polyexponential", "vp"] + custom_scheduler_names + cls.scheduler_settings = {name: settings for name, settings in cls.custom_schedulers.get_scheduler_settings().items()} cls.custom_scheduler_defaults = cls.custom_schedulers.get_scheduler_defaults() @classmethod @@ -95,7 +271,7 @@ class SimpleSamplerCustom: height, **kwargs, ): - custom_sampler = SamplerCustom() + custom_sampler = _SamplerCustom() vae_decode = VAEDecode() text_encode = CLIPTextEncode() @@ -110,9 +286,7 @@ class SimpleSamplerCustom: sigmas = self.get_custom_sigmas(model, model_type, scheduler_name, steps, denoise) self.logger.print_sigmas_differences(scheduler_name, sigmas) - samples = custom_sampler.sample( - model, add_noise, noise_seed, cfg, positive_prompt, negative_prompt, sampler, sigmas, latent_image - ) + samples = custom_sampler.sample(model, add_noise, noise_seed, cfg, positive_prompt, negative_prompt, sampler, sigmas, latent_image) try: if image_output is not None and image_output["links"] != []: image = vae_decode.decode(vae, samples[1]) @@ -128,7 +302,7 @@ class SimpleSamplerCustom: if denoise == 1: return (None, sigmas) - return SplitSigmasDenoise.get_sigmas(self, sigmas, denoise) + return _SplitSigmasDenoise.get_sigmas(self, sigmas, denoise) def get_custom_sigmas(self, model, model_type, scheduler_name, steps, denoise): if denoise > 0 and scheduler_name not in comfy.samplers.KSampler.SCHEDULERS: @@ -139,18 +313,14 @@ class SimpleSamplerCustom: model_type = "SDXL" sigmas = AlignYourStepsScheduler.get_sigmas(self, model_type, steps, denoise)[0] elif scheduler_name == "polyexponential": - sigmas = self.get_denoised_sigmas( - PolyexponentialScheduler.get_sigmas(self, steps, 14.614642, 0.0291675, 1.0)[0], denoise - )[1] + sigmas = self.get_denoised_sigmas(_PolyexponentialScheduler.get_sigmas(self, steps, 14.614642, 0.0291675, 1.0)[0], denoise)[1] elif scheduler_name == "vp": - sigmas = self.get_denoised_sigmas(VPScheduler.get_sigmas(self, steps, 14.0, 0.05, 0.075)[0], denoise)[1] + sigmas = self.get_denoised_sigmas(_VPScheduler.get_sigmas(self, steps, 14.0, 0.05, 0.075)[0], denoise)[1] elif scheduler_name in comfy.samplers.KSampler.SCHEDULERS: - sigmas = BasicScheduler.get_sigmas(self, model, scheduler_name, steps, denoise)[0] + sigmas = _BasicScheduler.get_sigmas(self, model, scheduler_name, steps, denoise)[0] else: defaults = self.get_custom_scheduler_defaults(scheduler_name) - sigmas = self.get_denoised_sigmas( - self.get_custom_class_sigmas(scheduler_name, steps, *defaults.values())[0], denoise - )[1] + sigmas = self.get_denoised_sigmas(self.get_custom_class_sigmas(scheduler_name, steps, *defaults.values())[0], denoise)[1] return sigmas