From 5ed55d99d08379472a9db637d31641787e9e9065 Mon Sep 17 00:00:00 2001 From: Claire Vlases Date: Tue, 24 Feb 2026 09:31:01 -0500 Subject: [PATCH] readme update --- saliency_nodes.py | 624 +++++++++++++--------------------------------- 1 file changed, 170 insertions(+), 454 deletions(-) diff --git a/saliency_nodes.py b/saliency_nodes.py index 31deb08..6532c65 100644 --- a/saliency_nodes.py +++ b/saliency_nodes.py @@ -1,494 +1,210 @@ -""" -ComfyUI Saliency Nodes — Algorithmic Attention -For Mac (Apple Silicon) + comfyenv conda setup. +# Running the Saliency Art Node -Install: pip install saliency torchvision matplotlib -Place in: ComfyUI/custom_nodes/creative_code/saliency_nodes.py -""" +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. -import torch -import torch.nn.functional as F -import numpy as np +--- -# ── Lazy imports ────────────────────────────────────────────────────────────── -_saliency = None -_torchvision_models = None +## What is ComfyUI? -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") +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. -# ── 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) +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. -SUPPORTED_MODELS = [ - "vgg16", "vgg19", "resnet50", "resnet101", - "inception_v3", "mobilenet_v3_large", - # "efficientnet_b0", - "densenet121", -] +--- -_model_cache: dict = {} +## Running ComfyUI Locally -# ── ImageNet class labels ───────────────────────────────────────────────────── -_imagenet_labels: list = [] +### Requirements -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) +- 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)) - 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()) +### Installation +```bash +# Clone the ComfyUI repository +git clone https://github.com/comfyanonymous/ComfyUI.git +cd ComfyUI - _model_cache[model_name] = model - print(f"[SaliencyNodes] Loaded {model_name} on CPU (grad-safe)") - return _model_cache[model_name] +# Create and activate a conda environment +conda create -n comfyenv python=3.13 +conda activate comfyenv +# Install dependencies +pip install -r requirements.txt -# ── Image helpers ───────────────────────────────────────────────────────────── -def _comfy_to_np(image_tensor: torch.Tensor) -> np.ndarray: - return image_tensor[0].cpu().numpy().astype(np.float32) +# Start ComfyUI +python main.py +``` -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) +Then open `http://127.0.0.1:8188` in your browser. You'll see the node canvas. -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 -# ── 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 +### 1. Install Python dependencies - # 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) +With your ComfyUI conda environment active: +```bash +pip install saliency matplotlib torchvision +``` - 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 +### 2. Place the node files - # .clone() is required when escaping inference_mode - images_tensor = images_tensor.clone().detach().requires_grad_(True) +Copy the `creative_code/` folder into ComfyUI's custom nodes directory: +``` +ComfyUI/ +└── custom_nodes/ + └── creative_code/ + ├── __init__.py + └── saliency_nodes.py +``` - output = model(images_tensor) - if hasattr(output, "logits"): - output = output.logits +`__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 +``` - target = output[:, class_idx].sum() - model.zero_grad() - target.backward() +### 3. Restart ComfyUI +```bash +python main.py +``` - gradients = images_tensor.grad.detach().numpy() - gradients = np.transpose(gradients, (0, 2, 3, 1)) # BCHW→BHWC +You should see this in the startup log, confirming the node loaded: +``` +0.0 seconds: /path/to/ComfyUI/custom_nodes/creative_code +``` - 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 +The nodes will appear in the node menu under **creative-code/art** and **creative-code/explainability**. - return result +--- - return call_model_function +## Basic Workflow +The minimal working pipeline is three nodes: +``` +[Load Image] → [SaliencyArt] → [Preview Image] +``` -# ── Top class prediction ─────────────────────────────────────────────────────── -def _predict_top_class(model, img_np_hwc: np.ndarray) -> int: +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** + +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. + +--- + +## How the Node Was Built + +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. + +### Node structure + +A minimal ComfyUI node looks like this: +```python +class MyNode: + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "image": ("IMAGE",), + "my_param": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}), + }} + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("output",) + FUNCTION = "run" + CATEGORY = "my-category" + + def run(self, image, my_param): + # do something with image + return (processed_image,) + +NODE_CLASS_MAPPINGS = {"MyNode": MyNode} +NODE_DISPLAY_NAME_MAPPINGS = {"MyNode": "My Node (Display Name)"} +``` + +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. + +### The scaffolding problem + +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 +``` + +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 + +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. + +Three layers of fixes were required: + +**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() +``` + +**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.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()) + 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() +``` +**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. -# ── 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) +### 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 +``` -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) +## Technical Notes -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) +**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. +**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. -# ── 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 +**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. - 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) +**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. - 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) +## Dependencies - 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",), - "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"[SaliencyMap] 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 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: Saliency Art — the main art node -# ══════════════════════════════════════════════════════════════════════════════ -class SaliencyArtNode: - - METHODS = SaliencyMapNode.METHODS - - COLORMAPS = [ - "inferno", "plasma", "magma", "viridis", "hot", - "cool", "spring", "copper", "twilight", "ocean", - ] - - 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 - ] - - @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"[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 +| 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