1

.dockerignore

@@ -0,0 +1,13 @@
+.git
+.gitignore
+__pycache__
+*.pyc
+*.pyo
+*.pyd
+.Python
+env/
+venv/
+.env
+.venv
+tmp/
+*.log 

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.10-slim
+
+WORKDIR /code
+
+# 安装系统依赖
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# 复制依赖文件
+COPY requirements.txt .
+
+# 安装 Python 依赖
+RUN pip install --no-cache-dir -r requirements.txt
+
+# 复制应用代码
+COPY . .
+
+# 暴露端口
+EXPOSE 7860
+
+# 启动命令
+CMD ["python", "api.py"] 

api.py

@@ -0,0 +1,217 @@
+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)
+        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"
+    ) 

requirements.txt

@@ -20,6 +20,11 @@ git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057
 xformers==0.0.27.post2
 spconv-cu120==2.3.6
 transformers==4.46.3
+fastapi==0.110.0
+uvicorn[standard]==0.27.1
+python-multipart==0.0.9
+uvloop==0.19.0
+httptools==0.6.1
 https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.0.post2/flash_attn-2.7.0.post2+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
 https://huggingface.co/spaces/cavargas10/TRELLIS-Multiple3D/resolve/main/wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl?download=true
 https://huggingface.co/spaces/cavargas10/TRELLIS-Multiple3D/resolve/main/wheels/nvdiffrast-0.3.3-cp310-cp310-linux_x86_64.whl?download=true