mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 03:01:27 -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.
333 lines
10 KiB
Python
333 lines
10 KiB
Python
import asyncio
|
|
import copy
|
|
import json
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from py.utils import example_images_migration as migration_module
|
|
|
|
|
|
class FakeScanner:
|
|
def __init__(self, data_map, init_cycles=0):
|
|
self._data_map = data_map
|
|
self._init_cycles = init_cycles
|
|
self.update_calls = []
|
|
|
|
def is_initializing(self):
|
|
if self._init_cycles > 0:
|
|
self._init_cycles -= 1
|
|
return True
|
|
return False
|
|
|
|
def has_hash(self, hash_value):
|
|
return hash_value in self._data_map
|
|
|
|
async def get_cached_data(self):
|
|
payload = [copy.deepcopy(item) for item in self._data_map.values()]
|
|
return SimpleNamespace(raw_data=payload)
|
|
|
|
async def update_single_model_cache(self, *args):
|
|
self.update_calls.append(args)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_migrations_run_and_update_progress(tmp_path, monkeypatch):
|
|
example_root = tmp_path / "example_images"
|
|
library_root = example_root / "main"
|
|
missing_library = tmp_path / "missing"
|
|
library_root.mkdir(parents=True)
|
|
|
|
progress_path = library_root / ".download_progress.json"
|
|
progress_path.write_text(json.dumps({"naming_version": 0, "keep": "value"}))
|
|
|
|
migrating_hash = "a" * 64
|
|
migrating_folder = library_root / migrating_hash
|
|
migrating_folder.mkdir()
|
|
(migrating_folder / "image_1.png").write_bytes(b"one")
|
|
(migrating_folder / "image_2.png").write_bytes(b"two")
|
|
|
|
zero_based_hash = "b" * 64
|
|
zero_based_folder = library_root / zero_based_hash
|
|
zero_based_folder.mkdir()
|
|
(zero_based_folder / "image_0.png").write_bytes(b"zero")
|
|
(zero_based_folder / "image_1.png").write_bytes(b"one")
|
|
|
|
rename_only_hash = "c" * 64
|
|
rename_only_folder = library_root / rename_only_hash
|
|
rename_only_folder.mkdir()
|
|
(rename_only_folder / "image_1.png").write_bytes(b"needs rename")
|
|
|
|
metadata_path = tmp_path / "models" / "model.safetensors"
|
|
metadata_path.parent.mkdir()
|
|
|
|
metadata_entry = {
|
|
"sha256": migrating_hash,
|
|
"file_path": str(metadata_path),
|
|
"civitai": {
|
|
"images": [
|
|
{"url": "", "type": "image", "prompt": "custom"},
|
|
{"url": "https://example.com/remote.png", "type": "image"},
|
|
],
|
|
"customImages": [],
|
|
},
|
|
}
|
|
|
|
lora_scanner = FakeScanner({migrating_hash: metadata_entry}, init_cycles=1)
|
|
checkpoint_scanner = FakeScanner({}, init_cycles=0)
|
|
|
|
async def fake_get_lora_scanner(cls):
|
|
return lora_scanner
|
|
|
|
async def fake_get_checkpoint_scanner(cls):
|
|
return checkpoint_scanner
|
|
|
|
monkeypatch.setattr(
|
|
migration_module.ServiceRegistry,
|
|
"get_lora_scanner",
|
|
classmethod(fake_get_lora_scanner),
|
|
)
|
|
monkeypatch.setattr(
|
|
migration_module.ServiceRegistry,
|
|
"get_checkpoint_scanner",
|
|
classmethod(fake_get_checkpoint_scanner),
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
migration_module.settings,
|
|
"get",
|
|
lambda key, default=None: str(example_root) if key == "example_images_path" else default,
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
migration_module,
|
|
"iter_library_roots",
|
|
lambda: [("main", str(library_root)), ("missing", str(missing_library))],
|
|
)
|
|
|
|
saved_metadata = []
|
|
|
|
async def fake_save_metadata(path, metadata):
|
|
saved_metadata.append((path, metadata))
|
|
return True
|
|
|
|
monkeypatch.setattr(
|
|
migration_module.MetadataManager,
|
|
"save_metadata",
|
|
staticmethod(fake_save_metadata),
|
|
)
|
|
|
|
short_ids = iter(["short1234"])
|
|
monkeypatch.setattr(
|
|
migration_module.ExampleImagesProcessor,
|
|
"generate_short_id",
|
|
staticmethod(lambda: next(short_ids)),
|
|
)
|
|
|
|
sleep_calls = []
|
|
|
|
async def fake_sleep(delay):
|
|
sleep_calls.append(delay)
|
|
return None
|
|
|
|
monkeypatch.setattr(migration_module.asyncio, "sleep", fake_sleep)
|
|
|
|
scheduled_tasks = []
|
|
original_create_task = asyncio.create_task
|
|
|
|
def capture_create_task(coro, *args, **kwargs):
|
|
task = original_create_task(coro, *args, **kwargs)
|
|
scheduled_tasks.append(task)
|
|
return task
|
|
|
|
monkeypatch.setattr(migration_module.asyncio, "create_task", capture_create_task)
|
|
|
|
await migration_module.ExampleImagesMigration.check_and_run_migrations()
|
|
|
|
assert len(scheduled_tasks) == 1
|
|
await asyncio.gather(*scheduled_tasks)
|
|
|
|
progress_data = json.loads(progress_path.read_text())
|
|
assert progress_data["naming_version"] == migration_module.CURRENT_NAMING_VERSION
|
|
assert progress_data["keep"] == "value"
|
|
|
|
assert not (migrating_folder / "image_0.png").exists()
|
|
assert (migrating_folder / "custom_short1234.png").exists()
|
|
assert (migrating_folder / "image_1.png").exists()
|
|
|
|
assert (zero_based_folder / "image_0.png").exists()
|
|
assert (zero_based_folder / "image_1.png").exists()
|
|
|
|
assert (rename_only_folder / "image_0.png").exists()
|
|
assert not (rename_only_folder / "image_1.png").exists()
|
|
|
|
assert not missing_library.exists()
|
|
|
|
assert any(delay == 1 for delay in sleep_calls)
|
|
|
|
assert len(saved_metadata) == 1
|
|
saved_path, saved_payload = saved_metadata[0]
|
|
assert saved_path == str(metadata_path)
|
|
assert saved_payload["civitai"]["customImages"][0]["id"] == "short1234"
|
|
assert saved_payload["civitai"]["images"] == [
|
|
{"url": "https://example.com/remote.png", "type": "image"}
|
|
]
|
|
|
|
assert len(lora_scanner.update_calls) == 1
|
|
update_args = lora_scanner.update_calls[0]
|
|
assert update_args[0] == str(metadata_path)
|
|
assert update_args[2]["civitai"]["customImages"][0]["id"] == "short1234"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_v2_to_v3_migration_repairs_video_dimensions(tmp_path, monkeypatch):
|
|
"""Upgrading a library already at v2 backfills local video dimensions once.
|
|
|
|
This mirrors the real upgrade path for issue #1115: the naming migration is
|
|
already done, but imported videos still carry the 720x1280 placeholder.
|
|
"""
|
|
|
|
from tests.utils.test_video_dimension_probe import build_mp4
|
|
|
|
example_root = tmp_path / "example_images"
|
|
library_root = example_root / "main"
|
|
library_root.mkdir(parents=True)
|
|
|
|
progress_path = library_root / ".download_progress.json"
|
|
progress_path.write_text(json.dumps({"naming_version": 2}))
|
|
|
|
model_hash = "d" * 64
|
|
model_folder = library_root / model_hash
|
|
model_folder.mkdir()
|
|
# Landscape clip stored during the buggy import path.
|
|
(model_folder / "custom_land1.mp4").write_bytes(build_mp4(1280, 720))
|
|
|
|
model_file = tmp_path / "models" / "video.safetensors"
|
|
model_file.parent.mkdir()
|
|
model_file.write_text("weights", encoding="utf-8")
|
|
|
|
scanner = FakeScanner(
|
|
{
|
|
model_hash: {
|
|
"sha256": model_hash,
|
|
"file_path": str(model_file),
|
|
"civitai": {
|
|
"images": [
|
|
{"url": "https://example.com/remote.jpg", "type": "image", "width": 512, "height": 512}
|
|
],
|
|
"customImages": [
|
|
{"url": "", "id": "land1", "type": "video", "width": 720, "height": 1280}
|
|
],
|
|
},
|
|
}
|
|
}
|
|
)
|
|
|
|
async def fake_get_lora_scanner(cls):
|
|
return scanner
|
|
|
|
async def fake_get_checkpoint_scanner(cls):
|
|
return FakeScanner({})
|
|
|
|
monkeypatch.setattr(
|
|
migration_module.ServiceRegistry, "get_lora_scanner", classmethod(fake_get_lora_scanner)
|
|
)
|
|
monkeypatch.setattr(
|
|
migration_module.ServiceRegistry,
|
|
"get_checkpoint_scanner",
|
|
classmethod(fake_get_checkpoint_scanner),
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
migration_module.settings,
|
|
"get",
|
|
lambda key, default=None: str(example_root) if key == "example_images_path" else default,
|
|
)
|
|
monkeypatch.setattr(
|
|
migration_module,
|
|
"iter_library_roots",
|
|
lambda: [("main", str(library_root))],
|
|
)
|
|
|
|
saved_metadata = []
|
|
|
|
async def fake_save_metadata(path, metadata):
|
|
saved_metadata.append((path, metadata))
|
|
return True
|
|
|
|
async def fake_load_payload(path):
|
|
return {
|
|
"model_name": "Video",
|
|
"civitai": {
|
|
"images": [
|
|
{"url": "https://example.com/remote.jpg", "type": "image", "width": 512, "height": 512}
|
|
],
|
|
"customImages": [
|
|
{"url": "", "id": "land1", "type": "video", "width": 720, "height": 1280}
|
|
],
|
|
},
|
|
}
|
|
|
|
monkeypatch.setattr(
|
|
migration_module.MetadataManager, "save_metadata", staticmethod(fake_save_metadata)
|
|
)
|
|
monkeypatch.setattr(
|
|
migration_module.MetadataManager, "load_metadata_payload", staticmethod(fake_load_payload)
|
|
)
|
|
|
|
scheduled = []
|
|
original_create_task = asyncio.create_task
|
|
|
|
def capture_create_task(coro, *args, **kwargs):
|
|
task = original_create_task(coro, *args, **kwargs)
|
|
scheduled.append(task)
|
|
return task
|
|
|
|
monkeypatch.setattr(migration_module.asyncio, "create_task", capture_create_task)
|
|
|
|
await migration_module.ExampleImagesMigration.check_and_run_migrations()
|
|
await asyncio.gather(*scheduled)
|
|
|
|
assert len(saved_metadata) == 1
|
|
_path, payload = saved_metadata[0]
|
|
entry = payload["civitai"]["customImages"][0]
|
|
assert (entry["width"], entry["height"]) == (1280, 720)
|
|
# Remote-backed entry is untouched.
|
|
assert payload["civitai"]["images"][0]["width"] == 512
|
|
|
|
assert json.loads(progress_path.read_text())["naming_version"] == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_v3_migration_does_not_run_twice(tmp_path, monkeypatch):
|
|
"""The version gate keeps the repair off the startup path after one run."""
|
|
|
|
example_root = tmp_path / "example_images"
|
|
library_root = example_root / "main"
|
|
library_root.mkdir(parents=True)
|
|
(library_root / ".download_progress.json").write_text(json.dumps({"naming_version": 3}))
|
|
|
|
monkeypatch.setattr(
|
|
migration_module.settings,
|
|
"get",
|
|
lambda key, default=None: str(example_root) if key == "example_images_path" else default,
|
|
)
|
|
monkeypatch.setattr(
|
|
migration_module,
|
|
"iter_library_roots",
|
|
lambda: [("main", str(library_root))],
|
|
)
|
|
|
|
called = []
|
|
|
|
async def spy_run_migrations(*args, **kwargs):
|
|
called.append(args)
|
|
|
|
monkeypatch.setattr(
|
|
migration_module.ExampleImagesMigration, "run_migrations", staticmethod(spy_run_migrations)
|
|
)
|
|
|
|
await migration_module.ExampleImagesMigration.check_and_run_migrations()
|
|
|
|
assert called == []
|