diff --git a/src/transcribble/cli.py b/src/transcribble/cli.py index 7e17b9a..2bb68f8 100644 --- a/src/transcribble/cli.py +++ b/src/transcribble/cli.py @@ -10,6 +10,18 @@ class IgnoreHFWarning(logging.Filter): logger = logging.getLogger("huggingface_hub.utils._http") logger.addFilter(IgnoreHFWarning()) +class WhisperModel(argparse.Action): + def __call__(self, parser, namespace, value, option_string=None): + import whisper + models = whisper.available_models() + + if value == "?": + parser.exit(message="\n".join(models) + "\n") + if value not in models: + parser.error(f"invalid mode '{value}'\nchoose from:\n {'\n '.join(models)}") + + setattr(namespace, self.dest, value) + def main(): parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("input") @@ -17,28 +29,17 @@ def main(): help="Output subtitle (SRT) file path") parser.add_argument("-s", "--speakers", type=int, default=2, help="Number of speakers in the audio") - parser.add_argument("-m", "--model", default='turbo', + parser.add_argument("-m", "--model", default='turbo', action=WhisperModel, help="Whisper model to use ('?' to see available options)") parser.add_argument("--gpu", action="store_true", help="run on GPU if available") args = parser.parse_args() import torch - import whisper from .transcribe import transcribe from .diarize import diarize from .output import write_srt - models = whisper.available_models() - - if args.model == "?": - for model in models: - print(model) - return - - if args.model not in models: - parser.error(f"'{args.model}' is not a valid whisper model") - has_gpu = torch.cuda.is_available() if args.gpu and not has_gpu: print("GPU not available. Falling back to CPU")