mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-11 03:52:13 -03:00
Merge branch 'fix/example-images-failed-blacklist'
This commit is contained in:
@@ -241,8 +241,41 @@ class MetadataSyncService:
|
||||
raise MetadataPersistError(
|
||||
f"Failed to write metadata sidecar: {metadata_path}"
|
||||
)
|
||||
self._unblock_example_image_failures(local_metadata)
|
||||
return local_metadata
|
||||
|
||||
@staticmethod
|
||||
def _unblock_example_image_failures(local_metadata: Dict[str, Any]) -> None:
|
||||
"""Drop the model from the example-images failed list, if present.
|
||||
|
||||
Fresh provider metadata invalidates any earlier "no example images"
|
||||
verdict, so the next example-image run should retry the model. The
|
||||
unblock only edits the local progress file — no provider API calls —
|
||||
and must never break the metadata update itself.
|
||||
"""
|
||||
|
||||
sha256 = (local_metadata.get("sha256") or "").lower()
|
||||
if not sha256:
|
||||
return
|
||||
|
||||
try:
|
||||
# Lazy import: example_images_download_manager (via
|
||||
# example_images_metadata) imports this service at module level.
|
||||
from ..utils.example_images_download_manager import (
|
||||
unblock_failed_example_image_models,
|
||||
)
|
||||
|
||||
removed = unblock_failed_example_image_models([sha256])
|
||||
if removed:
|
||||
logger.info(
|
||||
"Re-enabled example image downloads for %s after metadata update",
|
||||
local_metadata.get("model_name", sha256),
|
||||
)
|
||||
except Exception as exc: # pragma: no cover - defensive logging
|
||||
logger.debug(
|
||||
"Failed to unblock example image failures for %s: %s", sha256, exc
|
||||
)
|
||||
|
||||
async def fetch_and_update_model(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -15,6 +15,7 @@ from ..utils.example_images_paths import (
|
||||
ExampleImagePathResolver,
|
||||
ensure_library_root_exists,
|
||||
get_example_images_root,
|
||||
get_library_root,
|
||||
is_hash_folder,
|
||||
uses_library_scoped_folders,
|
||||
)
|
||||
@@ -948,11 +949,26 @@ class DownloadManager:
|
||||
|
||||
return True # Return True to indicate a remote download happened
|
||||
else:
|
||||
# No civitai data or images available, mark as failed to avoid future attempts
|
||||
self._progress["failed_models"].add(model_hash)
|
||||
logger.debug(
|
||||
f"No civitai images available for model {model_name}, marking as failed"
|
||||
# No example images available. "Metadata not fetched yet"
|
||||
# (no civitai payload at all) is retryable — a later run only
|
||||
# re-reads local metadata, so no extra provider API calls.
|
||||
# Confirmed absences (provider metadata present but
|
||||
# image-less, or the model is known to be gone from CivitAI)
|
||||
# stay on the failed list to avoid future attempts.
|
||||
metadata_missing = not civitai_payload
|
||||
known_absent = bool((full_model or {}).get("civitai_deleted")) or (
|
||||
(full_model or {}).get("from_civitai") is False
|
||||
)
|
||||
if metadata_missing and not known_absent:
|
||||
logger.debug(
|
||||
"No civitai metadata yet for model %s, leaving it pending for a later run",
|
||||
model_name,
|
||||
)
|
||||
else:
|
||||
self._progress["failed_models"].add(model_hash)
|
||||
logger.debug(
|
||||
f"No civitai images available for model {model_name}, marking as failed"
|
||||
)
|
||||
|
||||
# Save progress periodically
|
||||
if (
|
||||
@@ -968,8 +984,8 @@ class DownloadManager:
|
||||
logger.error(error_msg, exc_info=True)
|
||||
self._progress["errors"].append(error_msg)
|
||||
self._progress["last_error"] = error_msg
|
||||
# Ensure model is marked as failed so we don't try again in this run
|
||||
self._progress["failed_models"].add(model_hash)
|
||||
# Transient failures (network, disk, ...) must not blacklist the
|
||||
# model: leave it pending so the next run retries it.
|
||||
return False
|
||||
|
||||
def _save_progress(self, output_dir):
|
||||
@@ -1546,6 +1562,98 @@ class DownloadManager:
|
||||
_default_download_manager: DownloadManager | None = None
|
||||
|
||||
|
||||
def _progress_file_candidates(library_name: str | None = None) -> List[str]:
|
||||
"""Return existing-progress-file locations for a library (plus the legacy root)."""
|
||||
|
||||
settings_manager = get_settings_manager()
|
||||
if not settings_manager.get("example_images_path"):
|
||||
return []
|
||||
|
||||
candidates: List[str] = []
|
||||
library_root = get_library_root(
|
||||
library_name or settings_manager.get_active_library_name()
|
||||
)
|
||||
if library_root:
|
||||
candidates.append(os.path.join(library_root, ".download_progress.json"))
|
||||
|
||||
if uses_library_scoped_folders():
|
||||
legacy_root = get_example_images_root()
|
||||
if legacy_root:
|
||||
legacy_file = os.path.join(legacy_root, ".download_progress.json")
|
||||
if legacy_file not in candidates:
|
||||
candidates.append(legacy_file)
|
||||
|
||||
return candidates
|
||||
|
||||
|
||||
def _remove_hashes_from_progress_file(progress_file: str, hashes: Set[str]) -> int:
|
||||
"""Remove *hashes* from the failed_models list of one progress file."""
|
||||
|
||||
if not os.path.exists(progress_file):
|
||||
return 0
|
||||
|
||||
try:
|
||||
with open(progress_file, "r", encoding="utf-8") as f:
|
||||
raw = f.read()
|
||||
except OSError:
|
||||
return 0
|
||||
|
||||
# Cheap pre-filter: skip the JSON parse and rewrite when none of the
|
||||
# hashes are even mentioned.
|
||||
if not any(h in raw for h in hashes):
|
||||
return 0
|
||||
|
||||
try:
|
||||
data = json.loads(raw)
|
||||
except ValueError:
|
||||
return 0
|
||||
|
||||
failed = data.get("failed_models")
|
||||
if not isinstance(failed, list) or not failed:
|
||||
return 0
|
||||
|
||||
remaining = [h for h in failed if (h or "").lower() not in hashes]
|
||||
removed = len(failed) - len(remaining)
|
||||
if not removed:
|
||||
return 0
|
||||
|
||||
data["failed_models"] = remaining
|
||||
try:
|
||||
with open(progress_file, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
except OSError:
|
||||
return 0
|
||||
|
||||
return removed
|
||||
|
||||
|
||||
def unblock_failed_example_image_models(
|
||||
model_hashes: Iterable[str], library_name: str | None = None
|
||||
) -> int:
|
||||
"""Remove hashes from the failed-models list so the next run retries them.
|
||||
|
||||
Called when fresh provider metadata lands for a model (metadata fetch or
|
||||
relink), which invalidates any earlier "no example images" verdict. Only
|
||||
rewrites the local progress file — no provider API calls. Also updates the
|
||||
live download manager's in-memory state when one exists.
|
||||
"""
|
||||
|
||||
normalized = {h.lower() for h in model_hashes if h}
|
||||
if not normalized:
|
||||
return 0
|
||||
|
||||
manager = _default_download_manager
|
||||
if manager is not None:
|
||||
failed = manager._progress.get("failed_models")
|
||||
if isinstance(failed, set):
|
||||
failed.difference_update(normalized)
|
||||
|
||||
removed = 0
|
||||
for progress_file in _progress_file_candidates(library_name):
|
||||
removed += _remove_hashes_from_progress_file(progress_file, normalized)
|
||||
return removed
|
||||
|
||||
|
||||
def get_default_download_manager(ws_manager) -> DownloadManager:
|
||||
"""Return the singleton download manager used by default routes."""
|
||||
|
||||
|
||||
@@ -795,3 +795,225 @@ async def test_explicit_targets_fill_partial_example_gaps(
|
||||
@pytest.fixture
|
||||
def settings_manager():
|
||||
return get_settings_manager()
|
||||
|
||||
|
||||
async def _run_single_model_download(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
settings_manager,
|
||||
model_metadata: dict[str, Any],
|
||||
*,
|
||||
get_updated_model=None,
|
||||
):
|
||||
"""Shared harness: run one download pass over a single lora model."""
|
||||
|
||||
ws_manager = RecordingWebSocketManager()
|
||||
manager = download_module.DownloadManager(ws_manager=ws_manager)
|
||||
|
||||
images_root = tmp_path / "examples"
|
||||
monkeypatch.setitem(settings_manager.settings, "example_images_path", str(images_root))
|
||||
monkeypatch.setitem(settings_manager.settings, "libraries", {"default": {}})
|
||||
monkeypatch.setitem(settings_manager.settings, "active_library", "default")
|
||||
|
||||
scanner = StubScanner([model_metadata.copy()])
|
||||
_patch_scanner(monkeypatch, scanner)
|
||||
|
||||
async def fake_process_local_examples(*_args, **_kwargs):
|
||||
return False
|
||||
|
||||
if get_updated_model is None:
|
||||
async def get_updated_model(_hash, _scanner):
|
||||
return model_metadata
|
||||
|
||||
async def fake_get_downloader():
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(
|
||||
download_module.ExampleImagesProcessor,
|
||||
"process_local_examples",
|
||||
staticmethod(fake_process_local_examples),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
download_module.MetadataUpdater,
|
||||
"get_updated_model",
|
||||
staticmethod(get_updated_model),
|
||||
)
|
||||
monkeypatch.setattr(download_module, "get_downloader", fake_get_downloader)
|
||||
|
||||
result = await manager.start_download({"model_types": ["lora"], "delay": 0})
|
||||
assert result["success"] is True
|
||||
if manager._download_task is not None:
|
||||
await asyncio.wait_for(manager._download_task, timeout=1)
|
||||
|
||||
return manager, images_root
|
||||
|
||||
|
||||
async def test_model_without_civitai_metadata_stays_pending(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
settings_manager,
|
||||
):
|
||||
"""Metadata not fetched yet is retryable: the model must NOT be blacklisted."""
|
||||
|
||||
model_hash = "c" * 64
|
||||
model_path = tmp_path / "model.safetensors"
|
||||
model_path.write_text("data", encoding="utf-8")
|
||||
|
||||
model_metadata = {
|
||||
"sha256": model_hash,
|
||||
"model_name": "No Metadata Yet",
|
||||
"file_path": str(model_path),
|
||||
"file_name": "model.safetensors",
|
||||
}
|
||||
|
||||
manager, images_root = await _run_single_model_download(
|
||||
monkeypatch, tmp_path, settings_manager, model_metadata
|
||||
)
|
||||
|
||||
assert model_hash not in manager._progress["failed_models"]
|
||||
assert model_hash not in manager._progress["processed_models"]
|
||||
|
||||
saved = json.loads((images_root / ".download_progress.json").read_text(encoding="utf-8"))
|
||||
assert model_hash not in saved["failed_models"]
|
||||
|
||||
|
||||
async def test_model_with_imageless_civitai_metadata_marked_failed(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
settings_manager,
|
||||
):
|
||||
"""Provider metadata exists but carries no images: permanently failed."""
|
||||
|
||||
model_hash = "d" * 64
|
||||
model_path = tmp_path / "model.safetensors"
|
||||
model_path.write_text("data", encoding="utf-8")
|
||||
|
||||
model_metadata = {
|
||||
"sha256": model_hash,
|
||||
"model_name": "Imageless Entry",
|
||||
"file_path": str(model_path),
|
||||
"file_name": "model.safetensors",
|
||||
"civitai": {"id": 123, "name": "Imageless Entry"},
|
||||
}
|
||||
|
||||
manager, _ = await _run_single_model_download(
|
||||
monkeypatch, tmp_path, settings_manager, model_metadata
|
||||
)
|
||||
|
||||
assert model_hash in manager._progress["failed_models"]
|
||||
|
||||
|
||||
async def test_civitai_deleted_model_marked_failed(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
settings_manager,
|
||||
):
|
||||
"""Models confirmed absent from CivitAI stay on the failed list."""
|
||||
|
||||
model_hash = "e" * 64
|
||||
model_path = tmp_path / "model.safetensors"
|
||||
model_path.write_text("data", encoding="utf-8")
|
||||
|
||||
model_metadata = {
|
||||
"sha256": model_hash,
|
||||
"model_name": "Deleted From CivitAI",
|
||||
"file_path": str(model_path),
|
||||
"file_name": "model.safetensors",
|
||||
"civitai_deleted": True,
|
||||
"from_civitai": False,
|
||||
}
|
||||
|
||||
manager, _ = await _run_single_model_download(
|
||||
monkeypatch, tmp_path, settings_manager, model_metadata
|
||||
)
|
||||
|
||||
assert model_hash in manager._progress["failed_models"]
|
||||
|
||||
|
||||
async def test_processing_error_does_not_blacklist_model(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
settings_manager,
|
||||
):
|
||||
"""Transient errors must leave the model pending instead of blacklisting it."""
|
||||
|
||||
model_hash = "f" * 64
|
||||
model_path = tmp_path / "model.safetensors"
|
||||
model_path.write_text("data", encoding="utf-8")
|
||||
|
||||
model_metadata = {
|
||||
"sha256": model_hash,
|
||||
"model_name": "Transient Error",
|
||||
"file_path": str(model_path),
|
||||
"file_name": "model.safetensors",
|
||||
"civitai": {"images": [{"url": "https://example.com/one.png"}]},
|
||||
}
|
||||
|
||||
async def failing_get_updated_model(_hash, _scanner):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
manager, _ = await _run_single_model_download(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
settings_manager,
|
||||
model_metadata,
|
||||
get_updated_model=failing_get_updated_model,
|
||||
)
|
||||
|
||||
assert model_hash not in manager._progress["failed_models"]
|
||||
assert manager._progress["last_error"]
|
||||
|
||||
|
||||
def test_unblock_failed_example_image_models_updates_file_and_memory(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
settings_manager,
|
||||
):
|
||||
images_root = tmp_path / "examples"
|
||||
images_root.mkdir()
|
||||
monkeypatch.setitem(settings_manager.settings, "example_images_path", str(images_root))
|
||||
monkeypatch.setitem(settings_manager.settings, "libraries", {"default": {}})
|
||||
monkeypatch.setitem(settings_manager.settings, "active_library", "default")
|
||||
|
||||
model_hash = "a" * 64
|
||||
other_hash = "b" * 64
|
||||
progress_file = images_root / ".download_progress.json"
|
||||
progress_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"processed_models": [other_hash],
|
||||
"failed_models": [model_hash, other_hash],
|
||||
"rate_limited_models": [],
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
|
||||
manager._progress["failed_models"] = {model_hash, other_hash}
|
||||
monkeypatch.setattr(download_module, "_default_download_manager", manager)
|
||||
|
||||
removed = download_module.unblock_failed_example_image_models([model_hash.upper()])
|
||||
assert removed == 1
|
||||
|
||||
saved = json.loads(progress_file.read_text(encoding="utf-8"))
|
||||
assert saved["failed_models"] == [other_hash]
|
||||
assert saved["processed_models"] == [other_hash]
|
||||
assert manager._progress["failed_models"] == {other_hash}
|
||||
|
||||
# Second call is a no-op and does not rewrite the file
|
||||
assert download_module.unblock_failed_example_image_models([model_hash]) == 0
|
||||
|
||||
|
||||
def test_unblock_failed_example_image_models_without_progress_file(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
settings_manager,
|
||||
):
|
||||
images_root = tmp_path / "examples"
|
||||
images_root.mkdir()
|
||||
monkeypatch.setitem(settings_manager.settings, "example_images_path", str(images_root))
|
||||
monkeypatch.setitem(settings_manager.settings, "libraries", {"default": {}})
|
||||
monkeypatch.setitem(settings_manager.settings, "active_library", "default")
|
||||
|
||||
assert download_module.unblock_failed_example_image_models(["a" * 64]) == 0
|
||||
|
||||
@@ -300,6 +300,57 @@ async def test_update_model_metadata_raises_when_sidecar_write_fails():
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_model_metadata_unblocks_example_image_failures(monkeypatch):
|
||||
"""Persisting fresh provider metadata re-enables example image downloads."""
|
||||
helpers = build_service()
|
||||
|
||||
from py.utils import example_images_download_manager as download_module
|
||||
|
||||
calls: list[list[str]] = []
|
||||
|
||||
def fake_unblock(hashes, library_name=None):
|
||||
calls.append(list(hashes))
|
||||
return len(hashes)
|
||||
|
||||
monkeypatch.setattr(
|
||||
download_module, "unblock_failed_example_image_models", fake_unblock
|
||||
)
|
||||
|
||||
result = await helpers.service.update_model_metadata(
|
||||
"path/to/model.metadata.json",
|
||||
{"sha256": "A" * 64, "model_name": "Local"},
|
||||
{"source": "api", "model": {"name": "Remote"}, "images": []},
|
||||
helpers.default_provider,
|
||||
)
|
||||
|
||||
assert result["model_name"] == "Remote"
|
||||
assert calls == [["a" * 64]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_model_metadata_without_hash_skips_unblock(monkeypatch):
|
||||
helpers = build_service()
|
||||
|
||||
from py.utils import example_images_download_manager as download_module
|
||||
|
||||
def fail_unblock(hashes, library_name=None): # pragma: no cover - must not run
|
||||
raise AssertionError("unblock must not run without a sha256")
|
||||
|
||||
monkeypatch.setattr(
|
||||
download_module, "unblock_failed_example_image_models", fail_unblock
|
||||
)
|
||||
|
||||
result = await helpers.service.update_model_metadata(
|
||||
"path/to/model.metadata.json",
|
||||
{"model_name": "Local"},
|
||||
{"source": "api", "model": {"name": "Remote"}, "images": []},
|
||||
helpers.default_provider,
|
||||
)
|
||||
|
||||
assert result["model_name"] == "Remote"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_and_update_model_reports_sidecar_write_failure(tmp_path):
|
||||
"""Fetch fails the item (and leaves the cache untouched) when the sidecar
|
||||
|
||||
Reference in New Issue
Block a user