Create app.py

app.py

@@ -0,0 +1,78 @@
+# backend.py
+from fastapi import FastAPI, UploadFile, Form
+from pydantic import BaseModel
+from typing import List, Literal
+import os
+import shutil
+import torch
+import numpy as np
+from PIL import Image
+from trellis.pipelines import TrellisImageTo3DPipeline
+from trellis.utils import render_utils, postprocessing_utils
+
+app = FastAPI()
+
+MAX_SEED = np.iinfo(np.int32).max
+TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
+os.makedirs(TMP_DIR, exist_ok=True)
+
+pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+pipeline.cuda()
+
+
+@app.on_event("startup")
+def preload_model():
+    try:
+        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))  # Preload rembg
+    except:
+        pass
+
+
+@app.post("/image-to-3d")
+async def image_to_3d(
+    image: UploadFile,
+    seed: int = Form(...),
+    ss_guidance_strength: float = Form(...),
+    ss_sampling_steps: int = Form(...),
+    slat_guidance_strength: float = Form(...),
+    slat_sampling_steps: int = Form(...),
+    req_session: str = Form(...)
+):
+    user_dir = os.path.join(TMP_DIR, req_session)
+    os.makedirs(user_dir, exist_ok=True)
+    image_data = Image.open(image.file)
+
+    outputs = pipeline.run(
+        image_data,
+        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,
+        },
+    )
+
+    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+    video_path = os.path.join(user_dir, 'sample.mp4')
+    render_utils.save_video(video, video_path)
+    
+    torch.cuda.empty_cache()
+    return {"video_path": video_path}
+
+
+@app.post("/extract-glb")
+async def extract_glb(
+    mesh_simplify: float = Form(...),
+    texture_size: int = Form(...),
+    req_session: str = Form(...),
+):
+    user_dir = os.path.join(TMP_DIR, req_session)
+    glb_path = os.path.join(user_dir, 'sample.glb')
+    postprocessing_utils.export_glb(glb_path, simplify=mesh_simplify, texture_size=texture_size)
+
+    torch.cuda.empty_cache()
+    return {"glb_path": glb_path}