diff --git a/__init__.py b/__init__.py index 0973d58..b45edd2 100644 --- a/__init__.py +++ b/__init__.py @@ -293,6 +293,148 @@ class ModelLinkerExtension: ) return web.json_response({"error": str(e)}, status=500) + @routes.post("/model_linker/loaded") + async def get_loaded_models(request): + """Get all currently loaded models in the workflow.""" + try: + data = await request.json() + workflow_json = data.get("workflow") + + if not workflow_json: + return web.json_response( + {"error": "Workflow JSON is required"}, status=400 + ) + + # Import workflow analyzer to extract models + from .core.workflow_analyzer import ( + analyze_workflow_models, + try_resolve_model_path, + is_model_filename, + URN_REGEX, + URN_TYPE_MAP, + ) + + # Get available models for existence checking + available_models = get_model_files() + available_paths = {m.get("path") for m in available_models} + + # Analyze workflow to get all model references + all_model_refs = analyze_workflow_models(workflow_json) + + # Also extract from node.properties.models + nodes = list(workflow_json.get("nodes", [])) + definitions = workflow_json.get("definitions", {}) + subgraphs = definitions.get("subgraphs", []) + for subgraph in subgraphs: + nodes.extend(subgraph.get("nodes", [])) + + # Collect all loaded models with their values + loaded_models = [] + + # Process each model reference from analyze_workflow_models + for ref in all_model_refs: + original_path = ref.get("original_path", "") + node_id = ref.get("node_id") + widget_index = ref.get("widget_index") + node_type = ref.get("node_type", "") + category = ref.get("category", "unknown") + + # Determine model name and strength + model_name = original_path.split("/")[-1].split("\\")[-1] + strength = None + + # For LoraLoader nodes, strength is in next widget_value + if node_type in ["LoraLoader", "LoraLoaderModelOnly"]: + # Find the node in workflow to get strength value + for node in nodes: + if str(node.get("id")) == str(node_id): + widgets_values = node.get("widgets_values", []) + if len(widgets_values) > widget_index + 1: + try: + strength = float( + widgets_values[widget_index + 1] + ) + except (ValueError, TypeError): + strength = 1.0 + break + + # Check if model exists locally + exists = ref.get("exists", False) + + # If URN, resolve to display name + if ref.get("is_urn"): + urn = ref.get("urn", {}) + # Use model name from URN as display name + model_name = ( + f"urn:{urn.get('type', 'model')}:{urn.get('model_id')}" + ) + category = urn.get("type", category) + if category in URN_TYPE_MAP: + category = URN_TYPE_MAP[category] + + loaded_models.append( + { + "name": model_name, + "category": category, + "node_id": node_id, + "widget_index": widget_index, + "node_type": node_type, + "exists": exists, + "strength": strength, + "original_path": original_path, + "is_urn": ref.get("is_urn", False), + } + ) + + # Also check node.properties.models for embedded models + for node in nodes: + node_type = node.get("type", "") + properties = node.get("properties", {}) + models_list = properties.get("models", []) + + for model_info in models_list: + if isinstance(model_info, dict): + name = model_info.get("name", "") + url = model_info.get("url", "") + directory = model_info.get("directory", "") + + if name: + # Check if this model is already in loaded_models + existing = next( + ( + m + for m in loaded_models + if m.get("original_path") == name + ), + None, + ) + if not existing: + loaded_models.append( + { + "name": name.split("/")[-1].split("\\")[ + -1 + ], + "category": directory or "checkpoints", + "node_id": node.get("id"), + "widget_index": None, + "node_type": node_type, + "exists": True, # Embedded models are loaded + "strength": None, + "original_path": name, + "is_urn": False, + } + ) + + return web.json_response( + {"loaded_models": loaded_models, "total": len(loaded_models)} + ) + + except Exception as e: + self.logger.error( + f"Model Linker get_loaded_models error: {e}", exc_info=True + ) + return web.json_response({"error": str(e)}, status=500) + # ==================== DOWNLOAD ROUTES ==================== if download_available: diff --git a/web/linker.js b/web/linker.js index 139ff04..6bbe9fd 100644 --- a/web/linker.js +++ b/web/linker.js @@ -17,6 +17,7 @@ class LinkerManagerDialog extends ComfyDialog { this.missingModels = []; this.activeDownloads = {}; // Track active downloads this.boundHandleOutsideClick = this.handleOutsideClick.bind(this); + this.activeTab = 'missing'; // Default tab // Inject global styles for the redesigned UI this.injectStyles(); @@ -419,6 +420,87 @@ class LinkerManagerDialog extends ComfyDialog { background: linear-gradient(to top, var(--ml-bg) 0%, var(--ml-bg) 70%, transparent 100%); border-top: 1px solid var(--ml-border); } + + /* Tabs */ + .ml-tabs { + display: flex; + gap: 4px; + padding: 0 20px; + margin-bottom: 0; + background-color: var(--ml-bg, #222); + border-bottom: 1px solid var(--ml-border); + } + .ml-tab { + padding: 12px 20px; + background: transparent; + border: none; + color: var(--ml-text-muted, #888); + font-size: 13px; + font-weight: 500; + cursor: pointer; + border-bottom: 2px solid transparent; + transition: all 0.2s ease; + } + .ml-tab:hover { + color: var(--ml-text, #e0e0e0); + background: rgba(255,255,255,0.05); + } + .ml-tab.ml-tab-active { + color: var(--ml-accent, #4CAF50); + border-bottom-color: var(--ml-accent, #4CAF50); + } + + /* Loaded Models Tab */ + .ml-loaded-models { + padding: 0; + } + .ml-loaded-models-header { + padding: 16px 20px; + border-bottom: 1px solid var(--ml-border); + } + .ml-loaded-models-title { + font-size: 15px; + font-weight: 600; + color: var(--ml-text); + margin: 0 0 8px 0; + } + .ml-loaded-models-subtitle { + font-size: 12px; + color: var(--ml-text-muted); + } + .ml-models-list { + padding: 12px 20px; + display: flex; + flex-direction: column; + gap: 8px; + } + .ml-model-chip { + display: inline-flex; + align-items: center; + padding: 6px 12px; + background: #3a3a3a; + border-radius: 6px; + font-family: 'SF Mono', 'Consolas', 'Monaco', monospace; + font-size: 12px; + color: var(--ml-text); + } + .ml-model-chip-category { + padding: 2px 6px; + background: #555; + border-radius: 4px; + font-size: 10px; + color: #aaa; + margin-left: 8px; + text-transform: uppercase; + } + .ml-model-chip-strength { + padding: 2px 6px; + background: #4CAF50; + border-radius: 4px; + font-size: 10px; + color: white; + margin-left: 8px; + } /* No matches text */ .ml-no-matches { @@ -583,45 +665,223 @@ class LinkerManagerDialog extends ComfyDialog { } createHeader() { + // Create tabs + this.missingTab = $el("button.ml-tab.ml-tab-active", { + textContent: "Missing Models", + onclick: () => this.switchTab('missing') + }); + + this.loadedTab = $el("button.ml-tab", { + textContent: "Loaded Models", + onclick: () => this.switchTab('loaded') + }); + return $el("div", { style: { display: "flex", - justifyContent: "space-between", - alignItems: "center", - padding: "20px 20px 10px 20px", - borderBottom: "1px solid var(--border-color)", - backgroundColor: "var(--comfy-menu-bg, #202020)" + flexDirection: "column" } }, [ - $el("h2", { - textContent: "🔗 Model Linker", + $el("div", { style: { - margin: "0", - color: "var(--input-text)", - fontSize: "18px", - fontWeight: "600" - } - }), - $el("button", { - textContent: "×", - onclick: () => this.close(), - style: { - background: "none", - border: "none", - fontSize: "24px", - cursor: "pointer", - color: "var(--input-text)", - padding: "0", - width: "30px", - height: "30px", - borderRadius: "4px", display: "flex", + justifyContent: "space-between", alignItems: "center", - justifyContent: "center" + padding: "20px 20px 10px 20px", + borderBottom: "1px solid var(--border-color)", + backgroundColor: "var(--comfy-menu-bg, #202020)" } - }) + }, [ + $el("h2", { + textContent: "🔗 Model Linker", + style: { + margin: "0", + color: "var(--input-text)", + fontSize: "18px", + fontWeight: "600" + } + }), + $el("button", { + textContent: "×", + onclick: () => this.close(), + style: { + background: "none", + border: "none", + fontSize: "24px", + cursor: "pointer", + color: "var(--input-text)", + padding: "0", + width: "30px", + height: "30px", + borderRadius: "4px", + display: "flex", + alignItems: "center", + justifyContent: "center" + } + }) + ]), + $el("div.ml-tabs", {}, [ + this.missingTab, + this.loadedTab + ]) ]); } + + switchTab(tab) { + this.activeTab = tab; + + if (tab === 'missing') { + this.missingTab.classList.add('ml-tab-active'); + this.loadedTab.classList.remove('ml-tab-active'); + this.downloadAllButton.style.display = 'inline-flex'; + this.autoResolveButton.style.display = 'inline-flex'; + this.loadWorkflowData(); + } else { + this.missingTab.classList.remove('ml-tab-active'); + this.loadedTab.classList.add('ml-tab-active'); + this.downloadAllButton.style.display = 'none'; + this.autoResolveButton.style.display = 'none'; + this.loadLoadedModels(); + } + } + + async loadLoadedModels() { + if (!this.contentElement) return; + + this.contentElement.innerHTML = '
Loading loaded models...
'; + + try { + const workflow = this.getCurrentWorkflow(); + if (!workflow) { + this.contentElement.innerHTML = 'No workflow loaded. Please load a workflow first.
'; + return; + } + + const response = await api.fetchApi('/model_linker/loaded', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ workflow }) + }); + + if (!response.ok) { + throw new Error(`API error: ${response.status}`); + } + + const data = await response.json(); + this.displayLoadedModels(this.contentElement, data); + + } catch (error) { + console.error('Model Linker: Error loading loaded models:', error); + if (this.contentElement) { + this.contentElement.innerHTML = `Error: ${error.message}
`; + } + } + } + + displayLoadedModels(container, data) { + const loadedModels = data.loaded_models || []; + const total = data.total || 0; + + if (total === 0) { + container.innerHTML = this.renderStatusMessage('No models found in workflow.', 'info'); + return; + } + + // Category display name mapping (singular form) + 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' + }; + + // Group by category + const byCategory = {}; + for (const model of loadedModels) { + const cat = model.category || 'unknown'; + if (!byCategory[cat]) { + byCategory[cat] = []; + } + byCategory[cat].push(model); + } + + // Build display string in format:Models currently used in the workflow
+${allParts}
+
+