feat: Add proper --base-directory support for custom data locations

CLAUDE.md

@@ -62,9 +62,15 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 - **Modular Launcher**: Manages installation, configuration, and runtime with separated concerns
 - **Logging System**: Provides consistent, configurable logging across all components
 
+### Command Line Options
+- `--base-directory PATH`: Set data directory for models, input, output, custom_nodes (preferred method)
+- `--open`: Auto-open browser when server is ready
+- `--port=XXXX`: Run on specific port (default: 8188)
+- `--debug` or `--verbose`: Enable detailed debug logging
+
 ### Important Environment Variables
-- `COMFY_USER_DIR`: Persistent storage directory (default: `~/.config/comfy-ui`)
-- `COMFY_APP_DIR`: ComfyUI application directory
+- `COMFY_USER_DIR`: Persistent storage directory (default: `~/.config/comfy-ui`, use `--base-directory` instead)
+- `COMFY_APP_DIR`: ComfyUI application directory (fixed at `~/.config/comfy-ui/app`)
 - `COMFY_SAVE_PATH`: User save path for outputs
 - `CUDA_VERSION`: CUDA version for PyTorch (default: `cu124`, options: `cu118`, `cu121`, `cu124`, `cpu`)
 - `LD_LIBRARY_PATH`: (Linux) Set automatically to include system libraries and NVIDIA drivers
@@ -75,15 +81,20 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 - **macOS**: Detects Apple Silicon and configures MPS acceleration
 - **Library Paths**: Automatically includes `/run/opengl-driver/lib` on Linux for NVIDIA drivers
 
-### Data Persistence Structure (`~/.config/comfy-ui/`)
+### Data Persistence Structure
+**Fixed locations** (always in `~/.config/comfy-ui/`):
 ```
 app/           - ComfyUI application code (auto-updated when flake changes)
+venv/          - Python virtual environment
+```
+
+**Configurable locations** (default `~/.config/comfy-ui/`, or `--base-directory`):
+```
 models/        - Model files (checkpoints, loras, vae, controlnet, embeddings, upscale_models, clip, diffusers, etc.)
 output/        - Generated images and outputs
 user/          - User configuration and custom nodes
 input/         - Input files for processing
 custom_nodes/  - Persistent custom node installations
-venv/          - Python virtual environment
 ```
 
 ## CI/CD and Automation

README.md

@@ -41,16 +41,19 @@ nix run github:utensils/comfyui-nix/[commit-hash] -- --open
 
 - `--open`: Automatically opens ComfyUI in your browser when the server is ready
 - `--port=XXXX`: Run ComfyUI on a specific port (default: 8188)
+- `--base-directory PATH`: Set data directory for models, input, output, and custom_nodes (default: `~/.config/comfy-ui`). Quote paths with spaces: `--base-directory "/path/with spaces"`
 - `--debug` or `--verbose`: Enable detailed debug logging
 
 ### Environment Variables
 
 - `CUDA_VERSION`: CUDA version for PyTorch (default: `cu124`, options: `cu118`, `cu121`, `cu124`, `cpu`)
-- `COMFY_USER_DIR`: Override the default user data directory (default: `~/.config/comfy-ui`)
 
 ```bash
 # Example: Use CUDA 12.1
 CUDA_VERSION=cu121 nix run github:utensils/comfyui-nix
+
+# Example: Use custom data directory (e.g., on a separate drive)
+nix run github:utensils/comfyui-nix -- --base-directory ~/AI
 ```
 
 ### Development Shell
@@ -383,39 +386,32 @@ The workflow uses Nix to ensure reproducible builds and leverages the same build
 
 ## Useful Hints
 
-### Using External Model Directories
-
-Use `--extra-model-paths-config` to point ComfyUI at existing model directories:
-
-```yaml
-# ~/AI/extra_model_paths.yaml
-ai_models:
-    base_path: ~/AI/Models/
-    checkpoints: checkpoints
-    loras: loras
-    vae: vae
-    controlnet: controlnet
-    clip: clip
-    text_encoders: text_encoders
-    diffusion_models: |
-        diffusion_models
-        unet
-    upscale_models: upscale_models
-
-ai_io:
-    base_path: ~/AI/
-    input: Input
-    output: Output
+### Using a Custom Data Directory
+
+Use `--base-directory` to store all data (models, input, output, custom_nodes) on a separate drive:
+
+```bash
+nix run github:utensils/comfyui-nix -- --base-directory ~/AI
+```
+
+Expected structure in your base directory:
+```
+~/AI/
+├── models/          # checkpoints, loras, vae, text_encoders, etc.
+├── input/           # input files
+├── output/          # generated outputs
+├── custom_nodes/    # extensions
+└── user/            # workflows and settings
 ```
 
-Use pipe (`|`) syntax to map multiple directories to one model type.
+For advanced setups with non-standard directory structures, use `--extra-model-paths-config` with a YAML file to map custom paths.
 
 ### Flux 2 Dev on RTX 4090
 
 Run Flux 2 Dev without offloading using GGUF quantization:
 
 ```bash
-nix run github:utensils/comfyui-nix -- --listen 0.0.0.0 --use-pytorch-cross-attention --cuda-malloc --lowvram --extra-model-paths-config ~/AI/extra_model_paths.yaml
+nix run github:utensils/comfyui-nix -- --base-directory ~/AI --listen 0.0.0.0 --use-pytorch-cross-attention --cuda-malloc --lowvram
 ```
 
 **Models** (install ComfyUI-GGUF via Manager first):

scripts/config.sh

@@ -33,9 +33,135 @@ COMFY_PORT="8188"
 CUDA_VERSION="${CUDA_VERSION:-cu124}"
 
 # Directory structure
+# Check for --base-directory in args first, otherwise use default
+# This is parsed early so all path variables can use BASE_DIR
+# NOTE: We use echo for errors here because logger.sh hasn't been sourced yet.
+# This is intentional - BASE_DIR must be set before other scripts are sourced.
+
+# Function to parse --base-directory from arguments
+# Returns the parsed path via stdout, or empty if not found
+_parse_base_directory() {
+  local args=("$@")
+  local skip_next=false
+  local raw_path=""
+  local i
+
+  for i in "${!args[@]}"; do
+    if [[ "$skip_next" == "true" ]]; then
+      skip_next=false
+      continue
+    fi
+    case "${args[$i]}" in
+      --base-directory=*)
+        raw_path="${args[$i]#*=}"
+        ;;
+      --base-directory)
+        # Validate next argument exists and isn't another flag
+        if [[ $((i+1)) -lt ${#args[@]} ]] && [[ ! "${args[$((i+1))]}" =~ ^-- ]]; then
+          raw_path="${args[$((i+1))]}"
+          skip_next=true
+        else
+          echo "ERROR: --base-directory requires a path argument" >&2
+          exit 1
+        fi
+        ;;
+    esac
+  done
+
+  # Process the path if one was provided
+  if [[ -n "$raw_path" ]]; then
+    # Strip surrounding quotes if present (handles edge case of quoted paths)
+    raw_path="${raw_path%\"}"
+    raw_path="${raw_path#\"}"
+    raw_path="${raw_path%\'}"
+    raw_path="${raw_path#\'}"
+
+    # Expand tilde (handles ~, ~/path, but not ~user)
+    raw_path="${raw_path/#\~/$HOME}"
+
+    # Convert to absolute path
+    if command -v realpath &>/dev/null; then
+      # realpath -m allows non-existent paths
+      realpath -m "$raw_path" 2>/dev/null || echo "$raw_path"
+    else
+      # Fallback: make path absolute if it isn't already
+      if [[ "$raw_path" == /* ]]; then
+        echo "$raw_path"
+      else
+        echo "$PWD/$raw_path"
+      fi
+    fi
+  fi
+}
+
+# Parse and set BASE_DIR
 BASE_DIR="$HOME/.config/comfy-ui"
-CODE_DIR="$BASE_DIR/app"
-COMFY_VENV="$BASE_DIR/venv"
+_parsed_base_dir="$(_parse_base_directory "$@")"
+if [[ -n "$_parsed_base_dir" ]]; then
+  BASE_DIR="$_parsed_base_dir"
+elif [[ "$*" =~ --base-directory ]]; then
+  # User provided the flag but parsing returned empty - this is an error
+  echo "ERROR: Failed to parse --base-directory argument" >&2
+  exit 1
+fi
+unset _parsed_base_dir
+
+# Security: Validate path is not in dangerous system directories
+# Uses blocklist approach to allow custom mount points while blocking system paths
+_validate_base_dir() {
+  local dir="$1"
+  local blocked_prefixes=("/etc" "/bin" "/sbin" "/usr" "/lib" "/lib32" "/lib64" "/boot" "/sys" "/proc" "/dev" "/root")
+
+  # Require absolute path
+  if [[ "$dir" != /* ]]; then
+    echo "ERROR: --base-directory must be an absolute path: $dir" >&2
+    exit 1
+  fi
+
+  # Check against blocked system directories
+  for prefix in "${blocked_prefixes[@]}"; do
+    if [[ "$dir" == "$prefix" || "$dir" == "$prefix"/* ]]; then
+      echo "ERROR: --base-directory cannot be in system directory: $prefix" >&2
+      exit 1
+    fi
+  done
+}
+_validate_base_dir "$BASE_DIR"
+
+# Resolve symlinks to prevent symlink attacks
+# Check both BASE_DIR itself and its parent directory
+
+# First, resolve BASE_DIR if it's a symlink (e.g., /home/user/link -> /etc)
+if [[ -L "$BASE_DIR" ]]; then
+  _resolved_base="$(readlink -f "$BASE_DIR" 2>/dev/null || echo "$BASE_DIR")"
+  _validate_base_dir "$_resolved_base"
+  BASE_DIR="$_resolved_base"
+fi
+
+# Validate base directory parent exists and is writable
+_parent_dir="$(dirname "$BASE_DIR")"
+
+# Also resolve symlinks in parent directory
+if [[ -L "$_parent_dir" ]]; then
+  _resolved_parent="$(readlink -f "$_parent_dir" 2>/dev/null || echo "$_parent_dir")"
+  # Re-validate the resolved path
+  _validate_base_dir "$_resolved_parent/$(basename "$BASE_DIR")"
+  _parent_dir="$_resolved_parent"
+fi
+
+if [[ ! -d "$_parent_dir" ]]; then
+  echo "ERROR: Parent directory of BASE_DIR does not exist: $_parent_dir" >&2
+  echo "Please create it first or specify a valid path with --base-directory" >&2
+  exit 1
+elif [[ ! -w "$_parent_dir" ]]; then
+  echo "ERROR: No write permission for parent directory: $_parent_dir" >&2
+  exit 1
+fi
+unset _parent_dir _resolved_parent _resolved_base
+
+# App code and venv always live in .config (separate from data)
+CODE_DIR="$HOME/.config/comfy-ui/app"
+COMFY_VENV="$HOME/.config/comfy-ui/venv"
 COMFY_MANAGER_DIR="$BASE_DIR/custom_nodes/ComfyUI-Manager"
 MODEL_DOWNLOADER_PERSISTENT_DIR="$BASE_DIR/custom_nodes/model_downloader"
 CUSTOM_NODE_DIR="$CODE_DIR/custom_nodes"
@@ -78,24 +204,39 @@ declare -A DIRECTORIES=(
 )
 
 # Python packages to install (as arrays for proper handling)
+# shellcheck disable=SC2034  # Used in install.sh
 BASE_PACKAGES=(pyyaml pillow numpy requests)
 # Core packages needed for ComfyUI v0.4.0+
+# shellcheck disable=SC2034  # Used in install.sh
 ADDITIONAL_PACKAGES=(spandrel av GitPython toml rich safetensors pydantic pydantic-settings alembic)
 
 # PyTorch installation will be determined dynamically based on GPU availability
 # This is set in install.sh based on platform detection
 
 # Function to parse command line arguments
+# Filters out arguments handled by this launcher, passes rest to ComfyUI
 parse_arguments() {
   ARGS=()
+  local skip_next=false
   for arg in "$@"; do
+    if [[ "$skip_next" == "true" ]]; then
+      skip_next=false
+      continue
+    fi
     case "$arg" in
       --open)
         OPEN_BROWSER=true
         ;;
       --port=*)
         COMFY_PORT="${arg#*=}"
         ;;
+      --base-directory=*)
+        # Already handled in config, don't pass to ComfyUI
+        ;;
+      --base-directory)
+        # Skip this and next arg (value), already handled in config
+        skip_next=true
+        ;;
       --debug)
         export LOG_LEVEL=$DEBUG
         ;;
@@ -111,8 +252,11 @@ parse_arguments() {
 
 # Export the configuration
 export_config() {
+  # Set COMFY_APP_DIR to CODE_DIR for Python persistence module
+  COMFY_APP_DIR="$CODE_DIR"
+
   # Export all defined variables to make them available to sourced scripts
-  export COMFY_VERSION COMFY_PORT BASE_DIR CODE_DIR COMFY_VENV
+  export COMFY_VERSION COMFY_PORT BASE_DIR CODE_DIR COMFY_VENV COMFY_APP_DIR
   export COMFY_MANAGER_DIR MODEL_DOWNLOADER_PERSISTENT_DIR
   export CUSTOM_NODE_DIR MODEL_DOWNLOADER_APP_DIR
   export OPEN_BROWSER PYTHON_ENV

scripts/launcher.sh

@@ -19,7 +19,8 @@ if [[ "$SCRIPT_DIR" == *"/bin" ]]; then
 fi
 
 # Source the component scripts
-source "$SCRIPT_DIR/config.sh"
+# IMPORTANT: Pass "$@" to config.sh so it can parse --base-directory early
+source "$SCRIPT_DIR/config.sh" "$@"
 source "$SCRIPT_DIR/logger.sh"
 source "$SCRIPT_DIR/install.sh"
 source "$SCRIPT_DIR/persistence.sh"

src/persistence/persistence.py

@@ -11,7 +11,6 @@
 import logging
 import os
 import shutil
-import sys
 from pathlib import Path
 
 # Configure logging
@@ -207,10 +206,29 @@ def setup_persistence() -> str:
     if _state.initialized and _state.base_dir is not None:
         return _state.base_dir
 
-    # Create the persistent directory if it doesn't exist
+    # Get base directory from COMFY_USER_DIR environment variable.
+    # This is set by the bash launcher (config.sh) which parses --base-directory.
+    # We intentionally do NOT re-parse command line args here to avoid inconsistencies
+    # between bash and Python argument parsing (edge cases with quotes, spaces, etc.)
     base_dir = os.environ.get(
         "COMFY_USER_DIR", os.path.join(os.path.expanduser("~"), ".config", "comfy-ui")
     )
+
+    # Validate path - defense in depth in case bash validation was bypassed
+    if not os.path.isabs(base_dir):
+        logger.error("base_dir must be an absolute path: %s", base_dir)
+        raise ValueError(f"Invalid base_dir: {base_dir}")
+
+    # Block dangerous system directories (must match bash blocklist in config.sh)
+    blocked_dirs = [
+        "/etc", "/bin", "/sbin", "/usr", "/lib", "/lib32",
+        "/lib64", "/boot", "/sys", "/proc", "/dev", "/root",
+    ]
+    for blocked in blocked_dirs:
+        if base_dir == blocked or base_dir.startswith(blocked + "/"):
+            logger.error("base_dir cannot be in system directory: %s", base_dir)
+            raise ValueError(f"Unsafe base_dir: {base_dir}")
+
     logger.info("Using persistent directory: %s", base_dir)
 
     # Get ComfyUI path
@@ -267,10 +285,6 @@ def setup_persistence() -> str:
     # Set up environment
     os.environ["COMFY_SAVE_PATH"] = os.path.join(base_dir, "user")
 
-    # Set command line args if needed
-    if "--base-directory" not in sys.argv:
-        sys.argv.extend(["--base-directory", base_dir])
-
     # Patch folder_paths module
     patch_folder_paths(base_dir)