Clean up workflow sent to workers

__init__.py

@@ -0,0 +1,23 @@
+import os
+import sys
+
+# Add the directory to Python path
+sys.path.append(os.path.dirname(__file__))
+
+# Import everything needed from the main module
+from .multigpu import (
+    NODE_CLASS_MAPPINGS, 
+    NODE_DISPLAY_NAME_MAPPINGS, 
+    ensure_config_exists,
+    CONFIG_FILE
+)
+
+WEB_DIRECTORY = "./web"
+
+# Use the centralized config management from multigpu.py
+ensure_config_exists()
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
+
+print("[MultiGPU] Loaded Multi-GPU nodes.")
+print(f"[MultiGPU] Config file: {CONFIG_FILE}")

gpu_config.json

@@ -0,0 +1,34 @@
+{
+  "master": {
+    "port": 8080,
+    "cuda_device": 0,
+    "extra_args": "--mmap-torch-files --listen"
+  },
+  "workers": [
+    {
+      "id": 1,
+      "name": "GPU 1",
+      "cuda_device": 1,
+      "port": 8180,
+      "enabled": true,
+      "extra_args": "--mmap-torch-files --listen"
+    },
+    {
+      "id": 2,
+      "name": "GPU 2",
+      "cuda_device": 2,
+      "port": 8280,
+      "enabled": true,
+      "extra_args": "--mmap-torch-files --listen"
+    },
+    {
+      "id": 3,
+      "name": "GPU 3",
+      "cuda_device": 3,
+      "port": 8380,
+      "enabled": false,
+      "extra_args": "--mmap-torch-files --listen --reserve-vram 2"
+    }
+  ],
+  "settings": {}
+}

multigpu.py

@@ -0,0 +1,438 @@
+import torch
+import numpy as np
+from PIL import Image
+import folder_paths
+import os
+import json
+import asyncio
+import aiohttp
+from aiohttp import web
+import io
+import server
+from concurrent.futures import Future
+
+# --- Config Management ---
+CONFIG_FILE = os.path.join(os.path.dirname(__file__), "gpu_config.json")
+
+def get_default_config():
+    """Returns the default configuration dictionary. Single source of truth."""
+    return {
+        "master": {"port": 8188, "cuda_device": 0},
+        "workers": [],
+        "settings": {}
+    }
+
+def load_config():
+    """Loads the config, falling back to defaults if the file is missing or invalid."""
+    if os.path.exists(CONFIG_FILE):
+        try:
+            with open(CONFIG_FILE, 'r') as f:
+                return json.load(f)
+        except Exception as e:
+            print(f"[MultiGPU Config] Error loading config, using defaults: {e}")
+    return get_default_config()
+
+def save_config(config):
+    """Saves the configuration to file."""
+    try:
+        with open(CONFIG_FILE, 'w') as f:
+            json.dump(config, f, indent=2)
+        return True
+    except Exception as e:
+        print(f"[MultiGPU Config] Error saving config: {e}")
+        return False
+
+def ensure_config_exists():
+    """Creates default config file if it doesn't exist. Used by __init__.py"""
+    if not os.path.exists(CONFIG_FILE):
+        default_config = get_default_config()
+        if save_config(default_config):
+            print("[MultiGPU] Created default config file")
+        else:
+            print("[MultiGPU] Could not create default config file")
+
+# --- Helper Functions ---
+async def handle_api_error(request, error, status=500):
+    """Standardized error response handler"""
+    print(f"[MultiGPU] API Error: {error}")
+    return web.json_response({"status": "error", "message": str(error)}, status=status)
+
+# --- API Endpoints ---
+@server.PromptServer.instance.routes.get("/multigpu/config")
+async def get_config_endpoint(request):
+    config = load_config()
+    return web.json_response(config)
+
+@server.PromptServer.instance.routes.post("/multigpu/config/update_worker")
+async def update_worker_endpoint(request):
+    try:
+        data = await request.json()
+        worker_id = data.get("worker_id")
+        enabled = data.get("enabled", False)
+        
+        if worker_id is None:
+            return await handle_api_error(request, "Missing worker_id", 400)
+            
+        config = load_config()
+        worker_found = False
+        
+        for worker in config.get("workers", []):
+            if worker["id"] == worker_id:
+                worker["enabled"] = enabled
+                worker_found = True
+                break
+                
+        if not worker_found:
+            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
+            
+        if save_config(config):
+            return web.json_response({"status": "success"})
+        else:
+            return await handle_api_error(request, "Failed to save config")
+    except Exception as e:
+        return await handle_api_error(request, e, 400)
+
+@server.PromptServer.instance.routes.post("/multigpu/config/update_setting")
+async def update_setting_endpoint(request):
+    """Updates a specific key in the settings object."""
+    try:
+        data = await request.json()
+        key = data.get("key")
+        value = data.get("value")
+
+        if not key or value is None:
+            return await handle_api_error(request, "Missing 'key' or 'value' in request", 400)
+
+        config = load_config()
+        if 'settings' not in config:
+            config['settings'] = {}
+        
+        config['settings'][key] = value
+
+        if save_config(config):
+            return web.json_response({"status": "success", "message": f"Setting '{key}' updated."})
+        else:
+            return await handle_api_error(request, "Failed to save config")
+    except Exception as e:
+        return await handle_api_error(request, e, 400)
+
+
+@server.PromptServer.instance.routes.post("/multigpu/prepare_job")
+async def prepare_job_endpoint(request):
+    try:
+        data = await request.json()
+        multi_job_id = data.get('multi_job_id')
+        if not multi_job_id:
+            return await handle_api_error(request, "Missing multi_job_id", 400)
+
+        async with PENDING_JOBS_LOCK:
+            if multi_job_id not in PENDING_JOBS:
+                PENDING_JOBS[multi_job_id] = asyncio.Queue()
+        
+        print(f"[MultiGPU] Prepared queue for job {multi_job_id}")
+        return web.json_response({"status": "success"})
+    except Exception as e:
+        return await handle_api_error(request, e)
+
+@server.PromptServer.instance.routes.post("/multigpu/clear_memory")
+async def clear_memory_endpoint(request):
+    print("[MultiGPU] Received request to clear VRAM.")
+    try:
+        # Use ComfyUI's prompt server queue system like the /free endpoint does
+        if hasattr(server.PromptServer.instance, 'prompt_queue'):
+            server.PromptServer.instance.prompt_queue.set_flag("unload_models", True)
+            server.PromptServer.instance.prompt_queue.set_flag("free_memory", True)
+            print("[MultiGPU] Set queue flags for memory clearing.")
+        
+        # Wait a bit for the queue to process
+        await asyncio.sleep(0.5)
+        
+        # Also do direct cleanup as backup, but with error handling
+        import gc
+        import comfy.model_management as mm
+        
+        try:
+            mm.unload_all_models()
+        except AttributeError as e:
+            print(f"[MultiGPU] Warning during model unload: {e}")
+        
+        try:
+            mm.soft_empty_cache()
+        except Exception as e:
+            print(f"[MultiGPU] Warning during cache clear: {e}")
+        
+        for _ in range(3):
+            gc.collect()
+        
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+        
+        print("[MultiGPU] VRAM cleared successfully.")
+        return web.json_response({"status": "success", "message": "GPU memory cleared."})
+    except Exception as e:
+        # Even if there's an error, try to do basic cleanup
+        import gc
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        print(f"[MultiGPU] Partial VRAM clear completed with warning: {e}")
+        return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"})
+
+
+# --- Global State & Callback Endpoint ---
+PENDING_JOBS = {}
+PENDING_JOBS_LOCK = asyncio.Lock()
+
+@server.PromptServer.instance.routes.post("/multigpu/job_complete")
+async def job_complete_endpoint(request):
+    try:
+        data = await request.post()
+        multi_job_id = data.get('multi_job_id')
+        image_file = data.get('image')
+        worker_id = data.get('worker_id')
+        image_index = data.get('image_index')
+        is_last = data.get('is_last', 'False').lower() == 'true'
+
+        if not all([multi_job_id, image_file]):
+            return await handle_api_error(request, "Missing job_id or image data", 400)
+
+        # Process image
+        img_data = image_file.file.read()
+        img = Image.open(io.BytesIO(img_data)).convert("RGB")
+        img_np = np.array(img).astype(np.float32) / 255.0
+        tensor = torch.from_numpy(img_np)[None,]
+
+        async with PENDING_JOBS_LOCK:
+            if multi_job_id in PENDING_JOBS:
+                await PENDING_JOBS[multi_job_id].put({
+                    'tensor': tensor,
+                    'worker_id': worker_id,
+                    'image_index': int(image_index) if image_index else 0,
+                    'is_last': is_last
+                })
+                print(f"[MultiGPU] Received result for job {multi_job_id} from worker {worker_id} (last: {is_last})")
+                return web.json_response({"status": "success"})
+            else:
+                return await handle_api_error(request, "Job not found or already complete", 404)
+    except Exception as e:
+        return await handle_api_error(request, e)
+
+# --- Collector Node ---
+class MultiGPUCollectorNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": { "images": ("IMAGE",) },
+            "hidden": {
+                "multi_job_id": ("STRING", {"default": ""}),
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "master_url": ("STRING", {"default": ""}),
+                "enabled_worker_ids": ("STRING", {"default": "[]"}),
+                "worker_batch_size": ("INT", {"default": 1, "min": 1, "max": 1024}),
+                "worker_id": ("STRING", {"default": ""}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "run"
+    CATEGORY = "image"
+    
+    def run(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id=""):
+        if not multi_job_id:
+            return (images,)
+
+        main_loop = server.PromptServer.instance.loop
+        future = Future()
+
+        async def run_and_set_future():
+            try:
+                result = await self.execute(images, multi_job_id, is_worker, master_url, enabled_worker_ids, worker_batch_size, worker_id)
+                future.set_result(result)
+            except Exception as e:
+                future.set_exception(e)
+
+        asyncio.run_coroutine_threadsafe(run_and_set_future(), main_loop)
+        return future.result()
+
+    async def send_image_to_master(self, image_tensor, multi_job_id, master_url, image_index, worker_id, is_last=False):
+        """Helper method to send a single image to the master"""
+        img_np = (image_tensor.cpu().numpy() * 255.).astype(np.uint8)
+        img = Image.fromarray(img_np)
+        byte_io = io.BytesIO()
+        img.save(byte_io, format='PNG')
+        byte_io.seek(0)
+        
+        data = aiohttp.FormData()
+        data.add_field('multi_job_id', multi_job_id)
+        data.add_field('worker_id', str(worker_id))
+        data.add_field('image_index', str(image_index))
+        data.add_field('is_last', str(is_last))
+        data.add_field('image', byte_io, filename=f'image_{image_index}.png', content_type='image/png')
+
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(f"{master_url}/multigpu/job_complete", data=data) as response:
+                    response.raise_for_status()
+        except Exception as e:
+            print(f"[MultiGPU Worker] Failed to send image {image_index+1} to master: {e}")
+
+    async def execute(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id=""):
+        if is_worker:
+            # Worker mode: send images to master
+            print(f"[MultiGPU Worker] Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master")
+            
+            # Send all images, marking the last one
+            for i in range(images.shape[0]):
+                is_last = (i == images.shape[0] - 1)
+                await self.send_image_to_master(images[i], multi_job_id, master_url, i, worker_id, is_last)
+            
+            return (images,)
+        else:
+            # Master mode: collect images from workers
+            enabled_workers = json.loads(enabled_worker_ids)
+            num_workers = len(enabled_workers)
+            if num_workers == 0:
+                return (images,)
+            
+            images_on_cpu = images.cpu()
+            print(f"[MultiGPU Master] Job {multi_job_id}: Collecting images from {num_workers} workers...")
+            
+            # Store master images first
+            master_batch_size = images.shape[0]
+            
+            # Initialize storage for collected images
+            worker_images = {}  # Dict to store images by worker_id and index
+            
+            async with PENDING_JOBS_LOCK:
+                if multi_job_id not in PENDING_JOBS:
+                    PENDING_JOBS[multi_job_id] = asyncio.Queue()
+                q = PENDING_JOBS[multi_job_id]
+            
+            # Collect images until all workers report they're done
+            collected_count = 0
+            workers_done = set()
+            
+            # Use a reasonable timeout for the first image
+            timeout = 10.0
+            
+            while len(workers_done) < num_workers:
+                try:
+                    result = await asyncio.wait_for(q.get(), timeout=timeout)
+                    worker_id = result['worker_id']
+                    image_index = result['image_index']
+                    tensor = result['tensor']
+                    is_last = result.get('is_last', False)
+                    
+                    if worker_id not in worker_images:
+                        worker_images[worker_id] = {}
+                    worker_images[worker_id][image_index] = tensor
+                    
+                    collected_count += 1
+                    
+                    # Once we start receiving images, use shorter timeout
+                    timeout = 10.0
+                    
+                    if is_last:
+                        workers_done.add(worker_id)
+                        print(f"[MultiGPU Master] Worker {worker_id} done. Collected {len(worker_images[worker_id])} images")
+                    else:
+                        print(f"[MultiGPU Master] Collected image {image_index + 1} from worker {worker_id}")
+                    
+                except asyncio.TimeoutError:
+                    missing_workers = set(enabled_workers) - workers_done
+                    print(f"[MultiGPU Master] Timeout. Still waiting for workers: {list(missing_workers)}")
+                    break
+            
+            total_collected = sum(len(imgs) for imgs in worker_images.values())
+            print(f"[MultiGPU Master] Collection complete. Received {total_collected} images from {len(workers_done)} workers")
+            
+            # Clean up job queue
+            async with PENDING_JOBS_LOCK:
+                if multi_job_id in PENDING_JOBS:
+                    del PENDING_JOBS[multi_job_id]
+
+            # Reorder images according to seed distribution pattern
+            # Pattern: master img 1, master img 2, worker 1 img 1, worker 1 img 2, worker 2 img 1, worker 2 img 2, etc.
+            ordered_tensors = []
+            
+            # Add master images first
+            for i in range(master_batch_size):
+                ordered_tensors.append(images_on_cpu[i:i+1])
+            
+            # Add worker images in order
+            for worker_id in enabled_workers:
+                # Convert worker_id to string since it comes as string from form data
+                worker_id_str = str(worker_id)
+                if worker_id_str in worker_images:
+                    # Sort by image index for each worker
+                    for idx in sorted(worker_images[worker_id_str].keys()):
+                        ordered_tensors.append(worker_images[worker_id_str][idx])
+            
+            combined = torch.cat(ordered_tensors, dim=0)
+            print(f"[MultiGPU Master] Job {multi_job_id} complete. Combined {combined.shape[0]} images total.")
+            return (combined,)
+
+# --- Distributor Node ---
+class MultiGPUDistributor:
+    """
+    Distributes seed values across multiple GPUs.
+    On master: passes through the original seed.
+    On workers: adds offset based on worker ID.
+    """
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "seed": ("INT", {
+                    "default": 0, 
+                    "min": 0,
+                    "max": 1125899906842624,
+                    "forceInput": False  # Widget by default, can be converted to input
+                }),
+            },
+            "hidden": {
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "worker_id": ("STRING", {"default": ""}),
+            },
+        }
+    
+    RETURN_TYPES = ("INT",)
+    RETURN_NAMES = ("seed",)
+    FUNCTION = "distribute"
+    CATEGORY = "utils"
+    
+    def distribute(self, seed, is_worker=False, worker_id=""):
+        if not is_worker:
+            # Master node: pass through original values
+            print(f"[MultiGPU Distributor] Master: seed={seed}")
+            return (seed,)
+        else:
+            # Worker node: apply offset based on worker index
+            # Find worker index from enabled_worker_ids
+            try:
+                # Worker IDs are passed as "worker_0", "worker_1", etc.
+                if worker_id.startswith("worker_"):
+                    worker_index = int(worker_id.split("_")[1])
+                else:
+                    # Fallback: try to parse as direct index
+                    worker_index = int(worker_id)
+                
+                offset = worker_index + 1
+                new_seed = seed + offset
+                print(f"[MultiGPU Distributor] Worker {worker_index}: seed={seed} → {new_seed}")
+                return (new_seed,)
+            except (ValueError, IndexError) as e:
+                print(f"[MultiGPU Distributor] Error parsing worker_id '{worker_id}': {e}")
+                # Fallback: return original seed
+                return (seed,)
+
+NODE_CLASS_MAPPINGS = { 
+    "MultiGPUCollector": MultiGPUCollectorNode,
+    "MultiGPUDistributor": MultiGPUDistributor
+}
+NODE_DISPLAY_NAME_MAPPINGS = { 
+    "MultiGPUCollector": "Multi-GPU Collector",
+    "MultiGPUDistributor": "Multi-GPU Seed"
+}

web/gpupanel.js

@@ -0,0 +1,493 @@
+import { app } from "../../scripts/app.js";
+import { api } from "../../scripts/api.js";
+
+class MultiGPUExtension {
+    // Button styling constants
+    static BUTTON_STYLES = {
+        base: "width: 100%; padding: 8px; font-size: 12px; color: white; border: none; border-radius: 4px; cursor: pointer; transition: background-color 0.2s;",
+        clearMemory: "background-color: #555;",
+        success: "background-color: #3ca03c;",
+        error: "background-color: #c04c4c;"
+    };
+
+    constructor() {
+        this.config = null;
+        this.originalQueuePrompt = api.queuePrompt.bind(api);
+
+        this.loadConfig().then(() => {
+            this.registerSidebarTab();
+            this.setupInterceptor();
+        });
+    }
+
+    // --- State & Config Management (Single Source of Truth) ---
+
+    get enabledWorkers() {
+        return this.config?.workers?.filter(w => w.enabled) || [];
+    }
+
+    get isEnabled() {
+        return this.enabledWorkers.length > 0;
+    }
+
+    async loadConfig() {
+        try {
+            const response = await fetch(`${window.location.origin}/multigpu/config`);
+            if (response.ok) {
+                this.config = await response.json();
+                console.log("[MultiGPU] Loaded config:", this.config);
+            } else {
+                console.error("[MultiGPU] Failed to load config");
+                this.config = { workers: [], settings: {} };
+            }
+        } catch (error) {
+            console.error("[MultiGPU] Error loading config:", error);
+            this.config = { workers: [], settings: {} };
+        }
+    }
+
+    async updateWorkerEnabled(workerId, enabled) {
+        const worker = this.config.workers.find(w => w.id === workerId);
+        if (worker) {
+            worker.enabled = enabled;
+        }
+        
+        try {
+            await fetch(`${window.location.origin}/multigpu/config/update_worker`, {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                body: JSON.stringify({ worker_id: workerId, enabled: enabled })
+            });
+        } catch (error) {
+            console.error("[MultiGPU] Error updating worker:", error);
+        }
+    }
+    
+    async _updateSetting(key, value) {
+        try {
+            await fetch(`${window.location.origin}/multigpu/config/update_setting`, {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                body: JSON.stringify({ key, value })
+            });
+        } catch (error) {
+            console.error(`[MultiGPU] Error updating setting '${key}':`, error);
+        }
+    }
+
+    // --- UI Rendering ---
+
+    registerSidebarTab() {
+        app.extensionManager.registerSidebarTab({
+            id: "multiGPU",
+            icon: "pi pi-server",
+            title: "Multi-GPU",
+            tooltip: "Multi-GPU Control Panel",
+            type: "custom",
+            render: (el) => this.renderSidebarContent(el)
+        });
+    }
+
+    _createButton(text, onClick, style) {
+        const button = document.createElement("button");
+        button.textContent = text;
+        button.style.cssText = MultiGPUExtension.BUTTON_STYLES.base + style;
+        button.onclick = onClick;
+        return button;
+    }
+
+    async _handleClearMemory(button) {
+        const originalText = button.textContent;
+        button.textContent = "Clearing...";
+        button.disabled = true;
+        
+        try {
+            const urlsToClear = this.enabledWorkers.map(w => ({ 
+                name: w.name, 
+                url: `http://${w.host || window.location.hostname}:${w.port}` 
+            }));
+            
+            if (urlsToClear.length === 0) {
+                button.textContent = "No Workers";
+                button.style.backgroundColor = MultiGPUExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', '');
+                setTimeout(() => {
+                    button.textContent = originalText;
+                    button.style.cssText = MultiGPUExtension.BUTTON_STYLES.base + MultiGPUExtension.BUTTON_STYLES.clearMemory;
+                }, 3000);
+                return;
+            }
+            
+            const promises = urlsToClear.map(target =>
+                fetch(`${target.url}/multigpu/clear_memory`, { 
+                    method: 'POST', 
+                    mode: 'cors'
+                })
+                    .then(response => ({ ok: response.ok, name: target.name }))
+                    .catch(() => ({ ok: false, name: target.name }))
+            );
+            
+            const results = await Promise.all(promises);
+            const failures = results.filter(r => !r.ok);
+            
+            if (failures.length === 0) {
+                button.textContent = "Success!";
+                button.style.backgroundColor = MultiGPUExtension.BUTTON_STYLES.success.split(':')[1].trim().replace(';', '');
+            } else {
+                button.textContent = "Error! See Console";
+                button.style.backgroundColor = MultiGPUExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', '');
+                console.error("[MultiGPU] Failed to clear memory on:", failures.map(f => f.name).join(", "));
+            }
+            
+            setTimeout(() => {
+                button.textContent = originalText;
+                button.style.cssText = MultiGPUExtension.BUTTON_STYLES.base + MultiGPUExtension.BUTTON_STYLES.clearMemory;
+            }, 3000);
+        } finally {
+            button.disabled = false;
+        }
+    }
+
+
+    renderSidebarContent(el) {
+        el.innerHTML = '';
+        const container = document.createElement("div");
+        container.style.cssText = "padding: 15px; display: flex; flex-direction: column; height: 100%;";
+        
+        const title = document.createElement("h3");
+        title.textContent = "Multi-GPU Control";
+        title.style.cssText = "margin: 0 0 15px 0; padding-bottom: 10px; border-bottom: 1px solid #444;";
+        container.appendChild(title);
+        
+        const gpuSection = document.createElement("div");
+        gpuSection.style.cssText = "flex: 1; overflow-y: auto; margin-bottom: 15px;";
+        const gpuTitle = document.createElement("h4");
+        gpuTitle.textContent = "Available Workers";
+        gpuTitle.style.cssText = "margin: 0 0 10px 0; font-size: 14px;";
+        gpuSection.appendChild(gpuTitle);
+        
+        const gpuList = document.createElement("div");
+        (this.config?.workers || []).forEach(worker => {
+            const gpuDiv = document.createElement("div");
+            gpuDiv.style.cssText = "margin-bottom: 8px; padding: 8px; background: #2a2a2a; border-radius: 4px;";
+            const checkbox = document.createElement("input");
+            checkbox.type = "checkbox";
+            checkbox.id = `gpu-${worker.id}`;
+            checkbox.checked = worker.enabled;
+            checkbox.onchange = async (e) => {
+                await this.updateWorkerEnabled(worker.id, e.target.checked);
+                this.updateSummary();
+            };
+            const label = document.createElement("label");
+            label.htmlFor = `gpu-${worker.id}`;
+            label.style.cssText = "cursor: pointer; display: flex; align-items: center; gap: 8px;";
+            const gpuInfo = document.createElement("span");
+            const hostInfo = worker.host ? ` • ${worker.host}` : '';
+            const cudaInfo = worker.cuda_device !== undefined ? `CUDA ${worker.cuda_device} • ` : '';
+            gpuInfo.innerHTML = `<strong>${worker.name}</strong><br><small style="color: #888;">${cudaInfo}Port ${worker.port}${hostInfo}</small>`;
+            label.appendChild(checkbox);
+            label.appendChild(gpuInfo);
+            gpuDiv.appendChild(label);
+            gpuList.appendChild(gpuDiv);
+        });
+        gpuSection.appendChild(gpuList);
+        container.appendChild(gpuSection);
+        
+        const actionsSection = document.createElement("div");
+        actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;";
+        
+        const clearMemButton = this._createButton("Clear Worker VRAM", (e) => this._handleClearMemory(e.target), MultiGPUExtension.BUTTON_STYLES.clearMemory);
+        clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)";
+        actionsSection.appendChild(clearMemButton);
+        
+        container.appendChild(actionsSection);
+
+        const summarySection = document.createElement("div");
+        summarySection.style.cssText = "border-top: 1px solid #444; padding-top: 10px;";
+        const summary = document.createElement("div");
+        summary.id = "multigpu-summary";
+        summary.style.cssText = "font-size: 11px; color: #888;";
+        summarySection.appendChild(summary);
+        container.appendChild(summarySection);
+        el.appendChild(container);
+        this.updateSummary();
+    }
+
+    updateSummary() {
+        const summaryEl = document.getElementById('multigpu-summary');
+        if (summaryEl) {
+            const totalGPUs = this.enabledWorkers.length + 1;
+            if (this.isEnabled) {
+                summaryEl.textContent = `If Collector node is present, total generation = (${totalGPUs} GPUs × Batch Size)`;
+            } else {
+                summaryEl.textContent = "Only the master GPU will be used.";
+            }
+        }
+    }
+
+    // --- Helper Methods ---
+
+    // --- Core Logic & Execution ---
+
+    setupInterceptor() {
+        api.queuePrompt = async (number, prompt) => {
+            if (this.isEnabled && this.findNodesByClass(prompt.output, "MultiGPUCollector").length > 0) {
+                return await this.executeParallelMultiGPU(prompt);
+            }
+            return this.originalQueuePrompt(number, prompt);
+        };
+    }
+
+    async executeParallelMultiGPU(promptWrapper) {
+        try {
+            const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution
+            const enabledWorkers = this.enabledWorkers;
+            
+            // Find all collector nodes in the workflow
+            const collectorNodes = this.findNodesByClass(promptWrapper.output, "MultiGPUCollector");
+            
+            // Map original node IDs to truly unique job IDs for this specific run
+            const job_id_map = new Map(collectorNodes.map(node => [node.id, `${executionPrefix}_${node.id}`]));
+            
+            // Prepare a separate job queue on the backend for each unique job ID
+            const preparePromises = Array.from(job_id_map.values()).map(uniqueId => this._prepareMultiGpuJob(uniqueId));
+            await Promise.all(preparePromises);
+
+            const jobs = [];
+            const participants = ['master', ...enabledWorkers.map(w => w.id)];
+
+            for (const participantId of participants) {
+                const jobApiPrompt = this._prepareApiPromptForParticipant(
+                    promptWrapper.output, participantId,
+                    { 
+                        enabled_worker_ids: enabledWorkers.map(w => w.id), 
+                        workflow: promptWrapper.workflow,
+                        job_id_map: job_id_map // Pass the map of unique IDs
+                    }
+                );
+                
+                if (participantId === 'master') {
+                    jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } });
+                } else {
+                    const worker = this.config.workers.find(w => w.id === participantId);
+                    if (worker) jobs.push({ type: 'worker', worker, prompt: jobApiPrompt, workflow: promptWrapper.workflow });
+                }
+            }
+            
+            const result = await this._executeJobs(jobs);
+            return result;
+        } catch (error) {
+            console.error("[MultiGPU] Parallel execution failed:", error);
+            alert(`[MultiGPU] Parallel execution failed: ${error.message}`);
+            throw error;
+        }
+    }
+
+
+    async _executeJobs(jobs) {
+        let masterPromptId = null;
+        const promises = jobs.map(job => {
+            if (job.type === 'master') {
+                return this.originalQueuePrompt(0, job.promptWrapper).then(result => {
+                    masterPromptId = result;
+                    return result;
+                });
+            } else {
+                return this._dispatchToWorker(job.worker, job.prompt, job.workflow);
+            }
+        });
+        await Promise.all(promises);
+        return masterPromptId || { "prompt_id": "multi-gpu-job-dispatched" };
+    }
+    
+    // --- Helper Methods ---
+
+    findNodesByClass(apiPrompt, className) {
+        return Object.entries(apiPrompt)
+            .filter(([, nodeData]) => nodeData.class_type === className)
+            .map(([nodeId, nodeData]) => ({ id: nodeId, data: nodeData }));
+    }
+
+    /**
+     * Find all nodes connected to collector nodes (both upstream and downstream)
+     * @param {Object} apiPrompt - The workflow API prompt
+     * @param {Array<string>} collectorIds - Array of collector node IDs
+     * @returns {Set<string>} Set of node IDs connected to any collector
+     */
+    findCollectorConnectedNodes(apiPrompt, collectorIds) {
+        const connected = new Set(collectorIds); // Include all collectors
+        const toProcess = [...collectorIds];
+        
+        // First, build a map of which nodes output to which nodes
+        const outputMap = new Map(); // sourceNode -> Set of target nodes
+        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) {
+                        const sourceNodeId = String(inputValue[0]);
+                        if (!outputMap.has(sourceNodeId)) {
+                            outputMap.set(sourceNodeId, new Set());
+                        }
+                        outputMap.get(sourceNodeId).add(nodeId);
+                    }
+                }
+            }
+        }
+        
+        // Traverse both upstream and downstream
+        while (toProcess.length > 0) {
+            const nodeId = toProcess.pop();
+            const node = apiPrompt[nodeId];
+            
+            // Traverse upstream (inputs)
+            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);
+                        }
+                    }
+                }
+            }
+            
+            // Traverse downstream (outputs)
+            if (outputMap.has(nodeId)) {
+                for (const targetNodeId of outputMap.get(nodeId)) {
+                    if (!connected.has(targetNodeId)) {
+                        connected.add(targetNodeId);
+                        toProcess.push(targetNodeId);
+                    }
+                }
+            }
+        }
+        
+        return connected;
+    }
+
+    /**
+     * Prune workflow to only include nodes connected to collectors
+     * @param {Object} apiPrompt - The full workflow API prompt
+     * @param {Array} collectorNodes - Array of collector nodes (optional, will find if not provided)
+     * @returns {Object} Pruned API prompt with only required nodes
+     */
+    pruneWorkflowForWorker(apiPrompt, collectorNodes = null) {
+        // Find all collector nodes if not provided
+        if (!collectorNodes) {
+            collectorNodes = this.findNodesByClass(apiPrompt, "MultiGPUCollector");
+        }
+        
+        if (collectorNodes.length === 0) {
+            // No collectors, return full workflow
+            return apiPrompt;
+        }
+        
+        // Get all nodes connected to collectors (both upstream and downstream)
+        const collectorIds = collectorNodes.map(node => node.id);
+        const connectedNodes = this.findCollectorConnectedNodes(apiPrompt, collectorIds);
+        
+        console.log(`[MultiGPU] Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes`);
+        
+        // Create pruned prompt with only required nodes
+        const prunedPrompt = {};
+        for (const nodeId of connectedNodes) {
+            prunedPrompt[nodeId] = JSON.parse(JSON.stringify(apiPrompt[nodeId]));
+        }
+        
+        return prunedPrompt;
+    }
+
+
+    _prepareApiPromptForParticipant(baseApiPrompt, participantId, options = {}) {
+        let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt));
+        const isMaster = participantId === 'master';
+        
+        // Find collector nodes once (before pruning)
+        const collectorNodes = this.findNodesByClass(jobApiPrompt, "MultiGPUCollector");
+        
+        // For workers, prune the workflow to only include collector dependencies
+        if (!isMaster && collectorNodes.length > 0) {
+            jobApiPrompt = this.pruneWorkflowForWorker(jobApiPrompt, collectorNodes);
+        }
+        
+        // Handle MultiGPU distributor nodes
+        const distributorNodes = this.findNodesByClass(jobApiPrompt, "MultiGPUDistributor");
+        if (distributorNodes.length > 0) {
+            console.log(`[MultiGPU] Found ${distributorNodes.length} distributor node(s)`);
+        }
+        
+        for (const distributor of distributorNodes) {
+            const { inputs } = jobApiPrompt[distributor.id];
+            inputs.is_worker = !isMaster;
+            if (!isMaster) {
+                const workerIndex = options.enabled_worker_ids.indexOf(participantId);
+                inputs.worker_id = `worker_${workerIndex}`;
+                console.log(`[MultiGPU] Set distributor ${distributor.id} for worker ${workerIndex}`);
+            }
+        }
+        
+        // Handle MultiGPU collector nodes (already found above)
+        for (const collector of collectorNodes) {
+            const { inputs } = jobApiPrompt[collector.id];
+            
+            // Get the unique job ID from the map created for this execution
+            const uniqueJobId = options.job_id_map ? options.job_id_map.get(collector.id) : collector.id;
+            
+            // Use the truly unique ID for this execution
+            inputs.multi_job_id = uniqueJobId;
+            inputs.is_worker = !isMaster;
+            if (isMaster) {
+                inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
+            } else {
+                inputs.master_url = window.location.origin;
+                // Also make the worker_job_id unique to prevent potential caching issues
+                inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`;
+                inputs.worker_id = participantId;
+            }
+        }
+        
+        return jobApiPrompt;
+    }
+
+    async _prepareMultiGpuJob(multi_job_id) {
+        try {
+            await fetch(`${window.location.origin}/multigpu/prepare_job`, {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                body: JSON.stringify({ multi_job_id })
+            });
+        } catch (error) {
+            console.error("[MultiGPU] Error preparing job:", error);
+            throw error;
+        }
+    }
+
+    async _dispatchToWorker(worker, prompt, workflow) {
+        const workerUrl = `http://${worker.host || window.location.hostname}:${worker.port}`;
+        
+        const promptToSend = {
+            prompt,
+            extra_data: { extra_pnginfo: { workflow } },
+            client_id: api.clientId
+        };
+        
+        
+        try {
+            await fetch(`${workerUrl}/prompt`, { 
+                method: 'POST', 
+                headers: { 'Content-Type': 'application/json' }, 
+                mode: 'cors', 
+                body: JSON.stringify(promptToSend) 
+            });
+        } catch (e) {
+            console.error(`[MultiGPU] Failed to connect to worker ${worker.name} at ${workerUrl}`, e);
+        }
+    }
+}
+
+app.registerExtension({
+    name: "MultiGPU.Panel",
+    async setup() {
+        new MultiGPUExtension();
+    }
+});