Make batch_size configurable for WhisperX

app.py

@@ -38,9 +38,13 @@ def main():
         compute_type = os.environ.get('PYTORCH_COMPUTE_TYPE') or st.selectbox(
             'Select a compute type:', ['int8', 'float16']
         )
+        batch_size = os.environ.get('PYTORCH_BATCH_SIZE') or st.selectbox(
+            'Select a batch size:', [4, 8, 16, 32, 64]
+        )
     else:
         device = None
         compute_type = None
+        batch_size = None
 
     engine_api_key = os.environ.get(
         f'{engine_type.upper()}_API_KEY'
@@ -62,7 +66,7 @@ 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)
+                engine = get_engine(engine_type, api_key=engine_api_key, device=device, compute_type=compute_type, batch_size=batch_size)
                 transcription = api.transcribe(engine, language, uploaded_audio)
 
                 st.markdown(

engines.py

@@ -20,7 +20,7 @@ class AssemblyAI:
     transcript = 'https://api.assemblyai.com/v2/transcript'
     upload = 'https://api.assemblyai.com/v2/upload'
 
-    def __init__(self, api_key: str):
+    def __init__(self, api_key: str, **kwargs: Any):
         self.api_key = api_key
 
     def transcribe(self, language, audio_file: BytesIO) -> str: