Update inference.py

inference.py

@@ -8,6 +8,10 @@ import torch, face_detection
 from models import Wav2Lip
 import platform
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print('Using {} for inference.'.format(device))
+
+
 parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models')
 
 parser.add_argument('--checkpoint_path', type=str, 
@@ -166,17 +170,21 @@ def _load(checkpoint_path):
 	return checkpoint
 
 def load_model(path):
-	model = Wav2Lip()
-	print("Load checkpoint from: {}".format(path))
-	checkpoint = _load(path)
-	s = checkpoint["state_dict"]
-	new_s = {}
-	for k, v in s.items():
-		new_s[k.replace('module.', '')] = v
-	model.load_state_dict(new_s)
-
-	model = model.to(device)
-	return model.eval()
+    model = Wav2Lip()
+    print("Load checkpoint from: {}".format(path))
+
+    if device == 'cuda':
+        checkpoint = torch.load(path)
+    else:
+        checkpoint = torch.load(path, map_location=lambda storage, loc: storage)
+
+    s = checkpoint["state_dict"]
+    new_s = {k.replace('module.', ''): v for k, v in s.items()}
+    
+    model.load_state_dict(new_s)
+    model = model.to(device)
+    
+    return model.eval()
 
 def main():
 	if not os.path.isfile(args.face):