224 lines
8.1 KiB
Python
224 lines
8.1 KiB
Python
# Adapted from https://github.com/laksjdjf/cgem156-ComfyUI/blob/main/scripts/attention_couple/node.py
|
|
# by @laksjdjf
|
|
|
|
from __future__ import annotations
|
|
from typing import NamedTuple
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import math
|
|
from torch import Tensor, Size
|
|
from comfy.model_patcher import ModelPatcher
|
|
from comfy_api.latest import io
|
|
|
|
|
|
def downsample_mask(mask: Tensor, batch: int, target_size: int, original_shape: Size) -> Tensor:
|
|
h, w = original_shape[2], original_shape[3]
|
|
hm, wm = mask.shape[2], mask.shape[3]
|
|
if (h, w) == (hm, wm): # Mask is already in latent resolution
|
|
base_factor = 1
|
|
elif (h * 8, w * 8) == (hm, wm): # Mask is in image resolution, downsample by 8
|
|
base_factor = 8
|
|
else:
|
|
raise ValueError(f"Bad mask size. Expected {w}x{h}, got {wm}x{hm}.")
|
|
|
|
result = mask
|
|
for factor in [1, 2, 4, 8]:
|
|
size = (math.ceil(h / factor), math.ceil(w / factor))
|
|
if size[0] * size[1] == target_size and base_factor * factor > 1:
|
|
result = F.interpolate(mask, size=size, mode="nearest")
|
|
break
|
|
|
|
num_conds = mask.shape[0]
|
|
result = result.view(num_conds, target_size, 1)
|
|
result = result.repeat_interleave(batch, dim=0)
|
|
return result
|
|
|
|
|
|
def lcm(a: int, b: int):
|
|
return a * b // math.gcd(a, b)
|
|
|
|
|
|
def lcm_for_list(numbers: list[int]):
|
|
current_lcm = numbers[0]
|
|
for number in numbers[1:]:
|
|
current_lcm = lcm(current_lcm, number)
|
|
return current_lcm
|
|
|
|
|
|
class Region(NamedTuple):
|
|
previous: "Region" | None
|
|
mask: Tensor | None
|
|
conditioning: list
|
|
|
|
def preprocess(self):
|
|
result: list[Region] = []
|
|
current = self
|
|
while current is not None:
|
|
result.append(current)
|
|
current = current.previous
|
|
assert len(result) > 1, "At least 2 regions are required."
|
|
|
|
result = list(reversed(result))
|
|
if result[0].mask is None: # BackgroundRegion
|
|
masks_above = torch.stack([r.mask for r in result[1:]], dim=0)
|
|
accumulated = torch.sum(masks_above, dim=0)
|
|
result[0] = Region(None, 1.0 - accumulated, result[0].conditioning)
|
|
return result
|
|
|
|
|
|
Regions = io.Custom("Regions")
|
|
|
|
|
|
class BackgroundRegion(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_BackgroundRegion",
|
|
display_name="Background Region",
|
|
category="external_tooling/regions",
|
|
inputs=[io.Conditioning.Input("conditioning")],
|
|
outputs=[Regions.Output(display_name="regions")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, conditioning: list):
|
|
return (Region(None, None, conditioning),)
|
|
|
|
|
|
class DefineRegion(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_DefineRegion",
|
|
display_name="Define Region",
|
|
category="external_tooling/regions",
|
|
inputs=[
|
|
io.Mask.Input("mask"),
|
|
io.Conditioning.Input("conditioning"),
|
|
Regions.Input("regions", optional=True),
|
|
],
|
|
outputs=[Regions.Output(display_name="regions")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, mask: Tensor, conditioning: list, regions: Region | None = None):
|
|
if mask.dim() < 3:
|
|
mask = mask.unsqueeze(0)
|
|
return io.NodeOutput(Region(regions, mask, conditioning))
|
|
|
|
|
|
class ListRegionMasks(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_ListRegionMasks",
|
|
display_name="List Region Masks",
|
|
category="external_tooling/regions",
|
|
inputs=[Regions.Input("regions")],
|
|
outputs=[io.Mask.Output(display_name="masks")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, regions: Region):
|
|
return io.NodeOutput(torch.stack([r.mask for r in regions.preprocess()], dim=0))
|
|
|
|
|
|
class AttentionMask(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_AttentionMask",
|
|
display_name="Regions Attention Mask",
|
|
category="external_tooling/regions",
|
|
inputs=[io.Model.Input("model"), Regions.Input("regions")],
|
|
outputs=[io.Model.Output(display_name="model")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model: ModelPatcher, regions: Region):
|
|
return io.NodeOutput(AttentionMaskPatch.apply(model, regions))
|
|
|
|
|
|
class AttentionMaskPatch:
|
|
def __init__(self, region_list: list[Region]):
|
|
mask = torch.stack([r.mask for r in region_list], dim=0)
|
|
mask_sum = mask.sum(dim=0, keepdim=True)
|
|
assert mask_sum.sum() > 0, "There are areas that are zero in all masks."
|
|
self.mask = mask / mask_sum
|
|
self.conds = [r.conditioning[0][0] for r in region_list]
|
|
self.num_tokens = [cond.shape[1] for cond in self.conds]
|
|
self.num_conds = len(region_list)
|
|
self.batch_size = 0
|
|
|
|
@staticmethod
|
|
def apply(model: ModelPatcher, regions: Region):
|
|
patch = AttentionMaskPatch(regions.preprocess())
|
|
|
|
def attn2_patch(q: Tensor, k: Tensor, v: Tensor, extra_options: dict):
|
|
assert k.mean() == v.mean(), "k and v must be the same."
|
|
device, dtype = q.device, q.dtype
|
|
|
|
if patch.conds[0].device != device or patch.conds[0].dtype != dtype:
|
|
patch.conds = [cond.to(device, dtype=dtype) for cond in patch.conds]
|
|
if patch.mask.device != device or patch.mask.dtype != dtype:
|
|
patch.mask = patch.mask.to(device, dtype=dtype)
|
|
|
|
cond_or_unconds = extra_options["cond_or_uncond"]
|
|
num_chunks = len(cond_or_unconds)
|
|
patch.batch_size = q.shape[0] // num_chunks
|
|
q_chunks = q.chunk(num_chunks, dim=0)
|
|
k_chunks = k.chunk(num_chunks, dim=0)
|
|
lcm_tokens = lcm_for_list(patch.num_tokens + [k.shape[1]])
|
|
conds_tensor = [
|
|
cond.repeat(patch.batch_size, lcm_tokens // patch.num_tokens[i], 1)
|
|
for i, cond in enumerate(patch.conds)
|
|
]
|
|
conds_tensor = torch.cat(conds_tensor, dim=0)
|
|
|
|
qs, ks = [], []
|
|
for i, cond_or_uncond in reversed(list(enumerate(cond_or_unconds))):
|
|
if cond_or_uncond == 1: # uncond
|
|
k_target = k_chunks[i].repeat(1, lcm_tokens // k.shape[1], 1)
|
|
qs.insert(0, q_chunks[i])
|
|
ks.insert(0, k_target)
|
|
else:
|
|
qs.insert(0, q_chunks[i].repeat(patch.num_conds, 1, 1))
|
|
ks.insert(0, conds_tensor)
|
|
for _ in range(patch.num_conds - 1):
|
|
cond_or_unconds.insert(i, 0)
|
|
|
|
qs = torch.cat(qs, dim=0)
|
|
ks = torch.cat(ks, dim=0)
|
|
return qs, ks, ks
|
|
|
|
def attn2_output_patch(out: Tensor, extra_options: dict):
|
|
num_conds = patch.num_conds
|
|
cond_or_unconds = extra_options["cond_or_uncond"]
|
|
mask_downsample = downsample_mask(
|
|
patch.mask, patch.batch_size, out.shape[1], extra_options["original_shape"]
|
|
)
|
|
outputs: list[Tensor] = []
|
|
pos = 0
|
|
i = 0
|
|
while i < len(cond_or_unconds):
|
|
if cond_or_unconds[i] == 1: # uncond
|
|
outputs.append(out[pos : pos + patch.batch_size])
|
|
pos += patch.batch_size
|
|
else:
|
|
masked = out[pos : pos + num_conds * patch.batch_size] * mask_downsample
|
|
masked = masked.view(num_conds, patch.batch_size, out.shape[1], out.shape[2])
|
|
masked = masked.sum(dim=0)
|
|
outputs.append(masked)
|
|
pos += num_conds * patch.batch_size
|
|
for _ in range(num_conds - 1):
|
|
cond_or_unconds.pop(i)
|
|
i += 1
|
|
|
|
return torch.cat(outputs, dim=0)
|
|
|
|
new_model = model.clone()
|
|
new_model.set_model_attn2_patch(attn2_patch)
|
|
new_model.set_model_attn2_output_patch(attn2_output_patch)
|
|
new_model.set_attachments("etn_attention_mask", patch)
|
|
return new_model
|