nodes.py

import torch

try:
    from comfy.model_patcher import ModelPatcher

    # #BACKEND = "ComfyUI"
except ImportError:
    try:
        from ldm_patched.modules.model_patcher import ModelPatcher

        # #BACKEND = "reForge"
    except ImportError:
        from backend.patcher.base import ModelPatcher

        # #BACKEND = "Forge"

class MomentumBuffer:
    def __init__(self, momentum: float):
        self.momentum = momentum
        self.running_average = 0
        
    def update(self, update_value: torch.Tensor):
        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, 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)
        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 APG_ImYourCFGNow:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "momentum": ("FLOAT", {"default": 0.5, "min": -1.5, "max": 1.0, "step": 0.01, "round": 0.001}),
                "adaptive_momentum": ("FLOAT", {"default": 0.180, "min": 0.0, "max": 1.0, "step": 0.001, "round": 0.001}),
                "norm_threshold": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 50.0, "step": 0.05, "round": 0.01}),
                "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}),
                "guidance_limiter": ("BOOLEAN", {"default": False}),
                "guidance_sigma_start": ("FLOAT", {"default": 5.42, "min": -1.0, "max": 10000.0, "step": 0.01, "round": False}),
                "guidance_sigma_end": ("FLOAT", {"default": 0.28, "min": -1.0, "max": 10000.0, "step": 0.01, "round": False}),
                "print_data": ("BOOLEAN", {"default": False,}),
            },
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "patch"
    CATEGORY = "model_patches/unet"

    def patch(
        self,
        model: ModelPatcher,
        momentum: float = 0.5,
        adaptive_momentum: float = 0.180,
        norm_threshold: float = 15.0,
        eta: float = 1.0,
        guidance_limiter: bool = False,
        guidance_sigma_start: float = 5.42,
        guidance_sigma_end: float = 0.28,
        print_data = False,
        extras=[],
    ):
        momentum_buffer = MomentumBuffer(momentum)
        extras = [momentum_buffer, momentum, adaptive_momentum]

        def apg_function(args):
            cond = args["cond"]
            uncond = args["uncond"]
            sigma = args["sigma"]
            model = args["model"]
            cond_scale = args["cond_scale"]

            if guidance_limiter:
                if (guidance_sigma_start >= 0 and sigma[0] >  guidance_sigma_start) or \
                   (guidance_sigma_end   >= 0 and sigma[0] <= guidance_sigma_end):
                    if print_data:
                        print(f" guidance limiter active (sigma: {sigma[0]})")
                    return uncond + (cond - uncond)

            momentum_buffer=extras[0]
            momentum=extras[1]
            adaptive_momentum=extras[2]

            t = model.model_sampling.timestep(sigma)[0].item()
            
            if (torch.is_tensor(momentum_buffer.running_average) and (cond.shape[3]!=momentum_buffer.running_average.shape[3])) or t==999:
                momentum_buffer = MomentumBuffer(momentum)
                extras[0]=momentum_buffer
            else:
                signal_scale = momentum
                if adaptive_momentum > 0:
                    if momentum<0:
                        signal_scale += -momentum * (adaptive_momentum**4) * (1000 - t)
                        if signal_scale > 0:
                            signal_scale = 0
                    else:
                        signal_scale -= momentum * (adaptive_momentum**4) * (1000 - t)
                        if signal_scale < 0:
                            signal_scale = 0
                
                momentum_buffer.momentum = signal_scale
                
            if print_data:
                print(" momentum: ", momentum_buffer.momentum, " t: ", t)
            

            return normalized_guidance(cond, uncond, cond_scale, momentum_buffer, eta, norm_threshold)

        m = model.clone()
        m.set_model_sampler_cfg_function(apg_function, extras==extras)
        m.model_options["disable_cfg1_optimization"] = False

        return (m,)
        

NODE_CLASS_MAPPINGS = {
    "APG_ImYourCFGNow": APG_ImYourCFGNow,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "APG_ImYourCFGNow": "APG_ImYourCFGNow",
}