Spaces:
Runtime error
Runtime error
# 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() | |
def preload_model(): | |
try: | |
pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))) # Preload rembg | |
except: | |
pass | |
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} | |
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} | |