330 lines
14 KiB
Python
330 lines
14 KiB
Python
import torch
|
|
from torch import Tensor
|
|
import torch.nn.functional as F
|
|
from typing import List, Dict, Optional, Tuple
|
|
from einops import rearrange
|
|
import comfy.model_management as model_management
|
|
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 xformers.ops import memory_efficient_attention as xattention
|
|
import numpy as np
|
|
from PIL import Image
|
|
from functools import partial
|
|
|
|
class FluxAttentionControl:
|
|
has_xformers = False
|
|
try:
|
|
from xformers.ops import memory_efficient_attention as xattention
|
|
has_xformers = True
|
|
except ImportError:
|
|
print("\n" + "="*70)
|
|
print("\033[94mControlAltAI-Nodes: This node requires xformers to function.\033[0m")
|
|
print("\033[33mPlease check \"xformers_instructions.txt\" in ComfyUI\\custom_nodes\\ControlAltAI-Nodes for how to install XFormers\033[0m")
|
|
print("="*70 + "\n")
|
|
|
|
def __init__(self):
|
|
self.original_attention = comfy_attention.optimized_attention
|
|
self.original_flux_attention = flux_math.attention
|
|
self.original_flux_layers_attention = flux_layers.attention
|
|
print("FluxAttentionControl initialized")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"condition": ("CONDITIONING",),
|
|
"latent_dimensions": ("LATENT",),
|
|
"region1": ("REGION",),
|
|
"number_of_regions": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 4,
|
|
"step": 1,
|
|
"display": "Number of Regions"
|
|
}),
|
|
"enabled": ("BOOLEAN", {
|
|
"default": True,
|
|
"display": "Enable Regional Control"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"region2": ("REGION",),
|
|
"region3": ("REGION",),
|
|
"region4": ("REGION",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CONDITIONING",)
|
|
RETURN_NAMES = ("model", "conditioning",)
|
|
FUNCTION = "apply_attention_control"
|
|
CATEGORY = "conditioning/flux"
|
|
|
|
def generate_region_mask(self, region: Dict, width: int, height: int) -> Image.Image:
|
|
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)
|
|
mask_arr[int(y1*height):int(y2*height), int(x1*width):int(x2*width)] = 255
|
|
print(f'Generating masks with {width}x{height} and [{x1}, {y1}, {x2}, {y2}]')
|
|
return Image.fromarray(mask_arr)
|
|
elif region.get('mask') is not None:
|
|
mask = region['mask'][0].cpu().numpy()
|
|
mask = (mask * 255).astype(np.uint8)
|
|
mask = Image.fromarray(mask)
|
|
return mask.resize((width, height))
|
|
else:
|
|
raise Exception('Unknown region type')
|
|
|
|
def generate_test_mask(self, masks: List[Image.Image], height: int, width: int):
|
|
hH, hW = int(height) // 16, int(width) // 16
|
|
print(f'{width} {height} -> {hW} {hH}')
|
|
|
|
lin_masks = []
|
|
for mask in masks:
|
|
mask = mask.convert('L')
|
|
mask = torch.tensor(np.array(mask)).unsqueeze(0).unsqueeze(0) / 255
|
|
mask = F.interpolate(mask, (hH, hW), mode='nearest-exact').flatten()
|
|
lin_masks.append(mask)
|
|
return lin_masks, hH, hW
|
|
|
|
def prepare_attention_mask(self, lin_masks: List[torch.Tensor], reg_embeds: List[torch.Tensor],
|
|
Nx: int, emb_size: int, emb_len: int):
|
|
"""Prepare attention mask for specific numbers of regions (1, 2, or 3)"""
|
|
total_len = emb_len + Nx
|
|
cross_mask = torch.zeros(total_len, total_len)
|
|
q_scale = torch.ones(total_len)
|
|
k_scale = torch.ones(total_len)
|
|
|
|
n_regs = len(lin_masks)
|
|
|
|
# Indices for embeddings
|
|
main_prompt_start = 0
|
|
main_prompt_end = emb_size
|
|
|
|
# Subprompt indices
|
|
subprompt_starts = [emb_size * (i + 1) for i in range(n_regs)]
|
|
subprompt_ends = [emb_size * (i + 2) for i in range(n_regs)]
|
|
|
|
# Handle different cases based on the number of regions
|
|
if n_regs == 1:
|
|
# For one region
|
|
sp_start = subprompt_starts[0]
|
|
sp_end = subprompt_ends[0]
|
|
|
|
# Block attention between main prompt and subprompt
|
|
cross_mask[sp_start:sp_end, main_prompt_start:main_prompt_end] = 1
|
|
cross_mask[main_prompt_start:main_prompt_end, sp_start:sp_end] = 1
|
|
|
|
# Scale based on region size
|
|
scale = lin_masks[0].sum() / Nx
|
|
if scale > 1e-5:
|
|
q_scale[sp_start:sp_end] = 1 / scale
|
|
k_scale[sp_start:sp_end] = 1 / scale
|
|
|
|
# Create mask including tokens and positions
|
|
m_with_tokens = torch.cat([torch.ones(emb_len), lin_masks[0]])
|
|
mb = m_with_tokens > 0.5
|
|
|
|
# Block attention between positions not in mask and subprompt
|
|
cross_mask[~mb, sp_start:sp_end] = 1
|
|
cross_mask[sp_start:sp_end, ~mb] = 1
|
|
|
|
elif n_regs == 2:
|
|
# For two regions
|
|
for i in range(2):
|
|
sp_start = subprompt_starts[i]
|
|
sp_end = subprompt_ends[i]
|
|
|
|
# Block attention between main prompt and subprompts
|
|
cross_mask[sp_start:sp_end, main_prompt_start:main_prompt_end] = 1
|
|
cross_mask[main_prompt_start:main_prompt_end, sp_start:sp_end] = 1
|
|
|
|
# Block attention between subprompts
|
|
other_sp_start = subprompt_starts[1 - i]
|
|
other_sp_end = subprompt_ends[1 - i]
|
|
cross_mask[sp_start:sp_end, other_sp_start:other_sp_end] = 1
|
|
cross_mask[other_sp_start:other_sp_end, sp_start:sp_end] = 1
|
|
|
|
# Scale based on region size
|
|
scale = lin_masks[i].sum() / Nx
|
|
if scale > 1e-5:
|
|
q_scale[sp_start:sp_end] = 1 / scale
|
|
k_scale[sp_start:sp_end] = 1 / scale
|
|
|
|
# Create mask including tokens and positions
|
|
m_with_tokens = torch.cat([torch.ones(emb_len), lin_masks[i]])
|
|
mb = m_with_tokens > 0.5
|
|
|
|
# Block attention between positions not in mask and subprompt
|
|
cross_mask[~mb, sp_start:sp_end] = 1
|
|
cross_mask[sp_start:sp_end, ~mb] = 1
|
|
|
|
# Handle intersection between regions
|
|
intersect_idx = torch.logical_and(lin_masks[0] > 0.5, lin_masks[1] > 0.5)
|
|
if intersect_idx.any():
|
|
idx = intersect_idx.nonzero(as_tuple=True)[0] + emb_len
|
|
cross_mask[idx[:, None], idx[None, :]] = 1
|
|
|
|
elif n_regs == 3:
|
|
# For three regions
|
|
for i in range(3):
|
|
sp_start = subprompt_starts[i]
|
|
sp_end = subprompt_ends[i]
|
|
|
|
# Block attention between main prompt and subprompts
|
|
cross_mask[sp_start:sp_end, main_prompt_start:main_prompt_end] = 1
|
|
cross_mask[main_prompt_start:main_prompt_end, sp_start:sp_end] = 1
|
|
|
|
# Block attention between subprompts
|
|
for j in range(3):
|
|
if i != j:
|
|
other_sp_start = subprompt_starts[j]
|
|
other_sp_end = subprompt_ends[j]
|
|
cross_mask[sp_start:sp_end, other_sp_start:other_sp_end] = 1
|
|
cross_mask[other_sp_start:other_sp_end, sp_start:sp_end] = 1
|
|
|
|
# Scale based on region size
|
|
scale = lin_masks[i].sum() / Nx
|
|
if scale > 1e-5:
|
|
q_scale[sp_start:sp_end] = 1 / scale
|
|
k_scale[sp_start:sp_end] = 1 / scale
|
|
|
|
# Create mask including tokens and positions
|
|
m_with_tokens = torch.cat([torch.ones(emb_len), lin_masks[i]])
|
|
mb = m_with_tokens > 0.5
|
|
|
|
# Block attention between positions not in mask and subprompt
|
|
cross_mask[~mb, sp_start:sp_end] = 1
|
|
cross_mask[sp_start:sp_end, ~mb] = 1
|
|
|
|
# Handle intersections between regions
|
|
for i in range(3):
|
|
for j in range(i + 1, 3):
|
|
intersect_idx = torch.logical_and(lin_masks[i] > 0.5, lin_masks[j] > 0.5)
|
|
if intersect_idx.any():
|
|
idx = intersect_idx.nonzero(as_tuple=True)[0] + emb_len
|
|
cross_mask[idx[:, None], idx[None, :]] = 1
|
|
else:
|
|
raise ValueError("Number of regions must be 1, 2, or 3.")
|
|
|
|
# Ensure self-attention is allowed
|
|
cross_mask.fill_diagonal_(0)
|
|
|
|
# Prepare scales for GPU
|
|
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
|
|
|
|
|
|
def xformers_attention(self, q: Tensor, k: Tensor, v: Tensor, pe: Tensor,
|
|
attn_mask: Optional[Tensor] = None) -> Tensor:
|
|
q, k = flux_math.apply_rope(q, k, pe)
|
|
q = rearrange(q, "B H L D -> B L H D")
|
|
k = rearrange(k, "B H L D -> B L H D")
|
|
v = rearrange(v, "B H L D -> B L H D")
|
|
if attn_mask is not None:
|
|
x = xattention(q, k, v, attn_bias=attn_mask)
|
|
else:
|
|
x = xattention(q, k, v)
|
|
x = rearrange(x, "B L H D -> B L (H D)")
|
|
return x
|
|
|
|
def apply_attention_control(self,
|
|
model: object,
|
|
condition: List,
|
|
latent_dimensions: Dict,
|
|
region1: Dict,
|
|
number_of_regions: int,
|
|
enabled: bool,
|
|
region2: Optional[Dict] = None,
|
|
region3: Optional[Dict] = None,
|
|
region4: Optional[Dict] = None):
|
|
if not enabled:
|
|
# Restore original attention functions
|
|
flux_math.attention = self.original_flux_attention
|
|
flux_layers.attention = self.original_flux_layers_attention
|
|
print("Regional control disabled. Restored original attention functions.")
|
|
return (model, condition)
|
|
|
|
print(f'Region attention Node enabled: {enabled}, regions: {number_of_regions}')
|
|
|
|
# Extract dimensions and embeddings
|
|
latent = latent_dimensions["samples"]
|
|
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
|
|
|
|
# Process active regions
|
|
subprompts_embeds = []
|
|
masks = []
|
|
|
|
# Collect regions
|
|
regions = [region1, region2, region3, region4]
|
|
for idx, region in enumerate(regions[:number_of_regions]):
|
|
if region is not None and region.get('conditioning') is not None:
|
|
subprompts_embeds.append(region['conditioning'][0][0])
|
|
masks.append(self.generate_region_mask(region, iW, iH))
|
|
else:
|
|
print(f"Region {idx+1} is None or has no conditioning")
|
|
|
|
if not subprompts_embeds:
|
|
print("No active regions with conditioning found.")
|
|
# Restore original attention functions
|
|
flux_math.attention = self.original_flux_attention
|
|
flux_layers.attention = self.original_flux_layers_attention
|
|
return (model, condition)
|
|
|
|
n_regs = len(subprompts_embeds)
|
|
|
|
# Generate attention components
|
|
lin_masks, hH, hW = self.generate_test_mask(masks, iH, iW)
|
|
Nx = int(hH * hW)
|
|
emb_len = emb_size * (n_regs + 1) # +1 for main prompt
|
|
|
|
# Create attention mask
|
|
attn_mask, q_scale, k_scale = self.prepare_attention_mask(
|
|
lin_masks, subprompts_embeds, Nx, emb_size, emb_len)
|
|
|
|
# Format for xFormers
|
|
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'Applying attention masks: {attn_mask.shape}')
|
|
L = attn_mask.shape[0]
|
|
H = 24 # Number of heads in FLUX model
|
|
pad = (8 - L % 8) % 8 # Ensure pad is between 0 and 7
|
|
pad_L = L + pad
|
|
mask_out = torch.zeros([bs, H, pad_L, pad_L], dtype=attn_dtype, device=device)
|
|
mask_out[:, :, :L, :L] = attn_mask.to(device, dtype=attn_dtype)
|
|
attn_mask = mask_out[:, :, :pad_L, :pad_L]
|
|
|
|
# Prepare final mask
|
|
attn_mask_bool = attn_mask > 0.5
|
|
attn_mask.masked_fill_(attn_mask_bool, float('-inf'))
|
|
|
|
# Override attention
|
|
attn_mask_arg = attn_mask if enabled else None
|
|
override_attention = partial(self.xformers_attention, attn_mask=attn_mask_arg)
|
|
flux_math.attention = override_attention
|
|
flux_layers.attention = override_attention
|
|
|
|
# Create extended conditioning
|
|
extended_condition = torch.cat([text_emb] + subprompts_embeds, dim=1)
|
|
|
|
return (model, [[extended_condition, {'pooled_output': clip_emb}]])
|
|
|
|
# Node class mappings
|
|
NODE_CLASS_MAPPINGS = {
|
|
"FluxAttentionControl": FluxAttentionControl
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"FluxAttentionControl": "Flux Attention Control"
|
|
}
|