support sending videos to remote workers

2025-07-17 10:23:27 +10:00
parent 5703b62062
commit abcbab7574
3 changed files with 210 additions and 35 deletions
@@ -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)
@@ -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++) {
@@ -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");
                }
            }
        }