feat(api): add image with pytorch CUDA

ssube · Jan 7, 2023 · a721008 · a721008
1 parent 2ee0d27
commit a721008
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 3 deletions.
diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
@@ -104,12 +104,19 @@ build-gui-bundle:
 #     IMAGE_ARCH: alpine
 #     IMAGE_FILE: Containerfile.alpine
 
-package-api-oci-buster:
+package-api-oci-cpu-buster:
   extends:
     - .package-api-oci
   variables:
-    IMAGE_ARCH: buster
-    IMAGE_FILE: Containerfile.buster
+    IMAGE_ARCH: cpu-buster
+    IMAGE_FILE: Containerfile.cpu.buster
+
+package-api-oci-cuda-buster:
+  extends:
+    - .package-api-oci
+  variables:
+    IMAGE_ARCH: cuda-buster
+    IMAGE_FILE: Containerfile.cuda.buster
 
 package-gui-oci-nginx-alpine:
   extends:

diff --git a/api/Containerfile.alpine → api/Containerfile.cpu.alpine b/api/Containerfile.alpine → api/Containerfile.cpu.alpine
diff --git a/api/Containerfile.buster → api/Containerfile.cpu.buster b/api/Containerfile.buster → api/Containerfile.cpu.buster
diff --git a/api/Containerfile.cuda.buster b/api/Containerfile.cuda.buster
@@ -0,0 +1,15 @@
+FROM docker.io/python:3.10-buster
+
+WORKDIR /onnx-web
+
+# TODO: set up venv within the container
+
+RUN pip3 install torch torchvision torchaudio
+
+COPY requirements.txt /onnx-web/requirements.txt
+
+RUN pip install -r requirements.txt
+
+COPY onnx_web/ /onnx-web/onnx_web/
+
+ENTRYPOINT [ "flask", "--app=onnx_web.serve", "run", "--host=0.0.0.0" ]