import torch import torch.nn.functional as F from math import ceil, floor from copy import deepcopy import comfy.model_patcher from comfy.sampler_helpers import convert_cond from comfy.samplers import calc_cond_batch, encode_model_conds from comfy.ldm.modules.attention import optimized_attention_for_device from nodes import ConditioningConcat, ConditioningSetTimestepRange import comfy.model_management as model_management from comfy.latent_formats import SDXL as SDXL_Latent import os current_dir = os.path.dirname(os.path.realpath(__file__)) SDXL_Latent = SDXL_Latent() sdxl_latent_rgb_factors = SDXL_Latent.latent_rgb_factors ConditioningConcat = ConditioningConcat() ConditioningSetTimestepRange = ConditioningSetTimestepRange() default_attention = optimized_attention_for_device(model_management.get_torch_device()) default_device = model_management.get_torch_device() weighted_average = lambda tensor1, tensor2, weight1: (weight1 * tensor1 + (1 - weight1) * tensor2) selfnorm = lambda x: x / x.norm() minmaxnorm = lambda x: torch.nan_to_num((x - x.min()) / (x.max() - x.min()), nan=0.0, posinf=1.0, neginf=0.0) normlike = lambda x, y: x / x.norm() * y.norm() def get_sigma_min_max(model): model_sampling = model.model.model_sampling sigma_min = model_sampling.sigma(model_sampling.timestep(model_sampling.sigma_min)).item() sigma_max = model_sampling.sigma(model_sampling.timestep(model_sampling.sigma_max)).item() return sigma_min, sigma_max @torch.no_grad() def make_new_uncond_at_scale(cond,uncond,cond_scale,new_scale): new_scale_ratio = (new_scale - 1) / (cond_scale - 1) return cond * (1 - new_scale_ratio) + uncond * new_scale_ratio @torch.no_grad() def make_new_uncond_at_scale_co(conds_out,cond_scale,new_scale): new_scale_ratio = (new_scale - 1) / (cond_scale - 1) return conds_out[0] * (1 - new_scale_ratio) + conds_out[1] * new_scale_ratio @torch.no_grad() def get_denoised_at_scale(x_orig,cond,uncond,cond_scale): return x_orig - ((x_orig - uncond) + cond_scale * ((x_orig - cond) - (x_orig - uncond))) class pre_cfg_perp_neg: @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), "clip": ("CLIP",), "neg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 1/10, "round": 0.01}), "set_context_length" : ("BOOLEAN", {"default": False}), "context_length": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}), "start_at_sigma": ("FLOAT", {"default": 15, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), "end_at_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), # "cond_or_uncond": (["both","uncond"], {"default":"uncond"}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, clip, neg_scale, set_context_length, context_length, start_at_sigma, end_at_sigma, cond_or_uncond="uncond"): empty_cond, pooled = clip.encode_from_tokens(clip.tokenize(""), return_pooled=True) nocond = [[empty_cond, {"pooled_output": pooled}]] if context_length > 1 and set_context_length: short_nocond = deepcopy(nocond) for x in range(context_length - 1): (nocond,) = ConditioningConcat.concat(nocond, short_nocond) nocond = convert_cond(nocond) @torch.no_grad() def pre_cfg_perp_neg_function(args): conds_out = args["conds_out"] noise_pred_pos = conds_out[0] if args["sigma"][0] > start_at_sigma or args["sigma"][0] <= end_at_sigma or not torch.any(conds_out[1]): return conds_out noise_pred_neg = conds_out[1] model_options = args["model_options"] timestep = args["timestep"] model = args["model"] x = args["input"] nocond_processed = encode_model_conds(model.extra_conds, nocond, x, x.device, "negative") (noise_pred_nocond,) = calc_cond_batch(model, [nocond_processed], x, timestep, model_options) pos = noise_pred_pos - noise_pred_nocond neg = noise_pred_neg - noise_pred_nocond perp = neg - ((torch.mul(neg, pos).sum())/(torch.norm(pos)**2)) * pos perp_neg = perp * neg_scale if cond_or_uncond == "both": perp_p = pos - ((torch.mul(neg, pos).sum())/(torch.norm(neg)**2)) * neg perp_pos = perp_p * neg_scale conds_out[0] = noise_pred_nocond + perp_pos else: conds_out[0] = noise_pred_nocond + pos conds_out[1] = noise_pred_nocond + perp_neg return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(pre_cfg_perp_neg_function) return (m, ) class pre_cfg_re_negative: @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), "clip": ("CLIP",), "empty_proportion": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 1/20, "round": 0.01}), "progressive_scale" : ("BOOLEAN", {"default": False}), "set_context_length" : ("BOOLEAN", {"default": False}), "context_length": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}), "end_at_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, clip, empty_proportion, progressive_scale, set_context_length, context_length, end_at_sigma): sigma_min, sigma_max = get_sigma_min_max(model) empty_cond, pooled = clip.encode_from_tokens(clip.tokenize(""), return_pooled=True) nocond = [[empty_cond, {"pooled_output": pooled}]] if context_length > 1 and set_context_length: short_nocond = deepcopy(nocond) for x in range(context_length - 1): (nocond,) = ConditioningConcat.concat(nocond, short_nocond) nocond = convert_cond(nocond) @torch.no_grad() def pre_cfg_patch(args): conds_out = args["conds_out"] sigma = args["sigma"][0] # cond_scale = args["cond_scale"] if sigma <= end_at_sigma or not torch.any(conds_out[1]): return conds_out model_options = args["model_options"] timestep = args["timestep"] model = args["model"] x_orig = args["input"] nocond_processed = encode_model_conds(model.extra_conds, nocond, x_orig, x_orig.device, "negative") (noise_pred_nocond,) = calc_cond_batch(model, [nocond_processed], x_orig, timestep, model_options) if progressive_scale: progression = (sigma - sigma_min) / (sigma_max - sigma_min) current_scale = progression * empty_proportion + (1 - progression) * (1 - empty_proportion) current_scale = torch.clamp(current_scale, min=0, max=1) conds_out[1] = current_scale * noise_pred_nocond + conds_out[1] * (1 - current_scale) else: conds_out[1] = empty_proportion * noise_pred_nocond + conds_out[1] * (1 - empty_proportion) return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(pre_cfg_patch) return (m, ) @torch.no_grad() def normalize_adjust(a,b,strength=1): norm_a = torch.linalg.norm(a) a = selfnorm(a) b = selfnorm(b) res = b - a * (a * b).sum() if res.isnan().any(): res = torch.nan_to_num(res, nan=0.0) a = a - res * strength return a * norm_a class condDiffSharpeningNode: @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), "do_on": (["both","cond","uncond"], {"default": "both"},), "scale": ("FLOAT", {"default": 0.75, "min": -10.0, "max": 10.0, "step": 1/20, "round": 1/100}), "start_at_sigma": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 100.0, "step": 1/100, "round": 1/100}), "end_at_sigma": ("FLOAT", {"default": 01.0, "min": 0.0, "max": 100.0, "step": 1/100, "round": 1/100}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, do_on, scale, start_at_sigma, end_at_sigma): model_sampling = model.model.model_sampling sigma_max = model_sampling.sigma(model_sampling.timestep(model_sampling.sigma_max)).item() prev_cond = None prev_uncond = None @torch.no_grad() def sharpen_conds_pre_cfg(args): nonlocal prev_cond, prev_uncond conds_out = args["conds_out"] uncond = torch.any(conds_out[1]) sigma = args["sigma"][0].item() first_step = sigma > (sigma_max - 1) if first_step: prev_cond = None prev_uncond = None for b in range(len(conds_out[0])): for c in range(len(conds_out[0][b])): if not first_step and sigma > end_at_sigma and sigma <= start_at_sigma: if prev_cond is not None and do_on in ['both','cond']: conds_out[0][b][c] = normalize_adjust(conds_out[0][b][c], prev_cond[b][c], scale) if prev_uncond is not None and uncond and do_on in ['both','uncond']: conds_out[1][b][c] = normalize_adjust(conds_out[1][b][c], prev_uncond[b][c], scale) prev_cond = conds_out[0] if uncond: prev_uncond = conds_out[1] return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(sharpen_conds_pre_cfg) return (m, ) @torch.no_grad() def normalized_pow(t,p): t_norm = t.norm() t_sign = t.sign() t_pow = (t / t_norm).abs().pow(p) t_pow = selfnorm(t_pow) * t_norm * t_sign return t_pow class condExpNode: @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), "do_on": (["both","cond","uncond"], {"default": "both"},), "exponent": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 10.0, "step": 1/20, "round": 1/100}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, do_on, exponent): @torch.no_grad() def exponentiate_conds_pre_cfg(args): if args["sigma"][0] <= 1: return args["conds_out"] conds_out = args["conds_out"] uncond = torch.any(conds_out[1]) if do_on in ['both','uncond'] and not uncond: return conds_out for b in range(len(conds_out[0])): if do_on in ['both','cond']: conds_out[0][b] = normalized_pow(conds_out[0][b], exponent) if uncond and do_on in ['both','uncond']: conds_out[1][b] = normalized_pow(conds_out[1][b], exponent) return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(exponentiate_conds_pre_cfg) return (m, ) @torch.no_grad() def topk_average(latent, top_k=0.25, measure="average"): max_values = torch.topk(latent.flatten(), k=ceil(latent.numel()*top_k), largest=True ).values min_values = torch.topk(latent.flatten(), k=ceil(latent.numel()*top_k), largest=False).values value_range = measuring_methods[measure](max_values, min_values) return value_range apply_scaling_methods = { "individual": lambda c, m: c * torch.tensor(m).view(c.shape[0],1,1).to(c.device), "all_as_one": lambda c, m: c * m[0], "average_of_all_channels" : lambda c, m: c * (sum(m) / len(m)), "smallest_of_all_channels": lambda c, m: c * min(m), "biggest_of_all_channels" : lambda c, m: c * max(m), } measuring_methods = { "difference": lambda x, y: (x.mean() - y.mean()).abs() / 2, "average": lambda x, y: (x.mean() + y.abs().mean()) / 2, "biggest": lambda x, y: max(x.mean(), y.abs().mean()), } class automatic_pre_cfg: @classmethod def INPUT_TYPES(s): scaling_methods_names = [k for k in apply_scaling_methods] measuring_methods_names = [k for k in measuring_methods] return {"required": { "model": ("MODEL",), "scaling_method": (scaling_methods_names, {"default": scaling_methods_names[0]}), "min_max_method": ([m for m in measuring_methods], {"default": measuring_methods_names[1]}), "reference_CFG": ("FLOAT", {"default": 8, "min": 0.0, "max": 100, "step": 1/10, "round": 1/100}), "scale_multiplier": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 100, "step": 1/100, "round": 1/100}), "top_k": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 0.5, "step": 1/20, "round": 1/100}), }, "optional": { "channels_selection": ("CHANS",), } } RETURN_TYPES = ("MODEL","STRING",) RETURN_NAMES = ("MODEL","parameters",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, scaling_method, min_max_method="difference", reference_CFG=8, scale_multiplier=0.8, top_k=0.25, channels_selection=None): parameters_string = f"scaling_method: {scaling_method}\nmin_max_method: {min_max_method}" if channels_selection is not None: for x in range(channels_selection): parameters_string += f"\nchannel {x+1}: {channels_selection[x]}" scaling_methods_names = [k for k in apply_scaling_methods] @torch.no_grad() def automatic_pre_cfg(args): conds_out = args["conds_out"] cond_scale = args["cond_scale"] uncond = torch.any(conds_out[1]) if reference_CFG == 0: reference_scale = cond_scale else: reference_scale = reference_CFG if not uncond: return conds_out if channels_selection is None: channels = [True for _ in range(conds_out[0].shape[-3])] else: channels = channels_selection for b in range(len(conds_out[0])): chans = [] if scaling_method == scaling_methods_names[1]: if all(channels): mes = topk_average(reference_scale * conds_out[0][b] - (reference_scale - 1) * conds_out[1][b], top_k=top_k, measure=min_max_method) else: cond_for_measure = torch.stack([conds_out[0][b][j] for j in range(len(channels)) if channels[j]]) uncond_for_measure = torch.stack([conds_out[1][b][j] for j in range(len(channels)) if channels[j]]) mes = topk_average(reference_scale * cond_for_measure - (reference_scale - 1) * uncond_for_measure, top_k=top_k, measure=min_max_method) chans.append(scale_multiplier / max(mes,0.01)) else: for c in range(len(conds_out[0][b])): if not channels[c]: if scaling_method == scaling_methods_names[0]: chans.append(1) continue mes = topk_average(reference_scale * conds_out[0][b][c] - (reference_scale - 1) * conds_out[1][b][c], top_k=top_k, measure=min_max_method) new_scale = scale_multiplier / max(mes,0.01) chans.append(new_scale) conds_out[0][b] = apply_scaling_methods[scaling_method](conds_out[0][b],chans) conds_out[1][b] = apply_scaling_methods[scaling_method](conds_out[1][b],chans) return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(automatic_pre_cfg) return (m, parameters_string,) class channel_selection_node: CHANNELS_AMOUNT = 4 @classmethod def INPUT_TYPES(s): toggles = {f"channel_{x}" : ("BOOLEAN", {"default": True}) for x in range(s.CHANNELS_AMOUNT)} return {"required": toggles} RETURN_TYPES = ("CHANS",) FUNCTION = "exec" CATEGORY = "model_patches/Pre CFG/channels_selectors" def exec(self, **kwargs): chans = [] for k, v in kwargs.items(): if "channel_" in k: chans.append(v) return (chans, ) class individual_channel_selection_node: @classmethod def INPUT_TYPES(s): return {"required": { "exclude" : ("BOOLEAN", {"default": False}), "selected_channel": ("INT", {"default": 1, "min": 1, "max": 128}), "total_channels" : ("INT", {"default": 4, "min": 1, "max": 128}), } } RETURN_TYPES = ("CHANS",) FUNCTION = "exec" CATEGORY = "model_patches/Pre CFG/channels_selectors" def exec(self, exclude, selected_channel, total_channels): chans = [exclude for _ in range(total_channels)] chans[selected_channel - 1] = not exclude return (chans, ) class channel_multiplier_node: @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), "channel_1": ("FLOAT", {"default": 1, "min": -10.0, "max": 10.0, "step": 1/100, "round": 1/100}), "channel_2": ("FLOAT", {"default": 1, "min": -10.0, "max": 10.0, "step": 1/100, "round": 1/100}), "channel_3": ("FLOAT", {"default": 1, "min": -10.0, "max": 10.0, "step": 1/100, "round": 1/100}), "channel_4": ("FLOAT", {"default": 1, "min": -10.0, "max": 10.0, "step": 1/100, "round": 1/100}), "selection": (["both","cond","uncond"],), "start_at_sigma": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 100.0, "step": 1/100, "round": 1/100}), "end_at_sigma": ("FLOAT", {"default": 01.0, "min": 0.0, "max": 100.0, "step": 1/100, "round": 1/100}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, channel_1, channel_2, channel_3, channel_4, selection, start_at_sigma, end_at_sigma): chans = [channel_1, channel_2, channel_3, channel_4] @torch.no_grad() def channel_multiplier_function(args): conds_out = args["conds_out"] uncond = torch.any(conds_out[1]) sigma = args["sigma"] if sigma[0] <= end_at_sigma or sigma[0] > start_at_sigma: return conds_out for b in range(len(conds_out[0])): for c in range(len(conds_out[0][b])): if selection in ["both","cond"]: conds_out[0][b][c] *= chans[c] if uncond and selection in ["both","uncond"]: conds_out[1][b][c] *= chans[c] return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(channel_multiplier_function) return (m, ) class support_empty_uncond_pre_cfg_node: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "method": (["from cond","divide by CFG"],), }} RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, method): @torch.no_grad() def support_empty_uncond(args): conds_out = args["conds_out"] uncond = torch.any(conds_out[1]) cond_scale = args["cond_scale"] if not uncond and cond_scale > 1: if method == "divide by CFG": conds_out[0] /= cond_scale else: conds_out[1] = conds_out[0].clone() return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(support_empty_uncond) return (m, ) def replace_timestep(cond): cond = deepcopy(cond) cond[0]['timestep_start'] = 999999999.9 cond[0]['timestep_end'] = 0.0 return cond def check_if_in_timerange(conds,timestep_in): for c in conds: all_good = True if 'timestep_start' in c: timestep_start = c['timestep_start'] if timestep_in[0] > timestep_start: all_good = False if 'timestep_end' in c: timestep_end = c['timestep_end'] if timestep_in[0] < timestep_end: all_good = False if all_good: return True return False class zero_attention_pre_cfg_node: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "do_on": (["cond","uncond"], {"default": "uncond"},), "mix_scale": ("FLOAT", {"default": 1.5, "min": -2.0, "max": 2.0, "step": 1/2, "round": 1/100}), "start_at_sigma": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), "end_at_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), # "attention": (["both","self","cross"],), # "unet_block": (["input","middle","output"],), # "unet_block_id": ("INT", {"default": 8, "min": 0, "max": 20}), }} RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, do_on, mix_scale, start_at_sigma, end_at_sigma, attention="both", unet_block="input", unet_block_id=8): cond_index = 1 if do_on == "uncond" else 0 attn = {"both":["attn1","attn2"],"self":["attn1"],"cross":["attn2"]}[attention] def zero_attention_function(q, k, v, extra_options, mask=None): return torch.zeros_like(q) @torch.no_grad() def zero_attention_pre_cfg_patch(args): conds_out = args["conds_out"] sigma = args["sigma"][0].item() if sigma > start_at_sigma or sigma <= end_at_sigma: return conds_out conds = args["conds"] cond_to_process = conds[cond_index] cond_generated = torch.any(conds_out[cond_index]) if not cond_generated: cond_to_process = replace_timestep(cond_to_process) elif mix_scale == 1: print(" Mix scale at one!\nPrediction not generated.\nUse the node ConditioningSetTimestepRange to avoid generating if you want to use this node.") return conds_out model_options = deepcopy(args["model_options"]) for att in attn: model_options = comfy.model_patcher.set_model_options_patch_replace(model_options, zero_attention_function, att, unet_block, unet_block_id) (noise_pred,) = calc_cond_batch(args['model'], [cond_to_process], args['input'], args['timestep'], model_options) if mix_scale == 1 or not cond_generated: conds_out[cond_index] = noise_pred elif cond_generated: conds_out[cond_index] = weighted_average(noise_pred,conds_out[cond_index],mix_scale) return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(zero_attention_pre_cfg_patch) return (m, ) class perturbed_attention_guidance_pre_cfg_node: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "scale": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 10.0, "step": 1/20, "round": 1/100}), "start_at_sigma": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), "end_at_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), }} RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, scale, start_at_sigma, end_at_sigma, do_on="cond", attention="self", unet_block="middle", unet_block_id=0): cond_index = 1 if do_on == "uncond" else 0 attn = {"both":["attn1","attn2"],"self":["attn1"],"cross":["attn2"]}[attention] def perturbed_attention_guidance(q, k, v, extra_options, mask=None): return v @torch.no_grad() def perturbed_attention_guidance_pre_cfg_patch(args): conds_out = args["conds_out"] sigma = args["sigma"][0].item() if sigma > start_at_sigma or sigma <= end_at_sigma: return conds_out conds = args["conds"] cond_to_process = conds[cond_index] cond_generated = torch.any(conds_out[cond_index]) if not cond_generated: return conds_out model_options = deepcopy(args["model_options"]) for att in attn: model_options = comfy.model_patcher.set_model_options_patch_replace(model_options, perturbed_attention_guidance, att, unet_block, unet_block_id) (noise_pred,) = calc_cond_batch(args['model'], [cond_to_process], args['input'], args['timestep'], model_options) conds_out[cond_index] = conds_out[cond_index] + (conds_out[cond_index] - noise_pred) * scale return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(perturbed_attention_guidance_pre_cfg_patch) return (m, ) def sigma_to_percent(model_sampling, sigma_value): if sigma_value >= 999999999.9: return 0.0 if sigma_value <= 0.0: return 1.0 sigma_tensor = torch.tensor([sigma_value], dtype=torch.float32) timestep = model_sampling.timestep(sigma_tensor) percent = 1.0 - (timestep.item() / 999.0) return percent class ConditioningSetTimestepRangeFromSigma: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "conditioning": ("CONDITIONING", ), "sigma_start" : ("FLOAT", {"default": 15.0, "min": 0.0, "max": 10000.0, "step": 0.01}), "sigma_end" : ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10000.0, "step": 0.01}) }} RETURN_TYPES = ("CONDITIONING",) FUNCTION = "set_range" CATEGORY = "advanced/conditioning" def set_range(self, model, conditioning, sigma_start, sigma_end): model_sampling = model.model.model_sampling (c, ) = ConditioningSetTimestepRange.set_range(conditioning,sigma_to_percent(model_sampling, sigma_start),sigma_to_percent(model_sampling, sigma_end)) return (c, ) class ShapeAttentionNode: @classmethod def INPUT_TYPES(s): return { "required": { "model": ("MODEL",), "scale": ("FLOAT", {"default": 1.5, "min": 0.0, "max": 10.0, "step": 1/10, "round": 1/100}), # "start_at_sigma": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), # "end_at_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), # "enabled" : ("BOOLEAN", {"default": True}), # "attention": (["both","self","cross"],), # "unet_block": (["input","middle","output"],), # "unet_block_id": ("INT", {"default": 8, "min": 0, "max": 20}), # uncomment these lines if you want to have fun with the other layers } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches" def patch(self, model, scale, start_at_sigma=999999999.9, end_at_sigma=0.0, enabled=True, attention="self", unet_block="input", unet_block_id=8): attn = {"both":["attn1","attn2"],"self":["attn1"],"cross":["attn2"]}[attention] if scale == 1: print(" Shape attention disabled (scale is one)") if not enabled or scale == 1: return (model,) m = model.clone() def shape_attention(q, k, v, extra_options, mask=None): sigma = extra_options['sigmas'][0] if sigma > start_at_sigma or sigma <= end_at_sigma: return default_attention(q, k, v, extra_options['n_heads'], mask) if scale != 0: return default_attention(q, k, v, extra_options['n_heads'], mask) * scale else: return torch.zeros_like(q) for att in attn: m.model_options = comfy.model_patcher.set_model_options_patch_replace(m.model_options, shape_attention, att, unet_block, unet_block_id) return (m,) class ExlAttentionNode: @classmethod def INPUT_TYPES(s): return { "required": { "model": ("MODEL",), "scale": ("FLOAT", {"default": 2, "min": -1.0, "max": 10.0, "step": 1/10, "round": 1/100}), "enabled": ("BOOLEAN", {"default": True}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches" def patch(self, model, scale, enabled): if not enabled: return (model,) m = model.clone() def cross_patch(q, k, v, extra_options, mask=None): first_attention = default_attention(q, k, v, extra_options['n_heads'], mask) second_attention = normlike(q+(q-default_attention(first_attention, k, v, extra_options['n_heads'])), first_attention) * scale return second_attention m.model_options = comfy.model_patcher.set_model_options_patch_replace(m.model_options, cross_patch, "attn2", "middle", 0) return (m,) class PreCFGsubtractMeanNode: @classmethod def INPUT_TYPES(s): return { "required": { "model": ("MODEL",), # "per_channel" : ("BOOLEAN", {"default": False}), #It's just not good "start_at_sigma": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), "end_at_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), "enabled" : ("BOOLEAN", {"default": True}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, start_at_sigma, end_at_sigma, enabled, per_channel=False): if not enabled: return (model,) m = model.clone() def pre_cfg_function(args): conds_out = args["conds_out"] sigma = args["sigma"][0].item() if sigma > start_at_sigma or sigma <= end_at_sigma: return conds_out for x in range(len(conds_out)): if torch.any(conds_out[x]): for b in range(len(conds_out[x])): if per_channel: for c in range(len(conds_out[x][b])): conds_out[x][b][c] -= conds_out[x][b][c].mean() else: conds_out[x][b] -= conds_out[x][b].mean() return conds_out m.set_model_sampler_pre_cfg_function(pre_cfg_function) return (m,) class PostCFGsubtractMeanNode: @classmethod def INPUT_TYPES(s): return { "required": { "model": ("MODEL",), # "per_channel" : ("BOOLEAN", {"default": False}), #It's just not good "start_at_sigma": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), "end_at_sigma": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), "enabled" : ("BOOLEAN", {"default": True}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches" def patch(self, model, start_at_sigma, end_at_sigma, enabled, per_channel=False): if not enabled: return (model,) m = model.clone() def post_cfg_function(args): cfg_result = args["denoised"] sigma = args["sigma"][0].item() if sigma > start_at_sigma or sigma <= end_at_sigma: return cfg_result for b in range(len(cfg_result)): if per_channel: for c in range(len(cfg_result[b])): cfg_result[b][c] -= cfg_result[b][c].mean() else: cfg_result[b] -= cfg_result[b].mean() return cfg_result m.set_model_sampler_post_cfg_function(post_cfg_function) return (m,) class PostCFGDotNode: @classmethod def INPUT_TYPES(s): return { "required": { "model": ("MODEL",), "batch": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), "channel": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), "coord_x": ("INT", {"default": 64, "min": 0, "max": 1000, "step": 1}), "coord_y": ("INT", {"default": 64, "min": 0, "max": 1000, "step": 1}), "value": ("FLOAT", {"default": 1, "min": -10.0, "max": 10.0, "step": 1/10, "round": 1/100}), "start_at_sigma": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), "end_at_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), "enabled" : ("BOOLEAN", {"default": True}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches" def patch(self, model, batch, channel, coord_x, coord_y, value, start_at_sigma, end_at_sigma, enabled): if not enabled: return (model,) m = model.clone() def post_cfg_function(args): cfg_result = args["denoised"] sigma = args["sigma"][0].item() if sigma > start_at_sigma or sigma <= end_at_sigma: return cfg_result channel_norm = cfg_result[batch][channel].norm() cfg_result[batch][channel] /= channel_norm cfg_result[batch][channel][coord_y][coord_x] = value cfg_result[batch][channel] *= channel_norm return cfg_result m.set_model_sampler_post_cfg_function(post_cfg_function) return (m,) class uncondZeroPreCFGNode: @classmethod def INPUT_TYPES(s): scaling_methods_names = [k for k in apply_scaling_methods] return {"required": { "model": ("MODEL",), "scale": ("FLOAT", {"default": 0.75, "min": 0.0, "max": 10.0, "step": 1/20, "round": 0.01}), "start_at_sigma": ("FLOAT", {"default": 100, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), "end_at_sigma": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), "scaling_method": (scaling_methods_names, {"default": scaling_methods_names[2]}), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, scale, start_at_sigma, end_at_sigma, scaling_method): scaling_methods_names = [k for k in apply_scaling_methods] @torch.no_grad() def uncond_zero_pre_cfg(args): conds_out = args["conds_out"] uncond = torch.any(conds_out[1]) sigma = args["sigma"][0].item() if uncond or sigma <= end_at_sigma or sigma > start_at_sigma: return conds_out for b in range(len(conds_out[0])): chans = [] if scaling_method == scaling_methods_names[1]: mes = topk_average(8 * conds_out[0][b] - 7 * conds_out[1][b], measure="difference") for c in range(len(conds_out[0][b])): mes = topk_average(conds_out[0][b][c], measure="difference") ** 0.5 chans.append(scale / mes) conds_out[0][b] = apply_scaling_methods[scaling_method](conds_out[0][b],chans) return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(uncond_zero_pre_cfg) return (m, ) class variable_scale_pre_cfg_node: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "target_scale": ("FLOAT", {"default": 5.0, "min": 1.0, "max": 100.0, "step": 1/2, "round": 1/100}), "target_as_start": ("BOOLEAN", {"default": True}), "proportional_to": (["sigma","steps progression"],), }} RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, target_scale, target_as_start, proportional_to): model_sampling = model.model.model_sampling sigma_max = model_sampling.sigma(model_sampling.timestep(model_sampling.sigma_max)).item() @torch.no_grad() def variable_scale_pre_cfg_patch(args): conds_out = args["conds_out"] cond_scale = args["cond_scale"] sigma = args["sigma"][0].item() scales = [cond_scale,target_scale] if not torch.any(conds_out[1]): return conds_out if proportional_to == "steps progression": progression = sigma_to_percent(model_sampling, sigma) else: progression = 1 - sigma / sigma_max progression = max(min(progression, 1), 0) current_scale = scales[target_as_start] * (1 - progression) + scales[not target_as_start] * progression new_scale = (current_scale - 1) / (cond_scale - 1) conds_out[1] = weighted_average(conds_out[1], conds_out[0], new_scale) return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(variable_scale_pre_cfg_patch) return (m, ) class latent_noise_subtract_mean_node: def __init__(self): pass @classmethod def INPUT_TYPES(s): return {"required": { "latent_input": ("LATENT", {"forceInput": True}), "enabled" : ("BOOLEAN", {"default": True}), }} FUNCTION = "exec" RETURN_TYPES = ("LATENT",) CATEGORY = "latent" def exec(self, latent_input, enabled): if not enabled: return (latent_input,) new_latents = deepcopy(latent_input) for x in range(len(new_latents['samples'])): new_latents['samples'][x] -= torch.mean(new_latents['samples'][x]) return (new_latents,) class flip_flip_conds_pre_cfg_node: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "enabled" : ("BOOLEAN", {"default": True}) }} RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, enabled): @torch.no_grad() def pre_cfg_patch(args): conds_out = args["conds_out"] uncond = torch.any(conds_out[1]) if not uncond or not enabled: return conds_out conds_out[0], conds_out[1] = conds_out[1], conds_out[0] return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(pre_cfg_patch) return (m, ) class norm_uncond_to_cond_pre_cfg_node: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "enabled" : ("BOOLEAN", {"default": True}) }} RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, enabled): @torch.no_grad() def pre_cfg_patch(args): conds_out = args["conds_out"] uncond = torch.any(conds_out[1]) if not uncond or not enabled: return conds_out conds_out[1] = conds_out[1] / conds_out[1].norm() * conds_out[0].norm() return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(pre_cfg_patch) return (m, ) class replace_uncond_channel_pre_cfg_node: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "channel": ("INT", {"default": 1, "min": 1, "max": 128, "step": 1}), "enabled" : ("BOOLEAN", {"default": True}) }} RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, channel, enabled): @torch.no_grad() def pre_cfg_patch(args): conds_out = args["conds_out"] uncond = torch.any(conds_out[1]) if not uncond or not enabled: return conds_out for b in range(len(conds_out[0])): if len(conds_out[1][b]) < channel: print(F" WRONG CHANNEL SELECTED. THE LATENT SPACE ONLY HAS {len(conds_out[1][b])} CHANNELS") else: conds_out[1][b][channel - 1] = conds_out[0][b][channel - 1] return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(pre_cfg_patch) return (m, ) class merge_uncond_channel_pre_cfg_node: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "channel": ("INT", {"default": 1, "min": 1, "max": 128, "step": 1}), "CFG_scale": ("FLOAT", {"default": 5, "min": 2.0, "max": 100.0, "step": 1/2, "round": 1/100}), "start_at_sigma": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 100.0, "step": 1/100, "round": 1/100}), "end_at_sigma": ("FLOAT", {"default": 01.0, "min": 0.0, "max": 100.0, "step": 1/100, "round": 1/100}), "enabled" : ("BOOLEAN", {"default": True}) }} RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, channel, CFG_scale, start_at_sigma, end_at_sigma, enabled): if not enabled: return model, @torch.no_grad() def pre_cfg_patch(args): conds_out = args["conds_out"] cond_scale = args["cond_scale"] sigma = args["sigma"][0].item() if not torch.any(conds_out[1]) or sigma <= end_at_sigma or sigma > start_at_sigma: return conds_out for b in range(len(conds_out[0])): if len(conds_out[1][b]) < channel: print(F" WRONG CHANNEL SELECTED. THE LATENT SPACE ONLY HAS {len(conds_out[1][b])} CHANNELS") else: new_scale = (CFG_scale - 1) / (cond_scale - 1) conds_out[1][b][channel - 1] = weighted_average(conds_out[1][b][channel - 1], conds_out[0][b][channel - 1], new_scale) return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(pre_cfg_patch) return (m, ) class multiply_cond_pre_cfg_node: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "selection": (["both","cond","uncond"],), "value": ("FLOAT", {"default": 0, "min": -100.0, "max": 100.0, "step": 1/100, "round": 1/100}), "enabled" : ("BOOLEAN", {"default": True}) }} RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def patch(self, model, selection, value, enabled): @torch.no_grad() def pre_cfg_patch(args): conds_out = args["conds_out"] uncond = torch.any(conds_out[1]) if (not uncond and selection in ["both","uncond"]) or not enabled: return conds_out if selection in ["both","cond"]: conds_out[0] = conds_out[0] * value if selection in ["both","uncond"]: conds_out[1] = conds_out[1] * value return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(pre_cfg_patch) return (m, ) def generate_gradient_mask(tensor, horizontal=False): dim = 3 if horizontal else 2 gradient = torch.linspace(0, 1, steps=tensor.size(dim), device=tensor.device) if horizontal: merging_gradient = gradient.repeat(tensor.size(0), tensor.size(1), tensor.size(2), 1) else: merging_gradient = gradient.unsqueeze(1).repeat(tensor.size(0), tensor.size(1), 1, tensor.size(3)) return merging_gradient class gradient_scaling_pre_cfg_node: @classmethod def INPUT_TYPES(s): return {"required": {"model": ("MODEL",), "maximum_scale": ("FLOAT", {"default": 80, "min": 0.0, "max": 1000.0, "step": 1, "round": 1/100, "tooltip":"It is an equivalent to the CFG scale."}), "minimum_scale": ("FLOAT", {"default": 4.5, "min": 0.0, "max": 10.0, "step": 1/2, "round": 1/100, "tooltip":"It is an equivalent to the CFG scale."}), "strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 1/10, "round": 1/10}), "end_at_sigma": ("FLOAT", {"default": 0.28, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), # "free_scale" : ("BOOLEAN", {"default": False}), "converging_scales" : ("BOOLEAN", {"default": True}), # "noise_add_diff" : ("BOOLEAN", {"default": True}), # "split_channels" : ("BOOLEAN", {"default": False}), "invert_mask" : ("BOOLEAN", {"default": False}), # "no_input" : (["rand","rev","cond","uncond","swap","r_swap","diff","add_diff","rand_rev","rev_cond","rand_cond","rev_cond_sp","cond_rev_sp"],), # "start_at_sigma": ("FLOAT", {"default": 15, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), # "end_at_sigma": ("FLOAT", {"default": 0.28, "min": 0.0, "max": 1000.0, "step": 1/100, "round": 1/100}), }, "optional":{ "input_mask": ("MASK", {"tooltip":"If only a mask is connected the scale becomes a CFG scale of what is being masked.\nWhen a latent is connected the mask defines what will be modified by the node."},), "input_latent": ("LATENT", {"tooltip":"If a latent is connected the scale becomes the maximum scale allowed in which to seek similarity."},), } } RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "model_patches/Pre CFG" def get_latent_guidance_mask_channel(self,x_orig,cond,uncond,guide,minimum_scale,maximum_scale,noise_add_diff): scales = torch.zeros_like(x_orig, device=x_orig.device) for b in range(cond.shape[0]): for c in range(cond.shape[1]): scales[b][c] = self.get_latent_guidance_mask(x_orig[b][c],cond[b][c],uncond[b][c],guide[0][c],minimum_scale,maximum_scale,noise_add_diff) return scales @torch.no_grad() def get_latent_guidance_mask(self,x_orig,cond,uncond,guide,minimum_scale,maximum_scale,noise_add_diff): low_denoised = get_denoised_at_scale(x_orig,cond,uncond,minimum_scale) high_denoised = get_denoised_at_scale(x_orig,cond,uncond,maximum_scale) if noise_add_diff: guide = guide + (guide - (x_orig * guide.norm() / x_orig.norm())) guide = guide / guide.norm() low_diff = (low_denoised - guide * low_denoised.norm()).abs() high_diff = (high_denoised - guide * high_denoised.norm()).abs() return torch.clamp(low_diff / high_diff, min=0, max=1) def patch(self, model, maximum_scale, minimum_scale, invert_mask, strength, end_at_sigma, start_at_sigma=99999, no_input="swap", noise_add_diff=True, converging_scales=False, split_channels=False, free_scale=False, input_mask=None, input_latent=None): sigma_min, sigma_max = get_sigma_min_max(model) model_sampling = model.model.model_sampling scaling_function = self.get_latent_guidance_mask_channel if split_channels else self.get_latent_guidance_mask mask_as_weight = None latent_as_guidance = None random_guidance = False if input_mask is not None: mask_as_weight = input_mask.clone().to(device=default_device) if invert_mask: mask_as_weight = 1 - mask_as_weight if mask_as_weight.dim() == 3: mask_as_weight = mask_as_weight.unsqueeze(1) if input_latent is not None: latent_as_guidance = input_latent["samples"].clone().to(device=default_device) elif input_mask is None: random_guidance = True snc = lambda x: x / x.norm() trl = lambda x: torch.randn_like(x,device=x.device) no_input_operations = { "rand": lambda x, y, o, z, s: snc(trl(x)) * x.norm(), "rev": lambda x, y, o, z, s: x * -1, "cond": lambda x, y, o, z, s: snc(y) * x.norm(), "uncond": lambda x, y, o, z, s: snc(o) * x.norm() * -1, "swap": lambda x, y, o, z, s: no_input_operations["cond"](x, y, o, z, s) if s > 0.36 else no_input_operations["uncond"](x, y, o, z, s), "r_swap": lambda x, y, o, z, s: no_input_operations["cond"](x, y, o, z, s) if s <= 0.36 else no_input_operations["uncond"](x, y, o, z, s), "diff": lambda x, y, o, z, s: snc(y - o) * x.norm(), "add_diff": lambda x, y, o, z, s: snc(y + y - o) * x.norm(), "rand_rev": lambda x, y, o, z, s: snc(trl(x)) * x.norm(), "rev_cond": lambda x, y, o, z, s: (snc(x) * -1 + snc(y) * 0.5) * x.norm() / 1.5, "rand_cond": lambda x, y, o, z, s: (snc(x) * -1 + snc(trl(x)) * 0.5) * x.norm() / 1.5, "rev_cond_sp": lambda x, y, o, z, s: no_input_operations["rev"](x,y,z) * z + (1 - z) * no_input_operations["cond"](x,y,z), "cond_rev_sp": lambda x, y, o, z, s: no_input_operations["rev"](x,y,z) * (1 - z) + z * no_input_operations["cond"](x,y,z), } @torch.no_grad() def pre_cfg_patch(args): nonlocal mask_as_weight, latent_as_guidance conds_out = args["conds_out"] cond_scale = args["cond_scale"] x_orig = args['input'] sigma = args["sigma"][0] sp = min(1,max(0,sigma_to_percent(model_sampling, sigma - sigma_min * 3) + 1 / 100)) ** 2 if not torch.any(conds_out[1]) or sigma <= end_at_sigma or sigma > start_at_sigma or (converging_scales and sp == 1): return conds_out if converging_scales: current_maximum_scale = sp * cond_scale + (1 - sp) * maximum_scale current_minimum_scale = sp * cond_scale + (1 - sp) * minimum_scale else: current_maximum_scale = maximum_scale current_minimum_scale = minimum_scale if mask_as_weight is not None and mask_as_weight.shape[-2:] != conds_out[1].shape[-2:]: mask_as_weight = F.interpolate(mask_as_weight, size=(conds_out[1].shape[-2], conds_out[1].shape[-1]), mode='bilinear', align_corners=False) if random_guidance: latent_as_guidance = no_input_operations[no_input](x_orig.clone(),conds_out[0].clone(),conds_out[1].clone(),sp,sigma/sigma_max) if latent_as_guidance is not None: if latent_as_guidance.shape[-2:] != conds_out[1].shape[-2:]: latent_as_guidance = F.interpolate(latent_as_guidance, size=(conds_out[1].shape[-2], conds_out[1].shape[-1]), mode='bilinear', align_corners=False) scaling_weight = scaling_function(x_orig,conds_out[0],conds_out[1],latent_as_guidance.clone(),current_minimum_scale,current_maximum_scale,noise_add_diff) target_scales = scaling_weight * current_maximum_scale + (1 - scaling_weight) * current_minimum_scale if free_scale: target_scales = target_scales * cond_scale / target_scales.mean() global_multiplier = strength if mask_as_weight is not None: global_multiplier = global_multiplier * mask_as_weight target_scales = target_scales * global_multiplier + torch.full_like(target_scales, cond_scale) * (1 - global_multiplier) conds_out[1] = make_new_uncond_at_scale(conds_out[0],conds_out[1],cond_scale,target_scales) return conds_out else: target_scales = maximum_scale * mask_as_weight * strength + torch.full_like(conds_out[1], cond_scale) * (1 - mask_as_weight * strength) conds_out[1] = make_new_uncond_at_scale(conds_out[0],conds_out[1],cond_scale,target_scales) return conds_out m = model.clone() m.set_model_sampler_pre_cfg_function(pre_cfg_patch) return (m, ) class EmptyRGBImage: @classmethod def INPUT_TYPES(s): return {"required": { "width": ("INT", {"default": 1024, "min": 1, "max": 16384, "step": 1}), "height": ("INT", {"default": 1024, "min": 1, "max": 16384, "step": 1}), "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), "r": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), "g": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), "b": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), }, "optional": { "grayscale_to_color": ("IMAGE",), }} RETURN_TYPES = ("IMAGE",) FUNCTION = "generate" CATEGORY = "image" def generate(self, width, height, batch_size=1, r=0, g=0, b=0, grayscale_to_color=None): if grayscale_to_color is not None: grayscale_to_color = grayscale_to_color.permute(0, 3, 1, 2).mean(dim=1).unsqueeze(-1) height = grayscale_to_color.shape[1] width = grayscale_to_color.shape[2] r_normalized = torch.full([batch_size, height, width, 1], r / 255.0) g_normalized = torch.full([batch_size, height, width, 1], g / 255.0) b_normalized = torch.full([batch_size, height, width, 1], b / 255.0) rgb_image = torch.cat((r_normalized, g_normalized, b_normalized), dim=-1) if grayscale_to_color is not None: rgb_image = rgb_image * grayscale_to_color return (rgb_image,) gradient_patterns = { "linear": lambda x, y: x, "sine": lambda x, y: torch.sin(x * torch.pi * y), "triangle": lambda x, y: 2 * torch.abs(torch.round(x % (1 / max(y, 1)) * y) - (x % (1 / max(y, 1)) * y)), } class GradientRGBImage: @classmethod def INPUT_TYPES(s): return {"required": { "width": ("INT", {"default": 1024, "min": 0, "max": 16384, "step": 64}), "height": ("INT", {"default": 1024, "min": 0, "max": 16384, "step": 64}), "r1": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), "g1": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), "b1": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), "r2": ("INT", {"default": 255, "min": 0, "max": 255, "step": 1}), "g2": ("INT", {"default": 255, "min": 0, "max": 255, "step": 1}), "b2": ("INT", {"default": 255, "min": 0, "max": 255, "step": 1}), "axis" : (["vertical","horizontal","circular"],), "power_to": ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}), "reverse_power" : ("BOOLEAN", {"default": False}), }, "optional":{ "mask": ("MASK",), } } RETURN_TYPES = ("IMAGE","MASK",) FUNCTION = "generate" CATEGORY = "image" def get_gradient_mask(self,width,height,horizontal): if horizontal: return torch.linspace(0, 1, width).view(1, 1, width).repeat(1, height, 1) return torch.linspace(0, 1, height).view(1, height, 1).repeat(1, 1, width) def generate(self, width, height, batch_size=1, r1=0, g1=0, b1=0, r2=255, g2=255, b2=255, pattern_value=1, power_to=1, reverse_power=False, axis="vertical", mask=None): gradient = self.get_gradient_mask(width, height, axis in ["horizontal","circular"]) gradient = gradient_patterns["linear" if axis != "circular" else "sine"](gradient, pattern_value) if axis == "circular": gradient2 = self.get_gradient_mask(width, height, False) gradient2 = gradient_patterns["sine"](gradient2, pattern_value) gradient = gradient * gradient2 if power_to > 1: if reverse_power: gradient = 1 - gradient gradient = gradient ** power_to if reverse_power: gradient = 1 - gradient if mask is not None: if mask.shape != gradient.shape: mask = F.interpolate(mask.unsqueeze(1), size=(gradient.shape[-2], gradient.shape[-1]), mode='nearest').squeeze(1) gradient = gradient * mask gradient = gradient.squeeze(0).unsqueeze(-1) r_gradient = r1 / 255.0 + gradient * (r2 - r1) / 255.0 g_gradient = g1 / 255.0 + gradient * (g2 - g1) / 255.0 b_gradient = b1 / 255.0 + gradient * (b2 - b1) / 255.0 r_image = r_gradient.expand(batch_size, height, width, 1) g_image = g_gradient.expand(batch_size, height, width, 1) b_image = b_gradient.expand(batch_size, height, width, 1) rgb_image = torch.cat((r_image, g_image, b_image), dim=-1) mask_gradient = gradient.expand(1, height, width, 1).squeeze(-1) return (rgb_image,mask_gradient,)