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:
Will Miao
@@ -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."""
|
||||
|
||||
@@ -231,28 +377,20 @@ class MetadataUpdater:
|
||||
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
|
||||
@@ -322,13 +460,15 @@ class MetadataUpdater:
|
||||
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)
|
||||
|
||||
|
||||
@@ -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:
|
||||
@@ -438,3 +449,136 @@ class ExampleImagesMigration:
|
||||
migration_errors += 1
|
||||
|
||||
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_<id>`` → ``<id>``,
|
||||
``image_<index>`` → ``<index>``) 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,
|
||||
)
|
||||
@@ -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("<HH", payload[start + 3 : start + 7])
|
||||
return width & 0x3FFF, height & 0x3FFF
|
||||
|
||||
if fourcc == b"VP8L":
|
||||
payload = stream.read(5)
|
||||
if len(payload) < 5 or payload[0] != 0x2F:
|
||||
return None
|
||||
bits = int.from_bytes(payload[1:5], "little")
|
||||
return (bits & 0x3FFF) + 1, ((bits >> 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
|
||||
@@ -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."""
|
||||
|
||||
@@ -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:
|
||||
@@ -217,3 +218,127 @@ async def test_update_metadata_from_local_examples_generates_entries(monkeypatch
|
||||
)
|
||||
assert success is True
|
||||
assert model_data["civitai"]["images"]
|
||||
|
||||
|
||||
async def test_update_metadata_after_import_uses_real_video_dimensions(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path, patch_metadata_manager
|
||||
):
|
||||
"""Regression: imported videos must not fall back to the 720x1280 default.
|
||||
|
||||
See issue #1115 — landscape videos were stored as portrait, so the showcase
|
||||
viewer letterboxed them into a 9:16 container.
|
||||
"""
|
||||
|
||||
model_hash = "d" * 64
|
||||
model_file = tmp_path / "video-model.safetensors"
|
||||
model_file.write_text("content", encoding="utf-8")
|
||||
model_data = {
|
||||
"model_name": "VideoExample",
|
||||
"file_path": str(model_file),
|
||||
"civitai": {},
|
||||
}
|
||||
scanner = StubScanner([model_data])
|
||||
|
||||
video_path = tmp_path / "custom_abc.mp4"
|
||||
video_path.write_bytes(build_mp4(1280, 720))
|
||||
|
||||
monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None))
|
||||
|
||||
_regular, custom = await metadata_module.MetadataUpdater.update_metadata_after_import(
|
||||
model_hash,
|
||||
model_data,
|
||||
scanner,
|
||||
[(str(video_path), "abc")],
|
||||
)
|
||||
|
||||
assert custom[0]["type"] == "video"
|
||||
assert (custom[0]["width"], custom[0]["height"]) == (1280, 720)
|
||||
assert patch_metadata_manager[-1][1]["civitai"]["customImages"][0]["width"] == 1280
|
||||
|
||||
|
||||
async def test_update_metadata_after_import_uses_real_webm_dimensions(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path, patch_metadata_manager
|
||||
):
|
||||
model_hash = "e" * 64
|
||||
model_file = tmp_path / "webm-model.safetensors"
|
||||
model_file.write_text("content", encoding="utf-8")
|
||||
model_data = {
|
||||
"model_name": "WebmExample",
|
||||
"file_path": str(model_file),
|
||||
"civitai": {},
|
||||
}
|
||||
|
||||
video_path = tmp_path / "custom_def.webm"
|
||||
video_path.write_bytes(build_webm(480, 832))
|
||||
|
||||
monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None))
|
||||
|
||||
_regular, custom = await metadata_module.MetadataUpdater.update_metadata_after_import(
|
||||
model_hash,
|
||||
model_data,
|
||||
StubScanner([model_data]),
|
||||
[(str(video_path), "def")],
|
||||
)
|
||||
|
||||
assert (custom[0]["width"], custom[0]["height"]) == (480, 832)
|
||||
|
||||
|
||||
async def test_update_metadata_after_import_falls_back_for_unreadable_video(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path, patch_metadata_manager
|
||||
):
|
||||
"""An unparsable video keeps the legacy placeholder rather than failing."""
|
||||
|
||||
model_hash = "f" * 64
|
||||
model_file = tmp_path / "broken-model.safetensors"
|
||||
model_file.write_text("content", encoding="utf-8")
|
||||
model_data = {
|
||||
"model_name": "BrokenExample",
|
||||
"file_path": str(model_file),
|
||||
"civitai": {},
|
||||
}
|
||||
|
||||
video_path = tmp_path / "custom_ghi.mp4"
|
||||
video_path.write_bytes(b"\x00\x00\x00\x20ftypisom" + b"\xff" * 32)
|
||||
|
||||
monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None))
|
||||
|
||||
_regular, custom = await metadata_module.MetadataUpdater.update_metadata_after_import(
|
||||
model_hash,
|
||||
model_data,
|
||||
StubScanner([model_data]),
|
||||
[(str(video_path), "ghi")],
|
||||
)
|
||||
|
||||
assert (custom[0]["width"], custom[0]["height"]) == (720, 1280)
|
||||
|
||||
|
||||
async def test_update_metadata_from_local_examples_uses_real_video_dimensions(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path
|
||||
):
|
||||
model_hash = "1" * 64
|
||||
model_dir = tmp_path / model_hash
|
||||
model_dir.mkdir()
|
||||
(model_dir / "clip.mp4").write_bytes(build_mp4(1920, 1080))
|
||||
model_data: Dict[str, Any] = {
|
||||
"model_name": "LocalVideo",
|
||||
"civitai": {},
|
||||
"file_path": str(tmp_path / "model.safetensors"),
|
||||
}
|
||||
|
||||
async def fake_save(path, metadata):
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(metadata_module.MetadataManager, "save_metadata", staticmethod(fake_save))
|
||||
|
||||
success = await metadata_module.MetadataUpdater.update_metadata_from_local_examples(
|
||||
model_hash,
|
||||
model_data,
|
||||
"lora",
|
||||
StubScanner([model_data]),
|
||||
str(model_dir),
|
||||
)
|
||||
|
||||
assert success is True
|
||||
entry = model_data["civitai"]["images"][0]
|
||||
assert entry["type"] == "video"
|
||||
assert (entry["width"], entry["height"]) == (1920, 1080)
|
||||
@@ -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