mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-20 20:41:26 -03:00
feat(download): expose per-file downloadedFiles in check-model-exists (#1058)
The version branch of check-model-exists now returns
downloadedFiles: [{fileId, fileName, filePath}] so clients (e.g. the
browser extension) can tell a partially downloaded version apart from a
fully downloaded one. Reuses ModelCivitaiHandler._match_downloaded_files
(D2 rule) against the local cache; unmatchable local files are reported
with fileId: None. No CivitAI API call added.
This commit is contained in:
@@ -56,6 +56,7 @@ from ...utils.constants import (
|
|||||||
)
|
)
|
||||||
from .hf_handlers import HfHandler
|
from .hf_handlers import HfHandler
|
||||||
from .agent_handlers import AgentHandler
|
from .agent_handlers import AgentHandler
|
||||||
|
from .model_handlers import ModelCivitaiHandler
|
||||||
from ...utils.civitai_utils import rewrite_preview_url
|
from ...utils.civitai_utils import rewrite_preview_url
|
||||||
from ...utils.example_images_paths import (
|
from ...utils.example_images_paths import (
|
||||||
find_non_compliant_items_in_example_images_root,
|
find_non_compliant_items_in_example_images_root,
|
||||||
@@ -2061,6 +2062,63 @@ class ModelLibraryHandler:
|
|||||||
enriched.append(entry)
|
enriched.append(entry)
|
||||||
return enriched
|
return enriched
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
async def _get_downloaded_files(
|
||||||
|
scanner: Any, model_version_id: int
|
||||||
|
) -> list[dict[str, Any]]:
|
||||||
|
"""Return per-file downloaded state for a version in the library.
|
||||||
|
|
||||||
|
This handler has no CivitAI version payload, so the remote file list
|
||||||
|
is taken from the local entries' cached ``civitai`` metadata (the
|
||||||
|
full version payload persisted at download time, see
|
||||||
|
``BaseModelMetadata.from_civitai_info``) and matched with the same
|
||||||
|
D2 rule used by ``get_civitai_versions`` (#1058). Local entries that
|
||||||
|
cannot be matched to a known remote file (e.g. missing metadata or
|
||||||
|
renamed files) are still reported with ``fileId`` set to None.
|
||||||
|
Returns ``[{fileId, fileName, filePath}]``.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
cache = await scanner.get_cached_data()
|
||||||
|
except Exception: # pragma: no cover - defensive fallback
|
||||||
|
logger.debug(
|
||||||
|
"Failed to read cache for downloaded files of version %s",
|
||||||
|
model_version_id,
|
||||||
|
exc_info=True,
|
||||||
|
)
|
||||||
|
return []
|
||||||
|
|
||||||
|
files_getter = getattr(cache, "get_files_by_version_id", None)
|
||||||
|
local_entries = files_getter(model_version_id) if files_getter else []
|
||||||
|
if not local_entries:
|
||||||
|
return []
|
||||||
|
|
||||||
|
version_payload: Mapping[str, Any] = {}
|
||||||
|
for entry in local_entries:
|
||||||
|
civitai = entry.get("civitai") if isinstance(entry, Mapping) else None
|
||||||
|
if isinstance(civitai, Mapping) and isinstance(civitai.get("files"), list):
|
||||||
|
version_payload = civitai
|
||||||
|
break
|
||||||
|
|
||||||
|
downloaded = ModelCivitaiHandler._match_downloaded_files(
|
||||||
|
version_payload, local_entries
|
||||||
|
)
|
||||||
|
|
||||||
|
# Surface local files that D2 could not map to a known remote file
|
||||||
|
matched_paths = {item.get("filePath") for item in downloaded}
|
||||||
|
for entry in local_entries:
|
||||||
|
if not isinstance(entry, Mapping):
|
||||||
|
continue
|
||||||
|
if entry.get("file_path") in matched_paths:
|
||||||
|
continue
|
||||||
|
downloaded.append(
|
||||||
|
{
|
||||||
|
"fileId": None,
|
||||||
|
"fileName": entry.get("file_name"),
|
||||||
|
"filePath": entry.get("file_path"),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return downloaded
|
||||||
|
|
||||||
async def check_model_exists(self, request: web.Request) -> web.Response:
|
async def check_model_exists(self, request: web.Request) -> web.Response:
|
||||||
try:
|
try:
|
||||||
model_id_str = request.query.get("modelId")
|
model_id_str = request.query.get("modelId")
|
||||||
@@ -2096,9 +2154,11 @@ class ModelLibraryHandler:
|
|||||||
|
|
||||||
exists = False
|
exists = False
|
||||||
model_type = None
|
model_type = None
|
||||||
|
matched_scanner = None
|
||||||
if await lora_scanner.check_model_version_exists(model_version_id):
|
if await lora_scanner.check_model_version_exists(model_version_id):
|
||||||
exists = True
|
exists = True
|
||||||
model_type = "lora"
|
model_type = "lora"
|
||||||
|
matched_scanner = lora_scanner
|
||||||
elif (
|
elif (
|
||||||
checkpoint_scanner
|
checkpoint_scanner
|
||||||
and await checkpoint_scanner.check_model_version_exists(
|
and await checkpoint_scanner.check_model_version_exists(
|
||||||
@@ -2107,6 +2167,7 @@ class ModelLibraryHandler:
|
|||||||
):
|
):
|
||||||
exists = True
|
exists = True
|
||||||
model_type = "checkpoint"
|
model_type = "checkpoint"
|
||||||
|
matched_scanner = checkpoint_scanner
|
||||||
elif (
|
elif (
|
||||||
embedding_scanner
|
embedding_scanner
|
||||||
and await embedding_scanner.check_model_version_exists(
|
and await embedding_scanner.check_model_version_exists(
|
||||||
@@ -2115,6 +2176,7 @@ class ModelLibraryHandler:
|
|||||||
):
|
):
|
||||||
exists = True
|
exists = True
|
||||||
model_type = "embedding"
|
model_type = "embedding"
|
||||||
|
matched_scanner = embedding_scanner
|
||||||
|
|
||||||
if exists:
|
if exists:
|
||||||
return web.json_response(
|
return web.json_response(
|
||||||
@@ -2123,6 +2185,9 @@ class ModelLibraryHandler:
|
|||||||
"exists": True,
|
"exists": True,
|
||||||
"modelType": model_type,
|
"modelType": model_type,
|
||||||
"hasBeenDownloaded": False,
|
"hasBeenDownloaded": False,
|
||||||
|
"downloadedFiles": await self._get_downloaded_files(
|
||||||
|
matched_scanner, model_version_id
|
||||||
|
),
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -2144,6 +2209,7 @@ class ModelLibraryHandler:
|
|||||||
"exists": False,
|
"exists": False,
|
||||||
"modelType": history_type,
|
"modelType": history_type,
|
||||||
"hasBeenDownloaded": has_been_downloaded,
|
"hasBeenDownloaded": has_been_downloaded,
|
||||||
|
"downloadedFiles": [],
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -908,6 +908,33 @@ class FakeExistenceScanner:
|
|||||||
return []
|
return []
|
||||||
|
|
||||||
|
|
||||||
|
class FakeVersionFilesCache:
|
||||||
|
"""Cache stub exposing the multi-valued version->files index (#1058)."""
|
||||||
|
|
||||||
|
def __init__(self, entries):
|
||||||
|
self._entries = entries
|
||||||
|
|
||||||
|
def get_files_by_version_id(self, _version_id):
|
||||||
|
return list(self._entries)
|
||||||
|
|
||||||
|
|
||||||
|
class FakeDownloadedFilesScanner:
|
||||||
|
"""Scanner stub with one known version and per-file cache entries."""
|
||||||
|
|
||||||
|
def __init__(self, version_id, entries):
|
||||||
|
self._version_id = version_id
|
||||||
|
self._cache = FakeVersionFilesCache(entries)
|
||||||
|
|
||||||
|
async def check_model_version_exists(self, version_id):
|
||||||
|
return version_id == self._version_id
|
||||||
|
|
||||||
|
async def get_model_versions_by_id(self, _model_id):
|
||||||
|
return []
|
||||||
|
|
||||||
|
async def get_cached_data(self):
|
||||||
|
return self._cache
|
||||||
|
|
||||||
|
|
||||||
class FakeMetadataProvider:
|
class FakeMetadataProvider:
|
||||||
async def get_model_versions(self, _model_id):
|
async def get_model_versions(self, _model_id):
|
||||||
return {"modelVersions": [], "name": "", "type": "lora"}
|
return {"modelVersions": [], "name": "", "type": "lora"}
|
||||||
@@ -1524,9 +1551,125 @@ async def test_check_model_exists_returns_download_history_when_file_missing():
|
|||||||
"exists": False,
|
"exists": False,
|
||||||
"modelType": "checkpoint",
|
"modelType": "checkpoint",
|
||||||
"hasBeenDownloaded": True,
|
"hasBeenDownloaded": True,
|
||||||
|
"downloadedFiles": [],
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_check_model_exists_returns_downloaded_files():
|
||||||
|
"""The version branch exposes per-file downloaded state (#1058)."""
|
||||||
|
civitai_payload = {
|
||||||
|
"id": 100,
|
||||||
|
"files": [
|
||||||
|
{
|
||||||
|
"id": 1001,
|
||||||
|
"name": "model-fp16.safetensors",
|
||||||
|
"hashes": {"SHA256": "A" * 64},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"id": 1002,
|
||||||
|
"name": "model-fp32.safetensors",
|
||||||
|
"hashes": {"SHA256": "B" * 64},
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
entries = [
|
||||||
|
{
|
||||||
|
"file_name": "model-fp16",
|
||||||
|
"file_path": "/models/loras/model-fp16.safetensors",
|
||||||
|
"sha256": "a" * 64,
|
||||||
|
"civitai": civitai_payload,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
# Legacy entry without a hash: matched by extension-less name (D2)
|
||||||
|
"file_name": "model-fp32",
|
||||||
|
"file_path": "/models/loras/model-fp32.safetensors",
|
||||||
|
"sha256": "",
|
||||||
|
"civitai": civitai_payload,
|
||||||
|
},
|
||||||
|
]
|
||||||
|
lora_scanner = FakeDownloadedFilesScanner(100, entries)
|
||||||
|
|
||||||
|
async def lora_factory():
|
||||||
|
return lora_scanner
|
||||||
|
|
||||||
|
handler = ModelLibraryHandler(
|
||||||
|
ServiceRegistryAdapter(
|
||||||
|
get_lora_scanner=lora_factory,
|
||||||
|
get_checkpoint_scanner=fake_scanner_factory,
|
||||||
|
get_embedding_scanner=fake_scanner_factory,
|
||||||
|
get_downloaded_version_history_service=fake_download_history_service_factory,
|
||||||
|
),
|
||||||
|
metadata_provider_factory=fake_metadata_provider_factory,
|
||||||
|
)
|
||||||
|
|
||||||
|
response = await handler.check_model_exists(
|
||||||
|
FakeRequest(query={"modelId": "5", "modelVersionId": "100"}) # pyright: ignore[reportArgumentType]
|
||||||
|
)
|
||||||
|
payload = _json_payload(response)
|
||||||
|
|
||||||
|
assert payload == {
|
||||||
|
"success": True,
|
||||||
|
"exists": True,
|
||||||
|
"modelType": "lora",
|
||||||
|
"hasBeenDownloaded": False,
|
||||||
|
"downloadedFiles": [
|
||||||
|
{
|
||||||
|
"fileId": 1001,
|
||||||
|
"fileName": "model-fp16.safetensors",
|
||||||
|
"filePath": "/models/loras/model-fp16.safetensors",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"fileId": 1002,
|
||||||
|
"fileName": "model-fp32.safetensors",
|
||||||
|
"filePath": "/models/loras/model-fp32.safetensors",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_check_model_exists_downloaded_files_without_remote_metadata():
|
||||||
|
"""Unmatchable local files are still reported with fileId None (#1058)."""
|
||||||
|
entries = [
|
||||||
|
{
|
||||||
|
"file_name": "renamed-model",
|
||||||
|
"file_path": "/models/loras/renamed-model.safetensors",
|
||||||
|
"sha256": "c" * 64,
|
||||||
|
"civitai": {"id": 100}, # no remote files list cached
|
||||||
|
},
|
||||||
|
]
|
||||||
|
lora_scanner = FakeDownloadedFilesScanner(100, entries)
|
||||||
|
|
||||||
|
async def lora_factory():
|
||||||
|
return lora_scanner
|
||||||
|
|
||||||
|
handler = ModelLibraryHandler(
|
||||||
|
ServiceRegistryAdapter(
|
||||||
|
get_lora_scanner=lora_factory,
|
||||||
|
get_checkpoint_scanner=fake_scanner_factory,
|
||||||
|
get_embedding_scanner=fake_scanner_factory,
|
||||||
|
get_downloaded_version_history_service=fake_download_history_service_factory,
|
||||||
|
),
|
||||||
|
metadata_provider_factory=fake_metadata_provider_factory,
|
||||||
|
)
|
||||||
|
|
||||||
|
response = await handler.check_model_exists(
|
||||||
|
FakeRequest(query={"modelId": "5", "modelVersionId": "100"}) # pyright: ignore[reportArgumentType]
|
||||||
|
)
|
||||||
|
payload = _json_payload(response)
|
||||||
|
|
||||||
|
assert payload["success"] is True
|
||||||
|
assert payload["exists"] is True
|
||||||
|
assert payload["downloadedFiles"] == [
|
||||||
|
{
|
||||||
|
"fileId": None,
|
||||||
|
"fileName": "renamed-model",
|
||||||
|
"filePath": "/models/loras/renamed-model.safetensors",
|
||||||
|
}
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_model_version_download_status_endpoints():
|
async def test_model_version_download_status_endpoints():
|
||||||
history_service = FakeDownloadHistoryService({"lora": {123}})
|
history_service = FakeDownloadHistoryService({"lora": {123}})
|
||||||
|
|||||||
Reference in New Issue
Block a user