from fastapi import FastAPI, UploadFile, File, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse import uvicorn import os import shutil import torch import numpy as np from PIL import Image from typing import List, Optional from pydantic import BaseModel import imageio from trellis.pipelines import TrellisImageTo3DPipeline from trellis.representations import Gaussian, MeshExtractResult from trellis.utils import render_utils, postprocessing_utils from easydict import EasyDict as edict app = FastAPI(title="TRELLIS 3D API") # 添加 CORS 中间件 app.add_middleware( CORSMiddleware, allow_origins=["*"], # 在生产环境中应该设置具体的域名 allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # 配置 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp') os.makedirs(TMP_DIR, exist_ok=True) MAX_SEED = np.iinfo(np.int32).max # 初始化 pipeline pipeline = TrellisImageTo3DPipeline.from_pretrained("cavargas10/TRELLIS") pipeline.cuda() try: pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))) # 预加载 rembg except: pass class GenerationParams(BaseModel): seed: int = 0 ss_guidance_strength: float = 7.5 ss_sampling_steps: int = 12 slat_guidance_strength: float = 3.0 slat_sampling_steps: int = 12 multiimage_algo: str = "stochastic" class GLBParams(BaseModel): mesh_simplify: float = 0.95 texture_size: int = 1024 def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict: return { 'gaussian': { **gs.init_params, '_xyz': gs._xyz.cpu().numpy(), '_features_dc': gs._features_dc.cpu().numpy(), '_scaling': gs._scaling.cpu().numpy(), '_rotation': gs._rotation.cpu().numpy(), '_opacity': gs._opacity.cpu().numpy(), }, 'mesh': { 'vertices': mesh.vertices.cpu().numpy(), 'faces': mesh.faces.cpu().numpy(), }, } def unpack_state(state: dict) -> tuple[Gaussian, edict]: gs = Gaussian( aabb=state['gaussian']['aabb'], sh_degree=state['gaussian']['sh_degree'], mininum_kernel_size=state['gaussian']['mininum_kernel_size'], scaling_bias=state['gaussian']['scaling_bias'], opacity_bias=state['gaussian']['opacity_bias'], scaling_activation=state['gaussian']['scaling_activation'], ) gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda') gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda') gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda') gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda') gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda') mesh = edict( vertices=torch.tensor(state['mesh']['vertices'], device='cuda'), faces=torch.tensor(state['mesh']['faces'], device='cuda'), ) return gs, mesh @app.post("/generate") async def generate_3d( files: List[UploadFile] = File(...), params: GenerationParams = None ): if not params: params = GenerationParams() # 创建临时目录 session_id = str(np.random.randint(0, MAX_SEED)) user_dir = os.path.join(TMP_DIR, session_id) os.makedirs(user_dir, exist_ok=True) try: # 处理上传的图片 images = [] for file in files: image = Image.open(file.file) images.append(image) # 运行生成 outputs = pipeline.run_multi_image( images, seed=params.seed, formats=["gaussian", "mesh"], preprocess_image=False, sparse_structure_sampler_params={ "steps": params.ss_sampling_steps, "cfg_strength": params.ss_guidance_strength, }, slat_sampler_params={ "steps": params.slat_sampling_steps, "cfg_strength": params.slat_guidance_strength, }, mode=params.multiimage_algo, ) # 生成预览视频 video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color'] video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal'] video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))] video_path = os.path.join(user_dir, 'preview.mp4') imageio.mimsave(video_path, video, fps=15) # 保存状态 state = pack_state(outputs['gaussian'][0], outputs['mesh'][0]) state_path = os.path.join(user_dir, 'state.npz') np.savez(state_path, **state) return { "session_id": session_id, "preview_url": f"/preview/{session_id}", "state_url": f"/state/{session_id}" } except Exception as e: shutil.rmtree(user_dir) raise HTTPException(status_code=500, detail=str(e)) @app.post("/extract_glb") async def extract_glb( session_id: str, params: GLBParams = None ): if not params: params = GLBParams() user_dir = os.path.join(TMP_DIR, session_id) if not os.path.exists(user_dir): raise HTTPException(status_code=404, detail="Session not found") try: # 加载状态 state_path = os.path.join(user_dir, 'state.npz') state = np.load(state_path, allow_pickle=True, encoding='latin1') state = {k: state[k] for k in state.files} # 生成 GLB gs, mesh = unpack_state(state) glb = postprocessing_utils.to_glb( gs, mesh, simplify=params.mesh_simplify, texture_size=params.texture_size, verbose=False ) glb_path = os.path.join(user_dir, 'model.glb') glb.export(glb_path) return {"glb_url": f"/glb/{session_id}"} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.get("/preview/{session_id}") async def get_preview(session_id: str): preview_path = os.path.join(TMP_DIR, session_id, 'preview.mp4') if not os.path.exists(preview_path): raise HTTPException(status_code=404, detail="Preview not found") return FileResponse(preview_path) @app.get("/glb/{session_id}") async def get_glb(session_id: str): glb_path = os.path.join(TMP_DIR, session_id, 'model.glb') if not os.path.exists(glb_path): raise HTTPException(status_code=404, detail="GLB not found") return FileResponse(glb_path) @app.get("/state/{session_id}") async def get_state(session_id: str): state_path = os.path.join(TMP_DIR, session_id, 'state.npz') if not os.path.exists(state_path): raise HTTPException(status_code=404, detail="State not found") return FileResponse(state_path) if __name__ == "__main__": # 在 Hugging Face Spaces 中，我们需要使用 0.0.0.0 作为主机 # 端口 7860 是 Hugging Face Spaces 的默认端口 uvicorn.run( app, host="0.0.0.0", port=7860, # 添加以下配置以提高性能 workers=1, # 由于 GPU 限制，使用单工作进程 loop="uvloop", http="httptools" )