jhj0517 commited on
Commit
e48211d
·
1 Parent(s): 48382b6

add WhisperFactory class

Browse files
Files changed (1) hide show
  1. modules/whisper/whisper_factory.py +60 -0
modules/whisper/whisper_factory.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from argparse import Namespace
2
+ import os
3
+
4
+ from modules.whisper.faster_whisper_inference import FasterWhisperInference
5
+ from modules.whisper.whisper_Inference import WhisperInference
6
+ from modules.whisper.insanely_fast_whisper_inference import InsanelyFastWhisperInference
7
+ from modules.whisper.whisper_base import WhisperBase
8
+
9
+
10
+ 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
+ -------
38
+ WhisperBase
39
+ An instance of the appropriate whisper inference class based on the whisper_type.
40
+ """
41
+ # Temporal fix of the bug : https://github.com/jhj0517/Whisper-WebUI/issues/144
42
+ os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
43
+
44
+ whisper_type = whisper_type.lower().strip()
45
+
46
+ faster_whisper_typos = ["faster_whisper", "faster-whisper", "fasterwhisper"]
47
+ whisper_typos = ["whisper"]
48
+ insanely_fast_whisper_typos = [
49
+ "insanely_fast_whisper", "insanely-fast-whisper", "insanelyfastwhisper",
50
+ "insanely_faster_whisper", "insanely-faster-whisper", "insanelyfasterwhisper"
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)