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

add dir parameters

Browse files
Files changed (2) hide show
  1. app.py +3 -2
  2. modules/whisper/whisper_factory.py +41 -20
app.py CHANGED
@@ -18,9 +18,10 @@ class App:
18
  self.app = gr.Blocks(css=CSS, theme=self.args.theme)
19
  self.whisper_inf = WhisperFactory.create_whisper_inference(
20
  whisper_type=self.args.whisper_type,
21
- model_dir=self.args.faster_whisper_model_dir,
 
 
22
  output_dir=self.args.output_dir,
23
- args=self.args
24
  )
25
  print(f"Use \"{self.args.whisper_type}\" implementation")
26
  print(f"Device \"{self.whisper_inf.device}\" is detected")
 
18
  self.app = gr.Blocks(css=CSS, theme=self.args.theme)
19
  self.whisper_inf = WhisperFactory.create_whisper_inference(
20
  whisper_type=self.args.whisper_type,
21
+ whisper_model_dir=self.args.whisper_model_dir,
22
+ faster_whisper_model_dir=self.args.faster_whisper_model_dir,
23
+ insanely_fast_whisper_model_dir=self.args.insanely_fast_whisper_model_dir,
24
  output_dir=self.args.output_dir,
 
25
  )
26
  print(f"Use \"{self.args.whisper_type}\" implementation")
27
  print(f"Device \"{self.whisper_inf.device}\" is detected")
modules/whisper/whisper_factory.py CHANGED
@@ -1,4 +1,4 @@
1
- from argparse import Namespace
2
  import os
3
 
4
  from modules.whisper.faster_whisper_inference import FasterWhisperInference
@@ -11,27 +11,32 @@ class WhisperFactory:
11
  @staticmethod
12
  def create_whisper_inference(
13
  whisper_type: str,
14
- model_dir: str,
15
- output_dir: str,
16
- args: Namespace
 
 
17
  ) -> "WhisperBase":
18
  """
19
  Create a whisper inference class based on the provided whisper_type.
20
 
21
  Parameters
22
  ----------
23
- whisper_type: str
24
- The repository name of whisper inference to use. Supported values are:
25
- - "faster-whisper" from
26
- - "whisper"
27
- - insanely-fast-whisper", "insanely_fast_whisper", "insanelyfastwhisper",
28
- "insanely-faster-whisper", "insanely_faster_whisper", "insanelyfasterwhisper"
29
- model_dir: str
30
- The directory path where the whisper model is located.
31
- output_dir: str
32
- The directory path where the output files will be saved.
33
- args: Any
34
- Additional arguments to be passed to the whisper inference object.
 
 
 
35
 
36
  Returns
37
  -------
@@ -51,10 +56,26 @@ class WhisperFactory:
51
  ]
52
 
53
  if whisper_type in faster_whisper_typos:
54
- return FasterWhisperInference(model_dir, output_dir, args)
 
 
 
 
55
  elif whisper_type in whisper_typos:
56
- return WhisperInference(model_dir, output_dir, args)
 
 
 
 
57
  elif whisper_type in insanely_fast_whisper_typos:
58
- return InsanelyFastWhisperInference(model_dir, output_dir, args)
 
 
 
 
59
  else:
60
- return FasterWhisperInference(model_dir, output_dir, args)
 
 
 
 
 
1
+ from typing import Optional
2
  import os
3
 
4
  from modules.whisper.faster_whisper_inference import FasterWhisperInference
 
11
  @staticmethod
12
  def create_whisper_inference(
13
  whisper_type: str,
14
+ whisper_model_dir: Optional[str] = None,
15
+ faster_whisper_model_dir: Optional[str] = None,
16
+ insanely_fast_whisper_model_dir: Optional[str] = None,
17
+ diarization_model_dir: Optional[str] = None,
18
+ output_dir: Optional[str] = None,
19
  ) -> "WhisperBase":
20
  """
21
  Create a whisper inference class based on the provided whisper_type.
22
 
23
  Parameters
24
  ----------
25
+ whisper_type : str
26
+ The type of Whisper implementation to use. Supported values (case-insensitive):
27
+ - "faster-whisper": https://github.com/openai/whisper
28
+ - "whisper": https://github.com/openai/whisper
29
+ - "insanely-fast-whisper": https://github.com/Vaibhavs10/insanely-fast-whisper
30
+ whisper_model_dir : str
31
+ Directory path for the Whisper model.
32
+ faster_whisper_model_dir : str
33
+ Directory path for the Faster Whisper model.
34
+ insanely_fast_whisper_model_dir : str
35
+ Directory path for the Insanely Fast Whisper model.
36
+ diarization_model_dir : str
37
+ Directory path for the diarization model.
38
+ output_dir : str
39
+ Directory path where output files will be saved.
40
 
41
  Returns
42
  -------
 
56
  ]
57
 
58
  if whisper_type in faster_whisper_typos:
59
+ return FasterWhisperInference(
60
+ model_dir=faster_whisper_model_dir,
61
+ output_dir=output_dir,
62
+ diarization_model_dir=diarization_model_dir
63
+ )
64
  elif whisper_type in whisper_typos:
65
+ return WhisperInference(
66
+ model_dir=whisper_model_dir,
67
+ output_dir=output_dir,
68
+ diarization_model_dir=diarization_model_dir
69
+ )
70
  elif whisper_type in insanely_fast_whisper_typos:
71
+ return InsanelyFastWhisperInference(
72
+ model_dir=insanely_fast_whisper_model_dir,
73
+ output_dir=output_dir,
74
+ diarization_model_dir=diarization_model_dir
75
+ )
76
  else:
77
+ return FasterWhisperInference(
78
+ model_dir=faster_whisper_model_dir,
79
+ output_dir=output_dir,
80
+ diarization_model_dir=diarization_model_dir
81
+ )