Update src/facerender/animate.py

src/facerender/animate.py

@@ -122,66 +122,42 @@ class AnimateFromCoeff():
         
 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
-
-        # 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
+        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']
 
-        with autocast():
-             predictions_video = make_animation(source_image, source_semantics, target_semantics,
+        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]
 
-        # Create video
-        video = [np.transpose(img.data.cpu().numpy(), [1, 2, 0]).astype(np.float32) for img in predictions_video]
+        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)
 
         ### the generated video is 256x256, so we  keep the aspect ratio,