spike: add model folder debug counts

Amp-Thread-ID: https://ampcode.com/threads/T-019e5117-c707-729d-bf98-dce718fe64d5
Co-authored-by: Amp <amp@ampcode.com>

app/assets/api/routes.py

@@ -3,6 +3,7 @@ import functools
 import json
 import logging
 import os
+from pathlib import Path
 import urllib.parse
 import uuid
 from typing import Any
@@ -33,6 +34,7 @@ from app.assets.services import (
     get_asset_detail,
     list_assets_page,
     list_model_folder_counts,
+    list_model_folder_reference_paths,
     list_tags,
     remove_tags,
     resolve_asset_for_download,
@@ -238,11 +240,45 @@ def _build_asset_response(result: schemas.AssetDetailResult | schemas.UploadResu
 
 def _build_model_folders_response_payload(
     counts_by_folder: dict[str, int] | None = None,
+    reference_paths_by_folder: list[tuple[str, str]] | None = None,
 ) -> dict[str, Any]:
     counts_by_folder = counts_by_folder or {}
+    reference_paths_by_folder = reference_paths_by_folder or []
+
+    registered_model_folders = get_comfy_models_folders()
+    folder_count_lookup: dict[tuple[str, str], int] = {}
+    registered_by_name = {
+        name: [os.path.abspath(folder) for folder in folders]
+        for name, folders in registered_model_folders
+    }
+    for model_folder, file_path in reference_paths_by_folder:
+        file_path_abs = os.path.abspath(file_path)
+        candidates = [
+            folder
+            for folder in registered_by_name.get(model_folder, [])
+            if Path(file_path_abs).is_relative_to(folder)
+        ]
+        if not candidates:
+            continue
+        matched_folder = max(candidates, key=len)
+        folder_count_lookup[(model_folder, matched_folder)] = (
+            folder_count_lookup.get((model_folder, matched_folder), 0) + 1
+        )
+
     model_folders = [
-        {"name": name, "folders": folders, "count": counts_by_folder.get(name, 0)}
-        for name, folders in get_comfy_models_folders()
+        {
+            "name": name,
+            "folders": folders,
+            "count": counts_by_folder.get(name, 0),
+            "folder_counts": [
+                {
+                    "folder": folder,
+                    "count": folder_count_lookup.get((name, os.path.abspath(folder)), 0),
+                }
+                for folder in folders
+            ],
+        }
+        for name, folders in registered_model_folders
     ]
     return {"model_folders": model_folders, "total": len(model_folders)}
 
@@ -305,8 +341,10 @@ async def list_assets_route(request: web.Request) -> web.Response:
 @_require_assets_feature_enabled
 async def list_model_folders_route(request: web.Request) -> web.Response:
     """Debug endpoint for registered model folders known to the assets API."""
-    counts = list_model_folder_counts(owner_id=USER_MANAGER.get_request_user_id(request))
-    return web.json_response(_build_model_folders_response_payload(counts))
+    owner_id = USER_MANAGER.get_request_user_id(request)
+    counts = list_model_folder_counts(owner_id=owner_id)
+    reference_paths = list_model_folder_reference_paths(owner_id=owner_id)
+    return web.json_response(_build_model_folders_response_payload(counts, reference_paths))
 
 
 @ROUTES.get(f"/api/assets/{{id:{UUID_RE}}}")

app/assets/database/queries/__init__.py

@@ -32,6 +32,7 @@ from app.assets.database.queries.asset_reference import (
     get_reference_ids_by_ids,
     get_references_by_paths_and_asset_ids,
     get_references_for_prefixes,
+    list_model_reference_paths_by_folder,
     get_unenriched_references,
     get_unreferenced_unhashed_asset_ids,
     insert_reference,
@@ -108,6 +109,7 @@ __all__ = [
     "get_reference_tags",
     "get_references_by_paths_and_asset_ids",
     "get_references_for_prefixes",
+    "list_model_reference_paths_by_folder",
     "get_unenriched_references",
     "get_unreferenced_unhashed_asset_ids",
     "insert_reference",

app/assets/database/queries/asset_reference.py

@@ -367,6 +367,23 @@ def count_model_references_by_folder(
     return {model_folder: int(count) for model_folder, count in rows}
 
 
+def list_model_reference_paths_by_folder(
+    session: Session,
+    owner_id: str = "",
+) -> list[tuple[str, str]]:
+    """Return visible active model reference model_folder/file_path pairs."""
+    rows = session.execute(
+        select(AssetReference.model_folder, AssetReference.file_path)
+        .where(build_visible_owner_clause(owner_id))
+        .where(AssetReference.is_missing == False)  # noqa: E712
+        .where(AssetReference.deleted_at.is_(None))
+        .where(AssetReference.asset_type == "model")
+        .where(AssetReference.model_folder.isnot(None))
+        .where(AssetReference.file_path.isnot(None))
+    ).all()
+    return [(model_folder, file_path) for model_folder, file_path in rows]
+
+
 def fetch_reference_asset_and_tags(
     session: Session,
     reference_id: str,

app/assets/services/__init__.py

@@ -5,6 +5,7 @@ from app.assets.services.asset_management import (
     get_asset_detail,
     list_assets_page,
     list_model_folder_counts,
+    list_model_folder_reference_paths,
     resolve_asset_for_download,
     set_asset_preview,
     update_asset_metadata,
@@ -81,6 +82,7 @@ __all__ = [
     "get_size_and_mtime_ns",
     "list_assets_page",
     "list_model_folder_counts",
+    "list_model_folder_reference_paths",
     "list_files_recursively",
     "list_tags",
     "cleanup_unreferenced_assets",

app/assets/services/asset_management.py

@@ -19,6 +19,7 @@ from app.assets.database.queries import (
     list_all_file_paths_by_asset_id,
     list_references_by_asset_id,
     count_model_references_by_folder,
+    list_model_reference_paths_by_folder,
     set_reference_metadata,
     set_reference_preview,
     set_reference_tags,
@@ -290,6 +291,11 @@ def list_model_folder_counts(owner_id: str = "") -> dict[str, int]:
         return count_model_references_by_folder(session, owner_id=owner_id)
 
 
+def list_model_folder_reference_paths(owner_id: str = "") -> list[tuple[str, str]]:
+    with create_session() as session:
+        return list_model_reference_paths_by_folder(session, owner_id=owner_id)
+
+
 def resolve_hash_to_path(
     asset_hash: str,
     owner_id: str = "",

tests-unit/assets_test/queries/test_asset_info.py

@@ -18,6 +18,7 @@ from app.assets.database.queries import (
     fetch_reference_asset_and_tags,
     fetch_reference_and_asset,
     count_model_references_by_folder,
+    list_model_reference_paths_by_folder,
     update_reference_access_time,
     set_reference_metadata,
     delete_reference_by_id,
@@ -77,7 +78,14 @@ class TestModelFoldersDebugPayload:
         )
 
         payload = _build_model_folders_response_payload(
-            {"checkpoints": 3, "text_encoders/clip": 1}
+            {"checkpoints": 3, "text_encoders/clip": 1},
+            [
+                ("checkpoints", "/models/checkpoints/a.safetensors"),
+                ("checkpoints", "/models/checkpoints/nested/b.safetensors"),
+                ("checkpoints", "/extra/checkpoints/c.safetensors"),
+                ("checkpoints", "/unregistered/checkpoints/d.safetensors"),
+                ("text_encoders/clip", "/models/text_encoders/clip/clip.safetensors"),
+            ],
         )
 
         assert payload == {
@@ -86,16 +94,56 @@ class TestModelFoldersDebugPayload:
                     "name": "checkpoints",
                     "folders": ["/models/checkpoints", "/extra/checkpoints"],
                     "count": 3,
+                    "folder_counts": [
+                        {"folder": "/models/checkpoints", "count": 2},
+                        {"folder": "/extra/checkpoints", "count": 1},
+                    ],
                 },
                 {
                     "name": "text_encoders/clip",
                     "folders": ["/models/text_encoders/clip"],
                     "count": 1,
+                    "folder_counts": [
+                        {"folder": "/models/text_encoders/clip", "count": 1},
+                    ],
                 },
             ],
             "total": 2,
         }
 
+    def test_folder_counts_use_deepest_registered_physical_root(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        monkeypatch.setattr(
+            "app.assets.api.routes.get_comfy_models_folders",
+            lambda: [
+                (
+                    "checkpoints",
+                    ["/models/checkpoints", "/models/checkpoints/nested"],
+                ),
+            ],
+        )
+
+        payload = _build_model_folders_response_payload(
+            {"checkpoints": 2},
+            [
+                ("checkpoints", "/models/checkpoints/base.safetensors"),
+                ("checkpoints", "/models/checkpoints/nested/deeper.safetensors"),
+            ],
+        )
+
+        assert payload["model_folders"] == [
+            {
+                "name": "checkpoints",
+                "folders": ["/models/checkpoints", "/models/checkpoints/nested"],
+                "count": 2,
+                "folder_counts": [
+                    {"folder": "/models/checkpoints", "count": 1},
+                    {"folder": "/models/checkpoints/nested", "count": 1},
+                ],
+            }
+        ]
+
 
 def _make_asset(session: Session, hash_val: str | None = None, size: int = 1024) -> Asset:
     asset = Asset(hash=hash_val, size_bytes=size, mime_type="application/octet-stream")
@@ -860,6 +908,58 @@ class TestModelFolderCounts:
         }
         assert private.owner_id == "other-user"
 
+    def test_lists_visible_active_model_reference_paths_by_folder(self, session: Session):
+        asset = _make_asset(session, "hash-path-counts")
+        visible = _make_reference(
+            session,
+            asset,
+            name="checkpoint",
+            file_path="/models/checkpoints/a.safetensors",
+            asset_type="model",
+            model_folder="checkpoints",
+        )
+        _make_reference(
+            session,
+            asset,
+            name="pathless",
+            asset_type="model",
+            model_folder="checkpoints",
+        )
+        _make_reference(
+            session,
+            asset,
+            name="input",
+            file_path="/input/a.png",
+            asset_type="input",
+        )
+        missing = _make_reference(
+            session,
+            asset,
+            name="missing",
+            file_path="/models/checkpoints/missing.safetensors",
+            asset_type="model",
+            model_folder="checkpoints",
+        )
+        missing.is_missing = True
+        private = _make_reference(
+            session,
+            asset,
+            name="private",
+            owner_id="other-user",
+            file_path="/models/checkpoints/private.safetensors",
+            asset_type="model",
+            model_folder="checkpoints",
+        )
+        session.commit()
+
+        assert list_model_reference_paths_by_folder(session, owner_id="") == [
+            ("checkpoints", visible.file_path)
+        ]
+        assert set(list_model_reference_paths_by_folder(session, owner_id="other-user")) == {
+            ("checkpoints", visible.file_path),
+            ("checkpoints", private.file_path),
+        }
+
 
 class TestFetchReferenceAssetAndTags:
     def test_returns_none_for_nonexistent(self, session: Session):