Update src/facerender/modules/make_animation.py

src/facerender/modules/make_animation.py

@@ -99,94 +99,94 @@ def keypoint_transformation(kp_canonical, he, wo_exp=False):
     return {'value': kp_transformed}
 
 
-def make_animation(source_image, source_semantics, target_semantics,
-                            generator, kp_detector, he_estimator, mapping, 
-                            yaw_c_seq=None, pitch_c_seq=None, roll_c_seq=None,
-                            use_exp=True):
-    with torch.no_grad():
-        predictions = []
-
-        kp_canonical = kp_detector(source_image)
-        he_source = mapping(source_semantics)
-        kp_source = keypoint_transformation(kp_canonical, he_source)
-
-    
-        for frame_idx in tqdm(range(target_semantics.shape[1]), 'Face Renderer:'):
-            target_semantics_frame = target_semantics[:, frame_idx]
-            he_driving = mapping(target_semantics_frame)
-            if yaw_c_seq is not None:
-                he_driving['yaw_in'] = yaw_c_seq[:, frame_idx]
-            if pitch_c_seq is not None:
-                he_driving['pitch_in'] = pitch_c_seq[:, frame_idx] 
-            if roll_c_seq is not None:
-                he_driving['roll_in'] = roll_c_seq[:, frame_idx] 
-            
-            kp_driving = keypoint_transformation(kp_canonical, he_driving)
-                
-            #kp_norm = normalize_kp(kp_source=kp_source, kp_driving=kp_driving,
-                                   #kp_driving_initial=kp_driving_initial)
-            kp_norm = kp_driving
-            out = generator(source_image, kp_source=kp_source, kp_driving=kp_norm)
-            '''
-            source_image_new = out['prediction'].squeeze(1)
-            kp_canonical_new =  kp_detector(source_image_new)
-            he_source_new = he_estimator(source_image_new) 
-            kp_source_new = keypoint_transformation(kp_canonical_new, he_source_new, wo_exp=True)
-            kp_driving_new = keypoint_transformation(kp_canonical_new, he_driving, wo_exp=True)
-            out = generator(source_image_new, kp_source=kp_source_new, kp_driving=kp_driving_new)
-            '''
-            predictions.append(out['prediction'])
-            torch.cuda.empty_cache()
-        predictions_ts = torch.stack(predictions, dim=1)
-    return predictions_ts
-
-# import torch
-# from torch.cuda.amp import autocast
-
 # def make_animation(source_image, source_semantics, target_semantics,
-#                    generator, kp_detector, he_estimator, mapping, 
-#                    yaw_c_seq=None, pitch_c_seq=None, roll_c_seq=None,
-#                    use_exp=True):
-
-#     # device='cuda'
-#     # # Move inputs to GPU
-#     # source_image = source_image.to(device)
-#     # source_semantics = source_semantics.to(device)
-#     # target_semantics = target_semantics.to(device)
-
-#     with torch.no_grad():  # No gradients needed
+#                             generator, kp_detector, he_estimator, mapping, 
+#                             yaw_c_seq=None, pitch_c_seq=None, roll_c_seq=None,
+#                             use_exp=True):
+#     with torch.no_grad():
 #         predictions = []
+
 #         kp_canonical = kp_detector(source_image)
 #         he_source = mapping(source_semantics)
 #         kp_source = keypoint_transformation(kp_canonical, he_source)
 
-#         for frame_idx in tqdm(range(target_semantics.shape[1]), desc='Face Renderer:', unit='frame'):
+    
+#         for frame_idx in tqdm(range(target_semantics.shape[1]), 'Face Renderer:'):
 #             target_semantics_frame = target_semantics[:, frame_idx]
 #             he_driving = mapping(target_semantics_frame)
-
 #             if yaw_c_seq is not None:
 #                 he_driving['yaw_in'] = yaw_c_seq[:, frame_idx]
 #             if pitch_c_seq is not None:
-#                 he_driving['pitch_in'] = pitch_c_seq[:, frame_idx]
+#                 he_driving['pitch_in'] = pitch_c_seq[:, frame_idx] 
 #             if roll_c_seq is not None:
-#                 he_driving['roll_in'] = roll_c_seq[:, frame_idx]
-
+#                 he_driving['roll_in'] = roll_c_seq[:, frame_idx] 
+            
 #             kp_driving = keypoint_transformation(kp_canonical, he_driving)
+                
+#             #kp_norm = normalize_kp(kp_source=kp_source, kp_driving=kp_driving,
+#                                    #kp_driving_initial=kp_driving_initial)
 #             kp_norm = kp_driving
+#             out = generator(source_image, kp_source=kp_source, kp_driving=kp_norm)
+#             '''
+#             source_image_new = out['prediction'].squeeze(1)
+#             kp_canonical_new =  kp_detector(source_image_new)
+#             he_source_new = he_estimator(source_image_new) 
+#             kp_source_new = keypoint_transformation(kp_canonical_new, he_source_new, wo_exp=True)
+#             kp_driving_new = keypoint_transformation(kp_canonical_new, he_driving, wo_exp=True)
+#             out = generator(source_image_new, kp_source=kp_source_new, kp_driving=kp_driving_new)
+#             '''
+#             predictions.append(out['prediction'])
+#             torch.cuda.empty_cache()
+#         predictions_ts = torch.stack(predictions, dim=1)
+#     return predictions_ts
+
+import torch
+from torch.cuda.amp import autocast
+
+def make_animation(source_image, source_semantics, target_semantics,
+                   generator, kp_detector, he_estimator, mapping, 
+                   yaw_c_seq=None, pitch_c_seq=None, roll_c_seq=None,
+                   use_exp=True):
 
-#             # Use mixed precision for faster computation
-#             with autocast():
-#                 out = generator(source_image, kp_source=kp_source, kp_driving=kp_norm)
+    # device='cuda'
+    # # Move inputs to GPU
+    # source_image = source_image.to(device)
+    # source_semantics = source_semantics.to(device)
+    # target_semantics = target_semantics.to(device)
 
-#             predictions.append(out['prediction'])
+    with torch.no_grad():  # No gradients needed
+        predictions = []
+        kp_canonical = kp_detector(source_image)
+        he_source = mapping(source_semantics)
+        kp_source = keypoint_transformation(kp_canonical, he_source)
+
+        for frame_idx in tqdm(range(target_semantics.shape[1]), desc='Face Renderer:', unit='frame'):
+            target_semantics_frame = target_semantics[:, frame_idx]
+            he_driving = mapping(target_semantics_frame)
+
+            if yaw_c_seq is not None:
+                he_driving['yaw_in'] = yaw_c_seq[:, frame_idx]
+            if pitch_c_seq is not None:
+                he_driving['pitch_in'] = pitch_c_seq[:, frame_idx]
+            if roll_c_seq is not None:
+                he_driving['roll_in'] = roll_c_seq[:, frame_idx]
+
+            kp_driving = keypoint_transformation(kp_canonical, he_driving)
+            kp_norm = kp_driving
+
+            # Use mixed precision for faster computation
+            with autocast():
+                out = generator(source_image, kp_source=kp_source, kp_driving=kp_norm)
+
+            predictions.append(out['prediction'])
             
-#             # Optional: Explicitly synchronize (use only if necessary)
-#             torch.cuda.synchronize()
+            # Optional: Explicitly synchronize (use only if necessary)
+            torch.cuda.synchronize()
 
-#         # Stack predictions into a single tensor
-#         predictions_ts = torch.stack(predictions, dim=1)
+        # Stack predictions into a single tensor
+        predictions_ts = torch.stack(predictions, dim=1)
     
-#     return predictions_ts
+    return predictions_ts
 
 
 class AnimateModel(torch.nn.Module):