101 lines
3.2 KiB
Python
101 lines
3.2 KiB
Python
|
|
import torch
|
|
|
|
|
|
class MomentumBuffer:
|
|
def __init__(self, momentum: float):
|
|
self.momentum = momentum
|
|
self.running_average = 0
|
|
self.prev_sigma = None
|
|
|
|
def update(self, update_value: torch.Tensor, sigma: float):
|
|
if self.prev_sigma is not None and sigma > self.prev_sigma:
|
|
self.running_average = 0
|
|
self.prev_sigma = sigma
|
|
|
|
new_average = self.momentum * self.running_average
|
|
self.running_average = update_value + new_average
|
|
|
|
|
|
def project(v0: torch.Tensor, v1: torch.Tensor):
|
|
dtype = v0.dtype
|
|
v0, v1 = v0.double(), v1.double()
|
|
v1 = torch.nn.functional.normalize(v1, dim=[-1, -2, -3])
|
|
v0_parallel = (v0 * v1).sum(dim=[-1, -2, -3], keepdim=True) * v1
|
|
v0_orthogonal = v0 - v0_parallel
|
|
return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
|
|
|
|
|
|
def normalized_guidance(pred_cond: torch.Tensor,
|
|
pred_uncond: torch.Tensor,
|
|
guidance_scale: float,
|
|
sigma: float,
|
|
momentum_buffer: MomentumBuffer = None,
|
|
eta: float = 1.0,
|
|
norm_threshold: float = 0.0):
|
|
|
|
diff = pred_cond - pred_uncond
|
|
if momentum_buffer is not None:
|
|
momentum_buffer.update(diff, sigma)
|
|
diff = momentum_buffer.running_average
|
|
|
|
if norm_threshold > 0:
|
|
ones = torch.ones_like(diff)
|
|
diff_norm = diff.norm(p=2, dim=[-1, -2, -3], keepdim=True)
|
|
scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
|
|
diff = diff * scale_factor
|
|
|
|
diff_parallel, diff_orthogonal = project(diff, pred_cond)
|
|
normalized_update = diff_orthogonal + eta * diff_parallel
|
|
pred_guided = pred_cond + (guidance_scale - 1) * normalized_update
|
|
|
|
return pred_guided
|
|
|
|
|
|
class APGFunctionNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
|
"norm_threshold": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 10.0, "step": 0.01}),
|
|
"use_momentum": (["enable", "disable"],),
|
|
"momentum": ("FLOAT", {"default": 0.05, "min": -1.00, "max": 1.00, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "patch"
|
|
|
|
CATEGORY = "apg"
|
|
|
|
def patch(self, model, eta, norm_threshold, use_momentum, momentum):
|
|
m = model.clone()
|
|
|
|
momentum_buffer = MomentumBuffer(momentum) if use_momentum == "enable" else None
|
|
|
|
def cfg_function(args):
|
|
sigma = args["sigma"]
|
|
cond_pred = args["cond"]
|
|
uncond_pred = args["uncond"]
|
|
cfg = args["cond_scale"]
|
|
|
|
if uncond_pred is None:
|
|
return uncond_pred + (cond_pred - uncond_pred) * cfg
|
|
|
|
return normalized_guidance(
|
|
cond_pred,
|
|
uncond_pred,
|
|
cfg,
|
|
sigma.item(),
|
|
momentum_buffer=momentum_buffer,
|
|
eta=eta,
|
|
norm_threshold=norm_threshold
|
|
)
|
|
|
|
m.set_model_sampler_cfg_function(cfg_function)
|
|
|
|
return (m,)
|
|
|