Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 31 additions & 18 deletions comfy_extras/nodes_pag.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,25 +3,30 @@

#My modified one here is more basic but has less chances of breaking with ComfyUI updates.

from typing_extensions import override

import comfy.model_patcher
import comfy.samplers
from comfy_api.latest import ComfyExtension, io

class PerturbedAttentionGuidance:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 100.0, "step": 0.01, "round": 0.01}),
}
}

RETURN_TYPES = ("MODEL",)
FUNCTION = "patch"

CATEGORY = "model_patches/unet"
class PerturbedAttentionGuidance(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="PerturbedAttentionGuidance",
category="model_patches/unet",
inputs=[
io.Model.Input("model"),
io.Float.Input("scale", default=3.0, min=0.0, max=100.0, step=0.01, round=0.01),
],
outputs=[
io.Model.Output(),
],
)

def patch(self, model, scale):
@classmethod
def execute(cls, model, scale) -> io.NodeOutput:
unet_block = "middle"
unet_block_id = 0
m = model.clone()
Expand Down Expand Up @@ -49,8 +54,16 @@ def post_cfg_function(args):

m.set_model_sampler_post_cfg_function(post_cfg_function)

return (m,)
return io.NodeOutput(m)


class PAGExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [
PerturbedAttentionGuidance,
]


NODE_CLASS_MAPPINGS = {
"PerturbedAttentionGuidance": PerturbedAttentionGuidance,
}
async def comfy_entrypoint() -> PAGExtension:
return PAGExtension()
Loading