Merge pull request #219 from jhj0517/refactor/args

app.py

@@ -18,9 +18,10 @@ class App:
         self.app = gr.Blocks(css=CSS, theme=self.args.theme)
         self.whisper_inf = WhisperFactory.create_whisper_inference(
             whisper_type=self.args.whisper_type,
-            model_dir=self.args.faster_whisper_model_dir,
+            whisper_model_dir=self.args.whisper_model_dir,
+            faster_whisper_model_dir=self.args.faster_whisper_model_dir,
+            insanely_fast_whisper_model_dir=self.args.insanely_fast_whisper_model_dir,
             output_dir=self.args.output_dir,
-            args=self.args
         )
         print(f"Use \"{self.args.whisper_type}\" implementation")
         print(f"Device \"{self.whisper_inf.device}\" is detected")

modules/whisper/faster_whisper_inference.py

@@ -17,15 +17,20 @@ from modules.whisper.whisper_base import WhisperBase
 
 class FasterWhisperInference(WhisperBase):
     def __init__(self,
-                 model_dir: str,
-                 output_dir: str,
-                 args: Namespace
+                 model_dir: Optional[str] = None,
+                 diarization_model_dir: Optional[str] = None,
+                 output_dir: Optional[str] = None,
                  ):
         super().__init__(
             model_dir=model_dir,
-            output_dir=output_dir,
-            args=args
+            diarization_model_dir=diarization_model_dir,
+            output_dir=output_dir
         )
+        if model_dir is None:
+            model_dir = os.path.join("models", "Whisper", "faster-whisper")
+        self.model_dir = model_dir
+        os.makedirs(self.model_dir, exist_ok=True)
+
         self.model_paths = self.get_model_paths()
         self.device = self.get_device()
         self.available_models = self.model_paths.keys()

modules/whisper/insanely_fast_whisper_inference.py

@@ -17,15 +17,20 @@ from modules.whisper.whisper_base import WhisperBase
 
 class InsanelyFastWhisperInference(WhisperBase):
     def __init__(self,
-                 model_dir: str,
-                 output_dir: str,
-                 args: Namespace
+                 model_dir: Optional[str] = None,
+                 diarization_model_dir: Optional[str] = None,
+                 output_dir: Optional[str] = None,
                  ):
         super().__init__(
             model_dir=model_dir,
             output_dir=output_dir,
-            args=args
+            diarization_model_dir=diarization_model_dir
         )
+        if model_dir is None:
+            model_dir = os.path.join("models", "Whisper", "insanely-fast-whisper")
+        self.model_dir = model_dir
+        os.makedirs(self.model_dir, exist_ok=True)
+
         openai_models = whisper.available_models()
         distil_models = ["distil-large-v2", "distil-large-v3", "distil-medium.en", "distil-small.en"]
         self.available_models = openai_models + distil_models

modules/whisper/whisper_Inference.py

@@ -12,14 +12,14 @@ from modules.whisper.whisper_parameter import *
 
 class WhisperInference(WhisperBase):
     def __init__(self,
-                 model_dir: str,
-                 output_dir: str,
-                 args: Namespace
+                 model_dir: Optional[str] = None,
+                 diarization_model_dir: Optional[str] = None,
+                 output_dir: Optional[str] = None,
                  ):
         super().__init__(
             model_dir=model_dir,
             output_dir=output_dir,
-            args=args
+            diarization_model_dir=diarization_model_dir
         )
 
     def transcribe(self,

modules/whisper/whisper_base.py

@@ -6,7 +6,6 @@ from abc import ABC, abstractmethod
 from typing import BinaryIO, Union, Tuple, List
 import numpy as np
 from datetime import datetime
-from argparse import Namespace
 from faster_whisper.vad import VadOptions
 from dataclasses import astuple
 
@@ -20,26 +19,34 @@ from modules.vad.silero_vad import SileroVAD
 
 class WhisperBase(ABC):
     def __init__(self,
-                 model_dir: str,
-                 output_dir: str,
-                 args: Namespace
+                 model_dir: Optional[str] = None,
+                 diarization_model_dir: Optional[str] = None,
+                 output_dir: Optional[str] = None,
                  ):
-        self.model = None
-        self.current_model_size = None
+        if model_dir is None:
+            model_dir = os.path.join("models", "Whisper")
+        if diarization_model_dir is None:
+            diarization_model_dir = os.path.join("models", "Diarization")
+        if output_dir is None:
+            output_dir = os.path.join("outputs")
+
         self.model_dir = model_dir
         self.output_dir = output_dir
         os.makedirs(self.output_dir, exist_ok=True)
         os.makedirs(self.model_dir, exist_ok=True)
+        self.diarizer = Diarizer(
+            model_dir=diarization_model_dir
+        )
+        self.vad = SileroVAD()
+
+        self.model = None
+        self.current_model_size = None
         self.available_models = whisper.available_models()
         self.available_langs = sorted(list(whisper.tokenizer.LANGUAGES.values()))
         self.translatable_models = ["large", "large-v1", "large-v2", "large-v3"]
         self.device = self.get_device()
         self.available_compute_types = ["float16", "float32"]
         self.current_compute_type = "float16" if self.device == "cuda" else "float32"
-        self.diarizer = Diarizer(
-            model_dir=args.diarization_model_dir
-        )
-        self.vad = SileroVAD()
 
     @abstractmethod
     def transcribe(self,
@@ -47,6 +54,7 @@ class WhisperBase(ABC):
                    progress: gr.Progress,
                    *whisper_params,
                    ):
+        """Inference whisper model to transcribe"""
         pass
 
     @abstractmethod
@@ -55,6 +63,7 @@ class WhisperBase(ABC):
                      compute_type: str,
                      progress: gr.Progress
                      ):
+        """Initialize whisper model"""
         pass
 
     def run(self,

modules/whisper/whisper_factory.py

@@ -1,4 +1,4 @@
-from argparse import Namespace
+from typing import Optional
 import os
 
 from modules.whisper.faster_whisper_inference import FasterWhisperInference
@@ -11,27 +11,32 @@ class WhisperFactory:
     @staticmethod
     def create_whisper_inference(
         whisper_type: str,
-        model_dir: str,
-        output_dir: str,
-        args: Namespace
+        whisper_model_dir: Optional[str] = None,
+        faster_whisper_model_dir: Optional[str] = None,
+        insanely_fast_whisper_model_dir: Optional[str] = None,
+        diarization_model_dir: Optional[str] = None,
+        output_dir: Optional[str] = None,
     ) -> "WhisperBase":
         """
         Create a whisper inference class based on the provided whisper_type.
 
         Parameters
         ----------
-        whisper_type: str
-            The repository name of whisper inference to use. Supported values are:
-            - "faster-whisper" from
-            - "whisper"
-            - insanely-fast-whisper", "insanely_fast_whisper", "insanelyfastwhisper",
-              "insanely-faster-whisper", "insanely_faster_whisper", "insanelyfasterwhisper"
-        model_dir: str
-            The directory path where the whisper model is located.
-        output_dir: str
-            The directory path where the output files will be saved.
-        args: Any
-            Additional arguments to be passed to the whisper inference object.
+        whisper_type : str
+            The type of Whisper implementation to use. Supported values (case-insensitive):
+            - "faster-whisper": https://github.com/openai/whisper
+            - "whisper": https://github.com/openai/whisper
+            - "insanely-fast-whisper": https://github.com/Vaibhavs10/insanely-fast-whisper
+        whisper_model_dir : str
+            Directory path for the Whisper model.
+        faster_whisper_model_dir : str
+            Directory path for the Faster Whisper model.
+        insanely_fast_whisper_model_dir : str
+            Directory path for the Insanely Fast Whisper model.
+        diarization_model_dir : str
+            Directory path for the diarization model.
+        output_dir : str
+            Directory path where output files will be saved.
 
         Returns
         -------
@@ -51,10 +56,26 @@ class WhisperFactory:
         ]
 
         if whisper_type in faster_whisper_typos:
-            return FasterWhisperInference(model_dir, output_dir, args)
+            return FasterWhisperInference(
+                model_dir=faster_whisper_model_dir,
+                output_dir=output_dir,
+                diarization_model_dir=diarization_model_dir
+            )
         elif whisper_type in whisper_typos:
-            return WhisperInference(model_dir, output_dir, args)
+            return WhisperInference(
+                model_dir=whisper_model_dir,
+                output_dir=output_dir,
+                diarization_model_dir=diarization_model_dir
+            )
         elif whisper_type in insanely_fast_whisper_typos:
-            return InsanelyFastWhisperInference(model_dir, output_dir, args)
+            return InsanelyFastWhisperInference(
+                model_dir=insanely_fast_whisper_model_dir,
+                output_dir=output_dir,
+                diarization_model_dir=diarization_model_dir
+            )
         else:
-            return FasterWhisperInference(model_dir, output_dir, args)
+            return FasterWhisperInference(
+                model_dir=faster_whisper_model_dir,
+                output_dir=output_dir,
+                diarization_model_dir=diarization_model_dir
+            )