Update handler.py

4911116 verified 11 months ago

1.8 kB

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}

	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}