Upload folder using huggingface_hub

pitch.py

@@ -16,6 +16,8 @@ import torchvision.transforms as T
 import torchvision.transforms.functional as f
 from pydantic import BaseModel
 
+import tensorflow as tf
+
 import logging
 logger = logging.getLogger(__name__)
 
@@ -614,6 +616,9 @@ def inference_batch(frames, model, kp_threshold, device, batch_size=8):
     model_device = next(model.parameters()).device
     print(model_device)
     
+    @tf.function
+    def infer(x):
+        return model(x)
     # Process all frames in optimally-sized batches
     for i in range(0, num_frames, batch_size):
         current_batch_size = min(batch_size, num_frames - i)
@@ -626,8 +631,8 @@ def inference_batch(frames, model, kp_threshold, device, batch_size=8):
         # Move batch to model device
         batch = batch.to(model_device)
 
-        with torch.no_grad():
-            heatmaps = model(batch)
+        # with torch.no_grad():
+        heatmaps = infer(batch)
 
         # Ultra-fast keypoint extraction
         kp_coords = extract_keypoints_from_heatmap_fast(heatmaps[:,:-1,:,:], scale=2, max_keypoints=1)