SideQuestVR · emulamer · Dec 21, 2024 · Dec 21, 2024 · Dec 21, 2024 · Dec 21, 2024
diff --git a/Dockerfile b/Dockerfile
@@ -6,33 +6,22 @@ RUN apt-get update && \
 
 WORKDIR /app
 
-# Add the application files
-COPY . /app/
-
-# Initialize and update git submodules
-RUN cd /app && \
-    git init && \
-    git submodule init && \
-    git submodule update --init --recursive && \
-    git submodule update --recursive && \
-    rm -rf .git */.git **/.git  # Remove all .git directories
-
-# Setup conda and PyTorch
+# First, copy only the files needed for dependency installation
+COPY setup.sh ./
+
+# Setup conda and install dependencies
 RUN conda config --set always_yes true && conda init
 RUN conda install cuda=12.4 pytorch==2.4.0 torchvision==0.19.0 pytorch-cuda=12.4 -c pytorch -c nvidia
 
-# Install Kaolin dependencies first
+# Install Kaolin and other dependencies
 RUN conda run -n base pip install -r https://raw.githubusercontent.com/NVIDIAGameWorks/kaolin/v0.17.0/tools/build_requirements.txt \
     -r https://raw.githubusercontent.com/NVIDIAGameWorks/kaolin/v0.17.0/tools/viz_requirements.txt \
     -r https://raw.githubusercontent.com/NVIDIAGameWorks/kaolin/v0.17.0/tools/requirements.txt
 
-# Now install Kaolin with the correct version
 RUN conda run -n base pip install kaolin==0.17.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.4.0_cu124.html
 
-# Install diso and other dependencies
 RUN conda run -n base pip install diso
 
-# Verify Kaolin installation
 RUN conda run -n base python -c "import kaolin; print(kaolin.__version__)"
 
 # Create a g++ wrapper for JIT, since the include dirs are passed with -i rather than -I for some reason
@@ -60,6 +49,11 @@ RUN conda clean --all -f -y
 # This reduces the size of the image by a few hundred megs. Not great, but it's a start.
 RUN rdfind -makesymlinks true /opt/conda
 
+# Only after all dependencies are installed, copy the application code
+COPY ./trellis ./trellis
+COPY ./assets ./assets
+COPY ./extensions ./extensions
+
 # Final stage
 FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel AS final
 
@@ -86,6 +80,9 @@ RUN conda run -n base pip install fastapi uvicorn python-multipart
 COPY startup.sh /app/startup.sh
 RUN chmod +x /app/startup.sh
 
+COPY .gitignore CODE_OF_CONDUCT.md LICENSE README.md setup.sh startup.sh headless_app.py model_generator.py example.py ./
+RUN chmod +x /app/*.sh
+
 ENV PATH=/opt/conda/bin:$PATH
 
 # This script runs the post_install steps

diff --git a/Dockerfile.runpod b/Dockerfile.runpod
@@ -1,49 +1,95 @@
-FROM runpod/base:0.6.2-cuda12.4.1
+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 htop git python3-onnx rdfind
 
 WORKDIR /app
 
-# Install system dependencies
-RUN apt-get update && \
-    apt-get install -y ffmpeg build-essential git python3-onnx rdfind && \
-    rm -rf /var/lib/apt/lists/*
+# First, copy only the files needed for dependency installation
+COPY setup.sh ./
+
+# Setup conda and PyTorch
+RUN conda config --set always_yes true && conda init
+RUN conda install cuda=12.4 pytorch==2.4.0 torchvision==0.19.0 pytorch-cuda=12.4 -c pytorch -c nvidia
+
+# Install Kaolin dependencies first
+RUN conda run -n base pip install -r https://raw.githubusercontent.com/NVIDIAGameWorks/kaolin/v0.17.0/tools/build_requirements.txt \
+    -r https://raw.githubusercontent.com/NVIDIAGameWorks/kaolin/v0.17.0/tools/viz_requirements.txt \
+    -r https://raw.githubusercontent.com/NVIDIAGameWorks/kaolin/v0.17.0/tools/requirements.txt
 
-# Copy the application files
-COPY . .
+# Now install Kaolin with the correct version
+RUN conda run -n base pip install kaolin==0.17.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.4.0_cu124.html
 
-# Initialize and update git submodules
-RUN cd /app && \
-    git init && \
-    git submodule init && \
-    git submodule update --init --recursive && \
-    git submodule update --recursive && \
-    rm -rf .git */.git **/.git
+# Install diso and other dependencies
+RUN conda run -n base pip install diso
 
-# Create a g++ wrapper for JIT compilation
+# Verify Kaolin installation
+RUN conda run -n base python -c "import kaolin; print(kaolin.__version__)"
+
+# 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
 
-# Install Python dependencies
-RUN python3.11 -m pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu118
-RUN python3.11 -m pip install --no-cache-dir kaolin==0.17.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.0.0_cu118.html
-RUN python3.11 -m pip install --no-cache-dir diso plyfile utils3d flash_attn xformers
-RUN python3.11 -m pip install --no-cache-dir git+https://github.com/NVlabs/nvdiffrast.git
-
-# Run setup script with necessary components
-RUN ./setup.sh --basic --xformers --flash-attn --diffoctreerast --vox2seq --spconv --mipgaussian --kaolin --nvdiffrast
+# Run setup.sh - this won't install all the things, we'll need to install some later
+RUN conda run -n base ./setup.sh --basic --xformers --flash-attn --diffoctreerast --vox2seq --spconv --mipgaussian --kaolin --nvdiffrast --demo
 
-# Install RunPod
-RUN python3.11 -m pip install --no-cache-dir runpod
+# Now install additional Python packages
+# These ones work inside the builder
+RUN conda run -n base pip install plyfile utils3d flash_attn spconv-cu120 xformers
+RUN conda run -n base pip install git+https://github.com/NVlabs/nvdiffrast.git
 
-# Cleanup unnecessary files
-RUN apt-get remove -y ffmpeg git python3-onnx && \
+# Cleanup after builds are done
+RUN apt-get remove -y ffmpeg build-essential htop git python3-onnx && \
     apt-get autoremove -y && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# Copy and setup startup script
-COPY startup.runpod.sh /app/startup.sh
-RUN chmod +x /app/startup.sh
+RUN conda clean --all -f -y
+
+# Deduplicate with rdfind
+# This reduces the size of the image by a few hundred megs. Not great, but it's a start.
+RUN rdfind -makesymlinks true /opt/conda
+
+# Only after all dependencies are installed, copy the application code
+COPY ./trellis ./trellis
+COPY ./assets ./assets
+COPY ./extensions ./extensions
+
+# 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 the image size.
+RUN apt-get update && \
+    apt-get install -y build-essential \
+                       git \
+                       strace \
+                       vim && \
+    rm -rf /var/lib/apt/lists/*
+
+# Add FastAPI dependencies
+RUN conda run -n base pip install fastapi uvicorn python-multipart
+
+# Add RunPod dependencies
+RUN conda run -n base pip install runpod
+
+# Add the new startup script
+COPY startup.runpod.sh /app/startup.runpod.sh
+RUN chmod +x /app/startup.runpod.sh
+
+COPY .gitignore CODE_OF_CONDUCT.md LICENSE README.md setup.sh startup.runpod.sh rp_handler.py model_generator.py example.py ./
+RUN chmod +x /app/*.sh
+
+ENV PATH=/opt/conda/bin:$PATH
 
-# Use startup script as entrypoint
-ENTRYPOINT ["/app/startup.runpod.sh"] 
+# This script runs the post_install steps
+CMD ["/app/startup.runpod.sh"]
diff --git a/README.md b/README.md
@@ -79,14 +79,14 @@ We provide the following pretrained models:
 
 | Model | Description | #Params | Download |
 | --- | --- | --- | --- |
-| TRELLIS-image-large | Large image-to-3D model | 1.2B | [Download](https://huggingface.co/JeffreyXiang/TRELLIS-image-large) |
+| TRELLIS-image-large | Large image-to-3D model | 1.2B | [Download](https://huggingface.co/shakamone/trellis-large) |
 | TRELLIS-text-base | Base text-to-3D model | 342M | Coming Soon |
 | TRELLIS-text-large | Large text-to-3D model | 1.1B | Coming Soon |
 | TRELLIS-text-xlarge | Extra-large text-to-3D model | 2.0B | Coming Soon |
 
 The models are hosted on Hugging Face. You can directly load the models with their repository names in the code:
 ```python
-TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+TrellisImageTo3DPipeline.from_pretrained("shakamone/trellis-large")
 ```
 
 If you prefer loading the model from local, you can download the model files from the links above and load the model with the folder path (folder structure should be maintained):
@@ -114,7 +114,7 @@ from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.utils import render_utils, postprocessing_utils
 
 # Load a pipeline from a model folder or a Hugging Face model hub.
-pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+pipeline = TrellisImageTo3DPipeline.from_pretrained("shakamone/trellis-large")
 pipeline.cuda()
 
 # Load an image

diff --git a/app.py b/app.py
@@ -243,6 +243,6 @@ def deactivate_button() -> gr.Button:
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+    pipeline = TrellisImageTo3DPipeline.from_pretrained("shakamone/trellis-large")
     pipeline.cuda()
     demo.launch()
diff --git a/example.py b/example.py
@@ -10,7 +10,7 @@
 from trellis.utils import render_utils, postprocessing_utils
 
 # Load a pipeline from a model folder or a Hugging Face model hub.
-pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+pipeline = TrellisImageTo3DPipeline.from_pretrained("shakamone/trellis-large")
 pipeline.cuda()
 
 # Load an image