Update

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces
 # from gradio_litmodel3d import LitModel3D
 
 import os
@@ -23,6 +24,7 @@ def preprocess_image(image: Image.Image) -> Image.Image:
     return pipeline.preprocess_image(image)
 
 
+@spaces.GPU
 def image_to_3d(image: Image.Image) -> Tuple[dict, str]:
     """
     Convert an image to a 3D model.
@@ -44,6 +46,7 @@ def image_to_3d(image: Image.Image) -> Tuple[dict, str]:
     return model, video_path
 
 
+@spaces.GPU
 def extract_glb(model: dict, mesh_simplify: float, texture_size: int) -> Tuple[str, str]:
     """
     Extract a GLB file from the 3D model.

trellis/models/structured_latent_vae/decoder_mesh.py

@@ -102,8 +102,8 @@ class SLatMeshDecoder(SparseTransformerBase):
         )
         self.resolution = resolution
         self.rep_config = representation_config
-        mesh_extractor = SparseFeatures2Mesh('cpu', res=self.resolution*4, use_color=self.rep_config.get('use_color', False))
-        self.out_channels = mesh_extractor.feats_channels
+        self.mesh_extractor = SparseFeatures2Mesh(res=self.resolution*4, use_color=self.rep_config.get('use_color', False))
+        self.out_channels = self.mesh_extractor.feats_channels
         self.upsample = nn.ModuleList([
             SparseSubdivideBlock3d(
                 channels=model_channels,
@@ -153,9 +153,8 @@ class SLatMeshDecoder(SparseTransformerBase):
             list of representations
         """
         ret = []
-        mesh_extractor = SparseFeatures2Mesh(x.device, res=self.resolution*4, use_color=self.rep_config.get('use_color', False))
         for i in range(x.shape[0]):
-            mesh = mesh_extractor(x[i], training=self.training)
+            mesh = self.mesh_extractor(x[i], training=self.training)
             ret.append(mesh)
         return ret