From 52da54243e9591da84f5f05f52ba78bcf7c7898f Mon Sep 17 00:00:00 2001 From: ControlAltAI <6172927+gseth@users.noreply.github.com> Date: Sun, 8 Dec 2024 22:02:20 +0530 Subject: [PATCH] Fixed Import Error. Now Error will only occur when using Flux Attention, if no XFormers is installed --- __init__.py | 10 ------ flux_attention_control_node.py | 56 ++++++++++++++++++---------------- 2 files changed, 30 insertions(+), 36 deletions(-) diff --git a/__init__.py b/__init__.py index fb7eeb0..863777d 100644 --- a/__init__.py +++ b/__init__.py @@ -1,14 +1,4 @@ print("\n\033[32mInitializing ControlAltAI Nodes\033[0m") # Fixed green reset -try: - import xformers - print("\033[94mControlAltAI: xformers found!\033[0m") # Light blue -except ImportError: - print("\n" + "="*70) - print("\033[94mCONTROLALTAI-NODES: XFORMERS REQUIRED\033[0m") # Light blue - print("\033[33mSome nodes in this custom node pack require xformers to function properly.") # Start orange - print("\nPlease check \"import_error_solution.txt\" in ComfyUI\\custom_nodes\\ControlAltAI-Nodes for how to install XFormers\033[0m") # End orange - print("="*70 + "\n") - raise ImportError("xformers is required for ControlAltAI Nodes") from .flux_resolution_cal_node import FluxResolutionNode from .flux_sampler_node import FluxSampler diff --git a/flux_attention_control_node.py b/flux_attention_control_node.py index df3bf17..16504d9 100644 --- a/flux_attention_control_node.py +++ b/flux_attention_control_node.py @@ -7,18 +7,33 @@ 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, ImageFilter, ImageDraw from functools import partial class FluxAttentionControl: + has_xformers = False + xattention = None + 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") + def check_xformers(self): + if not self.has_xformers: + try: + from xformers.ops import memory_efficient_attention + self.__class__.xattention = memory_efficient_attention + self.__class__.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") + raise ImportError("xformers is required for this node") + @classmethod def INPUT_TYPES(cls): return { @@ -109,7 +124,6 @@ class FluxAttentionControl: return lin_masks, hH, hW def prepare_attention_mask(self, lin_masks: List[torch.Tensor], region_strengths: List[float], Nx: int, emb_size: int, emb_len: int): - """Prepare attention mask for three regions with per-region strengths.""" total_len = emb_len + Nx n_regs = len(lin_masks) @@ -127,7 +141,7 @@ class FluxAttentionControl: subprompt_ends = [emb_size * (i + 2) for i in range(n_regs)] # Initialize position masks - position_masks = torch.stack(lin_masks) # Shape: [n_regs, Nx] + position_masks = torch.stack(lin_masks) # Normalize masks so that overlapping areas sum to 1 position_masks_sum = position_masks.sum(dim=0) @@ -146,7 +160,7 @@ class FluxAttentionControl: # Create mask including tokens and positions m_with_tokens = torch.cat([torch.ones(emb_len), mask_i]) - mb = m_with_tokens > 0.0 # Include positions where mask > 0 + mb = m_with_tokens > 0.0 # Block attention between positions not in mask and subprompt cross_mask[~mb, sp_start:sp_end] = 1 @@ -176,14 +190,17 @@ class FluxAttentionControl: def xformers_attention(self, q: Tensor, k: Tensor, v: Tensor, pe: Tensor, attn_mask: Optional[Tensor] = None) -> Tensor: + if not self.has_xformers: + raise ImportError("xformers is required for this node") + 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) + x = self.xattention(q, k, v, attn_bias=attn_mask) else: - x = xattention(q, k, v) + x = self.xattention(q, k, v) x = rearrange(x, "B L H D -> B L (H D)") return x @@ -199,8 +216,10 @@ class FluxAttentionControl: feather_radius2: Optional[float] = 0.0, region3: Optional[Dict] = None, feather_radius3: Optional[float] = 0.0): + + self.check_xformers() # Only check xformers when node is used - # Extract dimensions and embeddings first (moved before enabled check) + # Extract dimensions and embeddings latent = latent_dimensions["samples"] bs_l, n_ch, lH, lW = latent.shape text_emb = condition[0][0].clone() @@ -213,18 +232,10 @@ class FluxAttentionControl: 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) # Return original condition when disabled + 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 = [] @@ -236,12 +247,10 @@ class FluxAttentionControl: for idx, region in enumerate(regions[:number_of_regions]): if region is not None and region.get('conditioning') is not None: - # Get 'strength' from region or default to 1.0 strength = region.get('strength', 1.0) region_strengths.append(strength) subprompt_emb = region['conditioning'][0][0] subprompts_embeds.append(subprompt_emb) - # Use per-region feather_radius feather_radius = feather_radii[idx] if feather_radii[idx] is not None else 0.0 masks.append(self.generate_region_mask(region, iW, iH, feather_radius)) else: @@ -249,7 +258,6 @@ class FluxAttentionControl: 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) @@ -273,17 +281,13 @@ class FluxAttentionControl: print(f'Applying attention masks: torch.Size([{attn_mask.shape[0]}, {attn_mask.shape[1]}])') 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 = (8 - L % 8) % 8 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] + attn_mask = mask_out - # Prepare final mask - attn_mask_bool = attn_mask > 0.5 - attn_mask.masked_fill_(attn_mask_bool, float('-inf')) - - # Override attention + # Override attention functions 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