ComfyUI-Lora-Manager/py/services/download_manager.py

import logging
import os
import asyncio
from typing import Dict
from ..utils.models import LoraMetadata, CheckpointMetadata
from ..utils.constants import CARD_PREVIEW_WIDTH, VALID_LORA_TYPES, CIVITAI_MODEL_TAGS
from ..utils.exif_utils import ExifUtils
from ..utils.metadata_manager import MetadataManager
from .service_registry import ServiceRegistry
from .settings_manager import settings

# Download to temporary file first
import tempfile

logger = logging.getLogger(__name__)

class DownloadManager:
    _instance = None
    _lock = asyncio.Lock()
    
    @classmethod
    async def get_instance(cls):
        """Get singleton instance of DownloadManager"""
        async with cls._lock:
            if cls._instance is None:
                cls._instance = cls()
            return cls._instance

    def __init__(self):
        # Check if already initialized for singleton pattern
        if hasattr(self, '_initialized'):
            return
        self._initialized = True
        
        self._civitai_client = None  # Will be lazily initialized

    async def _get_civitai_client(self):
        """Lazily initialize CivitaiClient from registry"""
        if self._civitai_client is None:
            self._civitai_client = await ServiceRegistry.get_civitai_client()
        return self._civitai_client
    
    async def _get_lora_scanner(self):
        """Get the lora scanner from registry"""
        return await ServiceRegistry.get_lora_scanner()
        
    async def _get_checkpoint_scanner(self):
        """Get the checkpoint scanner from registry"""
        return await ServiceRegistry.get_checkpoint_scanner()

    async def download_from_civitai(self, model_id: int, 
                                  model_version_id: int, save_dir: str = None, 
                                  relative_path: str = '', progress_callback=None, use_default_paths: bool = False) -> Dict:
        """Download model from Civitai
        
        Args:
            model_id: Civitai model ID
            model_version_id: Civitai model version ID (optional, if not provided, will download the latest version)
            save_dir: Directory to save the model to
            relative_path: Relative path within save_dir
            progress_callback: Callback function for progress updates
            use_default_paths: Flag to indicate whether to use default paths
            
        Returns:
            Dict with download result
        """
        try:
            # Check if model version already exists in library
            if model_version_id is not None:
                # Case 1: model_version_id is provided, check both scanners
                lora_scanner = await self._get_lora_scanner()
                checkpoint_scanner = await self._get_checkpoint_scanner()
                
                # Check lora scanner first
                if await lora_scanner.check_model_version_exists(model_id, model_version_id):
                    return {'success': False, 'error': 'Model version already exists in lora library'}
                
                # Check checkpoint scanner
                if await checkpoint_scanner.check_model_version_exists(model_id, model_version_id):
                    return {'success': False, 'error': 'Model version already exists in checkpoint library'}

            # Get civitai client
            civitai_client = await self._get_civitai_client()

            # Get version info based on the provided identifier
            version_info = await civitai_client.get_model_version(model_id, model_version_id)
            
            if not version_info:
                return {'success': False, 'error': 'Failed to fetch model metadata'}

            model_type_from_info = version_info.get('model', {}).get('type', '').lower()
            if model_type_from_info == 'checkpoint':
                model_type = 'checkpoint'
            elif model_type_from_info in VALID_LORA_TYPES:
                model_type = 'lora'
            else:
                return {'success': False, 'error': f'Model type "{model_type_from_info}" is not supported for download'}
            
            # Case 2: model_version_id was None, check after getting version_info
            if model_version_id is None:
                version_model_id = version_info.get('modelId')
                version_id = version_info.get('id')
                
                if model_type == 'lora':
                    # Check lora scanner
                    lora_scanner = await self._get_lora_scanner()
                    if await lora_scanner.check_model_version_exists(version_model_id, version_id):
                        return {'success': False, 'error': 'Model version already exists in lora library'}
                elif model_type == 'checkpoint':
                    # Check checkpoint scanner
                    checkpoint_scanner = await self._get_checkpoint_scanner()
                    if await checkpoint_scanner.check_model_version_exists(version_model_id, version_id):
                        return {'success': False, 'error': 'Model version already exists in checkpoint library'}
            
            # Handle use_default_paths
            if use_default_paths:
                # Set save_dir based on model type
                if model_type == 'checkpoint':
                    default_path = settings.get('default_checkpoint_root')
                    if not default_path:
                        return {'success': False, 'error': 'Default checkpoint root path not set in settings'}
                    save_dir = default_path
                else:  # model_type == 'lora'
                    default_path = settings.get('default_lora_root')
                    if not default_path:
                        return {'success': False, 'error': 'Default lora root path not set in settings'}
                    save_dir = default_path
                    
                # Set relative_path to version_info.baseModel/prioritized_tag
                base_model = version_info.get('baseModel', '')
                model_tags = version_info.get('model', {}).get('tags', [])
                
                if base_model:
                    # Find the first Civitai model tag that exists in model_tags
                    prioritized_tag = None
                    for civitai_tag in CIVITAI_MODEL_TAGS:
                        if civitai_tag in model_tags:
                            prioritized_tag = civitai_tag
                            break
                    
                    # If no Civitai model tag found, fallback to first tag
                    if prioritized_tag is None and model_tags:
                        prioritized_tag = model_tags[0]
                    
                    if prioritized_tag:
                        relative_path = os.path.join(base_model, prioritized_tag)
                    else:
                        relative_path = base_model

            # Update save directory with relative path if provided
            if relative_path:
                save_dir = os.path.join(save_dir, relative_path)
                # Create directory if it doesn't exist
                os.makedirs(save_dir, exist_ok=True)

            # Check if this is an early access model
            if version_info.get('earlyAccessEndsAt'):
                early_access_date = version_info.get('earlyAccessEndsAt', '')
                # Convert to a readable date if possible
                try:
                    from datetime import datetime
                    date_obj = datetime.fromisoformat(early_access_date.replace('Z', '+00:00'))
                    formatted_date = date_obj.strftime('%Y-%m-%d')
                    early_access_msg = f"This model requires early access payment (until {formatted_date}). "
                except:
                    early_access_msg = "This model requires early access payment. "
                
                early_access_msg += "Please ensure you have purchased early access and are logged in to Civitai."
                logger.warning(f"Early access model detected: {version_info.get('name', 'Unknown')}")
                
                # We'll still try to download, but log a warning and prepare for potential failure
                if progress_callback:
                    await progress_callback(1)  # Show minimal progress to indicate we're trying

            # Report initial progress
            if progress_callback:
                await progress_callback(0)

            # 2. Get file information
            file_info = next((f for f in version_info.get('files', []) if f.get('primary')), None)
            if not file_info:
                return {'success': False, 'error': 'No primary file found in metadata'}

            # 3. Prepare download
            file_name = file_info['name']
            save_path = os.path.join(save_dir, file_name)

            # 5. Prepare metadata based on model type
            if model_type == "checkpoint":
                metadata = CheckpointMetadata.from_civitai_info(version_info, file_info, save_path)
                logger.info(f"Creating CheckpointMetadata for {file_name}")
            else:
                metadata = LoraMetadata.from_civitai_info(version_info, file_info, save_path)
                logger.info(f"Creating LoraMetadata for {file_name}")
            
            # 6. Start download process
            result = await self._execute_download(
                download_url=file_info.get('downloadUrl', ''),
                save_dir=save_dir,
                metadata=metadata,
                version_info=version_info,
                relative_path=relative_path,
                progress_callback=progress_callback,
                model_type=model_type
            )

            return result

        except Exception as e:
            logger.error(f"Error in download_from_civitai: {e}", exc_info=True)
            # Check if this might be an early access error
            error_str = str(e).lower()
            if "403" in error_str or "401" in error_str or "unauthorized" in error_str or "early access" in error_str:
                return {'success': False, 'error': f"Early access restriction: {str(e)}. Please ensure you have purchased early access and are logged in to Civitai."}
            return {'success': False, 'error': str(e)}

    async def _execute_download(self, download_url: str, save_dir: str, 
                              metadata, version_info: Dict, 
                              relative_path: str, progress_callback=None,
                              model_type: str = "lora") -> Dict:
        """Execute the actual download process including preview images and model files"""
        try:
            civitai_client = await self._get_civitai_client()
            save_path = metadata.file_path
            metadata_path = os.path.splitext(save_path)[0] + '.metadata.json'

            # Download preview image if available
            images = version_info.get('images', [])
            if images:
                # Report preview download progress
                if progress_callback:
                    await progress_callback(1)  # 1% progress for starting preview download

                # Check if it's a video or an image
                is_video = images[0].get('type') == 'video'
                
                if (is_video):
                    # For videos, use .mp4 extension
                    preview_ext = '.mp4'
                    preview_path = os.path.splitext(save_path)[0] + preview_ext
                    
                    # Download video directly
                    if await civitai_client.download_preview_image(images[0]['url'], preview_path):
                        metadata.preview_url = preview_path.replace(os.sep, '/')
                        metadata.preview_nsfw_level = images[0].get('nsfwLevel', 0)
                else:
                    # For images, use WebP format for better performance
                    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file:
                        temp_path = temp_file.name
                    
                    # Download the original image to temp path
                    if await civitai_client.download_preview_image(images[0]['url'], temp_path):
                        # Optimize and convert to WebP
                        preview_path = os.path.splitext(save_path)[0] + '.webp'
                        
                        # Use ExifUtils to optimize and convert the image
                        optimized_data, _ = ExifUtils.optimize_image(
                            image_data=temp_path,
                            target_width=CARD_PREVIEW_WIDTH,
                            format='webp',
                            quality=85,
                            preserve_metadata=False
                        )
                        
                        # Save the optimized image
                        with open(preview_path, 'wb') as f:
                            f.write(optimized_data)
                            
                        # Update metadata
                        metadata.preview_url = preview_path.replace(os.sep, '/')
                        metadata.preview_nsfw_level = images[0].get('nsfwLevel', 0)
                        
                        # Remove temporary file
                        try:
                            os.unlink(temp_path)
                        except Exception as e:
                            logger.warning(f"Failed to delete temp file: {e}")

                # Report preview download completion
                if progress_callback:
                    await progress_callback(3)  # 3% progress after preview download

            # Download model file with progress tracking
            success, result = await civitai_client._download_file(
                download_url, 
                save_dir,
                os.path.basename(save_path),
                progress_callback=lambda p: self._handle_download_progress(p, progress_callback)
            )

            if not success:
                # Clean up files on failure
                for path in [save_path, metadata_path, metadata.preview_url]:
                    if path and os.path.exists(path):
                        os.remove(path)
                return {'success': False, 'error': result}

            # 4. Update file information (size and modified time)
            metadata.update_file_info(save_path)

            # 5. Final metadata update
            await MetadataManager.save_metadata(save_path, metadata, True)

            # 6. Update cache based on model type
            if model_type == "checkpoint":
                scanner = await self._get_checkpoint_scanner()
                logger.info(f"Updating checkpoint cache for {save_path}")
            else:
                scanner = await self._get_lora_scanner()
                logger.info(f"Updating lora cache for {save_path}")
                
            # Convert metadata to dictionary
            metadata_dict = metadata.to_dict()

            # Add model to cache and save to disk in a single operation
            await scanner.add_model_to_cache(metadata_dict, relative_path)

            # Report 100% completion
            if progress_callback:
                await progress_callback(100)

            return {
                'success': True
            }

        except Exception as e:
            logger.error(f"Error in _execute_download: {e}", exc_info=True)
            # Clean up partial downloads
            for path in [save_path, metadata_path]:
                if path and os.path.exists(path):
                    os.remove(path)
            return {'success': False, 'error': str(e)}

    async def _handle_download_progress(self, file_progress: float, progress_callback):
        """Convert file download progress to overall progress
        
        Args:
            file_progress: Progress of file download (0-100)
            progress_callback: Callback function for progress updates
        """
        if progress_callback:
            # Scale file progress to 3-100 range (after preview download)
            overall_progress = 3 + (file_progress * 0.97)  # 97% of progress for file download
            await progress_callback(round(overall_progress))