new

models/lrm/models/lrm_mesh.py

@@ -14,6 +14,7 @@
 
 import numpy as np
 import torch
+import spaces
 import torch.nn as nn
 import nvdiffrast.torch as dr
 from einops import rearrange, repeat
@@ -78,6 +79,7 @@ class PRM(nn.Module):
             samples_per_ray=rendering_samples_per_ray,
         )
 
+    @spaces.GPU
     def init_flexicubes_geometry(self, device, fovy=50.0):
         camera = PerspectiveCamera(fovy=fovy, device=device)
         renderer = NeuralRender(device, camera_model=camera)
@@ -103,6 +105,7 @@ class PRM(nn.Module):
 
         return planes
     
+    @spaces.GPU
     def get_sdf_deformation_prediction(self, planes):
         '''
         Predict SDF and deformation for tetrahedron vertices
@@ -195,6 +198,7 @@ class PRM(nn.Module):
         
         return v_list, f_list, imesh_list, sdf, deformation, v_deformed, (sdf_reg_loss, flexicubes_surface_reg, flexicubes_weight_reg)
     
+    @spaces.GPU
     def get_texture_prediction(self, planes, tex_pos, hard_mask=None, gb_normal=None, training=True):
         '''
         Predict Texture given triplanes
@@ -364,7 +368,8 @@ class PRM(nn.Module):
             'planes': planes,
             **out
         }
-
+    
+    @spaces.GPU
     def extract_mesh(
         self, 
         planes: torch.Tensor,