From 37de26ce2592e6743b6c2777a6f5442d61dc2515 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Fri, 30 May 2025 20:12:38 +0800 Subject: [PATCH] Enhance Lora code update handling for browser and desktop modes. Implement broadcast support for Lora Loader nodes and improve node ID management in the workflow. --- py/routes/misc_routes.py | 44 +++++++++++++++++++++++++++--------- static/js/utils/uiHelpers.js | 42 ++++++++++++++++++---------------- web/comfyui/lora_loader.js | 30 ++++++++++++++++++++++++ 3 files changed, 85 insertions(+), 31 deletions(-) diff --git a/py/routes/misc_routes.py b/py/routes/misc_routes.py index f86b4fab..fa2a999e 100644 --- a/py/routes/misc_routes.py +++ b/py/routes/misc_routes.py @@ -818,7 +818,7 @@ class MiscRoutes: Expects a JSON body with: { - "node_ids": [123, 456], # List of node IDs to update + "node_ids": [123, 456], # Optional - List of node IDs to update (for browser mode) "lora_code": "", # The Lora code to send "mode": "append" # or "replace" - whether to append or replace existing code } @@ -826,37 +826,59 @@ class MiscRoutes: try: # Parse the request body data = await request.json() - node_ids = data.get('node_ids', []) + node_ids = data.get('node_ids') lora_code = data.get('lora_code', '') mode = data.get('mode', 'append') - if not node_ids or not lora_code: + if not lora_code: return web.json_response({ 'success': False, - 'error': 'Missing node_ids or lora_code parameter' + 'error': 'Missing lora_code parameter' }, status=400) - # Send the lora code update to each node results = [] - for node_id in node_ids: + + # Desktop mode: no specific node_ids provided + if node_ids is None: try: - # Send the message to the frontend + # Send broadcast message with id=-1 to all Lora Loader nodes PromptServer.instance.send_sync("lora_code_update", { - "id": node_id, + "id": -1, "lora_code": lora_code, "mode": mode }) results.append({ - 'node_id': node_id, + 'node_id': 'broadcast', 'success': True }) except Exception as e: - logger.error(f"Error sending lora code to node {node_id}: {e}") + logger.error(f"Error broadcasting lora code: {e}") results.append({ - 'node_id': node_id, + 'node_id': 'broadcast', 'success': False, 'error': str(e) }) + else: + # Browser mode: send to specific nodes + for node_id in node_ids: + try: + # Send the message to the frontend + PromptServer.instance.send_sync("lora_code_update", { + "id": node_id, + "lora_code": lora_code, + "mode": mode + }) + results.append({ + 'node_id': node_id, + 'success': True + }) + except Exception as e: + logger.error(f"Error sending lora code to node {node_id}: {e}") + results.append({ + 'node_id': node_id, + 'success': False, + 'error': str(e) + }) return web.json_response({ 'success': True, diff --git a/static/js/utils/uiHelpers.js b/static/js/utils/uiHelpers.js index c86e3072..15b067cf 100644 --- a/static/js/utils/uiHelpers.js +++ b/static/js/utils/uiHelpers.js @@ -362,29 +362,31 @@ export function getNSFWLevelName(level) { */ export async function sendLoraToWorkflow(loraSyntax, replaceMode = false, syntaxType = 'lora') { try { + let loraNodes = []; + let isDesktopMode = false; + // Get the current workflow from localStorage const workflowData = localStorage.getItem('workflow'); - if (!workflowData) { - showToast('No active workflow found', 'error'); - return false; - } - - // Parse the workflow JSON - const workflow = JSON.parse(workflowData); - - // Find all Lora Loader (LoraManager) nodes - const loraNodes = []; - if (workflow.nodes && Array.isArray(workflow.nodes)) { - for (const node of workflow.nodes) { - if (node.type === "Lora Loader (LoraManager)") { - loraNodes.push(node.id); + if (workflowData) { + // Web browser mode - extract node IDs from workflow + const workflow = JSON.parse(workflowData); + + // Find all Lora Loader (LoraManager) nodes + if (workflow.nodes && Array.isArray(workflow.nodes)) { + for (const node of workflow.nodes) { + if (node.type === "Lora Loader (LoraManager)") { + loraNodes.push(node.id); + } } } - } - - if (loraNodes.length === 0) { - showToast('No Lora Loader nodes found in the workflow', 'warning'); - return false; + + if (loraNodes.length === 0) { + showToast('No Lora Loader nodes found in the workflow', 'warning'); + return false; + } + } else { + // ComfyUI Desktop mode - don't specify node IDs and let backend handle it + isDesktopMode = true; } // Call the backend API to update the lora code @@ -394,7 +396,7 @@ export async function sendLoraToWorkflow(loraSyntax, replaceMode = false, syntax 'Content-Type': 'application/json' }, body: JSON.stringify({ - node_ids: loraNodes, + node_ids: isDesktopMode ? undefined : loraNodes, lora_code: loraSyntax, mode: replaceMode ? 'replace' : 'append' }) diff --git a/web/comfyui/lora_loader.js b/web/comfyui/lora_loader.js index f8ee7332..431ce6c7 100644 --- a/web/comfyui/lora_loader.js +++ b/web/comfyui/lora_loader.js @@ -49,12 +49,42 @@ app.registerExtension({ // Handle lora code updates from Python handleLoraCodeUpdate(id, loraCode, mode) { + // Handle broadcast mode (for Desktop/non-browser support) + if (id === -1) { + // Find all Lora Loader nodes in the current graph + const loraLoaderNodes = []; + for (const nodeId in app.graph._nodes_by_id) { + const node = app.graph._nodes_by_id[nodeId]; + if (node.comfyClass === "Lora Loader (LoraManager)") { + loraLoaderNodes.push(node); + } + } + + // Update each Lora Loader node found + if (loraLoaderNodes.length > 0) { + loraLoaderNodes.forEach(node => { + this.updateNodeLoraCode(node, loraCode, mode); + }); + console.log(`Updated ${loraLoaderNodes.length} Lora Loader nodes in broadcast mode`); + } else { + console.warn("No Lora Loader nodes found in the workflow for broadcast update"); + } + + return; + } + + // Standard mode - update a specific node const node = app.graph.getNodeById(+id); if (!node || node.comfyClass !== "Lora Loader (LoraManager)") { console.warn("Node not found or not a LoraLoader:", id); return; } + this.updateNodeLoraCode(node, loraCode, mode); + }, + + // Helper method to update a single node's lora code + updateNodeLoraCode(node, loraCode, mode) { // Update the input widget with new lora code const inputWidget = node.widgets[0]; if (!inputWidget) return;