From 9a06dd3cde2fb24a6a83be033de3a72523389f5f Mon Sep 17 00:00:00 2001
From: Sora <atri@suzakuintsubaki.com>
Date: Mon, 30 Oct 2023 09:06:49 +0800
Subject: [PATCH] Update server_fastapi.py. Sort models by trained steps.
 (#131)

* 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>
---
 server_fastapi.py | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

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")