mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-04-18 18:33:30 +00:00
Add CFGZeroStar node.
Works on all models that use a negative prompt but is meant for rectified flow models.
This commit is contained in:
parent
8edc1f44c1
commit
84fdaf7b0e
45
comfy_extras/nodes_cfg.py
Normal file
45
comfy_extras/nodes_cfg.py
Normal file
@ -0,0 +1,45 @@
|
|||||||
|
import torch
|
||||||
|
|
||||||
|
# https://github.com/WeichenFan/CFG-Zero-star
|
||||||
|
def optimized_scale(positive, negative):
|
||||||
|
positive_flat = positive.reshape(positive.shape[0], -1)
|
||||||
|
negative_flat = negative.reshape(negative.shape[0], -1)
|
||||||
|
|
||||||
|
# Calculate dot production
|
||||||
|
dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
|
||||||
|
|
||||||
|
# Squared norm of uncondition
|
||||||
|
squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8
|
||||||
|
|
||||||
|
# st_star = v_cond^T * v_uncond / ||v_uncond||^2
|
||||||
|
st_star = dot_product / squared_norm
|
||||||
|
|
||||||
|
return st_star.reshape([positive.shape[0]] + [1] * (positive.ndim - 1))
|
||||||
|
|
||||||
|
class CFGZeroStar:
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(s):
|
||||||
|
return {"required": {"model": ("MODEL",),
|
||||||
|
}}
|
||||||
|
RETURN_TYPES = ("MODEL",)
|
||||||
|
RETURN_NAMES = ("patched_model",)
|
||||||
|
FUNCTION = "patch"
|
||||||
|
CATEGORY = "advanced/guidance"
|
||||||
|
|
||||||
|
def patch(self, model):
|
||||||
|
m = model.clone()
|
||||||
|
def cfg_zero_star(args):
|
||||||
|
guidance_scale = args['cond_scale']
|
||||||
|
x = args['input']
|
||||||
|
cond_p = args['cond_denoised']
|
||||||
|
uncond_p = args['uncond_denoised']
|
||||||
|
out = args["denoised"]
|
||||||
|
alpha = optimized_scale(x - cond_p, x - uncond_p)
|
||||||
|
|
||||||
|
return out + uncond_p * (alpha - 1.0) + guidance_scale * uncond_p * (1.0 - alpha)
|
||||||
|
m.set_model_sampler_post_cfg_function(cfg_zero_star)
|
||||||
|
return (m, )
|
||||||
|
|
||||||
|
NODE_CLASS_MAPPINGS = {
|
||||||
|
"CFGZeroStar": CFGZeroStar
|
||||||
|
}
|
Loading…
Reference in New Issue
Block a user