nodes/cvlases-creative-code-comfyui
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

minor tweak

Browse Source
This commit is contained in:
Claire Vlases
2026-02-24 09:34:03 -05:00
parent a8ae0d9745
commit cf5cbef001
1 changed files with 432 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
saliency_nodes.py
+432 -148
View File
@@ -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 |
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",
}