File size: 1,797 Bytes
ddfbfeb 4911116 ddfbfeb 4911116 ddfbfeb 4911116 ddfbfeb 4911116 ddfbfeb 4911116 ddfbfeb 4911116 | 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 | import torch
import base64
import io
import tempfile
from PIL import Image
# CORRECTED IMPORT: Package is 'trellis2', Class is 'Trellis2...'
from trellis2.pipelines import Trellis2ImageTo3DPipeline
class EndpointHandler:
def __init__(self, path=""):
# Load the official Microsoft Trellis 2 Model
print("Loading Trellis 2 (4B) Model...")
self.pipeline = Trellis2ImageTo3DPipeline.from_pretrained(
"microsoft/TRELLIS.2-4B",
torch_dtype=torch.float16,
use_safetensors=True
)
self.pipeline.cuda()
print("Trellis 2 Loaded!")
def __call__(self, data):
"""
Input: {"inputs": "base64_string"}
Output: {"glb": "base64_string"}
"""
# 1. Parse Input
inputs = data.pop("inputs", data)
if isinstance(inputs, dict) and "image" in inputs:
inputs = inputs["image"]
if isinstance(inputs, str):
image_data = base64.b64decode(inputs)
image = Image.open(io.BytesIO(image_data)).convert("RGB")
else:
image = inputs
# 2. Inference
# Trellis 2 returns a list of Mesh objects directly
outputs = self.pipeline.run(image, seed=42)
mesh_result = outputs[0]
# 3. Export to GLB
# The V2 mesh object has a direct export method
with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmp:
# export() takes a filepath string
mesh_result.export(tmp.name)
# Read back the bytes
with open(tmp.name, "rb") as f:
glb_bytes = f.read()
# 4. Return Base64
out_b64 = base64.b64encode(glb_bytes).decode('utf-8')
return {"glb": out_b64} |