Implement saving model example images locally. Fixes https://github.com/willmiao/ComfyUI-Lora-Manager/issues/88

py/lora_manager.py

@@ -6,10 +6,12 @@ from .routes.api_routes import ApiRoutes
 from .routes.recipe_routes import RecipeRoutes
 from .routes.checkpoints_routes import CheckpointsRoutes
 from .routes.update_routes import UpdateRoutes
-from .routes.usage_stats_routes import UsageStatsRoutes
+from .routes.misc_routes import MiscRoutes
 from .services.service_registry import ServiceRegistry
+from .services.settings_manager import settings
 import logging
 import sys
+import os
 
 logger = logging.getLogger(__name__)
 
@@ -29,6 +31,13 @@ class LoraManager:
 
         added_targets = set()  # Track already added target paths
         
+        # Add static route for example images if the path exists in settings
+        example_images_path = settings.get('example_images_path')
+        logger.info(f"Example images path: {example_images_path}")
+        if example_images_path and os.path.exists(example_images_path):
+            app.router.add_static('/example_images_static', example_images_path)
+            logger.info(f"Added static route for example images: /example_images_static -> {example_images_path}")
+        
         # Add static routes for each lora root
         for idx, root in enumerate(config.loras_roots, start=1):
             preview_path = f'/loras_static/root{idx}/preview'
@@ -102,7 +111,7 @@ class LoraManager:
         ApiRoutes.setup_routes(app)
         RecipeRoutes.setup_routes(app)
         UpdateRoutes.setup_routes(app)  
-        UsageStatsRoutes.setup_routes(app)  # Register usage stats routes
+        MiscRoutes.setup_routes(app)  # Register miscellaneous routes
         
         # Schedule service initialization 
         app.on_startup.append(lambda app: cls._initialize_services())

py/routes/api_routes.py

@@ -55,7 +55,6 @@ class ApiRoutes:
         app.router.add_get('/api/civitai/model/version/{modelVersionId}', routes.get_civitai_model_by_version)
         app.router.add_get('/api/civitai/model/hash/{hash}', routes.get_civitai_model_by_hash)
         app.router.add_post('/api/download-lora', routes.download_lora)
-        app.router.add_post('/api/settings', routes.update_settings)
         app.router.add_post('/api/move_model', routes.move_model)
         app.router.add_get('/api/lora-model-description', routes.get_lora_model_description)  # Add new route
         app.router.add_post('/api/loras/save-metadata', routes.save_metadata)
@@ -515,21 +514,6 @@ class ApiRoutes:
                 logger.error(f"Error downloading LoRA: {error_message}")
                 return web.Response(status=500, text=error_message)
 
-    async def update_settings(self, request: web.Request) -> web.Response:
-        """Update application settings"""
-        try:
-            data = await request.json()
-            
-            # Validate and update settings
-            if 'civitai_api_key' in data:
-                settings.set('civitai_api_key', data['civitai_api_key'])
-            if 'show_only_sfw' in data:
-                settings.set('show_only_sfw', data['show_only_sfw'])
-            
-            return web.json_response({'success': True})
-        except Exception as e:
-            logger.error(f"Error updating settings: {e}", exc_info=True)
-            return web.Response(status=500, text=str(e))
 
     async def move_model(self, request: web.Request) -> web.Response:
         """Handle model move request"""
@@ -1060,4 +1044,4 @@ class ApiRoutes:
             return web.json_response({
                 "success": False,
                 "error": str(e)
-            }, status=500)
+            }, status=500)

py/routes/misc_routes.py

@@ -0,0 +1,548 @@
+import logging
+import os
+import asyncio
+import json
+import time
+import tkinter as tk
+from tkinter import filedialog
+import aiohttp
+from aiohttp import web
+from ..services.settings_manager import settings
+from ..utils.usage_stats import UsageStats
+from ..services.service_registry import ServiceRegistry
+from ..utils.exif_utils import ExifUtils
+from ..utils.constants import EXAMPLE_IMAGE_WIDTH, SUPPORTED_MEDIA_EXTENSIONS
+
+logger = logging.getLogger(__name__)
+
+# Download status tracking
+download_task = None
+is_downloading = False
+download_progress = {
+    'total': 0,
+    'completed': 0,
+    'current_model': '',
+    'status': 'idle',  # idle, running, paused, completed, error
+    'errors': [],
+    'last_error': None,
+    'start_time': None,
+    'end_time': None,
+    'processed_models': set()  # Track models that have been processed
+}
+
+class MiscRoutes:
+    """Miscellaneous routes for various utility functions"""
+    
+    @staticmethod
+    def setup_routes(app):
+        """Register miscellaneous routes"""
+        app.router.add_post('/api/settings', MiscRoutes.update_settings)
+
+        # Usage stats routes
+        app.router.add_post('/api/update-usage-stats', MiscRoutes.update_usage_stats)
+        app.router.add_get('/api/get-usage-stats', MiscRoutes.get_usage_stats)
+        
+        # Example images download routes
+        app.router.add_post('/api/download-example-images', MiscRoutes.download_example_images)
+        app.router.add_get('/api/example-images-status', MiscRoutes.get_example_images_status)
+        app.router.add_post('/api/pause-example-images', MiscRoutes.pause_example_images)
+        app.router.add_post('/api/resume-example-images', MiscRoutes.resume_example_images)
+        
+        # Folder browser route
+        app.router.add_post('/api/browse-folder', MiscRoutes.browse_folder)
+
+    @staticmethod
+    async def update_settings(request):
+        """Update application settings"""
+        try:
+            data = await request.json()
+            
+            # Validate and update settings
+            for key, value in data.items():
+                # Special handling for example_images_path - verify path exists
+                if key == 'example_images_path' and value:
+                    if not os.path.exists(value):
+                        return web.json_response({
+                            'success': False,
+                            'error': f"Path does not exist: {value}"
+                        })
+                    
+                    # Add static route for example images path if it changed
+                    old_path = settings.get('example_images_path')
+                    if old_path != value:
+                        # We need to add the new static route
+                        # Note: we can't remove old routes in aiohttp, but we can create a new one
+                        app = request.app
+                        try:
+                            app.router.add_static('/example_images_static', value)
+                            logger.info(f"Added static route for example images: /example_images_static -> {value}")
+                        except Exception as e:
+                            logger.warning(f"Could not add static route for example images: {str(e)}")
+                
+                # Save to settings
+                settings.set(key, value)
+            
+            return web.json_response({'success': True})
+        except Exception as e:
+            logger.error(f"Error updating settings: {e}", exc_info=True)
+            return web.Response(status=500, text=str(e))
+    
+    @staticmethod
+    async def update_usage_stats(request):
+        """
+        Update usage statistics based on a prompt_id
+        
+        Expects a JSON body with:
+        {
+            "prompt_id": "string"
+        }
+        """
+        try:
+            # Parse the request body
+            data = await request.json()
+            prompt_id = data.get('prompt_id')
+            
+            if not prompt_id:
+                return web.json_response({
+                    'success': False,
+                    'error': 'Missing prompt_id'
+                }, status=400)
+            
+            # Call the UsageStats to process this prompt_id synchronously
+            usage_stats = UsageStats()
+            await usage_stats.process_execution(prompt_id)
+            
+            return web.json_response({
+                'success': True
+            })
+            
+        except Exception as e:
+            logger.error(f"Failed to update usage stats: {e}", exc_info=True)
+            return web.json_response({
+                'success': False,
+                'error': str(e)
+            }, status=500)
+    
+    @staticmethod
+    async def get_usage_stats(request):
+        """Get current usage statistics"""
+        try:
+            usage_stats = UsageStats()
+            stats = await usage_stats.get_stats()
+            
+            return web.json_response({
+                'success': True,
+                'data': stats
+            })
+            
+        except Exception as e:
+            logger.error(f"Failed to get usage stats: {e}", exc_info=True)
+            return web.json_response({
+                'success': False,
+                'error': str(e)
+            }, status=500)
+    
+    @staticmethod
+    async def download_example_images(request):
+        """
+        Download example images for models from Civitai
+        
+        Expects a JSON body with:
+        {
+            "output_dir": "path/to/output",  # Base directory to save example images
+            "optimize": true,                # Whether to optimize images (default: true)
+            "model_types": ["lora", "checkpoint"], # Model types to process (default: both)
+            "delay": 1.0                     # Delay between downloads to avoid rate limiting (default: 1.0)
+        }
+        """
+        global download_task, is_downloading, download_progress
+        
+        if is_downloading:
+            # Create a copy for JSON serialization
+            response_progress = download_progress.copy()
+            response_progress['processed_models'] = list(download_progress['processed_models'])
+            
+            return web.json_response({
+                'success': False,
+                'error': 'Download already in progress',
+                'status': response_progress
+            }, status=400)
+        
+        try:
+            # Parse the request body
+            data = await request.json()
+            output_dir = data.get('output_dir')
+            optimize = data.get('optimize', True)
+            model_types = data.get('model_types', ['lora', 'checkpoint'])
+            delay = float(data.get('delay', 1.0))
+            
+            if not output_dir:
+                return web.json_response({
+                    'success': False,
+                    'error': 'Missing output_dir parameter'
+                }, status=400)
+            
+            # Create the output directory
+            os.makedirs(output_dir, exist_ok=True)
+            
+            # Initialize progress tracking
+            download_progress['total'] = 0
+            download_progress['completed'] = 0
+            download_progress['current_model'] = ''
+            download_progress['status'] = 'running'
+            download_progress['errors'] = []
+            download_progress['last_error'] = None
+            download_progress['start_time'] = time.time()
+            download_progress['end_time'] = None
+            
+            # Get the processed models list from a file if it exists
+            progress_file = os.path.join(output_dir, '.download_progress.json')
+            if os.path.exists(progress_file):
+                try:
+                    with open(progress_file, 'r', encoding='utf-8') as f:
+                        saved_progress = json.load(f)
+                        download_progress['processed_models'] = set(saved_progress.get('processed_models', []))
+                        logger.info(f"Loaded previous progress, {len(download_progress['processed_models'])} models already processed")
+                except Exception as e:
+                    logger.error(f"Failed to load progress file: {e}")
+                    download_progress['processed_models'] = set()
+            else:
+                download_progress['processed_models'] = set()
+            
+            # Start the download task
+            is_downloading = True
+            download_task = asyncio.create_task(
+                MiscRoutes._download_all_example_images(
+                    output_dir, 
+                    optimize, 
+                    model_types,
+                    delay
+                )
+            )
+            
+            # Create a copy for JSON serialization
+            response_progress = download_progress.copy()
+            response_progress['processed_models'] = list(download_progress['processed_models'])
+            
+            return web.json_response({
+                'success': True,
+                'message': 'Download started',
+                'status': response_progress
+            })
+            
+        except Exception as e:
+            logger.error(f"Failed to start example images download: {e}", exc_info=True)
+            return web.json_response({
+                'success': False,
+                'error': str(e)
+            }, status=500)
+
+    @staticmethod
+    async def get_example_images_status(request):
+        """Get the current status of example images download"""
+        global download_progress
+        
+        # Create a copy of the progress dict with the set converted to a list for JSON serialization
+        response_progress = download_progress.copy()
+        response_progress['processed_models'] = list(download_progress['processed_models'])
+        
+        return web.json_response({
+            'success': True,
+            'is_downloading': is_downloading,
+            'status': response_progress
+        })
+
+    @staticmethod
+    async def pause_example_images(request):
+        """Pause the example images download"""
+        global download_progress
+        
+        if not is_downloading:
+            return web.json_response({
+                'success': False,
+                'error': 'No download in progress'
+            }, status=400)
+        
+        download_progress['status'] = 'paused'
+        
+        return web.json_response({
+            'success': True,
+            'message': 'Download paused'
+        })
+
+    @staticmethod
+    async def resume_example_images(request):
+        """Resume the example images download"""
+        global download_progress
+        
+        if not is_downloading:
+            return web.json_response({
+                'success': False,
+                'error': 'No download in progress'
+            }, status=400)
+        
+        if download_progress['status'] == 'paused':
+            download_progress['status'] = 'running'
+            
+            return web.json_response({
+                'success': True,
+                'message': 'Download resumed'
+            })
+        else:
+            return web.json_response({
+                'success': False,
+                'error': f"Download is in '{download_progress['status']}' state, cannot resume"
+            }, status=400)
+
+    @staticmethod
+    async def _download_all_example_images(output_dir, optimize, model_types, delay):
+        """Download example images for all models
+        
+        Args:
+            output_dir: Base directory to save example images
+            optimize: Whether to optimize images
+            model_types: List of model types to process
+            delay: Delay between downloads to avoid rate limiting
+        """
+        global is_downloading, download_progress
+        
+        # Get CivitAI client with proxy support
+        civitai_client = await ServiceRegistry.get_civitai_client()
+        session = await civitai_client.session
+        
+        try:
+            # Get the scanners
+            scanners = []
+            if 'lora' in model_types:
+                lora_scanner = await ServiceRegistry.get_lora_scanner()
+                scanners.append(('lora', lora_scanner))
+            
+            if 'checkpoint' in model_types:
+                checkpoint_scanner = await ServiceRegistry.get_checkpoint_scanner()
+                scanners.append(('checkpoint', checkpoint_scanner))
+            
+            # Get all models from all scanners
+            all_models = []
+            for scanner_type, scanner in scanners:
+                cache = await scanner.get_cached_data()
+                if cache and cache.raw_data:
+                    for model in cache.raw_data:
+                        # Only process models with images and a valid sha256
+                        if model.get('civitai') and model.get('civitai', {}).get('images') and model.get('sha256'):
+                            all_models.append((scanner_type, model))
+            
+            # Update total count
+            download_progress['total'] = len(all_models)
+            logger.info(f"Found {download_progress['total']} models with example images")
+            
+            # Process each model
+            for scanner_type, model in all_models:
+                # Check if download is paused
+                while download_progress['status'] == 'paused':
+                    await asyncio.sleep(1)
+                
+                # Check if download should continue
+                if download_progress['status'] != 'running':
+                    logger.info(f"Download stopped: {download_progress['status']}")
+                    break
+                
+                try:
+                    # Update current model info
+                    model_hash = model.get('sha256', '').lower()
+                    model_name = model.get('model_name', 'Unknown')
+                    download_progress['current_model'] = f"{model_name} ({model_hash[:8]})"
+                    
+                    # Skip if already processed
+                    if model_hash in download_progress['processed_models']:
+                        logger.debug(f"Skipping already processed model: {model_name}")
+                        download_progress['completed'] += 1
+                        continue
+                    
+                    # Create model directory
+                    model_dir = os.path.join(output_dir, model_hash)
+                    os.makedirs(model_dir, exist_ok=True)
+                    
+                    # Process images for this model
+                    images = model.get('civitai', {}).get('images', [])
+                    
+                    if not images:
+                        logger.debug(f"No images found for model: {model_name}")
+                        download_progress['processed_models'].add(model_hash)
+                        download_progress['completed'] += 1
+                        continue
+                    
+                    # Download example images
+                    for i, image in enumerate(images, 1):
+                        image_url = image.get('url')
+                        if not image_url:
+                            continue
+                        
+                        # Get image filename from URL
+                        image_filename = os.path.basename(image_url.split('?')[0])
+                        image_ext = os.path.splitext(image_filename)[1].lower()
+                        
+                        # Handle both images and videos
+                        is_image = image_ext in SUPPORTED_MEDIA_EXTENSIONS['images']
+                        is_video = image_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']
+                        
+                        if not (is_image or is_video):
+                            logger.debug(f"Skipping unsupported file type: {image_filename}")
+                            continue
+                        
+                        save_filename = f"image_{i}{image_ext}"
+                        
+                        # Check if already downloaded
+                        save_path = os.path.join(model_dir, save_filename)
+                        if os.path.exists(save_path):
+                            logger.debug(f"File already exists: {save_path}")
+                            continue
+                        
+                        # Download the file
+                        try:
+                            logger.debug(f"Downloading {save_filename} for {model_name}")
+                            
+                            # Direct download using session from CivitAI client
+                            async with session.get(image_url, timeout=60) as response:
+                                if response.status == 200:
+                                    if is_image and optimize:
+                                        # For images, optimize if requested
+                                        image_data = await response.read()
+                                        optimized_data, ext = ExifUtils.optimize_image(
+                                            image_data, 
+                                            target_width=EXAMPLE_IMAGE_WIDTH, 
+                                            format='webp', 
+                                            quality=85, 
+                                            preserve_metadata=False
+                                        )
+                                        
+                                        # Update save filename if format changed
+                                        if ext == '.webp':
+                                            save_filename = os.path.splitext(save_filename)[0] + '.webp'
+                                            save_path = os.path.join(model_dir, save_filename)
+                                        
+                                        # Save the optimized image
+                                        with open(save_path, 'wb') as f:
+                                            f.write(optimized_data)
+                                    else:
+                                        # For videos or unoptimized images, save directly
+                                        with open(save_path, 'wb') as f:
+                                            async for chunk in response.content.iter_chunked(8192):
+                                                if chunk:
+                                                    f.write(chunk)
+                                else:
+                                    logger.warning(f"Failed to download file: {image_url}, status code: {response.status}")
+                            
+                            # Add a delay between downloads
+                            await asyncio.sleep(delay)
+                        except Exception as e:
+                            error_msg = f"Error downloading file {image_url}: {str(e)}"
+                            logger.error(error_msg)
+                            download_progress['errors'].append(error_msg)
+                            download_progress['last_error'] = error_msg
+                            # Continue with next file
+                    
+                    # Mark model as processed
+                    download_progress['processed_models'].add(model_hash)
+                    
+                    # Save progress to file periodically
+                    if download_progress['completed'] % 10 == 0 or download_progress['completed'] == download_progress['total'] - 1:
+                        progress_file = os.path.join(output_dir, '.download_progress.json')
+                        with open(progress_file, 'w', encoding='utf-8') as f:
+                            json.dump({
+                                'processed_models': list(download_progress['processed_models']),
+                                'completed': download_progress['completed'],
+                                'total': download_progress['total'],
+                                'last_update': time.time()
+                            }, f, indent=2)
+                
+                except Exception as e:
+                    error_msg = f"Error processing model {model.get('model_name')}: {str(e)}"
+                    logger.error(error_msg, exc_info=True)
+                    download_progress['errors'].append(error_msg)
+                    download_progress['last_error'] = error_msg
+                
+                # Update progress
+                download_progress['completed'] += 1
+            
+            # Mark as completed
+            download_progress['status'] = 'completed'
+            download_progress['end_time'] = time.time()
+            logger.info(f"Example images download completed: {download_progress['completed']}/{download_progress['total']} models processed")
+        
+        except Exception as e:
+            error_msg = f"Error during example images download: {str(e)}"
+            logger.error(error_msg, exc_info=True)
+            download_progress['errors'].append(error_msg)
+            download_progress['last_error'] = error_msg
+            download_progress['status'] = 'error'
+            download_progress['end_time'] = time.time()
+        
+        finally:
+            # Save final progress to file
+            try:
+                progress_file = os.path.join(output_dir, '.download_progress.json')
+                with open(progress_file, 'w', encoding='utf-8') as f:
+                    json.dump({
+                        'processed_models': list(download_progress['processed_models']),
+                        'completed': download_progress['completed'],
+                        'total': download_progress['total'],
+                        'last_update': time.time(),
+                        'status': download_progress['status']
+                    }, f, indent=2)
+            except Exception as e:
+                logger.error(f"Failed to save progress file: {e}")
+            
+            # Set download status to not downloading
+            is_downloading = False
+
+    @staticmethod
+    async def browse_folder(request):
+        """Open a folder browser dialog and return the selected path
+        
+        Expects a JSON body with:
+        {
+            "initial_dir": "string" // Optional initial directory
+        }
+        """
+        try:
+            # Parse the request body
+            data = await request.json()
+            initial_dir = data.get('initial_dir', '')
+            
+            # If initial_dir doesn't exist, use a default path
+            if initial_dir and not os.path.isdir(initial_dir):
+                initial_dir = ''
+            
+            # Create a hidden root window for the dialog
+            root = tk.Tk()
+            root.withdraw()  # Hide the root window
+            
+            # Show the folder selection dialog
+            folder_path = filedialog.askdirectory(
+                initialdir=initial_dir,
+                title="Select folder for example images"
+            )
+            
+            # Destroy the root window
+            root.destroy()
+            
+            if folder_path:
+                # Convert to proper path format for the OS
+                folder_path = os.path.normpath(folder_path)
+                
+                return web.json_response({
+                    'success': True,
+                    'folder': folder_path
+                })
+            else:
+                return web.json_response({
+                    'success': False,
+                    'error': 'No folder selected'
+                })
+            
+        except Exception as e:
+            logger.error(f"Failed to browse folder: {e}", exc_info=True)
+            return web.json_response({
+                'success': False,
+                'error': str(e)
+            }, status=500)

py/routes/usage_stats_routes.py

@@ -1,69 +0,0 @@
-import logging
-from aiohttp import web
-from ..utils.usage_stats import UsageStats
-
-logger = logging.getLogger(__name__)
-
-class UsageStatsRoutes:
-    """Routes for handling usage statistics updates"""
-    
-    @staticmethod
-    def setup_routes(app):
-        """Register usage stats routes"""
-        app.router.add_post('/loras/api/update-usage-stats', UsageStatsRoutes.update_usage_stats)
-        app.router.add_get('/loras/api/get-usage-stats', UsageStatsRoutes.get_usage_stats)
-    
-    @staticmethod
-    async def update_usage_stats(request):
-        """
-        Update usage statistics based on a prompt_id
-        
-        Expects a JSON body with:
-        {
-            "prompt_id": "string"
-        }
-        """
-        try:
-            # Parse the request body
-            data = await request.json()
-            prompt_id = data.get('prompt_id')
-            
-            if not prompt_id:
-                return web.json_response({
-                    'success': False,
-                    'error': 'Missing prompt_id'
-                }, status=400)
-            
-            # Call the UsageStats to process this prompt_id synchronously
-            usage_stats = UsageStats()
-            await usage_stats.process_execution(prompt_id)
-            
-            return web.json_response({
-                'success': True
-            })
-            
-        except Exception as e:
-            logger.error(f"Failed to update usage stats: {e}", exc_info=True)
-            return web.json_response({
-                'success': False,
-                'error': str(e)
-            }, status=500)
-    
-    @staticmethod
-    async def get_usage_stats(request):
-        """Get current usage statistics"""
-        try:
-            usage_stats = UsageStats()
-            stats = await usage_stats.get_stats()
-            
-            return web.json_response({
-                'success': True,
-                'data': stats
-            })
-            
-        except Exception as e:
-            logger.error(f"Failed to get usage stats: {e}", exc_info=True)
-            return web.json_response({
-                'success': False,
-                'error': str(e)
-            }, status=500)

py/utils/constants.py

@@ -8,18 +8,16 @@ NSFW_LEVELS = {
 }
 
 # preview extensions
-PREVIEW_EXTENSIONS = [
-    '.webp',
-    '.preview.webp',
-    '.preview.png', 
-    '.preview.jpeg', 
-    '.preview.jpg', 
-    '.preview.mp4',
-    '.png', 
-    '.jpeg', 
-    '.jpg', 
-    '.mp4'
-]
+PREVIEW_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif']
 
 # Card preview image width
-CARD_PREVIEW_WIDTH = 480
+CARD_PREVIEW_WIDTH = 480
+
+# Width for optimized example images
+EXAMPLE_IMAGE_WIDTH = 832
+
+# Supported media extensions for example downloads
+SUPPORTED_MEDIA_EXTENSIONS = {
+    'images': ['.jpg', '.jpeg', '.png', '.webp', '.gif'],
+    'videos': ['.mp4', '.webm']
+}