From 2b6910bd5510b42cb189dcf8597522f447436f97 Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Thu, 9 Oct 2025 15:23:42 +0800
Subject: [PATCH] feat(misc): mark model versions in library for Civitai user
 models

---
 py/routes/handlers/misc_handlers.py |  6 +++---
 tests/routes/test_misc_routes.py    | 27 ++++++++++++++++++++++++++-
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/py/routes/handlers/misc_handlers.py b/py/routes/handlers/misc_handlers.py
index 4c41a06b..459d19d5 100644
--- a/py/routes/handlers/misc_handlers.py
+++ b/py/routes/handlers/misc_handlers.py
@@ -687,9 +687,6 @@ class ModelLibraryHandler:
                     except (TypeError, ValueError):
                         continue
 
-                    if await scanner.check_model_version_exists(version_id_int):
-                        continue
-
                     images = version.get("images") or []
                     thumbnail_url = None
                     if images and isinstance(images, list):
@@ -697,6 +694,8 @@ class ModelLibraryHandler:
                         if isinstance(first_image, dict):
                             thumbnail_url = first_image.get("url")
 
+                    in_library = await scanner.check_model_version_exists(version_id_int)
+
                     versions.append(
                         {
                             "modelId": model_id_int,
@@ -707,6 +706,7 @@ class ModelLibraryHandler:
                             "tags": tags,
                             "baseModel": version.get("baseModel"),
                             "thumbnailUrl": thumbnail_url,
+                            "inLibrary": in_library,
                         }
                     )
 
diff --git a/tests/routes/test_misc_routes.py b/tests/routes/test_misc_routes.py
index 0e597d35..3dca0515 100644
--- a/tests/routes/test_misc_routes.py
+++ b/tests/routes/test_misc_routes.py
@@ -365,7 +365,7 @@ async def test_misc_routes_bind_produces_expected_handlers():
 
 
 @pytest.mark.asyncio
-async def test_get_civitai_user_models_filters_versions():
+async def test_get_civitai_user_models_marks_library_versions():
     models = [
         {
             "id": 1,
@@ -475,6 +475,18 @@ async def test_get_civitai_user_models_filters_versions():
             "tags": ["style"],
             "baseModel": "Flux.1",
             "thumbnailUrl": "http://example.com/a1.jpg",
+            "inLibrary": False,
+        },
+        {
+            "modelId": 1,
+            "versionId": 101,
+            "modelName": "Model A",
+            "versionName": "v2",
+            "type": "LORA",
+            "tags": ["style"],
+            "baseModel": "Flux.1",
+            "thumbnailUrl": "http://example.com/a2.jpg",
+            "inLibrary": True,
         },
         {
             "modelId": 2,
@@ -485,6 +497,18 @@ async def test_get_civitai_user_models_filters_versions():
             "tags": ["embedding"],
             "baseModel": None,
             "thumbnailUrl": "http://example.com/e1.jpg",
+            "inLibrary": False,
+        },
+        {
+            "modelId": 2,
+            "versionId": 202,
+            "modelName": "Embedding",
+            "versionName": "v2",
+            "type": "TextualInversion",
+            "tags": ["embedding"],
+            "baseModel": None,
+            "thumbnailUrl": None,
+            "inLibrary": True,
         },
         {
             "modelId": 3,
@@ -495,6 +519,7 @@ async def test_get_civitai_user_models_filters_versions():
             "tags": ["checkpoint"],
             "baseModel": "SDXL",
             "thumbnailUrl": None,
+            "inLibrary": False,
         },
     ]