From 56bee4b747ba93f2377de5fc7495fa081ea97608 Mon Sep 17 00:00:00 2001 From: Bryce Date: Wed, 22 Apr 2026 09:16:42 -0700 Subject: [PATCH] Add implementation plan for Compass Image Loader custom node --- IMPLEMENTATION_PLAN.md | 240 +++++++++++++++++++++++++++++++++++++++++ SPEC.md | 227 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 467 insertions(+) create mode 100644 IMPLEMENTATION_PLAN.md create mode 100644 SPEC.md diff --git a/IMPLEMENTATION_PLAN.md b/IMPLEMENTATION_PLAN.md new file mode 100644 index 0000000..750bbbf --- /dev/null +++ b/IMPLEMENTATION_PLAN.md @@ -0,0 +1,240 @@ +# Implementation Plan: Compass Image Loader + +This document provides a step-by-step implementation plan for the `CompassImageLoader` ComfyUI custom node. Each task is a concrete, actionable item with a checkbox. + +--- + +## Phase 1: Project Bootstrap & File Structure + +**Goal:** Set up the custom node package so ComfyUI recognizes and registers the node. + +- [ ] **Create `__init__.py`** + - [ ] Import `CompassImageLoader` from `compass_image_loader`. + - [ ] Define `NODE_CLASS_MAPPINGS = {"CompassImageLoader": CompassImageLoader}`. + - [ ] Define `NODE_DISPLAY_NAME_MAPPINGS = {"CompassImageLoader": "Compass Image Loader"}`. + +- [ ] **Create `compass_image_loader.py`** + - [ ] Define the `CompassImageLoader` class. + - [ ] Set `CATEGORY = "image/loaders"`. + +--- + +## Phase 2: Backend Node Implementation + +**Goal:** Implement the Python class that ComfyUI executes. + +### 2.1 Input Types & Discovery + +- [ ] Implement `CompassImageLoader.INPUT_TYPES(cls)` classmethod. + - [ ] Implement `_discover_directories()` helper that: + - [ ] Gets the base path via `folder_paths.get_input_directory()`. + - [ ] Recursively walks the base path using `os.walk`. + - [ ] Identifies valid direction names: `{"n", "ne", "e", "se", "s", "sw", "w", "nw"}`. + - [ ] Identifies valid modality names: `{"image", "depth", "openpose"}`. + - [ ] For each directory, if it contains **any** valid direction subfolder **or** any valid modality subfolder, add its relative path (from `input/`) to the candidate list. + - [ ] Return sorted, deduplicated list of strings. + - [ ] Return dictionary with: + - `required`: + - `directory`: `(directories_list,)` + - `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})` + +### 2.2 Return Types & Function Signature + +- [ ] Set `RETURN_TYPES = ("IMAGE", "STRING", "INT", "INT", "INT")`. +- [ ] Set `RETURN_NAMES = ("IMAGE", "path", "width", "height", "frame_count")`. +- [ ] Set `FUNCTION = "load_images"`. +- [ ] Set `OUTPUT_NODE = False`. + +### 2.3 Core Loading Logic + +- [ ] Implement `load_images(self, directory, direction, modality, frame, width, height)`. + +#### Path Resolution +- [ ] Resolve `base_dir = folder_paths.get_input_directory()`. +- [ ] Construct `target_dir`: + - [ ] If `direction.strip()`: `os.path.join(base_dir, directory, direction, modality)`. + - [ ] Else: `os.path.join(base_dir, directory, modality)`. +- [ ] Verify `target_dir` exists; if not, raise `RuntimeError`. + +#### Image Discovery +- [ ] List files in `target_dir` with `os.listdir`. +- [ ] Filter to supported image extensions: `png`, `jpg`, `jpeg`, `webp`, `bmp`, `gif`, `tiff`. +- [ ] Sort alphabetically: `sorted(files)`. +- [ ] If list is empty, raise `RuntimeError`. + +#### Frame Selection +- [ ] If `frame` is `None` or `frame.strip() == ""`: + - [ ] Select all images (`selected_files = all_files`). + - [ ] Set `output_path = target_dir`. +- [ ] Else: + - [ ] Try `index = int(frame.strip())`. + - [ ] If `index < 0` or `index >= len(all_files)`, raise `RuntimeError`. + - [ ] Select single file: `selected_files = [all_files[index]]`. + - [ ] Set `output_path = os.path.join(target_dir, all_files[index])`. + +#### Image Processing Loop +- [ ] Initialize empty list `tensors = []`. +- [ ] For each file in `selected_files`: + - [ ] Open with `Image.open(path)`. + - [ ] Convert to RGB: `image = image.convert("RGB")`. + - [ ] Read original dimensions: `orig_w, orig_h = image.size`. + - [ ] Determine final dimensions and resize: + - [ ] If `width == 0 and height == 0`: + - [ ] `final_w, final_h = orig_w, orig_h` (no resize). + - [ ] Elif `width > 0 and height == 0`: + - [ ] `final_w = width`. + - [ ] `final_h = int(orig_h * (width / orig_w))`. + - [ ] `image = image.resize((final_w, final_h), Image.Resampling.LANCZOS)`. + - [ ] Elif `height > 0 and width == 0`: + - [ ] `final_h = height`. + - [ ] `final_w = int(orig_w * (height / orig_h))`. + - [ ] `image = image.resize((final_w, final_h), Image.Resampling.LANCZOS)`. + - [ ] Else (both non-zero): + - [ ] `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)`. + - [ ] Compute crop box: + - [ ] `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 = width, height`. + - [ ] Convert to NumPy array: + - [ ] `np_image = np.array(image).astype(np.float32) / 255.0` + - [ ] Convert to Torch tensor: + - [ ] `tensor = torch.from_numpy(np_image)[None,]` # shape [1, H, W, 3] + - [ ] Append to `tensors`. + +#### Batch Assembly +- [ ] If `len(tensors) == 1`: + - [ ] `image_batch = tensors[0]` +- [ ] Else: + - [ ] `image_batch = torch.cat(tensors, dim=0)` + +#### Return +- [ ] Return `(image_batch, output_path, final_w, final_h, len(selected_files))`. + +### 2.4 Change Detection + +- [ ] Implement `IS_CHANGED(self, directory, direction, modality, frame, width, height)`: + - [ ] Compute `target_dir` as above. + - [ ] If frame is empty, hash the sorted filenames list + mtime. + - [ ] If frame is specified, hash the single file content. + - [ ] Include `width` and `height` in the hash (so resizing triggers re-execution). + +--- + +## Phase 3: Frontend Extension + +**Goal:** Provide an enhanced UX for the `directory` combo widget (typing to filter / autocomplete). + +- [ ] **Create `js/compass_image_loader.js`** + - [ ] Import `app` from `scripts/app.js`. + - [ ] Import `api` from `scripts/api.js` (for event listeners if needed). + +- [ ] **Register Extension** + - [ ] Call `app.registerExtension({ name: "CompassImageLoader", ... })`. + +- [ ] **Override `beforeRegisterNodeDef`** + - [ ] Check `passedNodeData.name === "CompassImageLoader"`. + - [ ] Save reference to original `onNodeCreated`. + - [ ] Replace with a wrapper that: + - [ ] Calls the original `onNodeCreated`. + - [ ] Finds the `directory` widget by name. + - [ ] Stores the original full list of combo values on the widget instance (e.g., `widget._all_values`). + - [ ] Attaches a keyup/input listener to the widget's DOM element (`widget.inputEl` or equivalent) to implement filtering. + - [ ] When the user types, filter `_all_values` case-insensitively and update `widget.options.values`. + - [ ] If the current `widget.value` is not in the filtered list, leave it as-is (allow free-form or preserve selection). + +- [ ] **Handle Refresh** + - [ ] Listen for `fresh-node-defs` event from `api`. + - [ ] On event, iterate over all nodes in `app.graph._nodes`. + - [ ] For each `CompassImageLoader` node: + - [ ] Find the `directory` widget. + - [ ] Update `widget._all_values` and `widget.options.values` with the new list from the refreshed node definition. + - [ ] If the currently selected value still exists in the new list, keep it. + - [ ] If it does not exist, clear it to `""` (or first available). + +- [ ] **State Preservation** + - [ ] Ensure the directory value is correctly serialized in the workflow JSON (ComfyUI handles this automatically for widgets, but verify the filtered list does not override the saved value on `configure`). + +--- + +## Phase 4: Integration & Packaging + +- [ ] Verify file layout: + ``` + ComfyUI/custom_nodes/ComfyUI-compass-paths/ + __init__.py + compass_image_loader.py + js/ + compass_image_loader.js + ``` +- [ ] Ensure `__init__.py` is minimal and correct. +- [ ] Restart ComfyUI and verify the node appears in the node search under `image/loaders`. + +--- + +## Phase 5: Testing & Validation + +Run the following manual tests in the ComfyUI interface. + +- [ ] **Test 1: Single Frame Load** + - [ ] Set `directory`, `direction=s`, `modality=depth`, `frame=0`, `width=0`, `height=0`. + - [ ] Queue. Verify output `IMAGE` shape is `[1, H, W, 3]`, `frame_count=1`, `path` is the full file path. + +- [ ] **Test 2: Batch Load** + - [ ] Leave `frame` blank. + - [ ] Queue. Verify `IMAGE` shape is `[N, H, W, 3]` where `N > 1`, `frame_count=N`, `path` is the directory path. + +- [ ] **Test 3: Cover + Crop Resize** + - [ ] Set `width=512`, `height=768` on a non-square source image. + - [ ] Verify output `width=512`, `height=768`, and image appears center-cropped. + +- [ ] **Test 4: Proportional Resize (Width Only)** + - [ ] Set `width=512`, `height=0`. + - [ ] Verify output width=512 and height is proportionally scaled (no cropping). + +- [ ] **Test 5: Proportional Resize (Height Only)** + - [ ] Set `width=0`, `height=768`. + - [ ] Verify output height=768 and width is proportionally scaled (no cropping). + +- [ ] **Test 6: No Direction** + - [ ] Set `direction` to blank. + - [ ] Verify node loads from `input///` (no direction folder). + +- [ ] **Test 7: Error — Missing Directory** + - [ ] Set `directory` to a non-existent path. + - [ ] Queue. Verify a clear `RuntimeError` is raised. + +- [ ] **Test 8: Error — Empty Image Folder** + - [ ] Point to a valid directory/modality path that has no images. + - [ ] Queue. Verify a clear `RuntimeError` is raised. + +- [ ] **Test 9: Error — Out of Bounds Frame** + - [ ] Set `frame` to a number larger than the number of images minus one. + - [ ] Queue. Verify a clear `RuntimeError` is raised. + +- [ ] **Test 10: Refresh Updates Combo** + - [ ] Add a new valid folder under `input/` while ComfyUI is running. + - [ ] Click the ComfyUI Refresh button. + - [ ] Open the node's `directory` dropdown. Verify the new folder appears. + +- [ ] **Test 11: Autocomplete Filtering** + - [ ] Click the `directory` combo and type a substring (e.g., `walk`). + - [ ] Verify the dropdown filters to show only matching entries. + +--- + +## Future Enhancements (Out of Scope for Initial Implementation) + +- Support for recursive image loading across subdirectories. +- Caching of directory scans to improve performance on very large `input/` trees. +- Support for additional modalities beyond `image`, `depth`, `openpose` via a configuration file. +- Drag-and-drop folder path into the directory widget. diff --git a/SPEC.md b/SPEC.md new file mode 100644 index 0000000..17beb8a --- /dev/null +++ b/SPEC.md @@ -0,0 +1,227 @@ +# Specification: Compass Image Loader + +## Overview + +A ComfyUI custom node that loads images from a structured directory hierarchy under ComfyUI's `input/` directory. The node navigates via a user-selected subdirectory, an optional compass direction, and a modality folder, then outputs matching image(s) with optional resizing. + +## Directory Structure Convention + +``` +/input/ + / + / + [direction]/ + / + frame_0001.png + frame_0002.png + ... +``` + +When `direction` is left blank, the direction folder is omitted: + +``` +/input/ + / + / + / + frame_0001.png + ... +``` + +**Example resolved path (with direction):** `input/rosella/walk/s/depth/frame_0001.png` +**Example resolved path (without direction):** `input/rosella/walk/image/frame_0001.png` + +--- + +## Node Definition + +### Node Name +`CompassImageLoader` + +### Display Name +`Compass Image Loader` + +### Category +`image/loaders` + +--- + +## Widgets / Inputs + +| Name | Type | Widget | Default | Description | +|------|------|--------|---------|-------------| +| `directory` | `STRING` | Combo with autocomplete | `""` | Recursive subdirectory path under `input/`. Populated dynamically from all directories containing at least one valid direction subfolder (n, ne, e, se, s, sw, w, nw) **or** at least one valid modality subfolder (image, depth, openpose). Supports typing to filter. | +| `direction` | `STRING` | Combo | `""` | One of: `""`, `n`, `ne`, `e`, `se`, `s`, `sw`, `w`, `nw`. Empty string omits the direction folder entirely. | +| `modality` | `STRING` | Combo | `"image"` | One of: `image`, `depth`, `openpose` | +| `frame` | `STRING` | Text | `""` | Frame index as string. Empty string or whitespace-only = load all images as a batch. `"0"` = first image (0-based). | +| `width` | `INT` | Number | `0` | Target resize width in pixels. `0` = leave original width. | +| `height` | `INT` | Number | `0` | Target resize height in pixels. `0` = leave original height. | + +### Resize Behavior +- If `width == 0` and `height == 0`: no resize. +- If **only one** of `width` or `height` is non-zero: scale proportionally so that the non-zero dimension matches the target. The other dimension is computed to maintain aspect ratio. No cropping occurs. +- If **both** `width` and `height` are non-zero: use a **cover + center crop** strategy: + 1. Compute `scale = max(target_width / source_width, target_height / source_height)`. + 2. Resize by `scale` using Lanczos resampling. + 3. Center crop to exactly `width x height`. + +--- + +## Outputs + +| Name | Type | Description | +|------|------|-------------| +| `IMAGE` | `IMAGE` | Tensor of shape `(B, H, W, C)` where `B` = batch size. `B = 1` for a single frame; `B = N` when `frame` is empty (all images loaded as a batch). | +| `path` | `STRING` | The resolved file system path. For single frame: absolute path to the image file. For batch mode: absolute path to the directory containing the images. | +| `width` | `INT` | Final width in pixels after any resizing. | +| `height` | `INT` | Final height in pixels after any resizing. | +| `frame_count` | `INT` | Number of frames loaded. `1` for single frame, `N` for batch mode. | + +--- + +## Core Behavior + +### 1. Base Directory Resolution +- Base is always `folder_paths.get_input_directory()`. +- The user does **not** configure this. + +### 2. Directory Discovery (Autocomplete) +- The backend recursively scans the `input/` directory. +- A directory is eligible for the `directory` combo if it satisfies **either** of the following: + - It contains at least one valid direction subfolder (`n`, `ne`, `e`, `se`, `s`, `sw`, `w`, `nw`), **or** + - It contains at least one valid modality subfolder (`image`, `depth`, `openpose`). +- The list of eligible paths is provided as the combo values for the `directory` widget. +- The frontend extension enhances this combo with typing-to-filter autocomplete behavior. + +### 3. Path Construction +```python +if direction and direction.strip(): + full_dir = os.path.join(input_dir, directory, direction, modality) +else: + full_dir = os.path.join(input_dir, directory, modality) +``` + +### 4. Image Discovery +- List all files in `full_dir`. +- Filter to supported image content types. Use `folder_paths.filter_files_content_types` if available, otherwise filter by standard image extensions (`png`, `jpg`, `jpeg`, `webp`, `bmp`, `gif`, `tiff`). +- Sort alphabetically by filename. This sorted order defines frame indices (`0 = first file alphabetically`). + +### 5. Frame Selection +- If `frame` is empty or whitespace-only: select all images (batch mode). +- Else: parse `frame` as a base-10 integer. Select the single image at that 0-based index. +- If the parsed index is negative or `>= len(images)`, raise `RuntimeError` with a clear message. + +### 6. Image Loading & Resizing +- Open each selected file with PIL. +- Convert to RGB (`image.convert("RGB")`). +- If resize is required: + - **Proportional resize (one dimension zero):** + - If `width > 0` and `height == 0`: `new_height = int(source_height * (width / source_width))`. + - If `height > 0` and `width == 0`: `new_width = int(source_width * (height / source_height))`. + - Use `Image.Resampling.LANCZOS`. + - **Cover + Crop (both dimensions non-zero):** + - `scale = max(width / source_width, height / source_height)`. + - `new_size = (int(source_width * scale), int(source_height * scale))`. + - `image = image.resize(new_size, Image.Resampling.LANCZOS)`. + - Compute crop box to center the crop: + - `left = (new_size[0] - width) // 2` + - `top = (new_size[1] - height) // 2` + - `right = left + width` + - `bottom = top + height` + - `image = image.crop((left, top, right, bottom))`. + +### 7. Tensor Conversion +- `np.array(image).astype(np.float32) / 255.0` +- `torch.from_numpy(image)[None,]` (adds batch dimension of 1) +- For batch mode, `torch.cat(tensors, dim=0)` to merge into `(N, H, W, C)`. + +### 8. Return Values +- `path`: + - Single frame: `os.path.join(full_dir, filename)` (absolute path). + - Batch mode: `full_dir` (absolute path). +- `width`, `height`: Dimensions of the final tensor. +- `frame_count`: `len(images)`. + +--- + +## Error Handling + +| Scenario | Behavior | +|----------|----------| +| `directory` does not exist | `RuntimeError("Compass directory not found: {full_dir}")` | +| `full_dir` contains no supported images | `RuntimeError("No images found in: {full_dir}")` | +| `frame` index out of bounds | `RuntimeError("Frame index {n} out of bounds. Found {m} images in {full_dir}.")` | +| `frame` string is not a valid integer | `RuntimeError("Invalid frame number: '{frame}'. Must be an integer.")` | + +--- + +## Frontend Extension Requirements + +### File +`js/compass_image_loader.js` + +### Responsibilities +1. **Register Extension:** Use `app.registerExtension({ name: "CompassImageLoader", ... })`. +2. **Dynamic Combo Refresh:** Listen for the `fresh-node-defs` event fired by ComfyUI's `api.js` (which happens on startup and when the user clicks Refresh). When received, update the `directory` combo widget on existing node instances with the latest values from the node's definition. +3. **Autocomplete / Filtering:** Override the `directory` combo widget behavior to allow the user to type and filter the dropdown list in real-time (rgthree-style). The filtered list should be case-insensitive. +4. **State Preservation:** Ensure the user's selected `directory` value is preserved across page reloads and workflow loads. Do not reset the value when refreshing the combo list unless the previously selected value is no longer in the new list. + +--- + +## Backend API for Frontend Refresh + +To support the frontend refreshing the directory list without reloading the page, the backend exposes a lightweight HTTP endpoint. However, the standard ComfyUI pattern is to rely on the node definition refresh (clicking the Refresh button in the UI), which re-fetches `api.getNodeDefs()`. Therefore, the node must ensure its `INPUT_TYPES` classmethod dynamically rebuilds the directory list every time it is called. + +No custom API endpoint is required. The frontend simply refreshes the combo values from the re-fetched node definition. + +--- + +## Dependencies + +- `folder_paths` (ComfyUI builtin) +- `torch` +- `numpy` +- `PIL` (Pillow) +- `os` + +No external pip dependencies are required beyond what ComfyUI already provides. + +--- + +## Example Workflows + +### Example 1: Single Frame, with Direction, Resized +- `directory`: `rosella/walk` +- `direction`: `s` +- `modality`: `depth` +- `frame`: `0` +- `width`: `512` +- `height`: `768` + +**Result:** Loads `input/rosella/walk/s/depth/`, resizes to 512x768 with cover-crop + Lanczos. Returns `IMAGE` [1, 768, 512, 3], path string, width=512, height=768, frame_count=1. + +### Example 2: Batch Mode, No Direction, No Resize +- `directory`: `hero/run` +- `direction`: *(blank)* +- `modality`: `image` +- `frame`: *(blank)* +- `width`: `0` +- `height`: `0` + +**Result:** Loads all images from `input/hero/run/image/`, sorted alphabetically, as a batch tensor [N, H, W, 3]. Returns directory path, original dimensions, frame_count=N. + +--- + +## Testing Criteria + +| # | Test | Expected Result | +|---|------|-----------------| +| 1 | `frame="0"`, no resize | Loads first image alphabetically, outputs [1, H, W, 3] | +| 2 | `frame=""`, no resize | Loads all images, outputs [N, H, W, 3] | +| 3 | `width=512, height=768` | Image is resized to cover 512x768 and center-cropped | +| 4 | `width=512, height=0` | Image is proportionally scaled so width=512, height maintained | +| 5 | Direction blank | Path omits direction folder; loads from `input/dir/modality/` | +| 6 | Invalid directory | Clear `RuntimeError` | +| 7 | `frame` out of bounds | Clear `RuntimeError` | +| 8 | Refresh button clicked | Directory combo updates with newly added/removed folders | +| 9 | Type in directory combo | Dropdown filters to matching entries |