Update scripts/demo.py

scripts/demo.py

@@ -1,9 +1,9 @@
 import os
 import sys
-# os.environ["PYOPENGL_PLATFORM"] = "egl"
+
 os.environ['CUDA_VISIBLE_DEVICES'] = '0'
 sys.path.append(os.getcwd())
-os.environ["PYOPENGL_PLATFORM"] = 'osmesa'
+
 from transformers import Wav2Vec2Processor
 from glob import glob
 
@@ -24,8 +24,8 @@ from data_utils.rotation_conversion import rotation_6d_to_matrix, matrix_to_axis
 from data_utils.lower_body import part2full, pred2poses, poses2pred, poses2poses
 from visualise.rendering import RenderTool
 
-global device
-device = 'cpu'
+import time
+
 
 def init_model(model_name, model_path, args, config):
     if model_name == 's2g_face':
@@ -156,7 +156,7 @@ global_orient = torch.tensor([3.0747, -0.0158, -0.0152])
 
 
 def infer(g_body, g_face, smplx_model, rendertool, config, args):
-    betas = torch.zeros([1, 300], dtype=torch.float64).to(device)
+    betas = torch.zeros([1, 300], dtype=torch.float64).to('cuda')
     am = Wav2Vec2Processor.from_pretrained("vitouphy/wav2vec2-xls-r-300m-phoneme")
     am_sr = 16000
     num_sample = args.num_sample
@@ -165,7 +165,7 @@ def infer(g_body, g_face, smplx_model, rendertool, config, args):
     face = args.only_face
     stand = args.stand
     if face:
-        body_static = torch.zeros([1, 162], device=device)
+        body_static = torch.zeros([1, 162], device='cuda')
         body_static[:, 6:9] = torch.tensor([3.0747, -0.0158, -0.0152]).reshape(1, 3).repeat(body_static.shape[0], 1)
 
     result_list = []
@@ -179,7 +179,7 @@ def infer(g_body, g_face, smplx_model, rendertool, config, args):
                                       am=am,
                                       am_sr=am_sr
                                       )
-    pred_face = torch.tensor(pred_face).squeeze().to(device)
+    pred_face = torch.tensor(pred_face).squeeze().to('cuda')
     # pred_face = torch.zeros([gt.shape[0], 105])
 
     if config.Data.pose.convert_to_6d:
@@ -190,7 +190,7 @@ def infer(g_body, g_face, smplx_model, rendertool, config, args):
         pred_jaw = pred_face[:, :3]
         pred_face = pred_face[:, 3:]
 
-    id = torch.tensor([id], device=device)
+    id = torch.tensor([id], device='cuda')
 
     for i in range(num_sample):
         pred_res = g_body.infer_on_audio(cur_wav_file,
@@ -202,7 +202,7 @@ def infer(g_body, g_face, smplx_model, rendertool, config, args):
                                          fps=30,
                                          w_pre=False
                                          )
-        pred = torch.tensor(pred_res).squeeze().to(device)
+        pred = torch.tensor(pred_res).squeeze().to('cuda')
 
         if pred.shape[0] < pred_face.shape[0]:
             repeat_frame = pred[-1].unsqueeze(dim=0).repeat(pred_face.shape[0] - pred.shape[0], 1)
@@ -250,9 +250,8 @@ def infer(g_body, g_face, smplx_model, rendertool, config, args):
 def main():
     parser = parse_args()
     args = parser.parse_args()
-    # device = torch.device(args.gpu)
-    # torch.cuda.set_device(device)
-
+    device = torch.device(args.gpu)
+    torch.cuda.set_device(device)
 
     config = load_JsonConfig(args.config_file)
 
@@ -292,7 +291,7 @@ def main():
                         create_transl=False,
                         # gender='ne',
                         dtype=dtype, )
-    smplx_model = smpl.create(**model_params).to(device)
+    smplx_model = smpl.create(**model_params).to('cuda')
     print('init rendertool...')
     rendertool = RenderTool('visualise/video/' + config.Log.name)