mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
fix(example-images): read real dimensions for imported videos, fixes #1115
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.
This commit is contained in:
@@ -9,6 +9,7 @@ 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:
|
||||
@@ -216,4 +217,128 @@ async def test_update_metadata_from_local_examples_generates_entries(monkeypatch
|
||||
str(model_dir),
|
||||
)
|
||||
assert success is True
|
||||
assert model_data["civitai"]["images"]
|
||||
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)
|
||||
@@ -177,3 +177,156 @@ async def test_migrations_run_and_update_progress(tmp_path, monkeypatch):
|
||||
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 == []
|
||||
|
||||
@@ -0,0 +1,299 @@
|
||||
"""Tests for the one-shot repair of locally imported video dimensions (issue #1115)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict
|
||||
|
||||
import pytest
|
||||
|
||||
from py.utils import example_images_migration as migration_module
|
||||
from py.utils import example_images_metadata as metadata_module
|
||||
from tests.utils.test_video_dimension_probe import build_mp4
|
||||
|
||||
|
||||
def _metadata_payload(**civitai: Any) -> Dict[str, Any]:
|
||||
return {"model_name": "Example", "civitai": civitai}
|
||||
|
||||
|
||||
def test_repair_backfills_landscape_video_dimensions(tmp_path: Path):
|
||||
video = tmp_path / "custom_abc123.mp4"
|
||||
video.write_bytes(build_mp4(1280, 720))
|
||||
|
||||
payload = _metadata_payload(
|
||||
customImages=[
|
||||
{
|
||||
"url": "",
|
||||
"id": "abc123",
|
||||
"type": "video",
|
||||
"width": 720,
|
||||
"height": 1280,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
repaired = metadata_module.repair_local_video_dimensions(
|
||||
payload, {"abc123": str(video)}
|
||||
)
|
||||
|
||||
assert repaired == 1
|
||||
entry = payload["civitai"]["customImages"][0]
|
||||
assert (entry["width"], entry["height"]) == (1280, 720)
|
||||
|
||||
|
||||
def test_repair_handles_index_marked_images_array(tmp_path: Path):
|
||||
video = tmp_path / "image_3.mp4"
|
||||
video.write_bytes(build_mp4(1920, 1080))
|
||||
|
||||
payload = _metadata_payload(
|
||||
images=[
|
||||
{"url": "https://example.com/remote.png", "type": "image"},
|
||||
{"url": "", "type": "video", "width": 720, "height": 1280},
|
||||
{"url": "", "type": "video", "width": 720, "height": 1280},
|
||||
{"url": "", "type": "video", "width": 720, "height": 1280},
|
||||
]
|
||||
)
|
||||
|
||||
repaired = metadata_module.repair_local_video_dimensions(
|
||||
payload, {"3": str(video)}
|
||||
)
|
||||
|
||||
assert repaired == 1
|
||||
# Position 3 (index 3) is the one carrying the local file.
|
||||
assert payload["civitai"]["images"][3]["width"] == 1920
|
||||
assert payload["civitai"]["images"][3]["height"] == 1080
|
||||
# The remote entry keeps its API-provided shape.
|
||||
assert payload["civitai"]["images"][0].get("width") is None
|
||||
|
||||
|
||||
def test_repair_never_touches_remote_entries(tmp_path: Path):
|
||||
"""Remote entries keep API-provided dimensions even if a file exists."""
|
||||
|
||||
video = tmp_path / "custom_remote.mp4"
|
||||
video.write_bytes(build_mp4(1280, 720))
|
||||
|
||||
payload = _metadata_payload(
|
||||
customImages=[
|
||||
{
|
||||
"url": "https://civitai.com/1234.mp4",
|
||||
"id": "remote",
|
||||
"type": "video",
|
||||
"width": 720,
|
||||
"height": 1280,
|
||||
}
|
||||
]
|
||||
)
|
||||
before = json.dumps(payload, sort_keys=True)
|
||||
|
||||
repaired = metadata_module.repair_local_video_dimensions(
|
||||
payload, {"remote": str(video)}
|
||||
)
|
||||
|
||||
assert repaired == 0
|
||||
assert json.dumps(payload, sort_keys=True) == before
|
||||
|
||||
|
||||
def test_repair_is_idempotent(tmp_path: Path):
|
||||
video = tmp_path / "custom_abc.mp4"
|
||||
video.write_bytes(build_mp4(1280, 720))
|
||||
|
||||
payload = _metadata_payload(
|
||||
customImages=[{"url": "", "id": "abc", "type": "video", "width": 720, "height": 1280}]
|
||||
)
|
||||
files = {"abc": str(video)}
|
||||
|
||||
assert metadata_module.repair_local_video_dimensions(payload, files) == 1
|
||||
# Second run finds nothing to do and leaves the entry byte-identical.
|
||||
snapshot = json.dumps(payload, sort_keys=True)
|
||||
assert metadata_module.repair_local_video_dimensions(payload, files) == 0
|
||||
assert json.dumps(payload, sort_keys=True) == snapshot
|
||||
|
||||
|
||||
def test_repair_dry_run_does_not_mutate(tmp_path: Path):
|
||||
video = tmp_path / "custom_abc.mp4"
|
||||
video.write_bytes(build_mp4(1280, 720))
|
||||
|
||||
payload = _metadata_payload(
|
||||
customImages=[{"url": "", "id": "abc", "type": "video", "width": 720, "height": 1280}]
|
||||
)
|
||||
before = json.dumps(payload, sort_keys=True)
|
||||
|
||||
repaired = metadata_module.repair_local_video_dimensions(
|
||||
payload, {"abc": str(video)}, dry_run=True
|
||||
)
|
||||
|
||||
assert repaired == 1
|
||||
assert json.dumps(payload, sort_keys=True) == before
|
||||
|
||||
|
||||
def test_repair_skips_missing_file(tmp_path: Path):
|
||||
payload = _metadata_payload(
|
||||
customImages=[{"url": "", "id": "gone", "type": "video", "width": 720, "height": 1280}]
|
||||
)
|
||||
|
||||
repaired = metadata_module.repair_local_video_dimensions(
|
||||
payload, {"gone": str(tmp_path / "does-not-exist.mp4")}
|
||||
)
|
||||
|
||||
assert repaired == 0
|
||||
assert payload["civitai"]["customImages"][0]["width"] == 720
|
||||
|
||||
|
||||
def test_repair_leaves_correct_entries_untouched(tmp_path: Path):
|
||||
video = tmp_path / "custom_ok.mp4"
|
||||
video.write_bytes(build_mp4(1280, 720))
|
||||
|
||||
payload = _metadata_payload(
|
||||
customImages=[{"url": "", "id": "ok", "type": "video", "width": 1280, "height": 720}]
|
||||
)
|
||||
|
||||
assert metadata_module.repair_local_video_dimensions(payload, {"ok": str(video)}) == 0
|
||||
|
||||
|
||||
def test_local_file_map_keys_strip_naming_prefix(tmp_path: Path):
|
||||
(tmp_path / "custom_abc.mp4").write_bytes(build_mp4(1280, 720))
|
||||
(tmp_path / "image_2.png").write_bytes(b"not-a-real-image")
|
||||
(tmp_path / "notes.txt").write_text("ignore me", encoding="utf-8")
|
||||
|
||||
mapping = migration_module.ExampleImagesMigration._build_local_file_map(str(tmp_path))
|
||||
|
||||
assert set(mapping) == {"abc", "2"}
|
||||
|
||||
|
||||
async def test_migrate_to_v3_repairs_and_syncs_cache(tmp_path: Path, monkeypatch):
|
||||
model_hash = "a" * 64
|
||||
folder = tmp_path / model_hash
|
||||
folder.mkdir()
|
||||
(folder / "custom_xyz.mp4").write_bytes(build_mp4(1080, 1920))
|
||||
|
||||
model_file = tmp_path / "model.safetensors"
|
||||
model_file.write_text("weights", encoding="utf-8")
|
||||
|
||||
payload = _metadata_payload(
|
||||
customImages=[{"url": "", "id": "xyz", "type": "video", "width": 720, "height": 1280}]
|
||||
)
|
||||
saved: list[tuple[str, Dict[str, Any]]] = []
|
||||
|
||||
async def fake_load(file_path):
|
||||
return dict(payload, civitai=dict(payload["civitai"]))
|
||||
|
||||
async def fake_save(file_path, data):
|
||||
saved.append((file_path, data))
|
||||
return True
|
||||
|
||||
synced: list[tuple[str, Dict[str, Any]]] = []
|
||||
|
||||
async def fake_sync(scanner, file_path, data):
|
||||
synced.append((file_path, data))
|
||||
return True
|
||||
|
||||
class StubScanner:
|
||||
def has_hash(self, _hash):
|
||||
return True
|
||||
|
||||
async def get_cached_data(self):
|
||||
from types import SimpleNamespace
|
||||
|
||||
return SimpleNamespace(raw_data=[{"sha256": model_hash, "file_path": str(model_file)}])
|
||||
|
||||
monkeypatch.setattr(migration_module.MetadataManager, "load_metadata_payload", fake_load)
|
||||
monkeypatch.setattr(migration_module.MetadataManager, "save_metadata", fake_save)
|
||||
monkeypatch.setattr(migration_module, "update_cache_from_metadata", fake_sync)
|
||||
|
||||
async def fake_lora():
|
||||
return StubScanner()
|
||||
|
||||
async def fake_none():
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(migration_module.ServiceRegistry, "get_lora_scanner", fake_lora)
|
||||
monkeypatch.setattr(migration_module.ServiceRegistry, "get_checkpoint_scanner", fake_none)
|
||||
monkeypatch.setattr(migration_module.ServiceRegistry, "get_embedding_scanner", fake_none)
|
||||
|
||||
await migration_module.ExampleImagesMigration._migrate_to_v3(
|
||||
str(tmp_path), [str(folder)]
|
||||
)
|
||||
|
||||
assert len(saved) == 1
|
||||
saved_entry = saved[0][1]["civitai"]["customImages"][0]
|
||||
assert (saved_entry["width"], saved_entry["height"]) == (1080, 1920)
|
||||
assert len(synced) == 1
|
||||
assert synced[0][1]["civitai"]["customImages"][0]["width"] == 1080
|
||||
|
||||
|
||||
async def test_migrate_to_v3_skips_when_nothing_to_repair(tmp_path: Path, monkeypatch):
|
||||
model_hash = "b" * 64
|
||||
folder = tmp_path / model_hash
|
||||
folder.mkdir()
|
||||
(folder / "custom_ok.mp4").write_bytes(build_mp4(1080, 1920))
|
||||
|
||||
model_file = tmp_path / "model.safetensors"
|
||||
model_file.write_text("weights", encoding="utf-8")
|
||||
|
||||
payload = _metadata_payload(
|
||||
customImages=[{"url": "", "id": "ok", "type": "video", "width": 1080, "height": 1920}]
|
||||
)
|
||||
saved: list[Any] = []
|
||||
|
||||
async def fake_load(file_path):
|
||||
return dict(payload, civitai=dict(payload["civitai"]))
|
||||
|
||||
async def fake_save(file_path, data):
|
||||
saved.append(data)
|
||||
return True
|
||||
|
||||
class StubScanner:
|
||||
def has_hash(self, _hash):
|
||||
return True
|
||||
|
||||
async def get_cached_data(self):
|
||||
from types import SimpleNamespace
|
||||
|
||||
return SimpleNamespace(raw_data=[{"sha256": model_hash, "file_path": str(model_file)}])
|
||||
|
||||
monkeypatch.setattr(migration_module.MetadataManager, "load_metadata_payload", fake_load)
|
||||
monkeypatch.setattr(migration_module.MetadataManager, "save_metadata", fake_save)
|
||||
|
||||
async def fake_lora():
|
||||
return StubScanner()
|
||||
|
||||
async def fake_none():
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(migration_module.ServiceRegistry, "get_lora_scanner", fake_lora)
|
||||
monkeypatch.setattr(migration_module.ServiceRegistry, "get_checkpoint_scanner", fake_none)
|
||||
monkeypatch.setattr(migration_module.ServiceRegistry, "get_embedding_scanner", fake_none)
|
||||
|
||||
await migration_module.ExampleImagesMigration._migrate_to_v3(str(tmp_path), [str(folder)])
|
||||
|
||||
# Correctly-sized entries are never rewritten.
|
||||
assert saved == []
|
||||
|
||||
|
||||
async def test_migrate_to_v3_skips_unindexed_model(tmp_path: Path, monkeypatch):
|
||||
"""A folder whose model is absent from every scanner cache is skipped, not fatal."""
|
||||
|
||||
model_hash = "c" * 64
|
||||
folder = tmp_path / model_hash
|
||||
folder.mkdir()
|
||||
(folder / "custom_zzz.mp4").write_bytes(build_mp4(1080, 1920))
|
||||
|
||||
class EmptyScanner:
|
||||
def has_hash(self, _hash):
|
||||
return False
|
||||
|
||||
async def get_cached_data(self):
|
||||
from types import SimpleNamespace
|
||||
|
||||
return SimpleNamespace(raw_data=[])
|
||||
|
||||
async def fake_scanner():
|
||||
return EmptyScanner()
|
||||
|
||||
monkeypatch.setattr(migration_module.ServiceRegistry, "get_lora_scanner", fake_scanner)
|
||||
monkeypatch.setattr(migration_module.ServiceRegistry, "get_checkpoint_scanner", fake_scanner)
|
||||
monkeypatch.setattr(migration_module.ServiceRegistry, "get_embedding_scanner", fake_scanner)
|
||||
|
||||
# Must not raise.
|
||||
await migration_module.ExampleImagesMigration._migrate_to_v3(str(tmp_path), [str(folder)])
|
||||
@@ -0,0 +1,178 @@
|
||||
"""Tests for the container-level video dimension probe."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
|
||||
from py.utils.video_metadata import get_video_dimensions
|
||||
|
||||
|
||||
def _box(box_type: bytes, payload: bytes) -> bytes:
|
||||
return struct.pack(">I", len(payload) + 8) + box_type + payload
|
||||
|
||||
|
||||
def _full_box(box_type: bytes, payload: bytes) -> bytes:
|
||||
"""Build a box with a 4-byte version/flags header."""
|
||||
|
||||
return _box(box_type, b"\x00\x00\x00\x00" + payload)
|
||||
|
||||
|
||||
def build_mp4(width: int, height: int, *, with_stsd: bool = False) -> bytes:
|
||||
"""Build a minimal but structurally valid MP4 holding one video track."""
|
||||
|
||||
mvhd = _full_box(b"mvhd", b"\x00" * 96)
|
||||
|
||||
hdlr = _full_box(b"hdlr", b"\x00" * 4 + b"vide" + b"\x00" * 12)
|
||||
|
||||
tkhd_payload = struct.pack(">IIII", 0, 0, 0, 0) + b"\x00" * 52
|
||||
tkhd_payload += struct.pack(">II", width << 16, height << 16)
|
||||
tkhd = _full_box(b"tkhd", tkhd_payload)
|
||||
|
||||
stbl_children = b""
|
||||
if with_stsd:
|
||||
sample_entry = (
|
||||
b"\x00" * 6 + struct.pack(">H", 1) + struct.pack(">HH", width, height)
|
||||
)
|
||||
stsd = _full_box(b"stsd", struct.pack(">I", 1) + _box(b"avc1", sample_entry))
|
||||
stbl_children = stsd
|
||||
|
||||
minf = _box(b"minf", _box(b"stbl", stbl_children))
|
||||
mdia = _box(b"mdia", hdlr + minf)
|
||||
trak = _box(b"trak", tkhd + mdia)
|
||||
moov = _box(b"moov", mvhd + trak)
|
||||
ftyp = _box(b"ftyp", b"isom" + b"\x00\x00\x02\x00" + b"isomiso2avc1mp41")
|
||||
|
||||
return ftyp + moov
|
||||
|
||||
|
||||
def _ebml_vint(value: int) -> bytes:
|
||||
"""Encode a value as a minimal-length EBML variable length integer."""
|
||||
|
||||
for length in range(1, 9):
|
||||
if value < (1 << (7 * length)):
|
||||
encoded = value | (1 << (7 * length))
|
||||
return encoded.to_bytes(length, "big")
|
||||
raise ValueError("value too large for an EBML vint")
|
||||
|
||||
|
||||
def _ebml_element(element_id: bytes, payload: bytes) -> bytes:
|
||||
return element_id + _ebml_vint(len(payload)) + payload
|
||||
|
||||
|
||||
def _uint_element(element_id: int, value: int) -> bytes:
|
||||
length = max(1, (value.bit_length() + 7) // 8)
|
||||
return _ebml_element(
|
||||
element_id.to_bytes(2, "big") if element_id > 0xFF else element_id.to_bytes(1, "big"),
|
||||
value.to_bytes(length, "big"),
|
||||
)
|
||||
|
||||
|
||||
def build_webm(width: int, height: int, *, track_type: int = 1) -> bytes:
|
||||
"""Build a minimal WebM file holding one TrackEntry."""
|
||||
|
||||
video = _ebml_element(b"\xe0", _uint_element(0xB0, width) + _uint_element(0xBA, height))
|
||||
track_entry = _ebml_element(
|
||||
b"\xae", _uint_element(0x83, track_type) + video
|
||||
)
|
||||
tracks = _ebml_element(b"\x16\x54\xae\x6b", track_entry)
|
||||
|
||||
segment = _ebml_element(b"\x18\x53\x80\x67", tracks)
|
||||
ebml_header = _ebml_element(
|
||||
b"\x1a\x45\xdf\xa3",
|
||||
_uint_element(0x4286, 1) + _ebml_element(b"\x42\x82", b"webm"),
|
||||
)
|
||||
|
||||
return ebml_header + segment
|
||||
|
||||
|
||||
def test_mp4_dimensions_come_from_tkhd(tmp_path):
|
||||
video = tmp_path / "landscape.mp4"
|
||||
video.write_bytes(build_mp4(1280, 720))
|
||||
|
||||
assert get_video_dimensions(str(video)) == (1280, 720)
|
||||
|
||||
|
||||
def test_mp4_uses_stsd_when_tkhd_is_empty(tmp_path):
|
||||
video = tmp_path / "stsd-only.mp4"
|
||||
video.write_bytes(build_mp4(640, 480, with_stsd=True))
|
||||
|
||||
assert get_video_dimensions(str(video)) == (640, 480)
|
||||
|
||||
|
||||
def test_mp4_without_video_track_returns_none(tmp_path):
|
||||
# A moov whose only trak has no mdia box at all.
|
||||
tkhd = _full_box(b"tkhd", b"\x00" * 60)
|
||||
moov = _box(b"moov", _box(b"trak", tkhd))
|
||||
video = tmp_path / "audio-only.mp4"
|
||||
video.write_bytes(moov)
|
||||
|
||||
assert get_video_dimensions(str(video)) is None
|
||||
|
||||
|
||||
def test_webm_dimensions(tmp_path):
|
||||
video = tmp_path / "portrait.webm"
|
||||
video.write_bytes(build_webm(720, 1280))
|
||||
|
||||
assert get_video_dimensions(str(video)) == (720, 1280)
|
||||
|
||||
|
||||
def test_webm_non_video_track_is_ignored(tmp_path):
|
||||
video = tmp_path / "audio.webm"
|
||||
video.write_bytes(build_webm(720, 1280, track_type=2))
|
||||
|
||||
assert get_video_dimensions(str(video)) is None
|
||||
|
||||
|
||||
def test_container_signature_wins_over_extension(tmp_path):
|
||||
"""A WebM file named ``.mp4`` is still parsed as WebM."""
|
||||
|
||||
video = tmp_path / "actually-webm.mp4"
|
||||
video.write_bytes(build_webm(480, 832))
|
||||
|
||||
assert get_video_dimensions(str(video)) == (480, 832)
|
||||
|
||||
|
||||
def test_webp_renamed_to_mp4_is_read(tmp_path):
|
||||
"""Animated WebP examples are frequently saved with a video extension."""
|
||||
|
||||
vp8_payload = b"\x30\x36\x02" + b"\x9d\x01\x2a" + struct.pack("<HH", 450, 800)
|
||||
chunk = b"VP8 " + struct.pack("<I", len(vp8_payload)) + vp8_payload
|
||||
body = b"WEBP" + chunk
|
||||
riff = b"RIFF" + struct.pack("<I", len(body)) + body
|
||||
|
||||
video = tmp_path / "animated.mp4"
|
||||
video.write_bytes(riff)
|
||||
|
||||
assert get_video_dimensions(str(video)) == (450, 800)
|
||||
|
||||
|
||||
def test_webp_vp8x_canvas_dimensions(tmp_path):
|
||||
vp8x_payload = b"\x00" * 4 + (449).to_bytes(3, "little") + (799).to_bytes(3, "little")
|
||||
chunk = b"VP8X" + struct.pack("<I", len(vp8x_payload)) + vp8x_payload
|
||||
body = b"WEBP" + chunk
|
||||
riff = b"RIFF" + struct.pack("<I", len(body)) + body
|
||||
|
||||
video = tmp_path / "canvas.mp4"
|
||||
video.write_bytes(riff)
|
||||
|
||||
assert get_video_dimensions(str(video)) == (450, 800)
|
||||
|
||||
|
||||
def test_missing_file_returns_none(tmp_path):
|
||||
assert get_video_dimensions(str(tmp_path / "nope.mp4")) is None
|
||||
|
||||
|
||||
def test_corrupt_file_returns_none(tmp_path):
|
||||
video = tmp_path / "corrupt.mp4"
|
||||
video.write_bytes(b"\x00\x00\x00\x20ftypisom" + b"\xff" * 64)
|
||||
|
||||
assert get_video_dimensions(str(video)) is None
|
||||
|
||||
|
||||
def test_unsupported_extension_without_video_signature_returns_none(tmp_path):
|
||||
"""A non-video file is not probed just because of a video-like name."""
|
||||
|
||||
video = tmp_path / "clip.avi"
|
||||
video.write_bytes(b"RIFF\x00\x00\x00\x00AVI LIST\x00\x00\x00\x00")
|
||||
|
||||
assert get_video_dimensions(str(video)) is None
|
||||
Reference in New Issue
Block a user