alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-21 21:22:11 -03:00
Code Issues Packages Projects Releases Wiki Activity

Refactor example images handling by introducing migration logic, updating metadata structure, and enhancing image loading in the UI

Browse Source
This commit is contained in:
Will Miao
2025-06-18 17:14:49 +08:00
parent 022c6c157a
commit 3c047bee58
10 changed files with 412 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
py/lora_manager.py
View File

@@ -10,6 +10,7 @@ from .routes.misc_routes import MiscRoutes
from .routes.example_images_routes import ExampleImagesRoutes
from .services.service_registry import ServiceRegistry
from .services.settings_manager import settings
from .utils.example_images_migration import ExampleImagesMigration
import logging
import sys
import os
@@ -130,13 +131,13 @@ class LoraManager:
logging.getLogger('aiohttp.access').setLevel(logging.WARNING)
# Initialize CivitaiClient first to ensure it's ready for other services
civitai_client = await ServiceRegistry.get_civitai_client()
await ServiceRegistry.get_civitai_client()
# Register DownloadManager with ServiceRegistry
download_manager = await ServiceRegistry.get_download_manager()
await ServiceRegistry.get_download_manager()
# Initialize WebSocket manager
ws_manager = await ServiceRegistry.get_websocket_manager()
await ServiceRegistry.get_websocket_manager()
# Initialize scanners in background
lora_scanner = await ServiceRegistry.get_lora_scanner()
@@ -155,6 +156,8 @@ class LoraManager:
asyncio.create_task(lora_scanner.initialize_in_background(), name='lora_cache_init')
asyncio.create_task(checkpoint_scanner.initialize_in_background(), name='checkpoint_cache_init')
asyncio.create_task(recipe_scanner.initialize_in_background(), name='recipe_cache_init')
await ExampleImagesMigration.check_and_run_migrations()
logger.info("LoRA Manager: All services initialized and background tasks scheduled")

3
py/routes/lora_routes.py
View File

@@ -70,8 +70,7 @@ class LoraRoutes:
# It's initializing if the cache object doesn't exist yet,
# OR if the scanner explicitly says it's initializing (background task running).
is_initializing = (
self.scanner._cache is None or
(hasattr(self.scanner, '_is_initializing') and self.scanner._is_initializing)
self.scanner._cache is None or self.scanner.is_initializing()
)
if is_initializing:

4
py/services/model_scanner.py
View File

@@ -749,6 +749,10 @@ class ModelScanner:
"""Scan all model directories and return metadata"""
raise NotImplementedError("Subclasses must implement scan_all_models")
def is_initializing(self) -> bool:
"""Check if the scanner is currently initializing"""
return self._is_initializing
def get_model_roots(self) -> List[str]:
"""Get model root directories"""
raise NotImplementedError("Subclasses must implement get_model_roots")

70
py/utils/example_images_file_manager.py
View File

@@ -128,76 +128,6 @@ class ExampleImagesFileManager:
'is_video': file_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']
})
# Check if files use 1-based indexing (look for patterns like "image_1.jpg")
has_one_based = any(re.match(r'image_1\.\w+$', f['name']) for f in files)
has_zero_based = any(re.match(r'image_0\.\w+$', f['name']) for f in files)
# If there are 1-based indices and no 0-based indices, rename files
if has_one_based and not has_zero_based:
logger.info(f"Converting 1-based to 0-based indexing in {model_folder}")
# Sort files to ensure correct order
files.sort(key=lambda x: x['name'])
# First, create rename mapping to avoid conflicts
renames = []
for file in files:
match = re.match(r'image_(\d+)\.(\w+)$', file['name'])
if match:
index = int(match.group(1))
ext = match.group(2)
if index > 0: # Only rename if index is positive
new_name = f"image_{index-1}.{ext}"
renames.append((file['name'], new_name))
# Use temporary filenames to avoid conflicts
for old_name, new_name in renames:
old_path = os.path.join(model_folder, old_name)
temp_path = os.path.join(model_folder, f"temp_{old_name}")
try:
os.rename(old_path, temp_path)
except Exception as e:
logger.error(f"Failed to rename {old_path} to {temp_path}: {e}")
# Rename from temporary names to final names
for old_name, new_name in renames:
temp_path = os.path.join(model_folder, f"temp_{old_name}")
new_path = os.path.join(model_folder, new_name)
try:
os.rename(temp_path, new_path)
logger.debug(f"Renamed {old_name} to {new_name}")
# Update file list entry
for file in files:
if file['name'] == old_name:
file['name'] = new_name
file['path'] = f'/example_images_static/{model_hash}/{new_name}'
except Exception as e:
logger.error(f"Failed to rename {temp_path} to {new_path}: {e}")
# Refresh file list after renaming
files = []
for file in os.listdir(model_folder):
file_path = os.path.join(model_folder, file)
if os.path.isfile(file_path):
file_ext = os.path.splitext(file)[1].lower()
if (file_ext in SUPPORTED_MEDIA_EXTENSIONS['images'] or
file_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']):
files.append({
'name': file,
'path': f'/example_images_static/{model_hash}/{file}',
'extension': file_ext,
'is_video': file_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']
})
# Sort files by index for consistent order
def extract_index(filename):
match = re.match(r'image_(\d+)\.\w+$', filename)
if match:
return int(match.group(1))
return float('inf') # Place non-matching files at the end
files.sort(key=lambda x: extract_index(x['name']))
return web.json_response({
'success': True,
'files': files

31
py/utils/example_images_metadata.py
View File

@@ -198,29 +198,32 @@ class MetadataUpdater:
newly_imported_paths: List of paths to newly imported files
Returns:
list: Updated images array
tuple: (regular_images, custom_images) - Both image arrays
"""
try:
# Ensure civitai field exists in model_data
if not model_data.get('civitai'):
model_data['civitai'] = {}
# Ensure images array exists
if not model_data['civitai'].get('images'):
model_data['civitai']['images'] = []
# Ensure customImages array exists
if not model_data['civitai'].get('customImages'):
model_data['civitai']['customImages'] = []
# Get current images array
images = model_data['civitai']['images']
# Get current customImages array
custom_images = model_data['civitai']['customImages']
# Add new image entry for each imported file
for path in newly_imported_paths:
for path_tuple in newly_imported_paths:
path, short_id = path_tuple
# Determine if video or image
file_ext = os.path.splitext(path)[1].lower()
is_video = file_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']
# Create image metadata entry
image_entry = {
"url": "", # Empty URL as required
"url": "", # Empty URL as requested
"id": short_id,
"nsfwLevel": 0,
"width": 720, # Default dimensions
"height": 1280,
@@ -240,8 +243,8 @@ class MetadataUpdater:
# If PIL fails or is unavailable, use default dimensions
pass
# Append to existing images array
images.append(image_entry)
# Append to existing customImages array
custom_images.append(image_entry)
# Save metadata to .metadata.json file
file_path = model_data.get('file_path')
@@ -261,8 +264,12 @@ class MetadataUpdater:
if file_path:
await scanner.update_single_model_cache(file_path, file_path, model_data)
return images
# Get regular images array (might be None)
regular_images = model_data['civitai'].get('images', [])
# Return both image arrays
return regular_images, custom_images
except Exception as e:
logger.error(f"Failed to update metadata after import: {e}", exc_info=True)
return []
return [], []

318
py/utils/example_images_migration.py Normal file
View File

@@ -0,0 +1,318 @@
import asyncio
import logging
import os
import re
import json
from ..services.settings_manager import settings
from ..services.service_registry import ServiceRegistry
from ..utils.metadata_manager import MetadataManager
from ..utils.example_images_processor import ExampleImagesProcessor
from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS
logger = logging.getLogger(__name__)
CURRENT_NAMING_VERSION = 2 # Increment this when naming conventions change
class ExampleImagesMigration:
"""Handles migrations for example images naming conventions"""
@staticmethod
async def check_and_run_migrations():
"""Check if migrations are needed and run them in background"""
example_images_path = settings.get('example_images_path')
if not example_images_path or not os.path.exists(example_images_path):
logger.debug("No example images path configured or path doesn't exist, skipping migrations")
return
# Check current version from progress file
current_version = 0
progress_file = os.path.join(example_images_path, '.download_progress.json')
if os.path.exists(progress_file):
try:
with open(progress_file, 'r', encoding='utf-8') as f:
progress_data = json.load(f)
current_version = progress_data.get('naming_version', 0)
except Exception as e:
logger.error(f"Failed to load progress file for migration check: {e}")
# If current version is less than target version, start migration
if current_version < CURRENT_NAMING_VERSION:
logger.info(f"Starting example images naming migration from v{current_version} to v{CURRENT_NAMING_VERSION}")
# Start migration in background task
asyncio.create_task(
ExampleImagesMigration.run_migrations(example_images_path, current_version, CURRENT_NAMING_VERSION)
)
@staticmethod
async def run_migrations(example_images_path, from_version, to_version):
"""Run necessary migrations based on version difference"""
try:
# Get all model folders
model_folders = []
for item in os.listdir(example_images_path):
item_path = os.path.join(example_images_path, item)
if os.path.isdir(item_path) and len(item) == 64: # SHA256 hash is 64 chars
model_folders.append(item_path)
logger.info(f"Found {len(model_folders)} model folders to check for migration")
# Apply migrations sequentially
if from_version < 1 and to_version >= 1:
await ExampleImagesMigration._migrate_to_v1(model_folders)
if from_version < 2 and to_version >= 2:
await ExampleImagesMigration._migrate_to_v2(model_folders)
# Update version in progress file
progress_file = os.path.join(example_images_path, '.download_progress.json')
try:
progress_data = {}
if os.path.exists(progress_file):
with open(progress_file, 'r', encoding='utf-8') as f:
progress_data = json.load(f)
progress_data['naming_version'] = to_version
with open(progress_file, 'w', encoding='utf-8') as f:
json.dump(progress_data, f, indent=2)
logger.info(f"Example images naming migration to v{to_version} completed")
except Exception as e:
logger.error(f"Failed to update version in progress file: {e}")
except Exception as e:
logger.error(f"Error during migration: {e}", exc_info=True)
@staticmethod
async def _migrate_to_v1(model_folders):
"""Migrate from 1-based to 0-based indexing"""
count = 0
for folder in model_folders:
has_one_based = False
has_zero_based = False
files_to_rename = []
# Check naming pattern in this folder
for file in os.listdir(folder):
if re.match(r'image_1\.\w+$', file):
has_one_based = True
if re.match(r'image_0\.\w+$', file):
has_zero_based = True
# Only migrate folders with 1-based indexing and no 0-based
if has_one_based and not has_zero_based:
# Create rename mapping
for file in os.listdir(folder):
match = re.match(r'image_(\d+)\.(\w+)$', file)
if match:
index = int(match.group(1))
ext = match.group(2)
if index > 0: # Only rename if index is positive
files_to_rename.append((
file,
f"image_{index-1}.{ext}"
))
# Use temporary names to avoid conflicts
for old_name, new_name in files_to_rename:
old_path = os.path.join(folder, old_name)
temp_path = os.path.join(folder, f"temp_{old_name}")
try:
os.rename(old_path, temp_path)
except Exception as e:
logger.error(f"Failed to rename {old_path} to {temp_path}: {e}")
# Rename from temporary names to final names
for old_name, new_name in files_to_rename:
temp_path = os.path.join(folder, f"temp_{old_name}")
new_path = os.path.join(folder, new_name)
try:
os.rename(temp_path, new_path)
logger.debug(f"Renamed {old_name} to {new_name} in {folder}")
except Exception as e:
logger.error(f"Failed to rename {temp_path} to {new_path}: {e}")
count += 1
# Give other tasks a chance to run
if count % 10 == 0:
await asyncio.sleep(0)
logger.info(f"Migrated {count} folders from 1-based to 0-based indexing")
@staticmethod
async def _migrate_to_v2(model_folders):
"""
Migrate to v2 naming scheme:
- Move custom examples from images array to customImages array
- Rename files from image_<index>.<ext> to custom_<short_id>.<ext>
- Add id field to each custom image entry
"""
count = 0
updated_models = 0
migration_errors = 0
# Get scanner instances
lora_scanner = await ServiceRegistry.get_lora_scanner()
checkpoint_scanner = await ServiceRegistry.get_checkpoint_scanner()
# Wait until scanners are initialized
scanners = [lora_scanner, checkpoint_scanner]
for scanner in scanners:
if scanner.is_initializing():
logger.info("Waiting for scanners to complete initialization before starting migration...")
initialized = False
retry_count = 0
while not initialized and retry_count < 120: # Wait up to 120 seconds
await asyncio.sleep(1)
initialized = not scanner.is_initializing()
retry_count += 1
if not initialized:
logger.warning("Scanner initialization timeout - proceeding with migration anyway")
logger.info(f"Starting migration to v2 naming scheme for {len(model_folders)} model folders")
for folder in model_folders:
try:
# Extract model hash from folder name
model_hash = os.path.basename(folder)
if not model_hash or len(model_hash) != 64:
continue
# Find the model in scanner cache
model_data = None
scanner = None
for scan_obj in scanners:
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 or not scanner:
logger.debug(f"Model with hash {model_hash} not found in cache, skipping migration")
continue
# Clone model data to avoid modifying the cache directly
model_metadata = model_data.copy()
# Check if model has civitai metadata
if not model_metadata.get('civitai'):
continue
# Get images array
images = model_metadata.get('civitai', {}).get('images', [])
if not images:
continue
# Initialize customImages array if it doesn't exist
if not model_metadata['civitai'].get('customImages'):
model_metadata['civitai']['customImages'] = []
# Find custom examples (entries with empty url)
custom_indices = []
for i, image in enumerate(images):
if image.get('url') == "":
custom_indices.append(i)
if not custom_indices:
continue
logger.debug(f"Found {len(custom_indices)} custom examples in {model_hash}")
# Process each custom example
for index in custom_indices:
try:
image_entry = images[index]
# Determine media type based on the entry type
media_type = 'videos' if image_entry.get('type') == 'video' else 'images'
extensions_to_try = SUPPORTED_MEDIA_EXTENSIONS[media_type]
# Find the image file by trying possible extensions
old_path = None
old_filename = None
found = False
for ext in extensions_to_try:
test_path = os.path.join(folder, f"image_{index}{ext}")
if os.path.exists(test_path):
old_path = test_path
old_filename = f"image_{index}{ext}"
found = True
break
if not found:
logger.warning(f"Could not find file for index {index} in {model_hash}, skipping")
continue
# Generate short ID for the custom example
short_id = ExampleImagesProcessor.generate_short_id()
# Get file extension
file_ext = os.path.splitext(old_path)[1]
# Create new filename
new_filename = f"custom_{short_id}{file_ext}"
new_path = os.path.join(folder, new_filename)
# Rename the file
try:
os.rename(old_path, new_path)
logger.debug(f"Renamed {old_filename} to {new_filename} in {folder}")
except Exception as e:
logger.error(f"Failed to rename {old_path} to {new_path}: {e}")
continue
# Create a copy of the image entry with the id field
custom_entry = image_entry.copy()
custom_entry['id'] = short_id
# Add to customImages array
model_metadata['civitai']['customImages'].append(custom_entry)
count += 1
except Exception as e:
logger.error(f"Error migrating custom example at index {index} for {model_hash}: {e}")
# Remove custom examples from the original images array
model_metadata['civitai']['images'] = [
img for i, img in enumerate(images) if i not in custom_indices
]
# Save the updated metadata
file_path = model_data.get('file_path')
if file_path:
try:
# Create a copy of model data without 'folder' field
model_copy = model_metadata.copy()
model_copy.pop('folder', None)
# Save metadata to file
await MetadataManager.save_metadata(file_path, model_copy)
# Update scanner cache
await scanner.update_single_model_cache(file_path, file_path, model_metadata)
updated_models += 1
except Exception as e:
logger.error(f"Failed to save metadata for {model_hash}: {e}")
migration_errors += 1
# Give other tasks a chance to run
if count % 10 == 0:
await asyncio.sleep(0)
except Exception as e:
logger.error(f"Error migrating folder {folder}: {e}")
migration_errors += 1
logger.info(f"Migration to v2 complete: migrated {count} custom examples across {updated_models} models with {migration_errors} errors")

28
py/utils/example_images_processor.py
View File

@@ -2,14 +2,21 @@ import logging
import os
import re
import tempfile
import random
import string
from aiohttp import web
import asyncio
from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS
logger = logging.getLogger(__name__)
class ExampleImagesProcessor:
"""Processes and manipulates example images"""
@staticmethod
def generate_short_id(length=8):
"""Generate a short random alphanumeric identifier"""
chars = string.ascii_lowercase + string.digits
return ''.join(random.choice(chars) for _ in range(length))
@staticmethod
def get_civitai_optimized_url(image_url):
@@ -265,11 +272,6 @@ class ExampleImagesProcessor:
'error': f"Model with hash {model_hash} not found in cache"
}, status=404)
# Get the current number of images in the civitai.images array
civitai_data = model_data.get('civitai')
current_images = civitai_data.get('images', []) if civitai_data is not None else []
next_index = len(current_images)
# Create model folder
model_folder = os.path.join(example_images_path, model_hash)
os.makedirs(model_folder, exist_ok=True)
@@ -293,16 +295,17 @@ class ExampleImagesProcessor:
errors.append(f"Unsupported file type: {file_path}")
continue
# Generate new filename using sequential index starting from current image length
new_filename = f"image_{next_index}{file_ext}"
next_index += 1
# Generate new filename using short ID instead of UUID
short_id = ExampleImagesProcessor.generate_short_id()
new_filename = f"custom_{short_id}{file_ext}"
dest_path = os.path.join(model_folder, new_filename)
# Copy the file
import shutil
shutil.copy2(file_path, dest_path)
newly_imported_paths.append(dest_path)
# Store both the dest_path and the short_id
newly_imported_paths.append((dest_path, short_id))
# Add to imported files list
imported_files.append({
@@ -315,7 +318,7 @@ class ExampleImagesProcessor:
errors.append(f"Error importing {file_path}: {str(e)}")
# Update metadata with new example images
updated_images = await MetadataUpdater.update_metadata_after_import(
regular_images, custom_images = await MetadataUpdater.update_metadata_after_import(
model_hash,
model_data,
scanner,
@@ -328,7 +331,8 @@ class ExampleImagesProcessor:
(f' with {len(errors)} errors' if errors else ''),
'files': imported_files,
'errors': errors,
'updated_images': updated_images,
'regular_images': regular_images,
'custom_images': custom_images,
"model_file_path": model_data.get('file_path', ''),
})

20
py/utils/routes_common.py
View File

@@ -39,7 +39,23 @@ class ModelRouteUtils:
async def update_model_metadata(metadata_path: str, local_metadata: Dict,
civitai_metadata: Dict, client: CivitaiClient) -> None:
"""Update local metadata with CivitAI data"""
local_metadata['civitai'] = civitai_metadata
# Save existing trainedWords and customImages if they exist
existing_civitai = local_metadata.get('civitai', {})
existing_trained_words = existing_civitai.get('trainedWords', [])
# Create a new civitai metadata by updating existing with new
merged_civitai = existing_civitai.copy()
merged_civitai.update(civitai_metadata)
# Special handling for trainedWords - ensure we don't lose any existing trained words
new_trained_words = civitai_metadata.get('trainedWords', [])
if existing_trained_words:
# Use a set to combine words without duplicates, then convert back to list
merged_trained_words = list(set(existing_trained_words + new_trained_words))
merged_civitai['trainedWords'] = merged_trained_words
# Update local metadata with merged civitai data
local_metadata['civitai'] = merged_civitai
local_metadata['from_civitai'] = True
# Update model name if available
@@ -219,7 +235,7 @@ class ModelRouteUtils:
fields = [
"id", "modelId", "name", "createdAt", "updatedAt",
"publishedAt", "trainedWords", "baseModel", "description",
"model", "images", "creator"
"model", "images", "customImages", "creator"
]
return {k: data[k] for k in fields if k in data}

34
static/js/components/loraModal/ShowcaseView.js
View File

@@ -65,15 +65,21 @@ export function renderShowcaseContent(images, exampleFiles = []) {
// Find matching file in our list of actual files
let localFile = null;
if (exampleFiles.length > 0) {
// Try to find the corresponding file by index first
localFile = exampleFiles.find(file => {
const match = file.name.match(/image_(\d+)\./);
return match && parseInt(match[1]) === index;
});
// If not found by index, just use the same position in the array if available
if (!localFile && index < exampleFiles.length) {
localFile = exampleFiles[index];
if (img.id) {
// This is a custom image, find by custom_<id>
const customPrefix = `custom_${img.id}`;
localFile = exampleFiles.find(file => file.name.startsWith(customPrefix));
} else {
// This is a regular image from civitai, find by index
localFile = exampleFiles.find(file => {
const match = file.name.match(/image_(\d+)\./);
return match && parseInt(match[1]) === index;
});
// If not found by index, just use the same position in the array if available
if (!localFile && index < exampleFiles.length) {
localFile = exampleFiles[index];
}
}
}
@@ -301,8 +307,11 @@ async function handleImportFiles(files, modelHash, importContainer) {
const showcaseTab = document.getElementById('showcase-tab');
if (showcaseTab) {
// Get the updated images from the result
const updatedImages = result.updated_images || [];
showcaseTab.innerHTML = renderShowcaseContent(updatedImages, updatedFilesResult.files);
const regularImages = result.regular_images || [];
const customImages = result.custom_images || [];
// Combine both arrays for rendering
const allImages = [...regularImages, ...customImages];
showcaseTab.innerHTML = renderShowcaseContent(allImages, updatedFilesResult.files);
// Re-initialize showcase functionality
const carousel = showcaseTab.querySelector('.carousel');
@@ -321,7 +330,8 @@ async function handleImportFiles(files, modelHash, importContainer) {
// Create an update object with only the necessary properties
const updateData = {
civitai: {
images: updatedImages
images: regularImages,
customImages: customImages
}
};

13
static/js/components/loraModal/index.js
View File

@@ -192,17 +192,20 @@ export function showLoraModal(lora) {
// Load recipes for this Lora
loadRecipesForLora(lora.model_name, lora.sha256);
// Load example images asynchronously
loadExampleImages(lora.civitai?.images, lora.sha256, lora.file_path);
// Load example images asynchronously - merge regular and custom images
const regularImages = lora.civitai?.images || [];
const customImages = lora.civitai?.customImages || [];
// Combine images - regular images first, then custom images
const allImages = [...regularImages, ...customImages];
loadExampleImages(allImages, lora.sha256);
}
/**
* Load example images asynchronously
* @param {Array} images - Array of image objects
* @param {Array} images - Array of image objects (both regular and custom)
* @param {string} modelHash - Model hash for fetching local files
* @param {string} filePath - File path for fetching local files
*/
async function loadExampleImages(images, modelHash, filePath) {
async function loadExampleImages(images, modelHash) {
try {
const showcaseTab = document.getElementById('showcase-tab');
if (!showcaseTab) return;