Add cuda118 docker image build

hv0905 · Dec 27, 2023 · 4629fc2 · 4629fc2
1 parent 1694660
commit 4629fc2
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/.github/workflows/prod.yml b/.github/workflows/prod.yml
@@ -20,8 +20,16 @@ jobs:
             tags: |
               edgeneko/neko-image-gallery:latest
               edgeneko/neko-image-gallery:latest-cuda
+              edgeneko/neko-image-gallery:latest-cuda12.1
+            args: |
+              CUDA_VERSION=12.1
+          - dockerfile: "Dockerfile"
+            tags: edgeneko/neko-image-gallery:latest-cuda11.8
+            args: |
+              CUDA_VERSION=11.8
           - dockerfile: "cpu-only.Dockerfile"
             tags: edgeneko/neko-image-gallery:latest-cpu
+            args: ""
     steps:
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
@@ -36,3 +44,4 @@ jobs:
           file: ${{ matrix.configurations.dockerfile }}
           push: true
           tags: ${{ matrix.configurations.tags }}
+          args: ${{ matrix.configurations.args }}
diff --git a/Dockerfile b/Dockerfile
@@ -1,4 +1,6 @@
-FROM pytorch/pytorch:2.1.1-cuda12.1-cudnn8-runtime
+ARG TORCH_VERSION=2.1.2
+ARG CUDA_VERSION=12.1
+FROM pytorch/pytorch:${TORCH_VERSION}-cuda${CUDA_VERSION}-cudnn8-runtime
 
 WORKDIR /opt/NekoImageGallery