api.py

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"
    )