39 lines
1.7 KiB
Python
39 lines
1.7 KiB
Python
import argparse
|
||
import uvicorn
|
||
from fastapi_transformers import app
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--model_dir", type=str, default="/model",
|
||
help="模型目录（挂载到容器内的路径）")
|
||
parser.add_argument("--model_type", type=str, default=None,
|
||
help="可选，仅在自动推断失败时手动指定: whisper 或不填（CTC 类均不需要填）")
|
||
parser.add_argument("--use_gpu", action="store_true", default=True,
|
||
help="是否使用 GPU（CUDA）")
|
||
parser.add_argument("--warmup", action="store_true",
|
||
help="启动时用静音片段执行一次 warmup 推理")
|
||
parser.add_argument("--chunk_length_s", type=int, default=30,
|
||
help="长音频切片长度（秒），逐段推理，默认 30")
|
||
parser.add_argument("--fp16", action="store_true", default=False,
|
||
help="使用 float16 推理（默认 float32）。仅在确认硬件支持时开启，"
|
||
"注意 fp16/fp32 之间存在精度差异，跨卡对比时建议保持默认 fp32")
|
||
parser.add_argument("--port", type=int, default=8000,
|
||
help="FastAPI 服务端口，默认 8000")
|
||
|
||
args = parser.parse_args()
|
||
|
||
app.state.config = {
|
||
"model_dir": args.model_dir,
|
||
"model_type": args.model_type,
|
||
"use_gpu": args.use_gpu,
|
||
"warmup": args.warmup,
|
||
"chunk_length_s": args.chunk_length_s,
|
||
"fp16": args.fp16,
|
||
}
|
||
|
||
uvicorn.run("fastapi_transformers:app",
|
||
host="0.0.0.0",
|
||
port=args.port,
|
||
workers=1
|
||
)
|