Update handler.py

handler.py

@@ -1,71 +1,47 @@
-import os
-import subprocess
-import sys
 import torch
 import base64
-import io
+import os
 from PIL import Image
-from typing import Dict, Any
+from io import BytesIO
+from trellis.pipelines import TrellisImageTo3DPipeline
+from trellis.utils import postprocessing_utils
 
 class EndpointHandler:
-    def __init__(self, model_dir: str):
-        # 1. Run custom installation for complex CUDA kernels
-        self._setup()
-        
-        # 2. Now import Trellis (must be after _setup)
-        from trellis.pipelines import TrellisImageTo3DPipeline
-        from trellis.utils import postprocessing_utils
-        self.postprocessing_utils = postprocessing_utils
-
-        # 3. Initialize pipeline
-        os.environ['SPCONV_ALGO'] = 'native'
+    def __init__(self, model_dir):
+        # Load the pipeline from the local directory
         self.pipeline = TrellisImageTo3DPipeline.from_pretrained("microsoft/TRELLIS-image-large")
         self.pipeline.cuda()
-
-    def _setup(self):
-        """Install dependencies that require --no-build-isolation and specific order."""
-        try:
-            import trellis
-            import nvdiffrast
-            print("Dependencies already satisfied.")
-        except ImportError:
-            print("Installing custom CUDA extensions and TRELLIS...")
-            # Install nvdiffrast and trellis directly from GitHub
-            # --no-build-isolation is required so they can see the installed PyTorch
-            packages = [
-                "git+https://github.com/NVlabs/nvdiffrast.git",
-                "git+https://github.com/microsoft/TRELLIS.git"
-            ]
-            for pkg in packages:
-                subprocess.check_call([
-                    sys.executable, "-m", "pip", "install", pkg, "--no-build-isolation"
-                ])
-
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        
+    def __call__(self, data):
+        """
+        Args:
+            data (:obj:`dict`):
+                - "inputs": The base64 encoded image or URL.
+                - "params": Dictionary of generation parameters (optional).
+        """
         inputs = data.pop("inputs", data)
-        params = data.pop("params", {})
-
-        if isinstance(inputs, str):
-            image = Image.open(io.BytesIO(base64.b64decode(inputs))).convert("RGB")
-        else:
-            image = inputs
+        params = data.pop("parameters", {})
 
-        outputs = self.pipeline.run(
+        # Decode image
+        image = Image.open(BytesIO(base64.b64decode(inputs)))
+        
+        # Run Inference
+        # Note: You can adjust 'steps' or 'cfg' via params
+        outputs = self.pipeline(
             image,
-            seed=params.get("seed", 42),
-            sparse_structure_sampler_params=params.get("sparse_params", {"steps": 12, "cfg_strength": 7.5}),
-            slat_sampler_params=params.get("slat_params", {"steps": 12, "cfg_strength": 3.0})
+            num_samples=1,
+            return_flags=["mesh"],
+            **params
         )
 
-        glb = self.postprocessing_utils.to_glb(
-            outputs['gaussian'][0], 
-            outputs['mesh'][0],
-            simplify=params.get("simplify", 0.95),
-            texture_size=params.get("texture_size", 1024)
-        )
-
-        buffered = io.BytesIO()
-        glb.export(buffered)
-        glb_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-        return {"mesh_base64": glb_str, "format": "glb"}
+        # Process mesh to GLB
+        mesh = outputs['mesh'][0]
+        glb_io = BytesIO()
+        mesh.export(glb_io, file_type='glb')
+        glb_io.seek(0)
+
+        # Encode GLB to base64 for the response
+        return {
+            "mesh_base64": base64.b64encode(glb_io.getvalue()).decode("utf-8"),
+            "format": "glb"
+        }