Spaces:
Runtime error
Runtime error
File size: 2,292 Bytes
9a1a4ec |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
# 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}
|