jhj0517 commited on
Commit
d6e2447
·
1 Parent(s): 6960e0b

remove arg and add dir initialization

Browse files
Files changed (1) hide show
  1. 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
- output_dir: str,
25
- args: Namespace
26
  ):
27
- self.model = None
28
- self.current_model_size = None
 
 
 
 
 
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,