diff --git a/app/assets/api/routes.py b/app/assets/api/routes.py index 160b96d8b..7a0121e8b 100644 --- a/app/assets/api/routes.py +++ b/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}}}") diff --git a/app/assets/database/queries/__init__.py b/app/assets/database/queries/__init__.py index 0a8fbfca5..9d55faca5 100644 --- a/app/assets/database/queries/__init__.py +++ b/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", diff --git a/app/assets/database/queries/asset_reference.py b/app/assets/database/queries/asset_reference.py index 5ec8e058d..e7c0bf5db 100644 --- a/app/assets/database/queries/asset_reference.py +++ b/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, diff --git a/app/assets/services/__init__.py b/app/assets/services/__init__.py index d03469183..3f74fc71b 100644 --- a/app/assets/services/__init__.py +++ b/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", diff --git a/app/assets/services/asset_management.py b/app/assets/services/asset_management.py index 71096f0a1..6fc3b8d5e 100644 --- a/app/assets/services/asset_management.py +++ b/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 = "", diff --git a/tests-unit/assets_test/queries/test_asset_info.py b/tests-unit/assets_test/queries/test_asset_info.py index 48f432d4f..8fcc01d06 100644 --- a/tests-unit/assets_test/queries/test_asset_info.py +++ b/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):