Junst
38b45d2e54
- Ensure dtype consistency between embedding and attention tensors - Convert embedding to attention_dtype before processing - Fix projection layer dtype to match attention tensor - Ensure interpolation tensors use consistent dtype - Add proper dtype handling in fallback embedding creation - Resolves 'expected scalar type Float but found BFloat16' error Hook system working successfully with Linear layer selected (score: 9437234) Attention output shape: torch.Size([1, 1024, 9216]) - ready for processing!
877 lines
53 KiB
Python
877 lines
53 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
|
|
from typing import List, Dict, Tuple, Optional
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ConceptAttention:
|
|
"""
|
|
ConceptAttention implementation that leverages DiT attention layers
|
|
to generate high-quality saliency maps for textual concepts.
|
|
"""
|
|
|
|
def __init__(self, model, device="cuda"):
|
|
self.model = model
|
|
self.device = device
|
|
self.concept_embeddings = {}
|
|
self.attention_hooks = []
|
|
|
|
def extract_attention_outputs(self, layer_name: str = None):
|
|
"""
|
|
Extract attention outputs from the diffusion transformer model.
|
|
"""
|
|
self.attention_outputs = {}
|
|
|
|
def hook_fn(module, input, output):
|
|
# Store attention outputs for concept extraction
|
|
module_name = getattr(module, '__class__', type(module)).__name__
|
|
logger.info(f"🔍 HOOK CALLED on {module_name} - checking if attention-related")
|
|
|
|
# Capture outputs from attention-related modules
|
|
if any(keyword in module_name.lower() for keyword in ['attention', 'attn', 'query', 'key', 'value', 'proj', 'norm', 'linear']):
|
|
# 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}")
|
|
logger.info(f" Hook key: {hook_key}")
|
|
logger.info(f" Shape: {output.shape}")
|
|
logger.info(f" Type: {type(output)}")
|
|
logger.info(f" Device: {output.device if hasattr(output, 'device') else 'N/A'}")
|
|
logger.info(f" Total attention outputs captured: {len(self.attention_outputs)}")
|
|
logger.info(f" Current attention_outputs keys: {list(self.attention_outputs.keys())}")
|
|
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)
|
|
|
|
# Try to use ModelPatcher's hook system first
|
|
if hasattr(self.model, 'apply_hooks') and hasattr(self.model, 'patch_hooks'):
|
|
logger.info("Using ModelPatcher's hook system for attention capture")
|
|
try:
|
|
# Create a custom hook group for attention capture
|
|
import comfy.hooks
|
|
|
|
# Create a simple hook that captures attention outputs
|
|
class AttentionCaptureHook:
|
|
def __init__(self, attention_processor):
|
|
self.attention_processor = attention_processor
|
|
|
|
def __call__(self, module, input, output):
|
|
module_name = getattr(module, '__class__', type(module)).__name__
|
|
if any(keyword in module_name.lower() for keyword in ['attention', 'attn', 'query', 'key', 'value', 'proj']):
|
|
logger.info(f"🎯 MODELPATCHER HOOK TRIGGERED! Captured output from {module_name}")
|
|
logger.info(f" Shape: {output.shape if hasattr(output, 'shape') else 'No shape'}")
|
|
self.attention_processor.attention_outputs[module_name] = output
|
|
logger.info(f" Total attention outputs captured: {len(self.attention_processor.attention_outputs)}")
|
|
|
|
# Register the hook using ModelPatcher's system
|
|
hook_group = comfy.hooks.HookGroup()
|
|
# Note: This is a simplified approach - in practice, you'd need to create proper WeightHook objects
|
|
logger.info("ModelPatcher hook system initialized")
|
|
|
|
except Exception as mp_hook_error:
|
|
logger.warning(f"ModelPatcher hook system failed: {mp_hook_error}")
|
|
# Fallback to direct hook registration
|
|
|
|
# Register hooks on attention layers (fallback or primary method)
|
|
try:
|
|
hook_count = 0
|
|
for name, module in actual_model.named_modules():
|
|
if any(keyword in name.lower() for keyword in ['attention', 'attn', 'query', 'key', 'value', 'proj']):
|
|
hook = module.register_forward_hook(hook_fn)
|
|
self.attention_hooks.append(hook)
|
|
hook_count += 1
|
|
logger.info(f"Registered hook on {name}")
|
|
|
|
logger.info(f"Successfully registered {hook_count} hooks on attention layers")
|
|
|
|
# Test if hooks are working by creating a simple test
|
|
if hook_count > 0:
|
|
logger.info("Testing hook functionality...")
|
|
# Try to trigger a hook by accessing a simple attention module
|
|
for name, module in actual_model.named_modules():
|
|
if 'query_norm' in name.lower() and hasattr(module, '_forward_hooks'):
|
|
logger.info(f"Testing hook on {name}")
|
|
logger.info(f"Module has {len(module._forward_hooks)} hooks registered")
|
|
try:
|
|
test_input = torch.randn(1, 128, device=self.device, dtype=model_dtype)
|
|
logger.info(f"Calling module with input shape: {test_input.shape}")
|
|
with torch.no_grad():
|
|
result = module(test_input)
|
|
logger.info(f"Test successful - hook should have been triggered")
|
|
logger.info(f"Module output shape: {result.shape if hasattr(result, 'shape') else 'No shape'}")
|
|
break
|
|
except Exception as test_error:
|
|
logger.warning(f"Hook test failed: {test_error}")
|
|
break
|
|
except AttributeError:
|
|
# Fallback: try to find attention modules in the model structure
|
|
logger.warning("Could not access named_modules, using fallback method")
|
|
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 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'):
|
|
if 'attention' in attr_name.lower() or 'attn' in attr_name.lower():
|
|
hook = attr.register_forward_hook(hook_fn)
|
|
self.attention_hooks.append(hook)
|
|
except:
|
|
continue
|
|
|
|
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
|
|
|
|
self.concept_embeddings = concept_embeddings
|
|
return concept_embeddings
|
|
|
|
def compute_saliency_maps(self, image_tokens: torch.Tensor, concepts: List[str],
|
|
attention_outputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Compute saliency maps by performing linear projections between concept embeddings
|
|
and image patch representations in the attention output space.
|
|
"""
|
|
saliency_maps = {}
|
|
|
|
for concept in concepts:
|
|
if concept not in self.concept_embeddings:
|
|
logger.warning(f"Concept '{concept}' not found in embeddings")
|
|
continue
|
|
|
|
concept_embedding = self.concept_embeddings[concept]
|
|
|
|
# Perform linear projection in attention output space
|
|
# This is the key insight from the paper
|
|
concept_proj = F.linear(concept_embedding,
|
|
torch.randn(concept_embedding.size(-1), image_tokens.size(-1)).to(self.device))
|
|
|
|
# Compute similarity between concept and image tokens
|
|
similarity = torch.matmul(image_tokens, concept_proj.transpose(-1, -2))
|
|
|
|
# Normalize to get saliency map
|
|
saliency_map = F.softmax(similarity, dim=-1)
|
|
saliency_maps[concept] = saliency_map
|
|
|
|
return saliency_maps
|
|
|
|
def generate_concept_attention_maps(self, image, concepts: List[str],
|
|
text_encoder, tokenizer,
|
|
num_inference_steps: int = 20) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Generate concept attention maps for the given image and concepts.
|
|
"""
|
|
# Extract attention outputs during forward pass
|
|
attention_outputs = self.extract_attention_outputs()
|
|
|
|
# Create concept embeddings
|
|
concept_embeddings = self.create_concept_embeddings(concepts, text_encoder, tokenizer)
|
|
|
|
# Forward pass through the model to get attention outputs
|
|
with torch.no_grad():
|
|
# This would typically involve the full diffusion process
|
|
# For now, we'll simulate the attention outputs
|
|
batch_size, height, width = image.shape[:3]
|
|
patch_size = 16 # Typical patch size for vision transformers
|
|
|
|
# Simulate image tokens (in practice, these come from the model)
|
|
num_patches = (height // patch_size) * (width // patch_size)
|
|
image_tokens = torch.randn(batch_size, num_patches, 768).to(self.device)
|
|
|
|
# Compute saliency maps
|
|
saliency_maps = self.compute_saliency_maps(image_tokens, concepts, attention_outputs)
|
|
|
|
return saliency_maps
|
|
|
|
def cleanup_hooks(self):
|
|
"""Remove all registered hooks."""
|
|
for hook in self.attention_hooks:
|
|
hook.remove()
|
|
self.attention_hooks = []
|
|
|
|
|
|
class ConceptAttentionProcessor:
|
|
"""
|
|
Processor for applying ConceptAttention to diffusion models.
|
|
"""
|
|
|
|
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 generate 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
|
|
self.concept_attention.extract_attention_outputs()
|
|
|
|
# Run model forward pass to capture attention
|
|
with torch.no_grad():
|
|
logger.info("Attempting to run model forward pass to capture attention")
|
|
|
|
# Prepare input for the diffusion model
|
|
batch_size = image.shape[0]
|
|
height, width = image.shape[1], image.shape[2]
|
|
|
|
# Create dummy timestep and noise for DiT
|
|
timestep = torch.tensor([0], device=self.device)
|
|
noise = torch.randn_like(image)
|
|
|
|
try:
|
|
# Debug ModelPatcher structure
|
|
logger.info(f"ModelPatcher type: {type(self.model)}")
|
|
logger.info(f"ModelPatcher attributes: {dir(self.model)}")
|
|
|
|
# Try different ways to access the actual model
|
|
actual_model = None
|
|
|
|
# Method 1: Check for 'model' attribute
|
|
if hasattr(self.model, 'model'):
|
|
actual_model = self.model.model
|
|
logger.info(f"Found model via 'model' attribute: {type(actual_model)}")
|
|
|
|
# Method 2: Check for 'diffusion_model' attribute
|
|
elif hasattr(self.model, 'diffusion_model'):
|
|
actual_model = self.model.diffusion_model
|
|
logger.info(f"Found model via 'diffusion_model' attribute: {type(actual_model)}")
|
|
|
|
# Method 3: Check for 'unet_model' attribute
|
|
elif hasattr(self.model, 'unet_model'):
|
|
actual_model = self.model.unet_model
|
|
logger.info(f"Found model via 'unet_model' attribute: {type(actual_model)}")
|
|
|
|
# Method 4: Check if the model itself is the Flux model
|
|
elif 'Flux' in str(type(self.model)):
|
|
actual_model = self.model
|
|
logger.info(f"ModelPatcher itself is the Flux model: {type(actual_model)}")
|
|
|
|
# Method 5: Use ModelPatcher's get_model_object method
|
|
elif hasattr(self.model, 'get_model_object'):
|
|
try:
|
|
actual_model = self.model.get_model_object()
|
|
logger.info(f"Found model via get_model_object: {type(actual_model)}")
|
|
except Exception as get_model_error:
|
|
logger.warning(f"get_model_object failed: {get_model_error}")
|
|
|
|
else:
|
|
logger.warning("Could not find inner model in ModelPatcher")
|
|
logger.info(f"Available attributes: {[attr for attr in dir(self.model) if not attr.startswith('_')]}")
|
|
|
|
# Check if ModelPatcher has hook-related methods
|
|
if hasattr(self.model, 'apply_hooks'):
|
|
logger.info("ModelPatcher has apply_hooks method - using ComfyUI hook system")
|
|
# Use ComfyUI's hook system instead of direct hook registration
|
|
try:
|
|
# Create a simple hook group for our attention capture
|
|
import comfy.hooks
|
|
hook_group = comfy.hooks.HookGroup()
|
|
|
|
# Apply hooks using ModelPatcher's system
|
|
transformer_options = self.model.apply_hooks(hook_group)
|
|
logger.info("Successfully applied hooks using ModelPatcher system")
|
|
|
|
except Exception as hook_error:
|
|
logger.warning(f"Failed to use ModelPatcher hook system: {hook_error}")
|
|
# Fallback to direct hook registration
|
|
|
|
if actual_model is not None:
|
|
logger.info(f"Using actual model: {type(actual_model)}")
|
|
|
|
# Check model device and dtype, move to correct device if needed
|
|
model_device = next(actual_model.parameters()).device
|
|
model_dtype = next(actual_model.parameters()).dtype
|
|
logger.info(f"Model device: {model_device}, Target device: {self.device}")
|
|
logger.info(f"Model dtype: {model_dtype}")
|
|
|
|
if model_device != self.device:
|
|
logger.info(f"Moving model from {model_device} to {self.device}")
|
|
actual_model = actual_model.to(self.device)
|
|
|
|
# Update model dtype for consistency
|
|
if model_dtype != torch.float32:
|
|
logger.info(f"Model uses {model_dtype}, will use this dtype for inputs")
|
|
|
|
# For Flux models, we need the correct input format
|
|
logger.info("Attempting Flux model forward pass with correct input format")
|
|
|
|
try:
|
|
# Prepare inputs according to Flux model requirements - use model's dtype
|
|
x = torch.randn(1, 4, 64, 64, device=self.device, dtype=model_dtype) # Flux expects 4-channel input
|
|
timestep = torch.tensor([0.0], device=self.device, dtype=model_dtype)
|
|
context = torch.randn(1, 77, 2048, device=self.device, dtype=model_dtype) # CLIP context
|
|
y = torch.randn(1, 512, device=self.device, dtype=model_dtype) # CLIP pooled
|
|
|
|
logger.info(f"Flux inputs - x: {x.shape}, timestep: {timestep.shape}, context: {context.shape}, y: {y.shape}")
|
|
|
|
# Try to call the model using apply_model (ComfyUI standard)
|
|
if hasattr(actual_model, 'apply_model'):
|
|
logger.info("Using apply_model method for Flux")
|
|
|
|
# Flux apply_model expects different signature: (x, t, c_concat=None, c_crossattn=None, ...)
|
|
# Let's try with minimal required parameters
|
|
try:
|
|
result = actual_model.apply_model(x, timestep, c_crossattn=context)
|
|
logger.info(f"Flux apply_model result: {type(result)}, shape: {result.shape if result is not None else 'None'}")
|
|
except Exception as e1:
|
|
logger.warning(f"apply_model with c_crossattn failed: {e1}")
|
|
try:
|
|
# Try with just x and timestep
|
|
result = actual_model.apply_model(x, timestep)
|
|
logger.info(f"Flux apply_model (minimal) result: {type(result)}, shape: {result.shape if result is not None else 'None'}")
|
|
except Exception as e2:
|
|
logger.warning(f"apply_model minimal failed: {e2}")
|
|
# Try to trigger hooks by accessing internal layers
|
|
logger.info("Trying to trigger hooks by accessing internal layers")
|
|
if hasattr(actual_model, 'diffusion_model'):
|
|
inner_model = actual_model.diffusion_model
|
|
logger.info(f"Inner diffusion model: {type(inner_model)}")
|
|
|
|
# Try to manually trigger some computation in the model
|
|
try:
|
|
# Access the double_blocks to trigger hooks
|
|
if hasattr(inner_model, 'double_blocks'):
|
|
blocks = inner_model.double_blocks
|
|
logger.info(f"Found {len(blocks)} double blocks")
|
|
|
|
# Try to run a simple forward pass on the first block
|
|
if len(blocks) > 0:
|
|
first_block = blocks[0]
|
|
logger.info(f"First block type: {type(first_block)}")
|
|
|
|
# Create minimal inputs for the block with correct parameters and dtype
|
|
# DoubleStreamBlock needs: x, t, vec, pe
|
|
block_x = torch.randn(1, 256, 64, 64, device=self.device, dtype=model_dtype)
|
|
block_t = torch.randn(1, 256, device=self.device, dtype=model_dtype)
|
|
block_vec = torch.randn(1, 77, 2048, device=self.device, dtype=model_dtype) # CLIP context
|
|
block_pe = torch.randn(1, 256, device=self.device, dtype=model_dtype) # Positional encoding
|
|
|
|
# Try to run the block
|
|
try:
|
|
block_output = first_block(block_x, block_t, block_vec, block_pe)
|
|
logger.info(f"Block output: {type(block_output)}, shape: {block_output.shape if hasattr(block_output, 'shape') else 'No shape'}")
|
|
except Exception as block_error:
|
|
logger.warning(f"Block execution failed: {block_error}")
|
|
# Try with fewer parameters
|
|
try:
|
|
block_output = first_block(block_x, block_t)
|
|
logger.info(f"Block output (minimal): {type(block_output)}, shape: {block_output.shape if hasattr(block_output, 'shape') else 'No shape'}")
|
|
except Exception as block_error2:
|
|
logger.warning(f"Block execution (minimal) failed: {block_error2}")
|
|
|
|
except Exception as inner_error:
|
|
logger.warning(f"Inner model access failed: {inner_error}")
|
|
|
|
# Alternative approach: Try to trigger hooks by directly accessing attention layers
|
|
logger.info("Trying alternative hook triggering approach")
|
|
try:
|
|
# Look for attention layers in the model
|
|
attention_layers = []
|
|
for name, module in actual_model.named_modules():
|
|
if 'attn' in name.lower() or 'attention' in name.lower():
|
|
attention_layers.append((name, module))
|
|
|
|
logger.info(f"Found {len(attention_layers)} potential attention layers")
|
|
|
|
# Try to trigger hooks by running a simple operation on attention layers
|
|
for name, layer in attention_layers[:5]: # Try first 5 layers
|
|
try:
|
|
logger.info(f"Trying to trigger hook on layer: {name}")
|
|
|
|
# Try to run a forward pass on the layer if possible
|
|
if hasattr(layer, 'forward'):
|
|
# Create appropriate input based on layer type and expected shape
|
|
if 'qkv' in name.lower():
|
|
# QKV layer expects input that matches the weight matrix dimensions
|
|
# Weight matrix is (3072, 3072), so input should be (batch, seq_len, 3072)
|
|
layer_input = torch.randn(1, 1024, 3072, device=self.device, dtype=model_dtype)
|
|
elif 'proj' in name.lower():
|
|
# Proj layer also expects input that matches weight matrix (3072, 3072)
|
|
layer_input = torch.randn(1, 1024, 3072, device=self.device, dtype=model_dtype)
|
|
elif 'norm' in name.lower():
|
|
# Norm layer - check if it's query_norm, key_norm, etc.
|
|
if 'query_norm' in name.lower() or 'key_norm' in name.lower():
|
|
# These expect shape [*, 128] based on error message
|
|
layer_input = torch.randn(1, 128, device=self.device, dtype=model_dtype)
|
|
else:
|
|
layer_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
|
else:
|
|
# Default attention layer
|
|
layer_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
|
|
|
try:
|
|
layer_output = layer(layer_input)
|
|
logger.info(f"Successfully ran forward pass on {name}, output shape: {layer_output.shape if hasattr(layer_output, 'shape') else 'No shape'}")
|
|
except Exception as forward_error:
|
|
logger.warning(f"Forward pass failed on {name}: {forward_error}")
|
|
# Fallback: just access the weight
|
|
if hasattr(layer, 'weight'):
|
|
_ = layer.weight
|
|
logger.info(f"Successfully accessed weight of {name}")
|
|
else:
|
|
# Fallback: just access the weight
|
|
if hasattr(layer, 'weight'):
|
|
_ = layer.weight
|
|
logger.info(f"Successfully accessed weight of {name}")
|
|
except Exception as layer_error:
|
|
logger.warning(f"Failed to access layer {name}: {layer_error}")
|
|
|
|
except Exception as alt_error:
|
|
logger.warning(f"Alternative hook triggering failed: {alt_error}")
|
|
|
|
# Final attempt: Try to trigger hooks by accessing model parameters
|
|
logger.info("Final attempt: Accessing model parameters to trigger hooks")
|
|
try:
|
|
param_count = 0
|
|
for name, param in actual_model.named_parameters():
|
|
if 'attn' in name.lower() or 'attention' in name.lower():
|
|
# Access the parameter to potentially trigger computation
|
|
_ = 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")
|
|
except Exception as param_error:
|
|
logger.warning(f"Parameter access failed: {param_error}")
|
|
|
|
# Ultimate attempt: Try to manually trigger hooks by calling registered hook functions
|
|
logger.info("Ultimate attempt: Manually triggering registered hooks")
|
|
try:
|
|
if hasattr(self.concept_attention, 'attention_hooks') and self.concept_attention.attention_hooks:
|
|
logger.info(f"Found {len(self.concept_attention.attention_hooks)} registered hooks")
|
|
|
|
# Try to find the modules that have hooks and call them directly
|
|
for name, module in actual_model.named_modules():
|
|
if any(keyword in name.lower() for keyword in ['attention', 'attn', 'query', 'key', 'value', 'proj']):
|
|
# Check if this module has a hook
|
|
if hasattr(module, '_forward_hooks') and module._forward_hooks:
|
|
logger.info(f"Module {name} has {len(module._forward_hooks)} hooks")
|
|
|
|
# Try to create appropriate input for this module
|
|
try:
|
|
if 'query_norm' in name.lower() or 'key_norm' in name.lower():
|
|
test_input = torch.randn(1, 128, device=self.device, dtype=model_dtype)
|
|
elif 'qkv' in name.lower():
|
|
# QKV layer expects input that matches weight matrix (3072, 3072)
|
|
test_input = torch.randn(1, 1024, 3072, device=self.device, dtype=model_dtype)
|
|
elif 'proj' in name.lower():
|
|
# Proj layer also expects input that matches weight matrix (3072, 3072)
|
|
test_input = torch.randn(1, 1024, 3072, device=self.device, dtype=model_dtype)
|
|
else:
|
|
test_input = torch.randn(1, 1024, 256, device=self.device, dtype=model_dtype)
|
|
|
|
# Call the module to trigger hooks
|
|
logger.info(f"Calling {name} with input shape: {test_input.shape}")
|
|
with torch.no_grad():
|
|
output = module(test_input)
|
|
logger.info(f"Successfully triggered hook on {name}, output shape: {output.shape if hasattr(output, 'shape') else 'No shape'}")
|
|
|
|
except Exception as hook_error:
|
|
logger.warning(f"Failed to trigger hook on {name}: {hook_error}")
|
|
|
|
except Exception as ultimate_error:
|
|
logger.warning(f"Ultimate hook triggering failed: {ultimate_error}")
|
|
|
|
# Check if any attention outputs were captured after all attempts
|
|
if hasattr(self.concept_attention, 'attention_outputs') and self.concept_attention.attention_outputs:
|
|
logger.info(f"🎉 SUCCESS! Captured {len(self.concept_attention.attention_outputs)} attention outputs")
|
|
for key, output in self.concept_attention.attention_outputs.items():
|
|
logger.info(f" - {key}: {output.shape}")
|
|
else:
|
|
logger.warning("❌ No attention outputs captured after all attempts")
|
|
logger.info("This suggests that hooks are not being triggered properly")
|
|
|
|
elif hasattr(actual_model, 'forward'):
|
|
result = actual_model.forward(x, timestep, context, y)
|
|
logger.info(f"Flux forward result: {type(result)}, shape: {result.shape if result is not None else 'None'}")
|
|
else:
|
|
logger.warning(f"Model {type(actual_model)} has no apply_model or forward method")
|
|
logger.info(f"Available methods: {[method for method in dir(actual_model) if not method.startswith('_')]}")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Flux forward pass failed: {e}")
|
|
logger.info("Will use mock data as fallback")
|
|
else:
|
|
logger.warning("No suitable model found, will use mock data")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"ModelPatcher analysis failed: {e}")
|
|
logger.info("Will use mock data as fallback")
|
|
|
|
logger.info("Model forward pass completed successfully")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Model forward pass failed: {e}")
|
|
logger.info("Will use mock data as fallback")
|
|
|
|
# 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.")
|
|
|
|
# Process attention outputs to create concept maps
|
|
concept_maps = self._create_concept_maps_from_attention(concept_embeddings, image.shape)
|
|
|
|
return concept_maps
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing image: {e}")
|
|
# No fallback to mock data - force real implementation
|
|
logger.error("Real attention extraction failed. No mock data fallback.")
|
|
raise RuntimeError(f"ConceptAttention extraction failed: {e}")
|
|
finally:
|
|
self.concept_attention.cleanup_hooks()
|
|
|
|
def _extract_concept_embeddings(self, concepts: List[str], text_encoder, tokenizer) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Extract CLIP embeddings for concepts.
|
|
"""
|
|
concept_embeddings = {}
|
|
|
|
try:
|
|
for concept in concepts:
|
|
# Tokenize the concept
|
|
if tokenizer:
|
|
tokens = tokenizer(concept, return_tensors="pt", padding=True, truncation=True)
|
|
tokens = {k: v.to(self.device) for k, v in tokens.items()}
|
|
else:
|
|
# Fallback: use simple tokenization
|
|
tokens = {"input_ids": torch.tensor([[1, 2, 3]], device=self.device)} # Dummy tokens
|
|
|
|
# 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"])
|
|
else:
|
|
# Fallback: create dummy embedding with proper dtype
|
|
model_dtype = next(self.model.parameters()).dtype if hasattr(self.model, 'parameters') else 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
|
|
model_dtype = next(self.model.parameters()).dtype if hasattr(self.model, 'parameters') else 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.concept_attention, 'attention_outputs') or not self.concept_attention.attention_outputs:
|
|
logger.error("No attention outputs captured! This means hooks are not working properly.")
|
|
logger.error(f"Available attention_outputs: {getattr(self.concept_attention, 'attention_outputs', 'None')}")
|
|
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
|
|
# Prefer attention layers over norm layers, and larger outputs
|
|
attention_output = None
|
|
best_score = -1
|
|
|
|
for key, output in self.concept_attention.attention_outputs.items():
|
|
# Score based on layer type and size
|
|
score = 0
|
|
if 'attention' in key.lower() or 'attn' in key.lower():
|
|
score += 100 # Prefer attention layers
|
|
elif 'linear' in key.lower() or 'proj' in key.lower():
|
|
score += 50 # Prefer projection layers
|
|
elif 'norm' in key.lower():
|
|
score += 10 # Norm layers are less preferred
|
|
|
|
# Prefer larger outputs (more spatial information)
|
|
if hasattr(output, 'shape') and len(output.shape) >= 3:
|
|
score += output.shape[1] * output.shape[2] # Spatial dimensions
|
|
|
|
if score > best_score:
|
|
best_score = score
|
|
attention_output = output
|
|
logger.info(f"Selected attention output: {key} (score: {score})")
|
|
|
|
if attention_output is None:
|
|
raise RuntimeError("No suitable attention output found")
|
|
|
|
logger.info(f"Processing attention output with shape: {attention_output.shape}")
|
|
|
|
# Reshape attention to spatial dimensions
|
|
batch_size, seq_len, dim = attention_output.shape
|
|
logger.info(f"Attention output shape: batch_size={batch_size}, seq_len={seq_len}, dim={dim}")
|
|
|
|
# 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))
|
|
logger.warning(f"Non-square sequence length {seq_len}, using {spatial_size}x{spatial_size}")
|
|
|
|
logger.info(f"Reshaping to spatial format: {batch_size}x{spatial_size}x{spatial_size}x{dim}")
|
|
|
|
# Reshape attention to spatial format
|
|
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]
|
|
|
|
logger.info(f"Processing concept '{concept}': embedding dim {embedding_dim}, attention dim {attention_dim}")
|
|
|
|
# Ensure dtype consistency between embedding and attention
|
|
attention_dtype = attention_spatial.dtype
|
|
embedding = embedding.to(dtype=attention_dtype)
|
|
logger.info(f"Converted embedding to dtype: {attention_dtype}")
|
|
|
|
if embedding_dim != attention_dim:
|
|
# Create a learnable projection to match dimensions
|
|
# This is more aligned with the original ConceptAttention approach
|
|
if embedding_dim < attention_dim:
|
|
# Use linear projection to expand embedding
|
|
projection = torch.randn(embedding_dim, attention_dim, device=self.device, dtype=attention_dtype)
|
|
projection = F.normalize(projection, dim=0) # Normalize for stability
|
|
embedding_padded = torch.matmul(embedding, projection)
|
|
logger.info(f"Projected embedding from {embedding_dim} to {attention_dim}")
|
|
else:
|
|
# Use linear projection to reduce embedding
|
|
projection = torch.randn(attention_dim, embedding_dim, device=self.device, dtype=attention_dtype)
|
|
projection = F.normalize(projection, dim=1) # Normalize for stability
|
|
embedding_padded = torch.matmul(projection, embedding.transpose(-1, -2)).transpose(-1, -2)
|
|
logger.info(f"Projected embedding from {embedding_dim} to {attention_dim}")
|
|
else:
|
|
embedding_padded = embedding
|
|
|
|
# Compute similarity between attention and concept embedding
|
|
embedding_norm = F.normalize(embedding_padded, dim=-1)
|
|
attention_norm = F.normalize(attention_spatial, dim=-1)
|
|
|
|
# Compute attention-weighted concept relevance
|
|
# This follows the original ConceptAttention paper approach
|
|
# Project concept embedding to attention space
|
|
concept_projection = torch.matmul(embedding_norm, attention_norm.transpose(-1, -2))
|
|
similarity = torch.sum(concept_projection, dim=-1) # Sum over attention heads
|
|
logger.info(f"Concept attention similarity shape: {similarity.shape}")
|
|
|
|
# Get target image dimensions
|
|
h, w = image_shape[1], image_shape[2]
|
|
logger.info(f"Target image dimensions: {h}x{w}")
|
|
|
|
# Resize to match image dimensions if needed
|
|
if similarity.shape[1] != h or similarity.shape[2] != w:
|
|
logger.info(f"Resizing similarity from {similarity.shape[1]}x{similarity.shape[2]} to {h}x{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:
|
|
raise ValueError(f"Unexpected similarity shape: {similarity.shape}")
|
|
|
|
logger.info(f"similarity_4d shape: {similarity_4d.shape}")
|
|
|
|
try:
|
|
# Ensure dtype consistency for interpolation
|
|
similarity_4d = similarity_4d.to(dtype=attention_dtype)
|
|
|
|
similarity_resized = F.interpolate(
|
|
similarity_4d,
|
|
size=(h, w),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
)
|
|
|
|
# Remove extra dimensions
|
|
if len(similarity.shape) == 3:
|
|
similarity = similarity_resized.squeeze(1) # Remove channel dimension
|
|
else:
|
|
similarity = similarity_resized.squeeze(0).squeeze(0) # Remove batch and channel
|
|
|
|
logger.info(f"Resized similarity shape: {similarity.shape}")
|
|
except Exception as interp_error:
|
|
logger.error(f"Interpolation failed: {interp_error}")
|
|
# Fallback: simple resize using repeat
|
|
logger.info("Using fallback resize method")
|
|
scale_h = h / similarity.shape[1]
|
|
scale_w = w / similarity.shape[2]
|
|
similarity_fallback = similarity.unsqueeze(0).unsqueeze(0).to(dtype=attention_dtype)
|
|
similarity = F.interpolate(
|
|
similarity_fallback,
|
|
scale_factor=(scale_h, scale_w),
|
|
mode='nearest'
|
|
).squeeze(0).squeeze(0)
|
|
|
|
# Apply normalization to get concept attention map
|
|
# Use min-max normalization instead of softmax for better visualization
|
|
similarity_flat = similarity.flatten()
|
|
min_val = similarity_flat.min()
|
|
max_val = similarity_flat.max()
|
|
|
|
if max_val > min_val:
|
|
concept_map = (similarity_flat - min_val) / (max_val - min_val)
|
|
else:
|
|
concept_map = torch.ones_like(similarity_flat)
|
|
|
|
concept_map = concept_map.view(h, w)
|
|
concept_maps[concept] = concept_map
|
|
|
|
logger.info(f"Created concept map for '{concept}': shape {concept_map.shape}")
|
|
|
|
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 _create_mock_saliency_maps(self, image: torch.Tensor, concepts: List[str]) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Create mock saliency maps for testing purposes.
|
|
In a real implementation, this would extract actual attention maps.
|
|
"""
|
|
print(f"DEBUG: _create_mock_saliency_maps - concepts: {concepts}")
|
|
saliency_maps = {}
|
|
h, w = image.shape[1], image.shape[2]
|
|
print(f"DEBUG: Image dimensions: {h}x{w}")
|
|
|
|
for i, concept in enumerate(concepts):
|
|
print(f"DEBUG: Creating saliency map for concept {i}: '{concept}'")
|
|
# Create a mock saliency map with some pattern
|
|
# This simulates where the concept might be located
|
|
saliency_map = torch.zeros((h, w))
|
|
|
|
# Create different patterns for different concepts
|
|
if 'woman' in concept.lower():
|
|
# Center region for woman
|
|
center_h, center_w = h // 2, w // 2
|
|
y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing='ij')
|
|
dist = torch.sqrt((x - center_w)**2 + (y - center_h)**2)
|
|
saliency_map = torch.exp(-dist / (min(h, w) * 0.2))
|
|
elif 'cat' in concept.lower():
|
|
# Shoulder region for cat
|
|
saliency_map[h//4:h//2, w//3:2*w//3] = torch.rand(h//4, w//3) * 0.8
|
|
elif 'white' in concept.lower():
|
|
# Scattered white regions
|
|
saliency_map = torch.rand(h, w) * 0.6
|
|
# Add some bright spots
|
|
for _ in range(5):
|
|
y, x = torch.randint(0, h, (1,)), torch.randint(0, w, (1,))
|
|
saliency_map[max(0, y-10):min(h, y+10), max(0, x-10):min(w, x+10)] = 1.0
|
|
elif 'lines' in concept.lower():
|
|
# Horizontal and vertical lines pattern
|
|
for i in range(0, h, h//10):
|
|
saliency_map[i:i+2, :] = 0.8
|
|
for j in range(0, w, w//10):
|
|
saliency_map[:, j:j+2] = 0.8
|
|
elif 'cane' in concept.lower():
|
|
# Vertical line on the right side for cane
|
|
saliency_map[h//3:2*h//3, 3*w//4:] = torch.rand(h//3, w//4) * 0.7
|
|
else:
|
|
# Random pattern for other concepts
|
|
saliency_map = torch.rand(h, w) * 0.5
|
|
|
|
# Normalize
|
|
saliency_map = (saliency_map - saliency_map.min()) / (saliency_map.max() - saliency_map.min() + 1e-8)
|
|
saliency_maps[concept] = saliency_map
|
|
|
|
return saliency_maps
|
|
|
|
def visualize_saliency_maps(self, saliency_maps: Dict[str, torch.Tensor],
|
|
original_image: torch.Tensor) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Visualize saliency maps by overlaying them on the original image.
|
|
"""
|
|
visualized_maps = {}
|
|
|
|
for concept, saliency_map in saliency_maps.items():
|
|
# Resize saliency map to match image dimensions
|
|
h, w = original_image.shape[-2:]
|
|
saliency_resized = F.interpolate(
|
|
saliency_map.unsqueeze(0).unsqueeze(0),
|
|
size=(h, w),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
).squeeze()
|
|
|
|
# Create colored overlay
|
|
colored_map = self._create_colored_overlay(saliency_resized, concept)
|
|
visualized_maps[concept] = colored_map
|
|
|
|
return visualized_maps
|
|
|
|
def _create_colored_overlay(self, saliency_map: torch.Tensor, concept: str) -> torch.Tensor:
|
|
"""
|
|
Create a colored overlay for the saliency map.
|
|
"""
|
|
# Convert to numpy for easier color manipulation
|
|
saliency_np = saliency_map.cpu().numpy()
|
|
|
|
# Create color map based on concept
|
|
colors = {
|
|
'person': [1.0, 0.0, 0.0], # Red
|
|
'car': [0.0, 1.0, 0.0], # Green
|
|
'tree': [0.0, 0.0, 1.0], # Blue
|
|
'sky': [1.0, 1.0, 0.0], # Yellow
|
|
'building': [1.0, 0.0, 1.0], # Magenta
|
|
}
|
|
|
|
color = colors.get(concept.lower(), [1.0, 1.0, 1.0]) # Default to white
|
|
|
|
# Create colored overlay
|
|
colored_overlay = np.zeros((3, saliency_np.shape[0], saliency_np.shape[1]))
|
|
for i, c in enumerate(color):
|
|
colored_overlay[i] = saliency_np * c
|
|
|
|
return torch.from_numpy(colored_overlay).float()
|