feat: add latent_upscale_models category support and improve workflow model resolution

- core/downloader.py: add latent_upscale_model(s) download directory mappings - core/workflow_analyzer.py: - add node type hints for VAELoaderKJ, LatentUpscaleModelLoader, DualCLIPLoader - introduce MODEL_WIDGET_NAME_TO_CATEGORY and get_widget_category_hint() to infer model categories from widget/input metadata (best-effort fallback) - extend NESTED_MODEL_KEYS with clip_name / clip_name1 / clip_name2 - apply widget category hint inside get_node_model_info() for both plain filenames and URN-style references, matching resolution and error paths - web/linker.js: - add getCategoryDisplayName() and getCategoryTokenName() helpers (including latent_upscale_model entries) replacing duplicated inline maps - use helpers for download target category chips, model-group tokens, and summary/compact clipboard strings
2026-05-17 04:21:59 +02:00
parent 34e3d73f3f
commit 94abcdc05f
3 changed files with 138 additions and 52 deletions
@@ -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",
@@ -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

@@ -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 => `<option value="${category}" ${category === defaultCategory ? 'selected' : ''}>${category}</option>`)
-            .join('');
+        const subfolderListId = `download-subfolder-list-${missing.node_id}-${missing.widget_index}`;
+        const options = this.getDownloadCategoryOptions(defaultCategory)
+            .map(category => `<option value="${category}" ${category === defaultCategory ? 'selected' : ''}>${this.getCategoryDisplayName(category)}</option>`)
+            .join('');

        let html = `<div class="ml-download-target">`;
        html += `<div class="ml-download-target-grid">`;
@@ -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 {
            <div class="ml-models-list ml-models-list-pad">
        `;

-        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 {
            </div>`;
            
            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 += `<div class="ml-model-group">
                    <div class="ml-model-group-head">
@@ -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 += `<div>
                    <div class="ml-model-group-head">
@@ -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 += `<div class="ml-card-subtitle">`;
-        if (missing.category) {
-            html += `<span class="ml-category-chip">${missing.category}</span>`;
-        }
+        html += `<div class="ml-card-subtitle">`;
+        if (missing.category) {
+            html += `<span class="ml-category-chip" title="${missing.category}">${this.getCategoryDisplayName(missing.category)}</span>`;
+        }
        html += `<span id="${locateId}" class="${nodeChipClasses}"${nodeChipTitle ? ` title="${nodeChipTitle}"` : ''}>`;
        if (missing.is_top_level !== false) {
            html += this.getLocateIconHtml();