nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: normalize download category strings across backend and frontend

Browse Source 
Add a shared normalize_download_category helper (with CATEGORY_MAP) in Python
and use it in get_download_directory to prefer redirected paths and to rank
diffusion_models above legacy unet folders. Mirror the same normalization in
the JS downloadTargetMethods, applying it consistently when coalescing
downloadDirectories, building dropdown keys, and reading/writing selections.
Also fix display labels to use 'diffusion_models' instead of 'unet'.
This commit is contained in:
Dariusz L
2026-06-12 16:42:58 +02:00
parent bbfd50a075
commit 7a1473b58e
3 changed files with 123 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+6
View File
@@ -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",
core/downloader.py
+73 -30
View File
@@ -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}")
web/resolver/search/download_target_methods.js
+44 -9
View File
@@ -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 = `<div class="mr-download-target">`;
@@ -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,