jhj0517 commited on
Commit
0a34091
·
1 Parent(s): 2e08651

add dir initialization

Browse files
modules/whisper/faster_whisper_inference.py CHANGED
@@ -17,15 +17,20 @@ from modules.whisper.whisper_base import WhisperBase
17
 
18
  class FasterWhisperInference(WhisperBase):
19
  def __init__(self,
20
- model_dir: str,
21
- output_dir: str,
22
- args: Namespace
23
  ):
24
  super().__init__(
25
  model_dir=model_dir,
26
- output_dir=output_dir,
27
- args=args
28
  )
 
 
 
 
 
29
  self.model_paths = self.get_model_paths()
30
  self.device = self.get_device()
31
  self.available_models = self.model_paths.keys()
 
17
 
18
  class FasterWhisperInference(WhisperBase):
19
  def __init__(self,
20
+ model_dir: Optional[str] = None,
21
+ diarization_model_dir: Optional[str] = None,
22
+ output_dir: Optional[str] = None,
23
  ):
24
  super().__init__(
25
  model_dir=model_dir,
26
+ diarization_model_dir=diarization_model_dir,
27
+ output_dir=output_dir
28
  )
29
+ if model_dir is None:
30
+ model_dir = os.path.join("models", "Whisper", "faster-whisper")
31
+ self.model_dir = model_dir
32
+ os.makedirs(self.model_dir, exist_ok=True)
33
+
34
  self.model_paths = self.get_model_paths()
35
  self.device = self.get_device()
36
  self.available_models = self.model_paths.keys()
modules/whisper/insanely_fast_whisper_inference.py CHANGED
@@ -17,15 +17,20 @@ from modules.whisper.whisper_base import WhisperBase
17
 
18
  class InsanelyFastWhisperInference(WhisperBase):
19
  def __init__(self,
20
- model_dir: str,
21
- output_dir: str,
22
- args: Namespace
23
  ):
24
  super().__init__(
25
  model_dir=model_dir,
26
  output_dir=output_dir,
27
- args=args
28
  )
 
 
 
 
 
29
  openai_models = whisper.available_models()
30
  distil_models = ["distil-large-v2", "distil-large-v3", "distil-medium.en", "distil-small.en"]
31
  self.available_models = openai_models + distil_models
 
17
 
18
  class InsanelyFastWhisperInference(WhisperBase):
19
  def __init__(self,
20
+ model_dir: Optional[str] = None,
21
+ diarization_model_dir: Optional[str] = None,
22
+ output_dir: Optional[str] = None,
23
  ):
24
  super().__init__(
25
  model_dir=model_dir,
26
  output_dir=output_dir,
27
+ diarization_model_dir=diarization_model_dir
28
  )
29
+ if model_dir is None:
30
+ model_dir = os.path.join("models", "Whisper", "insanely-fast-whisper")
31
+ self.model_dir = model_dir
32
+ os.makedirs(self.model_dir, exist_ok=True)
33
+
34
  openai_models = whisper.available_models()
35
  distil_models = ["distil-large-v2", "distil-large-v3", "distil-medium.en", "distil-small.en"]
36
  self.available_models = openai_models + distil_models
modules/whisper/whisper_Inference.py CHANGED
@@ -12,14 +12,14 @@ from modules.whisper.whisper_parameter import *
12
 
13
  class WhisperInference(WhisperBase):
14
  def __init__(self,
15
- model_dir: str,
16
- output_dir: str,
17
- args: Namespace
18
  ):
19
  super().__init__(
20
  model_dir=model_dir,
21
  output_dir=output_dir,
22
- args=args
23
  )
24
 
25
  def transcribe(self,
 
12
 
13
  class WhisperInference(WhisperBase):
14
  def __init__(self,
15
+ model_dir: Optional[str] = None,
16
+ diarization_model_dir: Optional[str] = None,
17
+ output_dir: Optional[str] = None,
18
  ):
19
  super().__init__(
20
  model_dir=model_dir,
21
  output_dir=output_dir,
22
+ diarization_model_dir=diarization_model_dir
23
  )
24
 
25
  def transcribe(self,