Allow to customize the whisper model

app.py

@@ -41,10 +41,14 @@ def main():
         batch_size = os.environ.get('PYTORCH_BATCH_SIZE') or st.selectbox(
             'Select a batch size:', [4, 8, 16, 32, 64]
         )
+        whisper_model = os.environ.get('WHISPER_MODEL') or st.selectbox(
+            'Select a Whisper model:', ['large-v2', 'base']
+        )
     else:
         device = None
         compute_type = None
         batch_size = None
+        whisper_model = None
 
     engine_api_key = os.environ.get(
         f'{engine_type.upper()}_API_KEY'
@@ -69,7 +73,14 @@ def main():
         if uploaded_audio:
             if openai_api_key:
                 st.markdown('Transcribing the audio...')
-                engine = get_engine(engine_type, api_key=engine_api_key, device=device, compute_type=compute_type, batch_size=batch_size)
+                engine = get_engine(
+                    engine_type,
+                    api_key=engine_api_key,
+                    device=device,
+                    compute_type=compute_type,
+                    batch_size=batch_size,
+                    whisper_model=whisper_model,
+                )
                 transcription = api.transcribe(engine, language, uploaded_audio)
 
                 st.markdown(

engines.py

@@ -57,12 +57,12 @@ class AssemblyAI:
 
 
 class WhisperX:
-    def __init__(self, api_key: str, device: str = 'cuda', compute_type: str = 'int8', batch_size: int = 8):
+    def __init__(self, api_key: str, device: str = 'cuda', compute_type: str = 'int8', batch_size: int = 8, whisper_model: str = 'large-v2', **kwargs: Any):
         self.api_key = api_key  # HuggingFace API key
         self.device = device
         self.compute_type = compute_type
         self.batch_size = batch_size
-        _setup_whisperx(self.device, self.compute_type)
+        _setup_whisperx(self.device, self.compute_type, whisper_model=whisper_model)
 
     def transcribe(self, language, audio_file: BytesIO) -> str:
         global _whisperx_model
@@ -113,7 +113,7 @@ _whisperx_model = None
 _whisperx_model_a = None
 _whisperx_model_a_metadata = None
 
-def _setup_whisperx(device, compute_type):
+def _setup_whisperx(device, compute_type, whisper_model='large-v2'):
     global _whisperx_initialized, _whisperx_model, _whisperx_model_a, _whisperx_model_a_metadata
     if _whisperx_initialized:
         return
@@ -123,4 +123,4 @@ def _setup_whisperx(device, compute_type):
         dev = torch.device(device)
         torch.nn.functional.conv2d(torch.zeros(s, s, s, s, device=dev), torch.zeros(s, s, s, s, device=dev))
 
-    _whisperx_model = whisperx.load_model('large-v2', device, compute_type=compute_type)
+    _whisperx_model = whisperx.load_model(whisper_model, device, compute_type=compute_type)