Update src/facerender/modules/make_animation.py

src/facerender/modules/make_animation.py

@@ -99,98 +99,98 @@ 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 = []
-#         device = 'cuda'
-#         source_image = source_image.to(device)
-#         source_semantics = source_semantics.to(device)
-#         target_semantics = target_semantics.to(device)
-
-#         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 = []
+        device = 'cuda'
+        source_image = source_image.to(device)
+        source_semantics = source_semantics.to(device)
+        target_semantics = target_semantics.to(device)
+
         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
 
-            # Use mixed precision for faster computation
-            with autocast():
-                out = generator(source_image, kp_source=kp_source, kp_driving=kp_norm)
+# import torch
+# from torch.cuda.amp import autocast
 
-            predictions.append(out['prediction'])
+# 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
+#         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