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

@@ -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()  # 路径, 引用数
+        self.path2ids: Dict[str, Set[int]] = dict()  # 路径指向的model的id
-
-    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
 
     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 = os.path.realpath(model_path)
-            model_path=model_path,
+        if model_path not in self.path2ids.keys():
-            device=device,
+            self.path2ids[model_path] = {self.num}
-            language=language,
+            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
+        self.path2ids[model_path].remove(index)
-        assert self.paths[model_path] >= 0
+        if len(self.path2ids[model_path]) == 0:
-        if self.paths[model_path] == 0:
+            self.path2ids.pop(model_path)
-            # 引用数为零时予以清空
+            logger.success(f"删除模型{model_path}, id = {index}")
-            self.paths.pop(model_path)
+        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"
+    )

tools/log.py

@@ -0,0 +1,14 @@
+"""
+logger封装
+"""
+from loguru import logger
+import sys
+
+
+# 移除所有默认的处理器
+logger.remove()
+
+# 自定义格式并添加到标准输出
+log_format = "<g>{time:MM-DD HH:mm:ss}</g> [<lvl>{level}</lvl>] | {message}"
+
+logger.add(sys.stdout, format=log_format)

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 = []