Will Miao
0956126760
Video preview downloads keep streaming to disk, but the saved file now gets its extension sniffed from the actual bytes and renamed on mismatch, so a future Civitai transcode format change cannot leave e.g. WebM bytes in an .mp4 file. (#1155)
316 lines
12 KiB
Python
316 lines
12 KiB
Python
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import struct
|
|
from typing import Any
|
|
|
|
from .constants import (
|
|
CARD_PREVIEW_WIDTH,
|
|
DEFAULT_HASH_CHUNK_SIZE_MB,
|
|
INVALID_AUTOV3_EMPTY_HASH,
|
|
MAX_SAFETENSORS_HEADER_BYTES,
|
|
PREVIEW_EXTENSIONS,
|
|
SUPPORTED_MEDIA_EXTENSIONS,
|
|
)
|
|
from .exif_utils import ExifUtils
|
|
from ..services.settings_manager import get_settings_manager
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _get_hash_chunk_size_bytes() -> int:
|
|
"""Return the chunk size used for hashing, in bytes."""
|
|
|
|
settings_manager = get_settings_manager()
|
|
chunk_size_mb = settings_manager.get("hash_chunk_size_mb", DEFAULT_HASH_CHUNK_SIZE_MB)
|
|
try:
|
|
chunk_size_value = float(chunk_size_mb)
|
|
except (TypeError, ValueError):
|
|
chunk_size_value = float(DEFAULT_HASH_CHUNK_SIZE_MB)
|
|
|
|
if chunk_size_value <= 0:
|
|
chunk_size_value = float(DEFAULT_HASH_CHUNK_SIZE_MB)
|
|
|
|
return max(1, int(chunk_size_value * 1024 * 1024))
|
|
|
|
|
|
async def calculate_sha256(file_path: str) -> str:
|
|
"""Calculate SHA256 hash of a file (full file content).
|
|
|
|
The blocking read+hash loop runs in a thread executor so callers awaiting
|
|
this coroutine never stall the aiohttp event loop — a multi-GB checkpoint
|
|
hash would otherwise freeze every HTTP endpoint and websocket for the
|
|
duration of the hash.
|
|
|
|
Uses ``posix_fadvise`` with ``POSIX_FADV_DONTNEED`` to avoid polluting the OS page
|
|
cache — critical on WSL where cached file pages live inside the VM and are not
|
|
accounted for in guest ``used`` memory, causing VmmemWSL to balloon.
|
|
|
|
On Windows/macOS where ``posix_fadvise`` is not available the hint is silently
|
|
skipped.
|
|
"""
|
|
loop = asyncio.get_running_loop()
|
|
return await loop.run_in_executor(None, _calculate_sha256_sync, file_path)
|
|
|
|
|
|
def _calculate_sha256_sync(file_path: str) -> str:
|
|
sha256_hash = hashlib.sha256()
|
|
chunk_size = _get_hash_chunk_size_bytes()
|
|
with open(file_path, "rb") as f:
|
|
fd = f.fileno()
|
|
for byte_block in iter(lambda: f.read(chunk_size), b""):
|
|
sha256_hash.update(byte_block)
|
|
# Evict pages after reading so the data doesn't linger in the kernel page
|
|
# cache — on WSL this otherwise appears as unreclaimable VmmemWSL growth.
|
|
# Guard against platforms (Windows, macOS) that lack posix_fadvise.
|
|
if hasattr(os, "posix_fadvise") and hasattr(os, "POSIX_FADV_DONTNEED"):
|
|
os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
|
|
return sha256_hash.hexdigest()
|
|
|
|
|
|
def calculate_autov2(file_path: str) -> str:
|
|
"""Calculate CivitAI AutoV2 hash.
|
|
|
|
AutoV2 is the first 10 characters of the full file SHA256.
|
|
Used by CivitAI as a shortened file identifier.
|
|
|
|
Reference: https://developer.civitai.com/site/reference/model-versions
|
|
"""
|
|
full_hash = hashlib.sha256()
|
|
chunk_size = _get_hash_chunk_size_bytes()
|
|
with open(file_path, "rb") as f:
|
|
for byte_block in iter(lambda: f.read(chunk_size), b""):
|
|
full_hash.update(byte_block)
|
|
return full_hash.hexdigest()[:10]
|
|
|
|
|
|
def read_safetensors_metadata(file_path: str) -> dict[str, Any]:
|
|
"""Read the ``__metadata__`` dict from a safetensors file header.
|
|
|
|
Safetensors file format:
|
|
- 8 bytes: header length (little-endian 64-bit)
|
|
- N bytes: UTF-8 JSON header
|
|
- The header JSON contains a ``__metadata__`` key holding arbitrary metadata.
|
|
|
|
Returns an empty dict if the file is not a valid safetensors file or has no
|
|
metadata.
|
|
"""
|
|
try:
|
|
with open(file_path, "rb") as f:
|
|
header_len_bytes = f.read(8)
|
|
if len(header_len_bytes) < 8:
|
|
return {}
|
|
header_len = struct.unpack("<Q", header_len_bytes)[0]
|
|
if header_len > MAX_SAFETENSORS_HEADER_BYTES:
|
|
return {}
|
|
header_bytes = f.read(header_len)
|
|
if len(header_bytes) < header_len:
|
|
return {}
|
|
header = json.loads(header_bytes.decode("utf-8"))
|
|
return header.get("__metadata__", {})
|
|
except (OSError, json.JSONDecodeError, UnicodeDecodeError, struct.error, MemoryError, Exception):
|
|
return {}
|
|
|
|
|
|
def calculate_autov3(file_path: str) -> str | None:
|
|
"""Calculate CivitAI AutoV3 hash from a safetensors file.
|
|
|
|
AutoV3 is extracted from the safetensors file's embedded metadata, not
|
|
computed from the file bytes directly. The orchestrator reads the
|
|
``sshs_model_hash`` (kohya-ss format) or ``modelspec.hash_sha256`` field
|
|
from the safetensors header and stores the first 12 characters.
|
|
|
|
The embedded hash itself is the SHA256 of the file after skipping the
|
|
8-byte header length + JSON header (a.k.a. the addnet hash / tensor-only
|
|
hash).
|
|
|
|
Reference:
|
|
- CivitAI DB trigger: ``SUBSTRING(NEW.hash FROM 1 FOR 12)``
|
|
- https://developer.civitai.com/site/reference/model-versions
|
|
|
|
Returns ``None`` when no AutoV3 hash can be determined (e.g. the file is
|
|
not safetensors, or the metadata doesn't contain a recognised hash field).
|
|
"""
|
|
metadata = read_safetensors_metadata(file_path)
|
|
if not metadata:
|
|
return None
|
|
|
|
embedded_hash = metadata.get("sshs_model_hash") or metadata.get("modelspec.hash_sha256")
|
|
if embedded_hash and isinstance(embedded_hash, str):
|
|
# OneTrainer writes modelspec.hash_sha256 with a "0x" prefix.
|
|
embedded_hash = embedded_hash.strip().removeprefix("0x").removeprefix("0X")
|
|
if len(embedded_hash) >= 12:
|
|
autov3 = embedded_hash[:12].lower()
|
|
# The empty-string SHA256 placeholder written by some repackaging
|
|
# tools is not a real hash; treat it as unavailable so broken
|
|
# models never share one bogus value.
|
|
if autov3 != INVALID_AUTOV3_EMPTY_HASH:
|
|
return autov3
|
|
|
|
return None
|
|
|
|
|
|
def detect_media_extension(
|
|
content: bytes | None,
|
|
headers: dict | None = None,
|
|
fallback_url: str | None = None,
|
|
media_type_hint: str | None = None,
|
|
default: str = ".jpg",
|
|
) -> str:
|
|
"""Determine a file extension from content magic bytes or headers.
|
|
|
|
CDNs (Civitai's ``optimized=true`` renditions in particular) re-encode
|
|
media regardless of the extension in the URL path, so the URL extension
|
|
must only be used as a fallback after content and headers.
|
|
|
|
Args:
|
|
content: File content bytes (only the first bytes are inspected).
|
|
headers: HTTP response headers.
|
|
fallback_url: Original URL for extension extraction.
|
|
media_type_hint: Optional media type hint from metadata (e.g. "video" or "image").
|
|
default: Extension returned when nothing else matches. Pass ``""`` to
|
|
treat unrecognized content as a detection failure.
|
|
"""
|
|
# Check magic bytes for common formats
|
|
if content:
|
|
if content.startswith(b'\xFF\xD8\xFF'):
|
|
return '.jpg'
|
|
elif content.startswith(b'\x89PNG\r\n\x1A\n'):
|
|
return '.png'
|
|
elif content.startswith(b'GIF87a') or content.startswith(b'GIF89a'):
|
|
return '.gif'
|
|
elif content.startswith(b'RIFF') and b'WEBP' in content[:12]:
|
|
return '.webp'
|
|
elif len(content) >= 12 and content[4:8] == b'ftyp' and b'avif' in content[8:24]:
|
|
return '.avif'
|
|
elif content.startswith(b'\x00\x00\x00\x0cJXL \x0d\x0a\x87\x0a'):
|
|
return '.jxl'
|
|
elif content.startswith(b'\x00\x00\x00\x18ftypmp4') or content.startswith(b'\x00\x00\x00\x20ftypmp4'):
|
|
return '.mp4'
|
|
elif content.startswith(b'\x1A\x45\xDF\xA3'):
|
|
return '.webm'
|
|
|
|
# Check Content-Type header
|
|
if headers:
|
|
content_type = headers.get('content-type', '').lower()
|
|
type_map = {
|
|
'image/jpeg': '.jpg',
|
|
'image/png': '.png',
|
|
'image/gif': '.gif',
|
|
'image/webp': '.webp',
|
|
'image/avif': '.avif',
|
|
'image/jxl': '.jxl',
|
|
'video/mp4': '.mp4',
|
|
'video/webm': '.webm',
|
|
'video/quicktime': '.mov'
|
|
}
|
|
if content_type in type_map:
|
|
return type_map[content_type]
|
|
|
|
# Fallback to URL extension if available
|
|
if fallback_url:
|
|
filename = os.path.basename(fallback_url.split('?')[0])
|
|
ext = os.path.splitext(filename)[1].lower()
|
|
if ext in SUPPORTED_MEDIA_EXTENSIONS['images'] or ext in SUPPORTED_MEDIA_EXTENSIONS['videos']:
|
|
return ext
|
|
|
|
# Use media type hint from metadata if available
|
|
if media_type_hint == "video":
|
|
return '.mp4'
|
|
|
|
# Default fallback
|
|
return default
|
|
|
|
|
|
def reconcile_media_extension(file_path: str) -> str:
|
|
"""Rename *file_path* when its bytes indicate a different media format.
|
|
|
|
Returns the final path: the renamed path when the sniffed extension is a
|
|
supported media extension that differs from the current one, otherwise
|
|
*file_path* unchanged.
|
|
"""
|
|
current_ext = os.path.splitext(file_path)[1].lower()
|
|
try:
|
|
with open(file_path, "rb") as f:
|
|
head = f.read(64)
|
|
except OSError:
|
|
return file_path
|
|
detected = detect_media_extension(head, default="")
|
|
if not detected or detected == current_ext:
|
|
return file_path
|
|
if {detected, current_ext} <= {".jpg", ".jpeg"}:
|
|
return file_path
|
|
supported = (
|
|
SUPPORTED_MEDIA_EXTENSIONS["images"] + SUPPORTED_MEDIA_EXTENSIONS["videos"]
|
|
)
|
|
if detected not in supported:
|
|
return file_path
|
|
new_path = os.path.splitext(file_path)[0] + detected
|
|
try:
|
|
os.replace(file_path, new_path)
|
|
except OSError:
|
|
return file_path
|
|
return new_path
|
|
|
|
|
|
def find_preview_file(base_name: str, dir_path: str) -> str:
|
|
"""Find preview file for given base name in directory.
|
|
|
|
Performs an exact-case check first (fast path), then falls back to a
|
|
case-insensitive scan so that files like ``model.WEBP`` or ``model.Png``
|
|
are discovered on case-sensitive filesystems.
|
|
"""
|
|
|
|
temp_extensions = PREVIEW_EXTENSIONS.copy()
|
|
# Add example extension for compatibility
|
|
# https://github.com/willmiao/ComfyUI-Lora-Manager/issues/225
|
|
# The preview image will be optimized to lora-name.webp, so it won't affect other logic
|
|
temp_extensions.append(".example.0.jpeg")
|
|
|
|
# Fast path: exact-case match
|
|
for ext in temp_extensions:
|
|
full_pattern = os.path.join(dir_path, f"{base_name}{ext}")
|
|
if os.path.exists(full_pattern):
|
|
return full_pattern.replace(os.sep, "/")
|
|
|
|
# Slow path: case-insensitive match for systems with mixed-case extensions
|
|
# (e.g. .WEBP, .Png, .JPG placed manually or by external tools)
|
|
try:
|
|
dir_entries = os.listdir(dir_path)
|
|
except OSError:
|
|
return ""
|
|
|
|
base_lower = base_name.lower()
|
|
for ext in temp_extensions:
|
|
target = f"{base_lower}{ext}" # ext is already lowercase
|
|
for entry in dir_entries:
|
|
if entry.lower() == target:
|
|
return os.path.join(dir_path, entry).replace(os.sep, "/")
|
|
|
|
return ""
|
|
|
|
def get_preview_extension(preview_path: str) -> str:
|
|
"""Get the complete preview extension from a preview file path
|
|
|
|
Args:
|
|
preview_path: Path to the preview file
|
|
|
|
Returns:
|
|
str: The complete extension (e.g., '.preview.png', '.png', '.webp')
|
|
"""
|
|
preview_path_lower = preview_path.lower()
|
|
|
|
# Check for compound extensions first (longer matches first)
|
|
for ext in sorted(PREVIEW_EXTENSIONS, key=len, reverse=True):
|
|
if preview_path_lower.endswith(ext.lower()):
|
|
return ext
|
|
|
|
return os.path.splitext(preview_path)[1]
|
|
|
|
def normalize_path(path: str) -> str:
|
|
"""Normalize file path to use forward slashes"""
|
|
return path.replace(os.sep, "/") if path else path
|