nodes/Junst-ComfyUI-Concept-Diffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
570255e0ab030dae6de180d17f59aa6f44c8be3d
Junst-ComfyUI-Concept-Diffu…/nodes.py
T
Junst 570255e0ab Fix numpy array handling in visualization 
- Fix 'numpy.ndarray' object has no attribute 'float' error
- Ensure attention_map is properly converted to numpy array
- Add proper 2D shape validation for attention maps
- Normalize to 0-255 range for PIL Image.fromarray compatibility
- Improve tensor to numpy conversion and shape handling
2025-09-24 18:36:15 +09:00

482 lines
18 KiB
Python
Raw Blame History

"""
Simplified ComfyUI nodes for Concept Attention based on original structure.
"""
import torch
import numpy as np
from PIL import Image
import logging
from typing import Dict, List, Tuple, Any
from .concept_attention import ConceptAttentionProcessor
logger = logging.getLogger(__name__)
class ConceptAttentionNode:
"""
Simplified ComfyUI node for concept attention.
Based on original ConceptAttention structure.
"""
RETURN_TYPES = ("CONCEPT_MAPS", "IMAGE")
RETURN_NAMES = ("concept_maps", "visualized_image")
FUNCTION = "generate_concept_attention"
CATEGORY = "Concept Attention"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("MODEL",),
"clip": ("CLIP",),
"image": ("IMAGE",),
"prompt": ("STRING", {"multiline": True, "default": "A beautiful landscape"}),
"concept_list": ("STRING", {"multiline": True, "default": "woman, cat, white, lines, cane"}),
}
}
def __init__(self):
self.processor = None
def generate_concept_attention(self, model, clip, image, prompt, concept_list):
"""
Generate concept attention maps using simplified approach.
"""
try:
# Convert concept_list string to list
if isinstance(concept_list, str):
concept_list = [concept.strip() for concept in concept_list.split(',')]
elif isinstance(concept_list, list):
concept_list = [str(concept).strip() for concept in concept_list]
else:
concept_list = []
logger.info(f"DEBUG: Extracted concepts from prompt: {concept_list}")
# Initialize processor if not exists
if self.processor is None:
self.processor = ConceptAttentionProcessor(model, device="cuda")
# Process image to get concept maps
concept_maps = self.processor.process_image(image, concept_list, clip)
# Convert to ComfyUI format
saliency_maps = self._convert_to_comfyui_format(concept_maps)
# Create visualization
visualized_image = self._create_simple_visualization(saliency_maps, image)
logger.info(f"DEBUG: processor.process_image returned: {type(concept_maps)}")
logger.info(f"DEBUG: saliency_maps keys: {list(saliency_maps.keys()) if saliency_maps else None}")
return saliency_maps, visualized_image
except Exception as e:
logger.error(f"Error in SimpleConceptAttentionNode: {e}")
# Return empty results on error
return {}, image
def _convert_to_comfyui_format(self, concept_maps: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
"""
Convert concept maps to ComfyUI format.
"""
try:
if not concept_maps:
logger.warning("WARNING: saliency_maps is empty, returning empty ConceptMaps")
return {}
# Simple conversion - just return the concept maps
saliency_maps = {}
for concept, attention_map in concept_maps.items():
# Ensure tensor is on CPU and convert to numpy
if isinstance(attention_map, torch.Tensor):
attention_map = attention_map.cpu().numpy()
# Normalize to 0-1 range
attention_map = (attention_map - attention_map.min()) / (attention_map.max() - attention_map.min() + 1e-8)
saliency_maps[concept] = attention_map
logger.info(f"DEBUG: Converted concept_maps with keys: {list(saliency_maps.keys())}")
return saliency_maps
except Exception as e:
logger.error(f"Error converting to ComfyUI format: {e}")
return {}
def _create_simple_visualization(self, saliency_maps: Dict[str, torch.Tensor],
image: torch.Tensor) -> torch.Tensor:
"""
Create simple visualization of concept attention maps.
"""
try:
if not saliency_maps:
logger.warning("WARNING: saliency_maps is empty, returning empty ConceptMaps")
return image
# Convert image to numpy if needed
if isinstance(image, torch.Tensor):
img_np = image.squeeze().cpu().numpy()
if len(img_np.shape) == 3 and img_np.shape[0] == 3: # CHW format
img_np = np.transpose(img_np, (1, 2, 0))
else:
img_np = np.array(image)
# Ensure image is in correct format
if len(img_np.shape) == 2: # Grayscale
img_np = np.stack([img_np] * 3, axis=-1)
elif len(img_np.shape) == 3 and img_np.shape[-1] == 1: # Single channel
img_np = np.repeat(img_np, 3, axis=-1)
# Normalize image to 0-1 range
img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min() + 1e-8)
# Create overlay for each concept
overlay = img_np.copy()
colors = [
[1.0, 0.0, 0.0], # Red
[0.0, 1.0, 0.0], # Green
[0.0, 0.0, 1.0], # Blue
[1.0, 1.0, 0.0], # Yellow
[1.0, 0.0, 1.0], # Magenta
]
for i, (concept, attention_map) in enumerate(saliency_maps.items()):
if i < len(colors):
color = colors[i]
# Ensure attention_map is numpy array
if isinstance(attention_map, torch.Tensor):
attention_map = attention_map.cpu().numpy()
# Ensure attention_map is 2D
if len(attention_map.shape) > 2:
attention_map = attention_map.squeeze()
elif len(attention_map.shape) == 1:
attention_map = attention_map.reshape(-1, 1)
# Normalize to 0-255 range for PIL
attention_map = (attention_map - attention_map.min()) / (attention_map.max() - attention_map.min() + 1e-8)
attention_map = (attention_map * 255).astype(np.uint8)
# Resize attention map to image size
attention_resized = np.array(Image.fromarray(attention_map).resize(
(img_np.shape[1], img_np.shape[0]), Image.BILINEAR
))
# Normalize back to 0-1 range
attention_resized = attention_resized.astype(np.float32) / 255.0
# Create colored overlay
for c in range(3):
overlay[:, :, c] += attention_resized * color[c] * 0.3
# Clip to valid range
overlay = np.clip(overlay, 0, 1)
# Convert back to tensor
if len(overlay.shape) == 3:
overlay = np.transpose(overlay, (2, 0, 1)) # HWC to CHW
return torch.from_numpy(overlay).unsqueeze(0)
except Exception as e:
logger.error(f"Error creating visualization: {e}")
return image
def extract_concept_map(self, concept_maps: Dict[str, torch.Tensor],
concept_name: str) -> torch.Tensor:
"""
Extract specific concept map.
"""
try:
if not concept_maps:
raise ValueError(f"Concept '{concept_name}' not found in concept_maps. Available: None")
if concept_name not in concept_maps:
available = list(concept_maps.keys()) if concept_maps else "None"
raise ValueError(f"Concept '{concept_name}' not found in concept_maps. Available: {available}")
return torch.from_numpy(concept_maps[concept_name])
except Exception as e:
logger.error(f"Error extracting concept map: {e}")
raise ValueError(f"Failed to extract concept map for '{concept_name}': {e}")
def perform_segmentation(self, concept_maps: Dict[str, torch.Tensor],
concept_list: List[str], image: torch.Tensor) -> torch.Tensor:
"""
Perform segmentation based on concept maps.
"""
try:
if not concept_maps:
logger.warning("WARNING: saliency_maps is empty, returning empty ConceptMaps")
return image
# Create segmentation mask
segmentation_mask = torch.zeros_like(image[0, 0]) # Single channel mask
for i, concept in enumerate(concept_list):
if concept in concept_maps:
concept_map = torch.from_numpy(concept_maps[concept])
# Assign different values for different concepts
segmentation_mask += concept_map * (i + 1)
# Normalize segmentation mask
if segmentation_mask.max() > 0:
segmentation_mask = segmentation_mask / segmentation_mask.max()
return segmentation_mask.unsqueeze(0).unsqueeze(0)
except Exception as e:
logger.error(f"Error in segmentation: {e}")
return image
class ConceptSaliencyMapNode:
"""
Node for extracting saliency maps from concept attention.
"""
RETURN_TYPES = ("MASK", "IMAGE")
RETURN_NAMES = ("mask", "saliency_image")
FUNCTION = "extract_saliency_map"
CATEGORY = "Concept Attention"
@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}),
}
}
def __init__(self):
pass
def extract_saliency_map(self, concept_maps, concept_name, threshold=0.5):
"""
Extract saliency map for a specific concept.
"""
try:
if not concept_maps or concept_name not in concept_maps:
logger.warning(f"Concept '{concept_name}' not found in concept_maps")
return None, None
concept_map = concept_maps[concept_name]
# Apply threshold
saliency_map = (concept_map > threshold).float()
# Create visualization
saliency_image = self._create_saliency_visualization(concept_map, threshold)
return saliency_map, saliency_image
except Exception as e:
logger.error(f"Error extracting saliency map: {e}")
return None, None
def _create_saliency_visualization(self, concept_map, threshold):
"""
Create visualization of saliency map.
"""
try:
# Normalize concept map
normalized_map = (concept_map - concept_map.min()) / (concept_map.max() - concept_map.min() + 1e-8)
# Apply threshold
thresholded_map = (normalized_map > threshold).float()
# Create colored visualization
visualization = torch.stack([
thresholded_map, # Red channel
torch.zeros_like(thresholded_map), # Green channel
torch.zeros_like(thresholded_map) # Blue channel
], dim=0)
return visualization
except Exception as e:
logger.error(f"Error creating saliency visualization: {e}")
return None
class ConceptSegmentationNode:
"""
Node for performing segmentation based on concept attention.
"""
RETURN_TYPES = ("MASK", "IMAGE")
RETURN_NAMES = ("segmentation_mask", "segmented_image")
FUNCTION = "perform_segmentation"
CATEGORY = "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"}),
}
}
def __init__(self):
pass
def perform_segmentation(self, concept_maps, image, concepts):
"""
Perform segmentation based on concept maps.
"""
try:
if not concept_maps:
logger.warning("No concept maps available for segmentation")
return None, image
# Create segmentation mask
segmentation_mask = torch.zeros_like(image[0, 0]) # Single channel mask
for i, concept in enumerate(concepts):
if concept in concept_maps:
concept_map = concept_maps[concept]
# Assign different values for different concepts
segmentation_mask += concept_map * (i + 1)
# Normalize segmentation mask
if segmentation_mask.max() > 0:
segmentation_mask = segmentation_mask / segmentation_mask.max()
# Create segmented image
segmented_image = self._create_segmented_image(image, segmentation_mask)
return segmentation_mask.unsqueeze(0).unsqueeze(0), segmented_image
except Exception as e:
logger.error(f"Error in segmentation: {e}")
return None, image
def _create_segmented_image(self, image, segmentation_mask):
"""
Create segmented image visualization.
"""
try:
# Convert to numpy if needed
if isinstance(image, torch.Tensor):
img_np = image.squeeze().cpu().numpy()
if img_np.shape[0] == 3: # CHW format
img_np = np.transpose(img_np, (1, 2, 0))
else:
img_np = np.array(image)
# Normalize image
img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min() + 1e-8)
# Apply segmentation mask
segmented = img_np * segmentation_mask.cpu().numpy()
# Convert back to tensor
if len(segmented.shape) == 3:
segmented = np.transpose(segmented, (2, 0, 1)) # HWC to CHW
return torch.from_numpy(segmented).unsqueeze(0)
except Exception as e:
logger.error(f"Error creating segmented image: {e}")
return image
class ConceptAttentionVisualizerNode:
"""
Node for visualizing concept attention maps.
"""
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("visualized_image",)
FUNCTION = "visualize_attention"
CATEGORY = "Concept Attention"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"concept_maps": ("CONCEPT_MAPS",),
"image": ("IMAGE",),
}
}
def __init__(self):
pass
def visualize_attention(self, concept_maps, image):
"""
Create visualization of concept attention maps.
"""
try:
if not concept_maps:
logger.warning("No concept maps available for visualization")
return image
# Convert image to numpy if needed
if isinstance(image, torch.Tensor):
img_np = image.squeeze().cpu().numpy()
if img_np.shape[0] == 3: # CHW format
img_np = np.transpose(img_np, (1, 2, 0))
else:
img_np = np.array(image)
# Normalize image
img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min() + 1e-8)
# Create overlay for each concept
overlay = img_np.copy()
colors = [
[1.0, 0.0, 0.0], # Red
[0.0, 1.0, 0.0], # Green
[0.0, 0.0, 1.0], # Blue
[1.0, 1.0, 0.0], # Yellow
[1.0, 0.0, 1.0], # Magenta
]
for i, (concept, attention_map) in enumerate(concept_maps.items()):
if i < len(colors):
color = colors[i]
# Resize attention map to image size
attention_resized = np.array(Image.fromarray(attention_map).resize(
(img_np.shape[1], img_np.shape[0]), Image.BILINEAR
))
# Create colored overlay
for c in range(3):
overlay[:, :, c] += attention_resized * color[c] * 0.3
# Clip to valid range
overlay = np.clip(overlay, 0, 1)
# Convert back to tensor
if len(overlay.shape) == 3:
overlay = np.transpose(overlay, (2, 0, 1)) # HWC to CHW
return torch.from_numpy(overlay).unsqueeze(0)
except Exception as e:
logger.error(f"Error creating attention visualization: {e}")
return image
# ComfyUI node registration
NODE_CLASS_MAPPINGS = {
"ConceptAttentionNode": ConceptAttentionNode,
"ConceptSaliencyMapNode": ConceptSaliencyMapNode,
"ConceptSegmentationNode": ConceptSegmentationNode,
"ConceptAttentionVisualizerNode": ConceptAttentionVisualizerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ConceptAttentionNode": "Concept Attention",
"ConceptSaliencyMapNode": "Concept Saliency Map",
"ConceptSegmentationNode": "Concept Segmentation",
"ConceptAttentionVisualizerNode": "Concept Attention Visualizer",
}
Reference in New Issue View Git Blame Copy Permalink