tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Update server_fastapi.py. Allowed to search models in "./" (#133)

Browse Source
This commit is contained in:
Sora
2023-10-30 10:43:34 +08:00
committed by GitHub
parent 9a06dd3cde
commit b13c0fad9d
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
server_fastapi.py
View File

@@ -2,7 +2,7 @@
api服务 多版本多模型 fastapi实现 api服务 多版本多模型 fastapi实现
""" """
import logging import logging
import gc
import utils import utils
from fastapi import FastAPI, Query from fastapi import FastAPI, Query
from fastapi.responses import Response, FileResponse from fastapi.responses import Response, FileResponse
@@ -139,6 +139,7 @@ class Models:
# 删除模型 # 删除模型
logger.success(f"卸载模型{model_path}, id = {index}") logger.success(f"卸载模型{model_path}, id = {index}")
self.models.pop(index) self.models.pop(index)
gc.collect()
if torch.cuda.is_available(): if torch.cuda.is_available():
torch.cuda.empty_cache() torch.cuda.empty_cache()
return index return index
@@ -291,7 +292,7 @@ if __name__ == "__main__":
def _get_all_models(root_dir: str = "Data", only_unloaded: bool = False): def _get_all_models(root_dir: str = "Data", only_unloaded: bool = False):
result: Dict[str, List[str]] = dict() result: Dict[str, List[str]] = dict()
files = os.listdir(root_dir) files = os.listdir(root_dir) + ["."]
for file in files: for file in files:
if os.path.isdir(os.path.join(root_dir, file)): if os.path.isdir(os.path.join(root_dir, file)):
sub_dir = os.path.join(root_dir, file) sub_dir = os.path.join(root_dir, file)