From a7f51fc4785d1786cd44c4402c6df4a2f68a3c9a Mon Sep 17 00:00:00 2001
From: robertvoy <hello@robertvoy.com>
Date: Thu, 10 Jul 2025 13:14:32 +1000
Subject: [PATCH] Fix worker image saving conflicts

---
 web/gpupanel.js | 79 +++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 77 insertions(+), 2 deletions(-)

diff --git a/web/gpupanel.js b/web/gpupanel.js
index e2acb13..2d2d4ec 100644
--- a/web/gpupanel.js
+++ b/web/gpupanel.js
@@ -1036,6 +1036,39 @@ class DistributedExtension {
         return connected;
     }
 
+    /**
+     * Find only upstream nodes (inputs) for distributed collector nodes
+     * This is used for workers to avoid executing downstream nodes like SaveImage
+     * @param {Object} apiPrompt - The API prompt containing the workflow
+     * @param {Array<string>} collectorIds - Array of collector node IDs
+     * @returns {Set<string>} Set of node IDs that feed into collectors
+     */
+    findCollectorUpstreamNodes(apiPrompt, collectorIds) {
+        const connected = new Set(collectorIds); // Include all collectors
+        const toProcess = [...collectorIds];
+        
+        // Only traverse upstream (inputs)
+        while (toProcess.length > 0) {
+            const nodeId = toProcess.pop();
+            const node = apiPrompt[nodeId];
+            
+            // Traverse upstream (inputs) only
+            if (node && node.inputs) {
+                for (const [inputName, inputValue] of Object.entries(node.inputs)) {
+                    if (Array.isArray(inputValue) && inputValue.length === 2) {
+                        const sourceNodeId = String(inputValue[0]);
+                        if (!connected.has(sourceNodeId)) {
+                            connected.add(sourceNodeId);
+                            toProcess.push(sourceNodeId);
+                        }
+                    }
+                }
+            }
+        }
+        
+        return connected;
+    }
+
     /**
      * Prune workflow to only include nodes connected to distributed nodes
      * @param {Object} apiPrompt - The full workflow API prompt
@@ -1055,9 +1088,11 @@ class DistributedExtension {
             return apiPrompt;
         }
         
-        // Get all nodes connected to distributed nodes (both upstream and downstream)
+        // Get all nodes connected to distributed nodes
         const distributedIds = distributedNodes.map(node => node.id);
-        const connectedNodes = this.findCollectorConnectedNodes(apiPrompt, distributedIds);
+        
+        // For workers, only include upstream nodes to prevent downstream execution
+        const connectedNodes = this.findCollectorUpstreamNodes(apiPrompt, distributedIds);
         
         this.debugLog(`Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes`);
         
@@ -1067,6 +1102,46 @@ class DistributedExtension {
             prunedPrompt[nodeId] = JSON.parse(JSON.stringify(apiPrompt[nodeId]));
         }
         
+        // Check if any distributed node has downstream SaveImage nodes that were removed
+        // If so, add a PreviewImage node after the collector
+        for (const distNode of distributedNodes) {
+            const distNodeId = distNode.id;
+            
+            // Check if this distributed node had any downstream nodes in the original workflow
+            const originalOutputMap = new Map();
+            for (const [nodeId, node] of Object.entries(apiPrompt)) {
+                if (node.inputs) {
+                    for (const [inputName, inputValue] of Object.entries(node.inputs)) {
+                        if (Array.isArray(inputValue) && inputValue.length === 2 && String(inputValue[0]) === distNodeId) {
+                            if (!originalOutputMap.has(distNodeId)) {
+                                originalOutputMap.set(distNodeId, []);
+                            }
+                            originalOutputMap.get(distNodeId).push({nodeId, inputName});
+                        }
+                    }
+                }
+            }
+            
+            // If this distributed node had downstream nodes that were removed, add a PreviewImage
+            if (originalOutputMap.has(distNodeId) && originalOutputMap.get(distNodeId).length > 0) {
+                // Generate a unique ID for the preview node
+                const previewNodeId = `preview_${distNodeId}`;
+                
+                // Add PreviewImage node connected to the distributed node
+                prunedPrompt[previewNodeId] = {
+                    inputs: {
+                        images: [distNodeId, 0]  // Connect to first output of distributed node
+                    },
+                    class_type: "PreviewImage",
+                    _meta: {
+                        title: "Preview Image (auto-added)"
+                    }
+                };
+                
+                this.debugLog(`Added PreviewImage node ${previewNodeId} after distributed node ${distNodeId} for worker`);
+            }
+        }
+        
         return prunedPrompt;
     }