update

src/utils/frame_interpolation.py

@@ -5,6 +5,13 @@ import torch
 import bisect
 import shutil
 
+if torch.backends.mps.is_available():
+  device = "mps"
+elif torch.cuda.is_available():
+  device = "cuda"
+else:
+  device = "cpu"
+
 def init_frame_interpolation_model():
     print("Initializing frame interpolation model")
     checkpoint_name = os.path.join("./pretrained_model/film_net_fp16.pt")
@@ -12,7 +19,7 @@ def init_frame_interpolation_model():
     model = torch.load(checkpoint_name, map_location='cpu')
     model.eval()
     model = model.half()
-    model = model.to(device="cuda")
+    model = model.to(device=device)
     return model
 
 
@@ -54,8 +61,8 @@ def batch_images_interpolation_tool(input_file, model, fps, inter_frames=1):
 
             x0 = x0.half()
             x1 = x1.half()
-            x0 = x0.cuda()
-            x1 = x1.cuda()
+            x0 = x0.to(device)
+            x1 = x1.to(device)
 
             dt = x0.new_full((1, 1), (splits[remains[step]] - splits[idxes[start_i]])) / (splits[idxes[end_i]] - splits[idxes[start_i]])
 
@@ -87,4 +94,4 @@ def batch_images_interpolation_tool(input_file, model, fps, inter_frames=1):
     
     shutil.rmtree(image_save_dir)
     
-    return video_save_dir
+    return video_save_dir

src/utils/util.py

@@ -12,6 +12,13 @@ import torchvision
 from einops import rearrange
 from PIL import Image
 
+if torch.backends.mps.is_available():
+  device = "mps"
+elif torch.cuda.is_available():
+  device = "cuda"
+else:
+  device = "cpu"
+
 
 def seed_everything(seed):
     import random
@@ -19,7 +26,8 @@ def seed_everything(seed):
     import numpy as np
 
     torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
+    if device == "cuda":
+        torch.cuda.manual_seed_all(seed)
     np.random.seed(seed % (2**32))
     random.seed(seed)