diff --git a/__init__.py b/__init__.py index f2afa3a..dca9739 100644 --- a/__init__.py +++ b/__init__.py @@ -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__) diff --git a/prompt_control/attention_couple_ppm.py b/prompt_control/attention_couple_ppm.py index 95509a8..3e9aeb4 100644 --- a/prompt_control/attention_couple_ppm.py +++ b/prompt_control/attention_couple_ppm.py @@ -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) diff --git a/prompt_control/nodes_hooks.py b/prompt_control/nodes_hooks.py index d5a3751..e36706b 100644 --- a/prompt_control/nodes_hooks.py +++ b/prompt_control/nodes_hooks.py @@ -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)", }