api_server.py

import os
import threading
from pathlib import Path
from typing import Any

import torch
from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from trellis_api_service import GenerationParameters, TMP_DIR, TrellisAPIService


def _parse_allowed_origins(value: str) -> list[str]:
    value = value.strip()
    if not value:
        return ["*"]
    if value == "*":
        return ["*"]
    return [origin.strip() for origin in value.split(",") if origin.strip()]


service = TrellisAPIService()
app = FastAPI(
    title="TRELLIS.2 API",
    version="1.0.0",
    description="API service for TRELLIS.2 image-to-3D generation and GLB export.",
)

allowed_origins = _parse_allowed_origins(os.environ.get("TRELLIS_ALLOWED_ORIGINS", "*"))
app.add_middleware(
    CORSMiddleware,
    allow_origins=allowed_origins,
    allow_credentials=allowed_origins != ["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)
Path(TMP_DIR).mkdir(parents=True, exist_ok=True)
app.mount("/artifacts", StaticFiles(directory=str(TMP_DIR)), name="artifacts")


class GenerateResponse(BaseModel):
    request_id: str
    seed: int
    resolution: str
    glb_url: str
    request_url: str
    preview_urls: dict[str, list[str]]


@app.on_event("startup")
def preload_model() -> None:
    preload_setting = os.environ.get("TRELLIS_PRELOAD_MODEL", "auto").lower()
    preload = preload_setting in {"1", "true", "yes", "on"} or (
        preload_setting == "auto" and torch.cuda.is_available()
    )
    if preload:
        threading.Thread(target=service.ensure_initialized, daemon=True).start()


@app.get("/")
def root() -> dict[str, Any]:
    return {
        "service": "TRELLIS.2 API",
        "docs_url": "/docs",
        "health_url": "/healthz",
        "parameters_url": "/v1/parameters",
        "generate_url": "/v1/generate",
    }


@app.get("/healthz")
def healthz() -> dict[str, Any]:
    return service.health()


@app.get("/v1/parameters")
def get_parameters() -> dict[str, Any]:
    return service.get_parameter_schema()


@app.post("/v1/generate", response_model=GenerateResponse)
def generate(
    request: Request,
    image: UploadFile = File(...),
    resolution: str = Form("1024"),
    seed: int = Form(0),
    randomize_seed: bool = Form(True),
    preprocess_image: bool = Form(True),
    decimation_target: int = Form(300000),
    texture_size: int = Form(2048),
    include_previews: bool = Form(True),
    ss_guidance_strength: float = Form(7.5),
    ss_guidance_rescale: float = Form(0.7),
    ss_sampling_steps: int = Form(12),
    ss_rescale_t: float = Form(5.0),
    shape_slat_guidance_strength: float = Form(7.5),
    shape_slat_guidance_rescale: float = Form(0.5),
    shape_slat_sampling_steps: int = Form(12),
    shape_slat_rescale_t: float = Form(3.0),
    tex_slat_guidance_strength: float = Form(1.0),
    tex_slat_guidance_rescale: float = Form(0.0),
    tex_slat_sampling_steps: int = Form(12),
    tex_slat_rescale_t: float = Form(3.0),
) -> GenerateResponse:
    valid_resolutions = {"512", "1024", "1536"}
    if resolution not in valid_resolutions:
        raise HTTPException(status_code=400, detail=f"Unsupported resolution '{resolution}'. Choose from {sorted(valid_resolutions)}.")

    params = GenerationParameters(
        resolution=resolution,
        seed=seed,
        randomize_seed=randomize_seed,
        preprocess_image=preprocess_image,
        decimation_target=decimation_target,
        texture_size=texture_size,
        include_previews=include_previews,
        ss_guidance_strength=ss_guidance_strength,
        ss_guidance_rescale=ss_guidance_rescale,
        ss_sampling_steps=ss_sampling_steps,
        ss_rescale_t=ss_rescale_t,
        shape_slat_guidance_strength=shape_slat_guidance_strength,
        shape_slat_guidance_rescale=shape_slat_guidance_rescale,
        shape_slat_sampling_steps=shape_slat_sampling_steps,
        shape_slat_rescale_t=shape_slat_rescale_t,
        tex_slat_guidance_strength=tex_slat_guidance_strength,
        tex_slat_guidance_rescale=tex_slat_guidance_rescale,
        tex_slat_sampling_steps=tex_slat_sampling_steps,
        tex_slat_rescale_t=tex_slat_rescale_t,
    )
    image_bytes = image.file.read()
    result = service.generate(image_bytes, image.filename or "upload.png", params)

    glb_path = Path(result["glb_path"])
    request_path = Path(result["request_path"])
    request_id = result["request_id"]

    preview_urls = {
        mode: [
            str(request.url_for("artifacts", path=f"{request_id}/previews/{Path(path).name}"))
            for path in paths
        ]
        for mode, paths in result["preview_paths"].items()
    }

    return GenerateResponse(
        request_id=request_id,
        seed=result["seed"],
        resolution=result["resolution"],
        glb_url=str(request.url_for("artifacts", path=f"{request_id}/{glb_path.name}")),
        request_url=str(request.url_for("artifacts", path=f"{request_id}/{request_path.name}")),
        preview_urls=preview_urls,
    )