From a6e23a7630be18fdc14848dfbe0c64ff7ddc654b Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Tue, 9 Dec 2025 20:37:09 +0800 Subject: [PATCH] feat(example-images): add NSFW level setting endpoint Add new POST endpoint `/api/lm/example-images/set-nsfw-level` to allow updating NSFW classification for individual example images. The endpoint supports both regular and custom images, validates required parameters, and updates the corresponding model metadata. This enables users to manually adjust NSFW ratings for better content filtering. --- py/routes/example_images_route_registrar.py | 1 + py/routes/handlers/example_images_handlers.py | 4 + py/utils/example_images_processor.py | 109 +++++++++++++ static/css/components/lora-modal/showcase.css | 8 +- .../ContextMenu/CheckpointContextMenu.js | 8 +- .../ContextMenu/EmbeddingContextMenu.js | 6 +- .../components/ContextMenu/LoraContextMenu.js | 8 +- .../ContextMenu/ModelContextMenuMixin.js | 146 +++++------------ .../ContextMenu/RecipeContextMenu.js | 8 +- .../js/components/shared/NsfwLevelSelector.js | 126 +++++++++++++++ .../components/shared/showcase/MediaUtils.js | 148 +++++++++++++++++- .../shared/showcase/ShowcaseView.js | 7 + static/js/managers/BulkManager.js | 41 ++--- ...example_images_route_registrar_handlers.py | 5 + tests/routes/test_example_images_routes.py | 26 +++ 15 files changed, 486 insertions(+), 165 deletions(-) create mode 100644 static/js/components/shared/NsfwLevelSelector.js diff --git a/py/routes/example_images_route_registrar.py b/py/routes/example_images_route_registrar.py index 63bac34b..e3e4b564 100644 --- a/py/routes/example_images_route_registrar.py +++ b/py/routes/example_images_route_registrar.py @@ -29,6 +29,7 @@ ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = ( RouteDefinition("POST", "/api/lm/delete-example-image", "delete_example_image"), RouteDefinition("POST", "/api/lm/force-download-example-images", "force_download_example_images"), RouteDefinition("POST", "/api/lm/cleanup-example-image-folders", "cleanup_example_image_folders"), + RouteDefinition("POST", "/api/lm/example-images/set-nsfw-level", "set_example_image_nsfw_level"), ) diff --git a/py/routes/handlers/example_images_handlers.py b/py/routes/handlers/example_images_handlers.py index f7acf5ef..8fe293d6 100644 --- a/py/routes/handlers/example_images_handlers.py +++ b/py/routes/handlers/example_images_handlers.py @@ -113,6 +113,9 @@ class ExampleImagesManagementHandler: async def delete_example_image(self, request: web.Request) -> web.StreamResponse: return await self._processor.delete_custom_image(request) + async def set_example_image_nsfw_level(self, request: web.Request) -> web.StreamResponse: + return await self._processor.set_example_image_nsfw_level(request) + async def cleanup_example_image_folders(self, request: web.Request) -> web.StreamResponse: result = await self._cleanup_service.cleanup_example_image_folders() @@ -160,6 +163,7 @@ class ExampleImagesHandlerSet: "force_download_example_images": self.download.force_download_example_images, "import_example_images": self.management.import_example_images, "delete_example_image": self.management.delete_example_image, + "set_example_image_nsfw_level": self.management.set_example_image_nsfw_level, "cleanup_example_image_folders": self.management.cleanup_example_image_folders, "open_example_images_folder": self.files.open_example_images_folder, "get_example_image_files": self.files.get_example_image_files, diff --git a/py/utils/example_images_processor.py b/py/utils/example_images_processor.py index b979bdb7..f20ddcf1 100644 --- a/py/utils/example_images_processor.py +++ b/py/utils/example_images_processor.py @@ -593,5 +593,114 @@ class ExampleImagesProcessor: 'error': str(e) }, status=500) + @staticmethod + async def set_example_image_nsfw_level(request: web.Request) -> web.StreamResponse: + """ + Update the NSFW level for a single example image (regular or custom). + """ + try: + data = await request.json() + except Exception: + return web.json_response({'success': False, 'error': 'Invalid JSON body'}, status=400) + + model_hash = data.get('model_hash') + raw_level = data.get('nsfw_level') + source = (data.get('source') or 'civitai').lower() + index = data.get('index') + image_id = data.get('id') + + if model_hash is None or raw_level is None: + return web.json_response( + {'success': False, 'error': 'Missing required parameters: model_hash and nsfw_level'}, + status=400, + ) + + try: + nsfw_level = int(raw_level) + except (TypeError, ValueError): + return web.json_response( + {'success': False, 'error': 'nsfw_level must be an integer'}, status=400 + ) + + if source == 'custom': + if not image_id: + return web.json_response( + {'success': False, 'error': 'Custom images require an id field'}, status=400 + ) + else: + try: + index = int(index) + except (TypeError, ValueError): + return web.json_response( + {'success': False, 'error': 'Regular images require a numeric index'}, status=400 + ) + + try: + lora_scanner = await ServiceRegistry.get_lora_scanner() + checkpoint_scanner = await ServiceRegistry.get_checkpoint_scanner() + embedding_scanner = await ServiceRegistry.get_embedding_scanner() + + model_data = None + scanner = None + + for scan_obj in [lora_scanner, checkpoint_scanner, embedding_scanner]: + if scan_obj.has_hash(model_hash): + cache = await scan_obj.get_cached_data() + for item in cache.raw_data: + if item.get('sha256') == model_hash: + model_data = item + scanner = scan_obj + break + if model_data: + break + + if not model_data: + return web.json_response( + {'success': False, 'error': f"Model with hash {model_hash} not found in cache"}, + status=404, + ) + + await MetadataManager.hydrate_model_data(model_data) + civitai_data = model_data.setdefault('civitai', {}) + regular_images = civitai_data.get('images') or [] + custom_images = civitai_data.get('customImages') or [] + + target_image = None + if source == 'custom': + for image in custom_images: + if image.get('id') == image_id: + target_image = image + break + else: + if 0 <= index < len(regular_images): + target_image = regular_images[index] + + if target_image is None: + return web.json_response( + {'success': False, 'error': 'Target image not found'}, status=404 + ) + + target_image['nsfwLevel'] = nsfw_level + civitai_data['images'] = regular_images + civitai_data['customImages'] = custom_images + + file_path = model_data.get('file_path') + if file_path: + model_copy = model_data.copy() + model_copy.pop('folder', None) + await MetadataManager.save_metadata(file_path, model_copy) + await scanner.update_single_model_cache(file_path, file_path, model_data) + + return web.json_response({ + 'success': True, + 'regular_images': regular_images, + 'custom_images': custom_images, + 'model_file_path': model_data.get('file_path', ''), + 'nsfw_level': nsfw_level + }) + except Exception as exc: + logger.error("Failed to update example image NSFW level: %s", exc, exc_info=True) + return web.json_response({'success': False, 'error': str(exc)}, status=500) + diff --git a/static/css/components/lora-modal/showcase.css b/static/css/components/lora-modal/showcase.css index 252c927a..dc3467fd 100644 --- a/static/css/components/lora-modal/showcase.css +++ b/static/css/components/lora-modal/showcase.css @@ -129,6 +129,12 @@ border-color: var(--lora-accent); } +.media-control-btn.set-nsfw-btn:hover { + background: var(--warning-color, #f0ad4e); + color: #fff; + border-color: var(--warning-color, #f0ad4e); +} + .media-control-btn.example-delete-btn:hover:not(.disabled) { background: var(--lora-error); color: white; @@ -475,4 +481,4 @@ /* For dark theme */ [data-theme="dark"] .import-container { background: rgba(255, 255, 255, 0.03); -} \ No newline at end of file +} diff --git a/static/js/components/ContextMenu/CheckpointContextMenu.js b/static/js/components/ContextMenu/CheckpointContextMenu.js index cc9661e9..0375784c 100644 --- a/static/js/components/ContextMenu/CheckpointContextMenu.js +++ b/static/js/components/ContextMenu/CheckpointContextMenu.js @@ -11,11 +11,7 @@ export class CheckpointContextMenu extends BaseContextMenu { this.modelType = 'checkpoint'; this.resetAndReload = resetAndReload; - // Initialize NSFW Level Selector events only if not already initialized - if (this.nsfwSelector && !this.nsfwSelector.dataset.initialized) { - this.initNSFWSelector(); - this.nsfwSelector.dataset.initialized = 'true'; - } + this.initNSFWSelector(); } // Implementation needed by the mixin @@ -65,4 +61,4 @@ export class CheckpointContextMenu extends BaseContextMenu { } // Mix in shared methods -Object.assign(CheckpointContextMenu.prototype, ModelContextMenuMixin); \ No newline at end of file +Object.assign(CheckpointContextMenu.prototype, ModelContextMenuMixin); diff --git a/static/js/components/ContextMenu/EmbeddingContextMenu.js b/static/js/components/ContextMenu/EmbeddingContextMenu.js index ab26f1a9..24440b28 100644 --- a/static/js/components/ContextMenu/EmbeddingContextMenu.js +++ b/static/js/components/ContextMenu/EmbeddingContextMenu.js @@ -11,11 +11,7 @@ export class EmbeddingContextMenu extends BaseContextMenu { this.modelType = 'embedding'; this.resetAndReload = resetAndReload; - // Initialize NSFW Level Selector events only if not already initialized - if (this.nsfwSelector && !this.nsfwSelector.dataset.initialized) { - this.initNSFWSelector(); - this.nsfwSelector.dataset.initialized = 'true'; - } + this.initNSFWSelector(); } // Implementation needed by the mixin diff --git a/static/js/components/ContextMenu/LoraContextMenu.js b/static/js/components/ContextMenu/LoraContextMenu.js index bce9c8cc..0dbaa7b7 100644 --- a/static/js/components/ContextMenu/LoraContextMenu.js +++ b/static/js/components/ContextMenu/LoraContextMenu.js @@ -12,11 +12,7 @@ export class LoraContextMenu extends BaseContextMenu { this.modelType = 'lora'; this.resetAndReload = resetAndReload; - // Initialize NSFW Level Selector events only if not already initialized - if (this.nsfwSelector && !this.nsfwSelector.dataset.initialized) { - this.initNSFWSelector(); - this.nsfwSelector.dataset.initialized = 'true'; - } + this.initNSFWSelector(); } // Use the saveModelMetadata implementation from loraApi @@ -78,4 +74,4 @@ export class LoraContextMenu extends BaseContextMenu { } // Mix in shared methods -Object.assign(LoraContextMenu.prototype, ModelContextMenuMixin); \ No newline at end of file +Object.assign(LoraContextMenu.prototype, ModelContextMenuMixin); diff --git a/static/js/components/ContextMenu/ModelContextMenuMixin.js b/static/js/components/ContextMenu/ModelContextMenuMixin.js index 5aeca29a..b91dded1 100644 --- a/static/js/components/ContextMenu/ModelContextMenuMixin.js +++ b/static/js/components/ContextMenu/ModelContextMenuMixin.js @@ -5,96 +5,38 @@ import { getModelApiClient, resetAndReload } from '../../api/modelApiFactory.js' import { bulkManager } from '../../managers/BulkManager.js'; import { MODEL_CONFIG } from '../../api/apiConfig.js'; import { translate } from '../../utils/i18nHelpers.js'; +import { getNsfwLevelSelector } from '../shared/NsfwLevelSelector.js'; // Mixin with shared functionality for LoraContextMenu and CheckpointContextMenu export const ModelContextMenuMixin = { // NSFW Selector methods initNSFWSelector() { - // Remove any existing event listeners by cloning and replacing elements - // This is a simple way to ensure we don't have duplicate event listeners - const closeBtn = this.nsfwSelector.querySelector('.close-nsfw-selector'); - const newCloseBtn = closeBtn.cloneNode(true); - closeBtn.parentNode.replaceChild(newCloseBtn, closeBtn); - newCloseBtn.addEventListener('click', () => { - this.nsfwSelector.style.display = 'none'; - this.resetNSFWSelectorState(); - }); - - // Level buttons - const levelButtons = this.nsfwSelector.querySelectorAll('.nsfw-level-btn'); - levelButtons.forEach(btn => { - // Remove any existing event listeners by cloning and replacing the button - const newBtn = btn.cloneNode(true); - btn.parentNode.replaceChild(newBtn, btn); - - newBtn.addEventListener('click', async () => { - const level = parseInt(newBtn.dataset.level); - const mode = this.nsfwSelector.dataset.mode || 'single'; - - if (mode === 'bulk') { - let bulkFilePaths = []; - if (this.nsfwSelector.dataset.bulkFilePaths) { - try { - bulkFilePaths = JSON.parse(this.nsfwSelector.dataset.bulkFilePaths); - } catch (error) { - console.warn('Failed to parse bulk file paths for content rating', error); - } - } - - const success = await bulkManager.setBulkContentRating(level, bulkFilePaths); - if (success) { - this.nsfwSelector.style.display = 'none'; - this.resetNSFWSelectorState(); - } - return; - } - - const filePath = this.nsfwSelector.dataset.cardPath; - - if (!filePath) return; - - try { - await this.saveModelMetadata(filePath, { preview_nsfw_level: level }); - - showToast('toast.contextMenu.contentRatingSet', { level: getNSFWLevelName(level) }, 'success'); - this.nsfwSelector.style.display = 'none'; - this.resetNSFWSelectorState(); - } catch (error) { - showToast('toast.contextMenu.contentRatingFailed', { message: error.message }, 'error'); - } - }); - }); - - // Close when clicking outside - use a named function so we can remove it later - const outsideClickListener = (e) => { - if (this.nsfwSelector.style.display === 'block' && - !this.nsfwSelector.contains(e.target) && - !e.target.closest('.context-menu-item[data-action="set-nsfw"], .context-menu-item[data-action="set-content-rating"]')) { - this.nsfwSelector.style.display = 'none'; - this.resetNSFWSelectorState(); - } - }; - - // Remove previous listener if it exists - if (this._outsideClickListener) { - document.removeEventListener('click', this._outsideClickListener); + if (this._nsfwSelectorInitialized) { + return; } - - // Store and add new listener - this._outsideClickListener = outsideClickListener; - document.addEventListener('click', this._outsideClickListener); + + const selector = getNsfwLevelSelector(); + if (!selector) { + console.warn('NSFW selector element not found'); + return; + } + + this._nsfwSelectorInitialized = true; + this._nsfwSelector = selector; }, resetNSFWSelectorState() { - if (!this.nsfwSelector) return; - delete this.nsfwSelector.dataset.bulkFilePaths; - delete this.nsfwSelector.dataset.mode; - delete this.nsfwSelector.dataset.cardPath; + // maintained for compatibility; no-op with shared selector }, showNSFWLevelSelector(x, y, card) { - const selector = document.getElementById('nsfwLevelSelector'); - const currentLevelEl = document.getElementById('currentNSFWLevel'); + this.initNSFWSelector(); + const selector = this._nsfwSelector || getNsfwLevelSelector(); + + if (!selector) { + console.warn('NSFW selector not available'); + return; + } // Get current NSFW level let currentLevel = 0; @@ -104,44 +46,28 @@ export const ModelContextMenuMixin = { // Update if we have no recorded level but have a dataset attribute if (!currentLevel && card.dataset.nsfwLevel) { - currentLevel = parseInt(card.dataset.nsfwLevel) || 0; + currentLevel = parseInt(card.dataset.nsfwLevel, 10) || 0; } } catch (err) { console.error('Error parsing metadata:', err); } - currentLevelEl.textContent = getNSFWLevelName(currentLevel); - - // Position the selector - if (x && y) { - const viewportWidth = document.documentElement.clientWidth; - const viewportHeight = document.documentElement.clientHeight; - const selectorRect = selector.getBoundingClientRect(); - - // Center the selector if no coordinates provided - let finalX = (viewportWidth - selectorRect.width) / 2; - let finalY = (viewportHeight - selectorRect.height) / 2; - - selector.style.left = `${finalX}px`; - selector.style.top = `${finalY}px`; - } - - // Highlight current level button - selector.querySelectorAll('.nsfw-level-btn').forEach(btn => { - if (parseInt(btn.dataset.level) === currentLevel) { - btn.classList.add('active'); - } else { - btn.classList.remove('active'); - } + const filePath = card.dataset.filepath; + selector.show({ + currentLevel, + onSelect: async (level) => { + if (!filePath) return false; + try { + await this.saveModelMetadata(filePath, { preview_nsfw_level: level }); + showToast('toast.contextMenu.contentRatingSet', { level: getNSFWLevelName(level) }, 'success'); + return true; + } catch (error) { + showToast('toast.contextMenu.contentRatingFailed', { message: error.message }, 'error'); + return false; + } + }, + onClose: () => this.resetNSFWSelectorState(), }); - - // Store reference to current card - selector.dataset.mode = 'single'; - selector.dataset.cardPath = card.dataset.filepath; - delete selector.dataset.bulkFilePaths; - - // Show selector - selector.style.display = 'block'; }, // Civitai re-linking methods diff --git a/static/js/components/ContextMenu/RecipeContextMenu.js b/static/js/components/ContextMenu/RecipeContextMenu.js index e5bb53b8..f9cb9719 100644 --- a/static/js/components/ContextMenu/RecipeContextMenu.js +++ b/static/js/components/ContextMenu/RecipeContextMenu.js @@ -11,11 +11,7 @@ export class RecipeContextMenu extends BaseContextMenu { this.nsfwSelector = document.getElementById('nsfwLevelSelector'); this.modelType = 'recipe'; - // Initialize NSFW Level Selector events only if not already initialized - if (this.nsfwSelector && !this.nsfwSelector.dataset.initialized) { - this.initNSFWSelector(); - this.nsfwSelector.dataset.initialized = 'true'; - } + this.initNSFWSelector(); } // Use the updateRecipeMetadata implementation from recipeApi @@ -286,4 +282,4 @@ export class RecipeContextMenu extends BaseContextMenu { } // Mix in shared methods from ModelContextMenuMixin -Object.assign(RecipeContextMenu.prototype, ModelContextMenuMixin); \ No newline at end of file +Object.assign(RecipeContextMenu.prototype, ModelContextMenuMixin); diff --git a/static/js/components/shared/NsfwLevelSelector.js b/static/js/components/shared/NsfwLevelSelector.js new file mode 100644 index 00000000..098ebdaf --- /dev/null +++ b/static/js/components/shared/NsfwLevelSelector.js @@ -0,0 +1,126 @@ +import { getNSFWLevelName } from '../../utils/uiHelpers.js'; +import { translate } from '../../utils/i18nHelpers.js'; + +let selectorController = null; + +function buildController(selectorElement) { + if (!selectorElement) return null; + + const levelButtons = Array.from(selectorElement.querySelectorAll('.nsfw-level-btn')); + const closeBtn = selectorElement.querySelector('.close-nsfw-selector'); + const currentLevelEl = selectorElement.querySelector('#currentNSFWLevel'); + + let onSelect = null; + let onClose = null; + let isOpen = false; + let ignoreNextOutside = false; + + const setLabel = (level, multipleLabel) => { + if (!currentLevelEl) return; + if (multipleLabel) { + currentLevelEl.textContent = multipleLabel; + return; + } + currentLevelEl.textContent = getNSFWLevelName(level); + }; + + const hide = () => { + selectorElement.style.display = 'none'; + isOpen = false; + if (typeof onClose === 'function') { + onClose(); + } + onSelect = null; + onClose = null; + }; + + const show = ({ currentLevel = 0, onSelect: selectCb, onClose: closeCb, multipleLabel = '' } = {}) => { + onSelect = selectCb || null; + onClose = closeCb || null; + isOpen = true; + ignoreNextOutside = true; // ignore the click that triggered open + + // Position near center of viewport + const viewportWidth = document.documentElement.clientWidth; + const viewportHeight = document.documentElement.clientHeight; + const rect = selectorElement.getBoundingClientRect(); + const finalX = Math.max((viewportWidth - rect.width) / 2, 0); + const finalY = Math.max((viewportHeight - rect.height) / 2, 0); + + selectorElement.style.left = `${finalX}px`; + selectorElement.style.top = `${finalY}px`; + + setLabel(currentLevel, multipleLabel); + + // Highlight selected level (if not showing multiple values) + levelButtons.forEach((btn) => { + const btnLevel = parseInt(btn.dataset.level || '0', 10); + if (multipleLabel) { + btn.classList.remove('active'); + } else if (btnLevel === currentLevel) { + btn.classList.add('active'); + } else { + btn.classList.remove('active'); + } + }); + + selectorElement.style.display = 'block'; + }; + + if (closeBtn) { + closeBtn.addEventListener('click', hide); + } + + document.addEventListener('click', (e) => { + if (!isOpen) return; + if (ignoreNextOutside) { + ignoreNextOutside = false; + return; + } + if (!selectorElement.contains(e.target)) { + hide(); + } + }); + + levelButtons.forEach((btn) => { + btn.addEventListener('click', async () => { + if (!isOpen) return; + const level = parseInt(btn.dataset.level || '0', 10); + if (typeof onSelect === 'function') { + try { + const result = await onSelect(level); + if (result === false) { + return; + } + } catch (error) { + console.error('NSFW selector onSelect failed', error); + return; + } + } + hide(); + }); + }); + + const showMultiple = (labelKey = 'modals.contentRating.multiple') => { + const fallback = 'Multiple values'; + const text = translate(labelKey, {}, fallback); + show({ multipleLabel: text }); + }; + + return { + show, + hide, + showMultiple, + isOpen: () => isOpen, + }; +} + +export function getNsfwLevelSelector() { + if (selectorController) { + return selectorController; + } + + const element = document.getElementById('nsfwLevelSelector'); + selectorController = buildController(element); + return selectorController; +} diff --git a/static/js/components/shared/showcase/MediaUtils.js b/static/js/components/shared/showcase/MediaUtils.js index 7a56f5ab..deae96e7 100644 --- a/static/js/components/shared/showcase/MediaUtils.js +++ b/static/js/components/shared/showcase/MediaUtils.js @@ -3,9 +3,11 @@ * Media-specific utility functions for showcase components * (Moved from uiHelpers.js to better organize code) */ -import { showToast, copyToClipboard } from '../../../utils/uiHelpers.js'; +import { showToast, copyToClipboard, getNSFWLevelName } from '../../../utils/uiHelpers.js'; import { state } from '../../../state/index.js'; import { getModelApiClient } from '../../../api/modelApiFactory.js'; +import { NSFW_LEVELS } from '../../../utils/constants.js'; +import { getNsfwLevelSelector } from '../NsfwLevelSelector.js'; /** * Try to load local image first, fall back to remote if local fails @@ -471,6 +473,9 @@ export function initMediaControlHandlers(container) { // Initialize set preview buttons initSetPreviewHandlers(container); + + // Initialize NSFW level buttons + initSetNsfwHandlers(container); // Media control visibility is now handled in initMetadataPanelHandlers // Any click handlers or other functionality can still be added here @@ -590,4 +595,143 @@ export function positionAllMediaControls(container) { mediaWrappers.forEach(wrapper => { positionMediaControlsInMediaRect(wrapper); }); -} \ No newline at end of file +} + +function applyNsfwLevelChange(mediaWrapper, nsfwLevel) { + if (!mediaWrapper) return; + + const mediaElement = mediaWrapper.querySelector('img, video'); + if (mediaElement) { + mediaElement.dataset.nsfwLevel = String(nsfwLevel); + } + mediaWrapper.dataset.nsfwLevel = String(nsfwLevel); + + const shouldBlur = state.settings.blur_mature_content && nsfwLevel > NSFW_LEVELS.PG13; + let overlay = mediaWrapper.querySelector('.nsfw-overlay'); + let toggleBtn = mediaWrapper.querySelector('.toggle-blur-btn'); + + if (shouldBlur) { + mediaWrapper.classList.add('nsfw-media-wrapper'); + if (mediaElement) { + mediaElement.classList.add('blurred'); + } + if (!overlay) { + overlay = document.createElement('div'); + overlay.className = 'nsfw-overlay'; + overlay.innerHTML = ` +
Mature Content
+ +