diff --git a/distributed.py b/distributed.py index 6e7bd2e..7b53c41 100644 --- a/distributed.py +++ b/distributed.py @@ -1464,7 +1464,7 @@ if not hasattr(prompt_server, 'distributed_pending_jobs'): @server.PromptServer.instance.routes.post("/distributed/load_image") async def load_image_endpoint(request): - """Load an image or video file and return it as base64 data.""" + """Load an image or video file and return it as base64 data with hash.""" try: data = await request.json() image_path = data.get("image_path") @@ -1475,6 +1475,7 @@ async def load_image_endpoint(request): import folder_paths import base64 import io + import hashlib # Use ComfyUI's folder paths to find the file full_path = folder_paths.get_annotated_filepath(image_path) @@ -1482,6 +1483,13 @@ async def load_image_endpoint(request): if not os.path.exists(full_path): return await handle_api_error(request, f"File not found: {image_path}", 404) + # Calculate file hash + hash_md5 = hashlib.md5() + with open(full_path, 'rb') as f: + for chunk in iter(lambda: f.read(4096), b""): + hash_md5.update(chunk) + file_hash = hash_md5.hexdigest() + # Check if it's a video file video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.webm'} file_ext = os.path.splitext(full_path)[1].lower() @@ -1505,7 +1513,8 @@ async def load_image_endpoint(request): video_base64 = base64.b64encode(file_data).decode('utf-8') return web.json_response({ "status": "success", - "image_data": f"data:{mime_type};base64,{video_base64}" + "image_data": f"data:{mime_type};base64,{video_base64}", + "hash": file_hash }) else: # For images, use PIL @@ -1524,12 +1533,55 @@ async def load_image_endpoint(request): return web.json_response({ "status": "success", - "image_data": f"data:image/png;base64,{img_base64}" + "image_data": f"data:image/png;base64,{img_base64}", + "hash": file_hash }) except Exception as e: return await handle_api_error(request, e, 500) +@server.PromptServer.instance.routes.post("/distributed/check_file") +async def check_file_endpoint(request): + """Check if a file exists and matches the given hash.""" + try: + data = await request.json() + filename = data.get("filename") + expected_hash = data.get("hash") + + if not filename or not expected_hash: + return await handle_api_error(request, "Missing filename or hash", 400) + + import folder_paths + import hashlib + + # Use ComfyUI's folder paths to find the file + full_path = folder_paths.get_annotated_filepath(filename) + + if not os.path.exists(full_path): + return web.json_response({ + "status": "success", + "exists": False + }) + + # Calculate file hash + hash_md5 = hashlib.md5() + with open(full_path, 'rb') as f: + for chunk in iter(lambda: f.read(4096), b""): + hash_md5.update(chunk) + file_hash = hash_md5.hexdigest() + + # Check if hash matches + hash_matches = file_hash == expected_hash + + return web.json_response({ + "status": "success", + "exists": True, + "hash_matches": hash_matches + }) + + except Exception as e: + return await handle_api_error(request, e, 500) + @server.PromptServer.instance.routes.get("/distributed/system_info") async def system_info_endpoint(request): """Return system information including platform details""" diff --git a/web/executionUtils.js b/web/executionUtils.js index 56ad8c2..4c116cd 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -325,20 +325,27 @@ async function dispatchToWorker(extension, worker, prompt, workflow, imageRefere // Handle image uploads for remote workers if (imageReferences && imageReferences.size > 0) { - extension.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`, "debug"); + // Check if this is a local worker (same host as master) + const isLocalWorker = workerUrl.includes('127.0.0.1') || workerUrl.includes('localhost'); - try { - // Load images from master - const images = await loadImagesForWorker(extension, imageReferences); + if (isLocalWorker) { + extension.log(`[Distributed] Skipping image processing for local worker ${worker.name} (shares filesystem with master)`, "debug"); + } else { + extension.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`, "debug"); - // Upload images to worker - if (images.length > 0) { - await uploadImagesToWorker(extension, workerUrl, images); - extension.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`, "debug"); + try { + // Load images from master + const images = await loadImagesForWorker(extension, imageReferences); + + // Upload images to worker + if (images.length > 0) { + await uploadImagesToWorker(extension, workerUrl, images); + extension.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`, "debug"); + } + } catch (error) { + extension.log(`Failed to process images for worker ${worker.name}: ${error.message}`, "error"); + // Continue with workflow execution even if image upload fails } - } catch (error) { - extension.log(`Failed to process images for worker ${worker.name}: ${error.message}`, "error"); - // Continue with workflow execution even if image upload fails } } @@ -391,7 +398,8 @@ export async function loadImagesForWorker(extension, imageReferences) { const data = await extension.api.loadImage(filename); const imageData = { name: filename, - image: data.image_data + image: data.image_data, + hash: data.hash // Include hash from the response }; images.push(imageData); @@ -421,6 +429,32 @@ export async function loadImagesForWorker(extension, imageReferences) { export async function uploadImagesToWorker(extension, workerUrl, images) { // Upload images to worker's ComfyUI instance for (const imageData of images) { + // Check if file already exists with matching hash + if (imageData.hash) { + try { + const checkResponse = await fetch(`${workerUrl}/distributed/check_file`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + mode: 'cors', + body: JSON.stringify({ + filename: imageData.name, + hash: imageData.hash + }) + }); + + if (checkResponse.ok) { + const result = await checkResponse.json(); + if (result.exists && result.hash_matches) { + extension.log(`File ${imageData.name} already exists on worker with matching hash, skipping upload`, "debug"); + continue; + } + } + } catch (error) { + // If check fails, proceed with upload + extension.log(`Failed to check file existence for ${imageData.name}: ${error.message}`, "debug"); + } + } + const formData = new FormData(); // Detect MIME type from base64 header (supports both image and video)