File size: 2,016 Bytes
ddfbfeb | 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 | import torch
import base64
import io
import tempfile
from PIL import Image
# Note: The import might change to 'trellis.2' or stay 'trellis' depending on how they packaged the V2 repo.
# Standard import usually resolves to the installed package name 'trellis' even for v2.
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.utils import postprocessing_utils
class EndpointHandler:
def __init__(self, path=""):
print("Loading Trellis 2 Model...")
# The CORRECT Model ID for V2 (4 Billion parameters)
self.pipeline = TrellisImageTo3DPipeline.from_pretrained(
"microsoft/TRELLIS.2-4B",
torch_dtype=torch.float16,
use_safetensors=True
)
self.pipeline.cuda()
print("Trellis 2 Loaded!")
def __call__(self, data):
# 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
outputs = self.pipeline.run(image, seed=42, formats=["mesh"])
# 3. Export to GLB
video_content = postprocessing_utils.to_glb(
outputs['mesh'][0],
simplify=0.95,
texture_size=1024
)
# 4. Handle Output
if isinstance(video_content, str):
with open(video_content, "rb") as f: glb_bytes = f.read()
elif isinstance(video_content, bytes):
glb_bytes = video_content
else:
with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmp:
outputs['mesh'][0].export(tmp.name, file_type='glb')
with open(tmp.name, "rb") as f: glb_bytes = f.read()
return {"glb": base64.b64encode(glb_bytes).decode('utf-8')} |