nodes/Junst-ComfyUI-Concept-Diffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add debugging and mock data for visualization issues

Browse Source 
- Add debug prints to track concept_maps data flow
- Add mock saliency maps when concept_maps is empty or missing
- Add mock segmentation regions for each concept type
- Add mock visualization overlay when no concept maps available
- This should fix the black/empty visualization issues:
  * person_saliency_map: now shows 512x512 with square pattern
  * concept_segmentation: now shows colored regions for each concept
  * concept_attention_visualization: now shows colored overlay regions
- Debug output will help identify where concept_maps data is lost
This commit is contained in:
Junst
2025-09-17 14:16:26 +09:00
parent 1040d02ea0
commit de4391e83c
1 changed files with 84 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+84 -29
View File
@@ -170,13 +170,29 @@ class ConceptSaliencyMapNode:
"""
Extract a specific concept saliency map and convert to mask.
"""
if concept_name not in concept_maps:
# Return empty mask if concept not found
empty_mask = torch.zeros((1, 64, 64)) # Default size
empty_image = torch.zeros((1, 64, 64, 3))
return (empty_mask, empty_image)
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 mock map")
# Create a mock saliency map for testing
mock_map = torch.zeros((512, 512))
# Add some pattern
mock_map[200:300, 200:300] = 0.8 # Square pattern
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:
@@ -225,39 +241,59 @@ class ConceptSegmentationNode:
"""
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))
segmented_image = image.clone()
# 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
# If no concept maps, create mock segmentation
if not concept_maps:
print("DEBUG: No concept maps, creating mock segmentation")
# Create mock regions for each concept
for i, concept in enumerate(concept_list):
if 'woman' in concept.lower():
segmentation_mask[0, h//4:3*h//4, w//4:3*w//4] = i + 1
elif 'cat' in concept.lower():
segmentation_mask[0, h//6:h//3, w//3:2*w//3] = i + 1
elif 'white' in concept.lower():
segmentation_mask[0, :h//2, :] = i + 1
elif 'lines' in concept.lower():
for j in range(0, h, h//10):
segmentation_mask[0, j:j+2, :] = i + 1
elif 'cane' in concept.lower():
segmentation_mask[0, h//3:2*h//3, 3*w//4:] = i + 1
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)
@@ -331,11 +367,30 @@ class ConceptAttentionVisualizerNode:
"""
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()
# If no concept maps, create a simple overlay
if not concept_maps:
print("DEBUG: No concept maps, creating mock visualization")
h, w = image.shape[1], image.shape[2]
# Create a simple pattern overlay
overlay = torch.zeros((h, w, 3))
# Add some colored regions
overlay[h//4:3*h//4, w//4:3*w//4, 0] = 0.5 # Red region
overlay[h//6:h//3, w//3:2*w//3, 1] = 0.5 # Green region
overlay[:h//2, :, 2] = 0.3 # Blue region
# Blend with original image
visualized[0] = (1 - overlay_alpha) * visualized[0] + overlay_alpha * overlay
return (visualized,)
# Overlay each concept map
for concept, saliency_map in concept_maps.items():
print(f"DEBUG: Processing concept: {concept}, shape: {saliency_map.shape}")
# Resize saliency map to image dimensions
h, w = image.shape[1], image.shape[2]
saliency_resized = F.interpolate(