spike: add model folder debug counts
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled

Amp-Thread-ID: https://ampcode.com/threads/T-019e5117-c707-729d-bf98-dce718fe64d5
Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
Simon Pinfold 2026-06-10 15:22:23 +12:00
parent 505367b52b
commit 70de84cac7
6 changed files with 170 additions and 5 deletions

View File

@ -3,6 +3,7 @@ import functools
import json import json
import logging import logging
import os import os
from pathlib import Path
import urllib.parse import urllib.parse
import uuid import uuid
from typing import Any from typing import Any
@ -33,6 +34,7 @@ from app.assets.services import (
get_asset_detail, get_asset_detail,
list_assets_page, list_assets_page,
list_model_folder_counts, list_model_folder_counts,
list_model_folder_reference_paths,
list_tags, list_tags,
remove_tags, remove_tags,
resolve_asset_for_download, resolve_asset_for_download,
@ -238,11 +240,45 @@ def _build_asset_response(result: schemas.AssetDetailResult | schemas.UploadResu
def _build_model_folders_response_payload( def _build_model_folders_response_payload(
counts_by_folder: dict[str, int] | None = None, counts_by_folder: dict[str, int] | None = None,
reference_paths_by_folder: list[tuple[str, str]] | None = None,
) -> dict[str, Any]: ) -> dict[str, Any]:
counts_by_folder = counts_by_folder or {} 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 = [ 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)} 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 @_require_assets_feature_enabled
async def list_model_folders_route(request: web.Request) -> web.Response: async def list_model_folders_route(request: web.Request) -> web.Response:
"""Debug endpoint for registered model folders known to the assets API.""" """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)) owner_id = USER_MANAGER.get_request_user_id(request)
return web.json_response(_build_model_folders_response_payload(counts)) 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}}}") @ROUTES.get(f"/api/assets/{{id:{UUID_RE}}}")

View File

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

View File

@ -367,6 +367,23 @@ def count_model_references_by_folder(
return {model_folder: int(count) for model_folder, count in rows} 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( def fetch_reference_asset_and_tags(
session: Session, session: Session,
reference_id: str, reference_id: str,

View File

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

View File

@ -19,6 +19,7 @@ from app.assets.database.queries import (
list_all_file_paths_by_asset_id, list_all_file_paths_by_asset_id,
list_references_by_asset_id, list_references_by_asset_id,
count_model_references_by_folder, count_model_references_by_folder,
list_model_reference_paths_by_folder,
set_reference_metadata, set_reference_metadata,
set_reference_preview, set_reference_preview,
set_reference_tags, 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) 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( def resolve_hash_to_path(
asset_hash: str, asset_hash: str,
owner_id: str = "", owner_id: str = "",

View File

@ -18,6 +18,7 @@ from app.assets.database.queries import (
fetch_reference_asset_and_tags, fetch_reference_asset_and_tags,
fetch_reference_and_asset, fetch_reference_and_asset,
count_model_references_by_folder, count_model_references_by_folder,
list_model_reference_paths_by_folder,
update_reference_access_time, update_reference_access_time,
set_reference_metadata, set_reference_metadata,
delete_reference_by_id, delete_reference_by_id,
@ -77,7 +78,14 @@ class TestModelFoldersDebugPayload:
) )
payload = _build_model_folders_response_payload( 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 == { assert payload == {
@ -86,16 +94,56 @@ class TestModelFoldersDebugPayload:
"name": "checkpoints", "name": "checkpoints",
"folders": ["/models/checkpoints", "/extra/checkpoints"], "folders": ["/models/checkpoints", "/extra/checkpoints"],
"count": 3, "count": 3,
"folder_counts": [
{"folder": "/models/checkpoints", "count": 2},
{"folder": "/extra/checkpoints", "count": 1},
],
}, },
{ {
"name": "text_encoders/clip", "name": "text_encoders/clip",
"folders": ["/models/text_encoders/clip"], "folders": ["/models/text_encoders/clip"],
"count": 1, "count": 1,
"folder_counts": [
{"folder": "/models/text_encoders/clip", "count": 1},
],
}, },
], ],
"total": 2, "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: 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") 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" 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: class TestFetchReferenceAssetAndTags:
def test_returns_none_for_nonexistent(self, session: Session): def test_returns_none_for_nonexistent(self, session: Session):