From 3de77f6e24901fa0c423f1e298a8be9f89f312cc Mon Sep 17 00:00:00 2001
From: doctorpangloss <@hiddenswitch.com>
Date: Wed, 24 Sep 2025 11:01:29 -0700
Subject: [PATCH] ROCm support

---
 .github/workflows/test.yml | 34 ++++++++++++++++++++--------------
 amd.Dockerfile             | 10 ++++------
 pyproject.toml             |  8 +++++---
 3 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 15b87ec6f..414c8fe82 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -79,33 +79,39 @@ jobs:
             friendly_name: "Python 3.12 ROCm 7.0 Torch 2.7.1"
             container: "rocm/pytorch:rocm7.0_ubuntu24.04_py3.12_pytorch_release_2.7.1"
     steps:
-      - run: |
-          apt-get update
-          # required for opencv
-          apt-get install --no-install-recommends -y ffmpeg libsm6 libxext6
-        name: Prepare Python
       - run: |
           curl -LsSf https://astral.sh/uv/install.sh | sh
         name: Install uv
+      - name: Align ROCm to Image
+        run: |
+          # our testing infrastructure uses RX 7XXX GPUs, this includes express support for gfx1102, gfx1101 which will resolve crashes
+          export UV_PRERELEASE=allow
+          uv pip install --index-url https://rocm.nightlies.amd.com/v2/gfx110X-dgpu/ "rocm[libraries]<7.9" "torch==2.7.1+rocm7.0.0rc20250922" "torchaudio==2.7.1a0+rocm7.0.0rc20250922" "torchvision==0.22.1+rocm7.0.0rc20250922" "pytorch_triton_rocm" --pre
       - uses: actions/checkout@v4
         name: Checkout git repo
+        with:
+          path: src
+      - run: |
+          apt install libsndfile1 -y --no-install-recommends
+        name: sndfile Dependency
+      - name: Write Overrides
+        run: |
+          echo ${GITHUB_WORKSPACE}
+          touch overrides.txt
+          uv pip freeze | grep -E 'torch|rocm|numpy|opencv' >> overrides.txt
+          echo "sentry-sdk; python_version < '0'" >> overrides.txt
       - name: Install ComfyUI
         run: |
-          export UV_BREAK_SYSTEM_PACKAGES=true
-          export UV_SYSTEM_PYTHON=true
-          pip freeze | grep nvidia >> overrides.txt; pip freeze | grep torch >> overrides.txt; pip freeze | grep opencv >> overrides.txt; pip freeze | grep numpy >> overrides.txt; echo "sentry-sdk; python_version < '0'" >> overrides.txt
+          export UV_PRERELEASE=allow
           export UV_OVERRIDE=overrides.txt
-          export UV_TORCH_BACKEND=auto
-          
-          # our testing infrastructure uses RX 7600, this includes express support for gfx1102
-          uv pip install --no-deps --index-url https://rocm.nightlies.amd.com/v2/gfx110X-dgpu/ "rocm[libraries,devel]"
-          
-          uv pip install --torch-backend=auto ".[rocm]"
+          uv pip install "comfyui[rocm,dev]@${GITHUB_WORKSPACE}/src"
       - name: Lint for errors
         run: |
+          cd ${GITHUB_WORKSPACE}/src
           pylint --rcfile=.pylintrc comfy/ comfy_extras/ comfy_api/ comfy_api_nodes/
       - name: Run unit tests
         run: |
+          cd ${GITHUB_WORKSPACE}/src
           pytest -v tests/unit
   build_and_execute_macos:
     environment: "Testing"
diff --git a/amd.Dockerfile b/amd.Dockerfile
index e2fbca4b1..19f743e81 100644
--- a/amd.Dockerfile
+++ b/amd.Dockerfile
@@ -4,8 +4,6 @@ ENV TZ="Etc/UTC"
 
 ENV UV_COMPILE_BYTECODE=1
 ENV UV_NO_CACHE=1
-ENV UV_SYSTEM_PYTHON=1
-ENV UV_BREAK_SYSTEM_PACKAGES=1
 ENV PIP_DISABLE_PIP_VERSION_CHECK=1
 ENV PIP_NO_CACHE_DIR=1
 ENV DEBIAN_FRONTEND=noninteractive
@@ -13,15 +11,15 @@ ENV DEBIAN_FRONTEND=noninteractive
 ENV LANG=C.UTF-8
 ENV LC_ALL=C.UTF-8
 
-RUN pip freeze | grep nvidia >> /overrides.txt; pip freeze | grep torch >> /overrides.txt; pip freeze | grep opencv >> /overrides.txt; pip freeze | grep numpy >> /overrides.txt; echo "sentry-sdk; python_version < '0'" >> /overrides.txt
+RUN uv pip freeze | grep torch >> /overrides.txt; uv pip freeze | grep opencv >> /overrides.txt; uv pip freeze | grep numpy >> /overrides.txt; uv pip freeze | grep rocm >> /overrides.txt; echo "sentry-sdk; python_version < '0'" >> /overrides.txt
 
 ENV UV_OVERRIDE=/overrides.txt
-ENV UV_TORCH_BACKEND=auto
+env UV_PRERELEASE=allow
 
 # mitigates AttributeError: module 'cv2.dnn' has no attribute 'DictValue' \
 # see https://github.com/facebookresearch/nougat/issues/40
 RUN apt-get update && \
-    apt-get install --no-install-recommends -y ffmpeg libsm6 libxext6 && \
+    apt-get install --no-install-recommends -y ffmpeg libsm6 libxext6 libsndfile1 && \
     pip install uv && uv --version && \
     apt-get purge -y && \
     rm -rf /var/lib/apt/lists/*
@@ -42,7 +40,7 @@ RUN uv pip install $SOURCES
 WORKDIR /workspace
 # addresses https://github.com/pytorch/pytorch/issues/104801
 # and issues reported by importing nodes_canny
-RUN comfyui --quick-test-for-ci --cpu --cwd /workspace
+RUN python -c "import torch; import transformers; from transformers import AutoProcessor, BatchFeature; import torchaudio; import cv2" && comfyui --quick-test-for-ci --cpu --cwd /workspace
 
 EXPOSE 8188
 CMD ["python", "-m", "comfy.cmd.main", "--listen"]
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 808b3da46..b8f94cb2a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,7 +27,7 @@ dependencies = [
     "torchsde>=0.2.6",
     "einops>=0.6.0",
     "open-clip-torch>=2.24.0",
-    "transformers!=4.53.0,!=4.53.1,!=4.53.2",
+    "transformers>=4.46.0,!=4.53.0,!=4.53.1,!=4.53.2",
     "tokenizers>=0.13.3",
     "sentencepiece",
     "peft>=0.10.0",
@@ -161,7 +161,8 @@ cu128 = [
 rocm = [
     "torch",
     "torchvision",
-    "pytorch-triton-rocm"
+    "pytorch-triton-rocm",
+    "rocm[libraries,devel]",
 ]
 
 mps = [
@@ -204,6 +205,7 @@ no-build-isolation-package = [
     "sage-attention",
     "flash_attn",
 ]
+prerelease = "allow"
 
 [[tool.uv.index]]
 name = "pytorch-cpu"
@@ -224,7 +226,7 @@ explicit = true
 name = "pytorch-rocm"
 url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0/"
 format = "flat"
-explicit = true
+explicit = false
 
 [[tool.uv.index]]
 name = "pytorch-nightly-cpu"