From 35597f31eb3a739894271bfca95fd9987df2e3c1 Mon Sep 17 00:00:00 2001 From: Chris Tefer Date: Sun, 8 Dec 2024 11:50:42 -0600 Subject: [PATCH] Added some comments. Changed the scheduler to be universally a string because it works with my pipeline a little better. --- cfe_nodes.py | 156 +++++++++++++++++++++++++++++++++------------------ 1 file changed, 101 insertions(+), 55 deletions(-) diff --git a/cfe_nodes.py b/cfe_nodes.py index 66f4d6d..15751ad 100644 --- a/cfe_nodes.py +++ b/cfe_nodes.py @@ -6,13 +6,69 @@ import comfy.samplers import node_helpers +# NOTES: +## +# Artifacts +## +# Scheduler: this is the named scheduler, a string value +# Sigmas: this is the value being fed to the sampler which comprises of the scheduler (name), +# steps, denoise, and model + +def _flux_sampler(model, noise, cfg, cond, sampler, sigmas, latent_image): + """ + Simple sampler function for reuse + """ + 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 + + noise_mask = None + if "noise_mask" in latent: + noise_mask = latent["noise_mask"] + + + # update the conditioner value + condition = node_helpers.conditioning_set_values(cond, {"guidance": cfg}) + guider = comfy.samplers.CFGGuider(model) + guider.inner_set_conds({"positive": condition}) + + + disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED + samples = guider.sample(noise.generate_noise(latent), latent_image, sampler, sigmas, + denoise_mask=noise_mask, disable_pbar=disable_pbar, seed=noise.seed) + + samples = samples.to(comfy.model_management.intermediate_device()) + latent["samples"] = samples + + return latent + def _find_lowest_64(value): - # We use shifting to multiply by 64 and then divide by 64 - # not sure I understand this, but there's an additional requirement to be around 24 points - # of the corrected value, so we add 23 to ensure any time we're >24 points we end up in the - # corrected pixel space for most resolutions + """ + We use shifting to multiply by 64 and then divide by 64 + not sure I understand this, but there's an additional requirement to be around 24 points + of the corrected value, so we add 23 to ensure any time we're >24 points we end up in the + corrected pixel space for most resolutions + """ return (int(value + 23) >> 6) << 6 +def _build_sigmas(model, steps, scheduler, denoise): + """ + Builds the sigmas from the steps and the denoise. + """ + sigmas = None + if denoise < 1.0: + if denoise <= 0.0: + sigmas = torch.FloatTensor([]) + total_steps = int(steps/denoise) + + if sigmas is None: + sigmas = comfy.samplers.calculate_sigmas(model.get_model_object("model_sampling"), scheduler, total_steps).cpu() + sigmas = sigmas[-(steps + 1):] + + return sigmas + class CFE_Aspect_Ratio: """ @@ -121,7 +177,7 @@ class CFE_Flux_In_Pipe: "vae": ("VAE", {"tooltip": "The latent image to denoise."}), "cond": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}), "sampler": ("SAMPLER", {"tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output."}), - "sigmas": ("SIGMAS", {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}), + "scheduler": ("STRING", {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}), }, } @@ -131,8 +187,8 @@ class CFE_Flux_In_Pipe: FUNCTION = "in_pipe" CATEGORY = "CFE" - def in_pipe(self, model, clip, vae, cond, sampler, sigmas): - pipe = (model, clip, vae, cond, sampler, sigmas) + def in_pipe(self, model, clip, vae, cond, sampler, scheduler): + pipe = (model, clip, vae, cond, sampler, scheduler) return (pipe, ) @@ -148,17 +204,44 @@ class CFE_Flux_Out_Pipe: }, } - RETURN_TYPES = ("FLUX_PIPE", "MODEL", "CLIP", "VAE", "CONDITIONING", "SAMPLER", "SIGMAS") - RETURN_NAMES = ("pipe", "model", "clip", "vae", "cond", "sampler", "sigmas" ) + RETURN_TYPES = ("FLUX_PIPE", "MODEL", "CLIP", "VAE", "CONDITIONING", "SAMPLER", "STRING") + RETURN_NAMES = ("pipe", "model", "clip", "vae", "cond", "sampler", "scheduler" ) FUNCTION = "out_pipe" CATEGORY = "CFE" def out_pipe(self, pipe): - model, clip, vae, cond, sampler, sigmas = pipe - return (pipe, model, clip, vae, cond, sampler, sigmas) + model, clip, vae, cond, sampler, scheduler = pipe + return (pipe, model, clip, vae, cond, sampler, scheduler) +class CFE_Scheduler: + def __init__(self) -> None: + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required" : { + "model" : ("MODEL", {"tooltip" : "The model used for denoising the input latent"}), + "scheduler": ("STRING", {"default": "euler", "tooltip" : "The name of the scheduler being used"}), + "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",) + RETURN_NAMES = ("scheduler",) + FUNCTION = "output" + CATEGORY = "CFE" + + + + def output(self, model, scheduler, steps, denoise): + sigmas = _build_sigmas(model, steps, scheduler, denoise) + return (sigmas, ) + class CFE_Sigma_Sampler_Strings: def __init__(self) -> None: pass @@ -167,7 +250,7 @@ class CFE_Sigma_Sampler_Strings: def INPUT_TYPES(s): return { "required" : { - "sampler_select" : (comfy.samplers.SAMPLER_NAMES, {"tooltip" : "The name of the sampler being used"}), + "sampler" : (comfy.samplers.SAMPLER_NAMES, {"tooltip" : "The name of the sampler being used"}), "scheduler": (comfy.samplers.SCHEDULER_NAMES, {"tooltip" : "The name of the scheduler being used"}), } } @@ -184,7 +267,6 @@ class CFE_Sigma_Sampler_Strings: return (comfy.samplers.sampler_object(sampler_select), scheduler, ) - class CFE_Sigma_Sampler: def __init__(self) -> None: pass @@ -209,49 +291,11 @@ class CFE_Sigma_Sampler: def execute(self, model, sampler_select, scheduler, steps, denoise): - total_steps = steps - sigmas = None - if denoise < 1.0: - if denoise <= 0.0: - sigmas = torch.FloatTensor([]) - total_steps = int(steps/denoise) - - if sigmas is None: - sigmas = comfy.samplers.calculate_sigmas(model.get_model_object("model_sampling"), scheduler, total_steps).cpu() - sigmas = sigmas[-(steps + 1):] - + sigmas = _build_sigmas(model, steps, scheduler, denoise) return (comfy.samplers.sampler_object(sampler_select), sigmas, ) -def _flux_sampler(model, noise, cfg, cond, 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 - noise_mask = None - if "noise_mask" in latent: - noise_mask = latent["noise_mask"] - - - disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED - # update the conditioner value - - condition = node_helpers.conditioning_set_values(cond, {"guidance": cfg}) - guider = comfy.samplers.CFGGuider(model) - guider.inner_set_conds({"positive": condition}) - - - samples = guider.sample(noise.generate_noise(latent), latent_image, sampler, sigmas, - denoise_mask=noise_mask, disable_pbar=disable_pbar, seed=noise.seed) - - samples = samples.to(comfy.model_management.intermediate_device()) - - out = latent.copy() - out["samples"] = samples - - return out class CFE_FLUX_Pipe_Sampler: @@ -266,6 +310,8 @@ class CFE_FLUX_Pipe_Sampler: "noise":("NOISE", {"tooltip": "Noise for generation"},), "latent_image": ("LATENT", {"tooltip": "The latent image."}), "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}), + "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The amount of denoising applied, lower values will maintain the structure of the initial image allowing for image to image sampling."}), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}), }, } @@ -274,9 +320,9 @@ class CFE_FLUX_Pipe_Sampler: FUNCTION = "execute" CATEGORY = "CFE" - def execute(self, pipe, noise, latent_image, cfg): - model, _, vae, cond, sampler, sigmas = pipe - + def execute(self, pipe, noise, latent_image, cfg, steps, denoise): + model, _, vae, cond, sampler, scheduler = pipe + sigmas = _build_sigmas(model, steps, scheduler, denoise) out = _flux_sampler(model, noise, cfg, cond, sampler, sigmas, latent_image) return (pipe, out, vae)