Use official torch image (#1173)

siliconflow · Dec 27, 2024 · a919eba · a919eba
1 parent 144f9c2
commit a919eba
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 14 deletions.
diff --git a/.github/workflows/sd.yml b/.github/workflows/sd.yml
@@ -23,29 +23,25 @@ concurrency:
 
 jobs:
   build-oneflow:
-    name: Build sd ${{ matrix.image.repo }}:${{ matrix.image.suffix }}
+    name: Build sd ${{ matrix.image.repo }}:${{ matrix.image.tag }}
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       max-parallel: 5
       matrix:
         image:
           - {
-              BASE_IMAGE: "nvcr.io/nvidia/pytorch:23.08-py3",
+              BASE_IMAGE: "pytorch/pytorch:2.3.1-cuda11.8-cudnn8-runtime",
               suffix: "cu118",
+              tag: "torch2.3-cuda11.8",
               repo: "onediff",
             }
           - {
-              BASE_IMAGE: "nvcr.io/nvidia/pytorch:23.08-py3",
-              suffix: "cu121",
-              repo: "onediff-pro",
-            }
-          - {
-              BASE_IMAGE: "nvcr.io/nvidia/pytorch:23.08-py3",
+              BASE_IMAGE: "pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime",
               suffix: "cu122",
+              tag: "torch2.3-cuda12.1",
               repo: "onediff-pro",
             }
-
     env:
       DOCKER_BUILDKIT: 1
     steps:
@@ -79,8 +75,8 @@ jobs:
       # 1.2 Build and push image to ACR
       - name: Set env
         run: |
-          echo ACR_TAG="$REGISTRY/$NAMESPACE/${{ matrix.image.repo }}:${{ matrix.image.suffix }}" >> $GITHUB_ENV
-          echo DOCKERHUB_TAG="oneflowinc/${{ matrix.image.repo }}:${{ matrix.image.suffix }}" >> $GITHUB_ENV
+          echo ACR_TAG="$REGISTRY/$NAMESPACE/${{ matrix.image.repo }}:${{ matrix.image.tag }}" >> $GITHUB_ENV
+          echo DOCKERHUB_TAG="oneflowinc/${{ matrix.image.repo }}:${{ matrix.image.tag }}" >> $GITHUB_ENV
       - name: Set env (Pro)
         if: matrix.image.repo == 'onediff-pro'
         run: |
@@ -90,7 +86,7 @@ jobs:
         run: |
           curl "${{ env.ONEFLOW_PIP_INDEX }}"
           ONEFLOW_PIP_INDEX=${{ env.ONEFLOW_PIP_INDEX }}
-          DEV_VERSION=$(curl ${{ env.ONEFLOW_PIP_INDEX }} | rg '/oneflow-(\d.+?)%2B(.+?)-cp' -or '$1+${2}' | head -n 1)
+          DEV_VERSION=$(curl ${{ env.ONEFLOW_PIP_INDEX }} | rg '/oneflow-(\d.+?)%2B(.+?)-cp' -or '$1+${2}' | sort | tail -n 1)
           echo "DEV_VERSION: ${DEV_VERSION}"
           if [ -z "$DEV_VERSION" ]
           then

diff --git a/Dockerfile b/Dockerfile
@@ -4,8 +4,8 @@ FROM ${BASE_IMAGE}
 
 ARG ONEFLOW_PIP_INDEX
 ARG ONEFLOW_PACKAGE_NAME=oneflow
-RUN pip install -f ${ONEFLOW_PIP_INDEX} ${ONEFLOW_PACKAGE_NAME} "nvidia-cudnn-cu11>=8.9,<9.0"
-RUN python3 -m pip install "torch" "transformers==4.27.1" "diffusers[torch]==0.19.3" "huggingface-hub==0.23.2" nvidia-cuda-cupti-cu12 nvidia-nvjitlink-cu12
+RUN pip install -f ${ONEFLOW_PIP_INDEX} ${ONEFLOW_PACKAGE_NAME}
+RUN python3 -m pip install "torch" "transformers==4.27.1" "diffusers[torch]==0.19.3" "huggingface-hub==0.23.2"
 ADD . /src/onediff
 RUN python3 -m pip install -e /src/onediff
 RUN python3 -m pip install -e /src/onediff/onediff_diffusers_extensions