Merge pull request #154 from jhj0517/feature/add-local-model-path

app.py

@@ -17,8 +17,10 @@ class App:
         self.app = gr.Blocks(css=CSS, theme=self.args.theme)
         self.whisper_inf = WhisperInference() if self.args.disable_faster_whisper else FasterWhisperInference()
         if isinstance(self.whisper_inf, FasterWhisperInference):
+            self.whisper_inf.model_dir = args.faster_whisper_model_dir
             print("Use Faster Whisper implementation")
         else:
+            self.whisper_inf.model_dir = args.whisper_model_dir
             print("Use Open AI Whisper implementation")
         print(f"Device \"{self.whisper_inf.device}\" is detected")
         self.nllb_inf = NLLBInference()
@@ -296,6 +298,8 @@ parser.add_argument('--password', type=str, default=None, help='Gradio authentic
 parser.add_argument('--theme', type=str, default=None, help='Gradio Blocks theme')
 parser.add_argument('--colab', type=bool, default=False, nargs='?', const=True, help='Is colab user or not')
 parser.add_argument('--api_open', type=bool, default=False, nargs='?', const=True, help='enable api or not')
+parser.add_argument('--whisper_model_dir', type=str, default=os.path.join("models", "Whisper"), help='Directory path of the whisper model')
+parser.add_argument('--faster_whisper_model_dir', type=str, default=os.path.join("models", "Whisper", "faster-whisper"), help='Directory path of the faster-whisper model')
 _args = parser.parse_args()
 
 if __name__ == "__main__":

modules/faster_whisper_inference.py

@@ -32,7 +32,7 @@ class FasterWhisperInference(BaseInterface):
         self.available_compute_types = ctranslate2.get_supported_compute_types(
             "cuda") if self.device == "cuda" else ctranslate2.get_supported_compute_types("cpu")
         self.current_compute_type = "float16" if self.device == "cuda" else "float32"
-        self.default_beam_size = 1
+        self.model_dir = os.path.join("models", "Whisper", "faster-whisper")
 
     def transcribe_file(self,
                         files: list,
@@ -311,7 +311,7 @@ class FasterWhisperInference(BaseInterface):
         self.model = faster_whisper.WhisperModel(
             device=self.device,
             model_size_or_path=model_size,
-            download_root=os.path.join("models", "Whisper", "faster-whisper"),
+            download_root=self.model_dir,
             compute_type=self.current_compute_type
         )
 

modules/whisper_Inference.py

@@ -26,7 +26,7 @@ class WhisperInference(BaseInterface):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.available_compute_types = ["float16", "float32"]
         self.current_compute_type = "float16" if self.device == "cuda" else "float32"
-        self.default_beam_size = 1
+        self.model_dir = os.path.join("models", "Whisper")
 
     def transcribe_file(self,
                         files: list,
@@ -288,7 +288,7 @@ class WhisperInference(BaseInterface):
         self.model = whisper.load_model(
             name=model_size,
             device=self.device,
-            download_root=os.path.join("models", "Whisper")
+            download_root=self.model_dir
         )
 
     @staticmethod

user-start-webui.bat

@@ -10,9 +10,10 @@ set SHARE=
 set THEME=
 set DISABLE_FASTER_WHISPER=
 set API_OPEN=
+set WHISPER_MODEL_DIR=
+set FASTER_WHISPER_MODEL_DIR=
 
 
-:: Set args accordingly
 if not "%SERVER_NAME%"=="" (
     set SERVER_NAME_ARG=--server_name %SERVER_NAME%
 )
@@ -37,7 +38,13 @@ if /I "%DISABLE_FASTER_WHISPER%"=="true" (
 if /I "%API_OPEN%"=="true" (
     set API_OPEN=--api_open
 )
+if not "%WHISPER_MODEL_DIR%"=="" (
+    set WHISPER_MODEL_DIR_ARG=--whisper_model_dir "%WHISPER_MODEL_DIR%"
+)
+if not "%FASTER_WHISPER_MODEL_DIR%"=="" (
+    set FASTER_WHISPER_MODEL_DIR_ARG=--faster_whisper_model_dir "%FASTER_WHISPER_MODEL_DIR%"
+)
 
 :: Call the original .bat script with optional arguments
-start-webui.bat %SERVER_NAME_ARG% %SERVER_PORT_ARG% %USERNAME_ARG% %PASSWORD_ARG% %SHARE_ARG% %THEME_ARG% %DISABLE_FASTER_WHISPER_ARG% %API_OPEN%
+start-webui.bat %SERVER_NAME_ARG% %SERVER_PORT_ARG% %USERNAME_ARG% %PASSWORD_ARG% %SHARE_ARG% %THEME_ARG% %DISABLE_FASTER_WHISPER_ARG% %API_OPEN% %WHISPER_MODEL_DIR_ARG% %FASTER_WHISPER_MODEL_DIR_ARG%
 pause