Support longer audio files reducing memory usage with chunking

tests/test_audio.py

@@ -7,13 +7,13 @@
 
 def test_audio():
     audio_path = os.path.join(os.path.dirname(__file__), "jfk.flac")
-    audio = load_audio(audio_path)
+    audio = next(load_audio(audio_path))
     assert audio.ndim == 1
     assert SAMPLE_RATE * 10 < audio.shape[0] < SAMPLE_RATE * 12
     assert 0 < audio.std() < 1
 
-    mel_from_audio = log_mel_spectrogram(audio)
-    mel_from_file = log_mel_spectrogram(audio_path)
+    mel_from_audio = next(log_mel_spectrogram(audio))
+    mel_from_file = next(log_mel_spectrogram(audio_path))
 
     assert np.allclose(mel_from_audio, mel_from_file)
     assert mel_from_audio.max() - mel_from_audio.min() <= 2.0

whisper/audio.py

@@ -1,7 +1,8 @@
 import os
+import subprocess
 from functools import lru_cache
 from subprocess import CalledProcessError, run
-from typing import Optional, Union
+from typing import Generator, Optional, Union
 
 import numpy as np
 import torch
@@ -21,6 +22,7 @@
 FRAMES_PER_SECOND = exact_div(SAMPLE_RATE, HOP_LENGTH)  # 10ms per audio frame
 TOKENS_PER_SECOND = exact_div(SAMPLE_RATE, N_SAMPLES_PER_TOKEN)  # 20ms per audio token
 
+MAX_CHUNK_DURATION = 2 * 60 * 60 # 2 hour maximum chunk duration
 
 def load_audio(file: str, sr: int = SAMPLE_RATE):
     """
@@ -55,12 +57,16 @@ def load_audio(file: str, sr: int = SAMPLE_RATE):
     ]
     # fmt: on
     try:
-        out = run(cmd, capture_output=True, check=True).stdout
-    except CalledProcessError as e:
+        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+        while True:
+            out = process.stdout.read(MAX_CHUNK_DURATION * sr * 2)
+            if not out:
+                break
+            yield np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+    except Exception as e:
         raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
 
-    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
-
 
 def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1):
     """
@@ -108,7 +114,7 @@ def mel_filters(device, n_mels: int) -> torch.Tensor:
 
 
 def log_mel_spectrogram(
-    audio: Union[str, np.ndarray, torch.Tensor],
+    audio: Union[str, np.ndarray, torch.Tensor, Generator[np.ndarray, None, None]],
     n_mels: int = 80,
     padding: int = 0,
     device: Optional[Union[str, torch.device]] = None,
@@ -135,13 +141,26 @@ def log_mel_spectrogram(
     torch.Tensor, shape = (80, n_frames)
         A Tensor that contains the Mel spectrogram
     """
-    if not torch.is_tensor(audio):
-        if isinstance(audio, str):
-            audio = load_audio(audio)
-        audio = torch.from_numpy(audio)
-
-    if device is not None:
-        audio = audio.to(device)
+    if isinstance(audio, str):
+        audio = load_audio(audio)
+    elif isinstance(audio, np.ndarray):
+        audio = [audio]
+    elif isinstance(audio, torch.Tensor):
+        audio = [audio]
+
+    for chunk in audio:
+        if not isinstance(chunk, torch.Tensor):
+            chunk = torch.from_numpy(chunk)
+        if device is not None:
+            chunk = chunk.to(device)
+        yield _log_mel_spectrogram(chunk, n_mels, padding)
+
+
+def _log_mel_spectrogram(
+    audio: torch.Tensor,
+    n_mels: int = 80,
+    padding: int = 0,
+):
     if padding > 0:
         audio = F.pad(audio, (0, padding))
     window = torch.hann_window(N_FFT).to(audio.device)