feat: Enhance media handling by adding NSFW level support and improving preview image management

py/services/model_cache.py

@@ -32,12 +32,13 @@ class ModelCache:
             all_folders = set(l['folder'] for l in self.raw_data)
             self.folders = sorted(list(all_folders), key=lambda x: x.lower())
 
-    async def update_preview_url(self, file_path: str, preview_url: str) -> bool:
+    async def update_preview_url(self, file_path: str, preview_url: str, preview_nsfw_level: int) -> bool:
         """Update preview_url for a specific model in all cached data
         
         Args:
             file_path: The file path of the model to update
             preview_url: The new preview URL
+            preview_nsfw_level: The NSFW level of the preview
             
         Returns:
             bool: True if the update was successful, False if the model wasn't found
@@ -47,19 +48,9 @@ class ModelCache:
             for item in self.raw_data:
                 if item['file_path'] == file_path:
                     item['preview_url'] = preview_url
+                    item['preview_nsfw_level'] = preview_nsfw_level
                     break
             else:
                 return False  # Model not found
                     
-            # Update in sorted lists (references to the same dict objects)
-            for item in self.sorted_by_name:
-                if item['file_path'] == file_path:
-                    item['preview_url'] = preview_url
-                    break
-                    
-            for item in self.sorted_by_date:
-                if item['file_path'] == file_path:
-                    item['preview_url'] = preview_url
-                    break
-                    
             return True

py/services/model_scanner.py

@@ -1184,12 +1184,13 @@ class ModelScanner:
         """Get list of excluded model file paths"""
         return self._excluded_models.copy()
 
-    async def update_preview_in_cache(self, file_path: str, preview_url: str) -> bool:
+    async def update_preview_in_cache(self, file_path: str, preview_url: str, preview_nsfw_level: int) -> bool:
         """Update preview URL in cache for a specific lora
         
         Args:
             file_path: The file path of the lora to update
             preview_url: The new preview URL
+            preview_nsfw_level: The NSFW level of the preview
             
         Returns:
             bool: True if the update was successful, False if cache doesn't exist or lora wasn't found
@@ -1197,7 +1198,7 @@ class ModelScanner:
         if self._cache is None:
             return False
 
-        updated = await self._cache.update_preview_url(file_path, preview_url)
+        updated = await self._cache.update_preview_url(file_path, preview_url, preview_nsfw_level)
         if updated:
             # Save updated cache to disk
             await self._save_cache_to_disk()

py/utils/routes_common.py

@@ -409,6 +409,15 @@ class ModelRouteUtils:
                 raise ValueError("Expected 'model_path' field")
             model_path = (await field.read()).decode()
             
+            # Read NSFW level (new parameter)
+            nsfw_level = 0  # Default to 0 (unknown)
+            field = await reader.next()
+            if field and field.name == 'nsfw_level':
+                try:
+                    nsfw_level = int((await field.read()).decode())
+                except (ValueError, TypeError):
+                    logger.warning("Invalid NSFW level format, using default 0")
+            
             # Save preview file
             base_name = os.path.splitext(os.path.basename(model_path))[0]
             folder = os.path.dirname(model_path)
@@ -435,7 +444,7 @@ class ModelRouteUtils:
                 if os.path.exists(existing_preview):
                     try:
                         os.remove(existing_preview)
-                        logger.info(f"Deleted existing preview: {existing_preview}")
+                        logger.debug(f"Deleted existing preview: {existing_preview}")
                     except Exception as e:
                         logger.warning(f"Failed to delete existing preview {existing_preview}: {e}")
             
@@ -444,26 +453,28 @@ class ModelRouteUtils:
             with open(preview_path, 'wb') as f:
                 f.write(optimized_data)
             
-            # Update preview path in metadata
+            # Update preview path and NSFW level in metadata
             metadata_path = os.path.splitext(model_path)[0] + '.metadata.json'
             if os.path.exists(metadata_path):
                 try:
                     with open(metadata_path, 'r', encoding='utf-8') as f:
                         metadata = json.load(f)
                     
-                    # Update preview_url directly in the metadata dict
+                    # Update preview_url and preview_nsfw_level in the metadata dict
                     metadata['preview_url'] = preview_path
+                    metadata['preview_nsfw_level'] = nsfw_level
                     
                     await MetadataManager.save_metadata(model_path, metadata)
                 except Exception as e:
                     logger.error(f"Error updating metadata: {e}")
             
             # Update preview URL in scanner cache
-            await scanner.update_preview_in_cache(model_path, preview_path)
+            await scanner.update_preview_in_cache(model_path, preview_path, nsfw_level)
             
             return web.json_response({
                 "success": True,
-                "preview_url": config.get_preview_static_url(preview_path)
+                "preview_url": config.get_preview_static_url(preview_path),
+                "preview_nsfw_level": nsfw_level
             })
             
         except Exception as e:

static/css/components/lora-modal/showcase.css

@@ -94,7 +94,6 @@
   pointer-events: none;
 }
 
-.media-wrapper:hover .media-controls,
 .media-controls.visible {
   opacity: 1;
   transform: translateY(0);
@@ -115,6 +114,8 @@
   transition: all 0.2s ease;
   box-shadow: 0 2px 5px rgba(0, 0, 0, 0.15);
   padding: 0;
+  position: relative;
+  overflow: hidden;
 }
 
 .media-control-btn:hover {
@@ -128,18 +129,47 @@
   border-color: var(--lora-accent);
 }
 
-.media-control-btn.example-delete-btn:hover {
+.media-control-btn.example-delete-btn:hover:not(.disabled) {
   background: var(--lora-error);
   color: white;
   border-color: var(--lora-error);
 }
 
+/* Disabled state for delete button */
+.media-control-btn.example-delete-btn.disabled {
+  opacity: 0.5;
+  cursor: not-allowed;
+}
+
 /* Two-step confirmation for delete button */
+.media-control-btn.example-delete-btn .confirm-icon {
+  position: absolute;
+  top: 0;
+  left: 0;
+  right: 0;
+  bottom: 0;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  background: var(--lora-error);
+  color: white;
+  font-size: 1em;
+  opacity: 0;
+  transition: opacity 0.2s ease;
+}
+
+.media-control-btn.example-delete-btn.confirm .fa-trash-alt {
+  opacity: 0;
+}
+
+.media-control-btn.example-delete-btn.confirm .confirm-icon {
+  opacity: 1;
+}
+
 .media-control-btn.example-delete-btn.confirm {
   background: var(--lora-error);
   color: white;
   border-color: var(--lora-error);
-  animation: pulse 1.5s infinite;
 }
 
 @keyframes pulse {

static/js/api/baseModelApi.js

@@ -542,19 +542,17 @@ export async function excludeModel(filePath, modelType = 'lora') {
     }
 }
 
-// Private methods
-
 // Upload a preview image
-async function uploadPreview(filePath, file, modelType = 'lora') {
+export async function uploadPreview(filePath, file, modelType = 'lora', nsfwLevel = 0) {
     try {
         state.loadingManager.showSimpleLoading('Uploading preview...');
         
         const formData = new FormData();
         
-        // Use appropriate parameter names and endpoint based on model type
         // Prepare common form data
         formData.append('preview_file', file);
         formData.append('model_path', filePath);
+        formData.append('nsfw_level', nsfwLevel.toString());  // Add nsfw_level parameter
 
         // Set endpoint based on model type
         const endpoint = modelType === 'checkpoint' 
@@ -587,7 +585,8 @@ async function uploadPreview(filePath, file, modelType = 'lora') {
         }
 
         const updateData = {
-            preview_url: data.preview_url
+            preview_url: data.preview_url,
+            preview_nsfw_level: data.preview_nsfw_level // Include nsfw level in update data
         };
 
         state.virtualScroller.updateSingleItem(filePath, updateData);
@@ -601,6 +600,8 @@ async function uploadPreview(filePath, file, modelType = 'lora') {
     }
 }
 
+// Private methods
+
 // Private function to perform the delete operation
 async function performDelete(filePath, modelType = 'lora') {
     try {

static/js/components/checkpointModal/index.js

@@ -108,7 +108,7 @@ export function showCheckpointModal(checkpoint) {
                     </div>
                 </div>
 
-                <div class="showcase-section" data-checkpoint-id="${checkpoint.civitai?.modelId || ''}">
+                <div class="showcase-section" data-model-hash="${checkpoint.sha256 || ''}">
                     <div class="showcase-tabs">
                         <button class="tab-btn active" data-tab="showcase">Examples</button>
                         <button class="tab-btn" data-tab="description">Model Description</button>

static/js/components/loraModal/index.js

@@ -134,7 +134,7 @@ export function showLoraModal(lora) {
                     </div>
                 </div>
 
-                <div class="showcase-section" data-lora-id="${lora.civitai?.modelId || ''}">
+                <div class="showcase-section" data-model-hash="${lora.sha256 || ''}" data-filepath="${lora.file_path}">
                     <div class="showcase-tabs">
                         <button class="tab-btn active" data-tab="showcase">Examples</button>
                         <button class="tab-btn" data-tab="description">Model Description</button>

static/js/components/shared/showcase/MediaRenderers.js

@@ -16,8 +16,10 @@
  * @returns {string} HTML content
  */
 export function generateVideoWrapper(media, heightPercent, shouldBlur, nsfwText, metadataPanel, localUrl, remoteUrl, mediaControlsHtml = '') {
+    const nsfwLevel = media.nsfwLevel !== undefined ? media.nsfwLevel : 0;
+    
     return `
-        <div class="media-wrapper ${shouldBlur ? 'nsfw-media-wrapper' : ''}" style="padding-bottom: ${heightPercent}%" data-short-id="${media.id || ''}">
+        <div class="media-wrapper ${shouldBlur ? 'nsfw-media-wrapper' : ''}" style="padding-bottom: ${heightPercent}%" data-short-id="${media.id || ''}" data-nsfw-level="${nsfwLevel}">
             ${shouldBlur ? `
                 <button class="toggle-blur-btn showcase-toggle-btn" title="Toggle blur">
                     <i class="fas fa-eye"></i>
@@ -28,6 +30,7 @@ export function generateVideoWrapper(media, heightPercent, shouldBlur, nsfwText,
                 referrerpolicy="no-referrer" 
                 data-local-src="${localUrl || ''}"
                 data-remote-src="${remoteUrl}"
+                data-nsfw-level="${nsfwLevel}"
                 class="lazy ${shouldBlur ? 'blurred' : ''}">
                 <source data-local-src="${localUrl || ''}" data-remote-src="${remoteUrl}" type="video/mp4">
                 Your browser does not support video playback
@@ -58,8 +61,10 @@ export function generateVideoWrapper(media, heightPercent, shouldBlur, nsfwText,
  * @returns {string} HTML content
  */
 export function generateImageWrapper(media, heightPercent, shouldBlur, nsfwText, metadataPanel, localUrl, remoteUrl, mediaControlsHtml = '') {
+    const nsfwLevel = media.nsfwLevel !== undefined ? media.nsfwLevel : 0;
+    
     return `
-        <div class="media-wrapper ${shouldBlur ? 'nsfw-media-wrapper' : ''}" style="padding-bottom: ${heightPercent}%" data-short-id="${media.id || ''}">
+        <div class="media-wrapper ${shouldBlur ? 'nsfw-media-wrapper' : ''}" style="padding-bottom: ${heightPercent}%" data-short-id="${media.id || ''}" data-nsfw-level="${nsfwLevel}">
             ${shouldBlur ? `
                 <button class="toggle-blur-btn showcase-toggle-btn" title="Toggle blur">
                     <i class="fas fa-eye"></i>
@@ -68,6 +73,7 @@ export function generateImageWrapper(media, heightPercent, shouldBlur, nsfwText,
             ${mediaControlsHtml}
             <img data-local-src="${localUrl || ''}" 
                 data-remote-src="${remoteUrl}"
+                data-nsfw-level="${nsfwLevel}"
                 alt="Preview" 
                 crossorigin="anonymous" 
                 referrerpolicy="no-referrer"

static/js/components/shared/showcase/MediaUtils.js

@@ -5,6 +5,7 @@
  */
 import { showToast, copyToClipboard } from '../../../utils/uiHelpers.js';
 import { state } from '../../../state/index.js';
+import { uploadPreview } from '../../../api/baseModelApi.js';
 
 /**
  * Try to load local image first, fall back to remote if local fails
@@ -186,9 +187,10 @@ export function initMetadataPanelHandlers(container) {
     mediaWrappers.forEach(wrapper => {
         // Get the metadata panel and media element (img or video)
         const metadataPanel = wrapper.querySelector('.image-metadata-panel');
+        const mediaControls = wrapper.querySelector('.media-controls');
         const mediaElement = wrapper.querySelector('img, video');
         
-        if (!metadataPanel || !mediaElement) return;
+        if (!mediaElement) return;
         
         let isOverMetadataPanel = false;
         
@@ -210,24 +212,29 @@ export function initMetadataPanelHandlers(container) {
                 mouseY <= mediaRect.bottom
             );
             
-            // Show metadata panel when over media content or metadata panel itself
+            // Show metadata panel and controls when over media content or metadata panel itself
             if (isOverMedia || isOverMetadataPanel) {
-                metadataPanel.classList.add('visible');
+                if (metadataPanel) metadataPanel.classList.add('visible');
+                if (mediaControls) mediaControls.classList.add('visible');
             } else {
-                metadataPanel.classList.remove('visible');
+                if (metadataPanel) metadataPanel.classList.remove('visible');
+                if (mediaControls) mediaControls.classList.remove('visible');
             }
         });
         
         wrapper.addEventListener('mouseleave', () => {
             if (!isOverMetadataPanel) {
-                metadataPanel.classList.remove('visible');
+                if (metadataPanel) metadataPanel.classList.remove('visible');
+                if (mediaControls) mediaControls.classList.remove('visible');
             }
         });
         
         // Add mouse enter/leave events for the metadata panel itself
+        if (metadataPanel) {
             metadataPanel.addEventListener('mouseenter', () => {
                 isOverMetadataPanel = true;
                 metadataPanel.classList.add('visible');
+                if (mediaControls) mediaControls.classList.add('visible');
             });
             
             metadataPanel.addEventListener('mouseleave', () => {
@@ -247,6 +254,7 @@ export function initMetadataPanelHandlers(container) {
                 
                 if (!isOverMedia) {
                     metadataPanel.classList.remove('visible');
+                    if (mediaControls) mediaControls.classList.remove('visible');
                 }
             });
             
@@ -285,6 +293,7 @@ export function initMetadataPanelHandlers(container) {
                     e.stopPropagation();
                 }
             }, { passive: true });
+        }
     });
 }
 
@@ -356,13 +365,19 @@ export function initMediaControlHandlers(container) {
         
         btn.addEventListener('click', async function(e) {
             e.stopPropagation();
+            
+            // Explicitly check for disabled state
+            if (this.classList.contains('disabled')) {
+                return; // Don't do anything if button is disabled
+            }
+            
             const shortId = this.dataset.shortId;
-            const state = this.dataset.state;
+            const btnState = this.dataset.state;
             
             if (!shortId) return;
             
             // Handle two-step confirmation
-            if (state === 'initial') {
+            if (btnState === 'initial') {
                 // First click: show confirmation state
                 this.dataset.state = 'confirm';
                 this.classList.add('confirm');
@@ -381,14 +396,15 @@ export function initMediaControlHandlers(container) {
             }
             
             // Second click within 3 seconds: proceed with deletion
-            if (state === 'confirm') {
+            if (btnState === 'confirm') {
                 this.disabled = true;
+                this.classList.remove('confirm');
                 this.innerHTML = '<i class="fas fa-spinner fa-spin"></i>';
                 
                 // Get model hash from URL or data attribute
                 const mediaWrapper = this.closest('.media-wrapper');
-                const modelIdAttr = document.querySelector('.showcase-section')?.dataset;
+                const modelHashAttr = document.querySelector('.showcase-section')?.dataset;
-                const modelHash = modelIdAttr?.loraId || modelIdAttr?.checkpointId;
+                const modelHash = modelHashAttr?.modelHash;
                 
                 try {
                     // Call the API to delete the custom example
@@ -418,19 +434,15 @@ export function initMediaControlHandlers(container) {
                         // Show success toast
                         showToast('Example image deleted', 'success');
 
-                        // Update VirtualScroller if available
-                        if (state.virtualScroller && result.model_file_path) {
                         // Create an update object with only the necessary properties
                         const updateData = {
                             civitai: {
-                                    images: result.regular_images || [],
                                 customImages: result.custom_images || []
                             }
                         };
                         
                         // Update the item in the virtual scroller
                         state.virtualScroller.updateSingleItem(result.model_file_path, updateData);
-                        }
                     } else {
                         // Show error message
                         showToast(result.error || 'Failed to delete example image', 'error');
@@ -457,31 +469,79 @@ export function initMediaControlHandlers(container) {
         });
     });
     
-    // Find all media controls
+    // Initialize set preview buttons
-    const mediaControls = container.querySelectorAll('.media-controls');
+    initSetPreviewHandlers(container);
     
-    // Set up same visibility behavior as metadata panel
+    // Media control visibility is now handled in initMetadataPanelHandlers
-    mediaControls.forEach(controlsEl => {
+    // Any click handlers or other functionality can still be added here
-        const mediaWrapper = controlsEl.closest('.media-wrapper');
+}
+
+/**
+ * Initialize set preview button handlers
+ * @param {HTMLElement} container - Container with media wrappers
+ */
+function initSetPreviewHandlers(container) {
+    const previewButtons = container.querySelectorAll('.set-preview-btn');
+    const modelType = state.currentPageType == 'loras' ? 'lora' : 'checkpoint';
+    
+    previewButtons.forEach(btn => {
+        btn.addEventListener('click', async function(e) {
+            e.stopPropagation();
+            
+            // Show loading state
+            this.innerHTML = '<i class="fas fa-spinner fa-spin"></i>';
+            this.disabled = true;
+            
+            try {
+                // Get the model file path from showcase section data attribute
+                const showcaseSection = document.querySelector('.showcase-section');
+                const modelHash = showcaseSection?.dataset.modelHash;
+                const modelFilePath = showcaseSection?.dataset.filepath;
+                
+                if (!modelFilePath) {
+                    throw new Error('Could not determine model file path');
+                }
+                
+                // Get the media wrapper and media element
+                const mediaWrapper = this.closest('.media-wrapper');
                 const mediaElement = mediaWrapper.querySelector('img, video');
                 
-        // Media controls should be visible when metadata panel is visible
+                if (!mediaElement) {
-        const metadataPanel = mediaWrapper.querySelector('.image-metadata-panel');
+                    throw new Error('Media element not found');
-        if (metadataPanel) {
-            const observer = new MutationObserver(mutations => {
-                mutations.forEach(mutation => {
-                    if (mutation.type === 'attributes' && mutation.attributeName === 'class') {
-                        if (metadataPanel.classList.contains('visible')) {
-                            controlsEl.classList.add('visible');
-                        } else if (!mediaWrapper.matches(':hover')) {
-                            controlsEl.classList.remove('visible');
                 }
-                    }
-                });
-            });
                 
-            observer.observe(metadataPanel, { attributes: true });
+                // Get NSFW level from the wrapper or media element
+                const nsfwLevel = parseInt(mediaWrapper.dataset.nsfwLevel || mediaElement.dataset.nsfwLevel || '0', 10);
+                
+                // Get local file path if available
+                const useLocalFile = mediaElement.dataset.localSrc && !mediaElement.dataset.localSrc.includes('undefined');
+                
+                if (useLocalFile) {
+                    // We have a local file, use it directly
+                    const response = await fetch(mediaElement.dataset.localSrc);
+                    const blob = await response.blob();
+                    const file = new File([blob], 'preview.jpg', { type: blob.type });
+                    
+                    // Use the existing baseModelApi uploadPreview method with nsfw level
+                    await uploadPreview(modelFilePath, file, modelType, nsfwLevel);
+                } else {
+                    // We need to download the remote file first
+                    const response = await fetch(mediaElement.src);
+                    const blob = await response.blob();
+                    const file = new File([blob], 'preview.jpg', { type: blob.type });
+                    
+                    // Use the existing baseModelApi uploadPreview method with nsfw level
+                    await uploadPreview(modelFilePath, file, modelType, nsfwLevel);
                 }
+            } catch (error) {
+                console.error('Error setting preview:', error);
+                showToast('Failed to set preview image', 'error');
+            } finally {
+                // Restore button state
+                this.innerHTML = '<i class="fas fa-image"></i>';
+                this.disabled = false;
+            }
+        });
     });
 }
 

static/js/components/shared/showcase/ShowcaseView.js

@@ -145,13 +145,15 @@ function renderMediaItem(img, index, exampleFiles) {
     const mediaControlsHtml = `
         <div class="media-controls">
             <button class="media-control-btn set-preview-btn" title="Set as preview">
-                <i class="fas fa-star"></i>
+                <i class="fas fa-image"></i>
             </button>
-            ${isCustomImage ? `
+            <button class="media-control-btn example-delete-btn ${!isCustomImage ? 'disabled' : ''}" 
-                <button class="media-control-btn example-delete-btn" title="Delete this example" data-short-id="${img.id}">
+                    title="${isCustomImage ? 'Delete this example' : 'Only custom images can be deleted'}" 
+                    data-short-id="${img.id || ''}" 
+                    ${!isCustomImage ? 'disabled' : ''}>
                 <i class="fas fa-trash-alt"></i>
+                <i class="fas fa-check confirm-icon"></i>
             </button>
-            ` : ''}
         </div>
     `;