| import torch |
| from modules import config |
| from modules import generate_audio as generate |
|
|
| from functools import lru_cache |
| from typing import Callable |
|
|
| from modules.api.Api import APIManager |
|
|
| from modules.api.impl import ( |
| base_api, |
| tts_api, |
| ssml_api, |
| google_api, |
| openai_api, |
| refiner_api, |
| ) |
|
|
| torch._dynamo.config.cache_size_limit = 64 |
| torch._dynamo.config.suppress_errors = True |
| torch.set_float32_matmul_precision("high") |
|
|
|
|
| def create_api(): |
| api = APIManager() |
|
|
| base_api.setup(api) |
| tts_api.setup(api) |
| ssml_api.setup(api) |
| google_api.setup(api) |
| openai_api.setup(api) |
| refiner_api.setup(api) |
|
|
| return api |
|
|
|
|
| def conditional_cache(condition: Callable): |
| def decorator(func): |
| @lru_cache(None) |
| def cached_func(*args, **kwargs): |
| return func(*args, **kwargs) |
|
|
| def wrapper(*args, **kwargs): |
| if condition(*args, **kwargs): |
| return cached_func(*args, **kwargs) |
| else: |
| return func(*args, **kwargs) |
|
|
| return wrapper |
|
|
| return decorator |
|
|
|
|
| if __name__ == "__main__": |
| import argparse |
| import uvicorn |
|
|
| parser = argparse.ArgumentParser( |
| description="Start the FastAPI server with command line arguments" |
| ) |
| parser.add_argument( |
| "--host", type=str, default="0.0.0.0", help="Host to run the server on" |
| ) |
| parser.add_argument( |
| "--port", type=int, default=8000, help="Port to run the server on" |
| ) |
| parser.add_argument( |
| "--reload", action="store_true", help="Enable auto-reload for development" |
| ) |
| parser.add_argument("--compile", action="store_true", help="Enable model compile") |
| parser.add_argument( |
| "--lru_size", |
| type=int, |
| default=64, |
| help="Set the size of the request cache pool, set it to 0 will disable lru_cache", |
| ) |
| parser.add_argument( |
| "--cors_origin", |
| type=str, |
| default="*", |
| help="Allowed CORS origins. Use '*' to allow all origins.", |
| ) |
|
|
| args = parser.parse_args() |
|
|
| config.args = args |
|
|
| if args.compile: |
| print("Model compile is enabled") |
| config.enable_model_compile = True |
|
|
| def should_cache(*args, **kwargs): |
| spk_seed = kwargs.get("spk_seed", -1) |
| infer_seed = kwargs.get("infer_seed", -1) |
| return spk_seed != -1 and infer_seed != -1 |
|
|
| if args.lru_size > 0: |
| config.lru_size = args.lru_size |
| generate.generate_audio = conditional_cache(should_cache)( |
| generate.generate_audio |
| ) |
|
|
| api = create_api() |
| config.api = api |
|
|
| if args.cors_origin: |
| api.set_cors(allow_origins=[args.cors_origin]) |
|
|
| uvicorn.run(api.app, host=args.host, port=args.port, reload=args.reload) |
|
|
