Update src/facerender/animate.py

src/facerender/animate.py

@@ -119,45 +119,70 @@ class AnimateFromCoeff():
             optimizer_discriminator.load_state_dict(checkpoint['optimizer_discriminator'])
 
         return checkpoint['epoch']
-
+        
+from torch.cuda.amp import autocast
     def generate(self, x, video_save_dir, pic_path, crop_info, enhancer=None, background_enhancer=None, preprocess='crop'):
 
-        source_image=x['source_image'].type(torch.FloatTensor)
-        source_semantics=x['source_semantics'].type(torch.FloatTensor)
-        target_semantics=x['target_semantics_list'].type(torch.FloatTensor) 
-        source_image=source_image.to(self.device)
-        source_semantics=source_semantics.to(self.device)
-        target_semantics=target_semantics.to(self.device)
-        if 'yaw_c_seq' in x:
-            yaw_c_seq = x['yaw_c_seq'].type(torch.FloatTensor)
-            yaw_c_seq = x['yaw_c_seq'].to(self.device)
-        else:
-            yaw_c_seq = None
-        if 'pitch_c_seq' in x:
-            pitch_c_seq = x['pitch_c_seq'].type(torch.FloatTensor)
-            pitch_c_seq = x['pitch_c_seq'].to(self.device)
-        else:
-            pitch_c_seq = None
-        if 'roll_c_seq' in x:
-            roll_c_seq = x['roll_c_seq'].type(torch.FloatTensor) 
-            roll_c_seq = x['roll_c_seq'].to(self.device)
-        else:
-            roll_c_seq = None
-
+        # source_image=x['source_image'].type(torch.FloatTensor)
+        # source_semantics=x['source_semantics'].type(torch.FloatTensor)
+        # target_semantics=x['target_semantics_list'].type(torch.FloatTensor) 
+        # source_image=source_image.to(self.device)
+        # source_semantics=source_semantics.to(self.device)
+        # target_semantics=target_semantics.to(self.device)
+        # if 'yaw_c_seq' in x:
+        #     yaw_c_seq = x['yaw_c_seq'].type(torch.FloatTensor)
+        #     yaw_c_seq = x['yaw_c_seq'].to(self.device)
+        # else:
+        #     yaw_c_seq = None
+        # if 'pitch_c_seq' in x:
+        #     pitch_c_seq = x['pitch_c_seq'].type(torch.FloatTensor)
+        #     pitch_c_seq = x['pitch_c_seq'].to(self.device)
+        # else:
+        #     pitch_c_seq = None
+        # if 'roll_c_seq' in x:
+        #     roll_c_seq = x['roll_c_seq'].type(torch.FloatTensor) 
+        #     roll_c_seq = x['roll_c_seq'].to(self.device)
+        # else:
+        #     roll_c_seq = None
+
+        # frame_num = x['frame_num']
+
+        # predictions_video = make_animation(source_image, source_semantics, target_semantics,
+        #                                 self.generator, self.kp_extractor, self.he_estimator, self.mapping, 
+        #                                 yaw_c_seq, pitch_c_seq, roll_c_seq, use_exp = True)
+
+        # predictions_video = predictions_video.reshape((-1,)+predictions_video.shape[2:])
+        # predictions_video = predictions_video[:frame_num]
+
+        # video = []
+        # for idx in range(predictions_video.shape[0]):
+        #     image = predictions_video[idx]
+        #     image = np.transpose(image.data.cpu().numpy(), [1, 2, 0]).astype(np.float32)
+        #     video.append(image)
+        # result = img_as_ubyte(video)
+
+
+
+
+
+        source_image = x['source_image'].to(self.device).type(torch.FloatTensor)
+        source_semantics = x['source_semantics'].to(self.device).type(torch.FloatTensor)
+        target_semantics = x['target_semantics_list'].to(self.device).type(torch.FloatTensor)
+        yaw_c_seq = x.get('yaw_c_seq', None).to(self.device).type(torch.FloatTensor) if 'yaw_c_seq' in x else None
+        pitch_c_seq = x.get('pitch_c_seq', None).to(self.device).type(torch.FloatTensor) if 'pitch_c_seq' in x else None
+        roll_c_seq = x.get('roll_c_seq', None).to(self.device).type(torch.FloatTensor) if 'roll_c_seq' in x else None
         frame_num = x['frame_num']
 
-        predictions_video = make_animation(source_image, source_semantics, target_semantics,
+         with autocast():
+             predictions_video = make_animation(source_image, source_semantics, target_semantics,
                                         self.generator, self.kp_extractor, self.he_estimator, self.mapping, 
                                         yaw_c_seq, pitch_c_seq, roll_c_seq, use_exp = True)
 
-        predictions_video = predictions_video.reshape((-1,)+predictions_video.shape[2:])
+        predictions_video = predictions_video.reshape((-1,) + predictions_video.shape[2:])
         predictions_video = predictions_video[:frame_num]
 
-        video = []
-        for idx in range(predictions_video.shape[0]):
-            image = predictions_video[idx]
-            image = np.transpose(image.data.cpu().numpy(), [1, 2, 0]).astype(np.float32)
-            video.append(image)
+        # Create video
+        video = [np.transpose(img.data.cpu().numpy(), [1, 2, 0]).astype(np.float32) for img in predictions_video]
         result = img_as_ubyte(video)
 
         ### the generated video is 256x256, so we  keep the aspect ratio,