| import torch |
| from typing_extensions import override |
|
|
| from comfy_api.latest import ComfyExtension, io |
|
|
|
|
| def project(v0, v1): |
| 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, v0_orthogonal |
|
|
| class APG(io.ComfyNode): |
| @classmethod |
| def define_schema(cls) -> io.Schema: |
| return io.Schema( |
| node_id="APG", |
| display_name="Adaptive Projected Guidance", |
| category="sampling/custom_sampling", |
| inputs=[ |
| io.Model.Input("model"), |
| io.Float.Input( |
| "eta", |
| default=1.0, |
| min=-10.0, |
| max=10.0, |
| step=0.01, |
| tooltip="Controls the scale of the parallel guidance vector. Default CFG behavior at a setting of 1.", |
| ), |
| io.Float.Input( |
| "norm_threshold", |
| default=5.0, |
| min=0.0, |
| max=50.0, |
| step=0.1, |
| tooltip="Normalize guidance vector to this value, normalization disable at a setting of 0.", |
| ), |
| io.Float.Input( |
| "momentum", |
| default=0.0, |
| min=-5.0, |
| max=1.0, |
| step=0.01, |
| tooltip="Controls a running average of guidance during diffusion, disabled at a setting of 0.", |
| ), |
| ], |
| outputs=[io.Model.Output()], |
| ) |
|
|
| @classmethod |
| def execute(cls, model, eta, norm_threshold, momentum) -> io.NodeOutput: |
| running_avg = 0 |
| prev_sigma = None |
|
|
| def pre_cfg_function(args): |
| nonlocal running_avg, prev_sigma |
|
|
| if len(args["conds_out"]) == 1: return args["conds_out"] |
|
|
| cond = args["conds_out"][0] |
| uncond = args["conds_out"][1] |
| sigma = args["sigma"][0] |
| cond_scale = args["cond_scale"] |
|
|
| if prev_sigma is not None and sigma > prev_sigma: |
| running_avg = 0 |
| prev_sigma = sigma |
|
|
| guidance = cond - uncond |
|
|
| if momentum != 0: |
| if not torch.is_tensor(running_avg): |
| running_avg = guidance |
| else: |
| running_avg = momentum * running_avg + guidance |
| guidance = running_avg |
|
|
| if norm_threshold > 0: |
| guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True) |
| scale = torch.minimum( |
| torch.ones_like(guidance_norm), |
| norm_threshold / guidance_norm |
| ) |
| guidance = guidance * scale |
|
|
| guidance_parallel, guidance_orthogonal = project(guidance, cond) |
| modified_guidance = guidance_orthogonal + eta * guidance_parallel |
|
|
| modified_cond = (uncond + modified_guidance) + (cond - uncond) / cond_scale |
|
|
| return [modified_cond, uncond] + args["conds_out"][2:] |
|
|
| m = model.clone() |
| m.set_model_sampler_pre_cfg_function(pre_cfg_function) |
| return io.NodeOutput(m) |
|
|
|
|
| class ApgExtension(ComfyExtension): |
| @override |
| async def get_node_list(self) -> list[type[io.ComfyNode]]: |
| return [ |
| APG, |
| ] |
|
|
| async def comfy_entrypoint() -> ApgExtension: |
| return ApgExtension() |
|
|
