Create separate docker images for python 3.10 and 3.11

README.md

@@ -62,7 +62,7 @@ docker run \
     --network host \
     --name skynet-worker \
     --mount type=bind,source="$(pwd)",target=/root/target \
-    guilledk/skynet:runtime-cuda \
+    guilledk/skynet:runtime-cuda-py310 \
     skynet run dgpu
 
 # run worker with specific gpu
@@ -74,6 +74,6 @@ docker run \
     --network host \
     --name skynet-worker-1 \
     --mount type=bind,source="$(pwd)",target=/root/target \
-    guilledk/skynet:runtime-cuda \
+    guilledk/skynet:runtime-cuda-py310 \
     skynet run dgpu
 ```

build_docker.sh

@@ -1,3 +1,7 @@
 docker build \
-    -t guilledk/skynet:runtime-cuda \
-    -f docker/Dockerfile.runtime+cuda .
+    -t guilledk/skynet:runtime-cuda-py310 \
+    -f docker/Dockerfile.runtime+cuda-py310 .
+
+docker build \
+    -t guilledk/skynet:runtime-cuda-py311 \
+    -f docker/Dockerfile.runtime+cuda-py311 .

docker/Dockerfile.runtime+cuda-py310

@@ -0,0 +1,46 @@
+from nvidia/cuda:11.8.0-devel-ubuntu20.04
+from python:3.10
+
+env DEBIAN_FRONTEND=noninteractive
+
+run apt-get update && apt-get install -y \
+    git \
+    clang \
+    cmake \
+    ffmpeg \
+    libsm6 \
+    libxext6 \
+    ninja-build
+
+env CC /usr/bin/clang
+env CXX /usr/bin/clang++
+
+# install llvm10 as required by llvm-lite
+run git clone https://github.com/llvm/llvm-project.git -b llvmorg-10.0.1
+workdir /llvm-project
+# this adds a commit from 12.0.0 that fixes build on newer compilers
+run git cherry-pick -n b498303066a63a203d24f739b2d2e0e56dca70d1
+run cmake -S llvm -B build -G Ninja -DCMAKE_BUILD_TYPE=Release
+run ninja -C build install  # -j8
+
+run curl -sSL https://install.python-poetry.org | python3 -
+
+env PATH "/root/.local/bin:$PATH"
+
+copy . /skynet
+
+workdir /skynet
+
+env POETRY_VIRTUALENVS_PATH /skynet/.venv
+
+run poetry install --with=cuda -v
+
+workdir /root/target
+
+env PYTORCH_CUDA_ALLOC_CONF max_split_size_mb:128
+env NVIDIA_VISIBLE_DEVICES=all
+
+copy docker/entrypoint.sh /entrypoint.sh
+entrypoint ["/entrypoint.sh"]
+
+cmd ["skynet", "--help"]

skynet/utils.py

@@ -114,7 +114,6 @@ def pipeline_for(
         else:
             pipe_class = StableDiffusionPipeline
 
-    breakpoint()
     pipe = pipe_class.from_pretrained(
         model, **params)
 
@@ -123,6 +122,11 @@ def pipeline_for(
 
     pipe.enable_xformers_memory_efficient_attention()
 
+    if sys.version_info[1] < 11:
+        # torch.compile only supported on python < 3.11
+        pipe.unet = torch.compile(
+            pipe.unet, mode='reduce-overhead', fullgraph=True)
+
     if over_mem:
         if not image:
             pipe.enable_vae_slicing()