1327 lines
60 KiB
Python
1327 lines
60 KiB
Python
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(len(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,)
|