Uploading fastAPI

Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.9
+
+# Set the working directory
+WORKDIR /app
+
+# Install dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application files
+COPY . /app/
+
+# Expose the port FastAPI will run on
+EXPOSE 8000
+
+# Run FastAPI when the container starts
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

main.py

@@ -0,0 +1,78 @@
+from fastapi import FastAPI, File, UploadFile, HTTPException, Header
+from fastapi.responses import FileResponse
+import torch
+import os
+import numpy as np
+import cv2
+from PIL import Image
+from typing import List
+from trellis.pipelines import TrellisImageTo3DPipeline
+from trellis.utils import render_utils, postprocessing_utils
+from trellis.representations import Gaussian, MeshExtractResult
+import imageio
+
+# Define working directories
+TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp")
+os.makedirs(TMP_DIR, exist_ok=True)
+
+# ✅ Load pipeline once to avoid reloading for every request
+pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+pipeline.cuda()
+
+# ✅ Preload model (to prevent cold starts)
+try:
+    pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))
+except:
+    pass
+
+# ✅ API Key for private access
+HF_API_KEY = os.getenv("HF_API_KEY", "your-secure-api-key")
+
+app = FastAPI()
+
+def preprocess_image(image: Image.Image) -> Image.Image:
+    """Preprocess a single input image using the Trellis pipeline."""
+    return pipeline.preprocess_image(image)
+
+@app.post("/generate_3d/")
+async def generate_3d(
+    image: UploadFile = File(...),
+    authorization: str = Header(None)
+):
+    """Accepts an image upload, runs inference, and returns a GLB file."""
+    
+    # 🔒 API Key authentication
+    if authorization != f"Bearer {HF_API_KEY}":
+        raise HTTPException(status_code=403, detail="Invalid API key")
+
+    if not image.filename.lower().endswith(("png", "jpg", "jpeg")):
+        raise HTTPException(status_code=400, detail="Invalid image format. Upload a PNG or JPG.")
+
+    # Save the uploaded image
+    image_path = os.path.join(TMP_DIR, image.filename)
+    with open(image_path, "wb") as f:
+        f.write(image.file.read())
+
+    # Load and preprocess the image
+    img = Image.open(image_path).convert("RGBA")
+    processed_image = preprocess_image(img)
+
+    # Run the Trellis pipeline
+    outputs = pipeline.run(
+        processed_image,
+        seed=np.random.randint(0, np.iinfo(np.int32).max),
+        formats=["gaussian", "mesh"],
+        preprocess_image=False,
+        sparse_structure_sampler_params={"steps": 12, "cfg_strength": 7.5},
+        slat_sampler_params={"steps": 12, "cfg_strength": 3.0},
+    )
+
+    # Extract the GLB file
+    gs, mesh = outputs["gaussian"][0], outputs["mesh"][0]
+    glb = postprocessing_utils.to_glb(gs, mesh, simplify=0.95, texture_size=1024, verbose=False)
+    glb_path = os.path.join(TMP_DIR, "sample.glb")
+    glb.export(glb_path)
+
+    torch.cuda.empty_cache()
+
+    return FileResponse(glb_path, media_type="model/gltf-binary", filename="sample.glb")

requirements.txt

@@ -1,5 +1,7 @@
 --extra-index-url https://download.pytorch.org/whl/cu121
 
+fastapi
+uvicorn
 torch==2.4.0
 torchvision==0.19.0
 pillow==10.4.0
@@ -23,4 +25,4 @@ transformers==4.46.3
 gradio_litmodel3d==0.0.1
 https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.0.post2/flash_attn-2.7.0.post2+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
 https://huggingface.co/spaces/JeffreyXiang/TRELLIS/resolve/main/wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl?download=true
-https://huggingface.co/spaces/JeffreyXiang/TRELLIS/resolve/main/wheels/nvdiffrast-0.3.3-cp310-cp310-linux_x86_64.whl?download=true
+https://huggingface.co/spaces/JeffreyXiang/TRELLIS/resolve/main/wheels/nvdiffrast-0.3.3-cp310-cp310-linux_x86_64.whl?download=true