diff --git a/distributed.py b/distributed.py index 9664862..6e7bd2e 100644 --- a/distributed.py +++ b/distributed.py @@ -16,6 +16,8 @@ import atexit import signal import sys import shlex +import uuid +from multiprocessing import Queue from comfy.utils import ProgressBar # Import shared utilities @@ -258,6 +260,33 @@ async def get_network_info_endpoint(request): "recommended_ip": None }) +@server.PromptServer.instance.routes.get("/distributed/system_info") +async def get_system_info_endpoint(request): + """Get system information including machine ID for local worker detection.""" + try: + import socket + + return web.json_response({ + "status": "success", + "hostname": socket.gethostname(), + "machine_id": get_machine_id(), + "platform": { + "system": platform.system(), + "machine": platform.machine(), + "node": platform.node(), + "path_separator": os.sep, # Add path separator + "os_name": os.name # Add OS name (posix, nt, etc.) + }, + "is_docker": is_docker_environment(), + "is_runpod": is_runpod_environment(), + "runpod_pod_id": os.environ.get('RUNPOD_POD_ID') + }) + except Exception as e: + return web.json_response({ + "status": "error", + "message": str(e) + }, status=500) + @server.PromptServer.instance.routes.post("/distributed/config/update_worker") async def update_worker_endpoint(request): try: @@ -1156,6 +1185,106 @@ class WorkerProcessManager: # Create global instance worker_manager = WorkerProcessManager() +# Add IPC queue storage +worker_manager.queues = {} + +# Helper to get aiohttp session +async def get_client_session(): + """Get or create aiohttp client session.""" + if not hasattr(server.PromptServer.instance, '_distributed_session'): + server.PromptServer.instance._distributed_session = aiohttp.ClientSession() + return server.PromptServer.instance._distributed_session + +# Local worker detection functions +async def is_local_worker(worker_config): + """Check if a worker is running on the same machine as the master.""" + host = worker_config.get('host', 'localhost') + if host in ['localhost', '127.0.0.1', '0.0.0.0', ''] or worker_config.get('type') == 'local': + return True + + # For cloud workers, check if on same physical host + if worker_config.get('type') == 'cloud': + return await is_same_physical_host(worker_config) + + return False + +async def is_same_physical_host(worker_config): + """Compare machine IDs to determine if worker is on same physical host.""" + try: + # Get master machine ID + master_machine_id = get_machine_id() + + # Fetch worker's machine ID via API + host = worker_config.get('host', 'localhost') + port = worker_config.get('port', 8188) + + session = await get_client_session() + async with session.get( + f"http://{host}:{port}/distributed/system_info", + timeout=aiohttp.ClientTimeout(total=5) + ) as resp: + if resp.status == 200: + data = await resp.json() + worker_machine_id = data.get('machine_id') + return worker_machine_id == master_machine_id + else: + debug_log(f"Failed to get system info from worker: HTTP {resp.status}") + return False + except Exception as e: + debug_log(f"Error checking same physical host: {e}") + return False + +def get_machine_id(): + """Get a unique identifier for this machine.""" + # Try multiple methods to get a stable machine ID + try: + # Method 1: MAC address-based UUID + return str(uuid.getnode()) + except: + try: + # Method 2: Platform + hostname + import socket + return f"{platform.machine()}_{socket.gethostname()}" + except: + # Fallback + return platform.machine() + +def is_docker_environment(): + """Check if running inside Docker container.""" + return (os.path.exists('/.dockerenv') or + os.environ.get('DOCKER_CONTAINER', False) or + 'docker' in platform.node().lower()) + +def is_runpod_environment(): + """Check if running in Runpod environment.""" + return (os.environ.get('RUNPOD_POD_ID') is not None or + os.environ.get('RUNPOD_API_KEY') is not None) + +async def get_comms_channel(worker_id, worker_config): + """Get communication channel for a worker (HTTP URL or IPC Queue).""" + if await is_local_worker(worker_config): + comm_mode = worker_config.get('communicationMode', 'http') + + if comm_mode == 'ipc': + # Use multiprocessing Queue for IPC + if worker_id not in worker_manager.queues: + worker_manager.queues[worker_id] = Queue() + return worker_manager.queues[worker_id] + elif is_docker_environment(): + # Docker: use host.docker.internal + return f"http://host.docker.internal:{worker_config['port']}" + else: + # Local: use loopback + return f"http://127.0.0.1:{worker_config['port']}" + elif worker_config.get('type') == 'cloud' and is_runpod_environment(): + # Runpod same-host optimization (if detected) + host = worker_config.get('host', 'localhost') + return f"http://{host}:{worker_config['port']}" + else: + # Remote worker: use configured endpoint + host = worker_config.get('host', 'localhost') + return f"http://{host}:{worker_config['port']}" + # Auto-launch workers if enabled def auto_launch_workers(): """Launch enabled workers if auto_launch_workers is set to true.""" @@ -1335,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 file and return it as base64 data.""" + """Load an image or video file and return it as base64 data.""" try: data = await request.json() image_path = data.get("image_path") @@ -1345,30 +1474,58 @@ async def load_image_endpoint(request): import folder_paths import base64 - from PIL import Image import io - # Use ComfyUI's folder paths to find the image + # Use ComfyUI's folder paths to find the file 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) + return await handle_api_error(request, f"File 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') + # Check if it's a video file + video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.webm'} + file_ext = os.path.splitext(full_path)[1].lower() + + if file_ext in video_extensions: + # For video files, read the raw bytes + with open(full_path, 'rb') as f: + file_data = f.read() - # 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}" - }) + # Determine MIME type + mime_types = { + '.mp4': 'video/mp4', + '.avi': 'video/x-msvideo', + '.mov': 'video/quicktime', + '.mkv': 'video/x-matroska', + '.webm': 'video/webm' + } + mime_type = mime_types.get(file_ext, 'video/mp4') + + # Return base64 encoded video with data URL + video_base64 = base64.b64encode(file_data).decode('utf-8') + return web.json_response({ + "status": "success", + "image_data": f"data:{mime_type};base64,{video_base64}" + }) + else: + # For images, use PIL + from PIL import Image + + # 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) diff --git a/web/executionUtils.js b/web/executionUtils.js index bbf700b..56ad8c2 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -9,6 +9,12 @@ import { TIMEOUTS } from './constants.js'; * @returns {Object} The converted API prompt */ function convertPathsForPlatform(apiPrompt, targetSeparator) { + // Validate target separator + if (!targetSeparator || (targetSeparator !== '/' && targetSeparator !== '\\')) { + console.warn('[Distributed] Invalid target separator:', targetSeparator, '- skipping path conversion'); + return apiPrompt; + } + // Regex to identify likely file paths with extensions const isLikelyFilename = (value) => { return value.match(/\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)(\s*\[\w+\])?$/i); @@ -152,12 +158,15 @@ export async function prepareApiPromptForParticipant(extension, baseApiPrompt, p try { const workerUrl = extension.getWorkerUrl(workerInfo); const systemInfo = await getCachedWorkerSystemInfo(workerUrl); - const targetSeparator = systemInfo.platform.path_separator; + const targetSeparator = systemInfo?.platform?.path_separator; - // Convert paths to match worker's platform - jobApiPrompt = convertPathsForPlatform(jobApiPrompt, targetSeparator); - - extension.log(`Converted paths for ${systemInfo.platform.system} worker ${participantId} (separator: '${targetSeparator}')`, "debug"); + if (targetSeparator) { + // Convert paths to match worker's platform + jobApiPrompt = convertPathsForPlatform(jobApiPrompt, targetSeparator); + extension.log(`Converted paths for ${systemInfo.platform.system} worker ${participantId} (separator: '${targetSeparator}')`, "debug"); + } else { + extension.log(`No path separator found for worker ${participantId}, skipping path conversion`, "debug"); + } } catch (e) { extension.log(`Failed to get system info for worker ${participantId}: ${e.message}`, "warn"); // Continue without path conversion @@ -178,10 +187,10 @@ export async function prepareApiPromptForParticipant(extension, baseApiPrompt, p const isRemote = workerInfo && workerInfo.host; if (isRemote) { - // Find all image references in the pruned workflow + // Find all image/video references in the pruned workflow const imageReferences = findImageReferences(extension, jobApiPrompt); if (imageReferences.size > 0) { - extension.log(`Found ${imageReferences.size} image references for remote worker ${workerId}`, "debug"); + extension.log(`Found ${imageReferences.size} media references (images/videos) for remote worker ${workerId}`, "debug"); // Store image references for later processing options._imageReferences = imageReferences; } @@ -284,7 +293,7 @@ export async function executeJobs(extension, jobs) { } if (allImageReferences.size > 0) { - extension.log(`Pre-loading ${allImageReferences.size} unique image(s) for all workers`, "debug"); + extension.log(`Pre-loading ${allImageReferences.size} unique media file(s) for all workers`, "debug"); await loadImagesForWorker(extension, allImageReferences); } @@ -414,12 +423,12 @@ export async function uploadImagesToWorker(extension, workerUrl, images) { for (const imageData of images) { const formData = new FormData(); - // Detect MIME type from base64 header - const mimeMatch = imageData.image.match(/^data:(image\/\w+);base64,/); + // Detect MIME type from base64 header (supports both image and video) + const mimeMatch = imageData.image.match(/^data:((?:image|video)\/\w+);base64,/); const mimeType = mimeMatch ? mimeMatch[1] : 'image/png'; // Default to PNG if not detected // Convert base64 to blob - const base64Data = imageData.image.replace(/^data:image\/\w+;base64,/, ''); + const base64Data = imageData.image.replace(/^data:(?:image|video)\/\w+;base64,/, ''); const byteCharacters = atob(base64Data); const byteNumbers = new Array(byteCharacters.length); for (let i = 0; i < byteCharacters.length; i++) { diff --git a/web/workerUtils.js b/web/workerUtils.js index 63280f7..ef87542 100644 --- a/web/workerUtils.js +++ b/web/workerUtils.js @@ -88,21 +88,30 @@ export function findImageReferences(extension, apiPrompt) { // - 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; + // - Video files: "video.mp4", "animation.avi", etc. + const imageExtensions = /\.(png|jpg|jpeg|gif|webp|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$/i; for (const [nodeId, node] of Object.entries(apiPrompt)) { + // Check for both 'image' and 'video' inputs + const mediaInputs = []; if (node.inputs && node.inputs.image) { - const imageValue = node.inputs.image; - if (typeof imageValue === 'string') { + mediaInputs.push(node.inputs.image); + } + if (node.inputs && node.inputs.video) { + mediaInputs.push(node.inputs.video); + } + + for (const mediaValue of mediaInputs) { + if (typeof mediaValue === 'string') { // Clean special suffixes like [input] or [output] - const cleanValue = imageValue.replace(/\s*\[\w+\]$/, '').trim(); + const cleanValue = mediaValue.replace(/\s*\[\w+\]$/, '').trim(); if (imageExtensions.test(cleanValue)) { images.set(cleanValue, { nodeId, nodeType: node.class_type, - inputName: 'image' + inputName: 'image' // Keep as 'image' for compatibility }); - extension.log(`Found image reference: ${cleanValue} in node ${nodeId} (${node.class_type})`, "debug"); + extension.log(`Found media reference: ${cleanValue} in node ${nodeId} (${node.class_type})`, "debug"); } } }