Update src/facerender/modules/make_animation.py

src/facerender/modules/make_animation.py

@@ -99,50 +99,50 @@ 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)
+# 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)
+#         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] 
+#         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_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
+#             #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
@@ -200,49 +200,68 @@ 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, batch_size=8):
+
     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():
-        predictions = []
+    num_frames = target_semantics.shape[1]
+    predictions = []
+
+    with torch.no_grad():  # No gradients needed
         kp_canonical = kp_detector(source_image)
         he_source = mapping(source_semantics)
         kp_source = keypoint_transformation(kp_canonical, he_source)
 
-        num_frames = target_semantics.shape[1]
         for start_idx in tqdm(range(0, num_frames, batch_size), desc='Face Renderer:', unit='batch'):
             end_idx = min(start_idx + batch_size, num_frames)
             target_semantics_batch = target_semantics[:, start_idx:end_idx]
-            he_driving = mapping(target_semantics_batch)
 
-            if yaw_c_seq is not None:
-                he_driving['yaw_in'] = yaw_c_seq[:, start_idx:end_idx]
-            if pitch_c_seq is not None:
-                he_driving['pitch_in'] = pitch_c_seq[:, start_idx:end_idx]
-            if roll_c_seq is not None:
-                he_driving['roll_in'] = roll_c_seq[:, start_idx:end_idx]
+            batch_predictions = []
 
-            kp_driving = keypoint_transformation(kp_canonical, he_driving)
-            kp_norm = kp_driving
+            for frame_idx in range(target_semantics_batch.shape[1]):
+                target_semantics_frame = target_semantics_batch[:, frame_idx]
+                he_driving = mapping(target_semantics_frame)
 
-            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()
+                if yaw_c_seq is not None:
+                    he_driving['yaw_in'] = yaw_c_seq[:, start_idx:end_idx, frame_idx]
+                if pitch_c_seq is not None:
+                    he_driving['pitch_in'] = pitch_c_seq[:, start_idx:end_idx, frame_idx]
+                if roll_c_seq is not None:
+                    he_driving['roll_in'] = roll_c_seq[:, start_idx:end_idx, frame_idx]
+
+                kp_driving = keypoint_transformation(kp_canonical, he_driving)
+                kp_norm = kp_driving
+
+                # Ensure correct input shape for conv1d layers
+                source_image_reshaped = source_image.view(source_image.size(0), source_image.size(1), -1)
+                kp_source_reshaped = kp_source.view(kp_source.size(0), kp_source.size(1), -1)
+                kp_norm_reshaped = kp_norm.view(kp_norm.size(0), kp_norm.size(1), -1)
+
+                # Use mixed precision for faster computation
+                with autocast():
+                    out = generator(source_image_reshaped, kp_source=kp_source_reshaped, kp_driving=kp_norm_reshaped)
 
-        # Stack predictions into a single tensor
-        predictions_ts = torch.stack(predictions, dim=1)
+                batch_predictions.append(out['prediction'])
+
+                # Optional: Explicitly synchronize (use only if necessary)
+                torch.cuda.synchronize()
+
+            # Stack predictions for this batch
+            batch_predictions_ts = torch.stack(batch_predictions, dim=1)
+            predictions.append(batch_predictions_ts)
+
+        # Stack all batch predictions into a single tensor
+        predictions_ts = torch.cat(predictions, dim=1)
     
     return predictions_ts
 
 
 
 
+
 class AnimateModel(torch.nn.Module):
     """
     Merge all generator related updates into single model for better multi-gpu usage