501 lines
17 KiB
Python
501 lines
17 KiB
Python
import gc
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
from torch import Tensor
|
|
from comfy.ldm.modules import attention as comfy_attention
|
|
from comfy.ldm.flux import math as flux_math
|
|
from comfy.ldm.flux import layers as flux_layers
|
|
from comfy import model_management
|
|
|
|
from PIL import Image
|
|
from typing import List, Dict, Optional
|
|
from functools import partial
|
|
from einops import rearrange
|
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
orig_attention = comfy_attention.optimized_attention
|
|
|
|
|
|
def masked_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor,
|
|
attn_mask: Tensor = None, q_scale=None, k_scale=None, **kwargs) -> Tensor:
|
|
"""
|
|
Implementation using PyTorch's scaled_dot_product_attention
|
|
|
|
Args:
|
|
q: Query tensor of shape [B, H, L, D]
|
|
k: Key tensor of shape [B, H, L, D]
|
|
v: Value tensor of shape [B, H, L, D]
|
|
pe: Positional encoding tensor
|
|
attn_mask: Optional attention mask
|
|
q_scale: Optional query scaling factor
|
|
k_scale: Optional key scaling factor
|
|
|
|
Returns:
|
|
Output tensor of shape [B, L, H*D]
|
|
"""
|
|
# Apply rotary positional encoding
|
|
q, k = flux_math.apply_rope(q, k, pe)
|
|
|
|
# PyTorch's scaled_dot_product_attention
|
|
# The scaling is handled internally by the function
|
|
x = F.scaled_dot_product_attention(
|
|
q, k, v,
|
|
attn_mask=attn_mask,
|
|
dropout_p=0.0, # Assuming no dropout in the original
|
|
is_causal=False # Assuming not causal, adjust if needed
|
|
)
|
|
|
|
# Reshape back to the expected output format
|
|
x = rearrange(x, "B H L D -> B L (H D)")
|
|
|
|
return x
|
|
|
|
def prepare_attention_mask(lin_masks: List[Image.Image], reg_embeds: List[Tensor],
|
|
Nx: int, emb_size: int, emb_len: int,):
|
|
cross_mask = torch.zeros(emb_len + Nx, emb_len + Nx)
|
|
q_scale = torch.ones(emb_len + Nx)
|
|
k_scale = torch.ones(emb_len + Nx)
|
|
|
|
n_regs = len(lin_masks)
|
|
emb_cum_idx = 0
|
|
|
|
# Mask main prompt to subprompts
|
|
for j in range(n_regs):
|
|
t1, t2 = emb_cum_idx + (j+1) * emb_size, emb_cum_idx + (j+2) * emb_size
|
|
p1, p2 = emb_cum_idx, emb_cum_idx + emb_size
|
|
print(t1, t2, p1, p2)
|
|
|
|
cross_mask[t1 : t2, p1 : p2] = 1
|
|
cross_mask[p1 : p2, t1 : t2] = 1
|
|
|
|
emb_cum_idx += emb_size
|
|
|
|
for i, (m, emb) in enumerate(zip(lin_masks, reg_embeds)):
|
|
# mask text
|
|
for j in range(1, n_regs - i):
|
|
t1, t2 = emb_cum_idx + j * emb_size, emb_cum_idx + (j+1) * emb_size
|
|
p1, p2 = emb_cum_idx, emb_cum_idx + emb_size
|
|
print(t1, t2, p1, p2)
|
|
|
|
cross_mask[t1 : t2, p1 : p2] = 1
|
|
cross_mask[p1 : p2, t1 : t2] = 1
|
|
|
|
scale = m.sum() / Nx
|
|
print('m: ', m.shape, scale)
|
|
if scale > 1e-5:
|
|
q_scale[emb_cum_idx : emb_cum_idx+emb_size] = 1 / scale
|
|
k_scale[emb_cum_idx : emb_cum_idx+emb_size] = 1 / scale
|
|
|
|
# m (4096) -> (N_text * 256 + 4096)
|
|
m = torch.cat([torch.ones(emb_size * (n_regs+1)), m])
|
|
print(m.shape)
|
|
|
|
mb = m > 0.5
|
|
cross_mask[~mb, emb_cum_idx : emb_cum_idx + emb_size] = 1
|
|
cross_mask[emb_cum_idx : emb_cum_idx + emb_size, ~mb] = 1
|
|
emb_cum_idx += emb_size
|
|
|
|
# Image Self-Attention attention between different areas blocking
|
|
# Calculate pairwise masks between different areas with the kronecker product
|
|
for i in range(n_regs):
|
|
for j in range(i+1, n_regs):
|
|
# We need to calculate two kr.prod for preserving the symmetry of the matrix
|
|
kron1 = torch.kron(lin_masks[i].unsqueeze(0), lin_masks[j].unsqueeze(-1))
|
|
kron2 = torch.kron(lin_masks[j].unsqueeze(0), lin_masks[i].unsqueeze(-1))
|
|
# cross_mask[emb_cum_idx:, emb_cum_idx:] += kron1 + kron2
|
|
|
|
# We need to select interesecting regions and set the rows and columns which are intersecting to 0
|
|
|
|
# Get the intersecting regions
|
|
intersect_idx = torch.logical_and(lin_masks[i] > 0.5, lin_masks[j] > 0.5)
|
|
# Set the intersecting regions to 0
|
|
kron_sum = kron1 + kron2
|
|
kron_sum[intersect_idx, :] = 0
|
|
kron_sum[:, intersect_idx] = 0
|
|
|
|
# kron_sum[intersect_idx, intersect_idx] = 0
|
|
|
|
# Add the kronecker product to the cross mask
|
|
cross_mask[emb_cum_idx:, emb_cum_idx:] += kron_sum
|
|
|
|
# Clean up the diagonal
|
|
cross_mask.fill_diagonal_(0)
|
|
|
|
q_scale = q_scale.reshape(1, 1, -1, 1).cuda()
|
|
k_scale = k_scale.reshape(1, 1, -1, 1).cuda()
|
|
|
|
return cross_mask, q_scale, k_scale
|
|
|
|
|
|
test_payload = {
|
|
'prompt': {
|
|
'positive': 'An italian cafe',
|
|
'width': 1024,
|
|
'height': 1024,
|
|
'bboxes': [
|
|
{
|
|
'caption': 'An asian man with sombrero',
|
|
'x': 100,
|
|
'y': 200,
|
|
'width': 300,
|
|
'height': 700,
|
|
},
|
|
{
|
|
'caption': 'A redhair sexual woman',
|
|
'x': 500,
|
|
'y': 200,
|
|
'width': 300,
|
|
'height': 700,
|
|
}
|
|
],
|
|
}
|
|
}
|
|
|
|
|
|
def process_payload(payload):
|
|
bboxes = payload['prompt']['bboxes']
|
|
masks = []
|
|
subprompts = []
|
|
|
|
for i, bbox in enumerate(bboxes):
|
|
mask = Image.new('L', (payload['prompt']['width'], payload['prompt']['height']), 0)
|
|
mask_arr = np.array(mask)
|
|
|
|
# Draw the bounding box
|
|
mask_arr[bbox['y']:bbox['y']+bbox['height'], bbox['x']:bbox['x']+bbox['width']] = 255
|
|
mask = Image.fromarray(mask_arr)
|
|
|
|
# Debug save the mask
|
|
mask.save(f'mask_{i}.png')
|
|
|
|
masks.append(mask)
|
|
subprompts.append(bbox['caption'])
|
|
|
|
return masks, subprompts
|
|
|
|
|
|
def generate_test_mask(masks, height, width):
|
|
hH, hW = int(height) // 16, int(width) // 16
|
|
print(height, width, '->', hH, hW)
|
|
|
|
lin_masks = []
|
|
for mask in masks:
|
|
mask = mask.convert('L')
|
|
mask = torch.tensor(np.array(mask)).unsqueeze(0).unsqueeze(0) / 255
|
|
# Linearize mask
|
|
mask = torch.nn.functional.interpolate(mask, (hH, hW), mode='nearest-exact').flatten()
|
|
lin_masks.append(mask)
|
|
|
|
return lin_masks, hH, hW
|
|
|
|
def generate_region_mask(region, width, height):
|
|
if region.get('bbox') is not None:
|
|
x1, y1, x2, y2 = region['bbox']
|
|
mask = Image.new('L', (width, height), 0)
|
|
mask_arr = np.array(mask)
|
|
|
|
print(f'Generating masks with {width}x{height} and [{x1}, {y1}, {x2}, {y2}]')
|
|
|
|
# Draw the bounding box
|
|
mask_arr[int(y1*height):int(y2*height), int(x1*width):int(x2*width)] = 255
|
|
mask = Image.fromarray(mask_arr)
|
|
|
|
return mask
|
|
elif region.get('mask') is not None:
|
|
mask = region['mask'] # ComfyUI mask is tensor (bs x height x width)
|
|
print('MASK: ', mask)
|
|
mask = mask[0].cpu().numpy()
|
|
mask = (mask * 255).astype(np.uint8)
|
|
mask = Image.fromarray(mask)
|
|
mask = mask.resize((width, height))
|
|
|
|
return mask
|
|
else:
|
|
raise Exception('Unknown region type')
|
|
|
|
|
|
class RegionAttention:
|
|
RETURN_TYPES = ("MODEL", "CONDITIONING")
|
|
RETURN_NAMES = ("model", "condition")
|
|
FUNCTION = "go"
|
|
CATEGORY = "model_patches"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
# "clip": ("CLIP", {"tooltip": "The CLIP model used for encoding the text."}),
|
|
"condition": ("CONDITIONING",),
|
|
"samples": ("LATENT",),
|
|
"region1": ("REGION",),
|
|
"enabled": ("BOOLEAN", {"default": True},),
|
|
},
|
|
"optional": {
|
|
"region2": ("REGION",),
|
|
"region3": ("REGION",),
|
|
"region4": ("REGION",),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def go(cls, *, model: object, condition, samples, region1, enabled: bool,
|
|
region2=None, region3=None, region4=None):
|
|
print(f'Region attention Node enabled: {enabled}, model: {model}')
|
|
# masks, payload = process_payload(test_payload)
|
|
|
|
latent = samples['samples']
|
|
print('latent.shape', latent.shape)
|
|
bs_l, n_ch, lH, lW = latent.shape
|
|
text_emb = condition[0][0].clone()
|
|
clip_emb = condition[0][1]['pooled_output'].clone()
|
|
bs, emb_size, emb_dim = text_emb.shape
|
|
iH, iW = lH * 8, lW * 8
|
|
|
|
subprompts_embeds, masks = [region1['condition'][0][0],], [generate_region_mask(region1, iW, iH),]
|
|
masks[-1].save(f'mask_1.png')
|
|
if region2 is not None:
|
|
print('append region2')
|
|
sub_emb2 = region2['condition'][0][0]
|
|
masks.append(generate_region_mask(region2, iW, iH))
|
|
subprompts_embeds.append(sub_emb2)
|
|
masks[-1].save(f'mask_2.png')
|
|
if region3 is not None:
|
|
print('append region3')
|
|
sub_emb3 = region3['condition'][0][0]
|
|
masks.append(generate_region_mask(region3, iW, iH))
|
|
subprompts_embeds.append(sub_emb3)
|
|
if region4 is not None:
|
|
print('append region4')
|
|
sub_emb4 = region4['condition'][0][0]
|
|
masks.append(generate_region_mask(region4, iW, iH))
|
|
subprompts_embeds.append(sub_emb4)
|
|
|
|
lin_masks, hH, hW = generate_test_mask(masks, lH * 8, lW * 8)
|
|
Nx = int(hH * hW)
|
|
emb_len = (len(subprompts_embeds) + 1) * emb_size
|
|
extended_condition = torch.cat([text_emb, *subprompts_embeds], dim=1) if enabled else text_emb
|
|
|
|
attn_mask, q_scale, k_scale = prepare_attention_mask(lin_masks, subprompts_embeds, Nx, emb_size, emb_len)
|
|
|
|
# Visualize and save the attention mask
|
|
# fig, ax = plt.subplots(1, 1, figsize=(10, 10))
|
|
# ax.imshow(attn_mask.cpu().numpy())
|
|
# plt.savefig('attention_mask_1.png')
|
|
|
|
# Pad mask for xformers to reduce allocations during inference
|
|
device = torch.device('cuda')
|
|
attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16
|
|
if attn_mask is not None:
|
|
print(f'Aplying attention masks: {attn_mask.shape}')
|
|
L, _ = attn_mask.shape
|
|
H = 24 # 24 heads for FLUX models
|
|
pad = 8 - L % 8
|
|
|
|
# print(f'Attention mask memory padded by: {pad}')
|
|
if pad != 8:
|
|
# TODO: take dtype from memory_management computational_type
|
|
mask_out = torch.empty([bs, H, L + pad, L + pad],
|
|
dtype=torch.bfloat16, device=device)
|
|
mask_out[:, :, :L, :L] = attn_mask
|
|
# print(f'Attention mask memory padded to: {mask_out.shape}')
|
|
attn_mask = mask_out[:, :, :L, :L]
|
|
else:
|
|
mask_out = torch.empty([bs, H, L, L],
|
|
dtype=torch.bfloat16, device=device)
|
|
mask_out[:, :, :, :] = attn_mask
|
|
attn_mask = mask_out
|
|
|
|
attn_mask_bool = attn_mask > 0.5
|
|
attn_mask.masked_fill_(attn_mask_bool, float('-inf'))
|
|
|
|
attn_mask_arg: Tensor = attn_mask if enabled else None
|
|
|
|
# if attn_mask_arg is not None:
|
|
# fig, ax = plt.subplots(1, 1, figsize=(10, 10))
|
|
# ax.imshow(attn_mask[0][0].float().cpu().numpy())
|
|
# plt.savefig('attention_mask_2.png')
|
|
|
|
def region_attention(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False):
|
|
print(q.shape, k.shape, v.shape)
|
|
|
|
res = orig_attention(q, k, v, heads, mask=attn_mask, attn_precision=attn_precision, skip_reshape=skip_reshape)
|
|
|
|
return res
|
|
|
|
# comfy_attention.optimized_attention = orig_attention if not enabled else region_attention
|
|
|
|
def override_attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
|
|
|
|
q, k = flux_math.apply_rope(q, k, pe)
|
|
|
|
heads = q.shape[1]
|
|
x = region_attention(q, k, v, heads, skip_reshape=True)
|
|
return x
|
|
|
|
override_attention = partial(masked_attention, attn_mask=attn_mask_arg)
|
|
|
|
flux_math.attention = override_attention
|
|
flux_layers.attention = override_attention
|
|
|
|
del condition
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
new_condition = [[
|
|
extended_condition,
|
|
{'pooled_output': clip_emb},
|
|
]]
|
|
|
|
return (model, new_condition)
|
|
|
|
|
|
class FluxRegionMask:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mask": ("MASK",),
|
|
"condition": ("CONDITIONING",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("REGION",)
|
|
FUNCTION = "create_region"
|
|
|
|
def create_region(self, mask, condition):
|
|
return ({
|
|
"condition": condition,
|
|
"mask": mask,
|
|
},)
|
|
|
|
|
|
class FluxRegionBBOX:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"x1": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0}),
|
|
"y1": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0}),
|
|
"x2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
"y2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
"condition": ("CONDITIONING",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("REGION",)
|
|
FUNCTION = "create_region"
|
|
|
|
def create_region(self, x1, y1, x2, y2, condition):
|
|
return ({
|
|
"condition": condition,
|
|
"bbox": [x1, y1, x2, y2],
|
|
},)
|
|
|
|
|
|
class CLIPDebug:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP", {"tooltip": "The CLIP model used for encoding the text."}),
|
|
"condition": ("CONDITIONING",),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
OUTPUT_TOOLTIPS = ("A conditioning containing the embedded text used to guide the diffusion model.",)
|
|
FUNCTION = "debug"
|
|
|
|
def debug(self, clip, condition):
|
|
# print(clip)
|
|
print('len(condition)', len(condition))
|
|
print('len(condition[0]', len(condition[0]))
|
|
print('type(condition[0][1])', type(condition[0][1]))
|
|
print('condition[0][0].shape', condition[0][0].shape)
|
|
print('list(condition[0][1].keys())', list(condition[0][1].keys()))
|
|
print("condition[0][1]['pooled_output'].shape", condition[0][1]['pooled_output'].shape)
|
|
|
|
return (condition,)
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
|
|
class RegionBbox:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image_width": ("INT", {"default": 512, "min": 64, "max": 2048}),
|
|
"image_height": ("INT", {"default": 512, "min": 64, "max": 2048}),
|
|
"x1": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0}),
|
|
"y1": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0}),
|
|
"x2": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0}),
|
|
"y2": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BBOX",)
|
|
FUNCTION = "create_bbox"
|
|
|
|
def create_bbox(self, image_width, image_height, x1, y1, x2, y2):
|
|
bbox = {
|
|
"x1": int(x1 * image_width),
|
|
"y1": int(y1 * image_height),
|
|
"x2": int(x2 * image_width),
|
|
"y2": int(y2 * image_height),
|
|
}
|
|
return (bbox,)
|
|
|
|
class VisualizeBBoxesNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"bboxes": ("BBOX",),
|
|
"color": ("COLOR", {"default": "#FF0000"}),
|
|
"width": ("INT", {"default": 2, "min": 1, "max": 10}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "visualize_bboxes"
|
|
|
|
def visualize_bboxes(self, image, bboxes, color, width):
|
|
# Convert the PyTorch tensor to a PIL Image
|
|
pil_image = Image.fromarray((image[0].permute(1, 2, 0) * 255).byte().cpu().numpy())
|
|
draw = ImageDraw.Draw(pil_image)
|
|
|
|
for bbox in bboxes:
|
|
draw.rectangle([bbox["x1"], bbox["y1"], bbox["x2"], bbox["y2"]], outline=color, width=width)
|
|
|
|
# Convert back to PyTorch tensor
|
|
tensor_image = torch.from_numpy(np.array(pil_image)).float() / 255.0
|
|
tensor_image = tensor_image.permute(2, 0, 1).unsqueeze(0)
|
|
|
|
return (tensor_image,)
|
|
|
|
class BBoxToMaskNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image_width": ("INT", {"default": 512, "min": 64, "max": 2048}),
|
|
"image_height": ("INT", {"default": 512, "min": 64, "max": 2048}),
|
|
"bboxes": ("BBOX",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "create_mask"
|
|
|
|
def create_mask(self, image_width, image_height, bbox):
|
|
mask = torch.zeros((1, image_height, image_width))
|
|
|
|
mask[0, bbox["y1"]:bbox["y2"], bbox["x1"]:bbox["x2"]] = 1.0
|
|
|
|
return (mask,)
|