Add PCAttentionCoupleBatchNegative;
Revert some optimizations in AttentionCoupleHook
This commit is contained in:
Pam
+1
-1
@@ -29,7 +29,7 @@ cache_hack.init()
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
nodes = ["base", "lazy", "tools"]
|
||||
nodes = ["base", "lazy", "tools", "hooks"]
|
||||
|
||||
for node in nodes:
|
||||
mod = importlib.import_module(f".prompt_control.nodes_{node}", package=__name__)
|
||||
|
||||
@@ -2,33 +2,45 @@
|
||||
# Original implementation by laksjdjf, hako-mikan, Haoming02 licensed under GPL-3.0
|
||||
# https://github.com/laksjdjf/cgem156-ComfyUI/blob/1f5533f7f31345bafe4b833cbee15a3c4ad74167/scripts/attention_couple/node.py
|
||||
# https://github.com/Haoming02/sd-forge-couple/blob/e8e258e982a8d149ba59a4bc43b945467604311c/scripts/attention_couple.py
|
||||
import itertools
|
||||
import logging
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from comfy.hooks import TransformerOptionsHook, HookGroup, EnumHookScope, set_hooks_for_conditioning
|
||||
|
||||
from comfy.hooks import EnumHookScope, HookGroup, TransformerOptionsHook, set_hooks_for_conditioning
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
import logging
|
||||
|
||||
log = logging.getLogger("comfyui-prompt-control")
|
||||
|
||||
|
||||
def set_cond_attnmask(base_cond, extra_conds, fill=False):
|
||||
hook = AttentionCoupleHook(base_cond[0], extra_conds, fill=fill)
|
||||
hook = AttentionCoupleHook()
|
||||
hook.initialize_regions(base_cond[0], extra_conds, fill=fill)
|
||||
group = HookGroup()
|
||||
group.add(hook)
|
||||
return set_hooks_for_conditioning(base_cond, hooks=group)
|
||||
|
||||
|
||||
def lcm_for_list(numbers):
|
||||
def lcm_for_list(numbers: list[int]):
|
||||
current_lcm = numbers[0]
|
||||
for number in numbers[1:]:
|
||||
current_lcm = math.lcm(current_lcm, number)
|
||||
return current_lcm
|
||||
|
||||
|
||||
def get_mask(mask, batch_size, num_tokens, extra_options):
|
||||
activations_shape = extra_options["activations_shape"]
|
||||
size = activations_shape[-2:]
|
||||
|
||||
num_conds = mask.shape[0]
|
||||
mask_downsample = F.interpolate(mask, size=size, mode="nearest")
|
||||
mask_downsample_reshaped = mask_downsample.view(num_conds, num_tokens, 1).repeat_interleave(batch_size, dim=0)
|
||||
|
||||
return mask_downsample_reshaped
|
||||
|
||||
|
||||
class Proxy:
|
||||
def __init__(self, function):
|
||||
self.function = function
|
||||
@@ -42,76 +54,79 @@ class Proxy:
|
||||
|
||||
|
||||
class AttentionCoupleHook(TransformerOptionsHook):
|
||||
def __init__(self, base_cond, conds, fill):
|
||||
COND_UNCOND_COUPLE_OPTION = "cond_or_uncond_hook_couple"
|
||||
COND = 0
|
||||
UNCOND = 1
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(hook_scope=EnumHookScope.HookedOnly)
|
||||
|
||||
self.transformers_dict = {
|
||||
"patches": {
|
||||
"attn2_output_patch": [Proxy(self.attn2_output_patch)],
|
||||
"attn2_patch": [Proxy(self.attn2_patch)],
|
||||
}
|
||||
}
|
||||
self.has_negpip = False
|
||||
|
||||
# calculate later
|
||||
self.conds_kv: list[tuple[torch.Tensor, torch.Tensor]] = None
|
||||
self.num_tokens_k: list[int] = None
|
||||
self.num_tokens_v: list[int] = None
|
||||
self.lcm_tokens_k: int = None
|
||||
self.lcm_tokens_v: int = None
|
||||
|
||||
def initialize_regions(self, base_cond, conds, fill):
|
||||
self._base_cond = base_cond
|
||||
self._conds = conds
|
||||
self._fill = fill
|
||||
|
||||
self.num_conds = len(conds) + 1
|
||||
self.base_strength = base_cond[1].pop("strength", 1.0)
|
||||
self.base_strength = base_cond[1].get("strength", 1.0)
|
||||
self.strengths = [cond[1].get("strength", 1.0) for cond in conds]
|
||||
self.conds: list[torch.Tensor] = [base_cond[0]] + [cond[0] for cond in conds]
|
||||
base_mask = base_cond[1].pop("mask", None)
|
||||
masks = [cond[1].pop("mask") * cond[1].pop("mask_strength") for cond in conds]
|
||||
base_mask = base_cond[1].get("mask", None)
|
||||
masks = [cond[1].get("mask") * cond[1].get("mask_strength") for cond in conds]
|
||||
|
||||
if base_mask is None and not fill:
|
||||
raise ValueError("You must specify a base mask when fill=False")
|
||||
elif base_mask is None:
|
||||
if base_mask is None:
|
||||
if not fill:
|
||||
raise ValueError("You must specify a base mask when fill=False")
|
||||
sum = torch.stack(masks, dim=0).sum(dim=0)
|
||||
base_mask = torch.zeros_like(sum)
|
||||
base_mask[sum <= 0] = 1.0
|
||||
|
||||
mask = [base_mask] + masks
|
||||
mask = torch.stack(mask, dim=0)
|
||||
if mask.sum(dim=0).min() <= 0 and not fill:
|
||||
raise ValueError("Masks contain non-filled areas")
|
||||
|
||||
self.mask = mask / mask.sum(dim=0, keepdim=True)
|
||||
# calculate later
|
||||
self.conds_k_tensor = None
|
||||
self.conds_v_tensor = None
|
||||
|
||||
def on_apply_hooks(self, model: ModelPatcher, transformer_options: dict[str]):
|
||||
if self.conds_k_tensor is None:
|
||||
if self.conds_kv is None:
|
||||
attn_patches = model.model_options["transformer_options"].get("patches", {}).get("attn2_patch", [])
|
||||
has_negpip = any("negpip_attn" in i.__name__ for i in attn_patches)
|
||||
log.debug("AttentionCouple has_negpip=%s", has_negpip)
|
||||
self.has_negpip = any("negpip_attn" in i.__name__ for i in attn_patches)
|
||||
log.debug("AttentionCouple has_negpip=%s", self.has_negpip)
|
||||
|
||||
conds_kv = (
|
||||
self.conds_kv = (
|
||||
[(cond[:, 0::2], cond[:, 1::2]) for cond in self.conds]
|
||||
if has_negpip
|
||||
if self.has_negpip
|
||||
else [(cond, cond) for cond in self.conds]
|
||||
)
|
||||
|
||||
num_tokens_k = [cond[0].shape[1] for cond in conds_kv]
|
||||
num_tokens_v = [cond[1].shape[1] for cond in conds_kv]
|
||||
self.num_tokens_k = [cond[0].shape[1] for cond in self.conds_kv]
|
||||
self.num_tokens_v = [cond[1].shape[1] for cond in self.conds_kv]
|
||||
|
||||
lcm_tokens_k = lcm_for_list(num_tokens_k)
|
||||
lcm_tokens_v = lcm_for_list(num_tokens_v)
|
||||
# Skip the base cond here, which is always first
|
||||
self.conds_k_tensor = torch.cat(
|
||||
[
|
||||
cond[0].repeat(1, lcm_tokens_k // num_tokens_k[i + 1], 1) * self.strengths[i]
|
||||
for i, cond in enumerate(conds_kv[1:])
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
if has_negpip:
|
||||
self.conds_v_tensor = torch.cat(
|
||||
[
|
||||
cond[1].repeat(1, lcm_tokens_v // num_tokens_v[i + 1], 1) * self.strengths[i]
|
||||
for i, cond in enumerate(conds_kv[1:])
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
else:
|
||||
self.conds_v_tensor = self.conds_k_tensor
|
||||
self.lcm_tokens_k = lcm_for_list(self.num_tokens_k)
|
||||
self.lcm_tokens_v = lcm_for_list(self.num_tokens_v)
|
||||
|
||||
return super().on_apply_hooks(model, transformer_options)
|
||||
|
||||
def clone(self):
|
||||
c: AttentionCoupleHook = super().clone()
|
||||
c.initialize_regions(self._base_cond, self._conds, self._fill)
|
||||
return c
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
self.conds = [c.to(*args, **kwargs) for c in self.conds]
|
||||
self.mask = self.mask.to(*args, **kwargs)
|
||||
@@ -119,32 +134,77 @@ class AttentionCoupleHook(TransformerOptionsHook):
|
||||
|
||||
def attn2_patch(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, extra_options):
|
||||
cond_or_uncond = extra_options["cond_or_uncond"]
|
||||
num_chunks = len(cond_or_uncond) # should always be 1
|
||||
cond_or_uncond_couple = extra_options[self.COND_UNCOND_COUPLE_OPTION] = list(cond_or_uncond)
|
||||
num_chunks = len(cond_or_uncond)
|
||||
|
||||
q_chunks = q.chunk(num_chunks, dim=0)
|
||||
k_chunks = k.chunk(num_chunks, dim=0)
|
||||
v_chunks = v.chunk(num_chunks, dim=0)
|
||||
|
||||
bs = q.shape[0] // num_chunks
|
||||
# Skip the base cond here, which is always first
|
||||
conds_k_tensor = torch.cat(
|
||||
[
|
||||
cond[0].repeat(bs, self.lcm_tokens_k // self.num_tokens_k[i + 1], 1) * self.strengths[i]
|
||||
for i, cond in enumerate(self.conds_kv[1:])
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
if self.has_negpip:
|
||||
conds_v_tensor = torch.cat(
|
||||
[
|
||||
cond[1].repeat(bs, self.lcm_tokens_v // self.num_tokens_v[i + 1], 1) * self.strengths[i]
|
||||
for i, cond in enumerate(self.conds_kv[1:])
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
else:
|
||||
conds_v_tensor = conds_k_tensor
|
||||
|
||||
conds_k_tensor = self.conds_k_tensor.repeat(bs, 1, 1)
|
||||
conds_v_tensor = self.conds_v_tensor.repeat(bs, 1, 1)
|
||||
qs, ks, vs = [], [], []
|
||||
cond_or_uncond_couple.clear()
|
||||
|
||||
q = q.repeat(self.num_conds, 1, 1)
|
||||
k = k.repeat(1, self.conds_k_tensor.shape[1] // k.shape[1], 1)
|
||||
v = v.repeat(1, self.conds_v_tensor.shape[1] // v.shape[1], 1)
|
||||
for i, cond_type in enumerate(cond_or_uncond):
|
||||
q_target = q_chunks[i]
|
||||
k_target = k_chunks[i].repeat(1, self.lcm_tokens_k // k.shape[1], 1)
|
||||
v_target = v_chunks[i].repeat(1, self.lcm_tokens_v // v.shape[1], 1)
|
||||
if cond_type == self.UNCOND:
|
||||
qs.append(q_target)
|
||||
ks.append(k_target)
|
||||
vs.append(v_target)
|
||||
cond_or_uncond_couple.append(self.UNCOND)
|
||||
else:
|
||||
qs.append(q_target.repeat(self.num_conds, 1, 1))
|
||||
ks.append(torch.cat([k_target * self.base_strength, conds_k_tensor], dim=0))
|
||||
vs.append(torch.cat([v_target * self.base_strength, conds_v_tensor], dim=0))
|
||||
cond_or_uncond_couple.extend(itertools.repeat(self.COND, self.num_conds))
|
||||
|
||||
k = torch.cat([k * self.base_strength, conds_k_tensor], dim=0)
|
||||
v = torch.cat([v * self.base_strength, conds_v_tensor], dim=0)
|
||||
qs = torch.cat(qs, dim=0)
|
||||
ks = torch.cat(ks, dim=0)
|
||||
vs = torch.cat(vs, dim=0)
|
||||
|
||||
return q, k, v
|
||||
return qs, ks, vs
|
||||
|
||||
def attn2_output_patch(self, out, extra_options):
|
||||
# out has been extended to shape [num_conds*batch_size, TOKENS, N]
|
||||
# out is [b1c1 b1c2 ... b1cN, b2c1 b2c2 ... b2cn, ...]
|
||||
num_conds = self.mask.shape[0]
|
||||
bs = out.shape[0] // num_conds
|
||||
num_tokens = out.shape[1]
|
||||
mask_size = extra_options["activations_shape"][-2:]
|
||||
mask_downsample = F.interpolate(self.mask, size=mask_size, mode="nearest")
|
||||
mask_downsample = mask_downsample.view(num_conds, num_tokens, 1).repeat_interleave(bs, dim=0)
|
||||
cond_or_uncond = extra_options[self.COND_UNCOND_COUPLE_OPTION]
|
||||
bs = out.shape[0] // len(cond_or_uncond)
|
||||
mask_downsample = get_mask(self.mask, bs, out.shape[1], extra_options)
|
||||
outputs = []
|
||||
cond_outputs = []
|
||||
i_cond = 0
|
||||
for i, cond_type in enumerate(cond_or_uncond):
|
||||
pos, next_pos = i * bs, (i + 1) * bs
|
||||
|
||||
# cond_outputs is [num_conds*bs, tokens, N], output needs to be [bs, tokens, N]
|
||||
cond_outputs = out * mask_downsample
|
||||
cond_output = cond_outputs.view(num_conds, bs, out.shape[1], out.shape[2]).sum(0)
|
||||
return cond_output
|
||||
if cond_type == self.UNCOND:
|
||||
outputs.append(out[pos:next_pos])
|
||||
else:
|
||||
pos_cond, next_pos_cond = i_cond * bs, (i_cond + 1) * bs
|
||||
masked_output = out[pos:next_pos] * mask_downsample[pos_cond:next_pos_cond]
|
||||
cond_outputs.append(masked_output)
|
||||
i_cond += 1
|
||||
|
||||
if len(cond_outputs) > 0:
|
||||
cond_output = torch.stack(cond_outputs).sum(0)
|
||||
outputs.append(cond_output)
|
||||
|
||||
return torch.cat(outputs, dim=0)
|
||||
|
||||
@@ -1,9 +1,14 @@
|
||||
import logging
|
||||
import comfy.utils
|
||||
|
||||
import comfy.hooks
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
from .utils import consolidate_schedule
|
||||
from comfy.comfy_types.node_typing import IO, ComfyNodeABC, InputTypeDict
|
||||
from node_helpers import conditioning_set_values
|
||||
|
||||
from .attention_couple_ppm import AttentionCoupleHook
|
||||
from .parser import parse_prompt_schedules
|
||||
from .utils import consolidate_schedule
|
||||
|
||||
log = logging.getLogger("comfyui-prompt-control")
|
||||
|
||||
@@ -79,10 +84,45 @@ def lora_hooks_from_schedule(schedules, non_scheduled):
|
||||
return hooks
|
||||
|
||||
|
||||
class PCAttentionCoupleBatchNegative(ComfyNodeABC):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||||
return {
|
||||
"required": {
|
||||
"positive": (IO.CONDITIONING, {}),
|
||||
"negative": (IO.CONDITIONING, {}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (IO.CONDITIONING, IO.CONDITIONING)
|
||||
RETURN_NAMES = ("positive", "negative")
|
||||
CATEGORY = "promptcontrol/v2"
|
||||
FUNCTION = "batch"
|
||||
EXPERIMENTAL = True
|
||||
|
||||
# May cause side-effects?
|
||||
def batch(self, positive, negative):
|
||||
positive_hook_group: comfy.hooks.HookGroup = positive[0][1].get("hooks", comfy.hooks.HookGroup())
|
||||
attn_couple = [hook for hook in positive_hook_group.hooks if isinstance(hook, AttentionCoupleHook)]
|
||||
|
||||
negative_hook_group: comfy.hooks.HookGroup = negative[0][1].get("hooks", comfy.hooks.HookGroup())
|
||||
negative_hook_group_couple = negative_hook_group.clone()
|
||||
for hook in attn_couple:
|
||||
negative_hook_group_couple.add(hook)
|
||||
|
||||
if negative_hook_group_couple.hooks != positive_hook_group.hooks:
|
||||
negative_batch = conditioning_set_values(negative, {"hooks": negative_hook_group_couple})
|
||||
else:
|
||||
negative_batch = conditioning_set_values(negative, {"hooks": positive_hook_group})
|
||||
return (positive, negative_batch)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PCLoraHooksFromText": PCLoraHooksFromText,
|
||||
"PCAttentionCoupleBatchNegative": PCAttentionCoupleBatchNegative,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PCLoraHooksFromText": "PC: LoRA Hooks From Text (non-lazy)",
|
||||
"PCAttentionCoupleBatchNegative": "PC: AttentionCouple (batch negative)",
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user