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

Update server_fastapi.py. Sort models by trained steps. (#131)

Browse Source 
* Update server_fastapi.py. Sort models by trained steps.

* [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>
This commit is contained in:
Sora
2023-10-30 09:06:49 +08:00
committed by GitHub
parent a6d5348d5b
commit 9a06dd3cde
1 changed files with 19 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
server_fastapi.py
View File

@@ -298,14 +298,23 @@ if __name__ == "__main__":
# 搜索 "sub_dir" 、 "sub_dir/models" 两个路径 # 搜索 "sub_dir" 、 "sub_dir/models" 两个路径
result[file] = list() result[file] = list()
sub_files = os.listdir(sub_dir) sub_files = os.listdir(sub_dir)
model_files = []
for sub_file in sub_files: for sub_file in sub_files:
relpath = os.path.realpath(os.path.join(sub_dir, sub_file)) 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.paths.keys():
continue continue
if sub_file.endswith(".pth") and sub_file.startswith("G_"): if sub_file.endswith(".pth") and sub_file.startswith("G_"):
if os.path.isfile(relpath): if os.path.isfile(relpath):
result[file].append(sub_file) model_files.append(sub_file)
model_files = sorted(
model_files,
key=lambda pth: int(pth.lstrip("G_").rstrip(".pth"))
if pth.lstrip("G_").rstrip(".pth").isdigit()
else 10**10,
)
result[file] = model_files
models_dir = os.path.join(sub_dir, "models") models_dir = os.path.join(sub_dir, "models")
model_files = []
if os.path.isdir(models_dir): if os.path.isdir(models_dir):
sub_files = os.listdir(models_dir) sub_files = os.listdir(models_dir)
for sub_file in sub_files: for sub_file in sub_files:
@@ -314,9 +323,17 @@ if __name__ == "__main__":
continue continue
if sub_file.endswith(".pth") and sub_file.startswith("G_"): if sub_file.endswith(".pth") and sub_file.startswith("G_"):
if os.path.isfile(os.path.join(models_dir, sub_file)): if os.path.isfile(os.path.join(models_dir, sub_file)):
result[file].append(f"models/{sub_file}") model_files.append(f"models/{sub_file}")
model_files = sorted(
model_files,
key=lambda pth: int(pth.lstrip("models/G_").rstrip(".pth"))
if pth.lstrip("models/G_").rstrip(".pth").isdigit()
else 10**10,
)
result[file] += model_files
if len(result[file]) == 0: if len(result[file]) == 0:
result.pop(file) result.pop(file)
return result return result
@app.get("/models/get_unloaded") @app.get("/models/get_unloaded")