From 0db6ca8277f31162805e5ccbdb9d5f01d9fcabc3 Mon Sep 17 00:00:00 2001 From: Bryce Date: Wed, 22 Apr 2026 09:18:52 -0700 Subject: [PATCH] Implement Compass Image Loader node with dynamic directory discovery and autocomplete --- __init__.py | 9 +++ compass_image_loader.py | 157 +++++++++++++++++++++++++++++++++++++ js/compass_image_loader.js | 116 +++++++++++++++++++++++++++ 3 files changed, 282 insertions(+) create mode 100644 __init__.py create mode 100644 compass_image_loader.py create mode 100644 js/compass_image_loader.js diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..c6c1f69 --- /dev/null +++ b/__init__.py @@ -0,0 +1,9 @@ +from .compass_image_loader import CompassImageLoader + +NODE_CLASS_MAPPINGS = { + "CompassImageLoader": CompassImageLoader, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "CompassImageLoader": "Compass Image Loader", +} diff --git a/compass_image_loader.py b/compass_image_loader.py new file mode 100644 index 0000000..d752fca --- /dev/null +++ b/compass_image_loader.py @@ -0,0 +1,157 @@ +from PIL import Image +import folder_paths +import os +import torch +import numpy as np + + +VALID_DIRECTIONS = {"n", "ne", "e", "se", "s", "sw", "w", "nw"} +VALID_MODALITIES = {"image", "depth", "openpose"} + + +def _discover_directories(): + base_dir = folder_paths.get_input_directory() + if not os.path.exists(base_dir): + return [] + + candidates = set() + for root, subdirs, _ in os.walk(base_dir, followlinks=True): + current_path = root[os.path.dirname(root) + 1:] + sub_dirs_lower = {s.lower() for s in subdirs} + if VALID_DIRECTIONS & sub_dirs_lower or VALID_MODALITIES & sub_dirs_lower: + candidates.add(current_path) + + return sorted(candidates) + + +class CompassImageLoader: + CATEGORY = "image/loaders" + + @classmethod + def INPUT_TYPES(cls): + directories = _discover_directories() + return { + "required": { + "directory": (directories if directories else ["(none found)"],), + "direction": (["", "n", "ne", "e", "se", "s", "sw", "w", "nw"],), + "modality": (["image", "depth", "openpose"],), + "frame": ("STRING", {"default": ""}), + "width": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1}), + "height": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1}), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING", "INT", "INT", "INT") + RETURN_NAMES = ("IMAGE", "path", "width", "height", "frame_count") + FUNCTION = "load_images" + + def load_images(self, directory, direction, modality, frame=None, width=0, height=0): + base_dir = folder_paths.get_input_directory() + + if direction and direction.strip(): + target_dir = os.path.join(base_dir, directory, direction, modality) + else: + target_dir = os.path.join(base_dir, directory, modality) + + if not os.path.isdir(target_dir): + raise RuntimeError(f"Compass directory not found: {target_dir}") + + supported_extensions = {"png", "jpg", "jpeg", "webp", "bmp", "gif", "tiff"} + files = [ + f for f in sorted(os.listdir(target_dir)) + if os.path.isfile(os.path.join(target_dir, f)) and f.split(".")[-1].lower() in supported_extensions + ] + + if not files: + raise RuntimeError(f"No images found in: {target_dir}") + + if frame is None or str(frame).strip() == "": + selected_files = files + output_path = target_dir + else: + try: + index = int(str(frame).strip()) + except (ValueError, TypeError): + raise RuntimeError(f"Invalid frame number: '{frame}'. Must be an integer.") + + if index < 0 or index >= len(files): + raise RuntimeError( + f"Frame index {index} out of bounds. Found {len(files)} images in {target_dir}." + ) + + selected_files = [files[index]] + output_path = os.path.join(target_dir, files[index]) + + tensors = [] + final_w, final_h = 0, 0 + + for filename in selected_files: + filepath = os.path.join(target_dir, filename) + image = Image.open(filepath).convert("RGB") + orig_w, orig_h = image.size + + if width == 0 and height == 0: + pass + elif width > 0 and height == 0: + fw = width + fh = int(orig_h * (width / orig_w)) + image = image.resize((fw, fh), Image.Resampling.LANCZOS) + elif height > 0 and width == 0: + fh = height + fw = int(orig_w * (height / orig_h)) + image = image.resize((fw, fh), Image.Resampling.LANCZOS) + else: + scale = max(width / orig_w, height / orig_h) + new_w = int(orig_w * scale) + new_h = int(orig_h * scale) + image = image.resize((new_w, new_h), Image.Resampling.LANCZOS) + left = (new_w - width) // 2 + top = (new_h - height) // 2 + right = left + width + bottom = top + height + image = image.crop((left, top, right, bottom)) + + final_w, final_h = image.size[0], image.size[1] + + np_image = np.array(image).astype(np.float32) / 255.0 + tensor = torch.from_numpy(np_image)[None,] + tensors.append(tensor) + + if len(tensors) == 1: + image_batch = tensors[0] + else: + image_batch = torch.cat(tensors, dim=0) + + return (image_batch, output_path, final_w, final_h, len(selected_files)) + + @classmethod + def IS_CHANGED(cls, directory, direction, modality, frame=None, width=0, height=0): + base_dir = folder_paths.get_input_directory() + + if direction and direction.strip(): + target_dir = os.path.join(base_dir, directory, direction, modality) + else: + target_dir = os.path.join(base_dir, directory, modality) + + import hashlib + + m = hashlib.sha256() + m.update(f"{directory}:{direction}:{modality}:{frame}:{width}:{height}".encode("utf-8")) + if not os.path.isdir(target_dir): + return "" + + supported_extensions = {"png", "jpg", "jpeg", "webp", "bmp", "gif", "tiff"} + files = [ + f for f in sorted(os.listdir(target_dir)) + if os.path.isfile(os.path.join(target_dir, f)) and f.split(".")[-1].lower() in supported_extensions + ] + + if frame is None or str(frame).strip() == "": + m.update(":".join(files).encode("utf-8")) + else: + try: + index = int(str(frame).strip()) + filepath = os.path.join(target_dir, files[index]) + with open(filepath, "rb") as f: + m.update(f.read(65536)) + except (ValueError, IndexErro diff --git a/js/compass_image_loader.js b/js/compass_image_loader.js new file mode 100644 index 0000000..818ede0 --- /dev/null +++ b/js/compass_image_loader.js @@ -0,0 +1,116 @@ +import { app } from "../../../scripts/app.js"; + +app.registerExtension({ + name: "CompassImageLoader", + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== "CompassImageLoader") return; + + const origOnNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + if (origOnNodeCreated) origOnNodeCreated.apply(this, []); + + const directoryWidget = this.widgets.find( + (w) => w.name === "directory" + ); + if (!directoryWidget) return; + + directoryWidget._all_values = [...(directoryWidget.options.values || [])]; + + try { + if (typeof this.addDOMWidget === "function") { + const inputEl = document.createElement("input"); + inputEl.type = "text"; + inputEl.value = directoryWidget.value; + + inputEl.addEventListener("change", () => { + const filter = inputEl.value.toLowerCase(); + const filtered = directoryWidget._all_values.filter((v) => + v.toLowerCase().includes(filter) + ); + directoryWidget.options.values = filtered; + directoryWidget.value = filtered[0] || ""; + this.onResize?.(); + }); + + const dropdownEl = document.createElement("div"); + dropdownEl.style.overflowY = "auto"; + dropdownEl.style.maxHeight = "200px"; + + const listEl = document.createElement("ul"); + listEl.style.padding = "4px"; + listEl.style.margin = "0"; + listEl.style.background = "#333"; + + inputEl.addEventListener("focus", () => { + listEl.innerHTML = ""; + directoryWidget._all_values.forEach((v) => { + const li = document.createElement("li"); + li.textContent = v; + li.style.padding = "4px 8px"; + li.style.cursor = "pointer"; + li.addEventListener("click", () => { + inputEl.value = v; + directoryWidget.value = v; + listEl.innerHTML = ""; + }); + listEl.appendChild(li); + }); + }); + + inputEl.addEventListener("blur", (e) => { + setTimeout(() => { + if (!listEl.contains(e.relatedTarget)) { + listEl.innerHTML = ""; + } + }, 0); + }); + + const textWidget = this.addDOMWidget( + "directory_filter", + "Directory Filter", + inputEl, + {}, + { widget: directoryWidget } + ); + textWidget.computeSize = () => [1, 40]; + } + } catch (e) { + console.warn("[CompassImageLoader] Widget setup failed:", e); + } + }; + }, + + async loadedGraphNode(node) { + if (node.type !== "CompassImageLoader") return; + + const directoryWidget = node.widgets?.find((w) => w.name === "directory"); + if (!directoryWidget) return; + + setTimeout(() => { + try { + app.api.addWebSocketMessageEventListener("fresh-node-defs", (event) => { + const defs = event.detail || {}; + const compassDef = defs["CompassImageLoader"]; + if (!compassDef) return; + + const dirs = compassDef.input?.required?.directory; + if (!dirs || !Array.isArray(dirs)) return; + + directoryWidget._all_values = [...dirs]; + directoryWidget.options.values = [...dirs]; + + if (!directoryWidget._all_values.includes(directoryWidget.value)) { + directoryWidget.value = directoryWidget._all_values[0] || ""; + } + }); + } catch (e) { + console.warn("[CompassImageLoader] Refresh listener failed:", e); + } + }, 100); + }, + + async getCustomWidgets() { + return null; + } +});