Spaces:
Running
Running
jhj0517
commited on
Commit
·
d6e2447
1
Parent(s):
6960e0b
remove arg and add dir initialization
Browse files- modules/whisper/whisper_base.py +19 -10
modules/whisper/whisper_base.py
CHANGED
@@ -6,7 +6,6 @@ from abc import ABC, abstractmethod
|
|
6 |
from typing import BinaryIO, Union, Tuple, List
|
7 |
import numpy as np
|
8 |
from datetime import datetime
|
9 |
-
from argparse import Namespace
|
10 |
from faster_whisper.vad import VadOptions
|
11 |
from dataclasses import astuple
|
12 |
|
@@ -20,26 +19,34 @@ from modules.vad.silero_vad import SileroVAD
|
|
20 |
|
21 |
class WhisperBase(ABC):
|
22 |
def __init__(self,
|
23 |
-
model_dir: str,
|
24 |
-
|
25 |
-
|
26 |
):
|
27 |
-
|
28 |
-
|
|
|
|
|
|
|
|
|
|
|
29 |
self.model_dir = model_dir
|
30 |
self.output_dir = output_dir
|
31 |
os.makedirs(self.output_dir, exist_ok=True)
|
32 |
os.makedirs(self.model_dir, exist_ok=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
33 |
self.available_models = whisper.available_models()
|
34 |
self.available_langs = sorted(list(whisper.tokenizer.LANGUAGES.values()))
|
35 |
self.translatable_models = ["large", "large-v1", "large-v2", "large-v3"]
|
36 |
self.device = self.get_device()
|
37 |
self.available_compute_types = ["float16", "float32"]
|
38 |
self.current_compute_type = "float16" if self.device == "cuda" else "float32"
|
39 |
-
self.diarizer = Diarizer(
|
40 |
-
model_dir=args.diarization_model_dir
|
41 |
-
)
|
42 |
-
self.vad = SileroVAD()
|
43 |
|
44 |
@abstractmethod
|
45 |
def transcribe(self,
|
@@ -47,6 +54,7 @@ class WhisperBase(ABC):
|
|
47 |
progress: gr.Progress,
|
48 |
*whisper_params,
|
49 |
):
|
|
|
50 |
pass
|
51 |
|
52 |
@abstractmethod
|
@@ -55,6 +63,7 @@ class WhisperBase(ABC):
|
|
55 |
compute_type: str,
|
56 |
progress: gr.Progress
|
57 |
):
|
|
|
58 |
pass
|
59 |
|
60 |
def run(self,
|
|
|
6 |
from typing import BinaryIO, Union, Tuple, List
|
7 |
import numpy as np
|
8 |
from datetime import datetime
|
|
|
9 |
from faster_whisper.vad import VadOptions
|
10 |
from dataclasses import astuple
|
11 |
|
|
|
19 |
|
20 |
class WhisperBase(ABC):
|
21 |
def __init__(self,
|
22 |
+
model_dir: Optional[str] = None,
|
23 |
+
diarization_model_dir: Optional[str] = None,
|
24 |
+
output_dir: Optional[str] = None,
|
25 |
):
|
26 |
+
if model_dir is None:
|
27 |
+
model_dir = os.path.join("models", "Whisper")
|
28 |
+
if diarization_model_dir is None:
|
29 |
+
diarization_model_dir = os.path.join("models", "Diarization")
|
30 |
+
if output_dir is None:
|
31 |
+
output_dir = os.path.join("outputs")
|
32 |
+
|
33 |
self.model_dir = model_dir
|
34 |
self.output_dir = output_dir
|
35 |
os.makedirs(self.output_dir, exist_ok=True)
|
36 |
os.makedirs(self.model_dir, exist_ok=True)
|
37 |
+
self.diarizer = Diarizer(
|
38 |
+
model_dir=diarization_model_dir
|
39 |
+
)
|
40 |
+
self.vad = SileroVAD()
|
41 |
+
|
42 |
+
self.model = None
|
43 |
+
self.current_model_size = None
|
44 |
self.available_models = whisper.available_models()
|
45 |
self.available_langs = sorted(list(whisper.tokenizer.LANGUAGES.values()))
|
46 |
self.translatable_models = ["large", "large-v1", "large-v2", "large-v3"]
|
47 |
self.device = self.get_device()
|
48 |
self.available_compute_types = ["float16", "float32"]
|
49 |
self.current_compute_type = "float16" if self.device == "cuda" else "float32"
|
|
|
|
|
|
|
|
|
50 |
|
51 |
@abstractmethod
|
52 |
def transcribe(self,
|
|
|
54 |
progress: gr.Progress,
|
55 |
*whisper_params,
|
56 |
):
|
57 |
+
"""Inference whisper model to transcribe"""
|
58 |
pass
|
59 |
|
60 |
@abstractmethod
|
|
|
63 |
compute_type: str,
|
64 |
progress: gr.Progress
|
65 |
):
|
66 |
+
"""Initialize whisper model"""
|
67 |
pass
|
68 |
|
69 |
def run(self,
|