""" ComfyUI Saliency Nodes - Algorithmic Attention For Mac (Apple Silicon) + comfyenv conda setup. Install: pip install saliency torchvision matplotlib Place in: ComfyUI/custom_nodes/creative_code/saliency_nodes.py """ import torch import torch.nn.functional as F import numpy as np # Lazy imports _saliency = None _torchvision_models = None def _ensure_deps(): global _saliency, _torchvision_models if _saliency is None: try: import saliency.core as sal _saliency = sal except ImportError: raise ImportError("[SaliencyNodes] Run: pip install saliency") if _torchvision_models is None: try: import torchvision.models as models _torchvision_models = models except ImportError: raise ImportError("[SaliencyNodes] Run: pip install torchvision") IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" SUPPORTED_MODELS = [ "vgg16", "vgg19", "resnet50", "resnet101", "inception_v3", "mobilenet_v3_large", "densenet121", ] _model_cache: dict = {} _imagenet_labels: list = [] def _get_imagenet_label(class_idx: int) -> str: global _imagenet_labels if not _imagenet_labels: try: import torchvision.models as tv _imagenet_labels = tv.VGG16_Weights.DEFAULT.meta["categories"] except Exception: pass if _imagenet_labels and 0 <= class_idx < len(_imagenet_labels): return _imagenet_labels[class_idx] return f"class_{class_idx}" def _load_model(model_name: str): if model_name not in _model_cache: _ensure_deps() with torch.inference_mode(False): with torch.enable_grad(): constructor = getattr(_torchvision_models, model_name) try: import torchvision.models as tv weight_map = { "inception_v3": tv.Inception_V3_Weights.DEFAULT, "resnet50": tv.ResNet50_Weights.DEFAULT, "resnet101": tv.ResNet101_Weights.DEFAULT, "vgg16": tv.VGG16_Weights.DEFAULT, "vgg19": tv.VGG19_Weights.DEFAULT, "mobilenet_v3_large": tv.MobileNet_V3_Large_Weights.DEFAULT, "densenet121": tv.DenseNet121_Weights.DEFAULT, } model = constructor(weights=weight_map.get(model_name)) except AttributeError: model = constructor(pretrained=True) device = "cuda" if torch.cuda.is_available() else "cpu" model.eval().to(device) for param in model.parameters(): param.data = param.data.clone() for buf_name, buf in model.named_buffers(): if buf is not None: parts = buf_name.rsplit('.', 1) if len(parts) == 2: parent = dict(model.named_modules())[parts[0]] setattr(parent, parts[1], buf.clone()) else: setattr(model, buf_name, buf.clone()) _model_cache[model_name] = model print(f"[SaliencyNodes] Loaded {model_name} on {DEVICE} (grad-safe)") return _model_cache[model_name] def _comfy_to_np(image_tensor: torch.Tensor) -> np.ndarray: return image_tensor[0].cpu().numpy().astype(np.float32) def _np_to_comfy(arr: np.ndarray) -> torch.Tensor: if arr.ndim == 2: arr = np.stack([arr, arr, arr], axis=-1) return torch.from_numpy(np.clip(arr, 0.0, 1.0).astype(np.float32)).unsqueeze(0) def _resize_np(img: np.ndarray, h: int, w: int) -> np.ndarray: t = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0) t = F.interpolate(t, size=(h, w), mode="bilinear", align_corners=False) return t.squeeze(0).permute(1, 2, 0).numpy() def _make_call_model_fn(model, class_idx): def call_model_function(images, call_model_args=None, expected_keys=None): import saliency.core as saliency with torch.inference_mode(False): with torch.enable_grad(): images_np = np.array(images, dtype=np.float32) images_tensor = torch.from_numpy(images_np).permute(0, 3, 1, 2) mean = torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1).to(DEVICE) std = torch.tensor(IMAGENET_STD).view(1, 3, 1, 1).to(DEVICE) images_tensor = images_tensor.to(DEVICE) images_tensor = (images_tensor - mean) / std images_tensor = images_tensor.clone().detach().requires_grad_(True) output = model(images_tensor) if hasattr(output, "logits"): output = output.logits target = output[:, class_idx].sum() model.zero_grad() target.backward() gradients = images_tensor.grad.detach().cpu().numpy() gradients = np.transpose(gradients, (0, 2, 3, 1)) result = {} if saliency.INPUT_OUTPUT_GRADIENTS in (expected_keys or []): result[saliency.INPUT_OUTPUT_GRADIENTS] = gradients if saliency.CONVOLUTION_LAYER_VALUES in (expected_keys or []): result[saliency.CONVOLUTION_LAYER_VALUES] = gradients return result return call_model_function def _predict_top_class(model, img_np_hwc: np.ndarray) -> int: with torch.inference_mode(False): with torch.no_grad(): t = torch.from_numpy(img_np_hwc).permute(2, 0, 1).unsqueeze(0) mean = torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1).to(DEVICE) std = torch.tensor(IMAGENET_STD).view(1, 3, 1, 1).to(DEVICE) t = t.to(DEVICE) t = (t - mean) / std out = model(t) if hasattr(out, "logits"): out = out.logits return int(out.argmax(dim=1).item()) def _vis_grayscale(mask): if mask.ndim == 2: m = np.abs(mask) vmax = np.percentile(m, 99) return np.clip(m / (vmax + 1e-8), 0, 1) import saliency.core as saliency return saliency.VisualizeImageGrayscale(mask) def _vis_diverging(mask3d): import saliency.core as saliency div = saliency.VisualizeImageDiverging(mask3d) return (div - div.min()) / (div.max() - div.min() + 1e-8) def _apply_colormap(gray_2d, colormap_name): try: import matplotlib.cm as cm cmap = getattr(cm, colormap_name, cm.viridis) return cmap(gray_2d)[..., :3].astype(np.float32) except Exception: return np.stack([gray_2d, gray_2d * 0.5, np.zeros_like(gray_2d)], axis=-1) def _apply_heatmap(gray2d, original_hwc, colormap="viridis"): heatmap = _apply_colormap(gray2d, colormap) return np.clip(original_hwc * 0.45 + heatmap * 0.55, 0, 1) def _run_saliency_method(method, call_fn, img_np, use_smoothgrad, n_samples, noise, ig_steps): _ensure_deps() import saliency.core as saliency get_kw = dict(call_model_function=call_fn, call_model_args=None) smooth_kw = dict(call_model_function=call_fn, call_model_args=None, stdev_spread=noise, nsamples=n_samples) if method == "Vanilla Gradients": algo = saliency.GradientSaliency() return algo.GetSmoothedMask(img_np, **smooth_kw) if use_smoothgrad \ else algo.GetMask(img_np, **get_kw) elif method == "SmoothGrad": return saliency.GradientSaliency().GetSmoothedMask(img_np, **smooth_kw) elif method == "Integrated Gradients": algo = saliency.IntegratedGradients() baseline = np.zeros_like(img_np) if use_smoothgrad: return algo.GetSmoothedMask(img_np, **smooth_kw, x_baseline=baseline, x_steps=ig_steps) return algo.GetMask(img_np, **get_kw, x_baseline=baseline, x_steps=ig_steps) elif method == "Blur IG": algo = saliency.BlurIG() return algo.GetSmoothedMask(img_np, **smooth_kw, steps=ig_steps) \ if use_smoothgrad else algo.GetMask(img_np, **get_kw, steps=ig_steps) elif method == "Guided IG": algo = saliency.GuidedIG() baseline = np.zeros_like(img_np) if use_smoothgrad: return algo.GetSmoothedMask(img_np, **smooth_kw, x_baseline=baseline, x_steps=ig_steps) return algo.GetMask(img_np, **get_kw, x_baseline=baseline, x_steps=ig_steps) elif method == "XRAI": algo = saliency.XRAI() params = saliency.XRAIParameters() params.algorithm = "fast" return algo.GetMask(img_np, **get_kw, extra_parameters=params) raise ValueError(f"Unknown method: {method}") # Node 1: Gradient Saliency Visualizer class SaliencyMapNode: METHODS = ["Vanilla Gradients", "SmoothGrad", "Integrated Gradients", "Blur IG", "Guided IG", "XRAI"] OUTPUT_MODES = ["heatmap_overlay", "grayscale_mask", "diverging_mask", "overlay_and_mask"] @classmethod def INPUT_TYPES(cls): return {"required": { "image": ("IMAGE",), "model_name": (SUPPORTED_MODELS, {"default": "vgg16"}), "method": (cls.METHODS, {"default": "Integrated Gradients"}), "output_mode": (cls.OUTPUT_MODES, {"default": "heatmap_overlay"}), "class_index": ("INT", {"default": -1, "min": -1, "max": 999}), "smoothgrad_samples": ("INT", {"default": 25, "min": 5, "max": 100, "step": 5}), "smoothgrad_noise": ("FLOAT", {"default": 0.15, "min": 0.01, "max": 0.5, "step": 0.01}), "ig_steps": ("INT", {"default": 50, "min": 10, "max": 300, "step": 10}), "input_size": ("INT", {"default": 224, "min": 224, "max": 512, "tooltip": "Use 299 for InceptionV3"}), }, "optional": { "use_smoothgrad": ("BOOLEAN", {"default": False}), }} RETURN_TYPES = ("IMAGE", "IMAGE") RETURN_NAMES = ("saliency_output", "original_resized") FUNCTION = "compute_saliency" CATEGORY = "creative-code/explainability" def compute_saliency(self, image, model_name, method, output_mode, class_index, smoothgrad_samples, smoothgrad_noise, ig_steps, input_size, use_smoothgrad=False): model = _load_model(model_name) img_np = _comfy_to_np(image) img_np_r = _resize_np(img_np, input_size, input_size) if class_index < 0: class_index = _predict_top_class(model, img_np_r) print(f"[Gradient Saliency Visualizer] Auto class: {class_index}") call_fn = _make_call_model_fn(model, class_index) mask3d = _run_saliency_method(method, call_fn, img_np_r, use_smoothgrad, smoothgrad_samples, smoothgrad_noise, ig_steps) return (self._render(output_mode, mask3d, img_np_r), _np_to_comfy(img_np_r)) def _render(self, mode, mask3d, original): if mode == "heatmap_overlay": return _np_to_comfy(_apply_heatmap(_vis_grayscale(mask3d), original)) elif mode == "grayscale_mask": return _np_to_comfy(_vis_grayscale(mask3d)) elif mode == "diverging_mask": return _np_to_comfy(_vis_diverging(mask3d)) elif mode == "overlay_and_mask": gray = _vis_grayscale(mask3d) overlay = _apply_heatmap(gray, original) gray_rgb = np.stack([gray, gray, gray], axis=-1) return _np_to_comfy(np.concatenate([original, overlay, gray_rgb], axis=1)) raise ValueError(f"Unknown mode: {mode}") # Node 2: Saliency Method Comparison Grid class SaliencyComparisonNode: @classmethod def INPUT_TYPES(cls): return {"required": { "image": ("IMAGE",), "model_name": (SUPPORTED_MODELS, {"default": "vgg16"}), "class_index": ("INT", {"default": -1, "min": -1, "max": 999}), "ig_steps": ("INT", {"default": 50, "min": 10, "max": 200}), "input_size": ("INT", {"default": 224, "min": 224, "max": 512}), }} RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("comparison_grid",) FUNCTION = "run_comparison" CATEGORY = "creative-code/explainability" def run_comparison(self, image, model_name, class_index, ig_steps, input_size): model = _load_model(model_name) img_np = _comfy_to_np(image) img_np_r = _resize_np(img_np, input_size, input_size) if class_index < 0: class_index = _predict_top_class(model, img_np_r) call_fn = _make_call_model_fn(model, class_index) node = SaliencyMapNode() results = [] for method in SaliencyMapNode.METHODS: mask3d = _run_saliency_method(method, call_fn, img_np_r, False, 25, 0.15, ig_steps) out = node._render("heatmap_overlay", mask3d, img_np_r) results.append(out[0].numpy()) row1 = np.concatenate(results[0:3], axis=1) row2 = np.concatenate(results[3:6], axis=1) return (_np_to_comfy(np.concatenate([row1, row2], axis=0)),) # Node 3: Algorithmic Attention Art class ExplainableAIArtNode: METHODS = SaliencyMapNode.METHODS COLORMAPS = [ "inferno", "plasma", "magma", "viridis", "hot", "cool", "spring", "copper", "twilight", "ocean", ] COMPOSITIONS = [ "isolation", "overlay", "spotlight", "ghost", "invert", "cutout", "multiply", "screen", "triptych", "mask_only", "pure_white", ] @classmethod def INPUT_TYPES(cls): return {"required": { "image": ("IMAGE",), "model_name": (SUPPORTED_MODELS, {"default": "vgg16"}), "method": (cls.METHODS, {"default": "Integrated Gradients"}), "composition": (cls.COMPOSITIONS, {"default": "isolation"}), "colormap": (cls.COLORMAPS, {"default": "inferno"}), "intensity": ("FLOAT", {"default": 0.75, "min": 0.0, "max": 1.0, "step": 0.05}), "contrast": ("FLOAT", {"default": 1.5, "min": 0.5, "max": 4.0, "step": 0.1}), "threshold": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 0.95, "step": 0.05}), "class_index": ("INT", {"default": -1, "min": -1, "max": 999}), "input_size": ("INT", {"default": 224, "min": 224, "max": 512}), }, "optional": { "use_smoothgrad": ("BOOLEAN", {"default": True}), "smoothgrad_samples": ("INT", {"default": 25, "min": 5, "max": 50}), }} RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("artwork", "classification_label") FUNCTION = "create_art" CATEGORY = "creative-code/art" def create_art(self, image, model_name, method, composition, colormap, intensity, contrast, threshold, class_index, input_size, use_smoothgrad=True, smoothgrad_samples=25): model = _load_model(model_name) img_np = _comfy_to_np(image) img_np_r = _resize_np(img_np, input_size, input_size) if class_index < 0: class_index = _predict_top_class(model, img_np_r) label = _get_imagenet_label(class_index) print(f"[Algorithmic Attention Art] Visualizing class {class_index}: {label}") call_fn = _make_call_model_fn(model, class_index) mask3d = _run_saliency_method(method, call_fn, img_np_r, use_smoothgrad, smoothgrad_samples, 0.15, 50) artwork = self._compose(mask3d, img_np_r, composition, colormap, intensity, contrast, threshold) caption = f"{label} (class {class_index}), {model_name} / {method}" return (_np_to_comfy(artwork), caption) def _compose(self, mask3d, original, composition, colormap, intensity, contrast, threshold): gray = _vis_grayscale(mask3d) gray = np.clip(np.power(gray, 1.0 / max(contrast, 0.01)), 0, 1) heatmap = _apply_colormap(gray, colormap) mask = gray[..., None] if composition == "overlay": return np.clip(original * (1 - intensity) + heatmap * intensity, 0, 1) elif composition == "mask_only": return heatmap elif composition == "isolation": gray_img = np.mean(original, axis=-1, keepdims=True).repeat(3, axis=-1) return np.clip( gray_img * (1 - mask * intensity) + heatmap * mask * intensity, 0, 1 ) elif composition == "spotlight": darkened = original * (0.15 + 0.85 * mask) return np.clip(darkened * (1 - intensity) + heatmap * intensity, 0, 1) elif composition == "ghost": return np.clip(original * 0.25 + heatmap * intensity, 0, 1) elif composition == "invert": inv_heatmap = _apply_colormap(1 - gray, colormap) return np.clip(original * (1 - intensity) + inv_heatmap * intensity, 0, 1) elif composition == "cutout": binary = (gray > threshold).astype(np.float32)[..., None] return np.clip(original * binary + heatmap * (1 - binary) * 0.08, 0, 1) elif composition == "multiply": return np.clip(original * (1 - mask * intensity) + heatmap * intensity * 0.3, 0, 1) elif composition == "screen": return np.clip(1 - (1 - original) * (1 - heatmap * intensity), 0, 1) elif composition == "triptych": overlay = np.clip(original * (1 - intensity) + heatmap * intensity, 0, 1) return np.concatenate([original, heatmap, overlay], axis=1) elif composition == "pure_white": white = gray[..., None] * np.ones((1, 1, 3), dtype=np.float32) return np.clip(white, 0, 1) return heatmap # Registration NODE_CLASS_MAPPINGS = { "SaliencyMap": SaliencyMapNode, "SaliencyComparison": SaliencyComparisonNode, "SaliencyArt": ExplainableAIArtNode, } NODE_DISPLAY_NAME_MAPPINGS = { "SaliencyMap": "Gradient Saliency Visualizer", "SaliencyComparison": "Saliency Method Comparison Grid", "SaliencyArt": "Algorithmic Attention Art", }