From bc82a8ac9ec559f4395cbed150b093f4804a343a Mon Sep 17 00:00:00 2001
From: Ali Hamdi Ali Fadel <aliosm1997@gmail.com>
Date: Thu, 27 Jun 2024 15:15:28 +0000
Subject: [PATCH] Determine the file size before loading it to handle OOM
 issues

---
 src/audio_splitter.py | 27 +++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/src/audio_splitter.py b/src/audio_splitter.py
index bd9d55a..e512eb9 100644
--- a/src/audio_splitter.py
+++ b/src/audio_splitter.py
@@ -5,11 +5,11 @@
 from auditok import AudioRegion
 from auditok.core import split
 from pydub import AudioSegment
-from pydub.exceptions import CouldntDecodeError
 from pydub.generators import WhiteNoise
 from pydub.utils import mediainfo
 
 
+MAX_FILE_SIZE = 4 * 1024 * 1024
 MAX_FILE_DURATION = 1 * 60 * 60
 
 
@@ -24,7 +24,20 @@ def split(
     noise_seconds: int = 1,
     noise_amplitude: int = 0,
   ) -> list[tuple[bytes, float, float]]:
-    try:
+    file_info = mediainfo(file_path)
+    file_size = float(file_info['duration']) * int(file_info['sampling_rate']) * int(file_info['channels']) * 16 / 8
+
+    if file_size > MAX_FILE_SIZE:
+      return self._split_large_file(
+        file_path,
+        min_dur,
+        max_dur,
+        max_silence,
+        energy_threshold,
+        noise_seconds,
+        noise_amplitude,
+      )
+    else:
       segments = [
         (
           self._expand_segment_with_noise(segment, noise_seconds, noise_amplitude),
@@ -40,16 +53,6 @@ def split(
       ]
 
       return self._segments_to_data(segments)
-    except CouldntDecodeError:
-      return self._split_large_file(
-        file_path,
-        min_dur,
-        max_dur,
-        max_silence,
-        energy_threshold,
-        noise_seconds,
-        noise_amplitude,
-      )
 
   def _split_large_file(
     self,