Junst
6595d25bad
- Remove VAE output from CheckpointLoaderSimple (not needed for ConceptAttention) - Update example prompt to: 'Elderly woman character with white hair and a red dress holding a cane and accompanied by a small white cat on her shoulder. Cartoon style with smooth lines and exaggerated facial features offering a whimsical and friendly appearance. Her name is abae.' - Update concept segmentation to: woman, cat, white, lines, cane - Update concept saliency map default to: woman - Update mock saliency map patterns for new concepts: * woman: center region (red) * cat: shoulder region (green) * white: scattered bright spots (white) * lines: horizontal/vertical line patterns (blue) * cane: vertical line on right side (magenta) - Update color mappings in utils.py and nodes.py - Fix workflow link validation issues
275 lines
11 KiB
Python
275 lines
11 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.
|
|
"""
|
|
attention_outputs = {}
|
|
|
|
def hook_fn(module, input, output):
|
|
# Store attention outputs for concept extraction
|
|
if hasattr(module, 'attention'):
|
|
attention_outputs[layer_name] = output
|
|
|
|
# Get the actual model from ModelPatcher
|
|
actual_model = getattr(self.model, 'model', self.model)
|
|
|
|
# Register hooks on attention layers
|
|
try:
|
|
for name, module in actual_model.named_modules():
|
|
if 'attention' in name.lower() or 'attn' in name.lower():
|
|
hook = module.register_forward_hook(hook_fn)
|
|
self.attention_hooks.append(hook)
|
|
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 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:
|
|
# For now, create mock saliency maps for testing
|
|
# In a real implementation, this would use the actual model
|
|
saliency_maps = self._create_mock_saliency_maps(image, concepts)
|
|
return saliency_maps
|
|
except Exception as e:
|
|
logger.error(f"Error processing image: {e}")
|
|
return {}
|
|
finally:
|
|
self.concept_attention.cleanup_hooks()
|
|
|
|
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.
|
|
"""
|
|
saliency_maps = {}
|
|
h, w = image.shape[1], image.shape[2]
|
|
|
|
for i, concept in enumerate(concepts):
|
|
# 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()
|