File size: 5,149 Bytes
be473d8 e4cc66d be473d8 e4cc66d be473d8 e4cc66d be473d8 e4cc66d be473d8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | 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,
)
|