Update src/facerender/modules/make_animation.py

src/facerender/modules/make_animation.py

@@ -192,6 +192,55 @@ def make_animation(source_image, source_semantics, target_semantics,
     
 #     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):
+    device = 'cuda'
+    source_image = source_image.to(device)
+    source_semantics = source_semantics.to(device)
+    target_semantics = target_semantics.to(device)
+
+    with torch.no_grad():
+        predictions = []
+        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]
+
+            kp_driving = keypoint_transformation(kp_canonical, he_driving)
+            kp_norm = kp_driving
+
+            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
+
+
 
 
 class AnimateModel(torch.nn.Module):