diff --git a/__init__.py b/__init__.py index 5ea4754..6fafaf3 100644 --- a/__init__.py +++ b/__init__.py @@ -89,6 +89,7 @@ class ModelResolverExtension: get_all_progress, cancel_download, get_download_directory, + normalize_download_category, ) from .core.sources.popular import ( get_popular_model_url, @@ -759,6 +760,9 @@ class ModelResolverExtension: "checkpoints": "checkpoints", "vae": "vae", "controlnet": "controlnet", + "diffusion_model": "diffusion_models", + "diffusion_models": "diffusion_models", + "unet": "diffusion_models", "upscale_models": "upscale_models", } folder_type = category_map.get( @@ -1741,6 +1745,7 @@ class ModelResolverExtension: url = data.get("url", "") filename = data.get("filename", "") category = data.get("category", "checkpoints") + category = normalize_download_category(category) subfolder = data.get("subfolder", "") if not url: @@ -1935,6 +1940,7 @@ class ModelResolverExtension: "upscale_models": "upscale_models", "embeddings": "embeddings", "embedding": "embeddings", + "diffusion_model": "diffusion_models", "diffusion_models": "diffusion_models", "unet": "diffusion_models", "text_encoders": "text_encoders", diff --git a/core/downloader.py b/core/downloader.py index 86ecd57..4a2c287 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -35,6 +35,34 @@ 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" is legacy in ComfyUI; +# new diffusion/UNet/transformer files should download into diffusion_models. +CATEGORY_MAP = { + "checkpoint": "checkpoints", + "checkpoints": "checkpoints", + "lora": "loras", + "loras": "loras", + "vae": "vae", + "controlnet": "controlnet", + "clip": "clip", + "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).""" @@ -49,6 +77,23 @@ def format_bytes(bytes_value: int) -> str: 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) -> Optional[str]: """ Get the appropriate download directory for a model category. @@ -70,33 +115,7 @@ def get_download_directory(category: str) -> Optional[str]: except ImportError: return None - # Map common category names to folder_paths keys - category_map = { - "checkpoint": "checkpoints", - "checkpoints": "checkpoints", - "lora": "loras", - "loras": "loras", - "vae": "vae", - "controlnet": "controlnet", - "clip": "clip", - "clip_vision": "clip_vision", - "upscaler": "upscale_models", - "upscale_models": "upscale_models", - "embeddings": "embeddings", - "embedding": "embeddings", - "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 - } - - folder_key = category_map.get(category.lower(), category.lower()) + folder_key = normalize_download_category(category) def _normalize(path_value: str) -> str: return os.path.normcase(os.path.abspath(path_value)) @@ -107,11 +126,35 @@ def get_download_directory(category: str) -> Optional[str]: except Exception: return False - def _choose_preferred_path(paths: List[str]) -> Optional[str]: + 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 comfy_root: redirected_paths = [path for path in paths if not _is_within(path, comfy_root)] if redirected_paths: @@ -122,7 +165,7 @@ def get_download_directory(category: str) -> Optional[str]: try: paths = folder_paths.get_folder_paths(folder_key) if paths: - return _choose_preferred_path(paths) + 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() @@ -130,7 +173,7 @@ def get_download_directory(category: str) -> Optional[str]: # 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) + 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}") diff --git a/web/resolver/search/download_target_methods.js b/web/resolver/search/download_target_methods.js index 7c689b9..7478588 100644 --- a/web/resolver/search/download_target_methods.js +++ b/web/resolver/search/download_target_methods.js @@ -30,7 +30,17 @@ export const downloadTargetMethods = { const resp = await api.fetchApi('/model_resolver/directories'); if (!resp.ok) throw new Error(`HTTP ${resp.status}`); const directories = await resp.json(); - this.downloadDirectories = directories && typeof directories === 'object' ? directories : {}; + if (directories && typeof directories === 'object') { + this.downloadDirectories = Object.entries(directories).reduce((acc, [key, value]) => { + const normalizedKey = this.normalizeDownloadCategory(key); + if (normalizedKey && !acc[normalizedKey]) { + acc[normalizedKey] = value; + } + return acc; + }, {}); + } else { + this.downloadDirectories = {}; + } } catch (e) { console.warn('Model Resolver: could not load download directories', e); this.downloadDirectories = {}; @@ -71,7 +81,31 @@ export const downloadTargetMethods = { return Boolean(this.capabilities?.sources?.[source]); }, + normalizeDownloadCategory(category = '') { + const token = String(category || '') + .trim() + .toLowerCase() + .replace(/[\/\\\s-]+/g, '_') + .replace(/_+/g, '_') + .replace(/^_|_$/g, ''); + const categoryMap = { + checkpoint: 'checkpoints', + lora: 'loras', + embedding: 'embeddings', + upscaler: 'upscale_models', + unet: 'diffusion_models', + diffusion_model: 'diffusion_models', + diffusion_models: 'diffusion_models', + text_encoder: 'text_encoders', + text_encoders: 'text_encoders', + ip_adapter: 'ipadapter', + 'default': 'upscale_models' + }; + return categoryMap[token] || token || 'checkpoints'; + }, + getCategoryDisplayName(category = '') { + category = this.normalizeDownloadCategory(category); const displayNames = { 'checkpoints': 'checkpoint', 'loras': 'lora', @@ -80,7 +114,7 @@ export const downloadTargetMethods = { 'embeddings': 'embedding', 'upscale_models': 'upscale_model', 'latent_upscale_models': 'latent_upscale_model', - 'diffusion_models': 'unet', + 'diffusion_models': 'diffusion_models', 'text_encoders': 'text encoders', 'clip': 'clip', 'clip_vision': 'clip_vision', @@ -159,6 +193,7 @@ export const downloadTargetMethods = { }, getCategoryTokenName(category = '') { + category = this.normalizeDownloadCategory(category); const tokenNames = { 'checkpoints': 'checkpoint', 'loras': 'lora', @@ -167,7 +202,7 @@ export const downloadTargetMethods = { 'embeddings': 'embedding', 'upscale_models': 'upscale_model', 'latent_upscale_models': 'latent_upscale_model', - 'diffusion_models': 'unet', + 'diffusion_models': 'diffusion_model', 'text_encoders': 'text_encoders', 'clip': 'clip', 'clip_vision': 'clip_vision', @@ -179,10 +214,10 @@ export const downloadTargetMethods = { getDownloadCategoryOptions(defaultCategory = 'checkpoints') { const directories = this.downloadDirectories || {}; const keys = Object.keys(directories); - const preferred = defaultCategory || 'checkpoints'; + const preferred = this.normalizeDownloadCategory(defaultCategory || 'checkpoints'); const ordered = [ preferred, - ...keys.filter(key => key !== preferred) + ...keys.map(key => this.normalizeDownloadCategory(key)).filter(key => key !== preferred) ].filter((value, index, arr) => value && arr.indexOf(value) === index); return ordered.length > 0 ? ordered : [preferred]; @@ -197,7 +232,7 @@ export const downloadTargetMethods = { }, getAvailableSubfolders(category = '') { - return this.downloadSubfolders.get((category || '').toLowerCase()) || []; + return this.downloadSubfolders.get(this.normalizeDownloadCategory(category)) || []; }, isAutoFillSubfolderEnabled() { @@ -342,7 +377,7 @@ export const downloadTargetMethods = { }, async ensureDownloadSubfoldersLoaded(category = '') { - const key = (category || '').trim().toLowerCase(); + const key = this.normalizeDownloadCategory(category); if (!key) return []; if (key === 'unknown') { this.downloadSubfolders.set(key, []); @@ -373,7 +408,7 @@ export const downloadTargetMethods = { const categoryListId = `download-category-list-${missing.node_id}-${missing.widget_index}`; const subfolderListId = `download-subfolder-list-${missing.node_id}-${missing.widget_index}`; const saved = this.getSavedDownloadTargetSelection(missing); - const selectedCategory = saved?.category || defaultCategory || 'checkpoints'; + const selectedCategory = this.normalizeDownloadCategory(saved?.category || defaultCategory || 'checkpoints'); const selectedSubfolder = saved ? saved.subfolder || '' : ''; let html = `
`; @@ -399,7 +434,7 @@ export const downloadTargetMethods = { getDownloadTargetSelection(missing, fallbackCategory = 'checkpoints') { const categoryEl = this.contentElement?.querySelector(`#download-category-${missing.node_id}-${missing.widget_index}`); const subfolderEl = this.contentElement?.querySelector(`#download-subfolder-${missing.node_id}-${missing.widget_index}`); - const category = this.getDropdownValue(categoryEl) || fallbackCategory || 'checkpoints'; + const category = this.normalizeDownloadCategory(this.getDropdownValue(categoryEl) || fallbackCategory || 'checkpoints'); const subfolder = (subfolderEl?.value || '').trim(); this.saveDownloadTargetSelection(missing, { category,