From b35b397720497e7913254fcc95b524860f215197 Mon Sep 17 00:00:00 2001
From: "B. Bergeron" <me@bbergeron.xyz>
Date: Thu, 12 Feb 2026 00:59:50 -0500
Subject: [PATCH] Allow giving custom --extra-index-url, defaults to CU130

---
 Dockerfile    | 12 +++++++++++-
 README.md     |  7 +++++--
 compose.yaml  |  3 +++
 entrypoint.sh |  9 +++++----
 4 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index ccdaa4be3..764f0a6c0 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -40,6 +40,13 @@ USER comfyui
 RUN python -m venv .venv
 ENV PATH="/comfyui/.venv/bin:$PATH"
 
+# Take PIP_EXTRA_INDEX_URL as a required build argument and bake it into the
+# image as an environment variable. This variable is then used to select
+# hardware-specific flavors of PyTorch, enabling support for CPU, AMD, and newer
+# CUDA inference.
+ARG PIP_EXTRA_INDEX_URL
+ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+
 # Install ComfyUI's Python dependencies. Although dependency keeping is also
 # performed at startup, building ComfyUI's base dependencies into the image
 # significantly speeds up each containers' first run.
@@ -47,7 +54,10 @@ ENV PATH="/comfyui/.venv/bin:$PATH"
 # Since this step takes a long time to complete, it's performed early to take
 # advantage of Docker's build cache, thereby accelerating subsequent builds.
 COPY requirements.txt manager_requirements.txt ./
-RUN pip install --no-cache-dir --disable-pip-version-check \
+RUN pip install                                \
+    --no-cache-dir                             \
+    --disable-pip-version-check                \
+    --extra-index-url "${PIP_EXTRA_INDEX_URL}" \
     -r requirements.txt
 
 # Install ComfyUI
diff --git a/README.md b/README.md
index 77cd03596..0871bae06 100644
--- a/README.md
+++ b/README.md
@@ -358,8 +358,11 @@ For models compatible with Iluvatar Extension for PyTorch. Here's a step-by-step
 
 Start by installing Docker, Docker Compose, and the NVIDIA Container Toolkit on
 your host. Next, edit `compose.yaml` and update the `UID` and `GID` variables to
-match your user. Additional fields are documented in the file for further
-customization.
+match your host user's UID and GID. Also ensure that `PIP_EXTRA_INDEX_URL`
+points to the recommended non-experimental index URL for your hardware, as
+described in the beginning of the [manual installation
+section](#manual-install-windows-linux). Additional fields are documented in the
+file for further customization.
 
 Once ready, build and run the image locally:
 
diff --git a/compose.yaml b/compose.yaml
index d242b05a4..3c996612e 100644
--- a/compose.yaml
+++ b/compose.yaml
@@ -8,6 +8,9 @@ services:
       args:
         # Declare additional system dependencies for custom nodes
         APT_EXTRA_PACKAGES:
+        # Set the value of the --extra-index-url flag to select a specific CUDA,
+        # ROCm, or CPU version.
+        PIP_EXTRA_INDEX_URL: "https://download.pytorch.org/whl/cu130"
 
     ports:
       - 8188:8188
diff --git a/entrypoint.sh b/entrypoint.sh
index 37e576452..3cdce8bd9 100755
--- a/entrypoint.sh
+++ b/entrypoint.sh
@@ -48,10 +48,11 @@ done
 # PIP_EXTRA_PACKAGES.
 echo "[entrypoint] Updating Python dependencies..."
 su -c "
-   pip install                      \\
-        --no-cache-dir              \\
-        --disable-pip-version-check \\
-        -r requirements.txt         \\
+   pip install                                   \\
+        --no-cache-dir                           \\
+        --disable-pip-version-check              \\
+        --extra-index-url '$PIP_EXTRA_INDEX_URL' \\
+        -r requirements.txt                      \\
         $(find custom_nodes -mindepth 2 -maxdepth 2 -type f -name requirements.txt -printf "-r '%p' ") \\
         $PIP_EXTRA_PACKAGES
 " comfyui \