ComfyUI-compass-paths/IMPLEMENTATION_PLAN.md

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/<directory>/<modality>/ (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.