Create Docker headless version

Dockerfile

@@ -0,0 +1,92 @@
+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
+
+# 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
+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
+
+# 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
+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, 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
+
+# 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 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/*
+
+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
+
+# 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 the new startup script
+COPY startup.sh /app/startup.sh
+RUN chmod +x /app/startup.sh
+
+ENV PATH=/opt/conda/bin:$PATH
+
+# This script runs the post_install steps
+CMD ["/app/startup.sh"]

Dockerfile.runpod

@@ -0,0 +1,49 @@
+FROM runpod/base:0.6.2-cuda12.4.1
+
+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/*
+
+# Copy the application files
+COPY . .
+
+# 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
+
+# Create a g++ wrapper for JIT compilation
+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
+
+# Install RunPod
+RUN python3.11 -m pip install --no-cache-dir runpod
+
+# Cleanup unnecessary files
+RUN apt-get remove -y ffmpeg 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
+
+# Use startup script as entrypoint
+ENTRYPOINT ["/app/startup.runpod.sh"] 

headless_app.py

@@ -0,0 +1,182 @@
+import os
+from typing import *
+import torch
+import numpy as np
+import imageio
+import uuid
+import time
+from easydict import EasyDict as edict
+from PIL import Image
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import FileResponse
+from trellis.pipelines import TrellisImageTo3DPipeline
+from trellis.representations import Gaussian, MeshExtractResult
+from trellis.utils import render_utils, postprocessing_utils
+import json
+from fastapi.middleware.cors import CORSMiddleware
+
+app = FastAPI()
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],  # Allows all headers
+)
+
+MAX_SEED = np.iinfo(np.int32).max
+TMP_DIR = "/workspace/Trellis-demo"
+os.makedirs(TMP_DIR, exist_ok=True)
+
+def cleanup_old_files(directory: str, max_age_hours: int = 24):
+    """Clean up files older than max_age_hours"""
+    current_time = time.time()
+    for filename in os.listdir(directory):
+        filepath = os.path.join(directory, filename)
+        if os.path.isfile(filepath):
+            if (current_time - os.path.getmtime(filepath)) > (max_age_hours * 3600):
+                try:
+                    os.remove(filepath)
+                except OSError:
+                    pass
+
+@app.on_event("startup")
+async def startup_event():
+    """Run cleanup on startup"""
+    cleanup_old_files(TMP_DIR)
+
+# Initialize pipeline globally
+pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+pipeline.cuda()
+
+def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
+    trial_id = str(uuid.uuid4())
+    processed_image = pipeline.preprocess_image(image)
+    processed_image.save(f"{TMP_DIR}/{trial_id}.png")
+    return trial_id, processed_image
+
+def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
+    return {
+        'gaussian': {
+            **gs.init_params,
+            '_xyz': gs._xyz.cpu().numpy().tolist(),
+            '_features_dc': gs._features_dc.cpu().numpy().tolist(),
+            '_scaling': gs._scaling.cpu().numpy().tolist(),
+            '_rotation': gs._rotation.cpu().numpy().tolist(),
+            '_opacity': gs._opacity.cpu().numpy().tolist(),
+        },
+        'mesh': {
+            'vertices': mesh.vertices.cpu().numpy().tolist(),
+            'faces': mesh.faces.cpu().numpy().tolist(),
+        },
+        'trial_id': trial_id,
+    }
+
+@app.post("/process-image")
+async def process_image(
+    file: UploadFile = File(...),
+    seed: int = 0,
+    randomize_seed: bool = True,
+    ss_guidance_strength: float = 7.5,
+    ss_sampling_steps: int = 12,
+    slat_guidance_strength: float = 3.0,
+    slat_sampling_steps: int = 12
+):
+    # Read and process the uploaded image
+    image = Image.open(file.file)
+    trial_id, processed_image = preprocess_image(image)
+
+    # Generate 3D model
+    if randomize_seed:
+        seed = np.random.randint(0, MAX_SEED)
+
+    outputs = pipeline.run(
+        processed_image,
+        seed=seed,
+        formats=["gaussian", "mesh"],
+        preprocess_image=False,
+        sparse_structure_sampler_params={
+            "steps": ss_sampling_steps,
+            "cfg_strength": ss_guidance_strength,
+        },
+        slat_sampler_params={
+            "steps": slat_sampling_steps,
+            "cfg_strength": slat_guidance_strength,
+        },
+    )
+
+    # Generate preview video
+    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
+    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+    video_path = f"{TMP_DIR}/{trial_id}_preview.mp4"
+    imageio.mimsave(video_path, video, fps=15)
+
+    # Pack state and return results
+    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], trial_id)
+
+    # Save state file
+    with open(f"{TMP_DIR}/{trial_id}_state.json", 'w') as f:
+        json.dump(state, f)
+
+    return {
+        "trial_id": trial_id,
+        "state": state,
+        "preview_video": f"/preview/{trial_id}"
+    }
+
+@app.get("/preview/{trial_id}")
+async def get_preview(trial_id: str):
+    video_path = f"{TMP_DIR}/{trial_id}_preview.mp4"
+    return FileResponse(video_path)
+
+@app.post("/extract-glb/{trial_id}")
+async def extract_glb(
+    trial_id: str,
+    mesh_simplify: float = 0.95,
+    texture_size: int = 1024
+):
+    # Load the state file
+    state_path = f"{TMP_DIR}/{trial_id}_state.json"
+    if not os.path.exists(state_path):
+        return {"error": "Trial ID not found"}
+
+    # Add this line to load the state
+    with open(state_path, 'r') as f:
+        state = json.load(f)
+
+    # Generate GLB
+    glb_path = f"{TMP_DIR}/{trial_id}.glb"
+    gs = Gaussian(
+        aabb=state['gaussian']['aabb'],
+        sh_degree=state['gaussian']['sh_degree'],
+        mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
+        scaling_bias=state['gaussian']['scaling_bias'],
+        opacity_bias=state['gaussian']['opacity_bias'],
+        scaling_activation=state['gaussian']['scaling_activation'],
+    )
+    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
+    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
+    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
+    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
+    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
+
+    mesh = edict(
+        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
+        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
+    )
+
+    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
+    glb.export(glb_path)
+
+    return FileResponse(glb_path, filename=f"{trial_id}.glb")
+
+@app.get("/health")
+async def health_check():
+    return {"status": "ok"}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)