Update handler.py

handler.py

@@ -3,16 +3,14 @@ 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
+# CORRECTED IMPORT: Package is 'trellis2', Class is 'Trellis2...'
+from trellis2.pipelines import Trellis2ImageTo3DPipeline
 
 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(
+        # 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
@@ -21,6 +19,10 @@ class EndpointHandler:
         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:
@@ -33,23 +35,20 @@ class EndpointHandler:
             image = inputs
 
         # 2. Inference
-        outputs = self.pipeline.run(image, seed=42, formats=["mesh"])
+        # Trellis 2 returns a list of Mesh objects directly
+        outputs = self.pipeline.run(image, seed=42)
+        mesh_result = outputs[0]
         
         # 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()
+        # 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()
 
-        return {"glb": base64.b64encode(glb_bytes).decode('utf-8')}
+        # 4. Return Base64
+        out_b64 = base64.b64encode(glb_bytes).decode('utf-8')
+        return {"glb": out_b64}