feat: add cudaCapabilities option to NixOS module (#20)

* feat: add cudaCapabilities option to NixOS module

- Add optional cudaCapabilities setting that maps to nixpkgs.config
- Document common CUDA compute capability values with examples
- Improve AGENTS.md with comprehensive development guidelines
- Fix platform constraints for CUDA packages (x86_64-linux only)

Closes #19

* fix: restore dockerImage for aarch64-linux

The previous change incorrectly restricted dockerImage to x86_64-linux only.
This broke ARM64 Docker builds in CI. Now:
- dockerImage is available for all Linux platforms
- cuda and dockerImageCuda remain x86_64-linux only (as intended)

* docs: clarify cudaCapabilities doesn't affect pre-built PyTorch wheels

Address review feedback by adding explicit documentation that:
- Pre-built PyTorch wheels already support all GPU architectures
- This setting primarily benefits other system CUDA packages
- The setting affects global nixpkgs configuration

Thanks to the reviewer for the excellent and thorough review!

Author: jamesbrink

AGENTS.md

@@ -1,41 +1,135 @@
-# Repository Guidelines
+# Repository Guidelines for Agents
 
-## Project Structure & Module Organization
+## Purpose
 This repo is a Nix flake that packages ComfyUI plus curated custom nodes.
+Use this file as the single source of truth for build, lint, test, and style
+expectations when working in this repository.
+
+## Project Structure & Module Organization
 - `nix/`: flake helpers, package definitions, module, checks, and version pins.
 - `src/custom_nodes/`: bundled custom nodes (Python and frontend assets).
 - `src/patches/`: patches applied to upstream components.
 - `docs/`: additional documentation (setup guides, etc.).
 - `scripts/`: maintenance scripts (template inputs, Cachix, model downloads).
 - `data/`: local runtime data for container examples and demos.
 
-## Build, Test, and Development Commands
-- `nix run`: run ComfyUI (add `-- --open`, `--port=XXXX`, `--listen 0.0.0.0` as needed).
-- `nix run .#cuda`: CUDA build on Linux/NVIDIA; use only when GPU support is required.
+## Development Environment
+- Use `nix develop` to enter the dev shell when running linting or formatting.
+- The dev shell provides `ruff`, `pyright`, `nixfmt`, and `shellcheck`.
+- Python runtime target is 3.12 (see `pyproject.toml`).
+
+## Build, Lint, Format, and Test Commands
+### Core Nix commands
 - `nix build`: build the default package without running it.
-- `nix develop`: enter the dev shell with `ruff`, `pyright`, and `nixfmt`.
+- `nix run`: run ComfyUI (add `-- --open`, `--port=XXXX`, `--listen 0.0.0.0`).
+- `nix run .#cuda`: CUDA build on Linux/NVIDIA only.
 - `nix flake check`: run all checks (build, ruff, pyright, nixfmt, shellcheck).
 - `nix fmt`: format Nix files with `nixfmt-rfc-style`.
 - `nix run .#update`: check for ComfyUI version updates.
-- `nix run .#buildDocker` / `nix run .#buildDockerCuda`: build CPU/CUDA images locally.
+- `nix run .#buildDocker`: build CPU Docker image.
+- `nix run .#buildDockerCuda`: build CUDA Docker image.
+
+### Python linting and formatting
+- `ruff check src/`: run linting on all Python sources.
+- `ruff check src/path/to/file.py`: lint a single Python file.
+- `ruff format src/`: format Python sources.
+- `ruff format src/path/to/file.py`: format a single Python file.
+
+### Python type checking
+- `pyright src/`: run type checks for the source tree.
+- `pyright src/path/to/file.py`: run type checks on a single Python file.
+
+### Shell scripting
+- `shellcheck scripts/*.sh`: lint shell scripts.
+- `shellcheck scripts/path/to/script.sh`: lint a single script.
+
+### Tests
+- There are no unit-test suites in this repo today.
+- The primary quality gate is `nix flake check`.
+- If tests are added in the future, prefer `pytest path/to/test.py::test_name`
+  for a single test and `pytest path/to/test.py` for a single file.
+
+## Coding Style & Conventions
+### Python formatting
+- Python is formatted with `ruff format`.
+- Use 4-space indentation and a 100-character line length.
+- Use double quotes for strings unless escaping makes it noisy.
+- Prefer trailing commas in multi-line collections.
+
+### Imports
+- Order imports as standard library, third-party, then local.
+- Ruff enforces import ordering (isort).
+- Avoid unused imports except in `__init__.py` where they are allowed.
+- Star imports are allowed only in `src/custom_nodes/**` modules.
+
+### Naming
+- `snake_case` for functions, modules, variables.
+- `PascalCase` for classes.
+- `UPPER_SNAKE_CASE` for constants.
+- Avoid single-letter variable names unless the context is obvious.
+
+### Types
+- Python version is 3.12 (`pyright` configured accordingly).
+- Type checking mode is `basic`; be pragmatic with annotations.
+- Add type hints to new public functions and key data structures.
+- Prefer `typing` (e.g., `Iterable`, `Mapping`, `Sequence`) over concrete
+  container types in APIs.
+- Use `typing.cast` only when necessary and document intent.
+- Use `from __future__ import annotations` when it improves readability.
+
+### Docstrings
+- Google-style docstrings are preferred.
+- Document public functions and classes when behavior is non-obvious.
+- Keep docstrings concise and action-oriented.
+
+### Error handling
+- Raise specific exceptions with actionable messages.
+- Preserve context by chaining (`raise ... from err`) when wrapping errors.
+- Avoid bare `except:`; catch explicit exception types.
+- It is acceptable to use `print` for status updates (ComfyUI convention).
+
+### Filesystem usage
+- `os.path` is acceptable (the repo does not mandate `pathlib`).
+- Prefer `pathlib.Path` only when it improves clarity or APIs require it.
+- Avoid writing outside the repo unless required by a script.
+
+### Performance and readability
+- Keep functions small and focused.
+- Avoid deeply nested conditionals; early returns are encouraged.
+- Prefer comprehensions when they remain readable.
+- Do not micro-optimize unless it is on a hot path.
+
+## Repo-Specific Ruff and Pyright Notes
+- Ruff targets Python 3.12 with 100-character lines.
+- Ruff ignores `T20` (print statements) and `G004` (f-strings in logging).
+- Ruff allows unused imports in `__init__.py` and star imports in
+  `src/custom_nodes/**`.
+- Ruff complexity guard: `max-complexity = 15` (McCabe).
+- Pyright uses `basic` type checking with lenient missing import handling.
+- Pyright is configured to warn (not fail) on missing imports.
+- Pyright uses `typings/` as the stub path when available.
+
+## Nix Style
+- Format Nix files with `nix fmt`.
+- Keep attrsets aligned and avoid reformatting unrelated sections.
+- Prefer existing patterns in `nix/` for adding new derivations.
 
-## Coding Style & Naming Conventions
-- Python: 4-space indentation, 100-char lines, double quotes (`ruff format`).
-- Imports: standard library, third-party, then local; `ruff` enforces ordering.
-- Naming: `snake_case` for functions/vars, `PascalCase` for classes.
-- Nix: format with `nix fmt` (nixfmt RFC style).
-- Shell: scripts in `scripts/` should pass `shellcheck`.
+## Shell Style
+- Shell scripts should pass `shellcheck`.
+- Prefer `set -euo pipefail` for new scripts when appropriate.
+- Avoid bashisms if the script might be invoked by `sh`.
+- Keep scripts small and focused; prefer helpers in `scripts/`.
 
-## Testing Guidelines
-There are no unit-test suites in the repo today. Quality gates are:
-- `nix flake check` (ruff + pyright + nixfmt + shellcheck).
-- Optional manual runs: `ruff check src/` and `pyright src/` inside `nix develop`.
+## Cursor and Copilot Rules
+- No `.cursorrules` or `.cursor/rules/` found in this repo.
+- No `.github/copilot-instructions.md` found in this repo.
 
 ## Commit & Pull Request Guidelines
-- Commit style follows a conventional prefix, e.g. `docs: ...`, `ci: ...`, `refactor: ...`.
-- Keep messages short and specific to the change.
-- PRs should include a concise description, linked issues if applicable, and the checks run
-  (e.g. `nix flake check` or relevant `nix run`/`nix build` commands).
+- Commit style follows a conventional prefix, e.g. `docs: ...`, `ci: ...`,
+  `refactor: ...`.
+- Keep commit messages short and specific to the change.
+- PRs should include a concise description, linked issues if applicable, and
+  the checks run (e.g. `nix flake check`, `nix build`).
 
 ## Agent-Specific Instructions
 - After any file edits, run `git add` so the flake can see changes correctly.

README.md

@@ -302,6 +302,8 @@ nix profile add github:utensils/comfyui-nix#cuda
   services.comfyui = {
     enable = true;
     cuda = true;  # Enable NVIDIA GPU acceleration (recommended for most users)
+    # cudaCapabilities = [ "8.9" ];  # Optional: optimize system CUDA packages for RTX 40xx
+    #   Note: Pre-built PyTorch wheels already support all GPU architectures
     enableManager = true;  # Enable the built-in ComfyUI Manager
     port = 8188;
     listenAddress = "127.0.0.1";  # Use "0.0.0.0" for network access
@@ -319,6 +321,7 @@ nix profile add github:utensils/comfyui-nix#cuda
 | --------------- | -------------------- | ------------------------------------------------ |
 | `enable`        | `false`              | Enable the ComfyUI service                       |
 | `cuda`          | `false`              | Enable NVIDIA GPU acceleration                   |
+| `cudaCapabilities` | `null`           | Optional CUDA compute capability list            |
 | `enableManager` | `false`              | Enable the built-in ComfyUI Manager              |
 | `port`          | `8188`               | Port for the web interface                       |
 | `listenAddress` | `"127.0.0.1"`        | Listen address (`"0.0.0.0"` for network access)  |
@@ -332,6 +335,9 @@ nix profile add github:utensils/comfyui-nix#cuda
 | `customNodes`   | `{}`                 | Declarative custom nodes (see below)             |
 | `requiresMounts`| `[]`                 | Mount units to wait for before starting          |
 
+`cudaCapabilities` maps to `nixpkgs.config.cudaCapabilities`, so setting it will
+apply to other CUDA packages in the system configuration as well.
+
 **Note:** When `dataDir` is under `/home/`, `ProtectHome` is automatically disabled to allow access.
 
 ### Using a Home Directory

flake.nix

@@ -152,9 +152,11 @@
           dockerImageLinuxArm64 = linuxArm64Packages.dockerImage;
         }
         // pkgs.lib.optionalAttrs pkgs.stdenv.isLinux {
+          dockerImage = nativePackages.dockerImage;
+        }
+        // pkgs.lib.optionalAttrs (pkgs.stdenv.isLinux && pkgs.stdenv.isx86_64) {
           # CUDA package uses pre-built wheels (supports all GPU architectures)
           cuda = nativePackagesCuda.default;
-          dockerImage = nativePackages.dockerImage;
           dockerImageCuda = nativePackagesCuda.dockerImageCuda;
         };
       in
@@ -223,12 +225,12 @@
         comfyui-nix = self.legacyPackages.${final.system};
         comfyui = self.packages.${final.system}.default;
         comfy-ui = self.packages.${final.system}.default;
-        # CUDA variant (Linux only) - uses pre-built wheels supporting all GPU architectures
+        # CUDA variant (x86_64 Linux only) - uses pre-built wheels supporting all GPU architectures
         comfy-ui-cuda =
-          if final.stdenv.isLinux then
+          if final.stdenv.isLinux && final.stdenv.isx86_64 then
             self.packages.${final.system}.cuda
           else
-            throw "comfy-ui-cuda is only available on Linux";
+            throw "comfy-ui-cuda is only available on x86_64 Linux";
         # Add custom nodes to overlay
         comfyui-custom-nodes = self.legacyPackages.${final.system}.customNodes;
       };

nix/modules/comfyui.nix

@@ -81,6 +81,33 @@ in
       '';
     };
 
+    cudaCapabilities = lib.mkOption {
+      type = lib.types.nullOr (lib.types.listOf lib.types.str);
+      default = null;
+      description = ''
+        Optional list of CUDA compute capabilities to use for builds that honor
+        `nixpkgs.config.cudaCapabilities`. When set, this updates the global
+        nixpkgs configuration, so it affects other CUDA packages too.
+
+        Note: ComfyUI's pre-built PyTorch wheels already support all GPU
+        architectures (Pascal through Hopper). This setting is primarily useful
+        for optimizing other CUDA-enabled packages in your system configuration.
+
+        Example: [ "8.9" ] for Ada Lovelace (RTX 40xx) GPUs.
+
+        Common values:
+        - "6.1": Pascal (GTX 1080/1070)
+        - "7.0": Volta (V100)
+        - "7.5": Turing (RTX 20xx, GTX 16xx)
+        - "8.0": Ampere (A100)
+        - "8.6": Ampere (RTX 30xx)
+        - "8.9": Ada Lovelace (RTX 40xx)
+        - "9.0": Hopper (H100)
+
+        See: https://developer.nvidia.com/cuda-gpus
+      '';
+    };
+
     enableManager = lib.mkOption {
       type = lib.types.bool;
       default = false;
@@ -216,6 +243,10 @@ in
   };
 
   config = lib.mkIf cfg.enable {
+    nixpkgs.config = lib.mkIf (cfg.cudaCapabilities != null) {
+      cudaCapabilities = cfg.cudaCapabilities;
+    };
+
     # Create system user/group only when using default "comfyui" names
     users.users = lib.mkIf (cfg.createUser && isSystemUser) {
       ${cfg.user} = {