Update src/facerender/animate.py

src/facerender/animate.py

@@ -27,7 +27,7 @@ from src.utils.videoio import save_video_with_watermark
 class AnimateFromCoeff():
 
     def __init__(self, free_view_checkpoint, mapping_checkpoint,
-                   config_path, device):
+                   config_path, device='cuda'):
 
         with open(config_path) as f:
             config = yaml.safe_load(f)
@@ -88,7 +88,7 @@ class AnimateFromCoeff():
     def load_cpk_facevid2vid(self, checkpoint_path, generator=None, discriminator=None, 
                         kp_detector=None, he_estimator=None, optimizer_generator=None, 
                         optimizer_discriminator=None, optimizer_kp_detector=None, 
-                        optimizer_he_estimator=None, device="cpu"):
+                        optimizer_he_estimator=None, device="cuda"):
         checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
         
         def adjust_state_dict(state_dict, model):
@@ -135,7 +135,7 @@ class AnimateFromCoeff():
         return checkpoint['epoch']
     
     def load_cpk_mapping(self, checkpoint_path, mapping=None, discriminator=None,
-                 optimizer_mapping=None, optimizer_discriminator=None, device='cpu'):
+                 optimizer_mapping=None, optimizer_discriminator=None, device='cuda'):
         checkpoint = torch.load(checkpoint_path,  map_location=torch.device(device))
         
         def adjust_state_dict(state_dict, model):