jhj0517 commited on
Commit
e37e314
1 Parent(s): 533926d

add gpu annotation to load model function

Browse files
Files changed (2) hide show
  1. musepose_inference.py +1 -0
  2. pose_align.py +1 -0
musepose_inference.py CHANGED
@@ -191,6 +191,7 @@ class MusePoseInference:
191
  )
192
  return output_path, output_path_demo
193
 
 
194
  def init_model(self,
195
  weight_dtype: torch.dtype,
196
  infer_config: DictConfig
 
191
  )
192
  return output_path, output_path_demo
193
 
194
+ @spaces.GPU(duration=120)
195
  def init_model(self,
196
  weight_dtype: torch.dtype,
197
  infer_config: DictConfig
pose_align.py CHANGED
@@ -312,6 +312,7 @@ class PoseAlignmentInference:
312
  print('pose align done')
313
  return outfn_align_pose_video, outfn
314
 
 
315
  def init_model(self):
316
  if self.detector is None:
317
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
312
  print('pose align done')
313
  return outfn_align_pose_video, outfn
314
 
315
+ @spaces.GPU(duration=120)
316
  def init_model(self):
317
  if self.detector is None:
318
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')