Update script.py

script.py

@@ -22,14 +22,18 @@ def preprocess(file_like):
     container = av.open(file_like)
     frames = []
     every = 10
+    MAX_MEMORY = 100 * 1024 * 1024  ## 100 MB maximum - some videos are large
+    current_memory = 0
     for i, frame in enumerate(container.decode(video=0)):
         if i % every == 0:
             frame_array = frame.to_ndarray(format="rgb24")
             frame_tensor = torch.from_numpy(frame_array).permute(2, 0, 1).float()
             frames.append(frame_tensor)
-
-    video_tensor = torch.stack(frames)
-    return video_tensor
+            ## Memory check
+            frame_bytes = frame_tensor.numel() * 4  # float32 → 4 bytes
+            current_memory += frame_bytes
+            if current_memory >= MAX_MEMORY:
+                break
 
 
 class Model(torch.nn.Module):