update

models/DiMeR/models/DiMeR.py

@@ -109,6 +109,7 @@ class DiMeR(nn.Module):
 
         return planes
     
+    @spaces.GPU
     def get_sdf_prediction(self, planes):
         '''
         Predict SDF and deformation for tetrahedron vertices
@@ -126,6 +127,7 @@ class DiMeR(nn.Module):
 
         return sdf
 
+    @spaces.GPU
     def get_sdf_deformation_prediction(self, planes):
         '''
         Predict SDF and deformation for tetrahedron vertices
@@ -183,6 +185,7 @@ class DiMeR(nn.Module):
         deformation = torch.cat(final_def, dim=0)
         return sdf, deformation, sdf_reg_loss, weight
     
+    @spaces.GPU
     def get_geometry_prediction(self, planes=None):
         '''
         Function to generate mesh with give triplanes
@@ -235,6 +238,7 @@ class DiMeR(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
@@ -474,6 +478,7 @@ class DiMeR(nn.Module):
             **out
         }
 
+    @spaces.GPU
     def extract_mesh(
         self, 
         planes: torch.Tensor, 

models/DiMeR/models/renderer/synthesizer_mesh.py

@@ -9,6 +9,7 @@ import itertools
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import spaces
 
 from .utils.renderer import generate_planes, project_onto_planes, sample_from_planes
 
@@ -94,6 +95,7 @@ class OSGDecoder(nn.Module):
 
         return sdf, deformation, weight
     
+    @spaces.GPU
     def get_texture_prediction(self, sampled_features):
         _N, n_planes, _M, _C = sampled_features.shape
         sampled_features = sampled_features.permute(0, 2, 1, 3).reshape(_N, _M, n_planes*_C)
@@ -156,6 +158,7 @@ class TriplaneSynthesizer(nn.Module):
         sdf, deformation, weight = self.decoder.get_geometry_prediction(sampled_features, flexicubes_indices)
         return sdf, deformation, weight
     
+    @spaces.GPU
     def get_texture_prediction(self, planes, sample_coordinates):
         plane_axes = self.plane_axes.to(planes.device)
         sampled_features = sample_from_planes(