Adding support for Turbo model

chidiwilliams · Oct 19, 2024 · 373b0d6 · 373b0d6
1 parent 56d7835
commit 373b0d6
Show file tree

Hide file tree

Showing 6 changed files with 205 additions and 129 deletions.
diff --git a/buzz/model_loader.py b/buzz/model_loader.py
@@ -50,6 +50,7 @@ class WhisperModelSize(str, enum.Enum):
     LARGE = "large"
     LARGEV2 = "large-v2"
     LARGEV3 = "large-v3"
+    LARGEV3TURBO = "large-v3-turbo"
     CUSTOM = "custom"
 
     def to_faster_whisper_model_size(self) -> str:

diff --git a/buzz/transformers_whisper.py b/buzz/transformers_whisper.py
@@ -170,7 +170,7 @@ def transcribe(
             use_safetensors = len(safetensors_files) > 0
 
         model = AutoModelForSpeechSeq2Seq.from_pretrained(
-            self.model_id, torch_dtype=torch_dtype, use_safetensors=use_safetensors
+            self.model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=use_safetensors
         )
 
         model.generation_config.language = language

diff --git a/buzz/widgets/preferences_dialog/models_preferences_widget.py b/buzz/widgets/preferences_dialog/models_preferences_widget.py
@@ -181,6 +181,11 @@ def reset(self):
                     model_size == WhisperModelSize.CUSTOM):
                 continue
 
+            # Skip turbo model for Faste Whisper and Whisper.cpp
+            if (self.model.model_type in {ModelType.FASTER_WHISPER, ModelType.WHISPER_CPP} and
+                    model_size == WhisperModelSize.LARGEV3TURBO):
+                continue
+
             model = TranscriptionModel(
                 model_type=self.model.model_type,
                 whisper_model_size=WhisperModelSize(model_size),

diff --git a/buzz/widgets/transcriber/transcription_options_group_box.py b/buzz/widgets/transcriber/transcription_options_group_box.py
@@ -51,7 +51,7 @@ def __init__(
 
         self.whisper_model_size_combo_box = QComboBox(self)
         self.whisper_model_size_combo_box.addItems(
-            [size.value.title() for size in WhisperModelSize if size != WhisperModelSize.CUSTOM]
+            [size.value.title() for size in WhisperModelSize if size not in {WhisperModelSize.CUSTOM, WhisperModelSize.LARGEV3TURBO}]
         )
         self.whisper_model_size_combo_box.currentTextChanged.connect(
             self.on_whisper_model_size_changed
@@ -141,6 +141,10 @@ def reset_visible_rows(self):
                 and whisper_model_size == WhisperModelSize.CUSTOM),
             )
 
+        if model_type == ModelType.WHISPER:
+            if self.whisper_model_size_combo_box.findText(WhisperModelSize.LARGEV3TURBO.value.title()) == -1:
+                self.whisper_model_size_combo_box.addItem(WhisperModelSize.LARGEV3TURBO.value.title())
+
         custom_model_index = (self.whisper_model_size_combo_box
                               .findText(WhisperModelSize.CUSTOM.value.title()))
         if (model_type == ModelType.WHISPER