perf(whisper): reuse encoded audio features for word timestamps

whisper/mlx_whisper/timing.py

@@ -2,7 +2,7 @@
 
 import itertools
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, List, Optional
 
 import mlx.core as mx
 import numba
@@ -116,6 +116,7 @@ def find_alignment(
     mel: mx.array,
     num_frames: int,
     *,
+    audio_features: Optional[mx.array] = None,
     medfilt_width: int = 7,
     qk_scale: float = 1.0,
 ) -> List[WordTiming]:
@@ -131,7 +132,13 @@ def find_alignment(
         ]
     )
 
-    logits, cross_qk = model.forward_with_cross_qk(mel[None, :], tokens[None, :])
+    if audio_features is not None and audio_features.ndim == 2:
+        audio_features = audio_features[None, :]
+    logits, cross_qk = model.forward_with_cross_qk(
+        mel[None, :],
+        tokens[None, :],
+        audio_features=audio_features,
+    )
     # consider only the logits associated with predicting text
     sampled_logits = logits[0][len(tokenizer.sot_sequence) : -2, : tokenizer.eot]
     token_probs = mx.softmax(sampled_logits, precise=True, axis=-1)

whisper/mlx_whisper/transcribe.py

@@ -418,6 +418,7 @@ def next_words_segment(segments: List[dict]) -> Optional[dict]:
                         tokenizer=tokenizer,
                         mel=mel_segment,
                         num_frames=segment_size,
+                        audio_features=result.audio_features,
                         prepend_punctuations=prepend_punctuations,
                         append_punctuations=append_punctuations,
                         last_speech_timestamp=last_speech_timestamp,

whisper/mlx_whisper/whisper.py

@@ -247,8 +247,10 @@ def embed_audio(self, mel):
     def logits(self, tokens, audio_features):
         return self.decoder(tokens, audio_features)[0]
 
-    def forward_with_cross_qk(self, mel, tokens):
-        logits, _, cross_qk = self.decoder(tokens, self.encoder(mel))
+    def forward_with_cross_qk(self, mel, tokens, audio_features=None):
+        if audio_features is None:
+            audio_features = self.encoder(mel)
+        logits, _, cross_qk = self.decoder(tokens, audio_features)
         return logits, cross_qk
 
     def __call__(self, mel, tokens):