From d7752b00dafe33c42cca5040fffb877303cd4927 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Tue, 8 Jul 2025 13:46:46 +1000 Subject: [PATCH] Add image handling for remote workers --- distributed.py | 40 +++++++++ web/gpupanel.js | 232 ++++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 263 insertions(+), 9 deletions(-) diff --git a/distributed.py b/distributed.py index 54e0f59..299ede4 100644 --- a/distributed.py +++ b/distributed.py @@ -1132,6 +1132,46 @@ if not hasattr(prompt_server, 'distributed_pending_jobs'): prompt_server.distributed_pending_jobs = {} prompt_server.distributed_jobs_lock = asyncio.Lock() +@server.PromptServer.instance.routes.post("/distributed/load_image") +async def load_image_endpoint(request): + """Load an image file and return it as base64 data.""" + try: + data = await request.json() + image_path = data.get("image_path") + + if not image_path: + return await handle_api_error(request, "Missing image_path", 400) + + import folder_paths + import base64 + from PIL import Image + import io + + # Use ComfyUI's folder paths to find the image + full_path = folder_paths.get_annotated_filepath(image_path) + + if not os.path.exists(full_path): + return await handle_api_error(request, f"Image not found: {image_path}", 404) + + # Load and convert to base64 + with Image.open(full_path) as img: + # Convert to RGB if needed + if img.mode not in ('RGB', 'RGBA'): + img = img.convert('RGB') + + # Save to bytes + buffer = io.BytesIO() + img.save(buffer, format='PNG', compress_level=1) # Fast compression + img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8') + + return web.json_response({ + "status": "success", + "image_data": f"data:image/png;base64,{img_base64}" + }) + + except Exception as e: + return await handle_api_error(request, e, 500) + @server.PromptServer.instance.routes.post("/distributed/job_complete") async def job_complete_endpoint(request): try: diff --git a/web/gpupanel.js b/web/gpupanel.js index 1122820..03d0ff8 100644 --- a/web/gpupanel.js +++ b/web/gpupanel.js @@ -850,20 +850,35 @@ class DistributedExtension { const participants = ['master', ...activeWorkers.map(w => w.id)]; for (const participantId of participants) { + const options = { + enabled_worker_ids: activeWorkers.map(w => w.id), + workflow: promptWrapper.workflow, + job_id_map: job_id_map // Pass the map of unique IDs + }; + const jobApiPrompt = this._prepareApiPromptForParticipant( - promptWrapper.output, participantId, - { - enabled_worker_ids: activeWorkers.map(w => w.id), - workflow: promptWrapper.workflow, - job_id_map: job_id_map // Pass the map of unique IDs - } + promptWrapper.output, participantId, options ); if (participantId === 'master') { jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } }); } else { const worker = activeWorkers.find(w => w.id === participantId); - if (worker) jobs.push({ type: 'worker', worker, prompt: jobApiPrompt, workflow: promptWrapper.workflow }); + if (worker) { + const job = { + type: 'worker', + worker, + prompt: jobApiPrompt, + workflow: promptWrapper.workflow + }; + + // Add image references if found for remote workers + if (options._imageReferences) { + job.imageReferences = options._imageReferences; + } + + jobs.push(job); + } } } @@ -878,6 +893,23 @@ class DistributedExtension { async _executeJobs(jobs) { let masterPromptId = null; + + // Pre-load all unique images before dispatching to workers + const allImageReferences = new Map(); + for (const job of jobs) { + if (job.type === 'worker' && job.imageReferences) { + for (const [filename, info] of job.imageReferences) { + allImageReferences.set(filename, info); + } + } + } + + if (allImageReferences.size > 0) { + this.debugLog(`Pre-loading ${allImageReferences.size} unique image(s) for all workers`); + await this.loadImagesForWorker(allImageReferences); + } + + // Now dispatch jobs in parallel const promises = jobs.map(job => { if (job.type === 'master') { return this.originalQueuePrompt(0, job.promptWrapper).then(result => { @@ -885,7 +917,7 @@ class DistributedExtension { return result; }); } else { - return this._dispatchToWorker(job.worker, job.prompt, job.workflow); + return this._dispatchToWorker(job.worker, job.prompt, job.workflow, job.imageReferences); } }); await Promise.all(promises); @@ -899,6 +931,36 @@ class DistributedExtension { .filter(([, nodeData]) => nodeData.class_type === className) .map(([nodeId, nodeData]) => ({ id: nodeId, data: nodeData })); } + + /** + * Find all image references in the workflow + * Looks for inputs named "image" that contain filename strings + */ + findImageReferences(apiPrompt) { + const images = new Map(); + // Updated regex to handle: + // - Standard files: "image.png" + // - Subfolder files: "subfolder/image.png" + // - ComfyUI special format: "clipspace/file.png [input]" + const imageExtensions = /\.(png|jpg|jpeg|gif|webp|bmp)(\s*\[\w+\])?$/i; + + for (const [nodeId, node] of Object.entries(apiPrompt)) { + if (node.inputs && node.inputs.image) { + const imageValue = node.inputs.image; + // Check if it's a string filename (not an array connection) + if (typeof imageValue === 'string' && imageExtensions.test(imageValue)) { + images.set(imageValue, { + nodeId, + nodeType: node.class_type, + inputName: 'image' + }); + this.debugLog(`Found image reference: ${imageValue} in node ${nodeId} (${node.class_type})`); + } + } + } + + return images; + } /** * Find all nodes connected to collector nodes (both upstream and downstream) @@ -1007,6 +1069,24 @@ class DistributedExtension { jobApiPrompt = this.pruneWorkflowForWorker(jobApiPrompt, allDistributedNodes); } + // Handle image references for remote workers + if (!isMaster && options.enabled_worker_ids) { + // Check if this is a remote worker + const workerId = participantId; + const workerInfo = this.config.workers.find(w => w.id === workerId); + const isRemote = workerInfo && workerInfo.host; + + if (isRemote) { + // Find all image references in the pruned workflow + const imageReferences = this.findImageReferences(jobApiPrompt); + if (imageReferences.size > 0) { + this.debugLog(`Found ${imageReferences.size} image references for remote worker ${workerId}`); + // Store image references for later processing + options._imageReferences = imageReferences; + } + } + } + // Handle Distributed seed nodes const distributorNodes = this.findNodesByClass(jobApiPrompt, "DistributedSeed"); if (distributorNodes.length > 0) { @@ -1214,16 +1294,44 @@ class DistributedExtension { this.updateWorkerControls(worker.id); } - async _dispatchToWorker(worker, prompt, workflow) { + async _dispatchToWorker(worker, prompt, workflow, imageReferences) { const host = worker.host || "localhost"; const workerUrl = `http://${host}:${worker.port}`; + // Debug logging - always log to console for debugging + console.log(`[Distributed] === Dispatching to ${worker.name} (${worker.id}) ===`); + console.log('[Distributed] Worker URL:', workerUrl); + + // Handle image uploads for remote workers + if (imageReferences && imageReferences.size > 0) { + if (this.debugMode) { + console.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`); + } + + try { + // Load images from master + const images = await this.loadImagesForWorker(imageReferences); + + // Upload images to worker + if (images.length > 0) { + await this.uploadImagesToWorker(workerUrl, images); + if (this.debugMode) { + console.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`); + } + } + } catch (error) { + this.log(`Failed to process images for worker ${worker.name}: ${error.message}`); + // Continue with workflow execution even if image upload fails + } + } + const promptToSend = { prompt, extra_data: { extra_pnginfo: { workflow } }, client_id: api.clientId }; + console.log('[Distributed] Prompt data:', promptToSend); try { await fetch(`${workerUrl}/prompt`, { @@ -1237,6 +1345,112 @@ class DistributedExtension { } } + async loadImagesForWorker(imageReferences) { + const images = []; + + // Use a cache to avoid loading the same image multiple times + if (!this._imageCache) { + this._imageCache = new Map(); + } + + for (const [filename, info] of imageReferences) { + try { + // Check cache first + if (this._imageCache.has(filename)) { + images.push(this._imageCache.get(filename)); + this.debugLog(`Using cached image: ${filename}`); + continue; + } + + // Load image from master's filesystem via API + const response = await fetch(`${window.location.origin}/distributed/load_image`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ image_path: filename }) + }); + + if (response.ok) { + const data = await response.json(); + const imageData = { + name: filename, + image: data.image_data + }; + images.push(imageData); + + // Cache the image for future use + this._imageCache.set(filename, imageData); + this.debugLog(`Loaded and cached image: ${filename}`); + } else { + this.log(`Failed to load image ${filename}: ${response.statusText}`); + } + } catch (error) { + this.log(`Error loading image ${filename}: ${error.message}`); + } + } + + // Clear cache after a reasonable time to avoid memory issues + setTimeout(() => { + if (this._imageCache && this._imageCache.size > 0) { + this.debugLog(`Clearing image cache (${this._imageCache.size} images)`); + this._imageCache.clear(); + } + }, 30000); // Clear after 30 seconds + + return images; + } + + async uploadImagesToWorker(workerUrl, images) { + // Upload images to worker's ComfyUI instance + for (const imageData of images) { + const formData = new FormData(); + + // Convert base64 to blob + const base64Data = imageData.image.replace(/^data:image\/\w+;base64,/, ''); + const byteCharacters = atob(base64Data); + const byteNumbers = new Array(byteCharacters.length); + for (let i = 0; i < byteCharacters.length; i++) { + byteNumbers[i] = byteCharacters.charCodeAt(i); + } + const byteArray = new Uint8Array(byteNumbers); + const blob = new Blob([byteArray], { type: 'image/png' }); + + // Clean the filename - remove [input] suffix and handle subfolder + let cleanName = imageData.name; + let subfolder = ''; + + // Remove [input] or other suffixes + cleanName = cleanName.replace(/\s*\[\w+\]$/, ''); + + // Extract subfolder if present + if (cleanName.includes('/')) { + const parts = cleanName.split('/'); + subfolder = parts.slice(0, -1).join('/'); + cleanName = parts[parts.length - 1]; + } + + formData.append('image', blob, cleanName); + formData.append('type', 'input'); + formData.append('subfolder', subfolder); + formData.append('overwrite', 'true'); + + try { + const response = await fetch(`${workerUrl}/upload/image`, { + method: 'POST', + mode: 'cors', + body: formData + }); + + if (!response.ok) { + throw new Error(`Upload failed: ${response.statusText}`); + } + + this.debugLog(`Uploaded image to worker: ${imageData.name} -> ${subfolder}/${cleanName}`); + } catch (error) { + throw new Error(`Failed to upload ${imageData.name}: ${error.message}`); + } + } + } + async performPreflightCheck(workers) { if (workers.length === 0) return [];