Update src/facerender/modules/make_animation.py

src/facerender/modules/make_animation.py

@@ -144,131 +144,55 @@ def keypoint_transformation(kp_canonical, he, wo_exp=False):
 #         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
-#         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()
-
-#         # Stack predictions into a single tensor
-#         predictions_ts = torch.stack(predictions, dim=1)
-    
-#     return predictions_ts
-
 import torch
 from torch.cuda.amp import autocast
-from tqdm import tqdm
 
 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):
+                   use_exp=True):
 
-    device = 'cuda'
+    device='cuda'
     # Move inputs to GPU
     source_image = source_image.to(device)
     source_semantics = source_semantics.to(device)
     target_semantics = target_semantics.to(device)
-    print("source_image",source_image)
-    num_frames = target_semantics.shape[1]
-    predictions = []
 
     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)
-        print("kp_canonical",kp_canonical)
-        print("he_source",he_source)
-        print("kp_source",kp_source)
-
-        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]
 
-            batch_predictions = []
+        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)
 
-            for frame_idx in range(target_semantics_batch.shape[1]):
-                target_semantics_frame = target_semantics_batch[:, 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]
 
-                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
 
-                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)
 
-                print("source_image",source_image)
-                print("kp_source",kp_source)
-                print("kp_norm",kp_norm)
-
-                # 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)
-
-                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)
+            predictions.append(out['prediction'])
+            
+            # Optional: Explicitly synchronize (use only if necessary)
+            torch.cuda.synchronize()
 
-        # Stack all batch predictions into a single tensor
-        predictions_ts = torch.cat(predictions, dim=1)
+        # Stack predictions into a single tensor
+        predictions_ts = torch.stack(predictions, dim=1)
     
     return predictions_ts
 
 
-
-
-
 class AnimateModel(torch.nn.Module):
     """
     Merge all generator related updates into single model for better multi-gpu usage