support for cpu also

inference.py

@@ -14,6 +14,8 @@
 from modules.audio2kp import AudioModel3D
 import yaml,os,imageio
 
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
 def draw_annotation_box( image, rotation_vector, translation_vector, color=(255, 255, 255), line_width=2):
     """Draw a 3D box as annotation of pose"""
 
@@ -130,26 +132,26 @@ def audio2head(audio_path, img_path, model_path, save_path):
 
     img = np.array(img_as_float32(img))
     img = img.transpose((2, 0, 1))
-    img = torch.from_numpy(img).unsqueeze(0).cuda()
+    img = torch.from_numpy(img).unsqueeze(0).to(device)
 
 
     ref_pose_rot, ref_pose_trans = get_pose_from_audio(img, audio_feature, model_path)
-    torch.cuda.empty_cache()
+    # torch.cuda.empty_cache()
 
     config_file = r"./config/vox-256.yaml"
     with open(config_file) as f:
-        config = yaml.load(f)
+        config = yaml.safe_load(f)
     kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
                              **config['model_params']['common_params'])
     generator = OcclusionAwareGenerator(**config['model_params']['generator_params'],
                                         **config['model_params']['common_params'])
-    kp_detector = kp_detector.cuda()
-    generator = generator.cuda()
+    kp_detector = kp_detector.to(device)
+    generator = generator.to(device)
 
-    opt = argparse.Namespace(**yaml.load(open("./config/parameters.yaml")))
-    audio2kp = AudioModel3D(opt).cuda()
+    opt = argparse.Namespace(**yaml.safe_load(open("./config/parameters.yaml")))
+    audio2kp = AudioModel3D(opt).to(device)
 
-    checkpoint  = torch.load(model_path)
+    checkpoint  = torch.load(model_path, map_location=device)
     kp_detector.load_state_dict(checkpoint["kp_detector"])
     generator.load_state_dict(checkpoint["generator"])
     audio2kp.load_state_dict(checkpoint["audio2kp"])
@@ -190,8 +192,8 @@ def audio2head(audio_path, img_path, model_path, save_path):
     for bs_idx in range(bs):
         t = {}
 
-        t["audio"] = audio_f[:, bs_idx].cuda()
-        t["pose"] = poses[:, bs_idx].cuda()
+        t["audio"] = audio_f[:, bs_idx].to(device)
+        t["pose"] = poses[:, bs_idx].to(device)
         t["id_img"] = img
         kp_gen_source = kp_detector(img)
 

modules/audio2pose.py

@@ -3,6 +3,8 @@
 from modules.util import MyResNet34
 import numpy as np
 
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
 class audio2poseLSTM(nn.Module):
     def __init__(self):
         super(audio2poseLSTM,self).__init__()
@@ -35,9 +37,9 @@ def get_pose_from_audio(img,audio,model_path):
     maxv = np.array([0.411, 0.547, 0.433, 159.1, 116.5, 376.5], dtype=np.float32)
 
 
-    generator = audio2poseLSTM().cuda()
+    generator = audio2poseLSTM().to(device)
 
-    ckpt_para = torch.load(model_path)
+    ckpt_para = torch.load(model_path, map_location=device)
 
     generator.load_state_dict(ckpt_para["audio2pose"])
     generator.eval()
@@ -46,7 +48,7 @@ def get_pose_from_audio(img,audio,model_path):
     for i in range(num_frame):
         audio_seq.append(audio[i*4:i*4+4])
 
-    audio = torch.from_numpy(np.array(audio_seq,dtype=np.float32)).unsqueeze(0).cuda()
+    audio = torch.from_numpy(np.array(audio_seq,dtype=np.float32)).unsqueeze(0).to(device)
 
     x = {}
     x["img"] = img