diff --git a/__init__.py b/__init__.py index 1850a0d..8598d75 100644 --- a/__init__.py +++ b/__init__.py @@ -1,72 +1,3 @@ -import os -import server -from aiohttp import web -import folder_paths -import nodes +from .saliency_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -class MyNode: - @classmethod - def INPUT_TYPES(s): - return {"required": {}} - - RETURN_TYPES = () - RETURN_NAMES = () - FUNCTION = "test" - CATEGORY = "creative-code" - - def test(self): - return () - -NODE_CLASS_MAPPINGS = { - "MyNode": MyNode -} -NODE_DISPLAY_NAME_MAPPINGS = { - "MyNode": "My First Node" -} - -__all__ = ["NODE_CLASS_MAPPINGS"] - -# Define the path to our extension -workspace_path = os.path.dirname(__file__) -dist_path = os.path.join(workspace_path, "dist/example_ext") -dist_locales_path = os.path.join(workspace_path, "dist/locales") - -# Print the current paths for debugging -print(f"ComfyUI_example_frontend_extension workspace path: {workspace_path}") -print(f"Dist path: {dist_path}") -print(f"Dist locales path: {dist_locales_path}") -print(f"Locales exist: {os.path.exists(dist_locales_path)}") - -# Register the static route for serving our React app assets -if os.path.exists(dist_path): - # Add the routes for the extension - server.PromptServer.instance.app.add_routes([ - web.static("/example_ext/", dist_path), - ]) - - # Register the locale files route - if os.path.exists(dist_locales_path): - server.PromptServer.instance.app.add_routes([ - web.static("/locales/", dist_locales_path), - ]) - print(f"Registered locale files route at /locales/") - else: - print("WARNING: Locale directory not found!") - - # Also register the standard ComfyUI extension web directory - - project_name = os.path.basename(workspace_path) - - try: - # Method added in https://github.com/comfyanonymous/ComfyUI/pull/8357 - from comfy_config import config_parser - - project_config = config_parser.extract_node_configuration(workspace_path) - project_name = project_config.project.name - print(f"project name read from pyproject.toml: {project_name}") - except Exception as e: - print(f"Could not load project config, using default name '{project_name}': {e}") - - nodes.EXTENSION_WEB_DIRS[project_name] = os.path.join(workspace_path, "dist") -else: - print("ComfyUI Example React Extension: Web directory not found") \ No newline at end of file +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/saliency_nodes.py b/saliency_nodes.py new file mode 100644 index 0000000..e63b1bd --- /dev/null +++ b/saliency_nodes.py @@ -0,0 +1,464 @@ +""" +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") + +# ── 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", +] + +_model_cache: dict = {} + +# ── Model loading — always CPU for gradient compatibility ───────────────────── +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) + + 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()) + + _model_cache[model_name] = model + print(f"[SaliencyNodes] Loaded {model_name} on CPU (grad-safe)") + return _model_cache[model_name] + + +# ── Image helpers ───────────────────────────────────────────────────────────── +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() + + +# ── 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 + + # 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) + + 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 + + # .clone() is required when escaping inference_mode + 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().numpy() + gradients = np.transpose(gradients, (0, 2, 3, 1)) # BCHW→BHWC + + 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 + + +# ── 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()) + + +# ── Visualization helpers ────────────────────────────────────────────────────── +def _vis_grayscale(mask3d): + import saliency.core as saliency + return saliency.VisualizeImageGrayscale(mask3d) + +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) + + +# ── 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, x_steps=ig_steps) \ + if use_smoothgrad else algo.GetMask(img_np, **get_kw, x_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 + ] + + @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",) + RETURN_NAMES = ("artwork",) + 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) + print(f"[SaliencyArt] Visualizing 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, + 0.15, 50) + + artwork = self._compose(mask3d, img_np_r, composition, colormap, + intensity, contrast, threshold) + return (_np_to_comfy(artwork),) + + 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) + + 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