jhj0517 commited on
Commit
6db9d8d
·
1 Parent(s): 988860f

Use enum for string

Browse files
Files changed (2) hide show
  1. app.py +1 -1
  2. modules/whisper/whisper_factory.py +5 -3
app.py CHANGED
@@ -320,7 +320,7 @@ class App:
320
 
321
 
322
  parser = argparse.ArgumentParser()
323
- parser.add_argument('--whisper_type', type=str, default=WhisperImpl.FASTER_WHISPER,
324
  choices=[item.value for item in WhisperImpl],
325
  help='A type of the whisper implementation (Github repo name)')
326
  parser.add_argument('--share', type=str2bool, default=False, nargs='?', const=True, help='Gradio share value')
 
320
 
321
 
322
  parser = argparse.ArgumentParser()
323
+ parser.add_argument('--whisper_type', type=str, default=WhisperImpl.FASTER_WHISPER.value,
324
  choices=[item.value for item in WhisperImpl],
325
  help='A type of the whisper implementation (Github repo name)')
326
  parser.add_argument('--share', type=str2bool, default=False, nargs='?', const=True, help='Gradio share value')
modules/whisper/whisper_factory.py CHANGED
@@ -52,21 +52,23 @@ class WhisperFactory:
52
  # Temporal fix of the bug : https://github.com/jhj0517/Whisper-WebUI/issues/144
53
  os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
54
 
55
- if whisper_type == WhisperImpl.FASTER_WHISPER:
 
 
56
  return FasterWhisperInference(
57
  model_dir=faster_whisper_model_dir,
58
  output_dir=output_dir,
59
  diarization_model_dir=diarization_model_dir,
60
  uvr_model_dir=uvr_model_dir
61
  )
62
- elif whisper_type == WhisperImpl.WHISPER:
63
  return WhisperInference(
64
  model_dir=whisper_model_dir,
65
  output_dir=output_dir,
66
  diarization_model_dir=diarization_model_dir,
67
  uvr_model_dir=uvr_model_dir
68
  )
69
- elif whisper_type == WhisperImpl.INSANELY_FAST_WHISPER:
70
  return InsanelyFastWhisperInference(
71
  model_dir=insanely_fast_whisper_model_dir,
72
  output_dir=output_dir,
 
52
  # Temporal fix of the bug : https://github.com/jhj0517/Whisper-WebUI/issues/144
53
  os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
54
 
55
+ whisper_type = whisper_type.strip().lower()
56
+
57
+ if whisper_type == WhisperImpl.FASTER_WHISPER.value:
58
  return FasterWhisperInference(
59
  model_dir=faster_whisper_model_dir,
60
  output_dir=output_dir,
61
  diarization_model_dir=diarization_model_dir,
62
  uvr_model_dir=uvr_model_dir
63
  )
64
+ elif whisper_type == WhisperImpl.WHISPER.value:
65
  return WhisperInference(
66
  model_dir=whisper_model_dir,
67
  output_dir=output_dir,
68
  diarization_model_dir=diarization_model_dir,
69
  uvr_model_dir=uvr_model_dir
70
  )
71
+ elif whisper_type == WhisperImpl.INSANELY_FAST_WHISPER.value:
72
  return InsanelyFastWhisperInference(
73
  model_dir=insanely_fast_whisper_model_dir,
74
  output_dir=output_dir,