diff --git a/core/downloader.py b/core/downloader.py index b2a07a9..381e5c7 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -86,6 +86,8 @@ def get_download_directory(category: str) -> Optional[str]: "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", diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index cd75d6b..a90c283 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -62,6 +62,7 @@ NODE_TYPE_TO_CATEGORY_HINTS = { "CheckpointLoader": "checkpoints", "unCLIPCheckpointLoader": "checkpoints", "VAELoader": "vae", + "VAELoaderKJ": "vae", "LoraLoader": "loras", "LoraLoaderModelOnly": "loras", "LoraLoaderV2": "loras", @@ -69,6 +70,8 @@ NODE_TYPE_TO_CATEGORY_HINTS = { "Lora Stacker (LoraManager)": "loras", # LoraManager Stacker node "Power Lora Loader (rgthree)": "loras", # rgthree's Power Lora Loader "UNETLoader": "diffusion_models", + "LatentUpscaleModelLoader": "latent_upscale_models", + "DualCLIPLoader": "text_encoders", "ControlNetLoader": "controlnet", "ControlNetLoaderAdvanced": "controlnet", "CLIPVisionLoader": "clip_vision", @@ -81,6 +84,23 @@ NODE_TYPE_TO_CATEGORY_HINTS = { "LTXVGemmaCLIPModelLoader": "text_encoders", } +# Model category hints by widget/input name. Workflow JSON does not always preserve +# widget names, but when it does this catches custom loaders without a node-type entry. +MODEL_WIDGET_NAME_TO_CATEGORY = { + "ckpt_name": "checkpoints", + "checkpoint": "checkpoints", + "model_name": "diffusion_models", + "unet_name": "diffusion_models", + "vae_name": "vae", + "clip_name": "text_encoders", + "clip_name1": "text_encoders", + "clip_name2": "text_encoders", + "clip_vision_name": "clip_vision", + "lora_name": "loras", + "control_net_name": "controlnet", + "upscale_model_name": "upscale_models", +} + # Keys within dict-type widget values that contain model file references. # Some nodes (e.g. rgthree Power Lora Loader) store model info as objects like # {"on": true, "lora": "name.safetensors", "strength": 1.0} inside widgets_values. @@ -90,10 +110,42 @@ NESTED_MODEL_KEYS = { "ckpt_name": "checkpoints", "checkpoint": "checkpoints", "vae_name": "vae", + "clip_name": "text_encoders", + "clip_name1": "text_encoders", + "clip_name2": "text_encoders", "control_net_name": "controlnet", } +def get_widget_category_hint(node: Dict[str, Any], widget_index: int) -> Optional[str]: + """ + Best-effort category hint for a widget index based on saved input/widget names. + + ComfyUI workflows usually store only widget values, but some node/workflow formats + include input/widget metadata. Use it when present; otherwise callers fall back to + node-type hints and folder_paths resolution. + """ + candidates = [] + + for key in ("widgets", "inputs"): + items = node.get(key) + if isinstance(items, list) and widget_index < len(items): + item = items[widget_index] + if isinstance(item, dict): + candidates.extend( + str(item.get(name, "")).strip() + for name in ("name", "widget", "label") + if item.get(name) + ) + + for candidate in candidates: + category = MODEL_WIDGET_NAME_TO_CATEGORY.get(candidate) + if category: + return category + + return None + + def is_model_filename(value: Any) -> bool: """ Check if a value looks like a model filename or URN. @@ -305,10 +357,15 @@ def get_node_model_info( # Get category hints for this node type category_hint = NODE_TYPE_TO_CATEGORY_HINTS.get(node_type) - categories_to_try = [category_hint] if category_hint else None # For each widget value, check if it looks like a model file or URN for idx, value in enumerate(widgets_values): + widget_category_hint = get_widget_category_hint(node, idx) + effective_category_hint = category_hint or widget_category_hint + categories_to_try_for_widget = ( + [effective_category_hint] if effective_category_hint else None + ) + if not is_model_filename(value): # Check for dict-type widget values containing model references (e.g. Power Lora Loader) # Some nodes store model info as objects like {"on": true, "lora": "name.safetensors", "strength": 1.0} @@ -359,7 +416,11 @@ def get_node_model_info( urn_match = URN_REGEX.match(value_str) if urn_match: base, typ, provider, model_id, version_id = urn_match.groups() - category = URN_TYPE_MAP.get(typ.lower(), "unknown") + category = ( + category_hint + or widget_category_hint + or URN_TYPE_MAP.get(typ.lower(), "unknown") + ) model_refs.append( { @@ -386,13 +447,13 @@ def get_node_model_info( continue # Existing logic for local filenames - resolved = try_resolve_model_path(value_str, categories_to_try) + resolved = try_resolve_model_path(value_str, categories_to_try_for_widget) if resolved: category, full_path = resolved exists = os.path.exists(full_path) else: - category = category_hint or "unknown" + category = effective_category_hint or "unknown" full_path = None exists = False diff --git a/web/linker.js b/web/linker.js index bd50539..5ff4cf8 100644 --- a/web/linker.js +++ b/web/linker.js @@ -1781,16 +1781,52 @@ class LinkerManagerDialog extends ComfyDialog { } } - isSourceAvailable(source) { - if (!source || ['all', 'local', 'huggingface', 'civitai'].includes(source)) { - return true; - } - return Boolean(this.capabilities?.sources?.[source]); - } - - getDownloadCategoryOptions(defaultCategory = 'checkpoints') { - const directories = this.downloadDirectories || {}; - const keys = Object.keys(directories); + isSourceAvailable(source) { + if (!source || ['all', 'local', 'huggingface', 'civitai'].includes(source)) { + return true; + } + return Boolean(this.capabilities?.sources?.[source]); + } + + getCategoryDisplayName(category = '') { + const displayNames = { + 'checkpoints': 'checkpoint', + 'loras': 'lora', + 'vae': 'vae', + 'controlnet': 'controlnet', + 'embeddings': 'embedding', + 'upscale_models': 'upscale_model', + 'latent_upscale_models': 'latent_upscale_model', + 'diffusion_models': 'unet', + 'text_encoders': 'text_encoders (also scans clip)', + 'clip': 'clip', + 'clip_vision': 'clip_vision', + 'hypernetworks': 'hypernetwork' + }; + return displayNames[category] || category || 'unknown'; + } + + getCategoryTokenName(category = '') { + const tokenNames = { + 'checkpoints': 'checkpoint', + 'loras': 'lora', + 'vae': 'vae', + 'controlnet': 'controlnet', + 'embeddings': 'embedding', + 'upscale_models': 'upscale_model', + 'latent_upscale_models': 'latent_upscale_model', + 'diffusion_models': 'unet', + 'text_encoders': 'text_encoders', + 'clip': 'clip', + 'clip_vision': 'clip_vision', + 'hypernetworks': 'hypernetwork' + }; + return tokenNames[category] || category || 'unknown'; + } + + getDownloadCategoryOptions(defaultCategory = 'checkpoints') { + const directories = this.downloadDirectories || {}; + const keys = Object.keys(directories); const preferred = defaultCategory || 'checkpoints'; const ordered = [ preferred, @@ -1915,10 +1951,10 @@ class LinkerManagerDialog extends ComfyDialog { renderDownloadTargetControls(missing, defaultCategory = 'checkpoints') { const selectId = `download-category-${missing.node_id}-${missing.widget_index}`; const subfolderId = `download-subfolder-${missing.node_id}-${missing.widget_index}`; - const subfolderListId = `download-subfolder-list-${missing.node_id}-${missing.widget_index}`; - const options = this.getDownloadCategoryOptions(defaultCategory) - .map(category => ``) - .join(''); + const subfolderListId = `download-subfolder-list-${missing.node_id}-${missing.widget_index}`; + const options = this.getDownloadCategoryOptions(defaultCategory) + .map(category => ``) + .join(''); let html = `
`; html += `
`; @@ -2453,20 +2489,7 @@ class LinkerManagerDialog extends ComfyDialog { return; } - const categoryDisplayNames = { - 'checkpoints': 'checkpoint', - 'loras': 'lora', - 'vae': 'vae', - 'controlnet': 'controlnet', - 'embeddings': 'embedding', - 'upscale_models': 'upscale_model', - 'diffusion_models': 'unet', - 'clip': 'clip', - 'clip_vision': 'clip_vision', - 'hypernetworks': 'hypernetwork' - }; - - const byCategory = {}; + const byCategory = {}; for (const model of loadedModels) { const cat = model.category || 'unknown'; @@ -2497,9 +2520,9 @@ class LinkerManagerDialog extends ComfyDialog { const inactiveCount = Object.values(byCategory).reduce((sum, cat) => sum + cat.inactive.length, 0); const buildCategoryStrings = (filter) => { - const result = {}; - for (const [category, modelsObj] of Object.entries(byCategory)) { - const displayCat = categoryDisplayNames[category] || category; + const result = {}; + for (const [category, modelsObj] of Object.entries(byCategory)) { + const displayCat = this.getCategoryTokenName(category); const models = filter === 'active' ? modelsObj.active : filter === 'inactive' ? modelsObj.inactive : [...modelsObj.active, ...modelsObj.inactive]; const parts = models.map(model => { const fullName = model.name || model.original_path?.split(/[\/\\]/).pop() || 'Unknown'; @@ -2538,8 +2561,8 @@ class LinkerManagerDialog extends ComfyDialog {
`; - for (const [category, modelsObj] of Object.entries(byCategory)) { - const displayName = categoryDisplayNames[category] || category; + for (const [category, modelsObj] of Object.entries(byCategory)) { + const displayName = this.getCategoryDisplayName(category); const hasActive = modelsObj.active.length > 0; const hasInactive = modelsObj.inactive.length > 0; @@ -2551,12 +2574,12 @@ class LinkerManagerDialog extends ComfyDialog {
`; if (hasActive) { - const activeStr = modelsObj.active.map(m => { - const fullName = m.name || m.original_path?.split(/[\/\\]/).pop() || 'Unknown'; - let name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx|gguf)$/i, ''); - const strength = m.strength !== null && m.strength !== undefined ? m.strength.toFixed(2) : '1.00'; - return `<${displayName}:${name}:${strength}>`; - }).join(' '); + const activeStr = modelsObj.active.map(m => { + const fullName = m.name || m.original_path?.split(/[\/\\]/).pop() || 'Unknown'; + let name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx|gguf)$/i, ''); + const strength = m.strength !== null && m.strength !== undefined ? m.strength.toFixed(2) : '1.00'; + return `<${this.getCategoryTokenName(category)}:${name}:${strength}>`; + }).join(' '); html += `
@@ -2577,11 +2600,11 @@ class LinkerManagerDialog extends ComfyDialog { if (hasInactive) { const inactiveStr = modelsObj.inactive.map(m => { - const fullName = m.name || m.original_path?.split(/[\/\\]/).pop() || 'Unknown'; - let name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx|gguf)$/i, ''); - const strength = m.strength !== null && m.strength !== undefined ? m.strength.toFixed(2) : '1.00'; - return `<${displayName}:${name}:${strength}>`; - }).join(' '); + const fullName = m.name || m.original_path?.split(/[\/\\]/).pop() || 'Unknown'; + let name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx|gguf)$/i, ''); + const strength = m.strength !== null && m.strength !== undefined ? m.strength.toFixed(2) : '1.00'; + return `<${this.getCategoryTokenName(category)}:${name}:${strength}>`; + }).join(' '); html += `
@@ -3763,10 +3786,10 @@ class LinkerManagerDialog extends ComfyDialog { const nodeChipClasses = missing.is_top_level !== false ? 'ml-node-chip is-locatable' : 'ml-node-chip'; const nodeChipTitle = missing.is_top_level !== false ? 'Click to locate this node in the graph' : ''; - html += `
`; - if (missing.category) { - html += `${missing.category}`; - } + html += `
`; + if (missing.category) { + html += `${this.getCategoryDisplayName(missing.category)}`; + } html += ``; if (missing.is_top_level !== false) { html += this.getLocateIconHtml();