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

This commit is contained in:
Will Miao
2025-06-19 15:19:24 +08:00
parent a7304ccf47
commit 605a06317b
10 changed files with 238 additions and 136 deletions

View File

@@ -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

View File

@@ -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()

View File

@@ -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:

View File

@@ -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 {

View File

@@ -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 {

View File

@@ -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>

View File

@@ -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>

View File

@@ -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"

View File

@@ -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,81 +212,88 @@ 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
metadataPanel.addEventListener('mouseenter', () => {
isOverMetadataPanel = true;
metadataPanel.classList.add('visible');
});
metadataPanel.addEventListener('mouseleave', () => {
isOverMetadataPanel = false;
// Only hide if mouse is not over the media
const rect = wrapper.getBoundingClientRect();
const mediaRect = getRenderedMediaRect(mediaElement, rect.width, rect.height);
const mouseX = event.clientX - rect.left;
const mouseY = event.clientY - rect.top;
if (metadataPanel) {
metadataPanel.addEventListener('mouseenter', () => {
isOverMetadataPanel = true;
metadataPanel.classList.add('visible');
if (mediaControls) mediaControls.classList.add('visible');
});
const isOverMedia = (
mouseX >= mediaRect.left &&
mouseX <= mediaRect.right &&
mouseY >= mediaRect.top &&
mouseY <= mediaRect.bottom
);
if (!isOverMedia) {
metadataPanel.classList.remove('visible');
}
});
// Prevent events from bubbling
metadataPanel.addEventListener('click', (e) => {
e.stopPropagation();
});
// Handle copy prompt buttons
const copyBtns = metadataPanel.querySelectorAll('.copy-prompt-btn');
copyBtns.forEach(copyBtn => {
const promptIndex = copyBtn.dataset.promptIndex;
const promptElement = wrapper.querySelector(`#prompt-${promptIndex}`);
copyBtn.addEventListener('click', async (e) => {
e.stopPropagation();
metadataPanel.addEventListener('mouseleave', () => {
isOverMetadataPanel = false;
// Only hide if mouse is not over the media
const rect = wrapper.getBoundingClientRect();
const mediaRect = getRenderedMediaRect(mediaElement, rect.width, rect.height);
const mouseX = event.clientX - rect.left;
const mouseY = event.clientY - rect.top;
if (!promptElement) return;
const isOverMedia = (
mouseX >= mediaRect.left &&
mouseX <= mediaRect.right &&
mouseY >= mediaRect.top &&
mouseY <= mediaRect.bottom
);
try {
await copyToClipboard(promptElement.textContent, 'Prompt copied to clipboard');
} catch (err) {
console.error('Copy failed:', err);
showToast('Copy failed', 'error');
if (!isOverMedia) {
metadataPanel.classList.remove('visible');
if (mediaControls) mediaControls.classList.remove('visible');
}
});
});
// Prevent panel scroll from causing modal scroll
metadataPanel.addEventListener('wheel', (e) => {
const isAtTop = metadataPanel.scrollTop === 0;
const isAtBottom = metadataPanel.scrollHeight - metadataPanel.scrollTop === metadataPanel.clientHeight;
// Only prevent default if scrolling would cause the panel to scroll
if ((e.deltaY < 0 && !isAtTop) || (e.deltaY > 0 && !isAtBottom)) {
// Prevent events from bubbling
metadataPanel.addEventListener('click', (e) => {
e.stopPropagation();
}
}, { passive: true });
});
// Handle copy prompt buttons
const copyBtns = metadataPanel.querySelectorAll('.copy-prompt-btn');
copyBtns.forEach(copyBtn => {
const promptIndex = copyBtn.dataset.promptIndex;
const promptElement = wrapper.querySelector(`#prompt-${promptIndex}`);
copyBtn.addEventListener('click', async (e) => {
e.stopPropagation();
if (!promptElement) return;
try {
await copyToClipboard(promptElement.textContent, 'Prompt copied to clipboard');
} catch (err) {
console.error('Copy failed:', err);
showToast('Copy failed', 'error');
}
});
});
// Prevent panel scroll from causing modal scroll
metadataPanel.addEventListener('wheel', (e) => {
const isAtTop = metadataPanel.scrollTop === 0;
const isAtBottom = metadataPanel.scrollHeight - metadataPanel.scrollTop === metadataPanel.clientHeight;
// Only prevent default if scrolling would cause the panel to scroll
if ((e.deltaY < 0 && !isAtTop) || (e.deltaY > 0 && !isAtBottom)) {
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 modelHash = modelIdAttr?.loraId || modelIdAttr?.checkpointId;
const modelHashAttr = document.querySelector('.showcase-section')?.dataset;
const modelHash = modelHashAttr?.modelHash;
try {
// Call the API to delete the custom example
@@ -417,20 +433,16 @@ export function initMediaControlHandlers(container) {
// Show success toast
showToast('Example image deleted', 'success');
// Create an update object with only the necessary properties
const updateData = {
civitai: {
customImages: result.custom_images || []
}
};
// 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);
}
// 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
const mediaControls = container.querySelectorAll('.media-controls');
// Initialize set preview buttons
initSetPreviewHandlers(container);
// Set up same visibility behavior as metadata panel
mediaControls.forEach(controlsEl => {
const mediaWrapper = controlsEl.closest('.media-wrapper');
const mediaElement = mediaWrapper.querySelector('img, video');
// Media controls should be visible when metadata panel is visible
const metadataPanel = mediaWrapper.querySelector('.image-metadata-panel');
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');
}
}
});
});
// Media control visibility is now handled in initMetadataPanelHandlers
// Any click handlers or other functionality can still be added here
}
/**
* 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();
observer.observe(metadataPanel, { attributes: true });
}
// 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');
if (!mediaElement) {
throw new Error('Media element not found');
}
// 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;
}
});
});
}

View File

@@ -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>
<button class="media-control-btn example-delete-btn ${!isCustomImage ? 'disabled' : ''}"
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>
${isCustomImage ? `
<button class="media-control-btn example-delete-btn" title="Delete this example" data-short-id="${img.id}">
<i class="fas fa-trash-alt"></i>
</button>
` : ''}
</div>
`;