Update main.py

main.py

@@ -1,44 +1,46 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException, Header
-from fastapi.responses import FileResponse
-import torch
 import os
+# Must happen before Trellis / huggingface_hub is imported
+os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/huggingface_cache"
+os.makedirs("/tmp/huggingface_cache", exist_ok=True)
+
+import torch
 import numpy as np
-import cv2
 from PIL import Image
-from typing import List
+from fastapi import FastAPI, File, UploadFile, HTTPException, Header
+from fastapi.responses import FileResponse
+
+# Trellis pipeline imports
 from trellis.pipelines import TrellisImageTo3DPipeline
-from trellis.utils import render_utils, postprocessing_utils
-from trellis.representations import Gaussian, MeshExtractResult
-import imageio
+from trellis.utils import postprocessing_utils
 
-# Define working directories
+# Use /tmp/space_tmp for user data & avoid read-only /app
 TMP_DIR = "/tmp/space_tmp"
 os.makedirs(TMP_DIR, exist_ok=True)
 
-# ✅ Define a writable cache directory
-cache_dir = "/tmp/huggingface_cache"
-os.makedirs(cache_dir, exist_ok=True)
-
-# ✅ Manually specify cache directory when loading the model
-pipeline = TrellisImageTo3DPipeline.from_pretrained(
-    "JeffreyXiang/TRELLIS-image-large", 
-    cache_dir=cache_dir
-)
+# Load the pipeline (no extra args like cache_dir)
+pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
 pipeline.cuda()
 
-# ✅ Preload model (to prevent cold starts)
+# Preload the model (avoids cold-start latencies)
 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")
+# Read your HF_API_KEY from Secrets (set in Space settings)
+HF_API_KEY = os.getenv("HF_API_KEY")
+if not HF_API_KEY:
+    raise RuntimeError("No HF_API_KEY found. Please set a secret in your Space settings.")
 
+# FastAPI App
 app = FastAPI()
 
+# (Optional) Limit max input image size
+MAX_IMAGE_SIZE = (1024, 1024)
+
 def preprocess_image(image: Image.Image) -> Image.Image:
-    """Preprocess a single input image using the Trellis pipeline."""
+    """Resize large images to keep memory usage in check, then let Trellis do its own preprocessing."""
+    image.thumbnail(MAX_IMAGE_SIZE)
     return pipeline.preprocess_image(image)
 
 @app.post("/generate_3d/")
@@ -46,27 +48,27 @@ 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
+    """Accept an image upload and return a .glb file of the 3D model."""
+    # Enforce HF_API_KEY check
     if authorization != f"Bearer {HF_API_KEY}":
         raise HTTPException(status_code=403, detail="Invalid API key")
 
+    # Require PNG/JPG
     if not image.filename.lower().endswith(("png", "jpg", "jpeg")):
-        raise HTTPException(status_code=400, detail="Invalid image format. Upload a PNG or JPG.")
+        raise HTTPException(status_code=400, detail="Upload PNG or JPG images.")
 
-    # Save the uploaded image
+    # Save upload to /tmp
     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
+    # Preprocess the image
     img = Image.open(image_path).convert("RGBA")
-    processed_image = preprocess_image(img)
+    processed = preprocess_image(img)
 
-    # Run the Trellis pipeline
+    # Run Trellis pipeline
     outputs = pipeline.run(
-        processed_image,
+        processed,
         seed=np.random.randint(0, np.iinfo(np.int32).max),
         formats=["gaussian", "mesh"],
         preprocess_image=False,
@@ -74,12 +76,20 @@ async def generate_3d(
         slat_sampler_params={"steps": 12, "cfg_strength": 3.0},
     )
 
-    # Extract the GLB file
+    # Extract and save the GLB
     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)
 
+    # Clear GPU memory
     torch.cuda.empty_cache()
 
+    # Return the GLB to the client
     return FileResponse(glb_path, media_type="model/gltf-binary", filename="sample.glb")
+
+# If you want to run locally or override CMD in Docker:
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.environ.get("PORT", "7860"))
+    uvicorn.run(app, host="0.0.0.0", port=port)