diff --git a/__init__.py b/__init__.py index 2073ef0..7299882 100644 --- a/__init__.py +++ b/__init__.py @@ -6,6 +6,7 @@ @description: Extension for relinking missing models and downloading from HuggingFace/CivitAI """ +import asyncio import logging from .core.log_system.log_funcs import ( log_debug, @@ -31,6 +32,7 @@ class ModelLinkerExtension: def __init__(self): self.routes_setup = False self.logger = logging.getLogger(__name__) + self.analysis_progress = {} def initialize(self): """Initialize the extension and set up API routes.""" @@ -118,14 +120,41 @@ class ModelLinkerExtension: try: data = await request.json() workflow_json = data.get("workflow") + analysis_id = str(data.get("analysis_id") or "").strip() if not workflow_json: return web.json_response( {"error": "Workflow JSON is required"}, status=400 ) + if analysis_id: + self.analysis_progress[analysis_id] = { + "status": "starting", + "stage": "starting", + "message": "Starting analysis...", + "current": 0, + "total": 0, + } + + def update_analysis_progress(payload): + if not analysis_id: + return + self.analysis_progress[analysis_id] = { + **self.analysis_progress.get(analysis_id, {}), + **payload, + "status": "running" + if payload.get("stage") != "completed" + else "completed", + } + # Analyze and find matches - result = analyze_and_find_matches(workflow_json) + result = await asyncio.to_thread( + analyze_and_find_matches, + workflow_json, + 0.0, + 10, + update_analysis_progress if analysis_id else None, + ) # Filter out LoraManager lorAs that already exist locally (exists=True) # These should not appear in missing models at all @@ -254,11 +283,52 @@ class ModelLinkerExtension: # now done on-demand via /model_linker/search endpoint # when user clicks "Search Online" button, not automatically + if analysis_id: + self.analysis_progress[analysis_id] = { + **self.analysis_progress.get(analysis_id, {}), + "status": "completed", + "stage": "completed", + "message": "Analysis complete", + "current": result.get("total_missing", 0), + "total": result.get("total_missing", 0), + } + return web.json_response(result) except Exception as e: + if "analysis_id" in locals() and analysis_id: + self.analysis_progress[analysis_id] = { + "status": "error", + "stage": "error", + "message": str(e), + "current": 0, + "total": 0, + } self.logger.error(f"Model Linker analyze error: {e}", exc_info=True) return web.json_response({"error": str(e)}, status=500) + @routes.get("/model_linker/analyze-progress/{analysis_id}") + async def get_analyze_progress(request): + """Get workflow analysis progress.""" + analysis_id = request.match_info.get("analysis_id", "").strip() + if not analysis_id: + return web.json_response( + {"error": "Analysis ID is required"}, status=400 + ) + + progress = self.analysis_progress.get(analysis_id) + if not progress: + return web.json_response( + { + "status": "unknown", + "stage": "unknown", + "message": "No analysis progress available", + "current": 0, + "total": 0, + } + ) + + return web.json_response(progress) + @routes.post("/model_linker/resolve") async def resolve_models(request): """Apply model resolution and return updated workflow.""" diff --git a/core/linker.py b/core/linker.py index 3529a66..8356fec 100644 --- a/core/linker.py +++ b/core/linker.py @@ -7,7 +7,7 @@ Integrates all components to provide high-level API for model linking. import os import re import json -from typing import Dict, Any, List, Optional, Tuple +from typing import Dict, Any, List, Optional, Tuple, Callable from urllib.parse import unquote from .log_system.log_funcs import ( @@ -227,6 +227,7 @@ def analyze_and_find_matches( workflow_json: Dict[str, Any], similarity_threshold: float = 0.0, max_matches_per_model: int = 10, + progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, ) -> Dict[str, Any]: """ Main entry point: analyze workflow and find matches for missing models. @@ -264,10 +265,30 @@ def analyze_and_find_matches( 'total_models_analyzed': count of all models in workflow } """ + if progress_callback: + progress_callback( + { + "stage": "extracting", + "message": "Extracting workflow model references...", + "current": 0, + "total": 0, + } + ) + # Extract URLs from workflow (node.properties.models + regex) workflow_urls = extract_workflow_urls(workflow_json) log_debug(f"Extracted {len(workflow_urls)} URLs from workflow") + if progress_callback: + progress_callback( + { + "stage": "scanning", + "message": "Scanning local model index...", + "current": 0, + "total": 0, + } + ) + # Analyze workflow to find all model references # Get available models available_models = get_model_files() @@ -280,11 +301,31 @@ def analyze_and_find_matches( ordered_candidates_cache: Dict[str, List[Dict[str, Any]]] = {} + if progress_callback: + progress_callback( + { + "stage": "analyzing", + "message": "Analyzing workflow nodes...", + "current": 0, + "total": 0, + } + ) + # Analyze workflow using the same already-scanned model list all_model_refs = analyze_workflow_models( workflow_json, available_models=available_models ) + if progress_callback: + progress_callback( + { + "stage": "identifying", + "message": "Identifying missing models...", + "current": 0, + "total": len(all_model_refs), + } + ) + # Identify missing models missing_models = identify_missing_models(all_model_refs, available_models) @@ -310,9 +351,32 @@ def analyze_and_find_matches( missing["urn_version_id"] = urn.get("version_id") missing["urn_type"] = urn.get("type", "") + if progress_callback: + progress_callback( + { + "stage": "matching", + "message": "Matching missing models...", + "current": 0, + "total": len(missing_models), + } + ) + # Find matches for each missing model missing_with_matches = [] - for missing in missing_models: + total_missing = len(missing_models) + for index, missing in enumerate(missing_models, start=1): + if progress_callback: + progress_callback( + { + "stage": "matching", + "message": f"Analyzing model {index} of {total_missing}", + "current": index, + "total": total_missing, + "model_name": missing.get("name") + or missing.get("original_path", ""), + } + ) + # Skip LoraManager lorAs that already exist locally (exists=True means no linking needed) is_lora_v2 = missing.get("is_lora_v2") exists = missing.get("exists") @@ -390,12 +454,24 @@ def analyze_and_find_matches( missing_with_matches.append({**missing, "matches": deduplicated_matches}) - return { + result = { "missing_models": missing_with_matches, "total_missing": len(missing_with_matches), "total_models_analyzed": len(all_model_refs), } + if progress_callback: + progress_callback( + { + "stage": "completed", + "message": "Analysis complete", + "current": total_missing, + "total": total_missing, + } + ) + + return result + def apply_resolution( workflow_json: Dict[str, Any], resolutions: List[Dict[str, Any]] diff --git a/web/linker.js b/web/linker.js index 233d223..b361788 100644 --- a/web/linker.js +++ b/web/linker.js @@ -27,6 +27,7 @@ class LinkerManagerDialog extends ComfyDialog { this.fullscreen = false; this._dragging = false; this._dragStart = null; + this._analysisProgressToken = null; // Inject global styles for the redesigned UI this.injectStyles(); @@ -192,16 +193,9 @@ class LinkerManagerDialog extends ComfyDialog { if (!container) return; const state = this.getSearchState(missing); - const buttons = container.querySelectorAll('.ml-search-source-btn'); - buttons.forEach(btn => { - const isActive = btn.dataset.searchSource === state.selectedSource; - btn.classList.toggle('ml-btn-primary', isActive); - btn.classList.toggle('ml-btn-secondary', !isActive); - }); - - const infoEl = container.querySelector(`#search-source-info-${missing.node_id}-${missing.widget_index}`); - if (infoEl) { - infoEl.textContent = `Selected source: ${this.getSearchSourceLabel(state.selectedSource)}`; + const selectEl = container.querySelector(`#search-source-select-${missing.node_id}-${missing.widget_index}`); + if (selectEl) { + selectEl.value = state.selectedSource; } } @@ -1781,10 +1775,37 @@ class LinkerManagerDialog extends ComfyDialog { } .ml-search-source-bar { display: flex; - flex-wrap: wrap; - gap: 5px; + align-items: center; + justify-content: space-between; + gap: 8px; margin-bottom: 8px; } + .ml-search-source-picker { + display: flex; + align-items: center; + gap: 8px; + min-width: 0; + margin-left: auto; + } + .ml-search-source-picker-label { + color: var(--ml-text-muted); + font-size: 11px; + white-space: nowrap; + } + .ml-search-source-select { + min-height: 30px; + padding: 5px 10px; + background: rgba(255,255,255,0.03); + color: var(--ml-text); + border: 1px solid rgba(255,255,255,0.1); + border-radius: 10px; + font-size: 11px; + outline: none; + } + .ml-search-source-select:focus { + border-color: rgba(255,255,255,0.16); + box-shadow: inset 0 1px 0 rgba(255,255,255,0.02); + } .ml-search-results { margin-top: 8px; display: none; @@ -2910,6 +2931,45 @@ class LinkerManagerDialog extends ComfyDialog { `; } + + renderAnalysisProgress(progress = {}) { + const current = Number(progress.current) || 0; + const total = Number(progress.total) || 0; + const percent = total > 0 ? Math.max(0, Math.min(100, Math.round((current / total) * 100))) : 0; + const message = progress.message || 'Analyzing workflow...'; + const modelName = progress.model_name ? this.escapeHtml(String(progress.model_name)) : ''; + const detail = total > 0 ? `${current} / ${total}` : 'Preparing...'; + + return ` +
+
+ ${this.getStatusBadge('Analyzing', 'info')} + ${message} +
+ ${this.renderProgressBar(percent, detail, `${percent}%`)} + ${modelName ? `
${modelName}
` : ''} +
+ `; + } + + async pollAnalysisProgress(analysisId, token) { + while (this._analysisProgressToken === token) { + try { + const response = await api.fetchApi(`/model_linker/analyze-progress/${analysisId}`); + if (response.ok && this.contentElement && this._analysisProgressToken === token) { + const progress = await response.json(); + this.contentElement.innerHTML = this.renderAnalysisProgress(progress); + if (progress.status === 'completed' || progress.status === 'error') { + return; + } + } + } catch (error) { + console.warn('Model Linker: analysis progress polling failed', error); + } + + await new Promise(resolve => setTimeout(resolve, 250)); + } + } /** * Handle clicks outside the dialog @@ -4410,8 +4470,16 @@ class LinkerManagerDialog extends ComfyDialog { async loadWorkflowData(workflow = null) { if (!this.contentElement) return; + const analysisId = `an-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + this._analysisProgressToken = analysisId; + // Show loading state - this.contentElement.innerHTML = '

Analyzing workflow...

'; + this.contentElement.innerHTML = this.renderAnalysisProgress({ + status: 'starting', + message: 'Starting analysis...', + current: 0, + total: 0 + }); try { // Use provided workflow, or get current workflow from ComfyUI @@ -4420,16 +4488,20 @@ class LinkerManagerDialog extends ComfyDialog { } if (!workflow) { + this._analysisProgressToken = null; this.contentElement.innerHTML = '

No workflow loaded. Please load a workflow first.

'; return; } // Call analyze endpoint + const progressPromise = this.pollAnalysisProgress(analysisId, analysisId); const response = await api.fetchApi('/model_linker/analyze', { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ workflow }) + body: JSON.stringify({ workflow, analysis_id: analysisId }) }); + this._analysisProgressToken = null; + await progressPromise; if (!response.ok) { throw new Error(`API error: ${response.status}`); @@ -4443,6 +4515,7 @@ class LinkerManagerDialog extends ComfyDialog { this.reconnectActiveDownloads(); } catch (error) { + this._analysisProgressToken = null; console.error('Model Linker: Error loading workflow data:', error); if (this.contentElement) { this.contentElement.innerHTML = `

Error: ${error.message}

`; @@ -4687,13 +4760,11 @@ class LinkerManagerDialog extends ComfyDialog { }); } - const sourceButtons = container.querySelectorAll(`#search-sources-${missing.node_id}-${missing.widget_index} .ml-search-source-btn`); - sourceButtons.forEach(btn => { - btn.addEventListener('click', () => { - this.setSearchSource(missing, btn.dataset.searchSource, container); + const sourceSelect = container.querySelector(`#search-source-select-${missing.node_id}-${missing.widget_index}`); + if (sourceSelect) { + sourceSelect.addEventListener('change', () => { + this.setSearchSource(missing, sourceSelect.value, container); }); - }); - if (sourceButtons.length > 0) { this.syncSearchSourceUi(missing, container); } @@ -4943,17 +5014,21 @@ class LinkerManagerDialog extends ComfyDialog { // No known download - offer search html += `
`; const searchSourcesId = `search-sources-${missing.node_id}-${missing.widget_index}`; - const searchInfoId = `search-source-info-${missing.node_id}-${missing.widget_index}`; + const searchSourceSelectId = `search-source-select-${missing.node_id}-${missing.widget_index}`; html += `
`; - html += ``; - html += ``; - html += ``; - html += ``; - html += `
`; html += ``; - html += `
Selected source: Everything
`; + html += `
`; + html += ``; + html += ``; + html += `
`; + html += `
`; html += this.renderDownloadTargetControls(missing, missing.category || 'checkpoints'); html += ``; html += `
`;