suvadityamuk commited on
Commit
c17a338
·
1 Parent(s): 1a4c05f

chore: update function decorators to use spaces.gpu

Browse files

Signed-off-by: Suvaditya Mukherjee <suvadityamuk@gmail.com>

Files changed (1) hide show
  1. demos/pipeline_fn.py +6 -0
demos/pipeline_fn.py CHANGED
@@ -1,5 +1,6 @@
1
  import os
2
  import os.path as osp
 
3
  import gc
4
  import trimesh
5
  from PIL import Image
@@ -49,6 +50,7 @@ def file_sha256(path: str, chunk_size: int = 1 << 20) -> str:
49
  h.update(chunk)
50
  return h.hexdigest()
51
 
 
52
  def init_partfield(obj_path):
53
  torch.manual_seed(0)
54
  random.seed(0)
@@ -71,6 +73,7 @@ def init_partfield(obj_path):
71
  partfield_model.eval()
72
  return partfield_model
73
 
 
74
  def partfield_pipeline_predict(obj_path, output_dir):
75
 
76
  log.info("Extracting PartField feature planes...")
@@ -115,6 +118,7 @@ class GuideFlow3dPipeline:
115
  self.cfg = config
116
  return self
117
 
 
118
  def preprocess(
119
  self,
120
  structure_mesh: str,
@@ -225,6 +229,7 @@ class GuideFlow3dPipeline:
225
  "voxel_dir": voxel_dir
226
  }
227
 
 
228
  def run_appearance(
229
  self,
230
  structure_mesh: str,
@@ -408,6 +413,7 @@ class GuideFlow3dPipeline:
408
  return None, None
409
  return output_mesh_path, output_video_path
410
 
 
411
  def run_self_similarity(
412
  self,
413
  structure_mesh: str,
 
1
  import os
2
  import os.path as osp
3
+ import spaces
4
  import gc
5
  import trimesh
6
  from PIL import Image
 
50
  h.update(chunk)
51
  return h.hexdigest()
52
 
53
+ @spaces.GPU(duration=360)
54
  def init_partfield(obj_path):
55
  torch.manual_seed(0)
56
  random.seed(0)
 
73
  partfield_model.eval()
74
  return partfield_model
75
 
76
+ @spaces.GPU(duration=360)
77
  def partfield_pipeline_predict(obj_path, output_dir):
78
 
79
  log.info("Extracting PartField feature planes...")
 
118
  self.cfg = config
119
  return self
120
 
121
+ @spaces.GPU(duration=360)
122
  def preprocess(
123
  self,
124
  structure_mesh: str,
 
229
  "voxel_dir": voxel_dir
230
  }
231
 
232
+ @spaces.GPU(duration=360)
233
  def run_appearance(
234
  self,
235
  structure_mesh: str,
 
413
  return None, None
414
  return output_mesh_path, output_video_path
415
 
416
+ @spaces.GPU(duration=360)
417
  def run_self_similarity(
418
  self,
419
  structure_mesh: str,