From cf5cbef001c7bc2f4daf85fdef81df3bbaed2186 Mon Sep 17 00:00:00 2001 From: Claire Vlases Date: Tue, 24 Feb 2026 09:34:03 -0500 Subject: [PATCH] minor tweak --- saliency_nodes.py | 580 ++++++++++++++++++++++++++++++++++------------ 1 file changed, 432 insertions(+), 148 deletions(-) diff --git a/saliency_nodes.py b/saliency_nodes.py index 6532c65..80df1f0 100644 --- a/saliency_nodes.py +++ b/saliency_nodes.py @@ -1,210 +1,494 @@ -# Running the Saliency Art Node +""" +ComfyUI Saliency Nodes — Algorithmic Attention +For Mac (Apple Silicon) + comfyenv conda setup. -This tool is a custom node for [ComfyUI](https://github.com/comfyanonymous/ComfyUI), an open-source node-based interface for running generative AI pipelines locally. The node interrogates pretrained image classifiers using gradient-based saliency methods, visualizing what the model attends to when it classifies an image. +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 -## What is ComfyUI? +# ── Lazy imports ────────────────────────────────────────────────────────────── +_saliency = None +_torchvision_models = None -ComfyUI is a local, node-based interface for building AI image pipelines. Unlike web-based tools, it runs entirely on your machine — no API keys, no cloud, no data leaving your computer. You build workflows by connecting nodes: each node does one thing (load an image, run a model, save output), and you wire them together visually. +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") -This project adds a custom node set to ComfyUI that exposes saliency visualization as a first-class pipeline step — meaning you can chain it with other nodes, batch process images, and experiment with parameters in real time. +# ── Constants ───────────────────────────────────────────────────────────────── +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) ---- +SUPPORTED_MODELS = [ + "vgg16", "vgg19", "resnet50", "resnet101", + "inception_v3", "mobilenet_v3_large", + "efficientnet_b0", + "densenet121", +] -## Running ComfyUI Locally +_model_cache: dict = {} -### Requirements +# ── ImageNet class labels ───────────────────────────────────────────────────── +_imagenet_labels: list = [] -- Python 3.10+ (this project used **Python 3.13.5** via Anaconda) -- A conda or virtual environment is strongly recommended -- Mac Apple Silicon (M-series), Windows, or Linux -- GPU optional but recommended — this node runs on CPU by design (see [Technical Notes](#technical-notes)) +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() + # Must load outside inference_mode — ComfyUI wraps everything in it, + # and tensors created inside inference_mode can't be used with autograd. + 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, + # "efficientnet_b0": tv.EfficientNet_B0_Weights.DEFAULT, + "densenet121": tv.DenseNet121_Weights.DEFAULT, + } + model = constructor(weights=weight_map.get(model_name)) + except AttributeError: + model = constructor(pretrained=True) -### Installation -```bash -# Clone the ComfyUI repository -git clone https://github.com/comfyanonymous/ComfyUI.git -cd ComfyUI + model.eval().cpu() + # Deep clone all params and buffers to ensure no inference tensors remain + for param in model.parameters(): + param.data = param.data.clone() + for buf_name, buf in model.named_buffers(): + if buf is not None: + # set_buffer isn't directly available, use the parent module + 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()) -# Create and activate a conda environment -conda create -n comfyenv python=3.13 -conda activate comfyenv + _model_cache[model_name] = model + print(f"[SaliencyNodes] Loaded {model_name} on CPU (grad-safe)") + return _model_cache[model_name] -# Install dependencies -pip install -r requirements.txt -# Start ComfyUI -python main.py -``` +# ── Image helpers ───────────────────────────────────────────────────────────── +def _comfy_to_np(image_tensor: torch.Tensor) -> np.ndarray: + return image_tensor[0].cpu().numpy().astype(np.float32) -Then open `http://127.0.0.1:8188` in your browser. You'll see the node canvas. +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() -## Installing the Saliency Node -### 1. Install Python dependencies +# ── call_model_function for PAIR saliency ───────────────────────────────────── +def _make_call_model_fn(model, class_idx): + """ + ComfyUI runs node execution inside torch.inference_mode() which disables + autograd entirely. We must escape that context to compute gradients. + Everything runs on CPU to avoid MPS autograd issues. + """ + def call_model_function(images, call_model_args=None, expected_keys=None): + import saliency.core as saliency -With your ComfyUI conda environment active: -```bash -pip install saliency matplotlib torchvision -``` + # Double-escape any no_grad / inference_mode context from ComfyUI + 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) -### 2. Place the node files + mean = torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1) + std = torch.tensor(IMAGENET_STD).view(1, 3, 1, 1) + images_tensor = (images_tensor - mean) / std -Copy the `creative_code/` folder into ComfyUI's custom nodes directory: -``` -ComfyUI/ -└── custom_nodes/ - └── creative_code/ - ├── __init__.py - └── saliency_nodes.py -``` + # .clone() is required when escaping inference_mode + images_tensor = images_tensor.clone().detach().requires_grad_(True) -`__init__.py` tells ComfyUI this folder is a node package. It should contain: -```python -from .saliency_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -``` + output = model(images_tensor) + if hasattr(output, "logits"): + output = output.logits -### 3. Restart ComfyUI -```bash -python main.py -``` + target = output[:, class_idx].sum() + model.zero_grad() + target.backward() -You should see this in the startup log, confirming the node loaded: -``` -0.0 seconds: /path/to/ComfyUI/custom_nodes/creative_code -``` + gradients = images_tensor.grad.detach().numpy() + gradients = np.transpose(gradients, (0, 2, 3, 1)) # BCHW→BHWC -The nodes will appear in the node menu under **creative-code/art** and **creative-code/explainability**. + 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 -## Basic Workflow + return call_model_function -The minimal working pipeline is three nodes: -``` -[Load Image] → [SaliencyArt] → [Preview Image] -``` -1. Right-click the canvas → Add Node → image → **Load Image** -2. Right-click → Add Node → creative-code/art → **Saliency Art — Algorithmic Attention** -3. Right-click → Add Node → image → **Preview Image** -4. Connect **Load Image** `IMAGE` output → **SaliencyArt** `image` input -5. Connect **SaliencyArt** `artwork` output → **Preview Image** `images` input -6. Upload an image in the Load Image node -7. Click **Queue Prompt** +# ── Top class prediction ─────────────────────────────────────────────────────── +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) + std = torch.tensor(IMAGENET_STD).view(1, 3, 1, 1) + t = (t - mean) / std + out = model(t) + if hasattr(out, "logits"): + out = out.logits + return int(out.argmax(dim=1).item()) -The first run downloads the model weights (~500MB for VGG16) automatically. Subsequent runs use the cached weights. -The node also outputs a `classification_label` string — wire this to a **Show Text** node to see what the model classified the image as and which class is being visualized. +# ── Visualization helpers ────────────────────────────────────────────────────── +def _vis_grayscale(mask): + """Handle both 3D (H,W,3) gradient masks and 2D (H,W) XRAI region masks.""" + if mask.ndim == 2: + # XRAI returns 2D directly — just normalize it + 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) -## How the Node Was Built +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) -ComfyUI custom nodes are Python classes that follow a specific interface. Each node declares its inputs, outputs, and a function that gets called when the pipeline runs. +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) -### Node structure -A minimal ComfyUI node looks like this: -```python -class MyNode: +# ── Shared saliency method runner ───────────────────────────────────────────── +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: Saliency Map +# ══════════════════════════════════════════════════════════════════════════════ +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",), - "my_param": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}), + "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",) - RETURN_NAMES = ("output",) - FUNCTION = "run" - CATEGORY = "my-category" + RETURN_TYPES = ("IMAGE", "IMAGE") + RETURN_NAMES = ("saliency_output", "original_resized") + FUNCTION = "compute_saliency" + CATEGORY = "creative-code/explainability" - def run(self, image, my_param): - # do something with image - return (processed_image,) + 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) -NODE_CLASS_MAPPINGS = {"MyNode": MyNode} -NODE_DISPLAY_NAME_MAPPINGS = {"MyNode": "My Node (Display Name)"} -``` + if class_index < 0: + class_index = _predict_top_class(model, img_np_r) + print(f"[SaliencyMap] Auto class: {class_index}") -ComfyUI reads `NODE_CLASS_MAPPINGS` at startup and registers each class as a node. The `INPUT_TYPES` dict defines what appears in the node's UI — dropdowns for string lists, sliders for floats, toggles for booleans. `RETURN_TYPES` defines the output ports. + 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) -### The scaffolding problem + return (self._render(output_mode, mask3d, img_np_r), _np_to_comfy(img_np_r)) -The recommended way to create a new node is via a cookiecutter template: -```bash -pip install cookiecutter -cookiecutter https://github.com/bronkula/comfyui-node-template -``` + 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}") -However, the template has a known issue: it generates `__init.py__` (wrong) instead of `__init__.py` (correct), which causes ComfyUI to silently skip the node at startup with no error. The fix is to rename the file manually before doing anything else. -### The key technical challenge: gradients inside ComfyUI +# ══════════════════════════════════════════════════════════════════════════════ +# Node 2: Saliency 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}), + }} -The core difficulty in building this node was that **ComfyUI wraps all node execution inside `torch.inference_mode()`** — a PyTorch context that disables the autograd engine entirely for performance. This is fine for diffusion models that only need forward passes, but saliency methods require computing gradients (backward passes), which inference mode makes impossible. + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("comparison_grid",) + FUNCTION = "run_comparison" + CATEGORY = "creative-code/explainability" -Three layers of fixes were required: + 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) -**1. Load models outside inference mode.** Model weights created inside `inference_mode()` are permanently marked as "inference tensors" and can't participate in gradient computation even if you later try to enable gradients. The solution is to wrap the entire model loading step in `torch.inference_mode(False)`: -```python -with torch.inference_mode(False): - with torch.enable_grad(): - model = torchvision.models.vgg16(weights=...) - model.eval().cpu() -``` + if class_index < 0: + class_index = _predict_top_class(model, img_np_r) -**2. Escape inference mode at call time.** Even with models loaded correctly, the `call_model_function` that the saliency library calls repeatedly during computation runs inside ComfyUI's inference mode context. Each call needs to explicitly break out: -```python -def call_model_function(images, ...): - with torch.inference_mode(False): - with torch.enable_grad(): - images_tensor = torch.from_numpy(images).clone().requires_grad_(True) - output = model(images_tensor) - output[:, class_idx].sum().backward() - gradients = images_tensor.grad.detach().numpy() -``` + call_fn = _make_call_model_fn(model, class_index) + node = SaliencyMapNode() + results = [] -**3. Force CPU for gradient computation.** On Apple Silicon (MPS backend), PyTorch's MPS implementation drops `requires_grad` silently during certain operations, making backward passes fail even when autograd is enabled. All gradient computation runs on CPU; the classification forward pass (which doesn't need gradients) can use MPS for speed. + 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()) -### File structure -``` -creative_code/ -├── __init__.py # registers nodes with ComfyUI -└── saliency_nodes.py # all node logic - ├── _load_model() # loads + caches torchvision models - ├── _make_call_model_fn() # builds the saliency library callback - ├── _predict_top_class() # auto-detects ImageNet class - ├── _run_saliency_method() # dispatches to PAIR saliency methods - ├── _compose() # applies composition modes - ├── SaliencyMapNode # single method output - ├── SaliencyComparisonNode # all-methods grid - └── SaliencyArtNode # main art node -``` + 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)),) ---- -## Technical Notes +# ══════════════════════════════════════════════════════════════════════════════ +# Node 3: Saliency Art — the main art node +# ══════════════════════════════════════════════════════════════════════════════ +class SaliencyArtNode: -**Why CPU?** Gradient-based saliency on Apple Silicon (MPS) silently drops `requires_grad` during backward passes, producing incorrect results with no error message. All saliency computation runs on CPU to guarantee correctness. The classification forward pass (auto-detecting the top class) still uses MPS for speed since it doesn't require gradients. + METHODS = SaliencyMapNode.METHODS -**Why these specific library versions?** The PAIR saliency library's `BlurIG` method uses `steps` as its parameter name while all other IG-family methods use `x_steps` — an inconsistency in their API. The node handles this internally so you don't have to. + COLORMAPS = [ + "inferno", "plasma", "magma", "viridis", "hot", + "cool", "spring", "copper", "twilight", "ocean", + ] -**XRAI is slower.** XRAI uses a fundamentally different algorithm (image segmentation + region attribution) rather than pixel-level gradients. It typically takes several minutes on CPU versus seconds for gradient methods. + COMPOSITIONS = [ + "isolation", # color only where AI looks, grayscale elsewhere + "overlay", # classic heatmap blend + "spotlight", # radial fade from salient to dark + "ghost", # faded original + bright highlights + "invert", # show what the algorithm IGNORES + "cutout", # hard binary: only pixels above threshold + "multiply", # attention dims non-salient areas + "screen", # attention brightens salient areas + "triptych", # original | mask | composite side by side + "mask_only", # just the colored saliency map + "pure_white", # black background, white gradient pixels only + ] -**Model weights are cached.** The first time you select a model, PyTorch downloads the pretrained weights (~100–550MB depending on the model) to `~/.cache/torch/hub/checkpoints/`. Subsequent runs load from cache instantly. + @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" -## Dependencies + def create_art(self, image, model_name, method, composition, colormap, + intensity, contrast, threshold, class_index, input_size, + use_smoothgrad=True, smoothgrad_samples=25): -| Package | Purpose | -|---------|---------| -| `saliency` | Google PAIR saliency methods (GradCAM, IG, XRAI, etc.) | -| `torchvision` | Pretrained ImageNet models | -| `matplotlib` | Colormaps | -| `torch` | Included with ComfyUI | -| `numpy` | Included with ComfyUI | \ No newline at end of file + 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"[SaliencyArt] 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": + # Black background, white where the gradient fires + 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": SaliencyArtNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SaliencyMap": "Saliency Map (PAIR)", + "SaliencyComparison": "Saliency Comparison Grid (PAIR)", + "SaliencyArt": "Saliency Art — Algorithmic Attention", +} \ No newline at end of file