mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-23 14:12:11 -03:00
feat(preview): respect blur mature content setting
This commit is contained in:
@@ -9,6 +9,7 @@ from urllib.parse import urlparse
|
||||
from ..utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
|
||||
from ..utils.constants import CARD_PREVIEW_WIDTH, VALID_LORA_TYPES
|
||||
from ..utils.civitai_utils import rewrite_preview_url
|
||||
from ..utils.preview_selection import select_preview_media
|
||||
from ..utils.utils import sanitize_folder_name
|
||||
from ..utils.exif_utils import ExifUtils
|
||||
from ..utils.metadata_manager import MetadataManager
|
||||
@@ -495,10 +496,21 @@ class DownloadManager:
|
||||
if progress_callback:
|
||||
await progress_callback(1) # 1% progress for starting preview download
|
||||
|
||||
first_image = images[0] if isinstance(images[0], dict) else None
|
||||
preview_url = first_image.get('url') if first_image else None
|
||||
media_type = (first_image.get('type') or '').lower() if first_image else ''
|
||||
nsfw_level = first_image.get('nsfwLevel', 0) if first_image else 0
|
||||
settings_manager = get_settings_manager()
|
||||
blur_mature_content = bool(
|
||||
settings_manager.get('blur_mature_content', True)
|
||||
)
|
||||
selected_image, nsfw_level = select_preview_media(
|
||||
images,
|
||||
blur_mature_content=blur_mature_content,
|
||||
)
|
||||
|
||||
preview_url = selected_image.get('url') if selected_image else None
|
||||
media_type = (
|
||||
(selected_image.get('type') or '').lower()
|
||||
if selected_image
|
||||
else ''
|
||||
)
|
||||
|
||||
def _extension_from_url(url: str, fallback: str) -> str:
|
||||
try:
|
||||
|
||||
@@ -9,6 +9,8 @@ from urllib.parse import urlparse
|
||||
|
||||
from ..utils.constants import CARD_PREVIEW_WIDTH, PREVIEW_EXTENSIONS
|
||||
from ..utils.civitai_utils import rewrite_preview_url
|
||||
from ..utils.preview_selection import select_preview_media
|
||||
from .settings_manager import get_settings_manager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -43,7 +45,18 @@ class PreviewAssetService:
|
||||
if not images:
|
||||
return
|
||||
|
||||
first_preview = images[0]
|
||||
settings_manager = get_settings_manager()
|
||||
blur_mature_content = bool(
|
||||
settings_manager.get("blur_mature_content", True)
|
||||
)
|
||||
first_preview, nsfw_level = select_preview_media(
|
||||
images,
|
||||
blur_mature_content=blur_mature_content,
|
||||
)
|
||||
|
||||
if not first_preview:
|
||||
return
|
||||
|
||||
base_name = os.path.splitext(os.path.splitext(os.path.basename(metadata_path))[0])[0]
|
||||
preview_dir = os.path.dirname(metadata_path)
|
||||
is_video = first_preview.get("type") == "video"
|
||||
@@ -81,7 +94,7 @@ class PreviewAssetService:
|
||||
success, _ = await downloader.download_file(candidate, preview_path, use_auth=False)
|
||||
if success:
|
||||
local_metadata["preview_url"] = preview_path.replace(os.sep, "/")
|
||||
local_metadata["preview_nsfw_level"] = first_preview.get("nsfwLevel", 0)
|
||||
local_metadata["preview_nsfw_level"] = nsfw_level
|
||||
return
|
||||
else:
|
||||
rewritten_url, rewritten = rewrite_preview_url(preview_url, media_type="image")
|
||||
@@ -93,7 +106,7 @@ class PreviewAssetService:
|
||||
)
|
||||
if success:
|
||||
local_metadata["preview_url"] = preview_path.replace(os.sep, "/")
|
||||
local_metadata["preview_nsfw_level"] = first_preview.get("nsfwLevel", 0)
|
||||
local_metadata["preview_nsfw_level"] = nsfw_level
|
||||
return
|
||||
|
||||
extension = ".webp"
|
||||
@@ -124,7 +137,7 @@ class PreviewAssetService:
|
||||
return
|
||||
|
||||
local_metadata["preview_url"] = preview_path.replace(os.sep, "/")
|
||||
local_metadata["preview_nsfw_level"] = first_preview.get("nsfwLevel", 0)
|
||||
local_metadata["preview_nsfw_level"] = nsfw_level
|
||||
|
||||
async def replace_preview(
|
||||
self,
|
||||
|
||||
63
py/utils/preview_selection.py
Normal file
63
py/utils/preview_selection.py
Normal file
@@ -0,0 +1,63 @@
|
||||
"""Utilities for selecting preview media from Civitai image metadata."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Mapping, Optional, Sequence, Tuple
|
||||
|
||||
from .constants import NSFW_LEVELS
|
||||
|
||||
PreviewMedia = Mapping[str, object]
|
||||
|
||||
|
||||
def _extract_nsfw_level(entry: Mapping[str, object]) -> int:
|
||||
"""Return a normalized NSFW level value for the supplied media entry."""
|
||||
|
||||
value = entry.get("nsfwLevel", 0)
|
||||
try:
|
||||
return int(value) # type: ignore[return-value]
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def select_preview_media(
|
||||
images: Sequence[Mapping[str, object]] | None,
|
||||
*,
|
||||
blur_mature_content: bool,
|
||||
) -> Tuple[Optional[PreviewMedia], int]:
|
||||
"""Select the most appropriate preview media entry.
|
||||
|
||||
When ``blur_mature_content`` is enabled we first try to return the first media
|
||||
item with an ``nsfwLevel`` lower than :pydata:`NSFW_LEVELS["R"]`. If none are
|
||||
available we return the media entry with the lowest NSFW level. When the
|
||||
setting is disabled we simply return the first entry.
|
||||
"""
|
||||
|
||||
if not images:
|
||||
return None, 0
|
||||
|
||||
candidates = [item for item in images if isinstance(item, Mapping)]
|
||||
if not candidates:
|
||||
return None, 0
|
||||
|
||||
selected = candidates[0]
|
||||
selected_level = _extract_nsfw_level(selected)
|
||||
|
||||
if not blur_mature_content:
|
||||
return selected, selected_level
|
||||
|
||||
safe_threshold = NSFW_LEVELS.get("R", 4)
|
||||
for candidate in candidates:
|
||||
level = _extract_nsfw_level(candidate)
|
||||
if level < safe_threshold:
|
||||
return candidate, level
|
||||
|
||||
for candidate in candidates[1:]:
|
||||
level = _extract_nsfw_level(candidate)
|
||||
if level < selected_level:
|
||||
selected = candidate
|
||||
selected_level = level
|
||||
|
||||
return selected, selected_level
|
||||
|
||||
|
||||
__all__ = ["select_preview_media"]
|
||||
Reference in New Issue
Block a user