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
308 lines
8.9 KiB
Python
308 lines
8.9 KiB
Python
"""
|
|
Utility functions for ConceptAttention ComfyUI nodes
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
from typing import Dict, List, Tuple, Optional
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def resize_saliency_map(saliency_map: torch.Tensor, target_size: Tuple[int, int]) -> torch.Tensor:
|
|
"""
|
|
Resize saliency map to target size.
|
|
|
|
Args:
|
|
saliency_map: Input saliency map tensor
|
|
target_size: Target (height, width) tuple
|
|
|
|
Returns:
|
|
Resized saliency map
|
|
"""
|
|
if len(saliency_map.shape) == 2:
|
|
saliency_map = saliency_map.unsqueeze(0).unsqueeze(0)
|
|
elif len(saliency_map.shape) == 3:
|
|
saliency_map = saliency_map.unsqueeze(0)
|
|
|
|
resized = F.interpolate(
|
|
saliency_map,
|
|
size=target_size,
|
|
mode='bilinear',
|
|
align_corners=False
|
|
)
|
|
|
|
return resized.squeeze()
|
|
|
|
def normalize_saliency_map(saliency_map: torch.Tensor, method: str = "minmax") -> torch.Tensor:
|
|
"""
|
|
Normalize saliency map.
|
|
|
|
Args:
|
|
saliency_map: Input saliency map
|
|
method: Normalization method ("minmax", "zscore", "softmax")
|
|
|
|
Returns:
|
|
Normalized saliency map
|
|
"""
|
|
if method == "minmax":
|
|
min_val = saliency_map.min()
|
|
max_val = saliency_map.max()
|
|
if max_val > min_val:
|
|
return (saliency_map - min_val) / (max_val - min_val)
|
|
else:
|
|
return saliency_map
|
|
|
|
elif method == "zscore":
|
|
mean_val = saliency_map.mean()
|
|
std_val = saliency_map.std()
|
|
if std_val > 0:
|
|
return (saliency_map - mean_val) / std_val
|
|
else:
|
|
return saliency_map
|
|
|
|
elif method == "softmax":
|
|
return F.softmax(saliency_map.flatten(), dim=0).reshape(saliency_map.shape)
|
|
|
|
else:
|
|
raise ValueError(f"Unknown normalization method: {method}")
|
|
|
|
def create_concept_colors(concepts: List[str]) -> Dict[str, Tuple[float, float, float]]:
|
|
"""
|
|
Create color mapping for concepts.
|
|
|
|
Args:
|
|
concepts: List of concept names
|
|
|
|
Returns:
|
|
Dictionary mapping concept names to RGB colors
|
|
"""
|
|
# Predefined colors
|
|
predefined_colors = {
|
|
'woman': (1.0, 0.0, 0.0), # Red
|
|
'cat': (0.0, 1.0, 0.0), # Green
|
|
'white': (1.0, 1.0, 1.0), # White
|
|
'lines': (0.0, 0.0, 1.0), # Blue
|
|
'cane': (1.0, 0.0, 1.0), # Magenta
|
|
'person': (1.0, 0.0, 0.0), # Red (fallback)
|
|
'car': (0.0, 1.0, 0.0), # Green (fallback)
|
|
'tree': (0.0, 0.0, 1.0), # Blue (fallback)
|
|
'sky': (1.0, 1.0, 0.0), # Yellow (fallback)
|
|
'building': (1.0, 0.0, 1.0), # Magenta (fallback)
|
|
}
|
|
|
|
colors = {}
|
|
used_colors = set()
|
|
|
|
# Assign predefined colors first
|
|
for concept in concepts:
|
|
concept_lower = concept.lower()
|
|
if concept_lower in predefined_colors:
|
|
colors[concept] = predefined_colors[concept_lower]
|
|
used_colors.add(predefined_colors[concept_lower])
|
|
|
|
# Generate colors for remaining concepts
|
|
remaining_concepts = [c for c in concepts if c not in colors]
|
|
for i, concept in enumerate(remaining_concepts):
|
|
# Generate a unique color
|
|
hue = (i * 137.5) % 360 # Golden angle for good distribution
|
|
rgb = hsv_to_rgb(hue, 0.8, 0.9)
|
|
colors[concept] = rgb
|
|
|
|
return colors
|
|
|
|
def hsv_to_rgb(h: float, s: float, v: float) -> Tuple[float, float, float]:
|
|
"""
|
|
Convert HSV to RGB.
|
|
|
|
Args:
|
|
h: Hue (0-360)
|
|
s: Saturation (0-1)
|
|
v: Value (0-1)
|
|
|
|
Returns:
|
|
RGB tuple (0-1)
|
|
"""
|
|
h = h / 360.0
|
|
c = v * s
|
|
x = c * (1 - abs((h * 6) % 2 - 1))
|
|
m = v - c
|
|
|
|
if h < 1/6:
|
|
r, g, b = c, x, 0
|
|
elif h < 2/6:
|
|
r, g, b = x, c, 0
|
|
elif h < 3/6:
|
|
r, g, b = 0, c, x
|
|
elif h < 4/6:
|
|
r, g, b = 0, x, c
|
|
elif h < 5/6:
|
|
r, g, b = x, 0, c
|
|
else:
|
|
r, g, b = c, 0, x
|
|
|
|
return (r + m, g + m, b + m)
|
|
|
|
def apply_gaussian_blur(saliency_map: torch.Tensor, kernel_size: int = 5, sigma: float = 1.0) -> torch.Tensor:
|
|
"""
|
|
Apply Gaussian blur to saliency map.
|
|
|
|
Args:
|
|
saliency_map: Input saliency map
|
|
kernel_size: Size of Gaussian kernel
|
|
sigma: Standard deviation of Gaussian kernel
|
|
|
|
Returns:
|
|
Blurred saliency map
|
|
"""
|
|
if len(saliency_map.shape) == 2:
|
|
saliency_map = saliency_map.unsqueeze(0).unsqueeze(0)
|
|
elif len(saliency_map.shape) == 3:
|
|
saliency_map = saliency_map.unsqueeze(0)
|
|
|
|
# Create Gaussian kernel
|
|
kernel = create_gaussian_kernel(kernel_size, sigma, saliency_map.device)
|
|
|
|
# Apply convolution
|
|
blurred = F.conv2d(saliency_map, kernel, padding=kernel_size//2)
|
|
|
|
return blurred.squeeze()
|
|
|
|
def create_gaussian_kernel(kernel_size: int, sigma: float, device: torch.device) -> torch.Tensor:
|
|
"""
|
|
Create Gaussian kernel for blurring.
|
|
|
|
Args:
|
|
kernel_size: Size of kernel
|
|
sigma: Standard deviation
|
|
device: Device to create kernel on
|
|
|
|
Returns:
|
|
Gaussian kernel tensor
|
|
"""
|
|
# Create coordinate grids
|
|
x = torch.arange(kernel_size, dtype=torch.float32, device=device)
|
|
y = torch.arange(kernel_size, dtype=torch.float32, device=device)
|
|
x, y = torch.meshgrid(x, y, indexing='ij')
|
|
|
|
# Center the kernel
|
|
center = kernel_size // 2
|
|
x = x - center
|
|
y = y - center
|
|
|
|
# Calculate Gaussian
|
|
kernel = torch.exp(-(x**2 + y**2) / (2 * sigma**2))
|
|
kernel = kernel / kernel.sum()
|
|
|
|
# Reshape for conv2d
|
|
kernel = kernel.unsqueeze(0).unsqueeze(0)
|
|
|
|
return kernel
|
|
|
|
def compute_attention_statistics(saliency_maps: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, float]]:
|
|
"""
|
|
Compute statistics for saliency maps.
|
|
|
|
Args:
|
|
saliency_maps: Dictionary of concept -> saliency map
|
|
|
|
Returns:
|
|
Dictionary of statistics for each concept
|
|
"""
|
|
stats = {}
|
|
|
|
for concept, saliency_map in saliency_maps.items():
|
|
saliency_flat = saliency_map.flatten()
|
|
|
|
stats[concept] = {
|
|
'mean': saliency_flat.mean().item(),
|
|
'std': saliency_flat.std().item(),
|
|
'min': saliency_flat.min().item(),
|
|
'max': saliency_flat.max().item(),
|
|
'median': saliency_flat.median().item(),
|
|
'sum': saliency_flat.sum().item(),
|
|
}
|
|
|
|
return stats
|
|
|
|
def threshold_saliency_map(saliency_map: torch.Tensor, threshold: float, method: str = "absolute") -> torch.Tensor:
|
|
"""
|
|
Apply threshold to saliency map.
|
|
|
|
Args:
|
|
saliency_map: Input saliency map
|
|
threshold: Threshold value
|
|
method: Threshold method ("absolute", "percentile", "otsu")
|
|
|
|
Returns:
|
|
Thresholded saliency map
|
|
"""
|
|
if method == "absolute":
|
|
return (saliency_map > threshold).float()
|
|
|
|
elif method == "percentile":
|
|
threshold_val = torch.quantile(saliency_map.flatten(), threshold)
|
|
return (saliency_map > threshold_val).float()
|
|
|
|
elif method == "otsu":
|
|
# Simple Otsu thresholding
|
|
saliency_flat = saliency_map.flatten()
|
|
hist = torch.histc(saliency_flat, bins=256, min=0, max=1)
|
|
hist = hist / hist.sum()
|
|
|
|
# Find optimal threshold
|
|
best_threshold = 0
|
|
best_variance = 0
|
|
|
|
for t in range(256):
|
|
w0 = hist[:t+1].sum()
|
|
w1 = hist[t+1:].sum()
|
|
|
|
if w0 == 0 or w1 == 0:
|
|
continue
|
|
|
|
mu0 = (torch.arange(t+1, dtype=torch.float32) * hist[:t+1]).sum() / w0
|
|
mu1 = (torch.arange(t+1, 256, dtype=torch.float32) * hist[t+1:]).sum() / w1
|
|
|
|
variance = w0 * w1 * (mu0 - mu1) ** 2
|
|
|
|
if variance > best_variance:
|
|
best_variance = variance
|
|
best_threshold = t / 255.0
|
|
|
|
return (saliency_map > best_threshold).float()
|
|
|
|
else:
|
|
raise ValueError(f"Unknown threshold method: {method}")
|
|
|
|
def save_saliency_maps(saliency_maps: Dict[str, torch.Tensor], output_dir: str, prefix: str = "concept"):
|
|
"""
|
|
Save saliency maps to files.
|
|
|
|
Args:
|
|
saliency_maps: Dictionary of concept -> saliency map
|
|
output_dir: Output directory
|
|
prefix: Filename prefix
|
|
"""
|
|
import os
|
|
from PIL import Image
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
for concept, saliency_map in saliency_maps.items():
|
|
# Convert to PIL Image
|
|
saliency_np = saliency_map.cpu().numpy()
|
|
saliency_np = (saliency_np * 255).astype(np.uint8)
|
|
|
|
# Create PIL Image
|
|
img = Image.fromarray(saliency_np, mode='L')
|
|
|
|
# Save
|
|
filename = f"{prefix}_{concept}.png"
|
|
filepath = os.path.join(output_dir, filename)
|
|
img.save(filepath)
|
|
|
|
logger.info(f"Saved saliency map: {filepath}")
|
|
|