diff --git a/whisper_server.py b/whisper_server.py index 19583f4..f9837c5 100644 --- a/whisper_server.py +++ b/whisper_server.py @@ -34,8 +34,9 @@ @app.on_event("startup") async def load_whisper_model(): global model - model_name = os.environ.get("WHISPER_MODEL", "large-v3") - device = os.environ.get("WHISPER_DEVICE", "cuda" if torch.cuda.is_available() else "cpu") + model_name = os.environ.get("WHISPER_MODEL", "large-v3") + device = os.environ.get("WHISPER_DEVICE", "cuda" if torch.cuda.is_available() else "cpu") + compute_type = os.environ.get("WHISPER_COMPUTE_TYPE", "float16" if device != "cpu" else "int8") logging.basicConfig( level=logging.INFO, @@ -48,7 +49,7 @@ async def load_whisper_model(): model = WhisperModel( model_name, device=device, - compute_type="float16" + compute_type=compute_type ) app.register_model("whisper-1", model_name) @@ -301,4 +302,4 @@ def parse_args(argv=None): host=args.host, port=args.port, workers=args.workers, - ) \ No newline at end of file + )