Update comfyui-rocm-dockerfile

comfyui-rocm-dockerfile

@@ -1,45 +1,70 @@
-# Use the official ROCm base image
+# --- Stage 1: The Builder ---
+FROM ubuntu:24.04 AS builder
+
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt update && apt install -y \
+    git cmake ninja-build python3-pip python3-venv \
+    libnuma-dev libpci-dev libelf-dev libdrm-dev \
+    build-essential g++-13 gcc-13
+
+WORKDIR /src
+RUN git clone https://github.com/ROCm/TheRock.git
+WORKDIR /src/TheRock
+
+# Configure for gfx1100 (7900 XTX) on Aarch64
+# We use a venv to manage the build tools
+RUN python3 -m venv venv && \
+    ./venv/bin/pip install --upgrade pip && \
+    ./venv/bin/pip install -r requirements.txt
+
+# Compile a subset of ROCm (HIP + KFD) needed for runtime
+RUN cmake -S . -B build \
+    -G Ninja \
+    -DTHEROCK_AMDGPU_TARGETS="gfx1100" \
+    -DCMAKE_INSTALL_PREFIX=/opt/rocm \
+    -DTHEROCK_ENABLE_ALL=OFF \
+    -DTHEROCK_ENABLE_HIP=ON \
+    -DTHEROCK_ENABLE_ROCBLAS=ON 
+
+RUN cmake --build build --target install
+
+# --- Stage 2: Final Runtime ---
 FROM ubuntu:24.04
 
-# Set environment variables for 7900 XTX (gfx1100)
+# Core Environment Variables
 ENV HSA_OVERRIDE_GFX_VERSION=11.0.0
 ENV TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
-ENV LD_LIBRARY_PATH=/usr/lib/aarch64-linux-gnu:$LD_LIBRARY_PATH
+ENV ROCM_PATH=/opt/rocm
+ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/lib/aarch64-linux-gnu:$LD_LIBRARY_PATH
 ENV PIP_BREAK_SYSTEM_PACKAGES=1
 ENV PYTHONUNBUFFERED=1
 
+# Copy compiled ROCm from builder
+COPY --from=builder /opt/rocm /opt/rocm
+
 RUN apt update -y && apt install -y --no-install-recommends \
-    python3-pip \
-    python3-dev \
-    git \
-    libgl1 \
-    libglib2.0-0 \
+    python3-pip python3-dev git libgl1 libglib2.0-0 \
+    libnuma1 libpci3 libelf1 \
     && rm -rf /var/lib/apt/lists/*
 
-# Set working directory
 WORKDIR /app
 
 # Clone ComfyUI
 RUN git clone https://github.com/comfyanonymous/ComfyUI.git .
 
-# Install PyTorch with ROCm 6.2 support (matching the base image)
+# Install PyTorch Aarch64 ROCm Wheels
+# NOTE: As of 2026, check the PyTorch nightly/stable index for the Aarch64-specific ROCm wheels
 RUN pip3 install --no-cache-dir torch torchvision torchaudio \
     --index-url https://download.pytorch.org/whl/rocm6.2
 
 # Install ComfyUI dependencies
 RUN pip3 install --no-cache-dir -r requirements.txt
 
-# Set custom Nodes
+# Setup Manager
 WORKDIR /app/custom_nodes
-
-# Install ComfyUi Manager
 RUN git clone https://github.com/Comfy-Org/ComfyUI-Manager.git
 
-# Go back to root
 WORKDIR /app
-
-# Expose the default ComfyUI port
 EXPOSE 8188
 
-# Launch ComfyUI
 CMD ["python3", "main.py", "--listen", "0.0.0.0"]