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