From 25ea864b56ff022f33f3af96e9ac247e467266dd Mon Sep 17 00:00:00 2001 From: Sora Date: Fri, 3 Nov 2023 13:09:58 +0800 Subject: [PATCH] Update server_fastapi.py. Avoid loading the model repeatedly. Add logger (#145) * Update server_fastapi.py. Avoid loading the model repeatedly. Add loguru logger. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- server_fastapi.py | 131 ++++++++++++++++++++++++++------------------- tools/log.py | 14 +++++ tools/translate.py | 2 +- 3 files changed, 92 insertions(+), 55 deletions(-) create mode 100644 tools/log.py diff --git a/server_fastapi.py b/server_fastapi.py index 2106668..a8f50a6 100644 --- a/server_fastapi.py +++ b/server_fastapi.py @@ -6,7 +6,7 @@ import gc import random import utils -from fastapi import FastAPI, Query +from fastapi import FastAPI, Query, Request from fastapi.responses import Response, FileResponse from fastapi.staticfiles import StaticFiles from io import BytesIO @@ -16,9 +16,10 @@ import torch import webbrowser import psutil import GPUtil -from typing import Dict, Optional, List +from typing import Dict, Optional, List, Set import os -from loguru import logger +from tools.log import logger +from urllib.parse import unquote from infer import infer, get_net_g, latest_version import tools.translate as trans @@ -67,25 +68,7 @@ class Models: self.num = 0 # spkInfo[角色名][模型id] = 角色id self.spk_info: Dict[str, Dict[int, int]] = dict() - self.paths: Dict[str, int] = dict() # 路径, 引用数 - - def add_model(self, model: Model): - """添加一个模型""" - self.models[self.num] = model - # 添加角色信息 - for speaker, speaker_id in model.spk2id.items(): - if speaker not in self.spk_info.keys(): - self.spk_info[speaker] = {self.num: speaker_id} - else: - self.spk_info[speaker][self.num] = speaker_id - # 添加路径信息 - model_path = os.path.realpath(model.model_path) - if model_path not in self.paths.keys(): - self.paths[model_path] = 1 - else: - self.paths[model_path] += 1 - # 修改计数 - self.num += 1 + self.path2ids: Dict[str, Set[int]] = dict() # 路径指向的model的id def init_model( self, config_path: str, model_path: str, device: str, language: str @@ -98,26 +81,30 @@ class Models: :param device: 模型推理使用设备 :param language: 模型推理默认语言 """ - self.models[self.num] = Model( - config_path=config_path, - model_path=model_path, - device=device, - language=language, - ) + # 若路径中的模型已存在，则不添加模型，若不存在，则进行初始化。 + model_path = os.path.realpath(model_path) + if model_path not in self.path2ids.keys(): + self.path2ids[model_path] = {self.num} + self.models[self.num] = Model( + config_path=config_path, + model_path=model_path, + device=device, + language=language, + ) + logger.success(f"添加模型{model_path}，使用配置文件{os.path.realpath(config_path)}") + else: + # 获取一个指向id + m_id = next(iter(self.path2ids[model_path])) + self.models[self.num] = self.models[m_id] + self.path2ids[model_path].add(self.num) + logger.success("模型已存在，添加模型引用。") # 添加角色信息 for speaker, speaker_id in self.models[self.num].spk2id.items(): if speaker not in self.spk_info.keys(): self.spk_info[speaker] = {self.num: speaker_id} else: self.spk_info[speaker][self.num] = speaker_id - # 添加路径信息 - model_path = os.path.realpath(self.models[self.num].model_path) - if model_path not in self.paths.keys(): - self.paths[model_path] = 1 - else: - self.paths[model_path] += 1 # 修改计数 - logger.success(f"添加模型{model_path}，使用配置文件{os.path.realpath(config_path)}") self.num += 1 return self.num - 1 @@ -133,13 +120,13 @@ class Models: self.spk_info.pop(speaker) # 删除路径信息 model_path = os.path.realpath(self.models[index].model_path) - self.paths[model_path] -= 1 - assert self.paths[model_path] >= 0 - if self.paths[model_path] == 0: - # 引用数为零时予以清空 - self.paths.pop(model_path) + self.path2ids[model_path].remove(index) + if len(self.path2ids[model_path]) == 0: + self.path2ids.pop(model_path) + logger.success(f"删除模型{model_path}, id = {index}") + else: + logger.success(f"删除模型引用{model_path}, id = {index}") # 删除模型 - logger.success(f"卸载模型{model_path}, id = {index}") self.models.pop(index) gc.collect() if torch.cuda.is_available(): @@ -181,6 +168,7 @@ if __name__ == "__main__": @app.get("/voice") def voice( + request: Request, # fastapi自动注入 text: str = Query(..., description="输入文字"), model_id: int = Query(..., description="模型ID"), # 模型序号 speaker_name: str = Query( @@ -195,7 +183,9 @@ if __name__ == "__main__": auto_translate: bool = Query(False, description="自动翻译"), ): """语音接口""" - + logger.info( + f"{request.client.host}:{request.client.port}/voice { unquote(str(request.query_params) )}" + ) # 检查模型是否存在 if model_id not in loaded_models.models.keys(): return {"status": 10, "detail": f"模型model_id={model_id}未加载"} @@ -235,7 +225,7 @@ if __name__ == "__main__": return response @app.get("/models/info") - def get_loaded_models_info(): + def get_loaded_models_info(request: Request): """获取已加载模型信息""" result: Dict[str, Dict] = dict() @@ -244,9 +234,13 @@ if __name__ == "__main__": return result @app.get("/models/delete") - def delete_model(model_id: int = Query(..., description="删除模型id")): + def delete_model( + request: Request, model_id: int = Query(..., description="删除模型id") + ): """删除指定模型""" - + logger.info( + f"{request.client.host}:{request.client.port}/models/delete { unquote(str(request.query_params) )}" + ) result = loaded_models.del_model(model_id) if result is None: return {"status": 14, "detail": f"模型{model_id}不存在，删除失败"} @@ -254,6 +248,7 @@ if __name__ == "__main__": @app.get("/models/add") def add_model( + request: Request, model_path: str = Query(..., description="添加模型路径"), config_path: str = Query( None, description="添加模型配置文件路径，不填则使用./config.json或../config.json" @@ -261,7 +256,10 @@ if __name__ == "__main__": device: str = Query("cuda", description="推理使用设备"), language: str = Query("ZH", description="模型默认语言"), ): - """添加指定模型：允许重复添加相同路径模型，注意，当前实现中模型会重复加载，加载两次占用两份内存""" + """添加指定模型：允许重复添加相同路径模型，且不重复占用内存""" + logger.info( + f"{request.client.host}:{request.client.port}/models/add { unquote(str(request.query_params) )}" + ) if config_path is None: model_dir = os.path.dirname(model_path) if os.path.isfile(os.path.join(model_dir, "config.json")): @@ -296,6 +294,7 @@ if __name__ == "__main__": } def _get_all_models(root_dir: str = "Data", only_unloaded: bool = False): + """从root_dir搜索获取所有可用模型""" result: Dict[str, List[str]] = dict() files = os.listdir(root_dir) + ["."] for file in files: @@ -307,11 +306,12 @@ if __name__ == "__main__": model_files = [] for sub_file in sub_files: relpath = os.path.realpath(os.path.join(sub_dir, sub_file)) - if only_unloaded and relpath in loaded_models.paths.keys(): + if only_unloaded and relpath in loaded_models.path2count.keys(): continue if sub_file.endswith(".pth") and sub_file.startswith("G_"): if os.path.isfile(relpath): model_files.append(sub_file) + # 对模型文件按步数排序 model_files = sorted( model_files, key=lambda pth: int(pth.lstrip("G_").rstrip(".pth")) @@ -325,11 +325,12 @@ if __name__ == "__main__": sub_files = os.listdir(models_dir) for sub_file in sub_files: relpath = os.path.realpath(os.path.join(models_dir, sub_file)) - if only_unloaded and relpath in loaded_models.paths.keys(): + if only_unloaded and relpath in loaded_models.path2count.keys(): continue if sub_file.endswith(".pth") and sub_file.startswith("G_"): if os.path.isfile(os.path.join(models_dir, sub_file)): model_files.append(f"models/{sub_file}") + # 对模型文件按步数排序 model_files = sorted( model_files, key=lambda pth: int(pth.lstrip("models/G_").rstrip(".pth")) @@ -343,13 +344,23 @@ if __name__ == "__main__": return result @app.get("/models/get_unloaded") - def get_unloaded_models_info(root_dir: str = Query("Data", description="搜索根目录")): + def get_unloaded_models_info( + request: Request, root_dir: str = Query("Data", description="搜索根目录") + ): """获取未加载模型""" + logger.info( + f"{request.client.host}:{request.client.port}/models/get_unloaded { unquote(str(request.query_params) )}" + ) return _get_all_models(root_dir, only_unloaded=True) @app.get("/models/get_local") - def get_local_models_info(root_dir: str = Query("Data", description="搜索根目录")): + def get_local_models_info( + request: Request, root_dir: str = Query("Data", description="搜索根目录") + ): """获取全部本地模型""" + logger.info( + f"{request.client.host}:{request.client.port}/models/get_local { unquote(str(request.query_params) )}" + ) return _get_all_models(root_dir, only_unloaded=False) @app.get("/status") @@ -390,22 +401,30 @@ if __name__ == "__main__": @app.get("/tools/translate") def translate( + request: Request, texts: str = Query(..., description="待翻译文本"), to_language: str = Query(..., description="翻译目标语言"), ): """翻译""" + logger.info( + f"{request.client.host}:{request.client.port}/tools/translate { unquote(str(request.query_params) )}" + ) return {"texts": trans.translate(Sentence=texts, to_Language=to_language)} all_examples: Dict[str, Dict[str, List]] = dict() # 存放示例 @app.get("/tools/random_example") def random_example( + request: Request, language: str = Query(None, description="指定语言，未指定则随机返回"), root_dir: str = Query("Data", description="搜索根目录"), ): """ 获取一个随机音频+文本，用于对比，音频会从本地目录随机选择。 """ + logger.info( + f"{request.client.host}:{request.client.port}/tools/random_example { unquote(str(request.query_params) )}" + ) global all_examples # 数据初始化 if root_dir not in all_examples.keys(): @@ -417,7 +436,6 @@ if __name__ == "__main__": for root, directories, _files in os.walk("Data"): for file in _files: if file in ["train.list", "val.list"]: - print(file) with open( os.path.join(root, file), mode="r", encoding="utf-8" ) as f: @@ -478,7 +496,10 @@ if __name__ == "__main__": } @app.get("/tools/get_audio") - def get_audio(path: str = Query(..., description="本地音频路径")): + def get_audio(request: Request, path: str = Query(..., description="本地音频路径")): + logger.info( + f"{request.client.host}:{request.client.port}/tools/get_audio { unquote(str(request.query_params) )}" + ) if not os.path.isfile(path): return {"status": 18, "detail": "指定音频不存在"} if not path.endswith(".wav"): @@ -486,6 +507,8 @@ if __name__ == "__main__": return FileResponse(path=path) logger.warning("本地服务，请勿将服务端口暴露于外网") - print(f"api文档地址 http://127.0.0.1:{config.server_config.port}/docs") + logger.info(f"api文档地址 http://127.0.0.1:{config.server_config.port}/docs") webbrowser.open(f"http://127.0.0.1:{config.server_config.port}") - uvicorn.run(app, port=config.server_config.port, host="0.0.0.0") + uvicorn.run( + app, port=config.server_config.port, host="0.0.0.0", log_level="warning" + ) diff --git a/tools/log.py b/tools/log.py new file mode 100644 index 0000000..626b9e9 --- /dev/null +++ b/tools/log.py @@ -0,0 +1,14 @@ +""" +logger封装 +""" +from loguru import logger +import sys + + +# 移除所有默认的处理器 +logger.remove() + +# 自定义格式并添加到标准输出 +log_format = "{time:MM-DD HH:mm:ss} [{level}] | {message}" + +logger.add(sys.stdout, format=log_format) diff --git a/tools/translate.py b/tools/translate.py index a99140a..d2c5aa8 100644 --- a/tools/translate.py +++ b/tools/translate.py @@ -20,7 +20,7 @@ def translate(Sentence: str, to_Language: str = "jp", from_Language: str = ""): appid = config.translate_config.app_key key = config.translate_config.secret_key if appid == "" or key == "": - return "请在开发者config.yml中配置app_key与secret_key" + return "请开发者在config.yml中配置app_key与secret_key" url = "https://fanyi-api.baidu.com/api/trans/vip/translate" texts = Sentence.split("\n") outTexts = []