diff --git a/py/utils/example_images_metadata.py b/py/utils/example_images_metadata.py index 2fed61d1..8d1369c8 100644 --- a/py/utils/example_images_metadata.py +++ b/py/utils/example_images_metadata.py @@ -2,7 +2,7 @@ import inspect import logging import os import re -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict, Mapping, MutableMapping, Optional from ..recipes.constants import GEN_PARAM_KEYS from ..services.metadata_service import get_default_metadata_provider, get_metadata_provider @@ -13,9 +13,20 @@ from ..services.downloader import get_downloader from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS from ..utils.exif_utils import ExifUtils from ..utils.metadata_manager import MetadataManager +from ..utils.video_metadata import get_video_dimensions logger = logging.getLogger(__name__) +# Placeholder dimensions written when the real ones cannot be determined. +# Kept for backwards compatibility with pre-existing metadata entries. +_DEFAULT_MEDIA_WIDTH = 720 +_DEFAULT_MEDIA_HEIGHT = 1280 + +# Example metadata entries carry a marker: ``customImages`` use their ``id`` +# while ``images`` use the positional index. Either way the marker must be a +# plain filename-safe token, never a path fragment. +_ENTRY_MARKER_PATTERN = re.compile(r"^(?:custom_|image_)?([^./\\]+)$") + _preview_service = PreviewAssetService( metadata_manager=MetadataManager, downloader_factory=get_downloader, @@ -66,6 +77,141 @@ def _build_metadata_sync_service(settings_manager: "SettingsManager") -> Metadat ) +def _read_media_dimensions(path: str, is_video: bool) -> tuple[int, int]: + """Return ``(width, height)`` for an example image or video file. + + Videos are read from their container headers (PIL cannot open them) so the + showcase viewer sizes the gallery to the real aspect ratio. Falls back to + the legacy ``720x1280`` placeholder when the dimensions cannot be + determined — e.g. an unreadable file or an exotic codec — which only + affects the displayed aspect ratio, never the file itself. + """ + + dimensions = None + + if is_video: + dimensions = get_video_dimensions(path) + else: + try: + from PIL import Image + + if os.path.exists(path): + with Image.open(path) as img: + dimensions = img.size + except Exception: + dimensions = None + + if dimensions: + width, height = dimensions + if width > 0 and height > 0: + return int(width), int(height) + + return _DEFAULT_MEDIA_WIDTH, _DEFAULT_MEDIA_HEIGHT + + +def _is_video_entry(file_path: Optional[str], entry: Mapping[str, Any]) -> bool: + """Return True when an example entry points at a video file. + + The local file extension wins over the recorded ``type`` because files in + the wild are frequently mislabelled (animated WebP saved as ``.mp4``); + ``_read_media_dimensions`` handles that correctly either way. + """ + + if file_path: + ext = os.path.splitext(file_path)[1].lower() + if ext in SUPPORTED_MEDIA_EXTENSIONS["videos"]: + return True + if ext in SUPPORTED_MEDIA_EXTENSIONS["images"]: + return False + return str(entry.get("type", "")).lower() == "video" + + +def _resolve_local_file( + entry: Mapping[str, Any], + index: int, + local_files: Mapping[str, str], +) -> Optional[str]: + """Map a metadata entry onto its example file inside the model folder. + + Reads the entry's own marker (``id`` for ``customImages``, positional + ``index`` for ``images``) with an anchored regex, so the identifier can + never bleed into a neighbouring filename the way a prefix comparison can. + """ + + marker = entry.get("id") + if not isinstance(marker, str) or not marker: + marker = str(index) + + match = _ENTRY_MARKER_PATTERN.fullmatch(marker) + if not match: + return None + + return local_files.get(match.group(1)) + + +def repair_local_video_dimensions( + metadata: MutableMapping[str, Any], + local_files: Mapping[str, str], + *, + dry_run: bool = False, +) -> int: + """Backfill real video dimensions for an entry that has local files. + + Only entries with an empty ``url`` are considered: those have no remote + source, so the local file is the single source of truth for their size and + rewriting them cannot discard API-supplied data. Entries whose dimensions + already match the file are left byte-identical. + + Args: + metadata: Raw metadata payload (mutated in place unless ``dry_run``). + local_files: ``{identifier: path}`` for files present in the model's + example folder, where the identifier is the entry's ``id`` (for + ``customImages``) or its positional index (for ``images``). + dry_run: Count the fixes without mutating ``metadata``. + + Returns: + The number of entries that were (or would be) repaired. + """ + + civitai = metadata.get("civitai") + if not isinstance(civitai, dict): + return 0 + + repaired = 0 + + for key in ("customImages", "images"): + entries = civitai.get(key) + if not isinstance(entries, list) or not entries: + continue + + for index, entry in enumerate(entries): + if not isinstance(entry, dict): + continue + if entry.get("url", "") != "": + # Remote-backed entry: never rebuilt from local state. + continue + + file_path = _resolve_local_file(entry, index, local_files) + if not file_path or not os.path.isfile(file_path): + continue + + dimensions = _read_media_dimensions( + file_path, _is_video_entry(file_path, entry) + ) + width, height = dimensions + if width <= 0 or height <= 0: + continue + if entry.get("width") == width and entry.get("height") == height: + continue + + if not dry_run: + entry["width"] = width + entry["height"] = height + repaired += 1 + + return repaired + + def _get_metadata_sync_service() -> MetadataSyncService: """Return the shared metadata sync service, initialising it lazily.""" @@ -230,29 +376,21 @@ class MetadataUpdater: # Determine if video or image file_ext = os.path.splitext(path)[1].lower() is_video = file_ext in SUPPORTED_MEDIA_EXTENSIONS['videos'] - + + width, height = _read_media_dimensions(path, is_video) + # Create image metadata entry image_entry = { "url": "", # Empty URL as required "nsfwLevel": 0, - "width": 720, # Default dimensions - "height": 1280, + "width": width, + "height": height, "type": "video" if is_video else "image", "meta": None, "hasMeta": False, "hasPositivePrompt": False } - - # If it's an image, try to get actual dimensions (optional enhancement) - try: - from PIL import Image - if not is_video and os.path.exists(path): - with Image.open(path) as img: - image_entry["width"], image_entry["height"] = img.size - except: - # If PIL fails or is unavailable, use default dimensions - pass - + images.append(image_entry) # Update the model's civitai.images field @@ -321,14 +459,16 @@ class MetadataUpdater: # Determine if video or image file_ext = os.path.splitext(path)[1].lower() is_video = file_ext in SUPPORTED_MEDIA_EXTENSIONS['videos'] - + + width, height = _read_media_dimensions(path, is_video) + # Create image metadata entry image_entry = { "url": "", # Empty URL as requested "id": short_id, "nsfwLevel": 0, - "width": 720, # Default dimensions - "height": 1280, + "width": width, + "height": height, "type": "video" if is_video else "image", "meta": None, "hasMeta": False, @@ -353,16 +493,6 @@ class MetadataUpdater: except Exception as e: logger.warning(f"Failed to extract metadata from {os.path.basename(path)}: {e}") - # If it's an image, try to get actual dimensions - try: - from PIL import Image - if not is_video and os.path.exists(path): - with Image.open(path) as img: - image_entry["width"], image_entry["height"] = img.size - except: - # If PIL fails or is unavailable, use default dimensions - pass - # Append to existing customImages array custom_images.append(image_entry) diff --git a/py/utils/example_images_migration.py b/py/utils/example_images_migration.py index abf606a4..ac203a5d 100644 --- a/py/utils/example_images_migration.py +++ b/py/utils/example_images_migration.py @@ -15,12 +15,20 @@ from ..utils.example_images_paths import ( ) from ..utils.metadata_manager import MetadataManager from ..utils.example_images_processor import ExampleImagesProcessor -from ..utils.example_images_metadata import update_cache_from_metadata +from ..utils.example_images_metadata import ( + repair_local_video_dimensions, + update_cache_from_metadata, +) from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS logger = logging.getLogger(__name__) -CURRENT_NAMING_VERSION = 2 # Increment this when naming conventions change +CURRENT_NAMING_VERSION = 3 # Increment this when naming conventions change + +# Example files worth inspecting during the dimension repair. +_REPAIRABLE_EXTENSIONS = frozenset( + SUPPORTED_MEDIA_EXTENSIONS["images"] + SUPPORTED_MEDIA_EXTENSIONS["videos"] +) class _SettingsProxy: @@ -185,6 +193,9 @@ class ExampleImagesMigration: if from_version < 2 and to_version >= 2: await ExampleImagesMigration._migrate_to_v2(model_folders) + if from_version < 3 and to_version >= 3: + await ExampleImagesMigration._migrate_to_v3(example_images_path, model_folders) + # Update version in progress file progress_file = os.path.join(example_images_path, '.download_progress.json') try: @@ -437,4 +448,137 @@ class ExampleImagesMigration: logger.error(f"Error migrating folder {folder}: {e}") migration_errors += 1 - logger.info(f"Migration to v2 complete: migrated {count} custom examples across {updated_models} models with {migration_errors} errors") \ No newline at end of file + logger.info(f"Migration to v2 complete: migrated {count} custom examples across {updated_models} models with {migration_errors} errors") + + @staticmethod + def _build_local_file_map(folder): + """Map entry markers to their files inside a model's example folder. + + Keys are the marker alone (``custom_`` → ````, + ``image_`` → ````) so they line up with the metadata + entries' ``id``/positional index without any prefix ambiguity. + """ + + local_files = {} + try: + entries = os.listdir(folder) + except OSError as exc: + logger.debug("Could not list example folder %s: %s", folder, exc) + return local_files + + for name in entries: + stem, ext = os.path.splitext(name) + if ext.lower() not in _REPAIRABLE_EXTENSIONS: + continue + if stem.startswith("custom_"): + local_files[stem[len("custom_"):]] = os.path.join(folder, name) + elif stem.startswith("image_"): + local_files[stem[len("image_"):]] = os.path.join(folder, name) + + return local_files + + @staticmethod + async def _find_scanner_for_hash(model_hash): + """Return the scanner owning ``model_hash``, or ``None``.""" + + lora_scanner = await ServiceRegistry.get_lora_scanner() + checkpoint_scanner = await ServiceRegistry.get_checkpoint_scanner() + embedding_scanner = await ServiceRegistry.get_embedding_scanner() + + for scanner in (lora_scanner, checkpoint_scanner, embedding_scanner): + if scanner is None: + continue + try: + if scanner.has_hash(model_hash): + return scanner + except Exception as exc: # pragma: no cover - defensive + logger.debug("has_hash check failed for %s: %s", type(scanner).__name__, exc) + return None + + @staticmethod + async def _migrate_to_v3(example_images_path, model_folders): + """Backfill real dimensions for locally imported example videos. + + Imported videos were stored with a hardcoded ``720x1280`` placeholder + (issue #1115), so landscape clips were rendered inside a portrait + container. Only entries with an empty ``url`` are touched — those have + no remote source, which makes the local file authoritative and the + rewrite lossless. Entries already carrying the right size are left + untouched, so re-running this migration is a no-op. + + This runs once per library via the ``naming_version`` gate in + ``run_migrations``; it is deliberately not wired into any request path. + """ + + repaired_entries = 0 + updated_models = 0 + migration_errors = 0 + + logger.info( + "Starting v3 migration (local example video dimensions) for %d model folders", + len(model_folders), + ) + + for folder in model_folders: + try: + model_hash = os.path.basename(folder) + if not model_hash or len(model_hash) != 64: + continue + + local_files = ExampleImagesMigration._build_local_file_map(folder) + if not local_files: + continue + + scanner = await ExampleImagesMigration._find_scanner_for_hash(model_hash) + if scanner is None: + logger.debug( + "Model %s not found in any scanner cache, skipping dimension repair", + model_hash, + ) + continue + + cache = await scanner.get_cached_data() + model_data = None + for item in cache.raw_data: + if item.get("sha256") == model_hash: + model_data = item + break + + if not model_data: + continue + + file_path = model_data.get("file_path") + if not file_path: + continue + + payload = await MetadataManager.load_metadata_payload(file_path) + if not isinstance(payload, dict): + continue + + repaired = repair_local_video_dimensions(payload, local_files) + if repaired <= 0: + continue + + # The model cache shape differs from the on-disk payload, so + # persist the file first and let the cache sync re-read it. + await MetadataManager.save_metadata(file_path, payload) + await update_cache_from_metadata(scanner, file_path, payload) + + repaired_entries += repaired + updated_models += 1 + + except Exception as exc: + logger.error( + "Failed to repair example video dimensions for %s: %s", + folder, + exc, + ) + migration_errors += 1 + + logger.info( + "Migration to v3 complete: repaired %d example entr(ies) across %d model(s) " + "with %d error(s)", + repaired_entries, + updated_models, + migration_errors, + ) \ No newline at end of file diff --git a/py/utils/video_metadata.py b/py/utils/video_metadata.py new file mode 100644 index 00000000..7cea72fc --- /dev/null +++ b/py/utils/video_metadata.py @@ -0,0 +1,623 @@ +"""Read intrinsic dimensions from video containers without external tooling. + +PIL cannot open ``.mp4``/``.webm`` files, so example videos imported through +the "Add examples" flow used to fall back to a hardcoded ``720x1280`` (portrait) +entry, which forced the showcase viewer to letterbox landscape videos. + +This module reads the dimensions out of the container headers themselves: + +* ISO base media files (``.mp4``/``.mov``/``.m4v``) — ``moov/trak/tkhd``, + falling back to the sample description of the video track. +* WebM/Matroska (``.webm``/``.mkv``) — ``Segment/Tracks/TrackEntry/Video`` + ``PixelWidth``/``PixelHeight``. +* Animated WebP (``RIFF``/``WEBP``) — handled because users routinely save + animated examples with a video extension. + +The container signature decides which reader runs, so a mislabelled file +(a ``.mp4`` that is really WebM) still reports the right dimensions. + +Both readers stream over the file: only container headers are read, so a +multi-gigabyte ``mdat`` is never pulled into memory (it is seeked past). +""" + +from __future__ import annotations + +import functools +import logging +import os +import struct +from typing import BinaryIO, Iterator, Optional, Tuple + +logger = logging.getLogger(__name__) + +ISO_MEDIA_EXTENSIONS = frozenset({".mp4", ".m4v", ".mov"}) +EBML_MEDIA_EXTENSIONS = frozenset({".webm", ".mkv"}) + +_EBML_MAGIC = b"\x1a\x45\xdf\xa3" + +# Cap recursion into nesting containers so a crafted/corrupt file cannot blow +# the Python stack. +_MAX_BOX_DEPTH = 12 +_MAX_EBML_DEPTH = 12 + +# Header structs (``tkhd``, sample entries) are tiny; guard against a bogus +# size claiming the whole file. +_MAX_HEADER_PAYLOAD = 1024 * 1024 + +_WIDTH_HEIGHT_UNSET = (0, 0) + + +@functools.lru_cache(maxsize=4096) +def _get_video_dimensions_cached( + path: str, _mtime_ns: int, _size: int +) -> Optional[Tuple[int, int]]: + """Return ``(width, height)`` for ``path``, or ``None`` on any failure. + + ``_mtime_ns`` and ``_size`` participate in the cache key only so a replaced + file is re-probed; they are never read by the parser. + """ + try: + return _read_video_dimensions(path) + except Exception: + logger.debug("Failed to read video dimensions for %s", path, exc_info=True) + return None + + +def _read_video_dimensions(path: str) -> Optional[Tuple[int, int]]: + """Dispatch to the ISO or EBML reader based on the container's magic bytes. + + Real libraries contain files whose extension lies about their container + (a ``.mp4`` that is really WebM, typically), so the sniffed signature wins + and the extension is only a fallback. + """ + + ext = os.path.splitext(path)[1].lower() + file_size = os.path.getsize(path) + + with open(path, "rb") as stream: + magic = stream.read(12) + + if _looks_like_iso_media(magic): + return _read_iso_media_dimensions(stream, file_size) + if magic[:4] == _EBML_MAGIC: + return _read_ebml_dimensions(stream, file_size) + if magic[:4] == b"RIFF" and magic[8:12] == b"WEBP": + return _read_riff_webp_dimensions(stream, file_size) + + # Signature is inconclusive (truncated or unusual file): fall back to + # the extension. + if ext in EBML_MEDIA_EXTENSIONS: + return _read_ebml_dimensions(stream, file_size) + if ext in ISO_MEDIA_EXTENSIONS: + return _read_iso_media_dimensions(stream, file_size) + return None + + +def _looks_like_iso_media(magic: bytes) -> bool: + """Return True when the leading bytes are an ISO base media box header.""" + + return len(magic) >= 8 and magic[4:8] in { + b"ftyp", + b"moov", + b"mdat", + b"free", + b"skip", + b"wide", + } + + +def get_video_dimensions(path: str) -> Optional[Tuple[int, int]]: + """Return the intrinsic ``(width, height)`` of a local video file. + + Returns ``None`` when the extension is unsupported, the file is missing or + corrupt, or the dimensions cannot be determined. Never raises. + """ + if not path: + return None + try: + stat = os.stat(path) + except OSError: + return None + return _get_video_dimensions_cached(path, stat.st_mtime_ns, stat.st_size) + + +def _clear_video_dimensions_cache() -> None: + """Drop the dimension cache (used by tests).""" + + _get_video_dimensions_cached.cache_clear() + + +# --------------------------------------------------------------------------- # +# ISO base media (MP4 / MOV) +# --------------------------------------------------------------------------- # + + +def _iter_boxes( + stream: BinaryIO, end: int, depth: int = 0 +) -> Iterator[Tuple[bytes, int, int]]: + """Yield ``(type, payload_start, box_end)`` for boxes in ``[tell, end)``. + + The stream is left at the next box boundary after each yielded box. + """ + if depth > _MAX_BOX_DEPTH: + return + + while True: + start = stream.tell() + if start + 8 > end: + return + + header = stream.read(8) + if len(header) < 8: + return + + size, box_type = struct.unpack(">I4s", header) + header_size = 8 + + if size == 1: + # 64-bit ``largesize`` follows the type. + extended = stream.read(8) + if len(extended) < 8: + return + size = struct.unpack(">Q", extended)[0] + header_size = 16 + elif size == 0: + # Box extends to the end of the enclosing container. + size = end - start + + if size < header_size or start + size > end: + return + + yield box_type, start + header_size, start + size + stream.seek(start + size) + + +def _read_iso_media_dimensions( + stream: BinaryIO, file_size: int +) -> Optional[Tuple[int, int]]: + """Walk ``moov`` looking for the video track's dimensions.""" + + stream.seek(0) + moov: Optional[Tuple[int, int]] = None + for box_type, payload_start, box_end in _iter_boxes(stream, file_size): + if box_type == b"moov": + moov = (payload_start, box_end) + break + + if moov is None: + return None + + stream.seek(moov[0]) + for box_type, payload_start, box_end in _iter_boxes(stream, moov[1], depth=1): + if box_type != b"trak": + continue + dimensions = _read_trak_dimensions(stream, payload_start, box_end) + if dimensions is not None: + return dimensions + + return None + + +def _read_trak_dimensions( + stream: BinaryIO, trak_start: int, trak_end: int +) -> Optional[Tuple[int, int]]: + """Return the dimensions of a ``trak`` when it describes a video track.""" + + stream.seek(trak_start) + + is_video = False + tkhd_dimensions = _WIDTH_HEIGHT_UNSET + stsd_dimensions = _WIDTH_HEIGHT_UNSET + + for box_type, payload_start, box_end in _iter_boxes(stream, trak_end, depth=2): + if box_type == b"tkhd": + tkhd_dimensions = _parse_tkhd(stream, payload_start, box_end) + elif box_type == b"mdia": + stream.seek(payload_start) + media = _read_mdia_dimensions(stream, payload_start, box_end) + if media is not None: + is_video, stsd_dimensions = media + + if not is_video: + return None + + # ``tkhd`` is preferred: it is display space, and its 16.16 fixed point + # encoding keeps non-integer dimensions (odd crops produce those). + for width, height in (tkhd_dimensions, stsd_dimensions): + if width > 0 and height > 0: + return int(round(width)), int(round(height)) + return None + + +def _read_mdia_dimensions( + stream: BinaryIO, mdia_start: int, mdia_end: int +) -> Optional[Tuple[bool, Tuple[float, float]]]: + """Return ``(is_video, dimensions)`` for a ``mdia`` box.""" + + handler_type = b"" + stsd_dimensions = _WIDTH_HEIGHT_UNSET + + for box_type, payload_start, box_end in _iter_boxes(stream, mdia_end, depth=3): + if box_type == b"hdlr": + handler_type = _parse_handler_type(stream, payload_start, box_end) + elif box_type == b"minf": + stream.seek(payload_start) + stsd_dimensions = _read_minf_dimensions(stream, payload_start, box_end) + + return handler_type == b"vide", stsd_dimensions + + +def _read_minf_dimensions( + stream: BinaryIO, minf_start: int, minf_end: int +) -> Tuple[float, float]: + """Return the sample-entry dimensions declared under ``minf/stbl/stsd``.""" + + for box_type, payload_start, box_end in _iter_boxes(stream, minf_end, depth=4): + if box_type != b"stbl": + continue + stream.seek(payload_start) + for inner_type, inner_start, inner_end in _iter_boxes( + stream, box_end, depth=5 + ): + if inner_type == b"stsd": + return _parse_stsd(stream, inner_start, inner_end) + return _WIDTH_HEIGHT_UNSET + + +def _parse_tkhd( + stream: BinaryIO, payload_start: int, box_end: int +) -> Tuple[float, float]: + """Parse the 16.16 fixed point width/height trailer of a ``tkhd`` box.""" + + size = box_end - payload_start + if size < 8 or size > _MAX_HEADER_PAYLOAD: + return _WIDTH_HEIGHT_UNSET + + stream.seek(box_end - 8) + trailer = stream.read(8) + if len(trailer) < 8: + return _WIDTH_HEIGHT_UNSET + + width, height = struct.unpack(">II", trailer) + return width / 65536.0, height / 65536.0 + + +def _parse_handler_type( + stream: BinaryIO, payload_start: int, box_end: int +) -> bytes: + """Parse the handler type from an ``hdlr`` box. + + Layout: version/flags (4) + pre_defined (4) + handler_type (4). + """ + + if box_end - payload_start < 12: + return b"" + stream.seek(payload_start) + data = stream.read(12) + if len(data) < 12: + return b"" + return data[8:12] + + +def _parse_stsd( + stream: BinaryIO, payload_start: int, box_end: int +) -> Tuple[float, float]: + """Parse the visual sample entry dimensions from an ``stsd`` box. + + Only the first entry is inspected: video tracks are single-entry in every + container we import from. + """ + + if box_end - payload_start < 16: + return _WIDTH_HEIGHT_UNSET + + stream.seek(payload_start) + header = stream.read(8) # version/flags + entry_count + if len(header) < 8: + return _WIDTH_HEIGHT_UNSET + + entry_start = payload_start + 8 + if entry_start + 8 > box_end: + return _WIDTH_HEIGHT_UNSET + + stream.seek(entry_start) + entry_header = stream.read(8) + if len(entry_header) < 8: + return _WIDTH_HEIGHT_UNSET + + entry_size = struct.unpack(">I", entry_header[:4])[0] + header_size = 8 + + if entry_size == 1: + extended = stream.read(8) + if len(extended) < 8: + return _WIDTH_HEIGHT_UNSET + entry_size = struct.unpack(">Q", extended)[0] + header_size = 16 + elif entry_size == 0: + entry_size = box_end - entry_start + + if entry_size < header_size + 8 or entry_start + entry_size > box_end: + return _WIDTH_HEIGHT_UNSET + + # Visual sample entries: 6 bytes reserved + 2 bytes data_reference_index, + # then width (2) and height (2). + stream.seek(entry_start + header_size + 6 + 2) + dimensions = stream.read(4) + if len(dimensions) < 4: + return _WIDTH_HEIGHT_UNSET + + width, height = struct.unpack(">HH", dimensions) + return float(width), float(height) + + +# --------------------------------------------------------------------------- # +# WebM / Matroska (EBML) +# --------------------------------------------------------------------------- # + +# EBML element IDs (stored with their length marker, as they appear on disk). +_ID_SEGMENT = 0x18538067 +_ID_TRACKS = 0x1654AE6B +_ID_TRACK_ENTRY = 0xAE +_ID_TRACK_TYPE = 0x83 +_ID_VIDEO = 0xE0 +_ID_PIXEL_WIDTH = 0xB0 +_ID_PIXEL_HEIGHT = 0xBA + +# Nested containers we descend into while hunting for video dimensions. +_EBML_CONTAINER_IDS = frozenset({_ID_SEGMENT, _ID_TRACKS, _ID_TRACK_ENTRY}) + + +def _read_ebml_vint(stream: BinaryIO, *, keep_marker: bool) -> Optional[Tuple[int, int]]: + """Read an EBML variable-length integer. + + Returns ``(value, byte_length)``. For element IDs the marker bit is kept + (``keep_marker=True``) because IDs are compared in their on-disk form; for + sizes the marker is stripped to yield the actual payload length. + """ + + first = stream.read(1) + if not first: + return None + + first_byte = first[0] + if first_byte == 0: + return None + + length = 1 + mask = 0x80 + while not first_byte & mask: + mask >>= 1 + length += 1 + if length > 8: + return None + + value = first_byte if keep_marker else first_byte & (mask - 1) + remaining = length - 1 + + if remaining: + extra = stream.read(remaining) + if len(extra) < remaining: + return None + for byte in extra: + value = (value << 8) | byte + + return value, length + + +def _read_ebml_dimensions( + stream: BinaryIO, file_size: int +) -> Optional[Tuple[int, int]]: + """Parse ``Segment/Tracks`` for the first video ``TrackEntry``.""" + + stream.seek(0) + header = stream.read(4) + if header != _EBML_MAGIC: + return None + + return _walk_ebml(stream, 0, file_size, depth=0) + + +def _walk_ebml( + stream: BinaryIO, start: int, end: int, *, depth: int +) -> Optional[Tuple[int, int]]: + """Recursively scan EBML elements in ``[start, end)`` for video dimensions.""" + + if depth > _MAX_EBML_DEPTH: + return None + + stream.seek(start) + + while stream.tell() < end: + element_start = stream.tell() + + element_id = _read_ebml_vint(stream, keep_marker=True) + if element_id is None: + return None + element_id_value = element_id[0] + + size_field = _read_ebml_vint(stream, keep_marker=False) + if size_field is None: + return None + payload_size, size_length = size_field + + payload_start = element_start + element_id[1] + size_length + + # A size field of all-ones marks an unknown-size element, which is + # legal for Segment/Tracks; treat it as "until the parent ends". + unknown_size = payload_size == (1 << (7 * size_length)) - 1 + payload_end = end if unknown_size else payload_start + payload_size + + if payload_end > end: + return None + + if element_id_value == _ID_VIDEO: + dimensions = _read_ebml_video(stream, payload_start, min(payload_end, end)) + if dimensions is not None: + return dimensions + elif element_id_value == _ID_TRACK_ENTRY: + track = _read_ebml_track_entry( + stream, payload_start, min(payload_end, end) + ) + if track is not None: + return track + elif element_id_value in _EBML_CONTAINER_IDS: + found = _walk_ebml( + stream, payload_start, min(payload_end, end), depth=depth + 1 + ) + if found is not None: + return found + + if unknown_size: + # Cannot resume after an unknown-size element; its siblings cannot + # be located reliably, so stop scanning this level. + return None + + stream.seek(payload_end) + + return None + + +def _read_ebml_track_entry( + stream: BinaryIO, start: int, end: int +) -> Optional[Tuple[int, int]]: + """Return dimensions when a ``TrackEntry`` is a video track.""" + + track_type: Optional[int] = None + dimensions: Optional[Tuple[int, int]] = None + + stream.seek(start) + while stream.tell() < end: + element_start = stream.tell() + + element_id = _read_ebml_vint(stream, keep_marker=True) + if element_id is None: + return None + + size_field = _read_ebml_vint(stream, keep_marker=False) + if size_field is None: + return None + payload_size, size_length = size_field + + payload_start = element_start + element_id[1] + size_length + payload_end = min(payload_start + payload_size, end) + + if element_id[0] == _ID_TRACK_TYPE: + track_type = _read_ebml_uint(stream, payload_start, payload_end) + elif element_id[0] == _ID_VIDEO: + dimensions = _read_ebml_video(stream, payload_start, payload_end) + + stream.seek(payload_end) + + # Track type 1 is video. + if track_type == 1 and dimensions is not None: + return dimensions + return None + + +def _read_ebml_video( + stream: BinaryIO, start: int, end: int +) -> Optional[Tuple[int, int]]: + """Return ``PixelWidth``/``PixelHeight`` from a ``Video`` element.""" + + width: Optional[int] = None + height: Optional[int] = None + + stream.seek(start) + while stream.tell() < end: + element_start = stream.tell() + + element_id = _read_ebml_vint(stream, keep_marker=True) + if element_id is None: + return None + + size_field = _read_ebml_vint(stream, keep_marker=False) + if size_field is None: + return None + payload_size, size_length = size_field + + payload_start = element_start + element_id[1] + size_length + payload_end = min(payload_start + payload_size, end) + + if element_id[0] == _ID_PIXEL_WIDTH: + width = _read_ebml_uint(stream, payload_start, payload_end) + elif element_id[0] == _ID_PIXEL_HEIGHT: + height = _read_ebml_uint(stream, payload_start, payload_end) + + stream.seek(payload_end) + + if width and height and width > 0 and height > 0: + return width, height + return None + + +def _read_ebml_uint(stream: BinaryIO, start: int, end: int) -> Optional[int]: + """Read an unsigned big-endian integer element payload.""" + + length = end - start + if length <= 0 or length > 8: + return None + + stream.seek(start) + raw = stream.read(length) + if len(raw) < length: + return None + + value = 0 + for byte in raw: + value = (value << 8) | byte + return value + + +# --------------------------------------------------------------------------- # +# RIFF / WebP (animated examples are often renamed to ``.mp4``) +# --------------------------------------------------------------------------- # + + +def _read_riff_webp_dimensions( + stream: BinaryIO, file_size: int +) -> Optional[Tuple[int, int]]: + """Return dimensions from a WebP file's first dimension-bearing chunk.""" + + stream.seek(12) + + while stream.tell() + 8 <= file_size: + header = stream.read(8) + if len(header) < 8: + return None + + fourcc, chunk_size = struct.unpack("<4sI", header) + payload_start = stream.tell() + + if fourcc == b"VP8X": + payload = stream.read(10) + if len(payload) < 10: + return None + # Canvas size is stored minus one, as 24-bit little endian values. + width = int.from_bytes(payload[4:7], "little") + 1 + height = int.from_bytes(payload[7:10], "little") + 1 + return width, height + + if fourcc == b"VP8 ": + # Frame tag (3 bytes, bit 0 = key frame) then the key frame start + # code 0x9d 0x01 0x2a and the 16-bit dimensions. + payload = stream.read(10) + if len(payload) < 10: + return None + start = payload.find(b"\x9d\x01\x2a") + if start < 0 or start + 7 > len(payload): + return None + width, height = struct.unpack("> 14) & 0x3FFF) + 1 + + # Skip this chunk (payloads are padded to an even byte boundary). + stream.seek(payload_start + chunk_size + (chunk_size & 1)) + + return None diff --git a/tests/conftest.py b/tests/conftest.py index 6193dbf5..d4c44e5f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -333,6 +333,21 @@ def mock_websocket_manager(): return RecordingWebSocketManager() +@pytest.fixture(autouse=True) +def reset_media_dimension_caches(): + """Clear path-keyed dimension caches so files reused across tests re-probe.""" + from py.utils.exif_utils import _get_image_dimensions_cached + from py.utils.video_metadata import _clear_video_dimensions_cache + + _get_image_dimensions_cached.cache_clear() + _clear_video_dimensions_cache() + + yield + + _get_image_dimensions_cached.cache_clear() + _clear_video_dimensions_cache() + + @pytest.fixture(autouse=True) def reset_singletons(): """Reset all singletons before each test to ensure isolation.""" diff --git a/tests/utils/test_example_images_metadata.py b/tests/utils/test_example_images_metadata.py index 095212f8..c3c46c97 100644 --- a/tests/utils/test_example_images_metadata.py +++ b/tests/utils/test_example_images_metadata.py @@ -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"] \ No newline at end of file + 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) \ No newline at end of file diff --git a/tests/utils/test_example_images_migration.py b/tests/utils/test_example_images_migration.py index f9ac6957..bcd56616 100644 --- a/tests/utils/test_example_images_migration.py +++ b/tests/utils/test_example_images_migration.py @@ -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 == [] diff --git a/tests/utils/test_example_images_video_repair.py b/tests/utils/test_example_images_video_repair.py new file mode 100644 index 00000000..43a7d2ab --- /dev/null +++ b/tests/utils/test_example_images_video_repair.py @@ -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)]) diff --git a/tests/utils/test_video_dimension_probe.py b/tests/utils/test_video_dimension_probe.py new file mode 100644 index 00000000..3c6112ac --- /dev/null +++ b/tests/utils/test_video_dimension_probe.py @@ -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("