Remove unnecessary concept_attention_clean.py file
🧹 **Cleanup**:
- Removed temporary concept_attention_clean.py file
- Only keeping the main concept_attention.py file
- Cleaner repository structure
The main concept_attention.py file contains the latest
implementation with patches_replace system and direct
attention module execution.
This commit is contained in:
Junst
@@ -1,478 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional
|
||||
import math
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class ConceptAttention:
|
||||
"""
|
||||
ConceptAttention implementation for ComfyUI using patches_replace system.
|
||||
"""
|
||||
|
||||
def __init__(self, device="cuda"):
|
||||
self.device = device
|
||||
self.attention_outputs = {}
|
||||
self.model = None
|
||||
|
||||
def register_attention_hooks(self, model):
|
||||
"""
|
||||
Register attention capture using ComfyUI's patches_replace system.
|
||||
"""
|
||||
logger.info("🔍 Using ComfyUI patches_replace system for attention capture")
|
||||
|
||||
# Store model reference for patches_replace system
|
||||
self.model = model
|
||||
|
||||
# Try to find the actual model
|
||||
actual_model = None
|
||||
if hasattr(model, 'model'):
|
||||
actual_model = model.model
|
||||
elif hasattr(model, 'diffusion_model'):
|
||||
actual_model = model.diffusion_model
|
||||
elif hasattr(model, 'unet_model'):
|
||||
actual_model = model.unet_model
|
||||
else:
|
||||
actual_model = model
|
||||
|
||||
logger.info(f"🔍 Using model for attention capture: {type(actual_model)}")
|
||||
|
||||
return self.attention_outputs
|
||||
|
||||
def _create_attention_capture_wrapper(self, block_idx):
|
||||
"""
|
||||
Create a wrapper function for ComfyUI's patches_replace system to capture attention.
|
||||
"""
|
||||
def attention_capture_wrapper(args, original_block):
|
||||
"""
|
||||
Wrapper function that captures attention outputs and calls the original block.
|
||||
"""
|
||||
logger.info(f"🎯 Attention capture wrapper called for block {block_idx}")
|
||||
|
||||
# Extract arguments
|
||||
img = args["img"]
|
||||
txt = args["txt"]
|
||||
vec = args["vec"]
|
||||
pe = args["pe"]
|
||||
attn_mask = args.get("attn_mask")
|
||||
transformer_options = args.get("transformer_options", {})
|
||||
|
||||
# Call the original block
|
||||
img_out, txt_out = original_block(args)
|
||||
|
||||
# Try to capture attention outputs from the block
|
||||
try:
|
||||
# Get the actual block from the model
|
||||
if hasattr(self.model, 'model') and hasattr(self.model.model, 'double_blocks'):
|
||||
actual_block = self.model.model.double_blocks[block_idx]
|
||||
|
||||
# Try to access attention modules
|
||||
if hasattr(actual_block, 'img_attn'):
|
||||
img_attn = actual_block.img_attn
|
||||
|
||||
# Create test input for attention
|
||||
test_input = torch.randn(1, 1024, 256, device=img.device, dtype=img.dtype)
|
||||
|
||||
with torch.no_grad():
|
||||
# Try to capture qkv output
|
||||
if hasattr(img_attn, 'qkv'):
|
||||
qkv_output = img_attn.qkv(test_input)
|
||||
hook_key = f"patches_replace_qkv_block_{block_idx}_{id(img_attn)}"
|
||||
self.attention_outputs[hook_key] = qkv_output
|
||||
logger.info(f"📝 Captured qkv output from block {block_idx}: {qkv_output.shape}")
|
||||
|
||||
# Try to capture 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)
|
||||
hook_key = f"patches_replace_attention_block_{block_idx}_{id(img_attn)}"
|
||||
self.attention_outputs[hook_key] = attn_output
|
||||
logger.info(f"📝 Captured attention output from block {block_idx}: {attn_output.shape}")
|
||||
|
||||
# Apply projection
|
||||
proj_output = img_attn.proj(attn_output)
|
||||
hook_key = f"patches_replace_proj_block_{block_idx}_{id(img_attn)}"
|
||||
self.attention_outputs[hook_key] = proj_output
|
||||
logger.info(f"📝 Captured projection output from block {block_idx}: {proj_output.shape}")
|
||||
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to capture attention from block {block_idx}: {e}")
|
||||
|
||||
return {"img": img_out, "txt": txt_out}
|
||||
|
||||
return attention_capture_wrapper
|
||||
|
||||
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, concepts: List[str], concept_embeddings: Dict[str, torch.Tensor], image_shape: tuple):
|
||||
"""
|
||||
Create concept maps from captured attention outputs.
|
||||
"""
|
||||
concept_maps = {}
|
||||
|
||||
if not self.attention_outputs:
|
||||
logger.warning("No attention outputs available for concept map creation")
|
||||
return concept_maps
|
||||
|
||||
# Select the best attention output
|
||||
best_attention_output = self._select_best_attention_output()
|
||||
if best_attention_output is None:
|
||||
logger.warning("No suitable attention output found")
|
||||
return concept_maps
|
||||
|
||||
attention_output, attention_key = best_attention_output
|
||||
logger.info(f"Using attention output: {attention_key}, shape: {attention_output.shape}")
|
||||
|
||||
# Get model dtype for consistency
|
||||
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:
|
||||
try:
|
||||
embedding = concept_embeddings[concept]
|
||||
logger.info(f"Processing concept '{concept}': embedding dim {embedding.shape[-1]}, attention dim {attention_output.shape[-1]}")
|
||||
|
||||
# Ensure embedding is on correct device and dtype
|
||||
embedding = embedding.to(device=self.device, dtype=model_dtype)
|
||||
|
||||
# Handle different embedding shapes
|
||||
if embedding.dim() == 2:
|
||||
embedding = embedding.unsqueeze(1) # [1, 1, embedding_dim]
|
||||
elif embedding.dim() == 3:
|
||||
pass # Already [1, 1, embedding_dim] or [1, seq_len, embedding_dim]
|
||||
else:
|
||||
logger.warning(f"Unexpected embedding shape: {embedding.shape}")
|
||||
continue
|
||||
|
||||
# Ensure embedding is [1, 1, embedding_dim]
|
||||
if embedding.shape[1] > 1:
|
||||
embedding = embedding.mean(dim=1, keepdim=True) # Average over sequence length
|
||||
|
||||
# Align embedding dimension with attention output dimension
|
||||
embedding_dim = embedding.shape[-1]
|
||||
attention_dim = attention_output.shape[-1]
|
||||
|
||||
if embedding_dim != attention_dim:
|
||||
# Create a learnable projection layer
|
||||
if not hasattr(self, f'projection_{attention_dim}'):
|
||||
projection = torch.nn.Linear(embedding_dim, attention_dim, device=self.device, dtype=model_dtype)
|
||||
setattr(self, f'projection_{attention_dim}', projection)
|
||||
else:
|
||||
projection = getattr(self, f'projection_{attention_dim}')
|
||||
|
||||
# Project embedding to match attention dimension
|
||||
embedding = projection(embedding)
|
||||
logger.info(f"Projected embedding from {embedding_dim} to {attention_dim}")
|
||||
|
||||
# Reshape attention output for similarity computation
|
||||
if attention_output.dim() == 3:
|
||||
# [batch, seq_len, dim] -> [batch, seq_len]
|
||||
attention_flat = attention_output.view(attention_output.shape[0], -1, attention_output.shape[-1])
|
||||
else:
|
||||
attention_flat = attention_output
|
||||
|
||||
# Compute similarity between concept embedding and attention output
|
||||
# Use einsum-like operation similar to original ConceptAttention
|
||||
embedding_expanded = embedding.expand(attention_flat.shape[0], attention_flat.shape[1], -1)
|
||||
|
||||
# Compute concept-image attention map using einsum-like operation
|
||||
concept_attn_output = torch.matmul(embedding_expanded, attention_flat.transpose(-2, -1))
|
||||
concept_attn_output = concept_attn_output.squeeze(1) # [batch, seq_len]
|
||||
|
||||
# Reshape to spatial dimensions
|
||||
seq_len = concept_attn_output.shape[-1]
|
||||
spatial_h = int(math.sqrt(seq_len))
|
||||
spatial_w = seq_len // spatial_h
|
||||
|
||||
if spatial_h * spatial_w == seq_len:
|
||||
similarity = concept_attn_output.view(1, spatial_h, spatial_w)
|
||||
else:
|
||||
# Fallback: use 1D reshaping
|
||||
similarity = concept_attn_output.view(1, seq_len, 1)
|
||||
|
||||
logger.info(f"Concept attention similarity shape: {similarity.shape}")
|
||||
|
||||
# Resize to target image dimensions
|
||||
target_h, target_w = image_shape[:2]
|
||||
logger.info(f"Target image dimensions: {target_h}x{target_w}")
|
||||
|
||||
# Ensure similarity is in correct format for interpolation
|
||||
if similarity.dim() == 2:
|
||||
similarity = similarity.unsqueeze(0) # Add batch dimension
|
||||
elif similarity.dim() == 3:
|
||||
if similarity.shape[1] == 1: # [1, 1, seq_len]
|
||||
similarity = similarity.permute(0, 2, 1) # [1, seq_len, 1]
|
||||
|
||||
# Convert to 4D for interpolation: [N, C, H, W]
|
||||
if similarity.dim() == 3:
|
||||
if similarity.shape[2] == 1: # [1, H, 1]
|
||||
similarity_4d = similarity.unsqueeze(1) # [1, 1, H, 1]
|
||||
else: # [1, H, W]
|
||||
similarity_4d = similarity.unsqueeze(1) # [1, 1, H, W]
|
||||
else:
|
||||
similarity_4d = similarity
|
||||
|
||||
logger.info(f"similarity_4d shape: {similarity_4d.shape}")
|
||||
|
||||
# Interpolate to target size
|
||||
resized_similarity = F.interpolate(
|
||||
similarity_4d,
|
||||
size=(target_h, target_w),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
)
|
||||
|
||||
logger.info(f"Resized similarity shape: {resized_similarity.shape}")
|
||||
|
||||
# Convert to concept map
|
||||
concept_map = resized_similarity.squeeze(0).squeeze(0) # Remove batch and channel dimensions
|
||||
|
||||
# Normalize concept map using percentile-based normalization
|
||||
concept_map_flat = concept_map.flatten()
|
||||
concept_map_flat = concept_map_flat.to(torch.float32) # Convert to float32 for quantile
|
||||
|
||||
# Use percentile-based normalization
|
||||
p5 = torch.quantile(concept_map_flat, 0.05)
|
||||
p95 = torch.quantile(concept_map_flat, 0.95)
|
||||
|
||||
# Clamp and normalize
|
||||
concept_map = torch.clamp(concept_map, p5, p95)
|
||||
concept_map = (concept_map - p5) / (p95 - p5 + 1e-8)
|
||||
|
||||
# Apply gamma correction for better contrast
|
||||
concept_map = torch.pow(concept_map, 0.5)
|
||||
|
||||
# Convert to float32 for ComfyUI compatibility
|
||||
concept_map = concept_map.float()
|
||||
|
||||
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 map for '{concept}': {e}")
|
||||
continue
|
||||
|
||||
return concept_maps
|
||||
|
||||
def _select_best_attention_output(self):
|
||||
"""
|
||||
Select the best attention output from captured outputs.
|
||||
"""
|
||||
if not self.attention_outputs:
|
||||
return None
|
||||
|
||||
best_output = None
|
||||
best_key = None
|
||||
best_score = -1
|
||||
|
||||
for key, output in self.attention_outputs.items():
|
||||
score = 0
|
||||
|
||||
# Prioritize patches_replace outputs
|
||||
if 'patches_replace' in key:
|
||||
score += 1000
|
||||
|
||||
# Prioritize attention-related outputs
|
||||
if 'attention' in key.lower():
|
||||
score += 500
|
||||
elif 'qkv' in key.lower():
|
||||
score += 300
|
||||
elif 'proj' in key.lower():
|
||||
score += 200
|
||||
|
||||
# Prefer outputs with proper spatial structure
|
||||
if hasattr(output, 'shape') and len(output.shape) >= 3:
|
||||
if output.shape[-1] > 100: # Prefer outputs with reasonable dimension
|
||||
score += 100
|
||||
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_output = output
|
||||
best_key = key
|
||||
|
||||
return (best_output, best_key) if best_output is not None else None
|
||||
|
||||
def extract_attention_outputs(self):
|
||||
"""
|
||||
Extract attention outputs using patches_replace system.
|
||||
"""
|
||||
if not self.model:
|
||||
logger.warning("No model available for attention extraction")
|
||||
return self.attention_outputs
|
||||
|
||||
# Try to find 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
|
||||
else:
|
||||
actual_model = self.model
|
||||
|
||||
if not actual_model:
|
||||
logger.warning("No actual model found")
|
||||
return self.attention_outputs
|
||||
|
||||
# Get model dtype
|
||||
if hasattr(actual_model, 'parameters'):
|
||||
params = list(actual_model.parameters())
|
||||
model_dtype = params[0].dtype if params else torch.float32
|
||||
else:
|
||||
model_dtype = torch.float32
|
||||
|
||||
logger.info(f"🔍 Attempting attention capture using patches_replace system")
|
||||
logger.info(f"Using model: {type(actual_model)}")
|
||||
|
||||
try:
|
||||
# Try to use ComfyUI's patches_replace system to intercept attention
|
||||
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 patches_replace system")
|
||||
|
||||
# Create a custom transformer_options with patches_replace
|
||||
transformer_options = {
|
||||
"patches_replace": {
|
||||
"dit": {
|
||||
("double_block", 15): self._create_attention_capture_wrapper(15),
|
||||
("double_block", 16): self._create_attention_capture_wrapper(16),
|
||||
("double_block", 17): self._create_attention_capture_wrapper(17),
|
||||
("double_block", 18): self._create_attention_capture_wrapper(18),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with torch.no_grad():
|
||||
output = actual_model.apply_model(x, timestep, context, y, transformer_options=transformer_options)
|
||||
logger.info(f"🚀 apply_model with patches_replace completed, output shape: {output.shape}")
|
||||
|
||||
# Store the output as attention output
|
||||
hook_key = f"patches_replace_output_{id(actual_model)}"
|
||||
self.attention_outputs[hook_key] = output
|
||||
logger.info(f"📝 Stored patches_replace output: {hook_key}")
|
||||
|
||||
except Exception as patches_error:
|
||||
logger.debug(f"patches_replace system failed: {patches_error}")
|
||||
|
||||
# Log total attention outputs captured
|
||||
logger.info(f"📊 Total attention outputs captured: {len(self.attention_outputs)}")
|
||||
if self.attention_outputs:
|
||||
logger.info(f"📊 Attention output keys: {list(self.attention_outputs.keys())}")
|
||||
else:
|
||||
logger.warning("⚠️ No attention outputs captured!")
|
||||
|
||||
return self.attention_outputs
|
||||
|
||||
|
||||
class ConceptAttentionProcessor:
|
||||
"""
|
||||
Main processor for ConceptAttention functionality.
|
||||
"""
|
||||
|
||||
def __init__(self, model, device="cuda"):
|
||||
self.model = model
|
||||
self.device = device
|
||||
self.concept_attention = ConceptAttention(device=device)
|
||||
|
||||
def process_image(self, image: torch.Tensor, concepts: List[str], text_encoder=None, tokenizer=None):
|
||||
"""
|
||||
Process an image to extract concept attention maps.
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Processing image with concepts: {concepts}")
|
||||
|
||||
# Register attention capture system
|
||||
self.concept_attention.register_attention_hooks(self.model)
|
||||
|
||||
# Extract concept embeddings
|
||||
concept_embeddings = self.concept_attention._extract_concept_embeddings(concepts, text_encoder, tokenizer)
|
||||
logger.info(f"Extracted embeddings for concepts: {list(concept_embeddings.keys())}")
|
||||
|
||||
# Extract attention outputs
|
||||
attention_outputs = self.concept_attention.extract_attention_outputs()
|
||||
|
||||
# Create concept maps
|
||||
image_shape = image.shape[-2:] # Get height and width
|
||||
concept_maps = self.concept_attention._create_concept_maps_from_attention(
|
||||
concepts, concept_embeddings, image_shape
|
||||
)
|
||||
|
||||
return concept_maps
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in ConceptAttention: {e}")
|
||||
return {}
|
||||
Reference in New Issue
Block a user