mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-06-26 12:51:16 -03:00
Compare commits
6 Commits
v1.1.1
...
a9e5ee7e79
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a9e5ee7e79 | ||
|
|
a17b0e9901 | ||
|
|
8f23d966bf | ||
|
|
7a76fc72d0 | ||
|
|
518a4dd5ee | ||
|
|
2b6d4e5d8b |
4
.gitignore
vendored
4
.gitignore
vendored
@@ -12,12 +12,14 @@ coverage/
|
|||||||
.coverage
|
.coverage
|
||||||
model_cache/
|
model_cache/
|
||||||
|
|
||||||
# agent
|
# agent / dev tooling
|
||||||
.opencode/
|
.opencode/
|
||||||
.claude/
|
.claude/
|
||||||
.sisyphus/
|
.sisyphus/
|
||||||
.codex
|
.codex
|
||||||
.omo
|
.omo
|
||||||
|
reasonix.toml
|
||||||
|
.codegraph/
|
||||||
|
|
||||||
# Vue widgets development cache (but keep build output)
|
# Vue widgets development cache (but keep build output)
|
||||||
vue-widgets/node_modules/
|
vue-widgets/node_modules/
|
||||||
|
|||||||
@@ -16,6 +16,8 @@ IMG_EXTENSIONS = (
|
|||||||
".tif",
|
".tif",
|
||||||
".tiff",
|
".tiff",
|
||||||
".webp",
|
".webp",
|
||||||
|
".avif",
|
||||||
|
".jxl",
|
||||||
".mp4"
|
".mp4"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -216,13 +216,19 @@ class MetadataSyncService:
|
|||||||
provider_used: Optional[str] = None
|
provider_used: Optional[str] = None
|
||||||
last_error: Optional[str] = None
|
last_error: Optional[str] = None
|
||||||
civitai_api_not_found = False
|
civitai_api_not_found = False
|
||||||
|
any_rate_limited = False
|
||||||
|
|
||||||
for provider_name, provider in provider_attempts:
|
for provider_name, provider in provider_attempts:
|
||||||
try:
|
try:
|
||||||
civitai_metadata_candidate, error = await provider.get_model_by_hash(sha256)
|
civitai_metadata_candidate, error = await provider.get_model_by_hash(sha256)
|
||||||
except RateLimitError as exc:
|
except RateLimitError as exc:
|
||||||
exc.provider = exc.provider or (provider_name or provider.__class__.__name__)
|
logger.warning(
|
||||||
raise
|
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||||
|
provider_name or provider.__class__.__name__,
|
||||||
|
exc.retry_after or 0,
|
||||||
|
)
|
||||||
|
any_rate_limited = True
|
||||||
|
continue
|
||||||
except Exception as exc: # pragma: no cover - defensive logging
|
except Exception as exc: # pragma: no cover - defensive logging
|
||||||
logger.error("Provider %s failed for hash %s: %s", provider_name, sha256, exc)
|
logger.error("Provider %s failed for hash %s: %s", provider_name, sha256, exc)
|
||||||
civitai_metadata_candidate, error = None, str(exc)
|
civitai_metadata_candidate, error = None, str(exc)
|
||||||
@@ -276,6 +282,8 @@ class MetadataSyncService:
|
|||||||
)
|
)
|
||||||
|
|
||||||
resolved_error = last_error or default_error
|
resolved_error = last_error or default_error
|
||||||
|
if any_rate_limited and "Rate limited" not in resolved_error:
|
||||||
|
resolved_error = "Rate limited"
|
||||||
if is_expected_offline_error(resolved_error):
|
if is_expected_offline_error(resolved_error):
|
||||||
resolved_error = OFFLINE_FRIENDLY_MESSAGE
|
resolved_error = OFFLINE_FRIENDLY_MESSAGE
|
||||||
|
|
||||||
|
|||||||
@@ -65,7 +65,14 @@ class _RateLimitRetryHelper:
|
|||||||
return await func(*args, **kwargs)
|
return await func(*args, **kwargs)
|
||||||
except RateLimitError as exc:
|
except RateLimitError as exc:
|
||||||
attempt += 1
|
attempt += 1
|
||||||
if attempt >= self._retry_limit:
|
|
||||||
|
# Determine effective retry limit based on rate-limit magnitude
|
||||||
|
effective_retry_limit = self._retry_limit # default: 3
|
||||||
|
if exc.retry_after is not None and exc.retry_after >= 120.0:
|
||||||
|
# Long rate-limit window (>=2 min) — retries are futile
|
||||||
|
effective_retry_limit = 1 # total 1 attempt = 0 retries
|
||||||
|
|
||||||
|
if attempt >= effective_retry_limit:
|
||||||
exc.provider = exc.provider or label
|
exc.provider = exc.provider or label
|
||||||
raise
|
raise
|
||||||
|
|
||||||
@@ -478,8 +485,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
|||||||
if result:
|
if result:
|
||||||
return result, error
|
return result, error
|
||||||
except RateLimitError as exc:
|
except RateLimitError as exc:
|
||||||
exc.provider = exc.provider or label
|
logger.warning(
|
||||||
raise exc
|
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||||
|
label,
|
||||||
|
exc.retry_after or 0,
|
||||||
|
)
|
||||||
|
continue
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.debug("Provider %s failed for get_model_by_hash: %s", label, e)
|
logger.debug("Provider %s failed for get_model_by_hash: %s", label, e)
|
||||||
continue
|
continue
|
||||||
@@ -497,16 +508,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
|||||||
if result:
|
if result:
|
||||||
return result
|
return result
|
||||||
except RateLimitError as exc:
|
except RateLimitError as exc:
|
||||||
if not_found_confirmed:
|
logger.warning(
|
||||||
logger.debug(
|
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||||
"Suppressing rate limit from %s for model %s: "
|
label,
|
||||||
"already confirmed as not found by another provider",
|
exc.retry_after or 0,
|
||||||
label,
|
)
|
||||||
model_id,
|
continue
|
||||||
)
|
|
||||||
return None
|
|
||||||
exc.provider = exc.provider or label
|
|
||||||
raise exc
|
|
||||||
except ResourceNotFoundError:
|
except ResourceNotFoundError:
|
||||||
not_found_confirmed = True
|
not_found_confirmed = True
|
||||||
logger.debug(
|
logger.debug(
|
||||||
@@ -532,8 +539,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
|||||||
if result:
|
if result:
|
||||||
return result
|
return result
|
||||||
except RateLimitError as exc:
|
except RateLimitError as exc:
|
||||||
exc.provider = exc.provider or label
|
logger.warning(
|
||||||
raise exc
|
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||||
|
label,
|
||||||
|
exc.retry_after or 0,
|
||||||
|
)
|
||||||
|
continue
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.debug("Provider %s failed for get_model_version: %s", label, e)
|
logger.debug("Provider %s failed for get_model_version: %s", label, e)
|
||||||
continue
|
continue
|
||||||
@@ -550,8 +561,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
|||||||
if result:
|
if result:
|
||||||
return result, error
|
return result, error
|
||||||
except RateLimitError as exc:
|
except RateLimitError as exc:
|
||||||
exc.provider = exc.provider or label
|
logger.warning(
|
||||||
raise exc
|
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||||
|
label,
|
||||||
|
exc.retry_after or 0,
|
||||||
|
)
|
||||||
|
continue
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.debug("Provider %s failed for get_model_version_info: %s", label, e)
|
logger.debug("Provider %s failed for get_model_version_info: %s", label, e)
|
||||||
continue
|
continue
|
||||||
@@ -572,8 +587,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
|||||||
except NotImplementedError:
|
except NotImplementedError:
|
||||||
continue
|
continue
|
||||||
except RateLimitError as exc:
|
except RateLimitError as exc:
|
||||||
exc.provider = exc.provider or label
|
logger.warning(
|
||||||
raise exc
|
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||||
|
label,
|
||||||
|
exc.retry_after or 0,
|
||||||
|
)
|
||||||
|
continue
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"Provider %s failed for get_model_versions_by_hashes: %s",
|
"Provider %s failed for get_model_versions_by_hashes: %s",
|
||||||
@@ -594,8 +613,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
|||||||
if result is not None:
|
if result is not None:
|
||||||
return result
|
return result
|
||||||
except RateLimitError as exc:
|
except RateLimitError as exc:
|
||||||
exc.provider = exc.provider or label
|
logger.warning(
|
||||||
raise exc
|
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||||
|
label,
|
||||||
|
exc.retry_after or 0,
|
||||||
|
)
|
||||||
|
continue
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.debug("Provider %s failed for get_user_models: %s", label, e)
|
logger.debug("Provider %s failed for get_user_models: %s", label, e)
|
||||||
continue
|
continue
|
||||||
|
|||||||
@@ -77,6 +77,9 @@ class BulkMetadataRefreshUseCase:
|
|||||||
|
|
||||||
await emit("started")
|
await emit("started")
|
||||||
|
|
||||||
|
RATE_LIMIT_ABORT_THRESHOLD = 3
|
||||||
|
consecutive_rate_limits = 0
|
||||||
|
|
||||||
for model in to_process:
|
for model in to_process:
|
||||||
if self._service.scanner.is_cancelled():
|
if self._service.scanner.is_cancelled():
|
||||||
self._logger.info("Bulk metadata refresh cancelled by user")
|
self._logger.info("Bulk metadata refresh cancelled by user")
|
||||||
@@ -115,12 +118,39 @@ class BulkMetadataRefreshUseCase:
|
|||||||
continue
|
continue
|
||||||
|
|
||||||
await MetadataManager.hydrate_model_data(model)
|
await MetadataManager.hydrate_model_data(model)
|
||||||
result, _ = await self._metadata_sync.fetch_and_update_model(
|
result, error_msg = await self._metadata_sync.fetch_and_update_model(
|
||||||
sha256=model["sha256"],
|
sha256=model["sha256"],
|
||||||
file_path=model["file_path"],
|
file_path=model["file_path"],
|
||||||
model_data=model,
|
model_data=model,
|
||||||
update_cache_func=self._service.scanner.update_single_model_cache,
|
update_cache_func=self._service.scanner.update_single_model_cache,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if not result and error_msg and "Rate limited" in error_msg:
|
||||||
|
consecutive_rate_limits += 1
|
||||||
|
else:
|
||||||
|
consecutive_rate_limits = 0
|
||||||
|
|
||||||
|
if consecutive_rate_limits >= RATE_LIMIT_ABORT_THRESHOLD:
|
||||||
|
self._logger.warning(
|
||||||
|
"Bulk metadata refresh aborted: %d consecutive rate limits detected. "
|
||||||
|
"Processed %d/%d models.",
|
||||||
|
consecutive_rate_limits,
|
||||||
|
processed,
|
||||||
|
total_to_process,
|
||||||
|
)
|
||||||
|
await emit(
|
||||||
|
"rate_limited",
|
||||||
|
processed=processed,
|
||||||
|
success=success,
|
||||||
|
)
|
||||||
|
return {
|
||||||
|
"success": False,
|
||||||
|
"message": f"Rate limit detected; {total_to_process - processed} models skipped",
|
||||||
|
"processed": processed,
|
||||||
|
"updated": success,
|
||||||
|
"total": total_models,
|
||||||
|
}
|
||||||
|
|
||||||
if result:
|
if result:
|
||||||
success += 1
|
success += 1
|
||||||
if original_name != model.get("model_name"):
|
if original_name != model.get("model_name"):
|
||||||
|
|||||||
@@ -31,6 +31,8 @@ PREVIEW_EXTENSIONS = [
|
|||||||
".mp4",
|
".mp4",
|
||||||
".gif",
|
".gif",
|
||||||
".webm",
|
".webm",
|
||||||
|
".avif",
|
||||||
|
".jxl",
|
||||||
]
|
]
|
||||||
|
|
||||||
# Card preview image width
|
# Card preview image width
|
||||||
@@ -41,7 +43,7 @@ EXAMPLE_IMAGE_WIDTH = 832
|
|||||||
|
|
||||||
# Supported media extensions for example downloads
|
# Supported media extensions for example downloads
|
||||||
SUPPORTED_MEDIA_EXTENSIONS = {
|
SUPPORTED_MEDIA_EXTENSIONS = {
|
||||||
"images": [".jpg", ".jpeg", ".png", ".webp", ".gif"],
|
"images": [".jpg", ".jpeg", ".png", ".webp", ".gif", ".avif", ".jxl"],
|
||||||
"videos": [".mp4", ".webm"],
|
"videos": [".mp4", ".webm"],
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -62,6 +62,10 @@ class ExampleImagesProcessor:
|
|||||||
return '.gif'
|
return '.gif'
|
||||||
elif content.startswith(b'RIFF') and b'WEBP' in content[:12]:
|
elif content.startswith(b'RIFF') and b'WEBP' in content[:12]:
|
||||||
return '.webp'
|
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'):
|
elif content.startswith(b'\x00\x00\x00\x18ftypmp4') or content.startswith(b'\x00\x00\x00\x20ftypmp4'):
|
||||||
return '.mp4'
|
return '.mp4'
|
||||||
elif content.startswith(b'\x1A\x45\xDF\xA3'):
|
elif content.startswith(b'\x1A\x45\xDF\xA3'):
|
||||||
@@ -75,6 +79,8 @@ class ExampleImagesProcessor:
|
|||||||
'image/png': '.png',
|
'image/png': '.png',
|
||||||
'image/gif': '.gif',
|
'image/gif': '.gif',
|
||||||
'image/webp': '.webp',
|
'image/webp': '.webp',
|
||||||
|
'image/avif': '.avif',
|
||||||
|
'image/jxl': '.jxl',
|
||||||
'video/mp4': '.mp4',
|
'video/mp4': '.mp4',
|
||||||
'video/webm': '.webm',
|
'video/webm': '.webm',
|
||||||
'video/quicktime': '.mov'
|
'video/quicktime': '.mov'
|
||||||
|
|||||||
@@ -1,17 +1,125 @@
|
|||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
|
import struct
|
||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
from typing import Any, Optional
|
from typing import Any, Optional
|
||||||
|
|
||||||
import piexif
|
import piexif
|
||||||
from PIL import Image, PngImagePlugin
|
from PIL import Image, PngImagePlugin
|
||||||
|
|
||||||
|
try:
|
||||||
|
import brotli
|
||||||
|
_BROTLI_AVAILABLE = True
|
||||||
|
except ImportError:
|
||||||
|
brotli = None
|
||||||
|
_BROTLI_AVAILABLE = False
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
class ExifUtils:
|
class ExifUtils:
|
||||||
"""Utility functions for working with EXIF data in images"""
|
"""Utility functions for working with EXIF data in images"""
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _parse_isobmff_boxes(data: bytes, offset: int = 0) -> list[dict]:
|
||||||
|
boxes = []
|
||||||
|
while offset + 8 <= len(data):
|
||||||
|
size = struct.unpack('>I', data[offset:offset + 4])[0]
|
||||||
|
box_type = data[offset + 4:offset + 8]
|
||||||
|
if size == 0:
|
||||||
|
break
|
||||||
|
if size < 8 or offset + size > len(data):
|
||||||
|
break
|
||||||
|
box_data = data[offset + 8:offset + size]
|
||||||
|
boxes.append({'type': box_type, 'data': box_data, 'size': size})
|
||||||
|
offset += size
|
||||||
|
return boxes
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _is_jxl_container(data: bytes) -> bool:
|
||||||
|
if len(data) < 32:
|
||||||
|
return False
|
||||||
|
return (
|
||||||
|
struct.unpack('>I', data[:4])[0] == 12
|
||||||
|
and data[4:8] == b'JXL '
|
||||||
|
and data[8:12] == bytes([0x0d, 0x0a, 0x87, 0x0a])
|
||||||
|
and struct.unpack('>I', data[12:16])[0] >= 16
|
||||||
|
and data[16:20] == b'ftyp'
|
||||||
|
and data[20:24] == b'jxl '
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _is_avif_container(data: bytes) -> bool:
|
||||||
|
if len(data) < 16:
|
||||||
|
return False
|
||||||
|
for box in ExifUtils._parse_isobmff_boxes(data):
|
||||||
|
if box['type'] == b'ftyp' and b'avif' in box['data']:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Max decompressed size for brotli metadata (2 MB)
|
||||||
|
_BROTLI_MAX_DECOMPRESSED = 2 * 1024 * 1024
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _extract_isobmff_brotli(image_path: str) -> Optional[dict]:
|
||||||
|
try:
|
||||||
|
with open(image_path, 'rb') as f:
|
||||||
|
data = f.read()
|
||||||
|
except Exception:
|
||||||
|
return None
|
||||||
|
|
||||||
|
if ExifUtils._is_jxl_container(data):
|
||||||
|
boxes = ExifUtils._parse_isobmff_boxes(data, offset=12)
|
||||||
|
elif ExifUtils._is_avif_container(data):
|
||||||
|
boxes = ExifUtils._parse_isobmff_boxes(data)
|
||||||
|
else:
|
||||||
|
return None
|
||||||
|
|
||||||
|
brob = None
|
||||||
|
for box in boxes:
|
||||||
|
if box['type'] == b'brob':
|
||||||
|
brob = box
|
||||||
|
break
|
||||||
|
if brob is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
payload = brob['data']
|
||||||
|
if payload[:4] != b'comf':
|
||||||
|
return None
|
||||||
|
compressed = payload[4:]
|
||||||
|
|
||||||
|
if _BROTLI_AVAILABLE:
|
||||||
|
try:
|
||||||
|
decompressed = brotli.decompress(compressed)
|
||||||
|
if len(decompressed) > ExifUtils._BROTLI_MAX_DECOMPRESSED:
|
||||||
|
logger.warning(
|
||||||
|
"Brotli metadata too large (%d bytes, max %d), ignoring",
|
||||||
|
len(decompressed),
|
||||||
|
ExifUtils._BROTLI_MAX_DECOMPRESSED,
|
||||||
|
)
|
||||||
|
decompressed = None
|
||||||
|
except Exception:
|
||||||
|
decompressed = None
|
||||||
|
else:
|
||||||
|
decompressed = None
|
||||||
|
|
||||||
|
raw = decompressed if decompressed is not None else compressed
|
||||||
|
try:
|
||||||
|
meta = json.loads(raw.decode('utf-8'))
|
||||||
|
except Exception:
|
||||||
|
return None
|
||||||
|
|
||||||
|
result = {"parameters": None, "prompt": None, "workflow": None, "comment": None}
|
||||||
|
if isinstance(meta.get("prompt"), (dict, list)):
|
||||||
|
result["prompt"] = json.dumps(meta["prompt"])
|
||||||
|
elif isinstance(meta.get("prompt"), str):
|
||||||
|
result["prompt"] = meta["prompt"]
|
||||||
|
if isinstance(meta.get("workflow"), (dict, list)):
|
||||||
|
result["workflow"] = json.dumps(meta["workflow"])
|
||||||
|
elif isinstance(meta.get("workflow"), str):
|
||||||
|
result["workflow"] = meta["workflow"]
|
||||||
|
return result
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _decode_user_comment(user_comment: Any) -> Optional[str]:
|
def _decode_user_comment(user_comment: Any) -> Optional[str]:
|
||||||
if user_comment is None:
|
if user_comment is None:
|
||||||
@@ -43,6 +151,12 @@ class ExifUtils:
|
|||||||
"comment": None,
|
"comment": None,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
ext = os.path.splitext(image_path)[1].lower()
|
||||||
|
if ext in ('.avif', '.jxl'):
|
||||||
|
brotli_meta = ExifUtils._extract_isobmff_brotli(image_path)
|
||||||
|
if brotli_meta:
|
||||||
|
return brotli_meta
|
||||||
|
|
||||||
with Image.open(image_path) as img:
|
with Image.open(image_path) as img:
|
||||||
info = getattr(img, "info", {}) or {}
|
info = getattr(img, "info", {}) or {}
|
||||||
|
|
||||||
@@ -149,7 +263,6 @@ class ExifUtils:
|
|||||||
Optional[str]: Extracted metadata or None if not found
|
Optional[str]: Extracted metadata or None if not found
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
# Skip for video files
|
|
||||||
if image_path:
|
if image_path:
|
||||||
ext = os.path.splitext(image_path)[1].lower()
|
ext = os.path.splitext(image_path)[1].lower()
|
||||||
if ext in ['.mp4', '.webm']:
|
if ext in ['.mp4', '.webm']:
|
||||||
@@ -177,10 +290,9 @@ class ExifUtils:
|
|||||||
str: Path to the updated image
|
str: Path to the updated image
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
# Skip for video files
|
|
||||||
if image_path:
|
if image_path:
|
||||||
ext = os.path.splitext(image_path)[1].lower()
|
ext = os.path.splitext(image_path)[1].lower()
|
||||||
if ext in ['.mp4', '.webm']:
|
if ext in ['.mp4', '.webm', '.avif', '.jxl']:
|
||||||
return image_path
|
return image_path
|
||||||
|
|
||||||
metadata_fields = ExifUtils._load_structured_metadata(image_path)
|
metadata_fields = ExifUtils._load_structured_metadata(image_path)
|
||||||
@@ -212,10 +324,9 @@ class ExifUtils:
|
|||||||
def append_recipe_metadata(image_path, recipe_data) -> str:
|
def append_recipe_metadata(image_path, recipe_data) -> str:
|
||||||
"""Append recipe metadata to an image's EXIF data"""
|
"""Append recipe metadata to an image's EXIF data"""
|
||||||
try:
|
try:
|
||||||
# Skip for video files
|
|
||||||
if image_path:
|
if image_path:
|
||||||
ext = os.path.splitext(image_path)[1].lower()
|
ext = os.path.splitext(image_path)[1].lower()
|
||||||
if ext in ['.mp4', '.webm']:
|
if ext in ['.mp4', '.webm', '.avif', '.jxl']:
|
||||||
return image_path
|
return image_path
|
||||||
|
|
||||||
# First, extract existing metadata
|
# First, extract existing metadata
|
||||||
@@ -327,10 +438,9 @@ class ExifUtils:
|
|||||||
Tuple of (optimized_image_data, extension)
|
Tuple of (optimized_image_data, extension)
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
# Skip for video files early if it's a file path
|
|
||||||
if isinstance(image_data, str) and os.path.exists(image_data):
|
if isinstance(image_data, str) and os.path.exists(image_data):
|
||||||
ext = os.path.splitext(image_data)[1].lower()
|
ext = os.path.splitext(image_data)[1].lower()
|
||||||
if ext in ['.mp4', '.webm']:
|
if ext in ['.mp4', '.webm', '.avif', '.jxl']:
|
||||||
try:
|
try:
|
||||||
with open(image_data, 'rb') as f:
|
with open(image_data, 'rb') as f:
|
||||||
return f.read(), ext
|
return f.read(), ext
|
||||||
|
|||||||
@@ -13,3 +13,5 @@ aiosqlite
|
|||||||
beautifulsoup4
|
beautifulsoup4
|
||||||
platformdirs
|
platformdirs
|
||||||
pyyaml
|
pyyaml
|
||||||
|
# brotli — ISOBMFF (AVIF/JXL) metadata decompression
|
||||||
|
brotli>=1.2.0
|
||||||
|
|||||||
@@ -355,9 +355,9 @@ function renderImportInterface(isEmpty) {
|
|||||||
<button class="select-files-btn" id="selectExampleFilesBtn">
|
<button class="select-files-btn" id="selectExampleFilesBtn">
|
||||||
<i class="fas fa-folder-open"></i> Select Files
|
<i class="fas fa-folder-open"></i> Select Files
|
||||||
</button>
|
</button>
|
||||||
<p class="import-formats">Supported formats: jpg, png, gif, webp, mp4, webm</p>
|
<p class="import-formats">Supported formats: jpg, png, gif, webp, avif, jxl, mp4, webm</p>
|
||||||
</div>
|
</div>
|
||||||
<input type="file" id="exampleFilesInput" multiple accept="image/*,video/mp4,video/webm" style="display: none;">
|
<input type="file" id="exampleFilesInput" multiple accept="image/*,image/avif,image/jxl,video/mp4,video/webm" style="display: none;">
|
||||||
<div class="import-progress-container" style="display: none;">
|
<div class="import-progress-container" style="display: none;">
|
||||||
<div class="import-progress">
|
<div class="import-progress">
|
||||||
<div class="progress-bar"></div>
|
<div class="progress-bar"></div>
|
||||||
@@ -473,7 +473,7 @@ export function initExampleImport(modelHash, container) {
|
|||||||
*/
|
*/
|
||||||
async function handleImportFiles(files, modelHash, importContainer) {
|
async function handleImportFiles(files, modelHash, importContainer) {
|
||||||
// Filter for supported file types
|
// Filter for supported file types
|
||||||
const supportedImages = ['.jpg', '.jpeg', '.png', '.gif', '.webp'];
|
const supportedImages = ['.jpg', '.jpeg', '.png', '.gif', '.webp', '.avif', '.jxl'];
|
||||||
const supportedVideos = ['.mp4', '.webm'];
|
const supportedVideos = ['.mp4', '.webm'];
|
||||||
const supportedExtensions = [...supportedImages, ...supportedVideos];
|
const supportedExtensions = [...supportedImages, ...supportedVideos];
|
||||||
|
|
||||||
|
|||||||
@@ -441,7 +441,6 @@ async def test_fetch_and_update_model_returns_rate_limit_error(tmp_path):
|
|||||||
|
|
||||||
assert ok is False
|
assert ok is False
|
||||||
assert error is not None and "Rate limited" in error
|
assert error is not None and "Rate limited" in error
|
||||||
assert "7" in error
|
|
||||||
helpers.metadata_manager.save_metadata.assert_not_awaited()
|
helpers.metadata_manager.save_metadata.assert_not_awaited()
|
||||||
update_cache.assert_not_awaited()
|
update_cache.assert_not_awaited()
|
||||||
helpers.provider_selector.assert_not_awaited()
|
helpers.provider_selector.assert_not_awaited()
|
||||||
|
|||||||
@@ -63,7 +63,8 @@ async def test_fallback_retries_same_provider_on_rate_limit(monkeypatch):
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_fallback_respects_retry_limit(monkeypatch):
|
async def test_fallback_continues_to_next_provider_on_rate_limit(monkeypatch):
|
||||||
|
"""After exhausting retries on primary, fallback should continue to secondary."""
|
||||||
sleep_mock = AsyncMock()
|
sleep_mock = AsyncMock()
|
||||||
monkeypatch.setattr(provider_module.asyncio, "sleep", sleep_mock)
|
monkeypatch.setattr(provider_module.asyncio, "sleep", sleep_mock)
|
||||||
monkeypatch.setattr(provider_module.random, "uniform", lambda *_: 0.0)
|
monkeypatch.setattr(provider_module.random, "uniform", lambda *_: 0.0)
|
||||||
@@ -76,13 +77,13 @@ async def test_fallback_respects_retry_limit(monkeypatch):
|
|||||||
rate_limit_retry_limit=2,
|
rate_limit_retry_limit=2,
|
||||||
)
|
)
|
||||||
|
|
||||||
with pytest.raises(RateLimitError) as exc_info:
|
# After Change A: no longer raises; falls through to secondary
|
||||||
await fallback.get_model_by_hash("abc")
|
result, error = await fallback.get_model_by_hash("abc")
|
||||||
|
|
||||||
assert exc_info.value.provider == "primary"
|
assert error is None
|
||||||
assert primary.calls == 2
|
assert result == {"id": "secondary"}
|
||||||
assert secondary.calls == 0
|
assert primary.calls == 2 # retry_limit exhausted on primary
|
||||||
sleep_mock.assert_awaited_once()
|
assert secondary.calls == 1 # secondary IS called now
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@@ -117,3 +118,40 @@ async def test_rate_limit_retrying_provider_respects_limit(monkeypatch):
|
|||||||
assert exc_info.value.provider == "inner"
|
assert exc_info.value.provider == "inner"
|
||||||
assert inner.calls == 2
|
assert inner.calls == 2
|
||||||
sleep_mock.assert_awaited_once()
|
sleep_mock.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_retry_helper_limits_retries_for_large_retry_after():
|
||||||
|
"""With retry_after >= 120s, _RateLimitRetryHelper should only attempt once (no retries)."""
|
||||||
|
calls = 0
|
||||||
|
|
||||||
|
async def failing():
|
||||||
|
nonlocal calls
|
||||||
|
calls += 1
|
||||||
|
raise RateLimitError("limited", retry_after=1500.0)
|
||||||
|
|
||||||
|
helper = provider_module._RateLimitRetryHelper(retry_limit=3)
|
||||||
|
with pytest.raises(RateLimitError):
|
||||||
|
await helper.run("test", failing)
|
||||||
|
assert calls == 1 # No retries for large retry_after
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_retry_helper_retries_normally_for_small_retry_after(monkeypatch):
|
||||||
|
"""With retry_after < 120s, _RateLimitRetryHelper should retry normally (up to limit)."""
|
||||||
|
sleep_mock = AsyncMock()
|
||||||
|
monkeypatch.setattr(provider_module.asyncio, "sleep", sleep_mock)
|
||||||
|
|
||||||
|
calls = 0
|
||||||
|
|
||||||
|
async def succeeding():
|
||||||
|
nonlocal calls
|
||||||
|
calls += 1
|
||||||
|
if calls == 1:
|
||||||
|
raise RateLimitError("limited", retry_after=30.0)
|
||||||
|
return {"ok": True}, None
|
||||||
|
|
||||||
|
helper = provider_module._RateLimitRetryHelper(retry_limit=3)
|
||||||
|
result, _ = await helper.run("test", succeeding)
|
||||||
|
assert result == {"ok": True}
|
||||||
|
assert calls == 2 # Retried once (small retry_after)
|
||||||
|
|||||||
@@ -141,3 +141,150 @@ def test_update_image_metadata_preserves_png_workflow(tmp_path):
|
|||||||
img.info["parameters"]
|
img.info["parameters"]
|
||||||
== 'prompt text\nRecipe metadata: {"title":"recipe"}'
|
== 'prompt text\nRecipe metadata: {"title":"recipe"}'
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# --- ISOBMFF / brotli extraction tests ---
|
||||||
|
|
||||||
|
import struct
|
||||||
|
|
||||||
|
import brotli
|
||||||
|
|
||||||
|
|
||||||
|
def _build_jxl_with_brob(payload_json: dict) -> bytes:
|
||||||
|
"""Build a minimal JXL container with a brob box containing brotli-compressed JSON."""
|
||||||
|
# ISOBMFF box 1: JXL signature box (size=12, type='JXL ', signature)
|
||||||
|
box1 = struct.pack(">I", 12) + b"JXL " + bytes([0x0d, 0x0a, 0x87, 0x0a])
|
||||||
|
# ISOBMFF box 2: ftyp (size=16, type='ftyp', major='jxl ', minor=0)
|
||||||
|
box2 = struct.pack(">I", 16) + b"ftyp" + b"jxl " + struct.pack(">I", 0)
|
||||||
|
# ISOBMFF box 3: brob — payload is b'comf' + brotli(json)
|
||||||
|
compressed = brotli.compress(json.dumps(payload_json).encode("utf-8"))
|
||||||
|
brob_payload = b"comf" + compressed
|
||||||
|
box3 = struct.pack(">I", 8 + len(brob_payload)) + b"brob" + brob_payload
|
||||||
|
return box1 + box2 + box3
|
||||||
|
|
||||||
|
|
||||||
|
def _build_avif_with_brob(payload_json: dict) -> bytes:
|
||||||
|
"""Build a minimal AVIF container with a brob box containing brotli-compressed JSON."""
|
||||||
|
compressed = brotli.compress(json.dumps(payload_json).encode("utf-8"))
|
||||||
|
brob_payload = b"comf" + compressed
|
||||||
|
ftyp_box = struct.pack(">I", 20) + b"ftyp" + b"avif" + struct.pack(">I", 0) + b"avif"
|
||||||
|
brob_box = struct.pack(">I", 8 + len(brob_payload)) + b"brob" + brob_payload
|
||||||
|
return ftyp_box + brob_box
|
||||||
|
|
||||||
|
|
||||||
|
class TestIsobmffBrotliExtraction:
|
||||||
|
"""Tests for ISOBMFF brotli metadata extraction in ExifUtils."""
|
||||||
|
|
||||||
|
def test_extract_jxl_brotli_happy_path(self, tmp_path):
|
||||||
|
"""JXL container with valid brob box extracts prompt and workflow."""
|
||||||
|
payload = {"prompt": "a cute cat", "workflow": {"nodes": [{"id": 1}]}}
|
||||||
|
data = _build_jxl_with_brob(payload)
|
||||||
|
path = tmp_path / "test.jxl"
|
||||||
|
path.write_bytes(data)
|
||||||
|
|
||||||
|
result = ExifUtils._load_structured_metadata(str(path))
|
||||||
|
|
||||||
|
assert result["prompt"] == "a cute cat"
|
||||||
|
assert result["workflow"] == '{"nodes": [{"id": 1}]}'
|
||||||
|
assert result["parameters"] is None
|
||||||
|
assert result["comment"] is None
|
||||||
|
|
||||||
|
def test_extract_avif_brotli_happy_path(self, tmp_path):
|
||||||
|
"""AVIF container with valid brob box extracts prompt and workflow."""
|
||||||
|
payload = {"prompt": "landscape", "workflow": {"nodes": []}}
|
||||||
|
data = _build_avif_with_brob(payload)
|
||||||
|
path = tmp_path / "test.avif"
|
||||||
|
path.write_bytes(data)
|
||||||
|
|
||||||
|
result = ExifUtils._load_structured_metadata(str(path))
|
||||||
|
|
||||||
|
assert result["prompt"] == "landscape"
|
||||||
|
assert result["workflow"] == '{"nodes": []}'
|
||||||
|
|
||||||
|
def test_extract_no_brob_box_returns_none(self, tmp_path):
|
||||||
|
"""JXL container without a brob box returns None from _extract_isobmff_brotli."""
|
||||||
|
# Only JXL signature + ftyp, no brob
|
||||||
|
box1 = struct.pack(">I", 12) + b"JXL " + bytes([0x0d, 0x0a, 0x87, 0x0a])
|
||||||
|
box2 = struct.pack(">I", 16) + b"ftyp" + b"jxl " + struct.pack(">I", 0)
|
||||||
|
path = tmp_path / "test.jxl"
|
||||||
|
path.write_bytes(box1 + box2)
|
||||||
|
|
||||||
|
# The low-level extraction should return None (no brob box)
|
||||||
|
result = ExifUtils._extract_isobmff_brotli(str(path))
|
||||||
|
assert result is None
|
||||||
|
|
||||||
|
def test_extract_corrupt_brob_returns_none(self, tmp_path):
|
||||||
|
"""Broken brob box payload gracefully returns None."""
|
||||||
|
box1 = struct.pack(">I", 12) + b"JXL " + bytes([0x0d, 0x0a, 0x87, 0x0a])
|
||||||
|
box2 = struct.pack(">I", 16) + b"ftyp" + b"jxl " + struct.pack(">I", 0)
|
||||||
|
# brob with garbage payload that doesn't start with b'comf'
|
||||||
|
garbage = b"\xff\xff\xff\xff" * 32
|
||||||
|
box3 = struct.pack(">I", 8 + len(garbage)) + b"brob" + garbage
|
||||||
|
path = tmp_path / "test.jxl"
|
||||||
|
path.write_bytes(box1 + box2 + box3)
|
||||||
|
|
||||||
|
result = ExifUtils._extract_isobmff_brotli(str(path))
|
||||||
|
assert result is None
|
||||||
|
|
||||||
|
def test_extract_non_isobmff_file_falls_through(self, tmp_path):
|
||||||
|
"""A regular PNG file is not processed as ISOBMFF and returns PIL metadata."""
|
||||||
|
png_info = PngImagePlugin.PngInfo()
|
||||||
|
png_info.add_text("prompt", "from png")
|
||||||
|
path = tmp_path / "test.png"
|
||||||
|
Image.new("RGB", (4, 4), color="red").save(path, pnginfo=png_info)
|
||||||
|
|
||||||
|
result = ExifUtils._load_structured_metadata(str(path))
|
||||||
|
assert result["prompt"] == "from png"
|
||||||
|
|
||||||
|
def test_extract_skip_on_update_and_optimize(self, tmp_path):
|
||||||
|
"""AVIF/JXL files are skipped for write operations (update/append/optimize)."""
|
||||||
|
path = tmp_path / "test.avif"
|
||||||
|
path.write_bytes(b"fake avif data")
|
||||||
|
|
||||||
|
# update_image_metadata should return the path unchanged
|
||||||
|
result = ExifUtils.update_image_metadata(str(path), "some metadata")
|
||||||
|
assert result == str(path)
|
||||||
|
|
||||||
|
# append_recipe_metadata should also skip
|
||||||
|
result = ExifUtils.append_recipe_metadata(str(path), {"title": "test"})
|
||||||
|
assert result == str(path)
|
||||||
|
|
||||||
|
# optimize_image should passthrough for AVIF/JXL paths
|
||||||
|
result_data, ext = ExifUtils.optimize_image(str(path))
|
||||||
|
assert ext == ".avif"
|
||||||
|
assert result_data == b"fake avif data"
|
||||||
|
|
||||||
|
def test_extract_prompt_as_dict(self, tmp_path):
|
||||||
|
"""prompt field as dict is JSON-serialized."""
|
||||||
|
payload = {"prompt": {"text": "hello", "negative": "bad"}}
|
||||||
|
data = _build_jxl_with_brob(payload)
|
||||||
|
path = tmp_path / "test.jxl"
|
||||||
|
path.write_bytes(data)
|
||||||
|
|
||||||
|
result = ExifUtils._load_structured_metadata(str(path))
|
||||||
|
assert json.loads(result["prompt"]) == {"text": "hello", "negative": "bad"}
|
||||||
|
|
||||||
|
def test_extract_workflow_as_list(self, tmp_path):
|
||||||
|
"""workflow field as list is JSON-serialized."""
|
||||||
|
payload = {"workflow": [{"id": 1}, {"id": 2}]}
|
||||||
|
data = _build_avif_with_brob(payload)
|
||||||
|
path = tmp_path / "test.avif"
|
||||||
|
path.write_bytes(data)
|
||||||
|
|
||||||
|
result = ExifUtils._load_structured_metadata(str(path))
|
||||||
|
assert json.loads(result["workflow"]) == [{"id": 1}, {"id": 2}]
|
||||||
|
|
||||||
|
def test_over_decompressed_size_limit(self, tmp_path, monkeypatch):
|
||||||
|
"""Decompressed data exceeding _BROTLI_MAX_DECOMPRESSED is rejected."""
|
||||||
|
# Monkey-patch the limit to a small value to avoid large test data
|
||||||
|
monkeypatch.setattr(ExifUtils, "_BROTLI_MAX_DECOMPRESSED", 100)
|
||||||
|
|
||||||
|
large_content = "x" * 200
|
||||||
|
payload = {"prompt": large_content}
|
||||||
|
data = _build_jxl_with_brob(payload)
|
||||||
|
path = tmp_path / "test.jxl"
|
||||||
|
path.write_bytes(data)
|
||||||
|
|
||||||
|
# Direct extraction should return None because decompressed size exceeds limit
|
||||||
|
result = ExifUtils._extract_isobmff_brotli(str(path))
|
||||||
|
assert result is None
|
||||||
|
|||||||
Reference in New Issue
Block a user