From 5342a1f861acb646de12469fceefe9e29a4896cb Mon Sep 17 00:00:00 2001
From: Gareth Davidson <gaz@bitplane.net>
Date: Sat, 7 Dec 2024 17:39:38 +0000
Subject: [PATCH 1/3] add Dockerfile and some post build steps

Add a post-build script that runs the first time you execute it on
a machine with a GPU, since some stuff won't build in the builder

Changed app.py to serve on 0.0.0.0, 'cause otherwise we can't expose
the ports
---
 .dockerignore   |  4 +++
 Dockerfile      | 88 +++++++++++++++++++++++++++++++++++++++++++++++++
 app.py          |  2 +-
 onstart.sh      | 12 +++++++
 post_install.sh | 25 ++++++++++++++
 5 files changed, 130 insertions(+), 1 deletion(-)
 create mode 100644 .dockerignore
 create mode 100644 Dockerfile
 create mode 100755 onstart.sh
 create mode 100755 post_install.sh

diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 00000000..4c4c570a
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,4 @@
+.git
+.env
+*.pyc
+Dockerfile
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 00000000..13a3c91b
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,88 @@
+FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel AS builder
+
+# Install build dependencies
+RUN apt-get update && \
+    apt-get install -y ffmpeg build-essential git rdfind
+
+WORKDIR /app
+ADD setup.sh   /app/setup.sh
+
+# Setup conda
+RUN conda config --set always_yes true && \
+    conda init && \
+    conda config --add channels defaults
+
+# Use bash shell so we can source activate
+SHELL ["/bin/bash", "-c"]
+
+
+RUN conda install torchvision=0.19.0 \
+                  onnx==1.17.0 \
+                  pytorch=2.4.0 \
+                  pytorch-cuda=12.4 \
+              --force-reinstall \
+              -c pytorch \
+              -c nvidia
+
+# Create a g++ wrapper for JIT, since the include dirs are passed with -i rather than -I for some reason
+RUN printf '#!/usr/bin/env bash\nexec /usr/bin/g++ -I/usr/local/cuda/include -I/usr/local/cuda/include/crt "$@"\n' > /usr/local/bin/gxx-wrapper && \
+    chmod +x /usr/local/bin/gxx-wrapper
+ENV CXX=/usr/local/bin/gxx-wrapper
+
+
+# Run setup.sh - this won't install all the things due to missing GPU in the builder
+RUN conda run -n base ./setup.sh --basic --xformers --flash-attn --diffoctreerast --vox2seq --spconv --mipgaussian --kaolin --nvdiffrast --demo
+
+# Now install additional Python packages
+# These ones work inside the builder
+RUN conda run -n base pip install diso
+RUN conda run -n base pip install plyfile utils3d flash_attn spconv-cu120 xformers onnxscript
+RUN conda run -n base pip install kaolin -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.4.0_cu121.html
+RUN conda run -n base pip install git+https://github.com/NVlabs/nvdiffrast.git
+
+# Remove downloaded packages from conda and pip
+RUN conda clean --all -f -y
+RUN pip cache purge
+
+# Deduplicate with rdfind
+# This reduces the size of the image by a few hundred megs.
+RUN rdfind -makesymlinks true /opt/conda
+
+# Final stage
+FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel AS final
+
+WORKDIR /app
+COPY --from=builder /usr/local/bin/gxx-wrapper /usr/local/bin/gxx-wrapper
+COPY --from=builder /opt/conda /opt/conda
+COPY --from=builder /root /root
+COPY --from=builder /app /app
+
+# Reinstall any runtime tools needed
+# git and build-essential are needed for post_install.sh script.
+# vim and strace are useful for debugging, remove those if you want to.
+RUN apt-get update && \
+    apt-get install -y build-essential \
+                       git \
+                       strace \
+                       vim && \
+    rm -rf /var/lib/apt/lists/*
+
+# install these last, so we can experiment without excessive build times.
+COPY trellis         /app/trellis
+COPY app.py          /app/app.py
+COPY example.py      /app/example.py
+COPY extensions      /app/extensions
+COPY assets          /app/assets
+COPY onstart.sh      /app/onstart.sh
+COPY post_install.sh /app/post_install.sh
+
+ENV PATH=/opt/conda/bin:$PATH
+
+# This script runs the post_install steps
+
+# If you're pushing to a container registry, let this run once, run some
+# tests, then do `docker commit` to save the models along with the image.
+# This will ensure that it won't fail at runtime due to models being
+# unavailable, or network restrictions.
+CMD ["/app/onstart.sh"]
+
diff --git a/app.py b/app.py
index 45016df7..ad624e99 100644
--- a/app.py
+++ b/app.py
@@ -245,4 +245,4 @@ def deactivate_button() -> gr.Button:
 if __name__ == "__main__":
     pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
     pipeline.cuda()
-    demo.launch()
+    demo.launch(server_name="0.0.0.0")
diff --git a/onstart.sh b/onstart.sh
new file mode 100755
index 00000000..9226ee59
--- /dev/null
+++ b/onstart.sh
@@ -0,0 +1,12 @@
+#!/bin/bash
+
+cd /app
+
+echo "Doing post install steps"
+./post_install.sh
+
+export CXX=/usr/local/bin/gxx-wrapper
+
+echo "Launching app"
+python3 app.py
+
diff --git a/post_install.sh b/post_install.sh
new file mode 100755
index 00000000..43a6362f
--- /dev/null
+++ b/post_install.sh
@@ -0,0 +1,25 @@
+#!/bin/bash
+set -e
+
+# Check if post-install steps have already been run
+if [ -f /app/.post_install_done ]; then
+    echo "Post-install steps already completed."
+    exit 0
+fi
+
+cd /app
+
+echo "Install stuff that couldn't be installed without GPU"
+# Run the demo setup
+conda run -n base ./setup.sh --mipgaussian --diffoctreerast
+
+echo "Proving it actually works..."
+
+export CXX=/usr/local/bin/gxx-wrapper
+python example.py
+
+# Mark completion
+touch /app/.post_install_done
+
+echo "Post-install steps completed successfully."
+

From e94e6fcc88b7a37c8142fec28d81ec2aeaeb8284 Mon Sep 17 00:00:00 2001
From: Gareth Davidson <gaz@bitplane.net>
Date: Wed, 11 Dec 2024 08:02:22 +0000
Subject: [PATCH 2/3] do apt upgrade in the builder step

---
 Dockerfile | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 13a3c91b..f745b677 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -60,8 +60,9 @@ COPY --from=builder /app /app
 # Reinstall any runtime tools needed
 # git and build-essential are needed for post_install.sh script.
 # vim and strace are useful for debugging, remove those if you want to.
-RUN apt-get update && \
-    apt-get install -y build-essential \
+RUN apt update && \
+    apt upgrade && \
+    apt install -y build-essential \
                        git \
                        strace \
                        vim && \
@@ -84,5 +85,5 @@ ENV PATH=/opt/conda/bin:$PATH
 # tests, then do `docker commit` to save the models along with the image.
 # This will ensure that it won't fail at runtime due to models being
 # unavailable, or network restrictions.
-CMD ["/app/onstart.sh"]
+CMD ["bash", "/app/onstart.sh"]
 

From f9a4832fcded13424f69fdd742b75afd9de4f178 Mon Sep 17 00:00:00 2001
From: Gareth Davidson <gaz@bitplane.net>
Date: Thu, 12 Dec 2024 06:47:07 +0000
Subject: [PATCH 3/3] oops! missed -y

---
 Dockerfile | 2 +-
 onstart.sh | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index f745b677..66c3c84d 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -61,7 +61,7 @@ COPY --from=builder /app /app
 # git and build-essential are needed for post_install.sh script.
 # vim and strace are useful for debugging, remove those if you want to.
 RUN apt update && \
-    apt upgrade && \
+    apt upgrade -y && \
     apt install -y build-essential \
                        git \
                        strace \
diff --git a/onstart.sh b/onstart.sh
index 9226ee59..1abeef9d 100755
--- a/onstart.sh
+++ b/onstart.sh
@@ -10,3 +10,4 @@ export CXX=/usr/local/bin/gxx-wrapper
 echo "Launching app"
 python3 app.py
 
+echo "Something went wrong and it exited?"