add support to macOS GPU for inference

inference.py

@@ -1,8 +1,11 @@
-from glob import glob
+import os
 import shutil
+
+os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
+
 import torch
 from time import  strftime
-import os, sys, time
+import sys
 from argparse import ArgumentParser
 
 from src.utils.preprocess import CropAndExtract
@@ -136,10 +139,16 @@ def main(args):
 
     args = parser.parse_args()
 
-    if torch.cuda.is_available() and not args.cpu:
+    if args.cpu:
+        args.device = "cpu"
+    elif torch.backends.mps.is_available():
+        args.device = "mps"
+    elif torch.cuda.is_available():
         args.device = "cuda"
     else:
         args.device = "cpu"
 
+    print(f"Using device: {args.device}")
+
     main(args)
 

src/facerender/modules/dense_motion.py

@@ -35,23 +35,24 @@ def create_sparse_motions(self, feature, kp_driving, kp_source):
         bs, _, d, h, w = feature.shape
         identity_grid = make_coordinate_grid((d, h, w), type=kp_source['value'].type())
         identity_grid = identity_grid.view(1, 1, d, h, w, 3)
-        coordinate_grid = identity_grid - kp_driving['value'].view(bs, self.num_kp, 1, 1, 1, 3)
-        
+        coordinate_grid = identity_grid - kp_driving['value'].view(bs, self.num_kp, 1, 1, 1, 3).to(identity_grid.device)
+
         # if 'jacobian' in kp_driving:
         if 'jacobian' in kp_driving and kp_driving['jacobian'] is not None:
             jacobian = torch.matmul(kp_source['jacobian'], torch.inverse(kp_driving['jacobian']))
             jacobian = jacobian.unsqueeze(-3).unsqueeze(-3).unsqueeze(-3)
             jacobian = jacobian.repeat(1, 1, d, h, w, 1, 1)
             coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1))
-            coordinate_grid = coordinate_grid.squeeze(-1)                  
+            coordinate_grid = coordinate_grid.squeeze(-1)
 
 
-        driving_to_source = coordinate_grid + kp_source['value'].view(bs, self.num_kp, 1, 1, 1, 3)    # (bs, num_kp, d, h, w, 3)
+        driving_to_source = (coordinate_grid +
+                             kp_source['value'].view(bs, self.num_kp, 1, 1, 1, 3).to(coordinate_grid.device))    # (bs, num_kp, d, h, w, 3)
 
         #adding background feature
         identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1, 1)
         sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1)                #bs num_kp+1 d h w 3
-        
+
         # sparse_motions = driving_to_source
 
         return sparse_motions
@@ -90,7 +91,7 @@ def forward(self, feature, kp_driving, kp_source):
 
         heatmap = self.create_heatmap_representations(deformed_feature, kp_driving, kp_source)
 
-        input_ = torch.cat([heatmap, deformed_feature], dim=2)
+        input_ = torch.cat([heatmap.to(deformed_feature.device), deformed_feature], dim=2)
         input_ = input_.view(bs, -1, d, h, w)
 
         # input = deformed_feature.view(bs, -1, d, h, w)      # (bs, num_kp+1 * c, d, h, w)
@@ -102,9 +103,9 @@ def forward(self, feature, kp_driving, kp_source):
         mask = F.softmax(mask, dim=1)
         out_dict['mask'] = mask
         mask = mask.unsqueeze(2)                                   # (bs, num_kp+1, 1, d, h, w)
-        
-        zeros_mask = torch.zeros_like(mask)   
-        mask = torch.where(mask < 1e-3, zeros_mask, mask) 
+
+        zeros_mask = torch.zeros_like(mask)
+        mask = torch.where(mask < 1e-3, zeros_mask, mask).to(sparse_motion.device)
 
         sparse_motion = sparse_motion.permute(0, 1, 5, 2, 3, 4)    # (bs, num_kp+1, 3, d, h, w)
         deformation = (sparse_motion * mask).sum(dim=1)            # (bs, 3, d, h, w)

src/facerender/modules/keypoint_detector.py

@@ -47,7 +47,7 @@ def gaussian2kp(self, heatmap):
         """
         shape = heatmap.shape
         heatmap = heatmap.unsqueeze(-1)
-        grid = make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0)
+        grid = make_coordinate_grid(shape[2:], heatmap.type()).to(heatmap.device).unsqueeze_(0).unsqueeze_(0)
         value = (heatmap * grid).sum(dim=(2, 3, 4))
         kp = {'value': value}
 

src/facerender/modules/util.py

@@ -26,7 +26,7 @@ def kp2gaussian(kp, spatial_size, kp_variance):
     shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 3)
     mean = mean.view(*shape)
 
-    mean_sub = (coordinate_grid - mean)
+    mean_sub = (coordinate_grid - mean.to(coordinate_grid.device))
 
     out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance)
 
@@ -53,6 +53,9 @@ def make_coordinate_grid_2d(spatial_size, type):
 
 def make_coordinate_grid(spatial_size, type):
     d, h, w = spatial_size
+
+    if torch.backends.mps.is_available():
+        type = torch.float
     x = torch.arange(w).type(type)
     y = torch.arange(h).type(type)
     z = torch.arange(d).type(type)