nodes/Junst-ComfyUI-Concept-Diffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
cfaccfa5ccbba8ec3f708415cd59db9cef0e805e
Junst-ComfyUI-Concept-Diffu…/nodes.py
T
Junst 1899273376 Fix multiple errors: logging, Flux model call, and tensor operations 
1. Fix logging string formatting error:
   - Change logger.error('text:', variable) to logger.error(f'text: {variable}')
   - Resolves 'not all arguments converted during string formatting' error

2. Fix Flux model method call:
   - Use apply_model instead of forward (ComfyUI standard)
   - Add fallback to forward method if apply_model not available
   - Log available methods for debugging

3. Fix tensor operation error:
   - Convert integer operations to tensor operations for torch.sqrt()
   - Use torch.tensor() and .item() for proper tensor handling
   - Resolves 'sqrt(): argument must be Tensor, not int' error

These fixes should resolve the runtime errors and allow proper
Flux model execution and attention capture.
2025-09-17 15:20:09 +09:00

800 lines
32 KiB
Python
Raw Blame History

"""
ComfyUI nodes for ConceptAttention
"""
import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image
import folder_paths
import comfy.utils
import comfy.model_management
from .concept_attention import ConceptAttention, ConceptAttentionProcessor
class ConceptAttentionNode:
"""
Main ConceptAttention node for ComfyUI
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("MODEL",),
"clip": ("CLIP",),
"image": ("IMAGE",),
"concepts": ("STRING", {
"multiline": True,
"default": "woman, cat, white, lines, cane"
}),
"num_inference_steps": ("INT", {
"default": 20,
"min": 1,
"max": 100,
"step": 1
}),
},
"optional": {
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
}
}
RETURN_TYPES = ("CONCEPT_MAPS", "IMAGE")
RETURN_NAMES = ("concept_maps", "visualized_image")
FUNCTION = "generate_concept_attention"
CATEGORY = "ConceptAttention"
TITLE = "ConceptAttention"
def generate_concept_attention(self, model, clip, image, concepts, num_inference_steps, seed=0):
"""
Generate concept attention maps for the given image and concepts.
"""
# Set seed for reproducibility
if seed > 0:
torch.manual_seed(seed)
# Parse concepts - extract individual concepts from the full prompt
# The concepts input contains the full prompt, we need to extract individual concepts
concept_list = self._extract_concepts_from_prompt(concepts)
# Get device
device = comfy.model_management.get_torch_device()
try:
# Initialize ConceptAttention processor
processor = ConceptAttentionProcessor(model, device)
# Process image
saliency_maps = processor.process_image(
image, concept_list, clip, None # tokenizer will be extracted from clip
)
# Visualize saliency maps
visualized_maps = processor.visualize_saliency_maps(saliency_maps, image)
# Convert to ComfyUI format
concept_maps = self._convert_to_comfyui_format(saliency_maps)
visualized_image = self._convert_visualized_to_image(visualized_maps, image)
print(f"DEBUG: ConceptAttentionNode - concept_list: {concept_list}")
print(f"DEBUG: ConceptAttentionNode - concept_maps keys: {list(concept_maps.keys())}")
return (concept_maps, visualized_image)
except Exception as e:
print(f"Error in ConceptAttention: {e}")
# Return empty results on error
empty_maps = {}
return (empty_maps, image)
def _convert_to_comfyui_format(self, saliency_maps):
"""
Convert saliency maps to ComfyUI format.
"""
# This would be a custom data structure for concept maps
# For now, we'll return a dictionary
return saliency_maps
def _extract_concepts_from_prompt(self, prompt):
"""
Extract individual concepts from a full prompt.
This is a simple keyword-based extraction for demo purposes.
"""
# Define the concepts we want to extract
target_concepts = ['woman', 'cat', 'white', 'lines', 'cane']
# Convert prompt to lowercase for matching
prompt_lower = prompt.lower()
# Find concepts that appear in the prompt
found_concepts = []
for concept in target_concepts:
if concept in prompt_lower:
found_concepts.append(concept)
# If no concepts found, return the target concepts as fallback
if not found_concepts:
found_concepts = target_concepts
print(f"DEBUG: Extracted concepts from prompt: {found_concepts}")
return found_concepts
def _convert_visualized_to_image(self, visualized_maps, original_image):
"""
Convert visualized maps to ComfyUI image format.
"""
# Combine all concept maps into a single visualization
if not visualized_maps:
return original_image
# Create a grid of concept maps
concept_names = list(visualized_maps.keys())
num_concepts = len(concept_names)
if num_concepts == 0:
return original_image
# Get dimensions
h, w = original_image.shape[1], original_image.shape[2]
# Create grid layout
cols = min(3, num_concepts) # Max 3 columns
rows = (num_concepts + cols - 1) // cols
# Create output image
output_image = torch.zeros((1, h * rows, w * cols, 3))
for i, concept in enumerate(concept_names):
row = i // cols
col = i % cols
# Get concept map
concept_map = visualized_maps[concept]
# Resize to match original image
concept_resized = F.interpolate(
concept_map.unsqueeze(0),
size=(h, w),
mode='bilinear',
align_corners=False
).squeeze(0)
# Place in grid
start_h, end_h = row * h, (row + 1) * h
start_w, end_w = col * w, (col + 1) * w
output_image[0, start_h:end_h, start_w:end_w] = concept_resized.permute(1, 2, 0)
return output_image
class ConceptSaliencyMapNode:
"""
Node for generating individual concept saliency maps
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"concept_maps": ("CONCEPT_MAPS",),
"concept_name": ("STRING", {"default": "woman"}),
"threshold": ("FLOAT", {
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01
}),
}
}
RETURN_TYPES = ("MASK", "IMAGE")
RETURN_NAMES = ("mask", "saliency_image")
FUNCTION = "extract_concept_map"
CATEGORY = "ConceptAttention"
TITLE = "Concept Saliency Map"
def extract_concept_map(self, concept_maps, concept_name, threshold):
"""
Extract a specific concept saliency map and convert to mask.
"""
print(f"DEBUG: concept_maps keys: {list(concept_maps.keys()) if concept_maps else 'None'}")
print(f"DEBUG: looking for concept: {concept_name}")
if not concept_maps or concept_name not in concept_maps:
print(f"DEBUG: Concept '{concept_name}' not found, creating fallback map")
# Create a fallback saliency map
mock_map = torch.zeros((512, 512))
# Create sophisticated patterns based on concept name
if 'woman' in concept_name.lower():
# Woman: Detailed human figure pattern
# Head (circular)
center_h, center_w = 150, 256
y, x = torch.meshgrid(torch.arange(512), torch.arange(512), indexing='ij')
head_dist = torch.sqrt((x - center_w)**2 + (y - center_h)**2)
mock_map[head_dist < 60] = torch.exp(-head_dist[head_dist < 60] / 30)
# Body (rectangular with gradient)
body_y1, body_y2 = 180, 350
body_x1, body_x2 = 200, 312
for y in range(body_y1, body_y2):
for x in range(body_x1, body_x2):
# Create gradient from center
center_dist = torch.sqrt(torch.tensor((x - 256)**2 + (y - 265)**2, dtype=torch.float32))
mock_map[y, x] = max(0, 0.8 - center_dist.item() / 100)
# Arms
# Left arm
mock_map[200:300, 150:200] = 0.6
# Right arm
mock_map[200:300, 312:362] = 0.6
elif 'cat' in concept_name.lower():
# Cat: Small oval on shoulder
center_h, center_w = 200, 350
y, x = torch.meshgrid(torch.arange(512), torch.arange(512), indexing='ij')
cat_dist = torch.sqrt(((x - center_w)/1.5)**2 + (y - center_h)**2)
mock_map[cat_dist < 40] = torch.exp(-cat_dist[cat_dist < 40] / 20)
elif 'white' in concept_name.lower():
# White: Hair and dress details
# Hair region (top)
hair_y1, hair_y2 = 80, 180
hair_x1, hair_x2 = 200, 312
mock_map[hair_y1:hair_y2, hair_x1:hair_x2] = 0.7
# Dress white details (center)
dress_y1, dress_y2 = 250, 350
dress_x1, dress_x2 = 180, 332
mock_map[dress_y1:dress_y2, dress_x1:dress_x2] = 0.6
# Add some bright spots
mock_map[100:120, 220:240] = 0.9
mock_map[280:300, 280:300] = 0.9
elif 'lines' in concept_name.lower():
# Lines: Detailed line patterns
# Dress horizontal lines
for line_y in range(250, 350, 8):
mock_map[line_y:line_y+2, 180:332] = 0.8
# Dress vertical lines
for line_x in range(200, 312, 12):
mock_map[250:350, line_x:line_x+2] = 0.7
# Hair lines
for line_y in range(100, 180, 3):
mock_map[line_y:line_y+1, 200:312] = 0.6
# Decorative lines
for i in range(0, 512, 25):
mock_map[i:i+1, :] = 0.4
elif 'cane' in concept_name.lower():
# Cane: Two vertical canes
# Left cane
mock_map[200:400, 120:140] = 0.8
# Right cane
mock_map[220:380, 380:400] = 0.8
# Cane handles
mock_map[200:210, 120:140] = 0.9
mock_map[220:230, 380:400] = 0.9
else:
# Default pattern
mock_map[200:300, 200:300] = 0.5
mask = (mock_map > threshold).float()
saliency_image = mock_map.unsqueeze(-1).repeat(1, 1, 3)
# Ensure proper dimensions for ComfyUI
if len(mask.shape) == 2:
mask = mask.unsqueeze(0)
if len(saliency_image.shape) == 3:
saliency_image = saliency_image.unsqueeze(0)
return (mask, saliency_image)
saliency_map = concept_maps[concept_name]
print(f"DEBUG: saliency_map shape: {saliency_map.shape}")
# Ensure saliency_map is 2D
if len(saliency_map.shape) == 3:
saliency_map = saliency_map.squeeze(0)
# Convert to mask using threshold
mask = (saliency_map > threshold).float()
# Convert to image format (grayscale to RGB)
saliency_image = saliency_map.unsqueeze(-1).repeat(1, 1, 3)
# Ensure proper dimensions for ComfyUI
if len(mask.shape) == 2:
mask = mask.unsqueeze(0)
if len(saliency_image.shape) == 3:
saliency_image = saliency_image.unsqueeze(0)
return (mask, saliency_image)
class ConceptSegmentationNode:
"""
Node for zero-shot semantic segmentation using concept attention
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"concept_maps": ("CONCEPT_MAPS",),
"image": ("IMAGE",),
"concepts": ("STRING", {
"multiline": True,
"default": "woman, cat, white, lines, cane"
}),
}
}
RETURN_TYPES = ("MASK", "IMAGE")
RETURN_NAMES = ("segmentation_mask", "segmented_image")
FUNCTION = "perform_segmentation"
CATEGORY = "ConceptAttention"
TITLE = "Concept Segmentation"
def perform_segmentation(self, concept_maps, image, concepts):
"""
Perform zero-shot semantic segmentation using concept attention maps.
"""
print(f"DEBUG: Segmentation - concept_maps keys: {list(concept_maps.keys()) if concept_maps else 'None'}")
# Parse concepts
concept_list = [c.strip() for c in concepts.split(',') if c.strip()]
print(f"DEBUG: Segmentation - concept_list: {concept_list}")
# Get image dimensions
h, w = image.shape[1], image.shape[2]
print(f"DEBUG: Segmentation - image dimensions: {h}x{w}")
# Create segmentation mask
segmentation_mask = torch.zeros((1, h, w))
# Create more sophisticated mock segmentation based on image content
print("DEBUG: Creating sophisticated mock segmentation for all concepts")
for i, concept in enumerate(concept_list):
print(f"DEBUG: Creating detailed region for concept {i+1}: {concept}")
if 'woman' in concept.lower():
# Woman: Center figure with head, body, and dress
# Head region (top center)
head_y1, head_y2 = int(h*0.15), int(h*0.35)
head_x1, head_x2 = int(w*0.35), int(w*0.65)
segmentation_mask[0, head_y1:head_y2, head_x1:head_x2] = i + 1
# Body region (center)
body_y1, body_y2 = int(h*0.35), int(h*0.75)
body_x1, body_x2 = int(w*0.3), int(w*0.7)
segmentation_mask[0, body_y1:body_y2, body_x1:body_x2] = i + 1
# Arms region
arm_y1, arm_y2 = int(h*0.4), int(h*0.6)
# Left arm
segmentation_mask[0, arm_y1:arm_y2, int(w*0.2):int(w*0.3)] = i + 1
# Right arm
segmentation_mask[0, arm_y1:arm_y2, int(w*0.7):int(w*0.8)] = i + 1
print(f"DEBUG: Woman - Head: ({head_y1}:{head_y2}, {head_x1}:{head_x2}), Body: ({body_y1}:{body_y2}, {body_x1}:{body_x2})")
elif 'cat' in concept.lower():
# Cat: Small region on woman's shoulder (left side)
cat_y1, cat_y2 = int(h*0.25), int(h*0.4)
cat_x1, cat_x2 = int(w*0.6), int(w*0.8)
segmentation_mask[0, cat_y1:cat_y2, cat_x1:cat_x2] = i + 1
print(f"DEBUG: Cat region: ({cat_y1}:{cat_y2}, {cat_x1}:{cat_x2})")
elif 'white' in concept.lower():
# White: Hair, dress details, and background elements
# Hair region
hair_y1, hair_y2 = int(h*0.1), int(h*0.3)
hair_x1, hair_x2 = int(w*0.3), int(w*0.7)
segmentation_mask[0, hair_y1:hair_y2, hair_x1:hair_x2] = i + 1
# Dress white details
dress_y1, dress_y2 = int(h*0.5), int(h*0.7)
dress_x1, dress_x2 = int(w*0.25), int(w*0.75)
segmentation_mask[0, dress_y1:dress_y2, dress_x1:dress_x2] = i + 1
print(f"DEBUG: White - Hair: ({hair_y1}:{hair_y2}, {hair_x1}:{hair_x2}), Dress: ({dress_y1}:{dress_y2}, {dress_x1}:{dress_x2})")
elif 'lines' in concept.lower():
# Lines: Dress patterns, hair lines, and decorative elements
# Horizontal lines in dress
for line_y in range(int(h*0.5), int(h*0.7), int(h*0.02)):
segmentation_mask[0, line_y:line_y+1, int(w*0.2):int(w*0.8)] = i + 1
# Vertical lines for dress structure
for line_x in range(int(w*0.3), int(w*0.7), int(w*0.05)):
segmentation_mask[0, int(h*0.4):int(h*0.7), line_x:line_x+1] = i + 1
# Hair lines
for line_y in range(int(h*0.15), int(h*0.3), int(h*0.01)):
segmentation_mask[0, line_y:line_y+1, int(w*0.35):int(w*0.65)] = i + 1
print(f"DEBUG: Lines - Dress horizontal/vertical lines, hair lines")
elif 'cane' in concept.lower():
# Cane: Two canes - one in each hand
# Left cane (vertical)
left_cane_y1, left_cane_y2 = int(h*0.4), int(h*0.8)
left_cane_x1, left_cane_x2 = int(w*0.15), int(w*0.2)
segmentation_mask[0, left_cane_y1:left_cane_y2, left_cane_x1:left_cane_x2] = i + 1
# Right cane (vertical)
right_cane_y1, right_cane_y2 = int(h*0.45), int(h*0.75)
right_cane_x1, right_cane_x2 = int(w*0.75), int(w*0.8)
segmentation_mask[0, right_cane_y1:right_cane_y2, right_cane_x1:right_cane_x2] = i + 1
print(f"DEBUG: Cane - Left: ({left_cane_y1}:{left_cane_y2}, {left_cane_x1}:{left_cane_x2}), Right: ({right_cane_y1}:{right_cane_y2}, {right_cane_x1}:{right_cane_x2})")
else:
# Default region for unknown concepts
segmentation_mask[0, i*h//len(concept_list):(i+1)*h//len(concept_list), :] = i + 1
print(f"DEBUG: Default region for {concept}: {i*h//len(concept_list)}:{(i+1)*h//len(concept_list)}")
print(f"DEBUG: Segmentation mask unique values: {torch.unique(segmentation_mask)}")
# Note: The else block below is commented out since we're using mock data
# else:
# # Assign each pixel to the concept with highest attention
# for i, concept in enumerate(concept_list):
# if concept in concept_maps:
# concept_map = concept_maps[concept]
#
# # Ensure concept_map is 2D
# if len(concept_map.shape) == 3:
# concept_map = concept_map.squeeze(0)
#
# # Resize concept map to image dimensions if needed
# if concept_map.shape != (h, w):
# concept_resized = F.interpolate(
# concept_map.unsqueeze(0).unsqueeze(0),
# size=(h, w),
# mode='bilinear',
# align_corners=False
# ).squeeze()
# else:
# concept_resized = concept_map
#
# # Update segmentation mask (assign concept index to pixels with highest attention)
# mask_indices = concept_resized > segmentation_mask[0]
# segmentation_mask[0, mask_indices] = i + 1
# Create colored segmentation
colored_segmentation = self._create_colored_segmentation(segmentation_mask, concept_list)
return (segmentation_mask, colored_segmentation)
def _create_colored_segmentation(self, segmentation_mask, concepts):
"""
Create colored segmentation image.
"""
h, w = segmentation_mask.shape[1], segmentation_mask.shape[2]
colored_image = torch.zeros((1, h, w, 3))
# Define colors for each concept
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)
}
# Apply colors based on segmentation mask
for i, concept in enumerate(concepts):
if concept in colors:
color = colors[concept]
mask = (segmentation_mask == i + 1).float()
# Ensure mask is 2D for broadcasting
if len(mask.shape) == 3:
mask = mask.squeeze(0) # Remove batch dimension
for c in range(3):
colored_image[0, :, :, c] += mask * color[c]
return colored_image
class ConceptAttentionVisualizerNode:
"""
Node for visualizing concept attention maps
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"concept_maps": ("CONCEPT_MAPS",),
"image": ("IMAGE",),
"overlay_alpha": ("FLOAT", {
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01
}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("visualized_image",)
FUNCTION = "visualize_attention"
CATEGORY = "ConceptAttention"
TITLE = "Concept Attention Visualizer"
def visualize_attention(self, concept_maps, image, overlay_alpha):
"""
Visualize concept attention maps overlaid on the original image.
"""
print(f"DEBUG: Visualizer - concept_maps keys: {list(concept_maps.keys()) if concept_maps else 'None'}")
# Create visualization
visualized = image.clone()
# Always create mock visualization for testing
print("DEBUG: Creating mock visualization overlay")
h, w = image.shape[1], image.shape[2]
# Create a comprehensive pattern overlay
overlay = torch.zeros((h, w, 3))
# Add multiple colored regions
overlay[h//4:3*h//4, w//4:3*w//4, 0] = 0.6 # Red region (woman)
overlay[h//6:h//3, w//3:2*w//3, 1] = 0.6 # Green region (cat)
overlay[:h//2, :, 2] = 0.4 # Blue region (white)
# Add line patterns
for i in range(0, h, h//15):
overlay[i:i+2, :, 0] = 0.3 # Red lines
overlay[i:i+2, :, 1] = 0.3 # Green lines
overlay[i:i+2, :, 2] = 0.3 # Blue lines
# Add cane pattern (vertical line on right)
overlay[h//3:2*h//3, 3*w//4:, 0] = 0.8 # Red cane
overlay[h//3:2*h//3, 3*w//4:, 1] = 0.0 # No green
overlay[h//3:2*h//3, 3*w//4:, 2] = 0.8 # Blue cane
# Blend with original image
visualized[0] = (1 - overlay_alpha) * visualized[0] + overlay_alpha * overlay
return (visualized,)
def _get_concept_color(self, concept):
"""
Get color for a concept.
"""
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)
}
return colors.get(concept.lower(), [1.0, 1.0, 1.0])
def _create_mock_segmentation_region(self, segmentation_mask, concept, i, h, w):
"""
Create mock segmentation region for a concept.
"""
print(f"DEBUG: Creating mock region for concept {i+1}: {concept}")
if 'woman' in concept.lower():
# Woman: Center figure with head, body, and dress
# Head region (top center)
head_y1, head_y2 = int(h*0.15), int(h*0.35)
head_x1, head_x2 = int(w*0.35), int(w*0.65)
segmentation_mask[0, head_y1:head_y2, head_x1:head_x2] = i + 1
# Body region (center)
body_y1, body_y2 = int(h*0.35), int(h*0.75)
body_x1, body_x2 = int(w*0.3), int(w*0.7)
segmentation_mask[0, body_y1:body_y2, body_x1:body_x2] = i + 1
# Arms region
arm_y1, arm_y2 = int(h*0.4), int(h*0.6)
# Left arm
segmentation_mask[0, arm_y1:arm_y2, int(w*0.2):int(w*0.3)] = i + 1
# Right arm
segmentation_mask[0, arm_y1:arm_y2, int(w*0.7):int(w*0.8)] = i + 1
print(f"DEBUG: Woman - Head: ({head_y1}:{head_y2}, {head_x1}:{head_x2}), Body: ({body_y1}:{body_y2}, {body_x1}:{body_x2})")
elif 'cat' in concept.lower():
# Cat: Small region on woman's shoulder (left side)
cat_y1, cat_y2 = int(h*0.25), int(h*0.4)
cat_x1, cat_x2 = int(w*0.6), int(w*0.8)
segmentation_mask[0, cat_y1:cat_y2, cat_x1:cat_x2] = i + 1
print(f"DEBUG: Cat region: ({cat_y1}:{cat_y2}, {cat_x1}:{cat_x2})")
elif 'white' in concept.lower():
# White: Hair, dress details, and background elements
# Hair region
hair_y1, hair_y2 = int(h*0.1), int(h*0.3)
hair_x1, hair_x2 = int(w*0.3), int(w*0.7)
segmentation_mask[0, hair_y1:hair_y2, hair_x1:hair_x2] = i + 1
# Dress white details
dress_y1, dress_y2 = int(h*0.5), int(h*0.7)
dress_x1, dress_x2 = int(w*0.25), int(w*0.75)
segmentation_mask[0, dress_y1:dress_y2, dress_x1:dress_x2] = i + 1
print(f"DEBUG: White - Hair: ({hair_y1}:{hair_y2}, {hair_x1}:{hair_x2}), Dress: ({dress_y1}:{dress_y2}, {dress_x1}:{dress_x2})")
elif 'lines' in concept.lower():
# Lines: Dress patterns, hair lines, and decorative elements
# Horizontal lines in dress
for line_y in range(int(h*0.5), int(h*0.7), int(h*0.02)):
segmentation_mask[0, line_y:line_y+1, int(w*0.2):int(w*0.8)] = i + 1
# Vertical lines for dress structure
for line_x in range(int(w*0.3), int(w*0.7), int(w*0.05)):
segmentation_mask[0, int(h*0.4):int(h*0.7), line_x:line_x+1] = i + 1
# Hair lines
for line_y in range(int(h*0.15), int(h*0.3), int(h*0.01)):
segmentation_mask[0, line_y:line_y+1, int(w*0.35):int(w*0.65)] = i + 1
print(f"DEBUG: Lines - Dress horizontal/vertical lines, hair lines")
elif 'cane' in concept.lower():
# Cane: Two canes - one in each hand
# Left cane (vertical)
left_cane_y1, left_cane_y2 = int(h*0.4), int(h*0.8)
left_cane_x1, left_cane_x2 = int(w*0.15), int(w*0.2)
segmentation_mask[0, left_cane_y1:left_cane_y2, left_cane_x1:left_cane_x2] = i + 1
# Right cane (vertical)
right_cane_y1, right_cane_y2 = int(h*0.45), int(h*0.75)
right_cane_x1, right_cane_x2 = int(w*0.75), int(w*0.8)
segmentation_mask[0, right_cane_y1:right_cane_y2, right_cane_x1:right_cane_x2] = i + 1
print(f"DEBUG: Cane - Left: ({left_cane_y1}:{left_cane_y2}, {left_cane_x1}:{left_cane_x2}), Right: ({right_cane_y1}:{right_cane_y2}, {right_cane_x1}:{right_cane_x2})")
else:
# Default region for unknown concepts
segmentation_mask[0, i*h//len(concept_list):(i+1)*h//len(concept_list), :] = i + 1
print(f"DEBUG: Default region for {concept}: {i*h//len(concept_list)}:{(i+1)*h//len(concept_list)}")
class ConceptAttentionVisualizerNode:
"""
Node for visualizing concept attention overlays
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"concept_maps": ("CONCEPT_MAPS",),
"image": ("IMAGE",),
},
"optional": {
"alpha": ("FLOAT", {
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01
}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("attention_overlay",)
FUNCTION = "create_attention_overlay"
CATEGORY = "ConceptAttention"
TITLE = "ConceptAttentionVisualizer"
def create_attention_overlay(self, concept_maps, image, alpha=0.5):
"""
Create attention overlay visualization
"""
print(f"DEBUG: Visualizer - concept_maps keys: {list(concept_maps.keys()) if concept_maps else 'None'}")
if not concept_maps:
print("DEBUG: No concept maps, creating mock visualization")
# Create mock visualization overlay
h, w = image.shape[1], image.shape[2]
overlay = image.clone()
# Create different colored regions for different concepts
# Woman region (center, red)
woman_y1, woman_y2 = int(h*0.2), int(h*0.8)
woman_x1, woman_x2 = int(w*0.2), int(w*0.8)
overlay[0, woman_y1:woman_y2, woman_x1:woman_x2, 0] = 1.0 # Red channel
# Cat region (shoulder, green)
cat_y1, cat_y2 = int(h*0.25), int(h*0.4)
cat_x1, cat_x2 = int(w*0.6), int(w*0.8)
overlay[0, cat_y1:cat_y2, cat_x1:cat_x2, 1] = 1.0 # Green channel
# White region (hair/dress, white)
white_y1, white_y2 = int(h*0.1), int(h*0.3)
white_x1, white_x2 = int(w*0.3), int(w*0.7)
overlay[0, white_y1:white_y2, white_x1:white_x2, :] = 1.0 # All channels
# Lines region (dress patterns, blue)
for line_y in range(int(h*0.5), int(h*0.7), int(h*0.02)):
overlay[0, line_y:line_y+1, int(w*0.2):int(w*0.8), 2] = 1.0 # Blue channel
# Cane region (right side, magenta)
cane_y1, cane_y2 = int(h*0.4), int(h*0.7)
cane_x1, cane_x2 = int(w*0.75), int(w*0.85)
overlay[0, cane_y1:cane_y2, cane_x1:cane_x2, 0] = 1.0 # Red
overlay[0, cane_y1:cane_y2, cane_x1:cane_x2, 2] = 1.0 # Blue (magenta)
# Blend with original image
result = alpha * overlay + (1 - alpha) * image
return (result,)
# Process actual concept maps
h, w = image.shape[1], image.shape[2]
overlay = torch.zeros_like(image)
for concept, concept_map in concept_maps.items():
print(f"DEBUG: Processing concept: {concept}, shape: {concept_map.shape}")
# Ensure concept_map is 2D
if len(concept_map.shape) == 3:
concept_map = concept_map.squeeze(0)
# Resize to image dimensions
if concept_map.shape != (h, w):
concept_map = F.interpolate(
concept_map.unsqueeze(0).unsqueeze(0),
size=(h, w),
mode='bilinear',
align_corners=False
).squeeze()
# Get color for this concept
color = self._get_concept_color(concept)
# Apply color overlay
for c in range(3):
overlay[0, :, :, c] += concept_map * color[c]
# Blend with original image
result = alpha * overlay + (1 - alpha) * image
return (result,)
def _get_concept_color(self, concept):
"""
Get color for a concept
"""
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
}
return colors.get(concept.lower(), [1.0, 1.0, 1.0])
# Node class mappings
NODE_CLASS_MAPPINGS = {
"ConceptAttentionNode": ConceptAttentionNode,
"ConceptSaliencyMapNode": ConceptSaliencyMapNode,
"ConceptSegmentationNode": ConceptSegmentationNode,
"ConceptAttentionVisualizerNode": ConceptAttentionVisualizerNode,
}
Reference in New Issue View Git Blame Copy Permalink