import argparse
import uvicorn
from fastapi_funasr import app

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_dir", type=str, default="/model", help="model directory")
    parser.add_argument("--model_type", type=str, default="sensevoice", help="model type, e.g. sensevoice")
    parser.add_argument("--use_gpu", action="store_true", default=True)
    parser.add_argument("--warmup", action="store_true", help="whether do warmup when first initializing model")
    parser.add_argument("--port", type=int, default=8000, help="service port")

    args = parser.parse_args()

    # 将参数加到 app.state 中
    app.state.config = {
        "model_dir": args.model_dir,
        "model_type": args.model_type,
        "use_gpu": args.use_gpu, # True
        "warmup": args.warmup
    }

    uvicorn.run("fastapi_funasr:app",
        host="0.0.0.0",
        port=args.port,
        workers=1
    )