mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
Will Miao
7d963b27b5
Example videos added through the "Add examples" flow were stored with a hardcoded 720x1280 entry. The dimension probe next to it only ran for images (PIL cannot open .mp4/.webm files), so every video entry stayed portrait regardless of the source. The showcase viewer then sizes its container straight from that value (--media-aspect in showcase.css), so landscape clips were letterboxed inside a 9:16 box. CivitAI-sourced examples were unaffected because their dimensions come from the API. PIL cannot read video containers, so add a dependency-free reader that parses the container headers instead: moov/trak/tkhd for ISO base media (with the sample description as a fallback), Segment/Tracks/Pixel* for WebM/Matroska, and RIFF/WebP for animated examples saved with a video extension. The sniffed signature decides which reader runs, so a .mp4 that is really WebM still reports the right size; the extension is only a fallback. Both readers seek past mdat rather than reading it, so a large file costs the same as a small one. Imported entries now record the file's real size and keep the previous placeholder only when the file cannot be parsed. Existing libraries keep their wrong entries, so backfill them once via the existing naming migration: bump CURRENT_NAMING_VERSION to 3 and repair each model's empty-url entries from the files on disk, then sync the scanner cache. Only entries with no remote url are touched -- those have no other source, which makes the rewrite lossless -- and entries already carrying the right size are left byte-identical, so the pass is idempotent and a no-op for libraries that never imported a video.
344 lines
12 KiB
Python
344 lines
12 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.utils import example_images_metadata as metadata_module
|
|
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_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) |