""" Model Downloader Module Handles downloading models from various sources with progress tracking. """ import os import threading import time import requests from typing import Dict, Any, Optional, Callable, List from pathlib import Path from collections import deque from .log_system.log_funcs import ( log_debug, log_info, log_warn, log_error, log_exception, ) try: import folder_paths except ImportError: folder_paths = None # Download state tracking download_progress: Dict[str, Dict[str, Any]] = {} download_lock = threading.Lock() cancelled_downloads: set = set() # Speed calculation settings SPEED_HISTORY_SIZE = 5 # Number of samples for smoothing CHUNK_SIZE = 1024 * 1024 # 1MB chunks for faster downloads CLI_LOG_INTERVAL = 5 # Log progress to CLI every N seconds # Map common category names to folder_paths keys. "unet" and "clip" are legacy # in ComfyUI; new files should download into diffusion_models/text_encoders. CATEGORY_MAP = { "checkpoint": "checkpoints", "checkpoints": "checkpoints", "lora": "loras", "loras": "loras", "vae": "vae", "controlnet": "controlnet", "clip": "text_encoders", "clips": "text_encoders", "clip_vision": "clip_vision", "upscaler": "upscale_models", "upscale_models": "upscale_models", "embeddings": "embeddings", "embedding": "embeddings", "diffusion_model": "diffusion_models", "diffusion_models": "diffusion_models", "unet": "diffusion_models", "latent_upscale_model": "latent_upscale_models", "latent_upscale_models": "latent_upscale_models", "text_encoders": "text_encoders", "text_encoder": "text_encoders", "ipadapter": "ipadapter", "ip_adapter": "ipadapter", # Legacy/unsupported categories map to checkpoints as fallback "default": "upscale_models", # Model list uses "default" for upscalers } def format_bytes(bytes_value: int) -> str: """Format bytes to human readable string (e.g., 1.5 GB).""" if bytes_value == 0: return "0 B" k = 1024 sizes = ["B", "KB", "MB", "GB", "TB"] i = 0 while bytes_value >= k and i < len(sizes) - 1: bytes_value /= k i += 1 return f"{bytes_value:.1f} {sizes[i]}" def normalize_download_category(category: str) -> str: """Return the canonical ComfyUI folder_paths key for a download category.""" token = ( str(category or "") .strip() .lower() .replace("\\", "_") .replace("/", "_") .replace("-", "_") .replace(" ", "_") ) while "__" in token: token = token.replace("__", "_") token = token.strip("_") return CATEGORY_MAP.get(token, token or "checkpoints") def get_download_directory(category: str, preferred_base_directory: str = "") -> Optional[str]: """ Get the appropriate download directory for a model category. Args: category: Model category (e.g., 'checkpoints', 'loras', 'vae') preferred_base_directory: Optional configured base directory to use Returns: Absolute path to the download directory, or None if not found """ global folder_paths if folder_paths is None: # Try to import again - ComfyUI might have initialized since last check try: import folder_paths as fp folder_paths = fp except ImportError: return None folder_key = normalize_download_category(category) def _normalize(path_value: str) -> str: return os.path.normcase(os.path.abspath(path_value)) def _is_within(path_value: str, root_value: str) -> bool: try: return os.path.commonpath([_normalize(path_value), _normalize(root_value)]) == _normalize(root_value) except Exception: return False def _choose_preferred_path(paths: List[str], preferred_key: str = "") -> Optional[str]: if not paths: return None comfy_root = os.path.dirname(os.path.abspath(getattr(folder_paths, "__file__", ""))) def _basename(path_value: str) -> str: return os.path.basename(os.path.normpath(path_value)).lower() def _prefer_redirected(candidate_paths: List[str]) -> Optional[str]: if not candidate_paths: return None if comfy_root: redirected_paths = [path for path in candidate_paths if not _is_within(path, comfy_root)] if redirected_paths: return redirected_paths[0] return candidate_paths[0] if preferred_key == "diffusion_models": canonical_paths = [path for path in paths if _basename(path) == "diffusion_models"] preferred_path = _prefer_redirected(canonical_paths) if preferred_path: return preferred_path non_legacy_paths = [path for path in paths if _basename(path) != "unet"] preferred_path = _prefer_redirected(non_legacy_paths) if preferred_path: return preferred_path if preferred_key == "text_encoders": canonical_paths = [path for path in paths if _basename(path) == "text_encoders"] preferred_path = _prefer_redirected(canonical_paths) if preferred_path: return preferred_path non_legacy_paths = [path for path in paths if _basename(path) != "clip"] preferred_path = _prefer_redirected(non_legacy_paths) if preferred_path: return preferred_path if comfy_root: redirected_paths = [path for path in paths if not _is_within(path, comfy_root)] if redirected_paths: return redirected_paths[0] return paths[0] try: paths = folder_paths.get_folder_paths(folder_key) if paths: if preferred_base_directory: preferred_normalized = _normalize(preferred_base_directory) for path in paths: if _normalize(path) == preferred_normalized: return path return _choose_preferred_path(paths, folder_key) # If category not found, try to get any models directory as fallback all_names = folder_paths.get_folder_names() if all_names: # Fall back to first available directory fallback_paths = folder_paths.get_folder_paths(all_names[0]) if fallback_paths: return _choose_preferred_path(fallback_paths, all_names[0]) except Exception as e: log_debug(f"Could not get folder path for {folder_key}: {e}") return None def generate_download_id() -> str: """Generate a unique download ID.""" import uuid return str(uuid.uuid4())[:8] def download_file( url: str, dest_path: str, download_id: str, headers: Optional[Dict[str, str]] = None, chunk_size: int = None, progress_callback: Optional[Callable[[int, int], None]] = None, ) -> Dict[str, Any]: """ Download a file from URL with progress tracking and speed calculation. Args: url: URL to download from dest_path: Destination file path download_id: Unique ID for tracking this download headers: Optional HTTP headers (for auth tokens) chunk_size: Download chunk size in bytes (defaults to 1MB) progress_callback: Optional callback(downloaded_bytes, total_bytes) Returns: Result dictionary with status and info """ global download_progress, cancelled_downloads # Use default 1MB chunk size if not specified if chunk_size is None: chunk_size = CHUNK_SIZE result = { "success": False, "download_id": download_id, "path": dest_path, "error": None, "size": 0, } # Initialize progress tracking with speed calculation start_time = time.time() speed_history: deque = deque(maxlen=SPEED_HISTORY_SIZE) last_speed_update = start_time last_downloaded = 0 last_cli_log = start_time # Track when we last logged to CLI with download_lock: download_progress[download_id] = { "status": "starting", "progress": 0, "total_size": 0, "downloaded": 0, "filename": os.path.basename(dest_path), "path": dest_path, "directory": os.path.dirname(dest_path), "url": url, "error": None, "speed": 0, # bytes per second "start_time": start_time, } try: # Ensure destination directory exists os.makedirs(os.path.dirname(dest_path), exist_ok=True) # Verbose logging - what model and from where filename = os.path.basename(dest_path) source = ( "HuggingFace" if "huggingface.co" in url else "CivitAI" if "civitai.com" in url else "URL" ) log_info(f"Starting download: {filename}") log_info(f"Source: {source}") log_info(f"URL: {url}") # Start download response = requests.get(url, headers=headers, stream=True, timeout=30) response.raise_for_status() # Get total size total_size = int(response.headers.get("content-length", 0)) total_size_str = format_bytes(total_size) if total_size > 0 else "unknown" log_info(f"Size: {total_size_str}") with download_lock: download_progress[download_id]["total_size"] = total_size download_progress[download_id]["status"] = "downloading" downloaded = 0 # Download with progress and speed calculation cancelled = False with open(dest_path, "wb") as f: for chunk in response.iter_content(chunk_size=chunk_size): # Check for cancellation if download_id in cancelled_downloads: cancelled = True break if chunk: f.write(chunk) downloaded += len(chunk) # Calculate speed with smoothing current_time = time.time() time_delta = current_time - last_speed_update # Update speed every 0.5 seconds to avoid too frequent calculations if time_delta >= 0.5: bytes_delta = downloaded - last_downloaded instant_speed = ( bytes_delta / time_delta if time_delta > 0 else 0 ) speed_history.append(instant_speed) # Calculate smoothed speed (average of recent samples) smoothed_speed = ( sum(speed_history) / len(speed_history) if speed_history else 0 ) last_speed_update = current_time last_downloaded = downloaded # Update progress with speed with download_lock: download_progress[download_id]["downloaded"] = downloaded download_progress[download_id]["speed"] = int( smoothed_speed ) if total_size > 0: download_progress[download_id]["progress"] = int( (downloaded / total_size) * 100 ) # CLI progress logging (every CLI_LOG_INTERVAL seconds) if current_time - last_cli_log >= CLI_LOG_INTERVAL: last_cli_log = current_time progress_pct = ( int((downloaded / total_size) * 100) if total_size > 0 else 0 ) downloaded_str = format_bytes(downloaded) total_str = ( format_bytes(total_size) if total_size > 0 else "?" ) speed_str = format_bytes(int(smoothed_speed)) + "/s" log_info( f"Progress: {downloaded_str} / {total_str} ({progress_pct}%) - {speed_str}" ) else: # Just update downloaded bytes without recalculating speed with download_lock: download_progress[download_id]["downloaded"] = downloaded if total_size > 0: download_progress[download_id]["progress"] = int( (downloaded / total_size) * 100 ) if progress_callback: progress_callback(downloaded, total_size) # Handle cancellation after file is closed (so we can delete it on Windows) # Also check if cancellation was requested while we were finishing up if cancelled or download_id in cancelled_downloads: with download_lock: download_progress[download_id]["status"] = "cancelled" # Clean up partial/incomplete file try: if os.path.exists(dest_path): os.remove(dest_path) log_info(f"Cancelled: {filename} - incomplete file deleted") else: log_info(f"Cancelled: {filename} - no file to delete") except Exception as e: log_warn(f"Could not delete incomplete file {dest_path}: {e}") # Try harder on Windows - sometimes the file handle takes a moment to release try: time.sleep(0.5) # time is already imported at module level if os.path.exists(dest_path): os.remove(dest_path) log_info( f"Cancelled: {filename} - incomplete file deleted (delayed)" ) except Exception: pass result["error"] = "Download cancelled" cancelled_downloads.discard(download_id) return result # Success with download_lock: download_progress[download_id]["status"] = "completed" download_progress[download_id]["progress"] = 100 download_progress[download_id]["speed"] = 0 # Reset speed on completion result["success"] = True result["size"] = downloaded # CLI completion log elapsed = time.time() - start_time avg_speed = downloaded / elapsed if elapsed > 0 else 0 log_info(f"✓ Download complete: {filename}") log_info( f"Size: {format_bytes(downloaded)}, Time: {elapsed:.1f}s, Avg speed: {format_bytes(int(avg_speed))}/s" ) except requests.exceptions.RequestException as e: error_msg = str(e) # Check for specific HTTP errors if hasattr(e, "response") and e.response is not None: status_code = e.response.status_code if status_code in [401, 403]: if "huggingface.co" in url: error_msg = f"Unauthorized (HTTP {status_code}): HuggingFace token may be required." elif "civitai.com" in url: error_msg = f"Unauthorized (HTTP {status_code}): CivitAI API key may be required." else: error_msg = ( f"Unauthorized (HTTP {status_code}): Authentication required." ) elif status_code == 404: error_msg = "Model not found (HTTP 404): The file may have been moved or deleted." with download_lock: download_progress[download_id]["status"] = "error" download_progress[download_id]["error"] = error_msg result["error"] = error_msg # CLI error log log_error(f"✗ Download failed: {os.path.basename(dest_path)}") log_error(f"Error: {error_msg}") # Clean up partial file try: if os.path.exists(dest_path): os.remove(dest_path) except: pass except Exception as e: error_msg = str(e) with download_lock: download_progress[download_id]["status"] = "error" download_progress[download_id]["error"] = error_msg result["error"] = error_msg # CLI error log log_error(f"✗ Download failed: {os.path.basename(dest_path)}") log_error(f"Error: {error_msg}") log_error(f"Download error: {e}", exc_info=True) return result def download_model( url: str, filename: str, category: str, download_id: Optional[str] = None, headers: Optional[Dict[str, str]] = None, subfolder: str = "", base_directory: str = "", ) -> Dict[str, Any]: """ Download a model to the appropriate directory. Args: url: URL to download from filename: Filename to save as category: Model category for directory selection download_id: Optional download ID (generated if not provided) headers: Optional HTTP headers subfolder: Optional subfolder within category directory base_directory: Optional configured base directory to use Returns: Result dictionary """ if download_id is None: download_id = generate_download_id() # Get destination directory dest_dir = get_download_directory(category, base_directory) if not dest_dir: return { "success": False, "download_id": download_id, "error": f"Could not find directory for category: {category}", } # Add subfolder if specified if subfolder: dest_dir = os.path.join(dest_dir, subfolder) dest_path = os.path.join(dest_dir, filename) # Check if file already exists if os.path.exists(dest_path): return { "success": False, "download_id": download_id, "error": f"File already exists: {dest_path}", "path": dest_path, } return download_file(url, dest_path, download_id, headers) def get_progress(download_id: str) -> Optional[Dict[str, Any]]: """Get progress for a specific download.""" with download_lock: return download_progress.get(download_id, {}).copy() def get_all_progress() -> Dict[str, Dict[str, Any]]: """Get progress for all downloads.""" with download_lock: return {k: v.copy() for k, v in download_progress.items()} def cancel_download(download_id: str) -> bool: """Cancel a download in progress.""" cancelled_downloads.add(download_id) return True def clear_completed_downloads(): """Clear completed/failed downloads from progress tracking.""" with download_lock: to_remove = [ did for did, info in download_progress.items() if info.get("status") in ("completed", "error", "cancelled") ] for did in to_remove: del download_progress[did] cancelled_downloads.discard(did) def start_background_download( url: str, filename: str, category: str, headers: Optional[Dict[str, str]] = None, subfolder: str = "", base_directory: str = "", ) -> str: """ Start a download in a background thread. Returns: download_id for tracking progress """ download_id = generate_download_id() initial_directory = get_download_directory(category, base_directory) or "" if initial_directory and subfolder: initial_directory = os.path.join(initial_directory, subfolder) initial_path = os.path.join(initial_directory, filename) if initial_directory else "" # Pre-initialize progress dict so it's always available for polling # even if download fails before download_file is called with download_lock: download_progress[download_id] = { "status": "starting", "progress": 0, "total_size": 0, "downloaded": 0, "filename": filename, "path": initial_path, "directory": initial_directory, "url": url, "error": None, "speed": 0, "start_time": time.time(), } def run_download(): try: result = download_model( url, filename, category, download_id, headers, subfolder, base_directory ) if not result.get("success"): # Mark as error if download failed with download_lock: if download_id in download_progress: download_progress[download_id]["status"] = "error" download_progress[download_id]["error"] = result.get( "error", "Download failed" ) if result.get("path"): download_progress[download_id]["path"] = result["path"] download_progress[download_id]["directory"] = os.path.dirname( result["path"] ) except Exception as e: # Ensure any exception is captured and logged with download_lock: download_progress[download_id] = { "status": "error", "progress": 0, "total_size": 0, "downloaded": 0, "filename": filename, "path": "", "directory": "", "url": url, "error": str(e), "speed": 0, "start_time": time.time(), } thread = threading.Thread(target=run_download, daemon=True) thread.start() return download_id