From 8753c6cc3410e323d358c42c9c2b1cacbb1f23e0 Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sat, 21 Dec 2024 18:43:05 +0000 Subject: [PATCH 01/14] Add startup runpod file, update dockerfile for runpod --- Dockerfile.runpod | 69 +++++++++++++++++++++++++++++++++++++++++++++++ startup.runpod.sh | 14 ++++++++++ 2 files changed, 83 insertions(+) create mode 100644 Dockerfile.runpod create mode 100644 startup.runpod.sh diff --git a/Dockerfile.runpod b/Dockerfile.runpod new file mode 100644 index 00000000..9a169d4e --- /dev/null +++ b/Dockerfile.runpod @@ -0,0 +1,69 @@ +FROM runpod/base:0.6.2-cuda12.4.1 + +WORKDIR /app + +# Install system dependencies +RUN apt-get update && apt-get install -y \ + build-essential \ + git \ + ffmpeg \ + python3-onnx \ + rdfind \ + && rm -rf /var/lib/apt/lists/* + +# Copy 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 + +# Create 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 + +# Set environment variables for model caching +ENV HF_HOME=/workspace/models +ENV HF_CACHE_DIR=/workspace/models + +# Create workspace directories +RUN mkdir -p /workspace/models + +# Install basic dependencies and TRELLIS requirements +RUN python3.11 -m pip install --no-cache-dir \ + pillow \ + imageio \ + imageio-ffmpeg \ + tqdm \ + easydict \ + opencv-python-headless \ + scipy \ + ninja \ + rembg \ + onnxruntime \ + trimesh \ + xatlas \ + pyvista \ + pymeshfix \ + igraph \ + transformers \ + kaolin \ + runpod \ + fastapi \ + uvicorn \ + python-multipart + +# Run setup script for GPU-dependent packages +RUN ./setup.sh --mipgaussian --diffoctreerast --spconv --kaolin --nvdiffrast || exit 1 + +# Copy and set up the startup script +COPY startup.runpod.sh /app/startup.runpod.sh +RUN chmod +x /app/startup.runpod.sh + +# Set the default command +CMD ["/app/startup.runpod.sh"] \ No newline at end of file diff --git a/startup.runpod.sh b/startup.runpod.sh new file mode 100644 index 00000000..c20f7367 --- /dev/null +++ b/startup.runpod.sh @@ -0,0 +1,14 @@ +#!/bin/bash +set -e + +cd /app + +# Verify installation +export CXX=/usr/local/bin/gxx-wrapper +python3.11 example.py + +# 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 From 985540f72bacbdac141c623f78f98ff2680ebd8f Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sat, 21 Dec 2024 21:05:33 +0000 Subject: [PATCH 02/14] Add missing rp_handler --- Dockerfile | 8 ------ Dockerfile.runpod | 8 ------ rp_handler.py | 71 +++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 71 insertions(+), 16 deletions(-) create mode 100644 rp_handler.py diff --git a/Dockerfile b/Dockerfile index a4cce199..8225f78b 100644 --- a/Dockerfile +++ b/Dockerfile @@ -9,14 +9,6 @@ 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 diff --git a/Dockerfile.runpod b/Dockerfile.runpod index 9a169d4e..ee2660c4 100644 --- a/Dockerfile.runpod +++ b/Dockerfile.runpod @@ -14,14 +14,6 @@ RUN apt-get update && apt-get install -y \ # Copy 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 - # Create 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 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 From c5e84e8ff8391fcb74943bfafd758ec62d468c60 Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sat, 21 Dec 2024 21:11:55 +0000 Subject: [PATCH 03/14] update rp-handler --- rp_handler.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/rp_handler.py b/rp_handler.py index b4579bd4..42131182 100644 --- a/rp_handler.py +++ b/rp_handler.py @@ -20,15 +20,19 @@ def init_pipeline(): def handler(event): try: input_data = event["input"] - image_path = input_data.get("image_path") + image_base64 = input_data.get("image_base64") 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"} + if not image_base64: + return {"error": "Image base64 data not provided"} - # Load and process image - image = Image.open(image_path) + try: + # Decode base64 to image + image_data = base64.b64decode(image_base64) + image = Image.open(io.BytesIO(image_data)) + except Exception as e: + return {"error": f"Failed to decode base64 image: {str(e)}"} # Generate 3D model outputs = trellis_pipe.run( From abb0d7eab64d7d82476361124a253f719939f1b4 Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sat, 21 Dec 2024 21:29:00 +0000 Subject: [PATCH 04/14] try to use older version non runpod for docker runpod --- Dockerfile.runpod | 119 +++++++++++++++++++++++++++------------------- startup.runpod.sh | 20 ++++++-- 2 files changed, 87 insertions(+), 52 deletions(-) diff --git a/Dockerfile.runpod b/Dockerfile.runpod index ee2660c4..64002f63 100644 --- a/Dockerfile.runpod +++ b/Dockerfile.runpod @@ -1,61 +1,84 @@ -FROM runpod/base:0.6.2-cuda12.4.1 +FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel AS builder -WORKDIR /app +# Install build dependencies +RUN apt-get update && \ + apt-get install -y ffmpeg build-essential htop git python3-onnx rdfind -# Install system dependencies -RUN apt-get update && apt-get install -y \ - build-essential \ - git \ - ffmpeg \ - python3-onnx \ - rdfind \ - && rm -rf /var/lib/apt/lists/* +WORKDIR /app -# Copy the application files +# Add the application files COPY . /app/ -# Create g++ wrapper for JIT compilation +# 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 -# Set environment variables for model caching -ENV HF_HOME=/workspace/models -ENV HF_CACHE_DIR=/workspace/models - -# Create workspace directories -RUN mkdir -p /workspace/models - -# Install basic dependencies and TRELLIS requirements -RUN python3.11 -m pip install --no-cache-dir \ - pillow \ - imageio \ - imageio-ffmpeg \ - tqdm \ - easydict \ - opencv-python-headless \ - scipy \ - ninja \ - rembg \ - onnxruntime \ - trimesh \ - xatlas \ - pyvista \ - pymeshfix \ - igraph \ - transformers \ - kaolin \ - runpod \ - fastapi \ - uvicorn \ - python-multipart - -# Run setup script for GPU-dependent packages -RUN ./setup.sh --mipgaussian --diffoctreerast --spconv --kaolin --nvdiffrast || exit 1 - -# Copy and set up the startup script +# 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.runpod.sh /app/startup.runpod.sh RUN chmod +x /app/startup.runpod.sh -# Set the default command +ENV PATH=/opt/conda/bin:$PATH + +# This script runs the post_install steps CMD ["/app/startup.runpod.sh"] \ No newline at end of file diff --git a/startup.runpod.sh b/startup.runpod.sh index c20f7367..d671d9b0 100644 --- a/startup.runpod.sh +++ b/startup.runpod.sh @@ -3,12 +3,24 @@ set -e cd /app -# Verify installation -export CXX=/usr/local/bin/gxx-wrapper -python3.11 example.py +# 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 RunPod handler..." -python3.11 -u rp_handler.py \ No newline at end of file +python -u rp_handler.py \ No newline at end of file From 2c0af6f6a56eaa61bdc6d0847ee92c730e8ec5ec Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sat, 21 Dec 2024 22:32:20 +0000 Subject: [PATCH 05/14] Add more params to rp handler --- rp_handler.py | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/rp_handler.py b/rp_handler.py index 42131182..26d5c5f7 100644 --- a/rp_handler.py +++ b/rp_handler.py @@ -7,6 +7,7 @@ from trellis.utils import render_utils, postprocessing_utils import base64 import io +import numpy as np # Initialize pipeline def init_pipeline(): @@ -24,6 +25,14 @@ def handler(event): mesh_simplify = input_data.get("mesh_simplify", 0.95) texture_size = input_data.get("texture_size", 1024) + # Add missing parameters from headless version + seed = input_data.get("seed", 0) + randomize_seed = input_data.get("randomize_seed", True) + ss_guidance_strength = input_data.get("ss_guidance_strength", 7.5) + ss_sampling_steps = input_data.get("ss_sampling_steps", 12) + slat_guidance_strength = input_data.get("slat_guidance_strength", 3.0) + slat_sampling_steps = input_data.get("slat_sampling_steps", 12) + if not image_base64: return {"error": "Image base64 data not provided"} @@ -34,18 +43,22 @@ def handler(event): except Exception as e: return {"error": f"Failed to decode base64 image: {str(e)}"} + # Update seed if randomization is requested + if randomize_seed: + seed = np.random.randint(0, np.iinfo(np.int32).max) + # Generate 3D model outputs = trellis_pipe.run( image, - seed=42, + seed=seed, formats=["gaussian", "mesh"], sparse_structure_sampler_params={ - "steps": 12, - "cfg_strength": 7.5, + "steps": ss_sampling_steps, + "cfg_strength": ss_guidance_strength, }, slat_sampler_params={ - "steps": 12, - "cfg_strength": 3.0, + "steps": slat_sampling_steps, + "cfg_strength": slat_guidance_strength, } ) From 096019f0154df7207046716669a2cb531fa9ab81 Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sat, 21 Dec 2024 22:52:00 +0000 Subject: [PATCH 06/14] Simplify endpoints --- headless_app.py | 205 ++++++++++----------------------------------- model_generator.py | 55 ++++++++++++ rp_handler.py | 77 ++++------------- 3 files changed, 114 insertions(+), 223 deletions(-) create mode 100644 model_generator.py diff --git a/headless_app.py b/headless_app.py index 4e3419a9..7f1308f2 100644 --- a/headless_app.py +++ b/headless_app.py @@ -1,177 +1,58 @@ -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 +import base64 +from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware +from PIL import Image +from .model_generator import ModelGenerator +import io +from pydantic import BaseModel + +# Add new request model +class ImageRequest(BaseModel): + image_base64: str + 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 + mesh_simplify: float = 0.95 + texture_size: int = 1024 app = FastAPI() # Add CORS middleware app.add_middleware( CORSMiddleware, - allow_origins=["*"], # Allows all origins + allow_origins=["*"], allow_credentials=True, - allow_methods=["*"], # Allows all methods - allow_headers=["*"], # Allows all headers + allow_methods=["*"], + allow_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, - } +# Initialize generator globally +generator = ModelGenerator() @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") +async def process_image(request: ImageRequest): + try: + # Decode base64 image + image_data = base64.b64decode(request.image_base64) + image = Image.open(io.BytesIO(image_data)) + + # Generate 3D model + return generator.generate( + image=image, + seed=request.seed, + randomize_seed=request.randomize_seed, + ss_guidance_strength=request.ss_guidance_strength, + ss_sampling_steps=request.ss_sampling_steps, + slat_guidance_strength=request.slat_guidance_strength, + slat_sampling_steps=request.slat_sampling_steps, + mesh_simplify=request.mesh_simplify, + texture_size=request.texture_size + ) + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) @app.get("/health") async def health_check(): diff --git a/model_generator.py b/model_generator.py new file mode 100644 index 00000000..ba9c7f90 --- /dev/null +++ b/model_generator.py @@ -0,0 +1,55 @@ +import torch +import numpy as np +from PIL import Image +from trellis.pipelines import TrellisImageTo3DPipeline +from trellis.utils import postprocessing_utils +import io +import base64 + +class ModelGenerator: + def __init__(self): + self.pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large") + self.pipeline.cuda() + self.MAX_SEED = np.iinfo(np.int32).max + + def generate(self, image, seed=0, randomize_seed=True, + ss_guidance_strength=7.5, ss_sampling_steps=12, + slat_guidance_strength=3.0, slat_sampling_steps=12, + mesh_simplify=0.95, texture_size=1024): + + # Update seed if randomization is requested + if randomize_seed: + seed = np.random.randint(0, self.MAX_SEED) + + # Generate 3D model + outputs = self.pipeline.run( + image, + seed=seed, + formats=["gaussian", "mesh"], + 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 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 and convert to base64 + buffer = io.BytesIO() + glb.export(buffer) + glb_base64 = base64.b64encode(buffer.getvalue()).decode() + + return { + "glb_base64": glb_base64, + "seed": seed + } \ No newline at end of file diff --git a/rp_handler.py b/rp_handler.py index 26d5c5f7..a4300fa9 100644 --- a/rp_handler.py +++ b/rp_handler.py @@ -1,38 +1,18 @@ 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 -import numpy as np +from model_generator import ModelGenerator -# Initialize pipeline -def init_pipeline(): - pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large") - pipeline.cuda() - return pipeline - -# Global pipeline instance -trellis_pipe = init_pipeline() +# Initialize generator globally +generator = ModelGenerator() def handler(event): try: input_data = event["input"] - image_base64 = input_data.get("image_base64") - mesh_simplify = input_data.get("mesh_simplify", 0.95) - texture_size = input_data.get("texture_size", 1024) - # Add missing parameters from headless version - seed = input_data.get("seed", 0) - randomize_seed = input_data.get("randomize_seed", True) - ss_guidance_strength = input_data.get("ss_guidance_strength", 7.5) - ss_sampling_steps = input_data.get("ss_sampling_steps", 12) - slat_guidance_strength = input_data.get("slat_guidance_strength", 3.0) - slat_sampling_steps = input_data.get("slat_sampling_steps", 12) - + # Extract parameters from input + image_base64 = input_data.get("image_base64") if not image_base64: return {"error": "Image base64 data not provided"} @@ -43,44 +23,19 @@ def handler(event): except Exception as e: return {"error": f"Failed to decode base64 image: {str(e)}"} - # Update seed if randomization is requested - if randomize_seed: - seed = np.random.randint(0, np.iinfo(np.int32).max) - - # Generate 3D model - outputs = trellis_pipe.run( - image, - seed=seed, - formats=["gaussian", "mesh"], - 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 3D model using ModelGenerator + return generator.generate( + image=image, + seed=input_data.get("seed", 0), + randomize_seed=input_data.get("randomize_seed", True), + ss_guidance_strength=input_data.get("ss_guidance_strength", 7.5), + ss_sampling_steps=input_data.get("ss_sampling_steps", 12), + slat_guidance_strength=input_data.get("slat_guidance_strength", 3.0), + slat_sampling_steps=input_data.get("slat_sampling_steps", 12), + mesh_simplify=input_data.get("mesh_simplify", 0.95), + texture_size=input_data.get("texture_size", 1024) ) - # 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)} From 7327c85ef3f42721abb2affbed5c4d23b71c7a2c Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sun, 22 Dec 2024 00:16:29 +0000 Subject: [PATCH 07/14] update dockerfile --- Dockerfile | 14 +++++++------- headless_app.py | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/Dockerfile b/Dockerfile index 8225f78b..297b8180 100644 --- a/Dockerfile +++ b/Dockerfile @@ -6,25 +6,22 @@ RUN apt-get update && \ WORKDIR /app -# Add the application files -COPY . /app/ +# First, copy only the files needed for dependency installation +COPY setup.sh ./ -# Setup conda and PyTorch +# 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 @@ -52,6 +49,9 @@ 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 . . + # Final stage FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel AS final diff --git a/headless_app.py b/headless_app.py index 7f1308f2..3028882a 100644 --- a/headless_app.py +++ b/headless_app.py @@ -2,7 +2,7 @@ from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from PIL import Image -from .model_generator import ModelGenerator +from model_generator import ModelGenerator import io from pydantic import BaseModel From b39b9360ec834bfc0cfc1210d74f44692b260b34 Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sun, 22 Dec 2024 02:10:22 +0000 Subject: [PATCH 08/14] fix missing imports, make dockerfile better cached --- Dockerfile | 6 +++++- headless_app.py | 2 +- model_generator.py | 4 +++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/Dockerfile b/Dockerfile index 297b8180..9f34701b 100644 --- a/Dockerfile +++ b/Dockerfile @@ -50,7 +50,9 @@ RUN conda clean --all -f -y RUN rdfind -makesymlinks true /opt/conda # Only after all dependencies are installed, copy the application code -COPY . . +COPY ./trellis ./trellis +COPY ./assets ./assets +COPY ./extensions ./extensions # Final stage FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel AS final @@ -78,6 +80,8 @@ 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 ./ + ENV PATH=/opt/conda/bin:$PATH # This script runs the post_install steps diff --git a/headless_app.py b/headless_app.py index 3028882a..8b3eb67a 100644 --- a/headless_app.py +++ b/headless_app.py @@ -1,5 +1,5 @@ import base64 -from fastapi import FastAPI +from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from PIL import Image from model_generator import ModelGenerator diff --git a/model_generator.py b/model_generator.py index ba9c7f90..63446844 100644 --- a/model_generator.py +++ b/model_generator.py @@ -16,6 +16,8 @@ def generate(self, image, seed=0, randomize_seed=True, ss_guidance_strength=7.5, ss_sampling_steps=12, slat_guidance_strength=3.0, slat_sampling_steps=12, mesh_simplify=0.95, texture_size=1024): + + print(f"Generating 3D model with seed: {seed}") # Update seed if randomization is requested if randomize_seed: @@ -46,7 +48,7 @@ def generate(self, image, seed=0, randomize_seed=True, # Save GLB to bytes buffer and convert to base64 buffer = io.BytesIO() - glb.export(buffer) + glb.export(buffer, format="glb") glb_base64 = base64.b64encode(buffer.getvalue()).decode() return { From 719c476aa8da2aff88d93cab4237d706d7c51ce0 Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sun, 22 Dec 2024 02:21:27 +0000 Subject: [PATCH 09/14] mk copied shell files executable in dockerfile --- Dockerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/Dockerfile b/Dockerfile index 9f34701b..f1dfac7b 100644 --- a/Dockerfile +++ b/Dockerfile @@ -81,6 +81,7 @@ 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 From 8929bd9faff2265061defdc008036ae163598bbe Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sun, 22 Dec 2024 18:35:38 +0000 Subject: [PATCH 10/14] update model_generator --- model_generator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model_generator.py b/model_generator.py index 63446844..5744b7f5 100644 --- a/model_generator.py +++ b/model_generator.py @@ -48,7 +48,7 @@ def generate(self, image, seed=0, randomize_seed=True, # Save GLB to bytes buffer and convert to base64 buffer = io.BytesIO() - glb.export(buffer, format="glb") + glb.export(buffer, file_type="glb") glb_base64 = base64.b64encode(buffer.getvalue()).decode() return { From 29f838e4367dddc555b975712d9e36929114e202 Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sun, 22 Dec 2024 21:27:20 +0000 Subject: [PATCH 11/14] mk runpod dockerfile also use speed improvements --- Dockerfile.runpod | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/Dockerfile.runpod b/Dockerfile.runpod index 1959fde0..e3107574 100644 --- a/Dockerfile.runpod +++ b/Dockerfile.runpod @@ -6,8 +6,8 @@ RUN apt-get update && \ WORKDIR /app -# Add the application files -COPY . /app/ +# 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 @@ -52,6 +52,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 @@ -78,6 +83,9 @@ RUN conda run -n base pip install fastapi uvicorn python-multipart 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 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 From cd8c5c04e0849cdc4aaf2a0fcb17f79a7262d5d9 Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sun, 22 Dec 2024 21:52:05 +0000 Subject: [PATCH 12/14] update runpod dockerfile with missing files --- Dockerfile.runpod | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Dockerfile.runpod b/Dockerfile.runpod index e3107574..1138b6aa 100644 --- a/Dockerfile.runpod +++ b/Dockerfile.runpod @@ -83,7 +83,7 @@ RUN conda run -n base pip install fastapi uvicorn python-multipart 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 headless_app.py model_generator.py example.py ./ +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 From 3a7c9a4f4abaa1634da9ea0c89ffcf961b747b61 Mon Sep 17 00:00:00 2001 From: Alberto Taiuti Date: Sun, 22 Dec 2024 22:02:32 +0000 Subject: [PATCH 13/14] add missing runpod package --- Dockerfile.runpod | 3 +++ 1 file changed, 3 insertions(+) diff --git a/Dockerfile.runpod b/Dockerfile.runpod index 1138b6aa..142b8bcf 100644 --- a/Dockerfile.runpod +++ b/Dockerfile.runpod @@ -79,6 +79,9 @@ RUN apt-get update && \ # 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 From 19887020d00fc729564ea420079128b985256c34 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 21 May 2025 10:36:23 -0400 Subject: [PATCH 14/14] Update model references to new location --- README.md | 6 +++--- app.py | 2 +- example.py | 2 +- model_generator.py | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 1acea907..da676967 100644 --- 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 index 45016df7..ae77f902 100644 --- 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 index 155e6711..0b828f8a 100644 --- 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 diff --git a/model_generator.py b/model_generator.py index 5744b7f5..d56cc48c 100644 --- a/model_generator.py +++ b/model_generator.py @@ -8,7 +8,7 @@ class ModelGenerator: def __init__(self): - self.pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large") + self.pipeline = TrellisImageTo3DPipeline.from_pretrained("shakamone/trellis-large") self.pipeline.cuda() self.MAX_SEED = np.iinfo(np.int32).max