Add analysis progress reporting & UI polling

Add real-time analysis progress reporting and client-side polling during workflow analysis. __init__.py: import asyncio, track per-analysis progress in self.analysis_progress, pass an update callback to analyze endpoint (using asyncio.to_thread), set progress states on success/error, and expose a new GET /model_linker/analyze-progress/{analysis_id} endpoint. core/linker.py: add optional progress_callback parameter to analyze_and_find_matches and emit staged progress updates (extracting, scanning, analyzing, identifying, matching, per-model progress and completed). web/linker.js: introduce analysis progress UI rendering, polling loop (pollAnalysisProgress) keyed by an analysis token, send analysis_id with analyze request, show progress bar/message, and replace search-source buttons with a compact select picker. These changes improve UX by providing live feedback during potentially long analysis runs.
2026-04-25 02:38:37 +02:00
parent ff75374574
commit 391b05949a
3 changed files with 252 additions and 31 deletions
@@ -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."""
@@ -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]]
@@ -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 {
            </div>
        `;
    }
+
+    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 `
+            <div class="ml-download-section">
+                <div class="ml-status-inline">
+                    ${this.getStatusBadge('Analyzing', 'info')}
+                    <span class="ml-download-info">${message}</span>
+                </div>
+                ${this.renderProgressBar(percent, detail, `${percent}%`)}
+                ${modelName ? `<div class="ml-download-info">${modelName}</div>` : ''}
+            </div>
+        `;
+    }
+
+    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 = '<p>Analyzing workflow...</p>';
+        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 = '<p>No workflow loaded. Please load a workflow first.</p>';
                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 = `<p style="color: red;">Error: ${error.message}</p>`;
@@ -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 += `<div class="ml-download-section">`;
            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 += `<div id="${searchSourcesId}" class="ml-search-source-bar">`;
-            html += `<button type="button" class="ml-btn ml-btn-secondary ml-btn-sm ml-search-source-btn" data-search-source="local">Local DB</button>`;
-            html += `<button type="button" class="ml-btn ml-btn-secondary ml-btn-sm ml-search-source-btn" data-search-source="huggingface">HuggingFace</button>`;
-            html += `<button type="button" class="ml-btn ml-btn-secondary ml-btn-sm ml-search-source-btn" data-search-source="civitai">CivitAI</button>`;
-            html += `<button type="button" class="ml-btn ml-btn-secondary ml-btn-sm ml-search-source-btn" data-search-source="all">All</button>`;
-            html += `</div>`;
            html += `<button id="search-${missing.node_id}-${missing.widget_index}" class="ml-btn ml-btn-link">`;
            html += `<span class="ml-btn-icon">🔍</span> Search Online`;
            html += `</button>`;
-            html += `<div id="${searchInfoId}" class="ml-download-info">Selected source: Everything</div>`;
+            html += `<div class="ml-search-source-picker">`;
+            html += `<label class="ml-search-source-picker-label" for="${searchSourceSelectId}">Source</label>`;
+            html += `<select id="${searchSourceSelectId}" class="ml-search-source-select">`;
+            html += `<option value="all">Everything</option>`;
+            html += `<option value="local">Local Database</option>`;
+            html += `<option value="huggingface">HuggingFace</option>`;
+            html += `<option value="civitai">CivitAI</option>`;
+            html += `</select>`;
+            html += `</div>`;
+            html += `</div>`;
            html += this.renderDownloadTargetControls(missing, missing.category || 'checkpoints');
            html += `</div>`;
            html += `<div id="search-results-${missing.node_id}-${missing.widget_index}" class="ml-search-results"></div>`;