diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 00000000..a4cce199 --- /dev/null +++ b/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"] \ No newline at end of file diff --git a/Dockerfile.runpod b/Dockerfile.runpod new file mode 100644 index 00000000..af26a990 --- /dev/null +++ b/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"] \ No newline at end of file diff --git a/headless_app.py b/headless_app.py new file mode 100644 index 00000000..4e3419a9 --- /dev/null +++ b/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) diff --git a/rp_handler.py b/rp_handler.py new file mode 100644 index 00000000..b4579bd4 --- /dev/null +++ b/rp_handler.py @@ -0,0 +1,71 @@ +import runpod +import torch +from PIL import Image +import os +from datetime import datetime +from trellis.pipelines import TrellisImageTo3DPipeline +from trellis.utils import render_utils, postprocessing_utils +import base64 +import io + +# Initialize pipeline +def init_pipeline(): + pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large") + pipeline.cuda() + return pipeline + +# Global pipeline instance +trellis_pipe = init_pipeline() + +def handler(event): + try: + input_data = event["input"] + image_path = input_data.get("image_path") + mesh_simplify = input_data.get("mesh_simplify", 0.95) + texture_size = input_data.get("texture_size", 1024) + + if not image_path or not os.path.exists(image_path): + return {"error": "Image path not provided or file not found"} + + # Load and process image + image = Image.open(image_path) + + # Generate 3D model + outputs = trellis_pipe.run( + image, + seed=42, + formats=["gaussian", "mesh"], + sparse_structure_sampler_params={ + "steps": 12, + "cfg_strength": 7.5, + }, + slat_sampler_params={ + "steps": 12, + "cfg_strength": 3.0, + } + ) + + # Generate GLB file + glb = postprocessing_utils.to_glb( + outputs['gaussian'][0], + outputs['mesh'][0], + simplify=mesh_simplify, + texture_size=texture_size + ) + + # Save GLB to bytes buffer + buffer = io.BytesIO() + glb.export(buffer) + + # Convert to base64 + glb_base64 = base64.b64encode(buffer.getvalue()).decode() + + return { + "glb_base64": glb_base64 + } + + except Exception as e: + return {"error": str(e)} + +if __name__ == "__main__": + runpod.serverless.start({"handler": handler}) \ No newline at end of file diff --git a/startup.runpod.sh b/startup.runpod.sh new file mode 100644 index 00000000..edd3d263 --- /dev/null +++ b/startup.runpod.sh @@ -0,0 +1,26 @@ +#!/bin/bash +set -e + +cd /app + +# Run post-install steps if not already done +if [ ! -f /app/.post_install_done ]; then + echo "Running post-install steps..." + + # Install GPU-dependent packages + ./setup.sh --mipgaussian --diffoctreerast + + # Verify installation + export CXX=/usr/local/bin/gxx-wrapper + python3.11 example.py + + # Mark completion + touch /app/.post_install_done + echo "Post-install steps completed successfully." +fi + +# Set compiler wrapper for runtime +export CXX=/usr/local/bin/gxx-wrapper + +echo "Launching RunPod handler..." +python3.11 -u rp_handler.py \ No newline at end of file diff --git a/startup.sh b/startup.sh new file mode 100644 index 00000000..a143e668 --- /dev/null +++ b/startup.sh @@ -0,0 +1,33 @@ +#!/bin/bash +set -e + +cd /app + +# Run post-install steps if not already done +if [ ! -f /app/.post_install_done ]; then + echo "Running post-install steps..." + + # Install GPU-dependent packages + conda run -n base ./setup.sh --mipgaussian --diffoctreerast + + # Verify installation + export CXX=/usr/local/bin/gxx-wrapper + python example.py + + # Mark completion + touch /app/.post_install_done + echo "Post-install steps completed successfully." +fi + +# Set compiler wrapper for runtime +export CXX=/usr/local/bin/gxx-wrapper + +echo "Launching headless API server..." +# Check if headless_app.py exists in workspace and use that instead +if [ -f "/workspace/headless_app.py" ]; then + echo "Using headless_app.py from workspace" + python3 /workspace/headless_app.py +else + echo "Using built-in headless_app.py" + python3 headless_app.py +fi