fix: execstack

chidiwilliams · Jan 5, 2024 · 23e886d · 23e886d
1 parent 22e788b
commit 23e886d
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 7 deletions.
diff --git a/buzz/transformers_whisper.py b/buzz/transformers_whisper.py
@@ -1,10 +1,12 @@
+import sys
 from typing import Optional, Union
 
 import numpy as np
 from tqdm import tqdm
-from transformers import WhisperProcessor, WhisperForConditionalGeneration
 
-from buzz import whisper_audio
+if sys.platform != "linux":
+    import whisper
+    from transformers import WhisperProcessor, WhisperForConditionalGeneration
 
 
 def load_model(model_name_or_path: str):
@@ -14,8 +16,8 @@ def load_model(model_name_or_path: str):
 
 
 class TransformersWhisper:
-    SAMPLE_RATE = whisper_audio.SAMPLE_RATE
-    N_SAMPLES_IN_CHUNK = whisper_audio.N_SAMPLES
+    SAMPLE_RATE = whisper.audio.SAMPLE_RATE
+    N_SAMPLES_IN_CHUNK = whisper.audio.N_SAMPLES
 
     def __init__(
         self, processor: WhisperProcessor, model: WhisperForConditionalGeneration
@@ -34,7 +36,7 @@ def transcribe(
         verbose: Optional[bool] = None,
     ):
         if isinstance(audio, str):
-            audio = whisper_audio.load_audio(audio, sr=self.SAMPLE_RATE)
+            audio = whisper.load_audio(audio, sr=self.SAMPLE_RATE)
 
         self.model.config.forced_decoder_ids = self.processor.get_decoder_prompt_ids(
             task=task, language=language

diff --git a/buzz/widgets/model_type_combo_box.py b/buzz/widgets/model_type_combo_box.py
@@ -31,8 +31,15 @@ def __init__(
             ) or (
                 # Disable Whisper and Faster Whisper options
                 # on Linux due to execstack errors on Snap
-                model_type in (ModelType.WHISPER, ModelType.FASTER_WHISPER)
-                and sys.platform == "Linux"
+                (
+                    model_type
+                    in (
+                        ModelType.WHISPER,
+                        ModelType.FASTER_WHISPER,
+                        ModelType.HUGGING_FACE,
+                    )
+                )
+                and sys.platform == "linux"
             ):
                 continue
             self.addItem(model_type.value)

diff --git a/tests/transformers_whisper_test.py b/tests/transformers_whisper_test.py
@@ -1,6 +1,11 @@
+import sys
+
+import pytest
+
 from buzz.transformers_whisper import load_model
 
 
+@pytest.mark.skipif(sys.platform == "linux", reason="Not supported on Linux")
 class TestTransformersWhisper:
     def test_should_transcribe(self):
         model = load_model("openai/whisper-tiny")