SpongeBobFan2002/openaudio-s1-mini

1import re2from threading import Lock3 4import pyrootutils5import uvicorn6from kui.asgi import (7    Depends,8    FactoryClass,9    HTTPException,10    HttpRoute,11    Kui,12    OpenAPI,13    Routes,14)15from kui.cors import CORSConfig16from kui.openapi.specification import Info17from kui.security import bearer_auth18from loguru import logger19from typing_extensions import Annotated20 21pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)22 23from tools.server.api_utils import MsgPackRequest, parse_args24from tools.server.exception_handler import ExceptionHandler25from tools.server.model_manager import ModelManager26from tools.server.views import routes27 28 29class API(ExceptionHandler):30    def __init__(self):31        self.args = parse_args()32 33        def api_auth(endpoint):34            async def verify(token: Annotated[str, Depends(bearer_auth)]):35                if token != self.args.api_key:36                    raise HTTPException(401, None, "Invalid token")37                return await endpoint()38 39            async def passthrough():40                return await endpoint()41 42            if self.args.api_key is not None:43                return verify44            else:45                return passthrough46 47        self.routes = Routes(48            routes,  # keep existing routes49            http_middlewares=[api_auth],  # apply api_auth middleware50        )51 52        # OpenAPIの設定53        self.openapi = OpenAPI(54            Info(55                {56                    "title": "Fish Speech API",57                    "version": "1.5.0",58                }59            ),60        ).routes61 62        # Initialize the app63        self.app = Kui(64            routes=self.routes + self.openapi[1:],  # Remove the default route65            exception_handlers={66                HTTPException: self.http_exception_handler,67                Exception: self.other_exception_handler,68            },69            factory_class=FactoryClass(http=MsgPackRequest),70            cors_config=CORSConfig(),71        )72 73        # Add the state variables74        self.app.state.lock = Lock()75        self.app.state.device = self.args.device76        self.app.state.max_text_length = self.args.max_text_length77 78        # Associate the app with the model manager79        self.app.on_startup(self.initialize_app)80 81    async def initialize_app(self, app: Kui):82        # Make the ModelManager available to the views83        app.state.model_manager = ModelManager(84            mode=self.args.mode,85            device=self.args.device,86            half=self.args.half,87            compile=self.args.compile,88            asr_enabled=self.args.load_asr_model,89            llama_checkpoint_path=self.args.llama_checkpoint_path,90            decoder_checkpoint_path=self.args.decoder_checkpoint_path,91            decoder_config_name=self.args.decoder_config_name,92        )93 94        logger.info(f"Startup done, listening server at http://{self.args.listen}")95 96 97# Each worker process created by Uvicorn has its own memory space,98# meaning that models and variables are not shared between processes.99# Therefore, any variables (like `llama_queue` or `decoder_model`)100# will not be shared across workers.101 102# Multi-threading for deep learning can cause issues, such as inconsistent103# outputs if multiple threads access the same buffers simultaneously.104# Instead, it's better to use multiprocessing or independent models per thread.105 106if __name__ == "__main__":107    api = API()108 109    # IPv6 address format is [xxxx:xxxx::xxxx]:port110    match = re.search(r"\[([^\]]+)\]:(\d+)$", api.args.listen)111    if match:112        host, port = match.groups()  # IPv6113    else:114        host, port = api.args.listen.split(":")  # IPv4115 116    uvicorn.run(117        api.app,118        host=host,119        port=int(port),120        workers=api.args.workers,121        log_level="info",122    )123