Bread-F commited on
Commit
a49dbe4
·
verified ·
1 Parent(s): e41737c

Update utils/digital_human/realtime_inference.py

Browse files
utils/digital_human/realtime_inference.py CHANGED
@@ -111,8 +111,8 @@ def load_pose_model(model_dir):
111
  )
112
 
113
  config_file = r"./utils/digital_human/musetalk/utils/dwpose/rtmpose-l_8xb32-270e_coco-ubody-wholebody-384x288.py"
114
- pose_model = init_model(config_file, dw_pose_path, device="cuda")
115
-
116
  return pose_model
117
 
118
 
@@ -397,7 +397,8 @@ class Avatar:
397
  audio_feature_batch = self.model_handler.pe(audio_feature_batch)
398
  latent_batch = latent_batch.to(dtype=self.model_handler.unet.model.dtype)
399
 
400
- timesteps = torch.tensor([0], device="cuda")
 
401
  pred_latents = self.model_handler.unet.model(
402
  latent_batch, timesteps, encoder_hidden_states=audio_feature_batch
403
  ).sample
 
111
  )
112
 
113
  config_file = r"./utils/digital_human/musetalk/utils/dwpose/rtmpose-l_8xb32-270e_coco-ubody-wholebody-384x288.py"
114
+ pose_model = init_model(config_file, dw_pose_path, device="cpu")
115
+ #pose_model = init_model(config_file, dw_pose_path, device="cuda")
116
  return pose_model
117
 
118
 
 
397
  audio_feature_batch = self.model_handler.pe(audio_feature_batch)
398
  latent_batch = latent_batch.to(dtype=self.model_handler.unet.model.dtype)
399
 
400
+ # timesteps = torch.tensor([0], device="cuda")
401
+ timesteps = torch.tensor([0], device="cpu")
402
  pred_latents = self.model_handler.unet.model(
403
  latent_batch, timesteps, encoder_hidden_states=audio_feature_batch
404
  ).sample