diff --git a/server_fastapi.py b/server_fastapi.py index 972d63a..517500e 100644 --- a/server_fastapi.py +++ b/server_fastapi.py @@ -298,14 +298,23 @@ if __name__ == "__main__": # 搜索 "sub_dir" 、 "sub_dir/models" 两个路径 result[file] = list() sub_files = os.listdir(sub_dir) + 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(): continue if sub_file.endswith(".pth") and sub_file.startswith("G_"): 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") + model_files = [] if os.path.isdir(models_dir): sub_files = os.listdir(models_dir) for sub_file in sub_files: @@ -314,9 +323,17 @@ if __name__ == "__main__": continue if sub_file.endswith(".pth") and sub_file.startswith("G_"): 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: result.pop(file) + return result @app.get("/models/get_unloaded")