From 0a450d795dbb7fb80503cf51dbb00d74f971df82 Mon Sep 17 00:00:00 2001
From: James Brink <brink.james@gmail.com>
Date: Mon, 15 Dec 2025 22:27:03 -0700
Subject: [PATCH] fix: Resolve Docker container HOME environment and add
 cross-compilation support

- Fix $HOME not being set in Docker containers causing //.config error
- Add HOME=/root to both CPU and CUDA Docker image environments
- Update config.sh to use COMFY_USER_DIR with proper fallbacks
- Add --cpu flag to CPU Docker image to prevent CUDA init crash
- Refactor flake.nix with mkComfyPackages function for reusability
- Add cross-compilation support for building Linux images from macOS
- Add buildDockerLinux and buildDockerLinuxCuda apps for macOS users
- Remove outdated persistence warning from README
---
 README.md         |   2 -
 flake.nix         | 624 +++++++++++++++++++++++++++-------------------
 scripts/config.sh |   7 +-
 3 files changed, 368 insertions(+), 265 deletions(-)

diff --git a/README.md b/README.md
index a257b9d..4df1f6a 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,5 @@
 # ComfyUI Nix Flake
 
-**⚠️ NOTE: Model and workflow persistence should work but has not been thoroughly tested yet. Please report any issues.**
-
 A Nix flake for installing and running [ComfyUI](https://github.com/comfyanonymous/ComfyUI) with Python 3.12. Supports both macOS (Intel/Apple Silicon) and Linux with automatic GPU detection.
 
 ![ComfyUI Demo](comfyui-demo.gif)
diff --git a/flake.nix b/flake.nix
index 9415eab..16f842c 100644
--- a/flake.nix
+++ b/flake.nix
@@ -18,113 +18,106 @@
       comfyuiVersion = "0.4.0";
       comfyuiRev = "fc657f471a29d07696ca16b566000e8e555d67d1";
       comfyuiHash = "sha256-gq7/CfKqXGD/ti9ZeBVsHFPid+LTkpP4nTzt6NE/Jfo=";
-    in
-    flake-utils.lib.eachDefaultSystem (
-      system:
-      let
-        # Allow unfree packages
-        pkgs = import nixpkgs {
-          inherit system;
-          config = {
-            allowUnfree = true;
-            allowUnsupportedSystem = true;
+
+      # Function to build packages for a given pkgs
+      # This allows us to build for different target systems (cross-compilation)
+      mkComfyPackages =
+        {
+          pkgs,
+          forDocker ? false,
+        }:
+        let
+          # ComfyUI source
+          comfyui-src = pkgs.fetchFromGitHub {
+            owner = "comfyanonymous";
+            repo = "ComfyUI";
+            rev = comfyuiRev;
+            hash = comfyuiHash;
           };
-        };
 
-        # ComfyUI source
-        comfyui-src = pkgs.fetchFromGitHub {
-          owner = "comfyanonymous";
-          repo = "ComfyUI";
-          rev = comfyuiRev;
-          hash = comfyuiHash;
-        };
+          # Model downloader custom node
+          modelDownloaderDir = ./src/custom_nodes/model_downloader;
 
-        # Model downloader custom node
-        modelDownloaderDir = ./src/custom_nodes/model_downloader;
+          # Python environment with minimal dependencies for bootstrapping
+          pythonEnv = pkgs.python312.buildEnv.override {
+            extraLibs = with pkgs.python312Packages; [
+              setuptools
+              wheel
+              pip
+            ];
+            ignoreCollisions = true;
+          };
 
-        # Python environment with minimal dependencies for bootstrapping
-        # All ComfyUI dependencies are installed via pip in the virtual environment
-        pythonEnv = pkgs.python312.buildEnv.override {
-          extraLibs = with pkgs.python312Packages; [
-            setuptools
-            wheel
-            pip
-          ];
-          ignoreCollisions = true;
-        };
+          # Copy our persistence scripts to the nix store
+          persistenceScript = ./src/persistence/persistence.py;
+          persistenceMainScript = ./src/persistence/main.py;
+
+          # Process each script file individually
+          configScript = pkgs.substituteAll {
+            src = ./scripts/config.sh;
+            pythonEnv = pythonEnv;
+            comfyuiSrc = comfyui-src;
+            modelDownloaderDir = modelDownloaderDir;
+            persistenceScript = persistenceScript;
+            persistenceMainScript = persistenceMainScript;
+          };
 
-        # Copy our persistence scripts to the nix store
-        persistenceScript = ./src/persistence/persistence.py;
-        persistenceMainScript = ./src/persistence/main.py;
-
-        # Process each script file individually
-        configScript = pkgs.substituteAll {
-          src = ./scripts/config.sh;
-          pythonEnv = pythonEnv;
-          comfyuiSrc = comfyui-src;
-          modelDownloaderDir = modelDownloaderDir;
-          persistenceScript = persistenceScript;
-          persistenceMainScript = persistenceMainScript;
-        };
+          loggerScript = pkgs.substituteAll {
+            src = ./scripts/logger.sh;
+            pythonEnv = pythonEnv;
+          };
 
-        loggerScript = pkgs.substituteAll {
-          src = ./scripts/logger.sh;
-          pythonEnv = pythonEnv;
-        };
+          installScript = pkgs.substituteAll {
+            src = ./scripts/install.sh;
+            pythonEnv = pythonEnv;
+          };
 
-        installScript = pkgs.substituteAll {
-          src = ./scripts/install.sh;
-          pythonEnv = pythonEnv;
-        };
+          persistenceShScript = pkgs.substituteAll {
+            src = ./scripts/persistence.sh;
+            pythonEnv = pythonEnv;
+          };
 
-        persistenceShScript = pkgs.substituteAll {
-          src = ./scripts/persistence.sh;
-          pythonEnv = pythonEnv;
-        };
+          runtimeScript = pkgs.substituteAll {
+            src = ./scripts/runtime.sh;
+            pythonEnv = pythonEnv;
+          };
 
-        runtimeScript = pkgs.substituteAll {
-          src = ./scripts/runtime.sh;
-          pythonEnv = pythonEnv;
-        };
+          templateInputsScript = pkgs.substituteAll {
+            src = ./scripts/template_inputs.sh;
+            pythonEnv = pythonEnv;
+          };
 
-        templateInputsScript = pkgs.substituteAll {
-          src = ./scripts/template_inputs.sh;
-          pythonEnv = pythonEnv;
-        };
+          # Main launcher script with substitutions
+          launcherScript = pkgs.substituteAll {
+            src = ./scripts/launcher.sh;
+            pythonEnv = pythonEnv;
+            comfyuiSrc = comfyui-src;
+            modelDownloaderDir = modelDownloaderDir;
+            persistenceScript = persistenceScript;
+            persistenceMainScript = persistenceMainScript;
+            libPath = "${pkgs.stdenv.cc.cc.lib}/lib";
+          };
 
-        # Main launcher script with substitutions
-        launcherScript = pkgs.substituteAll {
-          src = ./scripts/launcher.sh;
-          pythonEnv = pythonEnv;
-          comfyuiSrc = comfyui-src;
-          modelDownloaderDir = modelDownloaderDir;
-          persistenceScript = persistenceScript;
-          persistenceMainScript = persistenceMainScript;
-          libPath = "${pkgs.stdenv.cc.cc.lib}/lib";
-        };
+          # Create a directory with all scripts
+          scriptDir = pkgs.runCommand "comfy-ui-scripts" { } ''
+            mkdir -p $out
+            cp ${configScript} $out/config.sh
+            cp ${loggerScript} $out/logger.sh
+            cp ${installScript} $out/install.sh
+            cp ${persistenceShScript} $out/persistence.sh
+            cp ${runtimeScript} $out/runtime.sh
+            cp ${templateInputsScript} $out/template_inputs.sh
+            cp ${launcherScript} $out/launcher.sh
+            chmod +x $out/*.sh
+          '';
 
-        # Create a directory with all scripts
-        scriptDir = pkgs.runCommand "comfy-ui-scripts" { } ''
-          mkdir -p $out
-          cp ${configScript} $out/config.sh
-          cp ${loggerScript} $out/logger.sh
-          cp ${installScript} $out/install.sh
-          cp ${persistenceShScript} $out/persistence.sh
-          cp ${runtimeScript} $out/runtime.sh
-          cp ${templateInputsScript} $out/template_inputs.sh
-          cp ${launcherScript} $out/launcher.sh
-          chmod +x $out/*.sh
-        '';
-
-        # Define all packages in one attribute set
-        packages = rec {
-          default = pkgs.stdenv.mkDerivation {
+          # Main comfy-ui package
+          comfyUiPackage = pkgs.stdenv.mkDerivation {
             pname = "comfy-ui";
             version = comfyuiVersion;
 
             src = comfyui-src;
 
-            # Passthru for scripting and testing
             passthru = {
               inherit comfyui-src;
               version = comfyuiVersion;
@@ -140,26 +133,18 @@
               pkgs.stdenv.cc.cc.lib
             ];
 
-            # Skip build and configure phases
             dontBuild = true;
             dontConfigure = true;
 
             installPhase = ''
-              # Create directories
               mkdir -p "$out/bin"
               mkdir -p "$out/share/comfy-ui"
 
-              # Copy ComfyUI files
               cp -r ${comfyui-src}/* "$out/share/comfy-ui/"
 
-              # Create scripts directory
               mkdir -p "$out/share/comfy-ui/scripts"
-
-              # Copy all script files
               cp -r ${scriptDir}/* "$out/share/comfy-ui/scripts/"
 
-              # Install the launcher script with wrapper for required tools
-              # curl and jq are needed for downloading workflow template input files
               makeWrapper "$out/share/comfy-ui/scripts/launcher.sh" "$out/bin/comfy-ui" \
                 --prefix PATH : "${
                   pkgs.lib.makeBinPath [
@@ -170,7 +155,6 @@
                   ]
                 }"
 
-              # Create alias for backwards compatibility
               ln -s "$out/bin/comfy-ui" "$out/bin/comfy-ui-launcher"
             '';
 
@@ -183,12 +167,11 @@
             };
           };
 
-          # Docker image for ComfyUI (CPU)
+          # Docker image (CPU)
           dockerImage = pkgs.dockerTools.buildImage {
             name = "comfy-ui";
             tag = "latest";
 
-            # Include essential utilities and core dependencies
             copyToRoot = pkgs.buildEnv {
               name = "root";
               paths = [
@@ -202,7 +185,7 @@
                 pkgs.libGL
                 pkgs.libGLU
                 pkgs.stdenv.cc.cc.lib
-                default
+                comfyUiPackage
               ];
               pathsToLink = [
                 "/bin"
@@ -212,14 +195,14 @@
               ];
             };
 
-            # Set up volumes and ports
             config = {
               Cmd = [
                 "/bin/bash"
                 "-c"
-                "export COMFY_USER_DIR=/data && mkdir -p /data && /bin/comfy-ui --listen 0.0.0.0"
+                "export COMFY_USER_DIR=/data && mkdir -p /data && /bin/comfy-ui --listen 0.0.0.0 --cpu"
               ];
               Env = [
+                "HOME=/root"
                 "COMFY_USER_DIR=/data"
                 "PATH=/bin:/usr/bin"
                 "PYTHONUNBUFFERED=1"
@@ -242,10 +225,10 @@
                   "localhost"
                   "8188"
                 ];
-                Interval = 30000000000; # 30 seconds in nanoseconds
-                Timeout = 5000000000; # 5 seconds in nanoseconds
+                Interval = 30000000000;
+                Timeout = 5000000000;
                 Retries = 3;
-                StartPeriod = 60000000000; # 60 seconds grace period for startup
+                StartPeriod = 60000000000;
               };
               Labels = {
                 "org.opencontainers.image.title" = "ComfyUI";
@@ -258,12 +241,11 @@
             };
           };
 
-          # Docker image for ComfyUI with CUDA support
+          # Docker image with CUDA support
           dockerImageCuda = pkgs.dockerTools.buildImage {
             name = "comfy-ui";
             tag = "cuda";
 
-            # Include essential utilities, core dependencies, and CUDA libraries
             copyToRoot = pkgs.buildEnv {
               name = "root";
               paths = [
@@ -277,7 +259,7 @@
                 pkgs.libGL
                 pkgs.libGLU
                 pkgs.stdenv.cc.cc.lib
-                default
+                comfyUiPackage
               ];
               pathsToLink = [
                 "/bin"
@@ -287,7 +269,6 @@
               ];
             };
 
-            # Set up volumes and ports
             config = {
               Cmd = [
                 "/bin/bash"
@@ -295,6 +276,7 @@
                 "export COMFY_USER_DIR=/data && mkdir -p /data && /bin/comfy-ui --listen 0.0.0.0"
               ];
               Env = [
+                "HOME=/root"
                 "COMFY_USER_DIR=/data"
                 "PATH=/bin:/usr/bin"
                 "PYTHONUNBUFFERED=1"
@@ -319,10 +301,10 @@
                   "localhost"
                   "8188"
                 ];
-                Interval = 30000000000; # 30 seconds in nanoseconds
-                Timeout = 5000000000; # 5 seconds in nanoseconds
+                Interval = 30000000000;
+                Timeout = 5000000000;
                 Retries = 3;
-                StartPeriod = 60000000000; # 60 seconds grace period for startup
+                StartPeriod = 60000000000;
               };
               Labels = {
                 "org.opencontainers.image.title" = "ComfyUI CUDA";
@@ -334,175 +316,301 @@
               };
             };
           };
+        in
+        {
+          default = comfyUiPackage;
+          inherit dockerImage dockerImageCuda;
+        };
+
+      # Map macOS systems to their Linux counterparts for cross-compilation
+      linuxSystemFor =
+        system:
+        if system == "aarch64-darwin" then
+          "aarch64-linux"
+        else if system == "x86_64-darwin" then
+          "x86_64-linux"
+        else
+          system;
+    in
+    flake-utils.lib.eachDefaultSystem (
+      system:
+      let
+        # Allow unfree packages
+        pkgs = import nixpkgs {
+          inherit system;
+          config = {
+            allowUnfree = true;
+            allowUnsupportedSystem = true;
+          };
+        };
+
+        # Build native packages
+        nativePackages = mkComfyPackages { inherit pkgs; };
+
+        # For macOS, also build Linux packages for Docker
+        linuxSystem = linuxSystemFor system;
+        isLinuxCrossCompile = system != linuxSystem;
+
+        linuxPkgs = import nixpkgs {
+          system = linuxSystem;
+          config = {
+            allowUnfree = true;
+            allowUnsupportedSystem = true;
+          };
+        };
+
+        linuxPackages = mkComfyPackages {
+          pkgs = linuxPkgs;
+          forDocker = true;
         };
+
+        # Python environment for dev shell (native only)
+        pythonEnv = pkgs.python312.buildEnv.override {
+          extraLibs = with pkgs.python312Packages; [
+            setuptools
+            wheel
+            pip
+          ];
+          ignoreCollisions = true;
+        };
+
+        # Define all packages
+        packages =
+          {
+            default = nativePackages.default;
+            dockerImage = nativePackages.dockerImage;
+            dockerImageCuda = nativePackages.dockerImageCuda;
+          }
+          // (
+            if isLinuxCrossCompile then
+              {
+                # Cross-compiled Linux Docker images (for macOS users)
+                dockerImageLinux = linuxPackages.dockerImage;
+                dockerImageLinuxCuda = linuxPackages.dockerImageCuda;
+              }
+            else
+              { }
+          );
       in
       {
         # Export packages
         inherit packages;
 
         # Define apps
-        apps = rec {
-          default = {
-            type = "app";
-            program = "${packages.default}/bin/comfy-ui";
-            meta = {
-              description = "Run ComfyUI with Nix";
-            };
-          };
-
-          # Add a buildDocker command
-          buildDocker =
-            let
-              script = pkgs.writeShellScriptBin "build-docker" ''
-                echo "Building Docker image for ComfyUI..."
-                # Load the Docker image directly
-                ${pkgs.docker}/bin/docker load < ${self.packages.${system}.dockerImage}
-                echo "Docker image built successfully! You can now run it with:"
-                echo "docker run -p 8188:8188 -v \$PWD/data:/data comfy-ui:latest"
-              '';
-            in
-            {
+        apps =
+          rec {
+            default = {
               type = "app";
-              program = "${script}/bin/build-docker";
+              program = "${packages.default}/bin/comfy-ui";
               meta = {
-                description = "Build ComfyUI Docker image (CPU)";
+                description = "Run ComfyUI with Nix";
               };
             };
 
-          # Add a buildDockerCuda command
-          buildDockerCuda =
-            let
-              script = pkgs.writeShellScriptBin "build-docker-cuda" ''
-                echo "Building Docker image for ComfyUI with CUDA support..."
-                # Load the Docker image directly
-                ${pkgs.docker}/bin/docker load < ${self.packages.${system}.dockerImageCuda}
-                echo "CUDA-enabled Docker image built successfully! You can now run it with:"
-                echo "docker run --gpus all -p 8188:8188 -v \$PWD/data:/data comfy-ui:cuda"
-                echo ""
-                echo "Note: Requires nvidia-container-toolkit and Docker GPU support."
-              '';
-            in
-            {
-              type = "app";
-              program = "${script}/bin/build-docker-cuda";
-              meta = {
-                description = "Build ComfyUI Docker image with CUDA support";
+            # Build native Docker image (matches host architecture)
+            buildDocker =
+              let
+                script = pkgs.writeShellScriptBin "build-docker" ''
+                  echo "Building Docker image for ComfyUI..."
+                  # Load the Docker image directly
+                  ${pkgs.docker}/bin/docker load < ${packages.dockerImage}
+                  echo "Docker image built successfully! You can now run it with:"
+                  echo "docker run -p 8188:8188 -v \$PWD/data:/data comfy-ui:latest"
+                '';
+              in
+              {
+                type = "app";
+                program = "${script}/bin/build-docker";
+                meta = {
+                  description = "Build ComfyUI Docker image (CPU)";
+                };
               };
-            };
 
-          # Update helper script
-          update = {
-            type = "app";
-            program = toString (
-              pkgs.writeShellScript "update-comfyui" ''
-                set -e
-                echo "Fetching latest ComfyUI release..."
-                LATEST=$(curl -s https://api.github.com/repos/comfyanonymous/ComfyUI/releases/latest | ${pkgs.jq}/bin/jq -r '.tag_name')
-                echo "Latest version: $LATEST"
-                echo ""
-                echo "To update, modify these values in flake.nix:"
-                echo "  comfyuiVersion = \"''${LATEST#v}\";"
-                echo ""
-                echo "Then run: nix flake update"
-                echo "And update the hash with: nix build 2>&1 | grep 'got:' | awk '{print \$2}'"
-              ''
-            );
-            meta = {
-              description = "Check for ComfyUI updates";
-            };
-          };
+            # Build native CUDA Docker image
+            buildDockerCuda =
+              let
+                script = pkgs.writeShellScriptBin "build-docker-cuda" ''
+                  echo "Building Docker image for ComfyUI with CUDA support..."
+                  # Load the Docker image directly
+                  ${pkgs.docker}/bin/docker load < ${packages.dockerImageCuda}
+                  echo "CUDA-enabled Docker image built successfully! You can now run it with:"
+                  echo "docker run --gpus all -p 8188:8188 -v \$PWD/data:/data comfy-ui:cuda"
+                  echo ""
+                  echo "Note: Requires nvidia-container-toolkit and Docker GPU support."
+                '';
+              in
+              {
+                type = "app";
+                program = "${script}/bin/build-docker-cuda";
+                meta = {
+                  description = "Build ComfyUI Docker image with CUDA support";
+                };
+              };
 
-          # Linting and formatting apps
-          lint =
-            let
-              script = pkgs.writeShellScriptBin "lint" ''
-                echo "Running ruff linter..."
-                ${pkgs.ruff}/bin/ruff check --no-cache src/
-              '';
-            in
-            {
+            # Update helper script
+            update = {
               type = "app";
-              program = "${script}/bin/lint";
+              program = toString (
+                pkgs.writeShellScript "update-comfyui" ''
+                  set -e
+                  echo "Fetching latest ComfyUI release..."
+                  LATEST=$(curl -s https://api.github.com/repos/comfyanonymous/ComfyUI/releases/latest | ${pkgs.jq}/bin/jq -r '.tag_name')
+                  echo "Latest version: $LATEST"
+                  echo ""
+                  echo "To update, modify these values in flake.nix:"
+                  echo "  comfyuiVersion = \"''${LATEST#v}\";"
+                  echo ""
+                  echo "Then run: nix flake update"
+                  echo "And update the hash with: nix build 2>&1 | grep 'got:' | awk '{print \$2}'"
+                ''
+              );
               meta = {
-                description = "Run ruff linter on Python code";
+                description = "Check for ComfyUI updates";
               };
             };
 
-          format =
-            let
-              script = pkgs.writeShellScriptBin "format" ''
-                echo "Formatting code with ruff..."
-                ${pkgs.ruff}/bin/ruff format --no-cache src/
-              '';
-            in
-            {
-              type = "app";
-              program = "${script}/bin/format";
-              meta = {
-                description = "Format Python code with ruff";
+            # Linting and formatting apps
+            lint =
+              let
+                script = pkgs.writeShellScriptBin "lint" ''
+                  echo "Running ruff linter..."
+                  ${pkgs.ruff}/bin/ruff check --no-cache src/
+                '';
+              in
+              {
+                type = "app";
+                program = "${script}/bin/lint";
+                meta = {
+                  description = "Run ruff linter on Python code";
+                };
               };
-            };
 
-          lint-fix =
-            let
-              script = pkgs.writeShellScriptBin "lint-fix" ''
-                echo "Running ruff linter with auto-fix..."
-                ${pkgs.ruff}/bin/ruff check --no-cache --fix src/
-              '';
-            in
-            {
-              type = "app";
-              program = "${script}/bin/lint-fix";
-              meta = {
-                description = "Run ruff linter with auto-fix";
+            format =
+              let
+                script = pkgs.writeShellScriptBin "format" ''
+                  echo "Formatting code with ruff..."
+                  ${pkgs.ruff}/bin/ruff format --no-cache src/
+                '';
+              in
+              {
+                type = "app";
+                program = "${script}/bin/format";
+                meta = {
+                  description = "Format Python code with ruff";
+                };
               };
-            };
 
-          type-check =
-            let
-              script = pkgs.writeShellScriptBin "type-check" ''
-                echo "Running pyright type checker..."
-                ${pkgs.pyright}/bin/pyright src/
-              '';
-            in
-            {
-              type = "app";
-              program = "${script}/bin/type-check";
-              meta = {
-                description = "Run pyright type checker on Python code";
+            lint-fix =
+              let
+                script = pkgs.writeShellScriptBin "lint-fix" ''
+                  echo "Running ruff linter with auto-fix..."
+                  ${pkgs.ruff}/bin/ruff check --no-cache --fix src/
+                '';
+              in
+              {
+                type = "app";
+                program = "${script}/bin/lint-fix";
+                meta = {
+                  description = "Run ruff linter with auto-fix";
+                };
               };
-            };
 
-          check-all =
-            let
-              script = pkgs.writeShellScriptBin "check-all" ''
-                echo "Running all checks..."
-                echo ""
-                echo "==> Running ruff linter..."
-                ${pkgs.ruff}/bin/ruff check --no-cache src/
-                RUFF_EXIT=$?
-                echo ""
-                echo "==> Running pyright type checker..."
-                ${pkgs.pyright}/bin/pyright src/
-                PYRIGHT_EXIT=$?
-                echo ""
-                if [ $RUFF_EXIT -eq 0 ] && [ $PYRIGHT_EXIT -eq 0 ]; then
-                  echo "All checks passed!"
-                  exit 0
-                else
-                  echo "Some checks failed."
-                  exit 1
-                fi
-              '';
-            in
-            {
-              type = "app";
-              program = "${script}/bin/check-all";
-              meta = {
-                description = "Run all Python code checks (ruff + pyright)";
+            type-check =
+              let
+                script = pkgs.writeShellScriptBin "type-check" ''
+                  echo "Running pyright type checker..."
+                  ${pkgs.pyright}/bin/pyright src/
+                '';
+              in
+              {
+                type = "app";
+                program = "${script}/bin/type-check";
+                meta = {
+                  description = "Run pyright type checker on Python code";
+                };
               };
-            };
-        };
+
+            check-all =
+              let
+                script = pkgs.writeShellScriptBin "check-all" ''
+                  echo "Running all checks..."
+                  echo ""
+                  echo "==> Running ruff linter..."
+                  ${pkgs.ruff}/bin/ruff check --no-cache src/
+                  RUFF_EXIT=$?
+                  echo ""
+                  echo "==> Running pyright type checker..."
+                  ${pkgs.pyright}/bin/pyright src/
+                  PYRIGHT_EXIT=$?
+                  echo ""
+                  if [ $RUFF_EXIT -eq 0 ] && [ $PYRIGHT_EXIT -eq 0 ]; then
+                    echo "All checks passed!"
+                    exit 0
+                  else
+                    echo "Some checks failed."
+                    exit 1
+                  fi
+                '';
+              in
+              {
+                type = "app";
+                program = "${script}/bin/check-all";
+                meta = {
+                  description = "Run all Python code checks (ruff + pyright)";
+                };
+              };
+          }
+          // (
+            # Add cross-compilation apps for macOS
+            if isLinuxCrossCompile then
+              {
+                # Build Linux Docker image from macOS
+                buildDockerLinux =
+                  let
+                    script = pkgs.writeShellScriptBin "build-docker-linux" ''
+                      echo "Building Linux Docker image for ComfyUI (cross-compiled from macOS)..."
+                      echo "Target architecture: ${linuxSystem}"
+                      # Load the Docker image directly
+                      ${pkgs.docker}/bin/docker load < ${packages.dockerImageLinux}
+                      echo ""
+                      echo "Linux Docker image built successfully! You can now run it with:"
+                      echo "docker run -p 8188:8188 -v \$PWD/data:/data comfy-ui:latest"
+                    '';
+                  in
+                  {
+                    type = "app";
+                    program = "${script}/bin/build-docker-linux";
+                    meta = {
+                      description = "Build Linux Docker image from macOS (cross-compile)";
+                    };
+                  };
+
+                # Build Linux CUDA Docker image from macOS
+                buildDockerLinuxCuda =
+                  let
+                    script = pkgs.writeShellScriptBin "build-docker-linux-cuda" ''
+                      echo "Building Linux CUDA Docker image for ComfyUI (cross-compiled from macOS)..."
+                      echo "Target architecture: ${linuxSystem}"
+                      # Load the Docker image directly
+                      ${pkgs.docker}/bin/docker load < ${packages.dockerImageLinuxCuda}
+                      echo ""
+                      echo "Linux CUDA Docker image built successfully! You can now run it with:"
+                      echo "docker run --gpus all -p 8188:8188 -v \$PWD/data:/data comfy-ui:cuda"
+                      echo ""
+                      echo "Note: Requires nvidia-container-toolkit and Docker GPU support."
+                    '';
+                  in
+                  {
+                    type = "app";
+                    program = "${script}/bin/build-docker-linux-cuda";
+                    meta = {
+                      description = "Build Linux CUDA Docker image from macOS (cross-compile)";
+                    };
+                  };
+              }
+            else
+              { }
+          );
 
         # Define development shell
         devShells.default = pkgs.mkShell {
@@ -558,7 +666,6 @@
                 cp -r $src source
                 chmod -R u+w source
                 cd source
-                # Disable cache to avoid permission issues in Nix sandbox
                 ${pkgs.ruff}/bin/ruff check --no-cache src/
                 touch $out
               '';
@@ -589,10 +696,7 @@
                 cp -r $src source
                 chmod -R u+w source
                 cd source/scripts
-                # Check launcher.sh with -x to follow all source statements
-                # This allows shellcheck to see variables defined in config.sh and used in install.sh
                 shellcheck -x launcher.sh
-                # Also check individual utility scripts
                 shellcheck logger.sh runtime.sh persistence.sh
                 touch $out
               '';
diff --git a/scripts/config.sh b/scripts/config.sh
index 03c5d70..0a61158 100644
--- a/scripts/config.sh
+++ b/scripts/config.sh
@@ -95,7 +95,8 @@ _parse_base_directory() {
 }
 
 # Parse and set BASE_DIR
-BASE_DIR="$HOME/.config/comfy-ui"
+# Priority: 1) --base-directory flag, 2) COMFY_USER_DIR env var, 3) $HOME/.config/comfy-ui default
+BASE_DIR="${COMFY_USER_DIR:-${HOME:-/root}/.config/comfy-ui}"
 _parsed_base_dir="$(_parse_base_directory "$@")"
 if [[ -n "$_parsed_base_dir" ]]; then
   BASE_DIR="$_parsed_base_dir"
@@ -160,8 +161,8 @@ 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"
+CODE_DIR="${HOME:-/root}/.config/comfy-ui/app"
+COMFY_VENV="${HOME:-/root}/.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"