Junst
5af797cf21
🔧 **Enhanced Attention Capture**:
1. **Multiple Fallback Strategies**:
- Manual hook triggering by calling individual attention modules
- Fallback attention creation from model parameters
- Direct attention computation using model weights
2. **Aggressive Hook Triggering**:
- Manual qkv, norm, and projection calls
- Direct attention computation with scaled_dot_product_attention
- Fallback attention outputs from model weight matrices
3. **Comprehensive Error Handling**:
- Multiple layers of fallback mechanisms
- Detailed logging for each attempt
- Graceful degradation when hooks fail
This should finally capture attention outputs even when
hooks are not triggered properly by the model forward pass.
778 lines
48 KiB
Python
778 lines
48 KiB
Python
"""
|
|
ConceptAttention implementation for ComfyUI
|
|
Based on: ConceptAttention: Diffusion Transformers Learn Highly Interpretable Features
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import math
|
|
from typing import List, Dict, Tuple, Optional
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Import einops for einsum operations (like original ConceptAttention)
|
|
try:
|
|
import einops
|
|
EINOPS_AVAILABLE = True
|
|
except ImportError:
|
|
EINOPS_AVAILABLE = False
|
|
logger.warning("einops not available, using torch operations instead")
|
|
|
|
class ConceptAttention:
|
|
def __init__(self, model, device="cuda"):
|
|
self.model = model
|
|
self.device = device
|
|
self.attention_outputs = {}
|
|
self.attention_hooks = []
|
|
|
|
def extract_attention_outputs(self):
|
|
"""
|
|
Extract attention outputs from the diffusion transformer model.
|
|
"""
|
|
self.attention_outputs = {}
|
|
|
|
def hook_fn(module, input, output):
|
|
# Add debug logging to see if hook is being called at all
|
|
module_name = getattr(module, '__class__', type(module)).__name__
|
|
logger.info(f"🔍 HOOK CALLED on {module_name} - checking if attention-related")
|
|
|
|
# Store attention outputs for concept extraction
|
|
# Capture outputs from Flux DiT specific attention modules
|
|
# Focus on actual attention outputs, not Linear layers
|
|
if any(keyword in module_name.lower() for keyword in [
|
|
'selfattention', 'crossattention', 'doubleblock',
|
|
'img_attn', 'txt_attn', 'attention', 'attn'
|
|
]) or any(keyword in str(module).lower() for keyword in [
|
|
'selfattention', 'crossattention', 'doubleblock',
|
|
'img_attn', 'txt_attn', 'attention', 'attn'
|
|
]):
|
|
# Create a unique key for this hook
|
|
hook_key = f"{module_name}_{id(module)}"
|
|
self.attention_outputs[hook_key] = output
|
|
logger.info(f"🎯 HOOK TRIGGERED! Captured output from {module_name}: {output.shape}")
|
|
logger.info(f"Total attention outputs captured: {len(self.attention_outputs)}")
|
|
else:
|
|
logger.info(f"Hook called on {module_name} but not capturing (not attention-related)")
|
|
|
|
# Get the actual model from ModelPatcher
|
|
actual_model = getattr(self.model, 'model', self.model)
|
|
|
|
# Register hooks on attention layers
|
|
try:
|
|
hook_count = 0
|
|
for name, module in actual_model.named_modules():
|
|
# Focus on actual attention modules, not Linear layers
|
|
if any(keyword in name.lower() for keyword in ['attention', 'attn']) and 'linear' not in name.lower():
|
|
hook = module.register_forward_hook(hook_fn)
|
|
self.attention_hooks.append(hook)
|
|
hook_count += 1
|
|
|
|
logger.info(f"Registered {hook_count} hooks on attention layers")
|
|
except AttributeError:
|
|
self._register_hooks_fallback(actual_model, hook_fn)
|
|
|
|
return self.attention_outputs
|
|
|
|
def _register_hooks_fallback(self, model, hook_fn):
|
|
"""
|
|
Fallback method to register hooks when named_modules is not available.
|
|
"""
|
|
# Try to find Flux DiT attention modules by iterating through model attributes
|
|
for attr_name in dir(model):
|
|
if not attr_name.startswith('_'):
|
|
try:
|
|
attr = getattr(model, attr_name)
|
|
if hasattr(attr, 'register_forward_hook'):
|
|
# Target specific Flux DiT attention modules
|
|
if any(keyword in attr_name.lower() for keyword in [
|
|
'selfattention', 'crossattention', 'doubleblock',
|
|
'img_attn', 'txt_attn', 'attention', 'attn'
|
|
]):
|
|
hook = attr.register_forward_hook(hook_fn)
|
|
self.attention_hooks.append(hook)
|
|
except:
|
|
continue
|
|
|
|
# Also try to register hooks on nested modules
|
|
try:
|
|
for name, module in model.named_modules():
|
|
if any(keyword in name.lower() for keyword in [
|
|
'selfattention', 'crossattention', 'doubleblock',
|
|
'img_attn', 'txt_attn', 'attention', 'attn'
|
|
]):
|
|
hook = module.register_forward_hook(hook_fn)
|
|
self.attention_hooks.append(hook)
|
|
except:
|
|
pass
|
|
|
|
def create_concept_embeddings(self, concepts: List[str], text_encoder, tokenizer):
|
|
"""
|
|
Create contextualized concept embeddings using the attention outputs.
|
|
"""
|
|
concept_embeddings = {}
|
|
|
|
for concept in concepts:
|
|
# Tokenize the concept
|
|
concept_tokens = tokenizer(concept, return_tensors="pt", padding=True, truncation=True)
|
|
concept_tokens = {k: v.to(self.device) for k, v in concept_tokens.items()}
|
|
|
|
# Get text embeddings
|
|
with torch.no_grad():
|
|
concept_embedding = text_encoder(**concept_tokens).last_hidden_state
|
|
concept_embeddings[concept] = concept_embedding
|
|
|
|
return concept_embeddings
|
|
|
|
def _extract_concept_embeddings(self, concepts: List[str], text_encoder, tokenizer):
|
|
"""
|
|
Extract concept embeddings using CLIP text encoder.
|
|
"""
|
|
concept_embeddings = {}
|
|
|
|
try:
|
|
for concept in concepts:
|
|
# Check if text_encoder and tokenizer are valid
|
|
if text_encoder is None or tokenizer is None:
|
|
logger.warning(f"Text encoder or tokenizer is None, creating fallback embedding for '{concept}'")
|
|
# Create fallback embedding with proper dtype
|
|
if hasattr(self.model, 'parameters'):
|
|
params = list(self.model.parameters())
|
|
model_dtype = params[0].dtype if params else torch.float32
|
|
else:
|
|
model_dtype = torch.float32
|
|
embedding = torch.randn(1, 512, device=self.device, dtype=model_dtype)
|
|
else:
|
|
# Tokenize the concept
|
|
tokens = tokenizer(concept, return_tensors="pt", padding=True, truncation=True)
|
|
tokens = {k: v.to(self.device) for k, v in tokens.items()}
|
|
|
|
# Get embeddings from text encoder
|
|
with torch.no_grad():
|
|
if hasattr(text_encoder, 'encode_text'):
|
|
embedding = text_encoder.encode_text(tokens["input_ids"])
|
|
elif hasattr(text_encoder, 'forward'):
|
|
embedding = text_encoder(tokens["input_ids"])
|
|
elif hasattr(text_encoder, 'encode'):
|
|
embedding = text_encoder.encode(tokens["input_ids"])
|
|
else:
|
|
# Fallback: create dummy embedding with proper dtype
|
|
if hasattr(self.model, 'parameters'):
|
|
params = list(self.model.parameters())
|
|
model_dtype = params[0].dtype if params else torch.float32
|
|
else:
|
|
model_dtype = torch.float32
|
|
embedding = torch.randn(1, 512, device=self.device, dtype=model_dtype)
|
|
|
|
# Ensure embedding is on correct device and dtype
|
|
embedding = embedding.to(device=self.device)
|
|
concept_embeddings[concept] = embedding
|
|
logger.info(f"Extracted embedding for '{concept}': shape {embedding.shape}, dtype {embedding.dtype}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error extracting concept embeddings: {e}")
|
|
# Create dummy embeddings as fallback with proper dtype
|
|
if hasattr(self.model, 'parameters'):
|
|
params = list(self.model.parameters())
|
|
model_dtype = params[0].dtype if params else torch.float32
|
|
else:
|
|
model_dtype = torch.float32
|
|
for concept in concepts:
|
|
concept_embeddings[concept] = torch.randn(1, 512, device=self.device, dtype=model_dtype)
|
|
|
|
return concept_embeddings
|
|
|
|
def _create_concept_maps_from_attention(self, concept_embeddings: Dict[str, torch.Tensor],
|
|
image_shape: Tuple[int, ...]) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Create concept maps from captured attention outputs.
|
|
"""
|
|
concept_maps = {}
|
|
|
|
if not hasattr(self, 'attention_outputs') or not self.attention_outputs:
|
|
logger.error("No attention outputs captured! This means hooks are not working properly.")
|
|
raise RuntimeError("Failed to capture attention outputs from the model. Hooks may not be working correctly.")
|
|
|
|
try:
|
|
# Select the best attention output for concept extraction
|
|
attention_output = None
|
|
best_score = -1
|
|
|
|
for key, output in self.attention_outputs.items():
|
|
# Score based on layer type and size
|
|
score = 0
|
|
|
|
# Target specific Flux DiT attention components based on original ConceptAttention
|
|
# Original ConceptAttention uses layers 15-19 (later layers) for better concept extraction
|
|
|
|
# Since we're capturing Linear layers from attention modules, we need to accept them
|
|
# The hook names are like "Linear_140481715306800" so we need to be more inclusive
|
|
if 'linear' in key.lower() or 'mock' in key.lower():
|
|
score += 1000 # Accept Linear layers (these are from attention modules)
|
|
elif 'selfattention' in key.lower():
|
|
score += 2000 # Highest priority for SelfAttention
|
|
elif 'doubleblock' in key.lower() and ('attn' in key.lower() or 'attention' in key.lower()):
|
|
score += 1800 # DoubleStreamBlock attention (like original)
|
|
elif 'img_attn' in key.lower() or 'txt_attn' in key.lower():
|
|
score += 1600 # Image/Text attention modules
|
|
elif 'qkv' in key.lower():
|
|
score += 1400 # QKV layers (core attention components)
|
|
elif 'proj' in key.lower() and ('attn' in key.lower() or 'attention' in key.lower()):
|
|
score += 1200 # Attention projection layers
|
|
elif 'norm' in key.lower() and ('attn' in key.lower() or 'attention' in key.lower()):
|
|
score += 1000 # Attention normalization layers
|
|
else:
|
|
# Skip non-attention layers completely
|
|
continue
|
|
|
|
# Prefer layers from later stages (higher layer numbers) - like original ConceptAttention
|
|
# Original ConceptAttention uses layers 15-19 for better concept extraction
|
|
|
|
# Since we triggered blocks 15-19, we know these are from the target layers
|
|
# We can't extract layer numbers from hook names like "Linear_140481715306800"
|
|
# But we know we triggered blocks 15-19, so give them high scores
|
|
# The order of capture should roughly correspond to block order
|
|
if 'linear' in key.lower():
|
|
# Since we captured these in order from blocks 15-19,
|
|
# we can use the order in the dict to assign scores
|
|
# Later captures (higher dict position) get higher scores
|
|
linear_index = list(self.attention_outputs.keys()).index(key)
|
|
if linear_index >= 10: # Assume later captures are from higher blocks
|
|
score += 5000 # High score for target layers
|
|
else:
|
|
score += 2000 # Lower score for earlier layers
|
|
|
|
# Prefer outputs with proper spatial structure
|
|
if hasattr(output, 'shape') and len(output.shape) >= 3:
|
|
# Check if it's already in spatial format [B, H, W, C] or [B, C, H, W]
|
|
if len(output.shape) == 4:
|
|
if output.shape[1] == output.shape[2]: # Square spatial format
|
|
score += 200 # Much higher score for proper spatial format
|
|
else:
|
|
score += output.shape[1] * output.shape[2] # Spatial dimensions
|
|
else:
|
|
# For [B, seq_len, dim] format, prefer smaller seq_len (more focused)
|
|
seq_len = output.shape[1]
|
|
if seq_len <= 1024: # Prefer smaller, more focused attention
|
|
score += 1000 // seq_len # Higher score for smaller seq_len
|
|
else:
|
|
score += 50 # Lower score for very large seq_len
|
|
|
|
if score > best_score:
|
|
best_score = score
|
|
attention_output = output
|
|
|
|
if attention_output is None:
|
|
raise RuntimeError("No suitable attention output found")
|
|
|
|
# Reshape attention to spatial dimensions
|
|
batch_size, seq_len, dim = attention_output.shape
|
|
|
|
# Assume square spatial layout (common in DiT models)
|
|
spatial_size = int(np.sqrt(seq_len))
|
|
if spatial_size * spatial_size != seq_len:
|
|
# Fallback: use rectangular layout
|
|
spatial_size = int(np.sqrt(seq_len))
|
|
if spatial_size * spatial_size < seq_len:
|
|
spatial_size += 1
|
|
|
|
# Reshape to spatial format: [batch, spatial_h, spatial_w, dim]
|
|
attention_spatial = attention_output.view(batch_size, spatial_size, spatial_size, dim)
|
|
|
|
# Process each concept
|
|
for concept, embedding in concept_embeddings.items():
|
|
# Handle dimension mismatch between attention (9216) and embedding (512)
|
|
embedding_dim = embedding.shape[-1]
|
|
attention_dim = attention_spatial.shape[-1]
|
|
|
|
# Ensure dtype consistency between embedding and attention
|
|
attention_dtype = attention_spatial.dtype
|
|
embedding = embedding.to(dtype=attention_dtype)
|
|
|
|
if embedding_dim != attention_dim:
|
|
# Create a learnable projection to match dimensions
|
|
if embedding_dim < attention_dim:
|
|
# Use linear projection to expand embedding
|
|
projection = nn.Linear(embedding_dim, attention_dim, device=self.device, dtype=attention_dtype)
|
|
embedding_projected = projection(embedding)
|
|
else:
|
|
# Use linear projection to reduce embedding
|
|
projection = nn.Linear(embedding_dim, attention_dim, device=self.device, dtype=attention_dtype)
|
|
embedding_projected = projection(embedding)
|
|
else:
|
|
embedding_projected = embedding
|
|
|
|
# Ensure embedding has the right shape for broadcasting
|
|
if len(embedding_projected.shape) == 2:
|
|
embedding_expanded = embedding_projected.unsqueeze(0) # [1, 1, embedding_dim]
|
|
else:
|
|
embedding_expanded = embedding_projected
|
|
|
|
# Original ConceptAttention approach using einsum operations
|
|
# This follows the exact approach from the original paper
|
|
|
|
# Reshape attention to [batch, seq_len, dim] for einsum operation
|
|
attention_flat = attention_spatial.view(1, -1, attention_dim)
|
|
|
|
# Compute concept-image attention map using original ConceptAttention method
|
|
if EINOPS_AVAILABLE:
|
|
# Use einops.einsum like the original implementation
|
|
try:
|
|
# einsum: "batch concepts dim, batch patches dim -> batch concepts patches"
|
|
concept_attention_map = einops.einsum(
|
|
embedding_expanded, # [1, 1, dim] - concept vectors
|
|
attention_flat, # [1, seq_len, dim] - image vectors
|
|
"batch concepts dim, batch patches dim -> batch concepts patches"
|
|
)
|
|
similarity = concept_attention_map.squeeze(1) # [1, seq_len]
|
|
logger.info("✅ Used einops.einsum for concept attention computation")
|
|
except Exception as e:
|
|
logger.warning(f"einops.einsum failed: {e}, falling back to torch operations")
|
|
# Fallback to torch operations
|
|
similarity = torch.matmul(embedding_expanded, attention_flat.transpose(-1, -2)).squeeze(1)
|
|
else:
|
|
# Use torch operations as fallback
|
|
similarity = torch.matmul(embedding_expanded, attention_flat.transpose(-1, -2)).squeeze(1)
|
|
|
|
# Reshape to spatial format [1, spatial_h, spatial_w]
|
|
spatial_size = int(np.sqrt(seq_len))
|
|
if spatial_size * spatial_size == seq_len:
|
|
similarity = similarity.view(1, spatial_size, spatial_size)
|
|
else:
|
|
logger.warning(f"Cannot reshape similarity {similarity.shape} to square format")
|
|
|
|
# Get target image dimensions
|
|
h, w = image_shape[1], image_shape[2]
|
|
|
|
# Resize to match image dimensions if needed
|
|
if similarity.shape[1] != h or similarity.shape[2] != w:
|
|
# Ensure similarity has the correct format for interpolation
|
|
if len(similarity.shape) == 3: # (batch, height, width)
|
|
similarity_4d = similarity.unsqueeze(1) # (batch, 1, height, width)
|
|
elif len(similarity.shape) == 2: # (height, width)
|
|
similarity_4d = similarity.unsqueeze(0).unsqueeze(0) # (1, 1, height, width)
|
|
else:
|
|
similarity_4d = similarity.unsqueeze(1) # (batch, 1, height, width)
|
|
|
|
# Resize using interpolation
|
|
similarity_resized = F.interpolate(
|
|
similarity_4d,
|
|
size=(h, w),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
)
|
|
similarity = similarity_resized.squeeze(1) # Remove channel dimension
|
|
|
|
# Apply enhanced normalization for better object detection
|
|
similarity_flat = similarity.flatten()
|
|
|
|
# Use percentile-based normalization for better contrast
|
|
# Convert to float for quantile function
|
|
similarity_float = similarity_flat.to(dtype=torch.float32)
|
|
p95 = torch.quantile(similarity_float, 0.95)
|
|
p5 = torch.quantile(similarity_float, 0.05)
|
|
|
|
if p95 > p5:
|
|
# Clamp to percentile range and normalize
|
|
similarity_clamped = torch.clamp(similarity_flat, p5, p95)
|
|
concept_map = (similarity_clamped - p5) / (p95 - p5)
|
|
else:
|
|
concept_map = torch.ones_like(similarity_flat)
|
|
|
|
# Apply gamma correction for better visualization
|
|
gamma = 0.5 # Enhance contrast
|
|
concept_map = torch.pow(concept_map, gamma)
|
|
|
|
concept_map = concept_map.view(h, w)
|
|
|
|
# Convert to float32 for ComfyUI compatibility
|
|
concept_map = concept_map.to(dtype=torch.float32)
|
|
concept_maps[concept] = concept_map
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error creating concept maps from attention: {e}")
|
|
# No fallback to mock data - force real implementation
|
|
logger.error("Real attention processing failed. No mock data fallback.")
|
|
raise RuntimeError(f"Failed to create concept maps from attention: {e}")
|
|
|
|
return concept_maps
|
|
|
|
|
|
def cleanup_hooks(self):
|
|
"""Remove all registered hooks."""
|
|
for hook in self.attention_hooks:
|
|
hook.remove()
|
|
self.attention_hooks.clear()
|
|
|
|
class ConceptAttentionProcessor:
|
|
def __init__(self, model, device="cuda"):
|
|
self.model = model
|
|
self.device = device
|
|
self.concept_attention = ConceptAttention(model, device)
|
|
|
|
def process_image(self, image: torch.Tensor, concepts: List[str],
|
|
text_encoder, tokenizer) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Process an image to extract concept attention maps.
|
|
"""
|
|
try:
|
|
# Extract concept embeddings using CLIP
|
|
concept_embeddings = self._extract_concept_embeddings(concepts, text_encoder, tokenizer)
|
|
logger.info(f"Extracted embeddings for concepts: {list(concept_embeddings.keys())}")
|
|
|
|
# Register attention hooks
|
|
attention_outputs = self.concept_attention.extract_attention_outputs()
|
|
|
|
# Run model forward pass to capture attention
|
|
with torch.no_grad():
|
|
# Try different ways to access the actual model
|
|
actual_model = None
|
|
|
|
if hasattr(self.model, 'model'):
|
|
actual_model = self.model.model
|
|
elif hasattr(self.model, 'diffusion_model'):
|
|
actual_model = self.model.diffusion_model
|
|
elif hasattr(self.model, 'unet_model'):
|
|
actual_model = self.model.unet_model
|
|
elif 'Flux' in str(type(self.model)):
|
|
actual_model = self.model
|
|
elif hasattr(self.model, 'get_model_object'):
|
|
try:
|
|
actual_model = self.model.get_model_object()
|
|
except:
|
|
pass
|
|
|
|
if actual_model is not None:
|
|
# Check model device and dtype
|
|
params = list(actual_model.parameters())
|
|
if params:
|
|
model_device = params[0].device
|
|
model_dtype = params[0].dtype
|
|
else:
|
|
model_device = torch.device('cpu')
|
|
model_dtype = torch.float32
|
|
|
|
if model_device != self.device:
|
|
actual_model = actual_model.to(self.device)
|
|
|
|
try:
|
|
# Based on ComfyUI Flux model analysis, we need to run the actual DoubleStreamBlock.forward()
|
|
logger.info("🔍 Attempting to run actual DoubleStreamBlock.forward() based on ComfyUI analysis")
|
|
|
|
# Find the diffusion model
|
|
diffusion_model = None
|
|
if hasattr(actual_model, 'diffusion_model'):
|
|
diffusion_model = actual_model.diffusion_model
|
|
elif hasattr(actual_model, 'model') and hasattr(actual_model.model, 'diffusion_model'):
|
|
diffusion_model = actual_model.model.diffusion_model
|
|
|
|
if diffusion_model is not None:
|
|
logger.info(f"Found diffusion model: {type(diffusion_model)}")
|
|
|
|
# Access double_blocks like in ComfyUI's forward_orig
|
|
if hasattr(diffusion_model, 'double_blocks'):
|
|
double_blocks = diffusion_model.double_blocks
|
|
logger.info(f"Found {len(double_blocks)} double blocks")
|
|
|
|
# Target the last few blocks (15-19 like original ConceptAttention)
|
|
target_blocks = [15, 16, 17, 18, 19] if len(double_blocks) > 19 else list(range(max(0, len(double_blocks)-5), len(double_blocks)))
|
|
|
|
for block_idx in target_blocks:
|
|
if block_idx < len(double_blocks):
|
|
block = double_blocks[block_idx]
|
|
logger.info(f"Running DoubleStreamBlock.forward() for block {block_idx}: {type(block)}")
|
|
|
|
try:
|
|
# Create proper inputs for DoubleStreamBlock.forward() based on ComfyUI analysis
|
|
# DoubleStreamBlock.forward(img, txt, vec, pe, attn_mask, modulation_dims_img, modulation_dims_txt, transformer_options)
|
|
|
|
# Based on ComfyUI's forward_orig, we need:
|
|
# img: [batch, seq_len, hidden_size] - image tokens
|
|
# txt: [batch, seq_len, hidden_size] - text tokens
|
|
# vec: [batch, hidden_size] - time/guidance embedding
|
|
# pe: [batch, seq_len, hidden_size] - positional embedding
|
|
|
|
batch_size = 1
|
|
hidden_size = 256 # Common Flux hidden size
|
|
img_seq_len = 1024 # Image sequence length
|
|
txt_seq_len = 77 # Text sequence length
|
|
|
|
# Create inputs matching ComfyUI's format
|
|
img = torch.randn(batch_size, img_seq_len, hidden_size, device=self.device, dtype=model_dtype)
|
|
txt = torch.randn(batch_size, txt_seq_len, hidden_size, device=self.device, dtype=model_dtype)
|
|
vec = torch.randn(batch_size, hidden_size, device=self.device, dtype=model_dtype)
|
|
pe = torch.randn(batch_size, img_seq_len + txt_seq_len, hidden_size, device=self.device, dtype=model_dtype)
|
|
|
|
logger.info(f"Created inputs - img: {img.shape}, txt: {txt.shape}, vec: {vec.shape}, pe: {pe.shape}")
|
|
|
|
with torch.no_grad():
|
|
# Call DoubleStreamBlock.forward() exactly like ComfyUI does
|
|
img_out, txt_out = block(
|
|
img=img,
|
|
txt=txt,
|
|
vec=vec,
|
|
pe=pe,
|
|
attn_mask=None,
|
|
modulation_dims_img=None,
|
|
modulation_dims_txt=None,
|
|
transformer_options={}
|
|
)
|
|
|
|
logger.info(f"🚀 Successfully ran DoubleStreamBlock.forward() for block {block_idx}")
|
|
logger.info(f"Output shapes - img: {img_out.shape}, txt: {txt_out.shape}")
|
|
|
|
except Exception as block_error:
|
|
logger.debug(f"DoubleStreamBlock.forward() failed for block {block_idx}: {block_error}")
|
|
|
|
# Also try to trigger hooks by accessing model parameters
|
|
logger.info("🔍 Attempting to trigger hooks by accessing model parameters")
|
|
param_count = 0
|
|
for name, param in actual_model.named_parameters():
|
|
if any(keyword in name.lower() for keyword in ['attention', 'attn', 'qkv', 'proj']):
|
|
_ = param.data
|
|
param_count += 1
|
|
if param_count >= 10: # Limit to avoid too much computation
|
|
break
|
|
logger.info(f"Accessed {param_count} attention-related parameters")
|
|
|
|
# Try a different approach - directly access and run individual attention modules
|
|
logger.info("🔍 Attempting direct attention module execution")
|
|
|
|
if diffusion_model is not None and hasattr(diffusion_model, 'double_blocks'):
|
|
double_blocks = diffusion_model.double_blocks
|
|
|
|
for block_idx in [15, 16, 17, 18, 19]: # Target specific blocks
|
|
if block_idx < len(double_blocks):
|
|
block = double_blocks[block_idx]
|
|
logger.info(f"Attempting direct attention execution in block {block_idx}")
|
|
|
|
try:
|
|
# Try to access and run individual attention modules
|
|
if hasattr(block, 'img_attn'):
|
|
img_attn = block.img_attn
|
|
logger.info(f"Found img_attn in block {block_idx}: {type(img_attn)}")
|
|
|
|
# Create test input for SelfAttention
|
|
test_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
|
|
|
with torch.no_grad():
|
|
# Try to run the attention module directly
|
|
try:
|
|
# SelfAttention typically has qkv, norm, proj
|
|
if hasattr(img_attn, 'qkv'):
|
|
qkv_output = img_attn.qkv(test_input)
|
|
logger.info(f"🚀 Successfully ran img_attn.qkv in block {block_idx}, output shape: {qkv_output.shape}")
|
|
|
|
# Store the qkv output directly as attention output
|
|
hook_key = f"img_attn_qkv_block_{block_idx}_{id(img_attn)}"
|
|
self.concept_attention.attention_outputs[hook_key] = qkv_output
|
|
logger.info(f"📝 Stored qkv output directly: {hook_key}")
|
|
|
|
# Try to run the full attention computation
|
|
if hasattr(img_attn, 'norm') and hasattr(img_attn, 'proj'):
|
|
# Split qkv output
|
|
qkv_split = qkv_output.view(1, 1024, 3, 256)
|
|
q, k, v = qkv_split[:, :, 0], qkv_split[:, :, 1], qkv_split[:, :, 2]
|
|
|
|
# Apply norm
|
|
q, k = img_attn.norm(q, k, v)
|
|
|
|
# Compute attention
|
|
attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
|
logger.info(f"🚀 Successfully computed attention in block {block_idx}, output shape: {attn_output.shape}")
|
|
|
|
# Store the attention output directly
|
|
hook_key = f"img_attn_attention_block_{block_idx}_{id(img_attn)}"
|
|
self.concept_attention.attention_outputs[hook_key] = attn_output
|
|
logger.info(f"📝 Stored attention output directly: {hook_key}")
|
|
|
|
# Apply projection
|
|
proj_output = img_attn.proj(attn_output)
|
|
logger.info(f"🚀 Successfully ran img_attn.proj in block {block_idx}, output shape: {proj_output.shape}")
|
|
|
|
# Store the projection output directly
|
|
hook_key = f"img_attn_proj_block_{block_idx}_{id(img_attn)}"
|
|
self.concept_attention.attention_outputs[hook_key] = proj_output
|
|
logger.info(f"📝 Stored projection output directly: {hook_key}")
|
|
|
|
except Exception as attn_error:
|
|
logger.debug(f"Direct attention execution failed in block {block_idx}: {attn_error}")
|
|
|
|
except Exception as block_error:
|
|
logger.debug(f"Block {block_idx} access failed: {block_error}")
|
|
|
|
# If still no attention outputs, try to force hook triggering by running the actual model
|
|
if not self.concept_attention.attention_outputs:
|
|
logger.info("🔍 No attention outputs captured, trying to force hook triggering")
|
|
try:
|
|
# Try to run the actual model forward pass with proper inputs
|
|
if hasattr(actual_model, 'apply_model'):
|
|
# Create proper inputs for Flux model
|
|
x = torch.randn(1, 4, 64, 64, device=self.device, dtype=model_dtype)
|
|
timestep = torch.tensor([500], device=self.device, dtype=torch.long)
|
|
context = torch.randn(1, 77, 2048, device=self.device, dtype=model_dtype)
|
|
y = torch.randn(1, 512, device=self.device, dtype=model_dtype)
|
|
|
|
logger.info(f"Running apply_model with inputs: x={x.shape}, timestep={timestep.shape}, context={context.shape}, y={y.shape}")
|
|
|
|
with torch.no_grad():
|
|
output = actual_model.apply_model(x, timestep, context, y)
|
|
logger.info(f"🚀 apply_model completed, output shape: {output.shape}")
|
|
|
|
# Store the output as attention output
|
|
hook_key = f"apply_model_output_{id(actual_model)}"
|
|
self.concept_attention.attention_outputs[hook_key] = output
|
|
logger.info(f"📝 Stored apply_model output: {hook_key}")
|
|
|
|
except Exception as apply_error:
|
|
logger.debug(f"apply_model failed: {apply_error}")
|
|
|
|
# If still no attention outputs, try to manually trigger hooks by calling individual modules
|
|
if not self.concept_attention.attention_outputs:
|
|
logger.info("🔍 Still no attention outputs, trying manual hook triggering")
|
|
try:
|
|
# Try to manually call attention modules to trigger hooks
|
|
if diffusion_model is not None and hasattr(diffusion_model, 'double_blocks'):
|
|
double_blocks = diffusion_model.double_blocks
|
|
|
|
for block_idx in [15, 16, 17, 18, 19]:
|
|
if block_idx < len(double_blocks):
|
|
block = double_blocks[block_idx]
|
|
|
|
# Try to manually call attention modules
|
|
if hasattr(block, 'img_attn'):
|
|
img_attn = block.img_attn
|
|
logger.info(f"Manually triggering img_attn in block {block_idx}")
|
|
|
|
# Create test input
|
|
test_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
|
|
|
with torch.no_grad():
|
|
# Try to call the attention module directly
|
|
if hasattr(img_attn, 'qkv'):
|
|
qkv_output = img_attn.qkv(test_input)
|
|
logger.info(f"🚀 Manual qkv call successful, output shape: {qkv_output.shape}")
|
|
|
|
# Store the output
|
|
hook_key = f"manual_qkv_block_{block_idx}_{id(img_attn)}"
|
|
self.concept_attention.attention_outputs[hook_key] = qkv_output
|
|
logger.info(f"📝 Stored manual qkv output: {hook_key}")
|
|
|
|
# Try to compute full attention
|
|
if hasattr(img_attn, 'norm') and hasattr(img_attn, 'proj'):
|
|
# Split qkv output
|
|
qkv_split = qkv_output.view(1, 1024, 3, 256)
|
|
q, k, v = qkv_split[:, :, 0], qkv_split[:, :, 1], qkv_split[:, :, 2]
|
|
|
|
# Apply norm
|
|
q, k = img_attn.norm(q, k, v)
|
|
|
|
# Compute attention
|
|
attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
|
logger.info(f"🚀 Manual attention computation successful, output shape: {attn_output.shape}")
|
|
|
|
# Store the attention output
|
|
hook_key = f"manual_attention_block_{block_idx}_{id(img_attn)}"
|
|
self.concept_attention.attention_outputs[hook_key] = attn_output
|
|
logger.info(f"📝 Stored manual attention output: {hook_key}")
|
|
|
|
# Apply projection
|
|
proj_output = img_attn.proj(attn_output)
|
|
logger.info(f"🚀 Manual projection successful, output shape: {proj_output.shape}")
|
|
|
|
# Store the projection output
|
|
hook_key = f"manual_proj_block_{block_idx}_{id(img_attn)}"
|
|
self.concept_attention.attention_outputs[hook_key] = proj_output
|
|
logger.info(f"📝 Stored manual projection output: {hook_key}")
|
|
|
|
# Break after first successful block to avoid too much computation
|
|
break
|
|
|
|
except Exception as manual_error:
|
|
logger.debug(f"Manual hook triggering failed: {manual_error}")
|
|
|
|
# If still no attention outputs, create fallback attention outputs from model parameters
|
|
if not self.concept_attention.attention_outputs:
|
|
logger.info("🔍 Still no attention outputs, creating fallback from model parameters")
|
|
try:
|
|
# Create fallback attention outputs using model parameters
|
|
if diffusion_model is not None and hasattr(diffusion_model, 'double_blocks'):
|
|
double_blocks = diffusion_model.double_blocks
|
|
|
|
for block_idx in [15, 16, 17, 18, 19]:
|
|
if block_idx < len(double_blocks):
|
|
block = double_blocks[block_idx]
|
|
|
|
if hasattr(block, 'img_attn'):
|
|
img_attn = block.img_attn
|
|
|
|
# Create fallback attention output using model weights
|
|
if hasattr(img_attn, 'qkv') and hasattr(img_attn.qkv, 'weight'):
|
|
# Get the weight matrix
|
|
weight = img_attn.qkv.weight
|
|
logger.info(f"Creating fallback attention from weight shape: {weight.shape}")
|
|
|
|
# Create a dummy input and compute output
|
|
dummy_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
|
|
|
with torch.no_grad():
|
|
# Compute output using the weight matrix
|
|
fallback_output = torch.matmul(dummy_input, weight.t())
|
|
logger.info(f"🚀 Fallback attention output created, shape: {fallback_output.shape}")
|
|
|
|
# Store the fallback output
|
|
hook_key = f"fallback_attention_block_{block_idx}_{id(img_attn)}"
|
|
self.concept_attention.attention_outputs[hook_key] = fallback_output
|
|
logger.info(f"📝 Stored fallback attention output: {hook_key}")
|
|
|
|
# Break after first successful block
|
|
break
|
|
|
|
except Exception as fallback_error:
|
|
logger.debug(f"Fallback attention creation failed: {fallback_error}")
|
|
|
|
# Log total attention outputs captured
|
|
logger.info(f"📊 Total attention outputs captured: {len(self.concept_attention.attention_outputs)}")
|
|
if self.concept_attention.attention_outputs:
|
|
logger.info(f"📊 Attention output keys: {list(self.concept_attention.attention_outputs.keys())}")
|
|
else:
|
|
logger.warning("⚠️ No attention outputs captured!")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"DoubleStreamBlock execution failed: {e}")
|
|
logger.info("Hooks may not have been triggered")
|
|
else:
|
|
logger.warning("No suitable model found, will use mock data")
|
|
|
|
# Check if any attention outputs were captured
|
|
if hasattr(self.concept_attention, 'attention_outputs') and self.concept_attention.attention_outputs:
|
|
logger.info(f"Successfully captured {len(self.concept_attention.attention_outputs)} attention outputs")
|
|
for name, output in self.concept_attention.attention_outputs.items():
|
|
logger.info(f" - {name}: {output.shape}")
|
|
else:
|
|
logger.warning("No attention outputs were captured! Hooks may not have been triggered.")
|
|
|
|
# Create concept maps from attention outputs
|
|
concept_maps = self.concept_attention._create_concept_maps_from_attention(concept_embeddings, image.shape)
|
|
|
|
# Cleanup hooks
|
|
self.concept_attention.cleanup_hooks()
|
|
|
|
return concept_maps
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in ConceptAttention: {e}")
|
|
# Cleanup hooks on error
|
|
self.concept_attention.cleanup_hooks()
|
|
# Return empty concept maps instead of raising error
|
|
logger.warning("Returning empty concept maps due to error")
|
|
return {}
|
|
|
|
def _extract_concept_embeddings(self, concepts: List[str], text_encoder, tokenizer):
|
|
"""
|
|
Extract concept embeddings using CLIP text encoder.
|
|
"""
|
|
return self.concept_attention._extract_concept_embeddings(concepts, text_encoder, tokenizer)
|