Will Miao
f9cb3988aa
- Add an atomic .lm-examples-complete marker; only a folder with media and the marker is skipped, so interrupted or legacy folders are re-checked by the per-image existence path. Clear the marker before each pass. - Treat processed/failed folders without a marker as pending so the auto-download pre-check repairs them, and clear markers when provider metadata is refreshed or custom images are imported. - Bound stale-metadata refreshes at 90s so one hung provider cannot park the serial run; count metadata-missing models as retryable. - Record failure reasons in the progress file, expose failure_summary, and render a no-media/unavailable/retry breakdown in the progress panel with a warning completion toast. - Add marker, partial-repair, timeout, failure-summary and UI tests.
396 lines
14 KiB
Python
396 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
import pytest
|
|
|
|
from py.services.settings_manager import get_settings_manager
|
|
from py.utils import example_images_metadata as metadata_module
|
|
from py.utils.example_images_paths import (
|
|
COMPLETION_MARKER_FILENAME,
|
|
get_model_folder,
|
|
)
|
|
from tests.utils.test_video_dimension_probe import build_mp4, build_webm
|
|
|
|
|
|
class StubScanner:
|
|
def __init__(self, cache_items: List[Dict[str, Any]]) -> None:
|
|
self.cache = SimpleNamespace(raw_data=cache_items)
|
|
self.updates: List[Tuple[str, str, Dict[str, Any]]] = []
|
|
self.sync_updates: List[Tuple[str, Dict[str, Any]]] = []
|
|
|
|
async def get_cached_data(self):
|
|
return self.cache
|
|
|
|
async def update_single_model_cache(self, old_path: str, new_path: str, metadata: Dict[str, Any]) -> bool:
|
|
self.updates.append((old_path, new_path, metadata))
|
|
return True
|
|
|
|
async def sync_cache_from_metadata(self, file_path: str, metadata: Dict[str, Any]) -> bool:
|
|
self.sync_updates.append((file_path, metadata))
|
|
return True
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def patch_metadata_manager(monkeypatch: pytest.MonkeyPatch):
|
|
saved: List[Tuple[str, Dict[str, Any]]] = []
|
|
|
|
async def fake_save(path: str, metadata: Dict[str, Any]) -> bool:
|
|
saved.append((path, metadata.copy()))
|
|
return True
|
|
|
|
class SimpleMetadata:
|
|
def __init__(self, payload: Dict[str, Any]) -> None:
|
|
self._payload = payload
|
|
self._unknown_fields: Dict[str, Any] = {}
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return self._payload.copy()
|
|
|
|
async def fake_load(path: str, *_args: Any, **_kwargs: Any):
|
|
metadata_path = path if path.endswith(".metadata.json") else f"{os.path.splitext(path)[0]}.metadata.json"
|
|
if os.path.exists(metadata_path):
|
|
data = json.loads(Path(metadata_path).read_text(encoding="utf-8"))
|
|
return SimpleMetadata(data), False
|
|
return None, False
|
|
|
|
monkeypatch.setattr(metadata_module.MetadataManager, "save_metadata", staticmethod(fake_save))
|
|
monkeypatch.setattr(metadata_module.MetadataManager, "load_metadata", staticmethod(fake_load))
|
|
return saved
|
|
|
|
|
|
async def test_update_metadata_after_import_enriches_entries(monkeypatch: pytest.MonkeyPatch, tmp_path, patch_metadata_manager):
|
|
model_hash = "a" * 64
|
|
model_file = tmp_path / "model.safetensors"
|
|
model_file.write_text("content", encoding="utf-8")
|
|
model_data = {
|
|
"model_name": "Example",
|
|
"file_path": str(model_file),
|
|
"civitai": {},
|
|
}
|
|
scanner = StubScanner([model_data])
|
|
|
|
image_path = tmp_path / "custom.png"
|
|
image_path.write_bytes(b"fakepng")
|
|
|
|
monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: "Prompt text Negative prompt: bad Steps: 20, Sampler: Euler"))
|
|
monkeypatch.setattr(metadata_module.MetadataUpdater, "_parse_image_metadata", staticmethod(lambda payload: {"prompt": "Prompt text", "negativePrompt": "bad", "parameters": {"Steps": "20"}}))
|
|
|
|
regular, custom = await metadata_module.MetadataUpdater.update_metadata_after_import(
|
|
model_hash,
|
|
model_data,
|
|
scanner,
|
|
[(str(image_path), "short-id")],
|
|
)
|
|
|
|
assert isinstance(custom, list)
|
|
assert custom[0]["id"] == "short-id"
|
|
assert custom[0]["meta"]["prompt"] == "Prompt text"
|
|
assert custom[0]["hasMeta"] is True
|
|
assert custom[0]["type"] == "image"
|
|
|
|
assert Path(patch_metadata_manager[0][0]) == model_file
|
|
assert scanner.sync_updates
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_metadata_after_import_preserves_existing_metadata(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path,
|
|
patch_metadata_manager,
|
|
):
|
|
model_hash = "b" * 64
|
|
model_file = tmp_path / "preserve.safetensors"
|
|
model_file.write_text("content", encoding="utf-8")
|
|
metadata_path = tmp_path / "preserve.metadata.json"
|
|
|
|
existing_payload: Dict[str, Any] = {
|
|
"model_name": "Example",
|
|
"file_path": str(model_file),
|
|
"civitai": {
|
|
"id": 42,
|
|
"modelId": 88,
|
|
"name": "Example",
|
|
"trainedWords": ["foo"],
|
|
"images": [{"url": "https://example.com/default.png", "type": "image"}],
|
|
"customImages": [
|
|
{"id": "existing-id", "type": "image", "url": "", "nsfwLevel": 0}
|
|
],
|
|
},
|
|
"extraField": "keep-me",
|
|
}
|
|
metadata_path.write_text(json.dumps(existing_payload), encoding="utf-8")
|
|
|
|
model_data = {
|
|
"sha256": model_hash,
|
|
"model_name": "Example",
|
|
"file_path": str(model_file),
|
|
"civitai": {
|
|
"id": 42,
|
|
"modelId": 88,
|
|
"name": "Example",
|
|
"trainedWords": ["foo"],
|
|
"customImages": [],
|
|
},
|
|
}
|
|
scanner = StubScanner([model_data])
|
|
|
|
image_path = tmp_path / "new.png"
|
|
image_path.write_bytes(b"fakepng")
|
|
|
|
monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None))
|
|
monkeypatch.setattr(metadata_module.MetadataUpdater, "_parse_image_metadata", staticmethod(lambda payload: None))
|
|
|
|
regular, custom = await metadata_module.MetadataUpdater.update_metadata_after_import(
|
|
model_hash,
|
|
model_data,
|
|
scanner,
|
|
[(str(image_path), "new-id")],
|
|
)
|
|
|
|
assert regular == existing_payload["civitai"]["images"]
|
|
assert any(entry["id"] == "new-id" for entry in custom)
|
|
|
|
saved_path, saved_payload = patch_metadata_manager[-1]
|
|
assert Path(saved_path) == model_file
|
|
assert saved_payload["extraField"] == "keep-me"
|
|
assert saved_payload["civitai"]["images"] == existing_payload["civitai"]["images"]
|
|
assert saved_payload["civitai"]["trainedWords"] == ["foo"]
|
|
assert {entry["id"] for entry in saved_payload["civitai"]["customImages"]} == {"existing-id", "new-id"}
|
|
|
|
assert scanner.sync_updates
|
|
updated_metadata = scanner.sync_updates[-1][1]
|
|
assert updated_metadata["civitai"]["images"] == existing_payload["civitai"]["images"]
|
|
assert {entry["id"] for entry in updated_metadata["civitai"]["customImages"]} == {"existing-id", "new-id"}
|
|
|
|
async def test_refresh_model_metadata_records_failures(monkeypatch: pytest.MonkeyPatch, tmp_path):
|
|
model_hash = "b" * 64
|
|
model_file = tmp_path / "model.safetensors"
|
|
model_file.write_text("content", encoding="utf-8")
|
|
cache_item = {"sha256": model_hash, "file_path": str(model_file)}
|
|
scanner = StubScanner([cache_item])
|
|
|
|
class StubMetadataSync:
|
|
async def fetch_and_update_model(self, **_kwargs):
|
|
return True, None
|
|
|
|
async def fake_hydrate(model_data: Dict[str, Any]) -> Dict[str, Any]:
|
|
model_data["hydrated"] = True
|
|
return model_data
|
|
|
|
monkeypatch.setattr(
|
|
metadata_module.MetadataManager,
|
|
"hydrate_model_data",
|
|
staticmethod(fake_hydrate),
|
|
)
|
|
|
|
monkeypatch.setattr(metadata_module, "_metadata_sync_service", StubMetadataSync())
|
|
|
|
result = await metadata_module.MetadataUpdater.refresh_model_metadata(
|
|
model_hash,
|
|
"Example",
|
|
"lora",
|
|
scanner,
|
|
{"refreshed_models": set(), "errors": [], "last_error": None},
|
|
)
|
|
assert result is True
|
|
assert cache_item["hydrated"] is True
|
|
|
|
|
|
async def test_refresh_model_metadata_clears_completion_marker(
|
|
monkeypatch: pytest.MonkeyPatch, tmp_path
|
|
):
|
|
"""Fresh provider metadata may add images; the folder must be re-checked."""
|
|
|
|
settings = get_settings_manager()
|
|
monkeypatch.setitem(settings.settings, "example_images_path", str(tmp_path))
|
|
monkeypatch.setitem(settings.settings, "libraries", {"default": {}})
|
|
monkeypatch.setitem(settings.settings, "active_library", "default")
|
|
|
|
model_hash = "c" * 64
|
|
model_file = tmp_path / "model.safetensors"
|
|
model_file.write_text("content", encoding="utf-8")
|
|
scanner = StubScanner([{"sha256": model_hash, "file_path": str(model_file)}])
|
|
|
|
model_dir = get_model_folder(model_hash, "default")
|
|
Path(model_dir).mkdir(parents=True, exist_ok=True)
|
|
(Path(model_dir) / "image_0.png").write_text("data", encoding="utf-8")
|
|
marker = Path(model_dir) / COMPLETION_MARKER_FILENAME
|
|
marker.write_text('{"version": 1}', encoding="utf-8")
|
|
|
|
class StubMetadataSync:
|
|
async def fetch_and_update_model(self, **_kwargs):
|
|
return True, None
|
|
|
|
async def fake_hydrate(model_data: Dict[str, Any]) -> Dict[str, Any]:
|
|
return model_data
|
|
|
|
monkeypatch.setattr(
|
|
metadata_module.MetadataManager,
|
|
"hydrate_model_data",
|
|
staticmethod(fake_hydrate),
|
|
)
|
|
monkeypatch.setattr(metadata_module, "_metadata_sync_service", StubMetadataSync())
|
|
|
|
result = await metadata_module.MetadataUpdater.refresh_model_metadata(
|
|
model_hash,
|
|
"Example",
|
|
"lora",
|
|
scanner,
|
|
{"refreshed_models": set(), "errors": [], "last_error": None},
|
|
)
|
|
|
|
assert result is True
|
|
assert marker.exists() is False
|
|
|
|
|
|
async def test_update_metadata_from_local_examples_generates_entries(monkeypatch: pytest.MonkeyPatch, tmp_path):
|
|
model_hash = "c" * 64
|
|
model_dir = tmp_path / model_hash
|
|
model_dir.mkdir()
|
|
(model_dir / "image.png").write_text("data", encoding="utf-8")
|
|
model_data: Dict[str, Any] = {"model_name": "Local", "civitai": {}, "file_path": str(tmp_path / "model.safetensors")}
|
|
|
|
async def fake_save(path, metadata):
|
|
return True
|
|
|
|
monkeypatch.setattr(metadata_module.MetadataManager, "save_metadata", staticmethod(fake_save))
|
|
monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None))
|
|
|
|
success = await metadata_module.MetadataUpdater.update_metadata_from_local_examples(
|
|
model_hash,
|
|
model_data,
|
|
"lora",
|
|
StubScanner([model_data]),
|
|
str(model_dir),
|
|
)
|
|
assert success is True
|
|
assert model_data["civitai"]["images"]
|
|
|
|
|
|
async def test_update_metadata_after_import_uses_real_video_dimensions(
|
|
monkeypatch: pytest.MonkeyPatch, tmp_path, patch_metadata_manager
|
|
):
|
|
"""Regression: imported videos must not fall back to the 720x1280 default.
|
|
|
|
See issue #1115 — landscape videos were stored as portrait, so the showcase
|
|
viewer letterboxed them into a 9:16 container.
|
|
"""
|
|
|
|
model_hash = "d" * 64
|
|
model_file = tmp_path / "video-model.safetensors"
|
|
model_file.write_text("content", encoding="utf-8")
|
|
model_data = {
|
|
"model_name": "VideoExample",
|
|
"file_path": str(model_file),
|
|
"civitai": {},
|
|
}
|
|
scanner = StubScanner([model_data])
|
|
|
|
video_path = tmp_path / "custom_abc.mp4"
|
|
video_path.write_bytes(build_mp4(1280, 720))
|
|
|
|
monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None))
|
|
|
|
_regular, custom = await metadata_module.MetadataUpdater.update_metadata_after_import(
|
|
model_hash,
|
|
model_data,
|
|
scanner,
|
|
[(str(video_path), "abc")],
|
|
)
|
|
|
|
assert custom[0]["type"] == "video"
|
|
assert (custom[0]["width"], custom[0]["height"]) == (1280, 720)
|
|
assert patch_metadata_manager[-1][1]["civitai"]["customImages"][0]["width"] == 1280
|
|
|
|
|
|
async def test_update_metadata_after_import_uses_real_webm_dimensions(
|
|
monkeypatch: pytest.MonkeyPatch, tmp_path, patch_metadata_manager
|
|
):
|
|
model_hash = "e" * 64
|
|
model_file = tmp_path / "webm-model.safetensors"
|
|
model_file.write_text("content", encoding="utf-8")
|
|
model_data = {
|
|
"model_name": "WebmExample",
|
|
"file_path": str(model_file),
|
|
"civitai": {},
|
|
}
|
|
|
|
video_path = tmp_path / "custom_def.webm"
|
|
video_path.write_bytes(build_webm(480, 832))
|
|
|
|
monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None))
|
|
|
|
_regular, custom = await metadata_module.MetadataUpdater.update_metadata_after_import(
|
|
model_hash,
|
|
model_data,
|
|
StubScanner([model_data]),
|
|
[(str(video_path), "def")],
|
|
)
|
|
|
|
assert (custom[0]["width"], custom[0]["height"]) == (480, 832)
|
|
|
|
|
|
async def test_update_metadata_after_import_falls_back_for_unreadable_video(
|
|
monkeypatch: pytest.MonkeyPatch, tmp_path, patch_metadata_manager
|
|
):
|
|
"""An unparsable video keeps the legacy placeholder rather than failing."""
|
|
|
|
model_hash = "f" * 64
|
|
model_file = tmp_path / "broken-model.safetensors"
|
|
model_file.write_text("content", encoding="utf-8")
|
|
model_data = {
|
|
"model_name": "BrokenExample",
|
|
"file_path": str(model_file),
|
|
"civitai": {},
|
|
}
|
|
|
|
video_path = tmp_path / "custom_ghi.mp4"
|
|
video_path.write_bytes(b"\x00\x00\x00\x20ftypisom" + b"\xff" * 32)
|
|
|
|
monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None))
|
|
|
|
_regular, custom = await metadata_module.MetadataUpdater.update_metadata_after_import(
|
|
model_hash,
|
|
model_data,
|
|
StubScanner([model_data]),
|
|
[(str(video_path), "ghi")],
|
|
)
|
|
|
|
assert (custom[0]["width"], custom[0]["height"]) == (720, 1280)
|
|
|
|
|
|
async def test_update_metadata_from_local_examples_uses_real_video_dimensions(
|
|
monkeypatch: pytest.MonkeyPatch, tmp_path
|
|
):
|
|
model_hash = "1" * 64
|
|
model_dir = tmp_path / model_hash
|
|
model_dir.mkdir()
|
|
(model_dir / "clip.mp4").write_bytes(build_mp4(1920, 1080))
|
|
model_data: Dict[str, Any] = {
|
|
"model_name": "LocalVideo",
|
|
"civitai": {},
|
|
"file_path": str(tmp_path / "model.safetensors"),
|
|
}
|
|
|
|
async def fake_save(path, metadata):
|
|
return True
|
|
|
|
monkeypatch.setattr(metadata_module.MetadataManager, "save_metadata", staticmethod(fake_save))
|
|
|
|
success = await metadata_module.MetadataUpdater.update_metadata_from_local_examples(
|
|
model_hash,
|
|
model_data,
|
|
"lora",
|
|
StubScanner([model_data]),
|
|
str(model_dir),
|
|
)
|
|
|
|
assert success is True
|
|
entry = model_data["civitai"]["images"][0]
|
|
assert entry["type"] == "video"
|
|
assert (entry["width"], entry["height"]) == (1920, 1080) |