Create handler.py

handler.py

@@ -0,0 +1,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')}