diff --git a/__init__.py b/__init__.py index 9bc2dde..18d9765 100644 --- a/__init__.py +++ b/__init__.py @@ -4,6 +4,36 @@ import sys # Add the directory to Python path sys.path.append(os.path.dirname(__file__)) +# Patch ComfyUI execution validation for DistributedVideoCollector +try: + import execution + from .distributed import ImageBatchDivider + + # Store original validate function if it exists + if hasattr(execution, 'validate_outputs'): + original_validate_outputs = execution.validate_outputs + + def patched_validate_outputs(executor, node_id, result, node_class): + if node_class == ImageBatchDivider: + return # Skip validation for our dynamic output node + return original_validate_outputs(executor, node_id, result, node_class) + + execution.validate_outputs = patched_validate_outputs + else: + # Fallback: patch the main execute method if validate_outputs doesn't exist + if hasattr(execution.PromptExecutor, 'execute'): + original_execute = execution.PromptExecutor.execute + + def patched_execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]): + # This is a more complex patch - for now just call original + # The ByPassTypeTuple should handle most validation issues + return original_execute(self, prompt, prompt_id, extra_data, execute_outputs) + + execution.PromptExecutor.execute = patched_execute + +except ImportError: + pass # ComfyUI execution module not available during import + # Import everything needed from the main module from .distributed import ( NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS, diff --git a/distributed.py b/distributed.py index 3189eab..c1d5c55 100644 --- a/distributed.py +++ b/distributed.py @@ -14,6 +14,11 @@ import platform import time import atexit import signal +import sys +import shlex +import uuid +from multiprocessing import Queue +from comfy.utils import ProgressBar # Import shared utilities from .utils.logging import debug_log, log @@ -25,7 +30,7 @@ from .utils.async_helpers import run_async_in_server_loop from .utils.constants import ( WORKER_JOB_TIMEOUT, PROCESS_TERMINATION_TIMEOUT, WORKER_CHECK_INTERVAL, STATUS_CHECK_INTERVAL, CHUNK_SIZE, LOG_TAIL_BYTES, WORKER_LOG_PATTERN, - WORKER_STARTUP_DELAY, PROCESS_WAIT_TIMEOUT, MEMORY_CLEAR_DELAY + WORKER_STARTUP_DELAY, PROCESS_WAIT_TIMEOUT, MEMORY_CLEAR_DELAY, MAX_BATCH ) # Try to import psutil for better process management @@ -95,6 +100,53 @@ async def get_network_info_endpoint(request): """Get network interfaces and recommend best IP for master.""" import socket + # Get CUDA device if available + cuda_device = None + cuda_device_count = 0 + physical_device_count = 0 + + if torch.cuda.is_available(): + try: + import os + import subprocess + + # Get visible device count (what PyTorch sees) + cuda_device_count = torch.cuda.device_count() + + # Try to get actual physical device info + cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES', '') + + # Method 1: Parse CUDA_VISIBLE_DEVICES + if cuda_visible and cuda_visible.strip(): + visible_devices = [int(d.strip()) for d in cuda_visible.split(',') if d.strip().isdigit()] + if visible_devices: + # Get the first visible device as the actual physical device + cuda_device = visible_devices[0] + + # Try to get total physical device count using nvidia-smi + try: + result = subprocess.run(['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'], + capture_output=True, text=True, timeout=5) + if result.returncode == 0: + physical_device_count = len(result.stdout.strip().split('\n')) + else: + physical_device_count = max(visible_devices) + 1 # Best guess + except: + physical_device_count = max(visible_devices) + 1 # Best guess + else: + cuda_device = 0 + physical_device_count = cuda_device_count + else: + # No CUDA_VISIBLE_DEVICES set, current device is actual device + cuda_device = torch.cuda.current_device() + physical_device_count = cuda_device_count + + except Exception as e: + debug_log(f"CUDA detection error: {e}") + cuda_device = None + cuda_device_count = 0 + physical_device_count = 0 + def get_network_ips(): """Get all network IPs, trying multiple methods.""" ips = [] @@ -195,6 +247,8 @@ async def get_network_info_endpoint(request): "hostname": hostname, "all_ips": all_ips, "recommended_ip": recommended_ip, + "cuda_device": cuda_device, + "cuda_device_count": physical_device_count if physical_device_count > 0 else cuda_device_count, "message": "Auto-detected network configuration" }) except Exception as e: @@ -206,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: @@ -248,6 +329,11 @@ async def update_worker_endpoint(request): worker.pop("extra_args", None) else: worker["extra_args"] = data["extra_args"] + + # Handle type field + if "type" in data: + worker["type"] = data["type"] + worker_found = True break @@ -261,7 +347,8 @@ async def update_worker_endpoint(request): "port": data["port"], "cuda_device": data["cuda_device"], "enabled": data.get("enabled", False), - "extra_args": data.get("extra_args", "") + "extra_args": data.get("extra_args", ""), + "type": data.get("type", "local") } if "workers" not in config: config["workers"] = [] @@ -347,6 +434,8 @@ async def update_master_endpoint(request): config['master'] = {} # Update all provided fields + if "name" in data: + config['master']['name'] = data['name'] if "host" in data: config['master']['host'] = data['host'] if "port" in data: @@ -522,6 +611,79 @@ async def get_managed_workers_endpoint(request): return await handle_api_error(request, e, 500) +@server.PromptServer.instance.routes.get("/distributed/local-worker-status") +async def get_local_worker_status_endpoint(request): + """Check status of all local workers (localhost/no host specified).""" + try: + config = load_config() + worker_statuses = {} + + for worker in config.get("workers", []): + # Only check local workers + if not worker.get("host") or worker.get("host") in ["localhost", "127.0.0.1"]: + worker_id = worker["id"] + port = worker["port"] + + # Check if worker is enabled + if not worker.get("enabled", False): + worker_statuses[worker_id] = { + "online": False, + "enabled": False, + "processing": False, + "queue_count": 0 + } + continue + + # Try to connect to worker + try: + session = await get_client_session() + async with session.get( + f"http://localhost:{port}/prompt", + timeout=aiohttp.ClientTimeout(total=2) + ) as resp: + if resp.status == 200: + data = await resp.json() + queue_remaining = data.get("exec_info", {}).get("queue_remaining", 0) + worker_statuses[worker_id] = { + "online": True, + "enabled": True, + "processing": queue_remaining > 0, + "queue_count": queue_remaining + } + else: + worker_statuses[worker_id] = { + "online": False, + "enabled": True, + "processing": False, + "queue_count": 0, + "error": f"HTTP {resp.status}" + } + except asyncio.TimeoutError: + worker_statuses[worker_id] = { + "online": False, + "enabled": True, + "processing": False, + "queue_count": 0, + "error": "Timeout" + } + except Exception as e: + worker_statuses[worker_id] = { + "online": False, + "enabled": True, + "processing": False, + "queue_count": 0, + "error": str(e) + } + + return web.json_response({ + "status": "success", + "worker_statuses": worker_statuses + }) + except Exception as e: + debug_log(f"Error checking local worker status: {e}") + return await handle_api_error(request, e, 500) + + @server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}") async def get_worker_log_endpoint(request): """Get log content for a specific worker.""" @@ -723,9 +885,20 @@ class WorkerProcessManager: raise RuntimeError(error_msg) - # Add any extra arguments + # Add any extra arguments safely if worker_config.get('extra_args'): - cmd.extend(worker_config['extra_args'].split()) + raw_args = worker_config['extra_args'].strip() + if raw_args: + # Safely split using shlex to handle quotes and spaces + extra_args_list = shlex.split(raw_args) + + # Validate: Block dangerous shell meta-characters (e.g., ; | & > < ` $) + forbidden_chars = set(';|>&<`$()[]{}*!?') + for arg in extra_args_list: + if any(c in forbidden_chars for c in arg): + raise ValueError(f"Invalid characters in extra_args: {arg}. Forbidden: {''.join(forbidden_chars)}") + + cmd.extend(extra_args_list) return cmd @@ -1085,6 +1258,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.""" @@ -1152,15 +1425,9 @@ def delayed_auto_launch(): timer.daemon = True timer.start() -# Call delayed auto-launch only if we're the master (not a worker) -if not os.environ.get('COMFYUI_MASTER_PID'): - delayed_auto_launch() -else: - debug_log("Running as worker, skipping auto-launch") - -# Register cleanup on exit - only clean up if setting is enabled -def cleanup_on_exit(signum=None, frame=None): - """Handle cleanup on exit or signal""" +# Async cleanup function for proper shutdown +async def async_cleanup_and_exit(signum=None): + """Async-friendly cleanup and exit.""" try: config = load_config() if config.get('settings', {}).get('stop_workers_on_master_exit', True): @@ -1168,24 +1435,95 @@ def cleanup_on_exit(signum=None, frame=None): worker_manager.cleanup_all() else: print("\n[Distributed] Master shutting down, workers will continue running") - # Still save the current state + worker_manager.save_processes() + except Exception as e: + print(f"[Distributed] Error during cleanup: {e}") + + # On Windows, we need to exit differently + if platform.system() == "Windows": + # Force exit on Windows + sys.exit(0) + else: + # On Unix, stop the event loop gracefully + loop = asyncio.get_running_loop() + loop.stop() + +def register_async_signals(): + """Register async signal handlers for graceful shutdown.""" + if platform.system() == "Windows": + # Windows doesn't support add_signal_handler, use traditional signal handling + def signal_handler(signum, frame): + # Schedule the async cleanup in the event loop + loop = server.PromptServer.instance.loop + if loop and loop.is_running(): + asyncio.run_coroutine_threadsafe(async_cleanup_and_exit(signum), loop) + else: + # Fallback to sync cleanup if loop isn't running + try: + config = load_config() + if config.get('settings', {}).get('stop_workers_on_master_exit', True): + print("\n[Distributed] Master shutting down, stopping all managed workers...") + worker_manager.cleanup_all() + else: + print("\n[Distributed] Master shutting down, workers will continue running") + worker_manager.save_processes() + except Exception as e: + print(f"[Distributed] Error during cleanup: {e}") + sys.exit(0) + + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + else: + # Unix-like systems support add_signal_handler + loop = server.PromptServer.instance.loop + + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, lambda s=sig: asyncio.create_task(async_cleanup_and_exit(s))) + + # SIGHUP is Unix-only + loop.add_signal_handler(signal.SIGHUP, lambda: asyncio.create_task(async_cleanup_and_exit(signal.SIGHUP))) + +def sync_cleanup(): + """Synchronous wrapper for atexit.""" + try: + # For atexit, we don't want to stop the loop or exit + config = load_config() + if config.get('settings', {}).get('stop_workers_on_master_exit', True): + print("\n[Distributed] Master shutting down, stopping all managed workers...") + worker_manager.cleanup_all() + else: + print("\n[Distributed] Master shutting down, workers will continue running") worker_manager.save_processes() except Exception as e: print(f"[Distributed] Error during cleanup: {e}") -# Register cleanup handlers -atexit.register(cleanup_on_exit) +# Register atexit handler for normal exits +atexit.register(sync_cleanup) -# Handle terminal window closing and Ctrl+C -try: - signal.signal(signal.SIGINT, cleanup_on_exit) - signal.signal(signal.SIGTERM, cleanup_on_exit) - - if platform.system() != "Windows": - # SIGHUP is sent when terminal closes on Unix - signal.signal(signal.SIGHUP, cleanup_on_exit) -except Exception as e: - print(f"[Distributed] Warning: Could not set signal handlers: {e}") +# Call delayed auto-launch only if we're the master (not a worker) +if not os.environ.get('COMFYUI_MASTER_PID'): + delayed_auto_launch() + try: + register_async_signals() # Register async signal handlers + except Exception as e: + print(f"[Distributed] Warning: Could not register async signal handlers: {e}") + # Fall back to basic signal handling + def basic_signal_handler(signum, frame): + print("\n[Distributed] Received signal, shutting down...") + try: + config = load_config() + if config.get('settings', {}).get('stop_workers_on_master_exit', True): + worker_manager.cleanup_all() + else: + worker_manager.save_processes() + except Exception as cleanup_error: + print(f"[Distributed] Error during cleanup: {cleanup_error}") + sys.exit(0) + + signal.signal(signal.SIGINT, basic_signal_handler) + signal.signal(signal.SIGTERM, basic_signal_handler) +else: + debug_log("Running as worker, skipping auto-launch") # --- Persistent State Storage --- # Store job queue on the persistent server instance to survive script reloads @@ -1199,7 +1537,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 with hash.""" try: data = await request.json() image_path = data.get("image_path") @@ -1209,71 +1547,244 @@ async def load_image_endpoint(request): import folder_paths import base64 - from PIL import Image import io + import hashlib - # 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') + # 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() + + 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') + # 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}", + "hash": file_hash + }) + 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}", + "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", - "image_data": f"data:image/png;base64,{img_base64}" + "exists": True, + "hash_matches": hash_matches }) 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: 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 multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) - debug_log(f"job_complete received - job_id: {multi_job_id}, worker: {worker_id}, index: {image_index}, is_last: {is_last}") - - if not all([multi_job_id, image_file]): - return await handle_api_error(request, "Missing job_id or image data", 400) - - # Process image with error handling - try: - img_data = image_file.file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - # Convert to tensor using utility function - tensor = pil_to_tensor(img) - # Ensure tensor is contiguous - tensor = ensure_contiguous(tensor) - except Exception as e: - log(f"Error processing image from worker {worker_id}: {e}") - return await handle_api_error(request, f"Image processing error: {e}", 400) + # Check for batch mode + batch_size = int(data.get('batch_size', 0)) + tensors = [] + + if batch_size > 0: + # Batch mode: Extract multiple images + debug_log(f"job_complete received batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}") + + # Check for JSON metadata + metadata_field = data.get('images_metadata') + metadata = None + if metadata_field: + # New JSON metadata format + try: + # Handle different types of metadata field + if hasattr(metadata_field, 'file'): + # File-like object + metadata_str = metadata_field.file.read().decode('utf-8') + elif isinstance(metadata_field, (bytes, bytearray)): + # Direct bytes/bytearray + metadata_str = metadata_field.decode('utf-8') + else: + # String + metadata_str = str(metadata_field) + + metadata = json.loads(metadata_str) + if len(metadata) != batch_size: + return await handle_api_error(request, "Metadata length mismatch", 400) + debug_log(f"Using JSON metadata for batch from worker {worker_id}") + except Exception as e: + log(f"Error parsing JSON metadata from worker {worker_id}: {e}") + return await handle_api_error(request, f"Metadata parsing error: {e}", 400) + else: + # Legacy format - log deprecation warning + debug_log(f"WARNING: Worker {worker_id} using legacy field format. Please update to use JSON metadata.") + + # Process images with per-item error handling + image_data_list = [] # List of tuples: (index, tensor) + + for i in range(batch_size): + image_field = data.get(f'image_{i}') + if image_field is None: + log(f"Missing image_{i} from worker {worker_id}, skipping") + continue + + try: + img_data = image_field.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + tensor = pil_to_tensor(img) + tensor = ensure_contiguous(tensor) + + # Get expected index from metadata or use position + if metadata and i < len(metadata): + expected_idx = metadata[i].get('index', i) + if i != expected_idx: + debug_log(f"Warning: Image order mismatch at position {i}, expected index {expected_idx}") + else: + expected_idx = i + + image_data_list.append((expected_idx, tensor)) + except Exception as e: + log(f"Error processing image {i} from worker {worker_id}: {e}, skipping this image") + # Continue processing other images instead of failing entire batch + continue + + # Check if we got any valid images + if not image_data_list: + return await handle_api_error(request, "No valid images in batch", 400) + + # Sort by index to ensure correct order + image_data_list.sort(key=lambda x: x[0]) + tensors = [tensor for _, tensor in image_data_list] + + # Validate final order + if metadata: + debug_log(f"Reordered {len(tensors)} images based on metadata indices") + else: + # Fallback for single-image mode (backward compatibility) + image_file = data.get('image') + image_index = data.get('image_index') + + debug_log(f"job_complete received single - job_id: {multi_job_id}, worker: {worker_id}, index: {image_index}, is_last: {is_last}") + + if not image_file: + return await handle_api_error(request, "Missing image data", 400) + + try: + img_data = image_file.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + tensor = pil_to_tensor(img) + tensor = ensure_contiguous(tensor) + tensors = [tensor] + except Exception as e: + log(f"Error processing image from worker {worker_id}: {e}") + return await handle_api_error(request, f"Image processing error: {e}", 400) + # Put batch into queue async with prompt_server.distributed_jobs_lock: debug_log(f"Current pending jobs: {list(prompt_server.distributed_pending_jobs.keys())}") if multi_job_id in prompt_server.distributed_pending_jobs: - await prompt_server.distributed_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 - }) - debug_log(f"Received result for job {multi_job_id} from worker {worker_id} (last: {is_last})") + if batch_size > 0: + # Put batch as single item + await prompt_server.distributed_pending_jobs[multi_job_id].put({ + 'worker_id': worker_id, + 'tensors': tensors, + 'is_last': is_last + }) + debug_log(f"Received batch result for job {multi_job_id} from worker {worker_id}, size={len(tensors)}") + else: + # Put single image (backward compat) + await prompt_server.distributed_pending_jobs[multi_job_id].put({ + 'tensor': tensors[0], + 'worker_id': worker_id, + 'image_index': int(image_index) if image_index else 0, + 'is_last': is_last + }) + debug_log(f"Received single result for job {multi_job_id} from worker {worker_id}") + debug_log(f"Queue size after put: {prompt_server.distributed_pending_jobs[multi_job_id].qsize()}") return web.json_response({"status": "success"}) else: @@ -1316,43 +1827,55 @@ class DistributedCollectorNode: ) return 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""" - # Ensure we handle the tensor shape correctly (H, W, C) - if image_tensor.dim() == 4: # Has batch dimension - img = tensor_to_pil(image_tensor, 0) - else: # Single image without batch dimension - # Add batch dimension temporarily for tensor_to_pil - img = tensor_to_pil(image_tensor.unsqueeze(0), 0) - byte_io = io.BytesIO() - # Use PNG with no compression for lossless transfer - img.save(byte_io, format='PNG', compress_level=0) - byte_io.seek(0) + async def send_batch_to_master(self, image_batch, multi_job_id, master_url, worker_id): + """Send image batch to master, chunked if large.""" + batch_size = image_batch.shape[0] + if batch_size == 0: + return - 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') + + for start in range(0, batch_size, MAX_BATCH): + chunk = image_batch[start:start + MAX_BATCH] + chunk_size = chunk.shape[0] + is_chunk_last = (start + chunk_size == batch_size) # True only for final chunk + + + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('is_last', str(is_chunk_last)) + data.add_field('batch_size', str(chunk_size)) + + # Chunk metadata: Absolute index from full batch + metadata = [{'index': start + j} for j in range(chunk_size)] + data.add_field('images_metadata', json.dumps(metadata), content_type='application/json') + + # Add chunk images + for j in range(chunk_size): + # Convert tensor slice to PIL + img = tensor_to_pil(chunk[j:j+1], 0) + byte_io = io.BytesIO() + img.save(byte_io, format='PNG', compress_level=0) + byte_io.seek(0) + data.add_field(f'image_{j}', byte_io, filename=f'image_{j}.png', content_type='image/png') + + try: + + session = await get_client_session() + url = f"{master_url}/distributed/job_complete" + async with session.post(url, data=data) as response: + response.raise_for_status() + except Exception as e: + log(f"Worker - Failed to send chunk to master: {e}") + debug_log(f"Worker - Full error details: URL={url}") + raise # Re-raise to handle at caller level - try: - session = await get_client_session() - async with session.post(f"{master_url}/distributed/job_complete", data=data) as response: - response.raise_for_status() - except Exception as e: - log(f"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 + # Worker mode: send images to master in a single batch debug_log(f"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) - + await self.send_batch_to_master(images, multi_job_id, master_url, worker_id) return (images,) else: # Master mode: collect images from workers @@ -1368,8 +1891,6 @@ class DistributedCollectorNode: # Ensure master images are contiguous images_on_cpu = ensure_contiguous(images_on_cpu) - # Debug: Check master tensor properties - debug_log(f"Master tensor - shape: {images_on_cpu.shape}, dtype: {images_on_cpu.dtype}, device: {images_on_cpu.device}") # Initialize storage for collected images worker_images = {} # Dict to store images by worker_id and index @@ -1390,14 +1911,15 @@ class DistributedCollectorNode: # Use a reasonable timeout for the first image timeout = WORKER_JOB_TIMEOUT - debug_log(f"Master - Starting collection loop, expecting {num_workers} workers") # Get queue size before starting async with prompt_server.distributed_jobs_lock: q = prompt_server.distributed_pending_jobs[multi_job_id] initial_size = q.qsize() - debug_log(f"Master - Queue size before collection: {initial_size}") - + + # NEW: Initialize progress bar for workers (total = num_workers) + p = ProgressBar(num_workers) + while len(workers_done) < num_workers: try: # Get the queue again each time to ensure we have the right reference @@ -1405,34 +1927,48 @@ class DistributedCollectorNode: q = prompt_server.distributed_pending_jobs[multi_job_id] current_size = q.qsize() - debug_log(f"Master - Waiting for queue item, timeout={timeout}s, queue size={current_size}") 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) - debug_log(f"Master - Got result from worker {worker_id}, image {image_index}, is_last={is_last}") - - if worker_id not in worker_images: - worker_images[worker_id] = {} - worker_images[worker_id][image_index] = tensor - - # Debug: Check worker tensor properties - if collected_count == 0: # Only print once - debug_log(f"Master - Worker tensor - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}") - - collected_count += 1 + # Check if batch mode + tensors = result.get('tensors', []) + if tensors: + # Batch mode + debug_log(f"Master - Got batch from worker {worker_id}, size={len(tensors)}, is_last={is_last}") + + if worker_id not in worker_images: + worker_images[worker_id] = {} + + # Add all tensors with sequential indices + for idx, tensor in enumerate(tensors): + worker_images[worker_id][idx] = tensor + + + collected_count += len(tensors) + else: + # Single image mode (backward compat) + image_index = result['image_index'] + tensor = result['tensor'] + + debug_log(f"Master - Got single result from worker {worker_id}, image {image_index}, is_last={is_last}") + + 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 = WORKER_JOB_TIMEOUT if is_last: workers_done.add(worker_id) - debug_log(f"Master - Worker {worker_id} done. Collected {len(worker_images[worker_id])} images") + p.update(1) # +1 per completed worker else: - debug_log(f"Master - Collected image {image_index + 1} from worker {worker_id}") + pass # Continue waiting for more results except asyncio.TimeoutError: missing_workers = set(str(w) for w in enabled_workers) - workers_done @@ -1443,7 +1979,6 @@ class DistributedCollectorNode: if multi_job_id in prompt_server.distributed_pending_jobs: final_q = prompt_server.distributed_pending_jobs[multi_job_id] final_size = final_q.qsize() - debug_log(f"Master - Queue size at timeout: {final_size}") # Try to drain any remaining items remaining_items = [] @@ -1455,31 +1990,41 @@ class DistributedCollectorNode: break if remaining_items: - debug_log(f"Master - Found {len(remaining_items)} items in queue after timeout!") # Process them for item in remaining_items: worker_id = item['worker_id'] - image_index = item['image_index'] - tensor = item['tensor'] is_last = item.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 + # Check if batch mode + tensors = item.get('tensors', []) + if tensors: + # Batch mode + if worker_id not in worker_images: + worker_images[worker_id] = {} + + for idx, tensor in enumerate(tensors): + worker_images[worker_id][idx] = tensor + + collected_count += len(tensors) + else: + # Single image mode + image_index = item['image_index'] + tensor = item['tensor'] + + if worker_id not in worker_images: + worker_images[worker_id] = {} + worker_images[worker_id][image_index] = tensor + + collected_count += 1 if is_last: workers_done.add(worker_id) - debug_log(f"Master - Worker {worker_id} done (found in timeout drain)") + p.update(1) # +1 here too else: log(f"Master - Queue {multi_job_id} no longer exists!") break total_collected = sum(len(imgs) for imgs in worker_images.values()) - debug_log(f"Master - Collection complete. Received {total_collected} images from {len(workers_done)} workers") - debug_log(f"Master - Workers done: {workers_done}, Enabled workers: {enabled_workers}") - debug_log(f"Master - Worker images keys: {list(worker_images.keys())}") # Clean up job queue async with prompt_server.distributed_jobs_lock: @@ -1519,7 +2064,6 @@ class DistributedCollectorNode: return (combined,) except Exception as e: log(f"Master - Error combining images: {e}") - debug_log(f"Master - Tensor shapes: {[t.shape for t in cpu_tensors]}") # Return just the master images as fallback return (images,) @@ -1578,11 +2122,87 @@ class DistributedSeed: # Fallback: return original seed return (seed,) +# Define ByPassTypeTuple for flexible return types +class AnyType(str): + def __ne__(self, __value: object) -> bool: + return False + +any_type = AnyType("*") + +class ByPassTypeTuple(tuple): + def __getitem__(self, index): + if index > 0: + index = 0 + item = super().__getitem__(index) + if isinstance(item, str): + return any_type + return item + +class ImageBatchDivider: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "divide_by": ("INT", { + "default": 2, + "min": 1, + "max": 10, + "step": 1, + "display": "number", + "tooltip": "Number of parts to divide the batch into" + }), + } + } + + RETURN_TYPES = ByPassTypeTuple(("IMAGE", )) # Flexible for variable outputs + RETURN_NAMES = ByPassTypeTuple(tuple([f"batch_{i+1}" for i in range(10)])) + FUNCTION = "divide_batch" + OUTPUT_NODE = True + CATEGORY = "image" + + def divide_batch(self, images, divide_by): + import torch + + # Use divide_by directly + total_splits = divide_by + + if total_splits > 10: + total_splits = 10 # Cap to max + + # Get total number of frames + total_frames = images.shape[0] + frames_per_split = total_frames // total_splits + remainder = total_frames % total_splits + + outputs = [] + start_idx = 0 + + for i in range(total_splits): + current_frames = frames_per_split + (1 if i < remainder else 0) + end_idx = start_idx + current_frames + split_frames = images[start_idx:end_idx] + outputs.append(split_frames) + start_idx = end_idx + + # Pad with empty tensors up to max (10) to match potential RETURN_TYPES len + empty_shape = (1, images.shape[1], images.shape[2], images.shape[3]) if total_frames > 0 else (1, 512, 512, 3) + empty_tensor = torch.zeros(empty_shape, dtype=images.dtype if total_frames > 0 else torch.float32, + device=images.device if total_frames > 0 else 'cpu') + + while len(outputs) < 10: + outputs.append(empty_tensor) + + return tuple(outputs) + + NODE_CLASS_MAPPINGS = { "DistributedCollector": DistributedCollectorNode, - "DistributedSeed": DistributedSeed + "DistributedSeed": DistributedSeed, + "ImageBatchDivider": ImageBatchDivider } NODE_DISPLAY_NAME_MAPPINGS = { "DistributedCollector": "Distributed Collector", - "DistributedSeed": "Distributed Seed" + "DistributedSeed": "Distributed Seed", + "ImageBatchDivider": "Image Batch Divider" } diff --git a/distributed_upscale.py b/distributed_upscale.py index 7bf702b..75a0eb9 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -1,3 +1,4 @@ +# Import statements (original + new) import torch import numpy as np from PIL import Image, ImageFilter, ImageDraw @@ -7,10 +8,15 @@ import aiohttp from aiohttp import web import io import server +import math +import os +import time from typing import List, Tuple +from functools import wraps # Import ComfyUI modules import comfy.samplers +import comfy.model_management # Import shared utilities from .utils.logging import debug_log, log @@ -20,25 +26,61 @@ from .utils.network import get_server_port, get_server_loop, get_client_session, from .utils.async_helpers import run_async_in_server_loop from .utils.constants import ( TILE_COLLECTION_TIMEOUT, TILE_WAIT_TIMEOUT, TILE_TRANSFER_TIMEOUT, - QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT + QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT, MAX_BATCH, HEARTBEAT_TIMEOUT ) -# Helper function to ensure persistent state is initialized -def ensure_tile_jobs_initialized(): - """Ensure tile job storage is initialized on the server instance.""" - prompt_server = server.PromptServer.instance - if not hasattr(prompt_server, 'distributed_pending_tile_jobs'): - debug_log("Initializing persistent tile job queue on server instance.") - prompt_server.distributed_pending_tile_jobs = {} - prompt_server.distributed_tile_jobs_lock = asyncio.Lock() - return prompt_server +# Import for controller support +from .utils.usdu_utils import crop_cond +from .utils.usdu_managment import ( + clone_conditioning, ensure_tile_jobs_initialized, + # Job management functions + _init_job_queue, _distribute_tasks, _get_next_task, _drain_results_queue, + _check_and_requeue_timed_out_workers, _get_completed_count, _mark_task_completed, + _send_heartbeat_to_master, _cleanup_job, + # Constants + JOB_COMPLETED_TASKS, JOB_WORKER_STATUS, JOB_PENDING_TASKS +) + + +# Note: MAX_BATCH and HEARTBEAT_TIMEOUT are imported from utils.constants +# They can be overridden via environment variables: +# - COMFYUI_MAX_BATCH (default: 20) +# - COMFYUI_HEARTBEAT_TIMEOUT (default: 120) + +# Sync wrapper decorator for async methods +def sync_wrapper(async_func): + """Decorator to wrap async methods for synchronous execution.""" + @wraps(async_func) + def sync_func(self, *args, **kwargs): + # Use run_async_in_server_loop for ComfyUI compatibility + return run_async_in_server_loop( + async_func(self, *args, **kwargs), + timeout=600.0 # 10 minute timeout for long operations + ) + return sync_func # Note: tensor_to_pil and pil_to_tensor are imported from utils.image class UltimateSDUpscaleDistributed: """ Distributed version of Ultimate SD Upscale (No Upscale). - Distributes tile processing across multiple workers. + + Supports three processing modes: + 1. Single GPU: No workers available, process everything locally + 2. Static Mode: Small batches, distributes tiles across workers (flattened) + 3. Dynamic Mode: Large batches, assigns whole images to workers dynamically + + Features: + - Multi-mode batch handling for efficient video/image upscaling + - Tiled VAE support for memory efficiency + - Dynamic load balancing for large batches + - Backward compatible with single-image workflows + + Environment Variables: + - COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20) + - COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB) + + Threshold: dynamic_threshold input controls mode switch (default 8) """ def __init__(self): @@ -67,6 +109,8 @@ class UltimateSDUpscaleDistributed: "padding": ("INT", {"default": 32, "min": 0, "max": 256, "step": 8}), "mask_blur": ("INT", {"default": 8, "min": 0, "max": 256}), "force_uniform_tiles": ("BOOLEAN", {"default": True}), + "tiled_decode": ("BOOLEAN", {"default": False}), + "static_distribution": ("BOOLEAN", {"default": False, "label": "Use Static Distribution (legacy)"}), }, "hidden": { "multi_job_id": ("STRING", {"default": ""}), @@ -75,6 +119,7 @@ class UltimateSDUpscaleDistributed: "enabled_worker_ids": ("STRING", {"default": "[]"}), "worker_id": ("STRING", {"default": ""}), "tile_indices": ("STRING", {"default": ""}), # Unused - kept for compatibility + "dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}), }, } @@ -84,253 +129,120 @@ class UltimateSDUpscaleDistributed: @classmethod def IS_CHANGED(cls, **kwargs): - """Force re-execution and ensure queues exist for the job.""" - multi_job_id = kwargs.get('multi_job_id', '') - if multi_job_id: - # Initialize queue for this job immediately - prompt_server = ensure_tile_jobs_initialized() - loop = get_server_loop() - - async def init_queue(): - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - prompt_server.distributed_pending_tile_jobs[multi_job_id] = asyncio.Queue() - debug_log(f"UltimateSDUpscaleDistributed - Pre-initialized queue for job {multi_job_id}") - - try: - asyncio.run_coroutine_threadsafe(init_queue(), loop).result(timeout=1.0) - except Exception as e: - debug_log(f"UltimateSDUpscaleDistributed - Error pre-initializing queue: {e}") - + """Force re-execution.""" return float("nan") # Always re-execute + def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, - mask_blur, force_uniform_tiles, multi_job_id="", is_worker=False, - master_url="", enabled_worker_ids="[]", worker_id="", tile_indices=""): + mask_blur, force_uniform_tiles, tiled_decode, static_distribution=False, + multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", + worker_id="", tile_indices="", dynamic_threshold=8): """Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler.""" if not multi_job_id: # No distributed processing, run single GPU version return self.process_single_gpu(upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, force_uniform_tiles) + tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode) if is_worker: # Worker mode: process tiles synchronously - return self.process_worker_tiles(upscaled_image, model, positive, negative, vae, - seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, multi_job_id, master_url, - worker_id, enabled_worker_ids) + return self.process_worker(upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_worker_ids, dynamic_threshold, static_distribution) else: # Master mode: distribute and collect synchronously return self.process_master(upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, multi_job_id, enabled_worker_ids) + force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, + dynamic_threshold, static_distribution) - def process_worker_tiles(self, upscaled_image, model, positive, negative, vae, - seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, multi_job_id, master_url, - worker_id, enabled_worker_ids): - """Worker processing - SYNCHRONOUS tile processing with async communication.""" - # Round tile dimensions - tile_width = self.round_to_multiple(tile_width) - tile_height = self.round_to_multiple(tile_height) + def process_worker(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_worker_ids, dynamic_threshold, static_distribution=False): + """Unified worker processing - handles both static and dynamic modes.""" + # Get batch size to determine mode + batch_size = upscaled_image.shape[0] - # Calculate all tile positions - _, height, width, _ = upscaled_image.shape - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height) + # Ensure mode consistency across master/workers via shared threshold + # Determine mode (must match master's logic) + enabled_workers = json.loads(enabled_worker_ids) + num_workers = len(enabled_workers) - # Get assigned tiles for this worker - assigned_tiles = self._get_worker_tiles(all_tiles, enabled_worker_ids, worker_id) - if not assigned_tiles: - return (upscaled_image,) - - debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} processing {len(assigned_tiles)} tiles") - - # Process tiles SYNCHRONOUSLY - processed_tiles = [] - for tile_idx in assigned_tiles: - x, y = all_tiles[tile_idx] + mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold) - # Extract tile - tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( - upscaled_image, x, y, tile_width, tile_height, padding - ) - - # Process tile through SD (SYNCHRONOUS - no async!) - processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, - seed + tile_idx, steps, cfg, sampler_name, - scheduler, denoise) - - processed_tiles.append({ - 'tile': processed_tile, - 'tile_idx': tile_idx, - 'x': x1, # Use extraction position - 'y': y1, - 'extracted_width': ew, - 'extracted_height': eh - }) + debug_log(f"UltimateSDUpscale Worker - Mode: {mode}, batch_size: {batch_size}") - # Send results back using the server's event loop - try: - run_async_in_server_loop( - self._async_send_tiles_to_master(processed_tiles, multi_job_id, master_url, - padding, worker_id), - timeout=TILE_SEND_TIMEOUT - ) - except Exception as e: - log(f"UltimateSDUpscale Worker - Error sending tiles: {e}") + if mode == "dynamic": + return self.process_worker_dynamic(upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_worker_ids, dynamic_threshold) - return (upscaled_image,) - - async def _async_send_tiles_to_master(self, processed_tiles, multi_job_id, master_url, - padding, worker_id): - """Async helper to send tiles to master.""" - for i, tile_data in enumerate(processed_tiles): - is_last = (i == len(processed_tiles) - 1) - await self.send_tile_to_master( - tile_data['tile'], multi_job_id, master_url, - tile_data['tile_idx'], tile_data['x'], tile_data['y'], - tile_data['extracted_width'], tile_data['extracted_height'], - padding, worker_id, is_last - ) + # Static mode - enhanced with health monitoring and retry logic + return self._process_worker_static_sync(upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_workers, static_distribution) def process_master(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, multi_job_id, enabled_worker_ids): - """Master processing - SYNCHRONOUS with async collection.""" + force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, + dynamic_threshold, static_distribution=False): + """Unified master processing with enhanced monitoring and failure handling.""" # Round tile dimensions tile_width = self.round_to_multiple(tile_width) tile_height = self.round_to_multiple(tile_height) - # Get image dimensions - _, height, width, _ = upscaled_image.shape + # Get image dimensions and batch size + batch_size, height, width, _ = upscaled_image.shape # Calculate all tiles - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height) - total_tiles = len(all_tiles) + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + num_tiles_per_image = len(all_tiles) - debug_log(f"UltimateSDUpscale Master - Total tiles: {total_tiles}") + debug_log(f"UltimateSDUpscale Master - Tiles per image: {num_tiles_per_image}") # Parse enabled workers enabled_workers = json.loads(enabled_worker_ids) num_workers = len(enabled_workers) - if num_workers == 0: + # Determine processing mode + mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold) + + debug_log(f"UltimateSDUpscale Master - Mode: {mode}, batch_size: {batch_size}, workers: {num_workers}") + + if mode == "single_gpu": # No workers, process all tiles locally return self.process_single_gpu(upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, force_uniform_tiles) + tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode) - # Initialize job queue FIRST before any processing - # Initialize the queue directly in the persistent server storage - # This ensures it's accessible from the web server's event loop - try: - run_async_in_server_loop( - self._init_job_queue(multi_job_id), - timeout=QUEUE_INIT_TIMEOUT - ) - debug_log(f"UltimateSDUpscale Master - Queue initialization complete") - except Exception as e: - debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}") - # Try synchronous initialization as fallback - prompt_server = ensure_tile_jobs_initialized() - if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - prompt_server.distributed_pending_tile_jobs[multi_job_id] = asyncio.Queue() - debug_log(f"UltimateSDUpscale Master - Queue initialized via fallback") + elif mode == "dynamic": + # Dynamic mode for large batches + return self.process_master_dynamic(upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers) - # Don't try to prepare via API since we're using local queues - # loop = asyncio.new_event_loop() - # asyncio.set_event_loop(loop) - - # Prepare result image - image_pil = tensor_to_pil(upscaled_image, 0) - result_image = image_pil.copy() - - # Mask will be created per tile for proper blending - - # Calculate master's tiles - master_tiles_indices = self._get_master_tiles(all_tiles, num_workers) - master_tiles = len(master_tiles_indices) - - debug_log(f"UltimateSDUpscale Master - Processing {master_tiles} tiles locally") - - # Process master tiles SYNCHRONOUSLY - for tile_idx in master_tiles_indices: - result_image = self._process_and_blend_tile( - tile_idx, all_tiles[tile_idx], upscaled_image, result_image, - model, positive, negative, vae, seed, steps, cfg, - sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height - ) - - # Collect worker tiles using async - worker_tiles_expected = total_tiles - master_tiles - if worker_tiles_expected > 0: - # Calculate tile distribution for active worker calculation - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) - - # Calculate which workers will actually have tiles - active_workers = [] - for i in range(num_workers): - worker_start_idx = master_tiles + (i * tiles_per_participant) - if i < remainder - 1: - worker_start_idx += i - worker_tile_count = tiles_per_participant + 1 - else: - worker_start_idx += remainder - 1 if remainder > 0 else 0 - worker_tile_count = tiles_per_participant - - if worker_start_idx < total_tiles and worker_tile_count > 0: - active_workers.append(enabled_workers[i]) - - debug_log(f"UltimateSDUpscale Master - Waiting for {worker_tiles_expected} tiles from {len(active_workers)} active workers") - debug_log(f"UltimateSDUpscale Master - Active worker IDs: {active_workers}") - debug_log(f"UltimateSDUpscale Master - All enabled worker IDs: {enabled_workers}") - - # Collect tiles from workers using the server's event loop - collected_tiles = run_async_in_server_loop( - self._async_collect_worker_tiles(multi_job_id, len(active_workers)), - timeout=TILE_COLLECTION_TIMEOUT - ) - - # Blend collected tiles SYNCHRONOUSLY - for tile_data in collected_tiles.values(): - x = tile_data['x'] - y = tile_data['y'] - tile_tensor = tile_data['tensor'] - tile_idx = tile_data['tile_idx'] - extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding) - extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding) - - # Convert and blend - tile_pil = tensor_to_pil(tile_tensor, 0) - # Get the original tile position from the tile index - orig_x, orig_y = all_tiles[tile_idx] - # Create mask for this specific tile - tile_mask = self.create_tile_mask(width, height, orig_x, orig_y, tile_width, tile_height, mask_blur) - # Use extraction position and size for blending - result_image = self.blend_tile(result_image, tile_pil, - x, y, (extracted_width, extracted_height), tile_mask, padding) - - # Convert back to tensor - result_tensor = pil_to_tensor(result_image) - if upscaled_image.is_cuda: - result_tensor = result_tensor.cuda() - - debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") - return (result_tensor,) + # Static mode - enhanced with unified job management + return self._process_master_static_sync(upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers, + all_tiles, num_tiles_per_image, static_distribution) - def _get_worker_tiles(self, all_tiles, enabled_worker_ids, worker_id): - """Calculate which tiles are assigned to a specific worker.""" - enabled_workers = json.loads(enabled_worker_ids) + def _get_worker_global_indices(self, total_tiles, enabled_workers, worker_id): + """Calculate which global tile indices are assigned to a specific worker in flattened mode.""" num_workers = len(enabled_workers) - total_tiles = len(all_tiles) debug_log(f"UltimateSDUpscale Worker - Worker ID: {worker_id}, Enabled workers: {enabled_workers}") @@ -340,34 +252,20 @@ class UltimateSDUpscaleDistributed: log(f"UltimateSDUpscale Worker - Worker {worker_id} not found in enabled list {enabled_workers}") return [] - # Calculate tile distribution - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) - master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) - - start_idx = master_tiles + (worker_index * tiles_per_participant) - if worker_index < remainder - 1: - start_idx += worker_index - end_idx = start_idx + tiles_per_participant + 1 - else: - start_idx += remainder - 1 if remainder > 0 else 0 - end_idx = start_idx + tiles_per_participant - - end_idx = min(end_idx, total_tiles) - return list(range(start_idx, end_idx)) + global_indices = list(range(total_tiles)) + all_assignments = _distribute_tasks(global_indices, num_workers + 1) + return all_assignments[worker_index + 1] # +1 because 0 is master - def _get_master_tiles(self, all_tiles, num_workers): - """Calculate which tiles are assigned to the master.""" - total_tiles = len(all_tiles) - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) - master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) - return list(range(master_tiles)) + def _get_master_global_indices(self, total_tiles, num_workers): + """Calculate which global tile indices are assigned to the master in flattened mode.""" + global_indices = list(range(total_tiles)) + all_assignments = _distribute_tasks(global_indices, num_workers + 1) + return all_assignments[0] def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, image_width, image_height): + padding, mask_blur, image_width, image_height, tiled_decode, batch_idx: int = 0): """Process a single tile and blend it into the result image.""" x, y = tile_pos @@ -378,7 +276,8 @@ class UltimateSDUpscaleDistributed: processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, seed + tile_idx, steps, cfg, sampler_name, - scheduler, denoise) + scheduler, denoise, tiled_decode, batch_idx=batch_idx, + region=(x1, y1, x1 + ew, y1 + eh), image_size=(image_width, image_height)) # Convert and blend processed_pil = tensor_to_pil(processed_tile, 0) @@ -390,92 +289,277 @@ class UltimateSDUpscaleDistributed: return result_image - async def _prepare_multigpu_job(self, multi_job_id): - """Prepare job via API endpoint to ensure server is ready.""" - session = await get_client_session() - # Use the actual server port (get from server instance) - port = get_server_port() - async with session.post(f"http://localhost:{port}/distributed/prepare_job", - json={"multi_job_id": multi_job_id}) as response: - if response.status != 200: - text = await response.text() - raise RuntimeError(f"Failed to prepare job: {response.status} - {text}") + async def _init_job_queue_unified(self, multi_job_id, mode='static', batch_size=None, + total_tiles=None, enabled_workers=None, num_tiles_per_image=None, + assigned_to_workers=None, all_indices=None, task_assignments=None): + """Unified job queue initialization for both static and dynamic modes.""" + if mode == 'dynamic': + # Dynamic mode initialization + await _init_job_queue(multi_job_id, 'dynamic', batch_size=batch_size, + all_indices=all_indices, + enabled_workers=list(assigned_to_workers.keys()) if assigned_to_workers else []) + # Add additional fields for backward compatibility + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + job_data['completed_images'] = {} + job_data['completed_tiles'] = {} # For uniformity + # Fix: The endpoint expects 'pending_images', not 'pending_tasks' + job_data['pending_images'] = job_data[JOB_PENDING_TASKS] + debug_log(f"UltimateSDUpscale Master Dynamic - Initialized job queue with {batch_size} pending images for all participants") + else: + # Static mode initialization + # If num_tiles_per_image not provided, calculate it + if num_tiles_per_image is None and total_tiles is not None and batch_size is not None: + num_tiles_per_image = total_tiles // batch_size if batch_size > 0 else total_tiles + + await _init_job_queue(multi_job_id, 'static', + batch_size=batch_size, + num_tiles_per_image=num_tiles_per_image, + all_indices=None, + enabled_workers=enabled_workers, + task_assignments=task_assignments) + + # Add completed_tiles and total_items for backward compatibility + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + job_data['completed_tiles'] = {} + if total_tiles is not None: + job_data['total_items'] = total_tiles + debug_log(f"UltimateSDUpscale Master - Initialized job queue for {multi_job_id}") - async def _init_job_queue(self, multi_job_id): + # Keep compatibility wrappers + async def _init_job_queue_dynamic(self, multi_job_id, batch_size, assigned_to_workers=None, worker_status=None, all_indices=None): + """Initialize the job queue for dynamic mode with pending images.""" + await self._init_job_queue_unified(multi_job_id, mode='dynamic', batch_size=batch_size, + assigned_to_workers=assigned_to_workers, all_indices=all_indices) + + async def _init_job_queue(self, multi_job_id, total_tiles, enabled_workers, batch_size=1, num_tiles_per_image=None, task_assignments=None): """Initialize the job queue for collecting tiles.""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Checking if {multi_job_id} exists in distributed_pending_tile_jobs") - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Current jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") - if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - prompt_server.distributed_pending_tile_jobs[multi_job_id] = asyncio.Queue() - debug_log(f"UltimateSDUpscale Master - Initialized job queue for {multi_job_id}") - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Created new queue for {multi_job_id}") - else: - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Queue already exists for {multi_job_id}") + await self._init_job_queue_unified(multi_job_id, mode='static', batch_size=batch_size, + total_tiles=total_tiles, enabled_workers=enabled_workers, + num_tiles_per_image=num_tiles_per_image, task_assignments=task_assignments) - async def _async_collect_worker_tiles(self, multi_job_id, num_workers): - """Async helper to collect tiles from workers.""" + async def _async_collect_results(self, multi_job_id, num_workers, mode='static', + remaining_to_collect=None, batch_size=None): + """Unified async helper to collect results from workers (tiles or images).""" # Get the already initialized queue prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: if multi_job_id not in prompt_server.distributed_pending_tile_jobs: raise RuntimeError(f"Job queue not initialized for {multi_job_id}") - q = prompt_server.distributed_pending_tile_jobs[multi_job_id] + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + raise RuntimeError("Invalid job data structure") + if job_data['mode'] != mode: + raise RuntimeError(f"Mode mismatch: expected {mode}, got {job_data['mode']}") + q = job_data['queue'] + + # For dynamic mode, get reference to completed_images + if mode == 'dynamic': + completed_images = job_data['completed_images'] + # Calculate expected count for logging + expected_count = remaining_to_collect or batch_size + else: + # Calculate expected count from job data for static mode + expected_count = job_data.get('total_items', 0) - len(job_data.get('completed_items', set())) - collected_tiles = {} + mode_str = "Dynamic" if mode == 'dynamic' else "Static" + item_type = "images" if mode == 'dynamic' else "tiles" + debug_log(f"UltimateSDUpscale Master {mode_str} - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers") + + collected_results = {} workers_done = set() - timeout = TILE_WAIT_TIMEOUT - - debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {num_workers} workers to complete") + timeout = 120.0 if mode == 'dynamic' else 60.0 # Longer timeout for full images + last_heartbeat_check = time.time() + collected_count = 0 while len(workers_done) < num_workers: - debug_log(f"UltimateSDUpscale Master - Loop status: {len(workers_done)}/{num_workers} workers done, {len(collected_tiles)} tiles collected, workers_done set: {workers_done}") + # Check for user interruption + if comfy.model_management.processing_interrupted(): + log("Processing interrupted by user") + raise comfy.model_management.InterruptProcessingException() + + # For dynamic mode with remaining_to_collect, check if we've collected enough + if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect: + break + try: - result = await asyncio.wait_for(q.get(), timeout=timeout) - tile_idx = result['tile_idx'] + # Shorter timeout for regular checks in dynamic mode + wait_timeout = 10.0 if mode == 'dynamic' else timeout + result = await asyncio.wait_for(q.get(), timeout=wait_timeout) worker_id = result['worker_id'] is_last = result.get('is_last', False) - collected_tiles[tile_idx] = result - debug_log(f"UltimateSDUpscale Master - Received tile {tile_idx} from worker '{worker_id}' (is_last={is_last})") + if mode == 'static': + # Handle tiles + tiles = result.get('tiles', []) + if tiles: + # Batch mode + debug_log(f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker '{worker_id}' (is_last={is_last})") + + for tile_data in tiles: + # Validate required fields + if 'batch_idx' not in tile_data: + log(f"UltimateSDUpscale Master - Missing batch_idx in tile data, skipping") + continue + + tile_idx = tile_data['tile_idx'] + # Use global_idx as key if available (for batch processing) + key = tile_data.get('global_idx', tile_idx) + + # Store the full tile data including metadata + collected_results[key] = { + 'tensor': tile_data['tensor'], + 'tile_idx': tile_idx, + 'x': tile_data['x'], + 'y': tile_data['y'], + 'extracted_width': tile_data['extracted_width'], + 'extracted_height': tile_data['extracted_height'], + 'padding': tile_data['padding'], + 'worker_id': worker_id, + 'batch_idx': tile_data.get('batch_idx', 0), + 'global_idx': tile_data.get('global_idx', tile_idx) + } + else: + # Single tile mode (backward compat) + tile_idx = result['tile_idx'] + collected_results[tile_idx] = result + debug_log(f"UltimateSDUpscale Master - Received single tile {tile_idx} from worker '{worker_id}' (is_last={is_last})") + + elif mode == 'dynamic': + # Handle full images + if 'image_idx' in result and 'image' in result: + image_idx = result['image_idx'] + image_pil = result['image'] + completed_images[image_idx] = image_pil + collected_results[image_idx] = image_pil + collected_count += 1 + debug_log(f"UltimateSDUpscale Master Dynamic - Received image {image_idx} from worker {worker_id}") if is_last: workers_done.add(worker_id) - debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed") - debug_log(f"UltimateSDUpscale Master - Worker '{worker_id}' marked as done. Total workers done: {len(workers_done)}/{num_workers}") - + debug_log(f"UltimateSDUpscale Master {mode_str} - Worker {worker_id} completed") + except asyncio.TimeoutError: - log(f"UltimateSDUpscale Master - Timeout waiting for tiles") - debug_log(f"UltimateSDUpscale Master - Final status: {len(workers_done)}/{num_workers} workers done, {len(collected_tiles)} tiles collected") - debug_log(f"UltimateSDUpscale Master - Workers that completed: {workers_done}") - break + if mode == 'dynamic': + # Check for worker timeouts periodically + current_time = time.time() + if current_time - last_heartbeat_check >= 10.0: + # Use the class method to check and requeue + requeued = await self._check_and_requeue_timed_out_workers(multi_job_id, batch_size) + if requeued > 0: + log(f"UltimateSDUpscale Master Dynamic - Requeued {requeued} images from timed out workers") + last_heartbeat_check = current_time + + # Check if we've been waiting too long overall + if current_time - last_heartbeat_check > timeout: + log(f"UltimateSDUpscale Master Dynamic - Overall timeout waiting for images") + break + else: + log(f"UltimateSDUpscale Master - Timeout waiting for {item_type}") + break - debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_tiles)} tiles from {len(workers_done)} workers") + debug_log(f"UltimateSDUpscale Master {mode_str} - Collection complete. Got {len(collected_results)} {item_type} from {len(workers_done)} workers") # Clean up job queue - prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: del prompt_server.distributed_pending_tile_jobs[multi_job_id] - return collected_tiles + return collected_results if mode == 'static' else completed_images + + # Keep compatibility wrappers for existing code + async def _async_collect_worker_tiles(self, multi_job_id, num_workers): + """Async helper to collect tiles from workers.""" + return await self._async_collect_results(multi_job_id, num_workers, mode='static') + + async def _mark_image_completed(self, multi_job_id, image_idx, image_pil): + """Mark an image as completed in the job data.""" + # Mark the image as completed with the image data + await _mark_task_completed(multi_job_id, image_idx, {'image': image_pil}) + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and 'completed_images' in job_data: + job_data['completed_images'][image_idx] = image_pil + + async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count): + """Collect remaining processed images from workers.""" + return await self._async_collect_results(multi_job_id, num_workers, mode='dynamic', + remaining_to_collect=remaining_to_collect, + batch_size=batch_size) def round_to_multiple(self, value: int, multiple: int = 8) -> int: """Round value to nearest multiple.""" return round(value / multiple) * multiple def calculate_tiles(self, image_width: int, image_height: int, - tile_width: int, tile_height: int) -> List[Tuple[int, int]]: + tile_width: int, tile_height: int, force_uniform_tiles: bool = True) -> List[Tuple[int, int]]: """Calculate tile positions for the image. Tiles are placed at grid positions without overlap in their placement. - The overlap happens during extraction with padding.""" - tiles = [] - for y in range(0, image_height, tile_height): - for x in range(0, image_width, tile_width): - tiles.append((x, y)) - return tiles + The overlap happens during extraction with padding. + + If force_uniform_tiles is True, tiles are distributed evenly across the image.""" + if force_uniform_tiles: + # Calculate number of tiles needed + num_tiles_x = math.ceil(image_width / tile_width) + num_tiles_y = math.ceil(image_height / tile_height) + + # Calculate uniform spacing + tiles = [] + + # Handle X dimension + if num_tiles_x == 1: + tiles_x = [0] + else: + # Calculate step to distribute tiles evenly + step_x = (image_width - tile_width) / (num_tiles_x - 1) + # Check for excessive overlap + if step_x < tile_width / 2: + debug_log(f"Warning: Uniform tiles would have excessive overlap (step_x={step_x}, tile_width={tile_width}). Falling back to non-uniform mode.") + force_uniform_tiles = False + else: + tiles_x = [round(i * step_x) for i in range(num_tiles_x)] + # Ensure tiles are aligned to 8-pixel boundaries + tiles_x = [self.round_to_multiple(x, 8) for x in tiles_x] + + # Handle Y dimension (if still using uniform tiles) + if force_uniform_tiles: + if num_tiles_y == 1: + tiles_y = [0] + else: + # Calculate step to distribute tiles evenly + step_y = (image_height - tile_height) / (num_tiles_y - 1) + # Check for excessive overlap + if step_y < tile_height / 2: + debug_log(f"Warning: Uniform tiles would have excessive overlap (step_y={step_y}, tile_height={tile_height}). Falling back to non-uniform mode.") + force_uniform_tiles = False + else: + tiles_y = [round(i * step_y) for i in range(num_tiles_y)] + # Ensure tiles are aligned to 8-pixel boundaries + tiles_y = [self.round_to_multiple(y, 8) for y in tiles_y] + + # If still using uniform tiles, generate positions + if force_uniform_tiles: + # Generate all tile positions + for y in tiles_y: + for x in tiles_x: + tiles.append((x, y)) + + return tiles + + # Fall through to non-uniform mode if needed + if not force_uniform_tiles: + # Original non-uniform tile calculation + tiles = [] + for y in range(0, image_height, tile_height): + for x in range(0, image_width, tile_width): + tiles.append((x, y)) + return tiles def extract_tile_with_padding(self, image: torch.Tensor, x: int, y: int, tile_width: int, tile_height: int, padding: int) -> Tuple[torch.Tensor, int, int, int, int]: @@ -514,29 +598,71 @@ class UltimateSDUpscaleDistributed: return tile_tensor, x1, y1, extracted_width, extracted_height def process_tile(self, tile_tensor: torch.Tensor, model, positive, negative, vae, - seed: int, steps: int, cfg: float, sampler_name: str, - scheduler: str, denoise: float) -> torch.Tensor: - """Process a single tile through SD sampling.""" + seed: int, steps: int, cfg: float, sampler_name: str, + scheduler: str, denoise: float, tiled_decode: bool = False, + batch_idx: int = 0, region: Tuple[int, int, int, int] = None, + image_size: Tuple[int, int] = None) -> torch.Tensor: + """Process a single tile through SD sampling. + Note: positive and negative should already be pre-sliced for the current batch_idx.""" + debug_log(f"[process_tile] Processing tile for batch_idx={batch_idx}, seed={seed}, region={region}") + + # Import here to avoid circular dependencies from nodes import common_ksampler, VAEEncode, VAEDecode + # Try to import tiled VAE nodes if available + try: + from nodes import VAEEncodeTiled, VAEDecodeTiled + tiled_vae_available = True + except ImportError: + tiled_vae_available = False + if tiled_decode: + debug_log("Tiled VAE nodes not available, falling back to standard VAE") + # Convert to PIL and back to ensure clean tensor without gradient tracking tile_pil = tensor_to_pil(tile_tensor, 0) clean_tensor = pil_to_tensor(tile_pil) + # Ensure tensor is detached and doesn't require gradients + clean_tensor = clean_tensor.detach() + if hasattr(clean_tensor, 'requires_grad_'): + clean_tensor.requires_grad_(False) + # Move to correct device if tile_tensor.is_cuda: clean_tensor = clean_tensor.cuda() + clean_tensor = clean_tensor.detach() # Detach again after device transfer - # Encode to latent + # Clone conditioning per tile (shares models, clones hints for cropping) + positive_tile = clone_conditioning(positive, clone_hints=True) + negative_tile = clone_conditioning(negative, clone_hints=True) + + # Crop conditioning to tile region if provided (assumes hints at image resolution) + if region is not None and image_size is not None: + init_size = image_size # (width, height) of full image + canvas_size = image_size + tile_size = (tile_tensor.shape[2], tile_tensor.shape[1]) # (width, height) + w_pad = 0 # No extra pad needed; region already includes padding + h_pad = 0 + positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) + negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) + else: + # No region cropping needed, use cloned conditioning as-is + positive_cropped = positive_tile + negative_cropped = negative_tile + + # Encode to latent (always non-tiled, matching original node) latent = VAEEncode().encode(vae, clean_tensor)[0] # Sample samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, - positive, negative, latent, denoise=denoise)[0] + positive_cropped, negative_cropped, latent, denoise=denoise)[0] # Decode back to image - image = VAEDecode().decode(vae, samples)[0] + if tiled_decode and tiled_vae_available: + image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0] + else: + image = VAEDecode().decode(vae, samples)[0] return image @@ -602,164 +728,1475 @@ class UltimateSDUpscaleDistributed: # Convert back to RGB return result.convert('RGB') + def _slice_conditioning(self, positive, negative, batch_idx): + """Helper to slice conditioning for a specific batch index.""" + # Clone and slice conditioning properly, including ControlNet hints + positive_sliced = clone_conditioning(positive) + negative_sliced = clone_conditioning(negative) + + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[batch_idx:batch_idx+1] + if 'control' in cond_dict: + control = cond_dict['control'] + while control is not None: + hint = control.cond_hint_original + if hint.shape[0] > 1: + control.cond_hint_original = hint[batch_idx:batch_idx+1] + control = control.previous_controlnet + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] + + return positive_sliced, negative_sliced - async def send_tile_to_master(self, tile_tensor, multi_job_id, master_url, - tile_idx, x, y, extracted_width, extracted_height, - padding, worker_id, is_last=False): - """Send processed tile to master with retry logic.""" - # Convert tensor to PIL image - img = tensor_to_pil(tile_tensor, 0) + async def _get_all_completed_tasks(self, multi_job_id): + """Helper to retrieve all completed tasks from the job data.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and JOB_COMPLETED_TASKS in job_data: + return dict(job_data[JOB_COMPLETED_TASKS]) # Return a copy + return {} + + async def _process_worker_static_async(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_workers): + """Enhanced static mode worker with health monitoring and retry logic.""" + # Round tile dimensions + tile_width = self.round_to_multiple(tile_width) + tile_height = self.round_to_multiple(tile_height) - # Store image bytes for retry - img_bytes = io.BytesIO() - img.save(img_bytes, format='PNG', compress_level=0) - img_data = img_bytes.getvalue() + # Calculate all tile positions + batch_size = upscaled_image.shape[0] + _, height, width, _ = upscaled_image.shape + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) - # Retry logic with exponential backoff - max_retries = 5 - retry_delay = 0.5 # Start with 500ms + # Calculate total tiles across all images + num_tiles_per_image = len(all_tiles) + total_tiles = batch_size * num_tiles_per_image - for attempt in range(max_retries): - try: - # Create fresh FormData for each attempt - data = aiohttp.FormData() - data.add_field('multi_job_id', multi_job_id) - data.add_field('worker_id', str(worker_id)) - data.add_field('tile_idx', str(tile_idx)) - data.add_field('x', str(x)) - data.add_field('y', str(y)) - data.add_field('extracted_width', str(extracted_width)) - data.add_field('extracted_height', str(extracted_height)) - data.add_field('padding', str(padding)) - data.add_field('is_last', str(is_last)) - data.add_field('image', io.BytesIO(img_data), filename=f'tile_{tile_idx}.png', content_type='image/png') + # Get initial assigned global tile indices for this worker + assigned_global_indices = self._get_worker_global_indices(total_tiles, enabled_workers, worker_id) + + debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} initially assigned {len(assigned_global_indices)} tiles") + + # Process initially assigned tiles + processed_tiles = [] + sliced_conditioning_cache = {} + + # Process assigned tiles with heartbeat after each + for global_idx in assigned_global_indices: + # Process the tile + processed_tile = await self._process_single_tile( + global_idx, num_tiles_per_image, upscaled_image, all_tiles, + model, positive, negative, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) - session = await get_client_session() - timeout = aiohttp.ClientTimeout(total=TILE_TRANSFER_TIMEOUT) - async with session.post(f"{master_url}/distributed/tile_complete", data=data, timeout=timeout) as response: - if response.status == 404 and attempt < max_retries - 1: - # Queue not ready yet, wait and retry - debug_log(f"UltimateSDUpscale Worker - Queue not ready for job {multi_job_id}, attempt {attempt + 1}/{max_retries}, retrying in {retry_delay}s...") - await asyncio.sleep(retry_delay) - retry_delay = min(retry_delay * 2, 5.0) # Exponential backoff with cap at 5s - continue - response.raise_for_status() - debug_log(f"UltimateSDUpscale Worker - Successfully sent tile {tile_idx} on attempt {attempt + 1}") - return # Success - except aiohttp.ClientResponseError as e: - if e.status == 404 and attempt < max_retries - 1: - debug_log(f"UltimateSDUpscale Worker - Got 404 for job {multi_job_id}, attempt {attempt + 1}/{max_retries}, retrying in {retry_delay}s...") - await asyncio.sleep(retry_delay) - retry_delay = min(retry_delay * 2, 5.0) + # Send heartbeat after each tile + try: + await _send_heartbeat_to_master(multi_job_id, master_url, worker_id) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send initial batch of processed tiles + if processed_tiles: + await self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id) + processed_tiles = [] + + # Check for requeued tiles from failed workers + debug_log(f"Worker {worker_id} checking for requeued tiles...") + while True: + # Request next task from pending queue + task_id = await _get_next_task(multi_job_id) + if task_id is None: + break + + debug_log(f"Worker {worker_id} processing requeued tile {task_id}") + + # Process the requeued tile + processed_tile = await self._process_single_tile( + task_id, num_tiles_per_image, upscaled_image, all_tiles, + model, positive, negative, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) + + # Send heartbeat + try: + await _send_heartbeat_to_master(multi_job_id, master_url, worker_id) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send tiles in batches + if len(processed_tiles) >= MAX_BATCH: + await self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id) + processed_tiles = [] + + # Send any remaining tiles + if processed_tiles: + await self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id) + + debug_log(f"Worker {worker_id} completed all assigned and requeued tiles") + return (upscaled_image,) + + def _process_worker_static_sync(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_workers, static_distribution=False): + """Worker static mode processing with optional dynamic queue pulling.""" + # Round tile dimensions + tile_width = self.round_to_multiple(tile_width) + tile_height = self.round_to_multiple(tile_height) + + # Get dimensions and calculate tiles + _, height, width, _ = upscaled_image.shape + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + num_tiles_per_image = len(all_tiles) + batch_size = upscaled_image.shape[0] + total_tiles = batch_size * num_tiles_per_image + + processed_tiles = [] + sliced_conditioning_cache = {} + + if static_distribution: + # Legacy mode: Use pre-assigned tiles + debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} using static tile assignment") + worker_global_indices = self._get_worker_global_indices(total_tiles, enabled_workers, worker_id) + debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} processing {len(worker_global_indices)} pre-assigned tiles") + + # Process assigned tiles synchronously + for global_idx in worker_global_indices: + batch_idx = global_idx // num_tiles_per_image + tile_idx = global_idx % num_tiles_per_image + + if batch_idx >= batch_size: continue - elif attempt == max_retries - 1: - log(f"UltimateSDUpscale Worker - Failed to send tile {tile_idx} after {max_retries} attempts: {e}") - raise - except Exception as e: - if attempt == max_retries - 1: - log(f"UltimateSDUpscale Worker - Failed to send tile {tile_idx} after {max_retries} attempts: {e}") - raise + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) else: - debug_log(f"UltimateSDUpscale Worker - Error on attempt {attempt + 1}: {e}, retrying...") - await asyncio.sleep(retry_delay) - retry_delay = min(retry_delay * 2, 5.0) + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Process tile synchronously + processed_tile = self._process_single_tile( + global_idx, num_tiles_per_image, upscaled_image, all_tiles, + model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) + + # Send heartbeat (async operation) + try: + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send tiles in batches + if len(processed_tiles) >= MAX_BATCH: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] + + # Send remaining tiles + if processed_tiles: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] + + # Check for requeued tiles from failed workers + debug_log(f"Worker {worker_id} checking for requeued tiles...") + while True: + # Request next task from pending queue (async operation) + task_id = run_async_in_server_loop( + _get_next_task(multi_job_id), + timeout=5.0 + ) + if task_id is None: + break + + debug_log(f"Worker {worker_id} processing requeued tile {task_id}") + + # Process the requeued tile synchronously + batch_idx = task_id // num_tiles_per_image + tile_idx = task_id % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + processed_tile = self._process_single_tile( + task_id, num_tiles_per_image, upscaled_image, all_tiles, + model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) + + # Send heartbeat + try: + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send tiles in batches + if len(processed_tiles) >= MAX_BATCH: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] + + else: + # Dynamic queue mode + debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} ready to process tiles dynamically") + processed_count = 0 + + # Poll for job readiness + max_poll_attempts = 20 + for attempt in range(max_poll_attempts): + ready = run_async_in_server_loop( + self._check_job_status(multi_job_id, master_url), + timeout=5.0 + ) + if ready: + debug_log(f"Job {multi_job_id} ready after {attempt} attempts") + break + time.sleep(1.0) + else: + log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting") + return (upscaled_image,) + + # Main processing loop - pull tiles from queue + while True: + # Request a tile to process + tile_idx, estimated_remaining = run_async_in_server_loop( + self._request_tile_from_master(multi_job_id, master_url, worker_id), + timeout=TILE_WAIT_TIMEOUT + ) + + if tile_idx is None: + debug_log(f"UltimateSDUpscale Worker - No more tiles to process") + break + + debug_log(f"UltimateSDUpscale Worker - Assigned tile {tile_idx} to worker {worker_id}") + processed_count += 1 + + # Calculate batch and local tile indices + batch_idx = tile_idx // num_tiles_per_image + local_tile_idx = tile_idx % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Process tile synchronously + processed_tile = self._process_single_tile( + tile_idx, num_tiles_per_image, upscaled_image, all_tiles, + model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) + + # Send heartbeat + try: + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send tiles in batches + if len(processed_tiles) >= MAX_BATCH: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] + + # Send any remaining tiles + if processed_tiles: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + + debug_log(f"Worker {worker_id} completed all assigned and requeued tiles") + return (upscaled_image,) + async def _async_collect_and_monitor_static(self, multi_job_id, total_tiles, expected_total): + """Async helper for collection and monitoring in static mode. + Returns collected tasks dict. Caller should check if all tasks are complete.""" + last_progress_log = time.time() + progress_interval = 5.0 + last_heartbeat_check = time.time() + stall_count = 0 + max_stall_count = 10 # Allow 10 iterations with no progress before returning + last_completed_count = 0 + + while True: + # Check for user interruption + if comfy.model_management.processing_interrupted(): + log("Processing interrupted by user") + raise comfy.model_management.InterruptProcessingException() + + # Drain any pending results + collected_count = await _drain_results_queue(multi_job_id) + + # Check and requeue timed-out workers periodically + current_time = time.time() + if current_time - last_heartbeat_check >= 10.0: + requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total) + if requeued_count > 0: + log(f"Requeued {requeued_count} tasks from timed-out workers") + stall_count = 0 # Reset stall count when we requeue + last_heartbeat_check = current_time + + # Get current completion count + completed_count = await _get_completed_count(multi_job_id) + + # Progress logging + if current_time - last_progress_log >= progress_interval: + log(f"Progress: {completed_count}/{expected_total} tasks completed") + last_progress_log = current_time + + # Check if all tasks are completed + if completed_count >= expected_total: + debug_log(f"All {expected_total} tasks completed") + break + + # Check if we're making progress + if completed_count > last_completed_count: + stall_count = 0 # Reset stall count on progress + last_completed_count = completed_count + else: + stall_count += 1 + + # If no progress for too long, check if we need to process locally + if stall_count >= max_stall_count: + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data: + pending_queue = job_data.get(JOB_PENDING_TASKS) + # If there are pending tasks but no active workers, or we're stalled + if pending_queue and not pending_queue.empty(): + log(f"Collection stalled with {expected_total - completed_count} tasks remaining. Returning for local processing.") + break + + # Wait a bit before next check + await asyncio.sleep(0.5) + + # Get all completed tasks for return + return await self._get_all_completed_tasks(multi_job_id) + def _process_master_static_sync(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers, + all_tiles, num_tiles_per_image, static_distribution=False): + """Static mode master processing with optional dynamic queue pulling.""" + batch_size = upscaled_image.shape[0] + _, height, width, _ = upscaled_image.shape + total_tiles = batch_size * num_tiles_per_image + + # Convert batch to PIL list for processing + result_images = [] + for b in range(batch_size): + image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) + result_images.append(image_pil.copy()) + + sliced_conditioning_cache = {} + master_indices = set() # Track which tiles master processed + master_processed_count = 0 # Track how many tiles master processed + + if static_distribution: + # Legacy mode: Pre-assign tiles to master and workers + debug_log(f"UltimateSDUpscale Master - Using static tile distribution") + + # Distribute tasks across master and workers + all_indices = list(range(total_tiles)) + num_participants = len(enabled_workers) + 1 + task_assignments = _distribute_tasks(all_indices, num_participants) + master_indices = set(task_assignments[0]) # Convert to set for efficient lookup + + # Initialize job queue with task assignments (don't populate pending queue) + run_async_in_server_loop( + self._init_job_queue(multi_job_id, total_tiles, enabled_workers, batch_size, num_tiles_per_image, task_assignments), + timeout=10.0 + ) + + debug_log(f"UltimateSDUpscale Master - Initialized static job queue for {total_tiles} tiles") + debug_log(f"UltimateSDUpscale Master - Processing {len(master_indices)} tiles locally") + master_processed_count = len(master_indices) + + # Process master's tiles synchronously + for global_idx in master_indices: + # Check for user interruption + comfy.model_management.throw_exception_if_processing_interrupted() + batch_idx = global_idx // num_tiles_per_image + tile_idx = global_idx % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Use unique seed per image + image_seed = seed + batch_idx * 1000 + + # Process tile synchronously - no async context for tensor operations + result_images[batch_idx] = self._process_and_blend_tile( + tile_idx, all_tiles[tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx + ) + + # Mark as completed (async operation) + run_async_in_server_loop( + _mark_task_completed(multi_job_id, global_idx, {'batch_idx': batch_idx, 'tile_idx': tile_idx}), + timeout=5.0 + ) + else: + # Dynamic queue mode: All tiles go to pending queue + debug_log(f"UltimateSDUpscale Master - Using dynamic tile distribution") + + # Initialize job queue with all tiles in pending queue for dynamic distribution + run_async_in_server_loop( + self._init_job_queue(multi_job_id, total_tiles, enabled_workers, batch_size, num_tiles_per_image), + timeout=10.0 + ) + + debug_log(f"UltimateSDUpscale Master - Initialized static job queue for {total_tiles} tiles with dynamic distribution") + debug_log(f"UltimateSDUpscale Master - Processing tiles dynamically from queue") + + # Process tiles by pulling from queue (similar to dynamic mode) + processed_count = 0 + consecutive_no_tile = 0 + max_consecutive_no_tile = 5 + + while processed_count < total_tiles: + # Check for user interruption + comfy.model_management.throw_exception_if_processing_interrupted() + + # Try to get a tile to process + tile_idx = run_async_in_server_loop( + self._get_next_tile_index(multi_job_id), + timeout=5.0 + ) + + if tile_idx is not None: + consecutive_no_tile = 0 + processed_count += 1 + + # Calculate batch and local tile indices + batch_idx = tile_idx // num_tiles_per_image + local_tile_idx = tile_idx % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + debug_log(f"Master processing tile {tile_idx} (batch {batch_idx}, tile {local_tile_idx})") + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Use unique seed per image + image_seed = seed + batch_idx * 1000 + + # Process tile synchronously + result_images[batch_idx] = self._process_and_blend_tile( + local_tile_idx, all_tiles[local_tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx + ) + + # Mark as completed + run_async_in_server_loop( + _mark_task_completed(multi_job_id, tile_idx, {'batch_idx': batch_idx, 'tile_idx': local_tile_idx}), + timeout=5.0 + ) + else: + consecutive_no_tile += 1 + if consecutive_no_tile >= max_consecutive_no_tile: + # No more tiles available, break to start collecting + debug_log(f"Master processed {processed_count} tiles, moving to collection phase") + break + # Small wait before retry + time.sleep(0.1) + + # Set master_processed_count for the collection phase + master_processed_count = processed_count + + # Continue processing any remaining tiles while collecting worker results + remaining_tiles = total_tiles - master_processed_count + if remaining_tiles > 0: + debug_log(f"Master waiting for {remaining_tiles} tiles from workers") + + # Collect worker results using async operations + try: + collected_tasks = run_async_in_server_loop( + self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles), + timeout=300.0 # 5 minutes to allow for retries and requeuing + ) + except comfy.model_management.InterruptProcessingException: + # Clean up job on interruption + run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0) + raise + + # Check if we need to process any remaining tasks locally after collection + completed_count = len(collected_tasks) + if completed_count < total_tiles: + log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures") + + # Process any remaining pending tasks + while True: + # Check for user interruption + comfy.model_management.throw_exception_if_processing_interrupted() + + # Get next task from pending queue + task_id = run_async_in_server_loop( + self._get_next_tile_index(multi_job_id), + timeout=5.0 + ) + + if task_id is None: + break + + batch_idx = task_id // num_tiles_per_image + local_tile_idx = task_id % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Use unique seed per image + image_seed = seed + batch_idx * 1000 + + # Process tile synchronously + result_images[batch_idx] = self._process_and_blend_tile( + local_tile_idx, all_tiles[local_tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx + ) + + # Mark as completed and store in collected_tasks + run_async_in_server_loop( + _mark_task_completed(multi_job_id, task_id, {'batch_idx': batch_idx, 'tile_idx': local_tile_idx}), + timeout=5.0 + ) + # Add to collected_tasks so we skip it in the blending loop + collected_tasks[task_id] = {'batch_idx': batch_idx, 'tile_idx': local_tile_idx} + else: + # Master processed all tiles + collected_tasks = run_async_in_server_loop( + self._get_all_completed_tasks(multi_job_id), + timeout=5.0 + ) + + # Blend worker tiles synchronously + for global_idx, tile_data in collected_tasks.items(): + # Skip tiles already processed by master + if static_distribution and global_idx in master_indices: + continue # Skip master's tiles in static mode + + # Skip tiles that don't have tensor data (already processed) + if 'tensor' not in tile_data: + continue + + batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image) + tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image) + + if batch_idx >= batch_size: + continue + + # Blend tile synchronously + x = tile_data.get('x', 0) + y = tile_data.get('y', 0) + tile_tensor = tile_data['tensor'] + tile_pil = tensor_to_pil(tile_tensor, 0) + orig_x, orig_y = all_tiles[tile_idx] + tile_mask = self.create_tile_mask(width, height, orig_x, orig_y, tile_width, tile_height, mask_blur) + extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding) + extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding) + result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil, + x, y, (extracted_width, extracted_height), tile_mask, padding) + + try: + # Convert back to tensor + if batch_size == 1: + result_tensor = pil_to_tensor(result_images[0]) + else: + result_tensors = [pil_to_tensor(img) for img in result_images] + result_tensor = torch.cat(result_tensors, dim=0) + + if upscaled_image.is_cuda: + result_tensor = result_tensor.cuda() + + log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") + return (result_tensor,) + finally: + # Cleanup (async operation) - always execute + run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0) + + def _process_single_tile(self, global_idx, num_tiles_per_image, upscaled_image, all_tiles, + model, positive, negative, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache): + """Process a single tile.""" + # Calculate which image and tile this corresponds to + batch_idx = global_idx // num_tiles_per_image + tile_idx = global_idx % num_tiles_per_image + + # Skip if batch_idx is out of range + if batch_idx >= upscaled_image.shape[0]: + debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}") + return None + + # Get or create sliced conditioning for this batch index + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + x, y = all_tiles[tile_idx] + + # Extract tile from the specific image in the batch + tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( + upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding + ) + + # Process tile through SD with unique seed + image_seed = seed + batch_idx * 1000 + processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae, + image_seed + tile_idx, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, batch_idx=batch_idx, + region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height)) + + return { + 'tile': processed_tile, + 'global_idx': global_idx, + 'batch_idx': batch_idx, + 'tile_idx': tile_idx, + 'x': x1, + 'y': y1, + 'extracted_width': ew, + 'extracted_height': eh + } + + async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, + padding, worker_id): + """Send all processed tiles to master, chunked if large.""" + if not processed_tiles: + return # Early exit if empty + + total_tiles = len(processed_tiles) + debug_log(f"UltimateSDUpscale Worker - Sending {total_tiles} tiles in chunks of max {MAX_BATCH}") + + for start in range(0, total_tiles, MAX_BATCH): + chunk = processed_tiles[start:start + MAX_BATCH] + chunk_size = len(chunk) + is_chunk_last = (start + chunk_size == total_tiles) # True only for final chunk + + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('is_last', str(is_chunk_last)) + data.add_field('batch_size', str(chunk_size)) + data.add_field('padding', str(padding)) + + # Chunk metadata: Keep absolute tile_idx and batch_idx if present + metadata = [] + for tile_data in chunk: + meta = { + 'tile_idx': tile_data['tile_idx'], # Original/absolute idx + 'x': tile_data['x'], + 'y': tile_data['y'], + 'extracted_width': tile_data['extracted_width'], + 'extracted_height': tile_data['extracted_height'] + } + # Add batch_idx and global_idx if present (for batch processing) + if 'batch_idx' in tile_data: + meta['batch_idx'] = tile_data['batch_idx'] + if 'global_idx' in tile_data: + meta['global_idx'] = tile_data['global_idx'] + metadata.append(meta) + + # Add JSON metadata + data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json') + + # Add chunk images as 'tile_{j}' where j=0 to chunk_size-1 + for j, tile_data in enumerate(chunk): + # Convert tensor to PIL + img = tensor_to_pil(tile_data['tile'], 0) + byte_io = io.BytesIO() + img.save(byte_io, format='PNG', compress_level=0) + byte_io.seek(0) + + # Add image field + data.add_field(f'tile_{j}', byte_io, filename=f'tile_{j}.png', content_type='image/png') + + + # Retry logic with exponential backoff + max_retries = 5 + retry_delay = 0.5 + + for attempt in range(max_retries): + try: + session = await get_client_session() + url = f"{master_url}/distributed/submit_tiles" + + async with session.post(url, data=data) as response: + response.raise_for_status() + break # Success, move to next chunk + + except Exception as e: + if attempt < max_retries - 1: + await asyncio.sleep(retry_delay) + retry_delay *= 2 # Exponential backoff + else: + log(f"UltimateSDUpscale Worker - Failed to send chunk after {max_retries} attempts: {e}") + raise + def process_single_gpu(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, force_uniform_tiles): + tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode): """Process all tiles on a single GPU (no distribution).""" # Round tile dimensions tile_width = self.round_to_multiple(tile_width) tile_height = self.round_to_multiple(tile_height) - # Get image dimensions - _, height, width, _ = upscaled_image.shape + # Get image dimensions and batch size + batch_size, height, width, _ = upscaled_image.shape # Calculate all tiles - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height) + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) - debug_log(f"UltimateSDUpscale - Processing {len(all_tiles)} tiles locally") + debug_log(f"UltimateSDUpscale - Processing {len(all_tiles)} tiles locally for batch of {batch_size}") - # Convert to PIL - image_pil = tensor_to_pil(upscaled_image, 0) - result_image = image_pil.copy() + # Convert batch to PIL list + result_images = [] + for b in range(batch_size): + image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) + # Ensure RGB mode for consistency + if image_pil.mode != 'RGB': + image_pil = image_pil.convert('RGB') + result_images.append(image_pil.copy()) - # Mask will be created per tile for proper blending - - # Process each tile - for idx, tile_pos in enumerate(all_tiles): - result_image = self._process_and_blend_tile( - idx, tile_pos, upscaled_image, result_image, - model, positive, negative, vae, seed, steps, cfg, - sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height - ) + # Process each image in the batch + for batch_idx in range(batch_size): + debug_log(f"[process_single_gpu] Processing batch_idx {batch_idx}") + # Pre-slice conditioning once per image (not per tile) + positive_sliced = clone_conditioning(positive) + negative_sliced = clone_conditioning(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[batch_idx:batch_idx+1] + if 'control' in cond_dict: + control = cond_dict['control'] + while control is not None: + hint = control.cond_hint_original + if hint.shape[0] > 1: + control.cond_hint_original = hint[batch_idx:batch_idx+1] + control = control.previous_controlnet + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] + + # Process each tile for this image with pre-sliced conditioning + for tile_idx, tile_pos in enumerate(all_tiles): + # Use unique seed per image + image_seed = seed + batch_idx * 1000 + result_images[batch_idx] = self._process_and_blend_tile( + tile_idx, tile_pos, upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx + ) # Convert back to tensor - result_tensor = pil_to_tensor(result_image) + if batch_size == 1: + result_tensor = pil_to_tensor(result_images[0]) + else: + result_tensors = [pil_to_tensor(img) for img in result_images] + result_tensor = torch.cat(result_tensors, dim=0) + if upscaled_image.is_cuda: result_tensor = result_tensor.cuda() return (result_tensor,) + + def process_master_dynamic(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers): + """Dynamic mode for large batches - assigns whole images to workers dynamically, including master.""" + # Get batch size and dimensions + batch_size, height, width, _ = upscaled_image.shape + num_workers = len(enabled_workers) + + debug_log(f"UltimateSDUpscale Master Dynamic - Processing batch of {batch_size} images with {num_workers} workers") + + # No fixed share - all images are dynamic + all_indices = list(range(batch_size)) + + log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.") + + # Calculate tiles for processing + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + + # Initialize job queue for communication + try: + assigned_to_workers = {w: [] for w in enabled_workers} + worker_status = {w: time.time() for w in enabled_workers} + + run_async_in_server_loop( + self._init_job_queue_dynamic(multi_job_id, batch_size, assigned_to_workers, worker_status, all_indices), + timeout=2.0 + ) + except Exception as e: + debug_log(f"UltimateSDUpscale Master Dynamic - Queue initialization error: {e}") + raise RuntimeError(f"Failed to initialize dynamic mode queue: {e}") + + # Convert batch to PIL list + result_images = [tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB').copy() for b in range(batch_size)] + + # Process images dynamically with master participating + prompt_server = ensure_tile_jobs_initialized() + processed_count = 0 + consecutive_retries = 0 + max_consecutive_retries = 10 + + # Process loop - master pulls from queue and processes synchronously + while processed_count < batch_size: + # Try to get an image to process + image_idx = run_async_in_server_loop( + self._get_next_image_index(multi_job_id), + timeout=5.0 # Short timeout to allow frequent checks + ) + if image_idx is not None: + # Reset retry counter and process locally + consecutive_retries = 0 + debug_log(f"Master processing image {image_idx} dynamically") + processed_count += 1 + + # Process locally + single_tensor = upscaled_image[image_idx:image_idx+1] + local_image = result_images[image_idx] + image_seed = seed + image_idx * 1000 + + # Pre-slice conditioning once per image (not per tile) + positive_sliced = clone_conditioning(positive) + negative_sliced = clone_conditioning(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[image_idx:image_idx+1] + if 'control' in cond_dict: + control = cond_dict['control'] + while control is not None: + hint = control.cond_hint_original + if hint.shape[0] > 1: + control.cond_hint_original = hint[image_idx:image_idx+1] + control = control.previous_controlnet + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1] + + for tile_idx, pos in enumerate(all_tiles): + local_image = self._process_and_blend_tile( + tile_idx, pos, single_tensor, local_image, + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode, batch_idx=image_idx + ) + + # Yield after each tile to minimize worker downtime + run_async_in_server_loop(self._async_yield(), timeout=0.1) + # Note: No per-tile drain here – that's what makes this "per-image" + + result_images[image_idx] = local_image + + # Mark as completed + run_async_in_server_loop( + self._mark_image_completed(multi_job_id, image_idx, local_image), + timeout=5.0 + ) + + # NEW: Drain after the full image is marked complete (catches workers who finished during master's processing) + drained_count = run_async_in_server_loop( + self._drain_worker_results_queue(multi_job_id), + timeout=5.0 + ) + + if drained_count > 0: + debug_log(f"Drained {drained_count} worker images after master's image {image_idx}") + + # NEW: Log overall progress (includes master's image + any drained workers) + completed_now = run_async_in_server_loop( + self._get_total_completed_count(multi_job_id), + timeout=1.0 + ) + log(f"Progress: {completed_now}/{batch_size} images completed") + + # Yield to allow workers to get new images after completing one + run_async_in_server_loop(self._async_yield(), timeout=0.1) + else: + # Queue empty: collect any queued worker results to update progress + drained_count = run_async_in_server_loop( + self._drain_worker_results_queue(multi_job_id), + timeout=5.0 + ) + run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after drain + + # Check for timed out workers and requeue their images + requeued_count = run_async_in_server_loop( + self._check_and_requeue_timed_out_workers(multi_job_id, batch_size), + timeout=5.0 + ) + run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after requeue + + if requeued_count > 0: + log(f"Requeued {requeued_count} images from timed out workers") + consecutive_retries = 0 # Reset since we have work to do + continue + + # Now check total completed (includes newly collected) + completed_now = run_async_in_server_loop( + self._get_total_completed_count(multi_job_id), + timeout=1.0 + ) + + log(f"Progress: {completed_now}/{batch_size} images completed") + + if completed_now >= batch_size: + break + + run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield before pending check + + # Check if there are pending images in the queue (could be requeued) + pending_count = run_async_in_server_loop( + self._get_pending_count(multi_job_id), + timeout=1.0 + ) + + if pending_count > 0: + consecutive_retries = 0 # Reset retries since there's work to do + continue + + consecutive_retries += 1 + if consecutive_retries >= max_consecutive_retries: + log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.") + break # Force exit to collection phase + + debug_log("Waiting for workers") + # Use async sleep to allow event loop to process worker requests + run_async_in_server_loop(asyncio.sleep(2), timeout=3.0) + + debug_log(f"Master processed {processed_count} images locally") + + # Get all completed images to check what needs to be collected + all_completed = run_async_in_server_loop( + self._get_all_completed_images(multi_job_id), + timeout=5.0 + ) + + # Calculate how many we still need to collect + remaining_to_collect = batch_size - len(all_completed) + + if remaining_to_collect > 0: + debug_log(f"Waiting for {remaining_to_collect} more images from workers") + collected_images = run_async_in_server_loop( + self._async_collect_dynamic_images(multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count), + timeout=TILE_COLLECTION_TIMEOUT * 2 + ) + + # Merge collected with already completed + all_completed.update(collected_images) + + # Update result images with all completed images + for idx, processed_img in all_completed.items(): + if idx < batch_size: + result_images[idx] = processed_img + + # Convert back to tensor + result_tensor = torch.cat([pil_to_tensor(img) for img in result_images], dim=0) if batch_size > 1 else pil_to_tensor(result_images[0]) + if upscaled_image.is_cuda: + result_tensor = result_tensor.cuda() + + debug_log(f"UltimateSDUpscale Master Dynamic - Job {multi_job_id} complete") + log(f"Completed processing all {batch_size} images") + return (result_tensor,) + + + async def _async_yield(self): + """Simple async yield to allow event loop processing.""" + await asyncio.sleep(0) + + async def _get_next_image_index(self, multi_job_id): + """Get next image index from pending queue for master.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data or 'pending_images' not in job_data: + return None + try: + image_idx = await asyncio.wait_for(job_data['pending_images'].get(), timeout=1.0) + return image_idx + except asyncio.TimeoutError: + return None + + async def _get_next_tile_index(self, multi_job_id): + """Get next tile index from pending queue for master in static mode.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data or JOB_PENDING_TASKS not in job_data: + return None + try: + tile_idx = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=1.0) + return tile_idx + except asyncio.TimeoutError: + return None + + + async def _get_total_completed_count(self, multi_job_id): + """Get total count of all completed images (master + workers).""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and 'completed_images' in job_data: + return len(job_data['completed_images']) + return 0 + + async def _get_all_completed_images(self, multi_job_id): + """Get all completed images.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and 'completed_images' in job_data: + return job_data['completed_images'].copy() + return {} + + async def _get_pending_count(self, multi_job_id): + """Get count of pending images in the queue.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and 'pending_images' in job_data: + return job_data['pending_images'].qsize() + return 0 + + async def _drain_worker_results_queue(self, multi_job_id): + """Drain pending worker results from queue and update completed_images. Returns count of drained images.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data or 'queue' not in job_data or 'completed_images' not in job_data: + return 0 + q = job_data['queue'] + completed_images = job_data['completed_images'] + + collected = 0 + while not q.empty(): + try: + result = await asyncio.wait_for(q.get(), timeout=0.1) + worker_id = result['worker_id'] + is_last = result.get('is_last', False) + + if 'image_idx' in result and 'image' in result: + image_idx = result['image_idx'] + image_pil = result['image'] + if image_idx not in completed_images: + completed_images[image_idx] = image_pil + collected += 1 + debug_log(f"Drained image {image_idx} from worker {worker_id}") + + if is_last: + # Optional: track worker completion if needed + pass + except asyncio.TimeoutError: + break # No more immediately available + + if collected > 0: + debug_log(f"Drained {collected} worker images during retry") + + return collected + + async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size): + """Check for timed out workers and requeue their assigned images. Returns count of requeued images.""" + # Use the original function from usdu_managment.py + return await _check_and_requeue_timed_out_workers(multi_job_id, batch_size) + + + def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_worker_ids, dynamic_threshold): + """Worker processing in dynamic mode - processes whole images.""" + # Round tile dimensions + tile_width = self.round_to_multiple(tile_width) + tile_height = self.round_to_multiple(tile_height) + + # Get dimensions + _, height, width, _ = upscaled_image.shape + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + + debug_log(f"UltimateSDUpscale Worker Dynamic - Worker {worker_id} ready to process images") + + # Keep track of processed images for is_last detection + processed_count = 0 + + # Poll for job readiness to avoid races during master init + max_poll_attempts = 20 # ~20s at 1s sleep + for attempt in range(max_poll_attempts): + ready = run_async_in_server_loop( + self._check_job_status(multi_job_id, master_url), + timeout=5.0 + ) + if ready: + debug_log(f"Job {multi_job_id} ready after {attempt} attempts") + break + time.sleep(1.0) # Poll every 1s + else: + log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting") + return (upscaled_image,) + + # Loop to request and process images + while True: + # Request an image to process + image_idx, estimated_remaining = run_async_in_server_loop( + self._request_image_from_master(multi_job_id, master_url, worker_id), + timeout=TILE_WAIT_TIMEOUT + ) + + if image_idx is None: + debug_log(f"UltimateSDUpscale Worker Dynamic - No more images to process") + break + + debug_log(f"UltimateSDUpscale Worker Dynamic - Mode: dynamic, assigned image {image_idx} to worker {worker_id}") + processed_count += 1 + + # Determine if this should be marked as last for this worker + is_last_for_worker = (estimated_remaining == 0) + + # Extract single image tensor + single_tensor = upscaled_image[image_idx:image_idx+1] + + # Convert to PIL for processing + local_image = tensor_to_pil(single_tensor, 0).copy() + + # Process all tiles for this image + image_seed = seed + image_idx * 1000 + + # Pre-slice conditioning once per image (not per tile) + positive_sliced = clone_conditioning(positive) + negative_sliced = clone_conditioning(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[image_idx:image_idx+1] + if 'control' in cond_dict: + control = cond_dict['control'] + while control is not None: + hint = control.cond_hint_original + if hint.shape[0] > 1: + control.cond_hint_original = hint[image_idx:image_idx+1] + control = control.previous_controlnet + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1] + + for tile_idx, pos in enumerate(all_tiles): + local_image = self._process_and_blend_tile( + tile_idx, pos, single_tensor, local_image, + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode, batch_idx=image_idx + ) + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + + # Send processed image back to master + try: + # Use the estimated remaining to determine if this is the last image + is_last = is_last_for_worker + run_async_in_server_loop( + self._send_full_image_to_master(local_image, image_idx, multi_job_id, + master_url, worker_id, is_last), + timeout=TILE_SEND_TIMEOUT + ) + # Send heartbeat after processing + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + if is_last: + break + except Exception as e: + log(f"UltimateSDUpscale Worker Dynamic - Error sending image {image_idx}: {e}") + # Continue processing other images + + # Send final is_last signal + debug_log(f"UltimateSDUpscale Worker Dynamic - Worker {worker_id} processed {processed_count} images, sending completion signal") + try: + run_async_in_server_loop( + self._send_worker_complete_signal(multi_job_id, master_url, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + except Exception as e: + log(f"UltimateSDUpscale Worker Dynamic - Error sending completion signal: {e}") + + return (upscaled_image,) + + async def _request_image_from_master(self, multi_job_id, master_url, worker_id): + """Request an image index to process from master in dynamic mode.""" + # Enhanced retries with 404-specific delay, backoff cap, and total timeout to handle init races + max_retries = 10 + retry_delay = 0.5 + start_time = time.time() + + for attempt in range(max_retries): + # Check total timeout + if time.time() - start_time > 30: + log(f"Total request timeout after 30s for worker {worker_id}") + return None, 0 + + try: + session = await get_client_session() + url = f"{master_url}/distributed/request_image" + + async with session.post(url, json={ + 'worker_id': str(worker_id), + 'multi_job_id': multi_job_id + }) as response: + if response.status == 200: + data = await response.json() + image_idx = data.get('image_idx') + estimated_remaining = data.get('estimated_remaining', 0) + return image_idx, estimated_remaining + elif response.status == 404: + # Special handling for 404 - job not found yet + text = await response.text() + debug_log(f"Job not found (404), will retry: {text}") + await asyncio.sleep(1.0) + else: + text = await response.text() + debug_log(f"Request image failed: {response.status} - {text}") + + except Exception as e: + if attempt < max_retries - 1: + debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") + await asyncio.sleep(retry_delay) + retry_delay *= 2 + retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s + else: + log(f"Failed to request image after {max_retries} attempts: {e}") + raise + + return None, 0 + + async def _request_tile_from_master(self, multi_job_id, master_url, worker_id): + """Request a tile index to process from master in static mode (reusing dynamic infrastructure).""" + # Reuse the same retry logic as dynamic mode + max_retries = 10 + retry_delay = 0.5 + start_time = time.time() + + for attempt in range(max_retries): + # Check total timeout + if time.time() - start_time > 30: + log(f"Total request timeout after 30s for worker {worker_id}") + return None, 0 + + try: + session = await get_client_session() + url = f"{master_url}/distributed/request_image" # Same endpoint + + async with session.post(url, json={ + 'worker_id': str(worker_id), + 'multi_job_id': multi_job_id + }) as response: + if response.status == 200: + data = await response.json() + tile_idx = data.get('tile_idx') + estimated_remaining = data.get('estimated_remaining', 0) + return tile_idx, estimated_remaining + elif response.status == 404: + # Special handling for 404 - job not found yet + text = await response.text() + debug_log(f"Job not found (404), will retry: {text}") + await asyncio.sleep(1.0) + else: + text = await response.text() + debug_log(f"Request tile failed: {response.status} - {text}") + + except Exception as e: + if attempt < max_retries - 1: + debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") + await asyncio.sleep(retry_delay) + retry_delay *= 2 + retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s + else: + log(f"Failed to request tile after {max_retries} attempts: {e}") + raise + + return None, 0 + + async def _check_job_status(self, multi_job_id, master_url): + """Check if job is ready on the master.""" + try: + session = await get_client_session() + url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}" + async with session.get(url) as response: + if response.status == 200: + data = await response.json() + return data.get('ready', False) + return False + except Exception as e: + debug_log(f"Job status check failed: {e}") + return False + + async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id, + master_url, worker_id, is_last): + """Send a processed full image back to master in dynamic mode.""" + # Serialize image to PNG + byte_io = io.BytesIO() + image_pil.save(byte_io, format='PNG', compress_level=0) + byte_io.seek(0) + + # Prepare form data + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('image_idx', str(image_idx)) + data.add_field('is_last', str(is_last)) + data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png', + content_type='image/png') + + # Retry logic + max_retries = 5 + retry_delay = 0.5 + + for attempt in range(max_retries): + try: + session = await get_client_session() + url = f"{master_url}/distributed/submit_image" + + async with session.post(url, data=data) as response: + response.raise_for_status() + debug_log(f"Successfully sent image {image_idx} to master") + return + + except Exception as e: + if attempt < max_retries - 1: + debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") + await asyncio.sleep(retry_delay) + retry_delay *= 2 + else: + log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}") + raise + + # Note: Using _send_heartbeat_to_master from usdu_managment.py instead of duplicate + + async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id): + """Send completion signal to master in dynamic mode.""" + # Send a dummy request with is_last=True + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('is_last', 'true') + # No image data - just completion signal + + session = await get_client_session() + url = f"{master_url}/distributed/submit_image" + + async with session.post(url, data=data) as response: + response.raise_for_status() + debug_log(f"Worker {worker_id} sent completion signal") + + def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str: + """Determines the processing mode based on batch size and worker count.""" + if num_workers == 0: + return "single_gpu" + if batch_size <= dynamic_threshold: + return "static" + return "dynamic" + + # Note: Using _distribute_tasks from usdu_managment.py instead of duplicate _distribute_items # Ensure initialization before registering routes ensure_tile_jobs_initialized() -# API Endpoint for tile completion -@server.PromptServer.instance.routes.post("/distributed/tile_complete") -async def tile_complete_endpoint(request): - """Endpoint for receiving completed tiles from workers.""" - try: - data = await request.post() - multi_job_id = data.get('multi_job_id') - image_file = data.get('image') - worker_id = data.get('worker_id') - tile_idx = int(data.get('tile_idx', 0)) - x = int(data.get('x', 0)) - y = int(data.get('y', 0)) - extracted_width = int(data.get('extracted_width', 512)) - extracted_height = int(data.get('extracted_height', 512)) - padding = int(data.get('padding', 32)) - 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,] - - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - debug_log(f"UltimateSDUpscale - tile_complete: Checking distributed_pending_tile_jobs for job {multi_job_id}") - debug_log(f"UltimateSDUpscale - tile_complete: Current jobs in distributed_pending_tile_jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - await prompt_server.distributed_pending_tile_jobs[multi_job_id].put({ - 'tensor': tensor, - 'worker_id': worker_id, - 'tile_idx': tile_idx, - 'x': x, - 'y': y, - 'extracted_width': extracted_width, - 'extracted_height': extracted_height, - 'padding': padding, - 'is_last': is_last - }) - debug_log(f"UltimateSDUpscale - Received tile {tile_idx} for job {multi_job_id} from worker {worker_id} (last: {is_last})") - debug_log(f"UltimateSDUpscale - tile_complete: Successfully queued tile {tile_idx}") - return web.json_response({"status": "success"}) - else: - debug_log(f"UltimateSDUpscale - API Error: Job {multi_job_id} not found in distributed_pending_tile_jobs") - debug_log(f"UltimateSDUpscale - API Error: Available jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") - return await handle_api_error(request, "Job not found or already complete", 404) - except Exception as e: - return await handle_api_error(request, e, 500) - - # Node registration NODE_CLASS_MAPPINGS = { "UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed, diff --git a/docs/model-download-script.md b/docs/model-download-script.md new file mode 100644 index 0000000..0ec346f --- /dev/null +++ b/docs/model-download-script.md @@ -0,0 +1,56 @@ +## Automating ComfyUI Model Downloads + +This guide will walk you through creating a shell script to automatically download the necessary models for your ComfyUI workflow, leveraging an advanced Large Language Model (LLM). + +1. Launched the Pod with these Environment Variables: + - CIVITAI_API_TOKEN ([get your token here](https://civitai.com/user/account)) + - HF_API_TOKEN ([get your token here](https://huggingface.co/settings/tokens)) +2. In ComfyUI, export your workflow as an API workflow +3. Copy the below prompt and upload the API workflow to a LLM **that has access to the internet** +``` +Create a sh script that will download the models from this workflow into the correct folders. For reference, these are the paths: + +base_path: /workspace/ComfyUI +checkpoints: models/checkpoints/ +clip: models/clip/ +clip_vision: models/clip_vision/ +controlnet: models/controlnet/ +diffusion_models: models/diffusion_models/ +embeddings: models/embeddings/ +florence2: models/florence2/ +ipadapter: models/ipadapter/ +loras: models/loras/ +style_models: models/style_models/ +text_encoders: models/text_encoders/ +unet: models/unet/ +upscale_models: models/upscale_models/ +vae: models/vae/ + +--- + +Important: +Make sure you find the correct URLs for the models online. +Use comfy cli to download the models: `comfy model download --url [--relative-path ] [--set-civitai-api-token ] [--set-hf-api-token ]` +Make sure you add `--set-civitai-api-token $CIVITAI_API_TOKEN` for CivitAI download and `--set-hf-api-token $HF_API_TOKEN` for Hugging Face downloads. + +--- + +Example: + +#!/bin/bash + +# Download from CivitAI +comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN + +# Download model from Hugging Face +comfy model download --url https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors --relative-path /workspace/ComfyUI/models/unet --set-hf-api-token $HF_API_TOKEN + +# If a model in the workflow was in a subfolder +comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints/SDXL --set-civitai-api-token $CIVITAI_API_TOKEN + +``` +4. Review the LLMs output to make sure all download links are correct and save it as a .sh file, for example `download_models.sh` +5. Upload it onto your Runpod instance, into `/workspace` +6. Then run these commands: + - `chmod 755 /workspace/download_models.sh` + - `/workspace/download_models.sh` diff --git a/docs/worker-setup-guides.md b/docs/worker-setup-guides.md new file mode 100644 index 0000000..0b2f242 --- /dev/null +++ b/docs/worker-setup-guides.md @@ -0,0 +1,108 @@ +## ComfyUI Distributed Worker Setup Guide + +ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs, whether they're in the same machine, on your local network, or in the cloud. + +**Master**: The main ComfyUI instance that coordinates and distributes work. This is where you load workflows, manage the queue, and view results. + +**Worker**: A ComfyUI instance that receives and processes tasks from the master. Workers handle just the GPU computation and send results back to the master. You can have multiple workers connected to a single master, each utilizing their own GPU. + +### Types of Workers + +- **Local workers**: Additional GPUs on the same machine as the master +- **Remote workers**: GPUs on different computers within your network +- **Cloud workers**: GPUs hosted on cloud services like Runpod + +## Local workers +> These are added automatically on first launch, but you can add them manually if you need to. + +1. **Open** the Distributed GPU panel. +2. **Click** "Add Worker" in the UI. +3. **Configure** your local worker: + - **Name**: A descriptive name for the worker (e.g., "Studio PC 1") + - **Port**: A unique port number for this worker (e.g., 8189, 8190...). + - **CUDA Device**: The GPU index from `nvidia-smi` (e.g., 0, 1). + - **Extra Args**: Optional ComfyUI arguments for this specific worker. +4. **Save** and launch the local worker. + +## Remote workers +> ComfyUI instances running on completely different computers on your network. These allow you to harness GPU power from other machines. Remote workers must be manually started on their respective computers and are connected via IP address. + +**On the Remote Worker Machine:** +1. **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!** + - This ComfyUI instance will serve as a worker for your main master. +2. *Optionally* add additional local workers on this machine if it has multiple GPUs: + - Access the Distributed GPU panel in this ComfyUI instance + - Add workers for any additional GPUs (if they haven't been added automatically) + - Make sure they have `--listen` set in `Extra Args` + - Launch them +3. **Open** the ComfyUI port (e.g., 8188) and any additional worker ports (e.g., 8189, 8190) in the firewall. + +**On the Main Machine:** +1. **Open** the Distributed GPU panel (sidebar on the left). +2. **Click** "Add Worker." +3. **Choose** "Remote". +4. **Configure** your remote worker: + - **Name**: A descriptive name for the worker (e.g., "Server Rack GPU 0") + - **Host**: The remote worker's IP address. + - **Port**: The port number used when launching ComfyUI on the remote master/worker (e.g., 8188). +5. **Save** the remote worker configuration. + +## Cloud workers +> ComfyUI instances running on a cloud service like [Runpod](https://get.runpod.io/0bw29uf3ug0p). + +### Deploy Cloud Worker on Runpod + +**On Runpod:** +1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account. +2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. +3. Now go to Pods and find a suitable GPU for your workflows. +4. Choose the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=ak218p52) template and make sure your network drive is mounted. +> To use the ComfyUI Distributed Pod template, you will need to filter instances by CUDA 12.8. +6. Launch your pod. +7. Access your pod using JupyterLabs. +8. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). +> You can use [this guide](model-download-script.md) to make this process easy for you. It will generate a shell script that will automatically download the models you need for a given workflow. +8. If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you git clone ComfyUI-Distributed into custom_nodes. ⚠️ **Required!** +9. Download any additional custom nodes using the ComfyUI Manager. + +**On the Main Machine:** +1. Launch a Cloudflare tunnel + - Download from here: [https://github.com/cloudflare/cloudflared/releases](https://github.com/cloudflare/cloudflared/releases) + - Then run, for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` +> Cloudflare tunnels create secure connections without exposing ports directly to the internet +2. Copy the Cloudflare address +3. **Open** the Distributed GPU panel (sidebar on the left). +4. **Edit** the Master's host address and replace it with the Cloudflare address. + - **Click** "Add Worker." + - **Choose** "Cloud". + - **Configure** your cloud worker: + - **Host**: The Runpod address + - **Port**: 443 + - **Save** the remote worker configuration. + +--- + +### Deploy Cloud Worker on Other Platforms + +**On the Cloud Worker machine:** + - Your cloud worker container needs to have the same models and custom nodes as the workflow you want to run on your local machine. + - If your cloud platform doesn't provide a secure connection, use Cloudflare to create a tunnel for the worker. Each GPU needs their own tunnel for their respective port. + - For example: `./cloudflared tunnel --url http://localhost:8188` +1. **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!** +2. **Add** workers in the UI panel if the remote machine has more than one GPU. + - Make sure that they also have `--listen` set in `Extra Args`. + - Then launch them. + +**On the Main Machine:** +1. **Launch** a Cloudflare tunnel on your local machine. + - Download from here: [https://github.com/cloudflare/cloudflared/releases](https://github.com/cloudflare/cloudflared/releases) + - Then run, for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` +2. **Copy** the Cloudflare address +3. **Open** the Distributed GPU panel (sidebar on the left). +4. **Edit** the Master's host address and replace it with the Cloudflare address. +5. **Click** "Add Worker." +6. **Choose** "Cloud". +7. **Configure** your cloud worker: + - **Host**: The remote worker's IP address/domain + - **Port**: 443 +8. **Save** the remote worker configuration. diff --git a/utils/constants.py b/utils/constants.py index a199f24..a4ca5bd 100644 --- a/utils/constants.py +++ b/utils/constants.py @@ -1,10 +1,11 @@ """ Shared constants for ComfyUI-Distributed. """ +import os # Timeouts (in seconds) -WORKER_JOB_TIMEOUT = 10.0 -TILE_COLLECTION_TIMEOUT = 60.0 +WORKER_JOB_TIMEOUT = 30.0 +TILE_COLLECTION_TIMEOUT = 30.0 TILE_WAIT_TIMEOUT = 30.0 PROCESS_TERMINATION_TIMEOUT = 5.0 @@ -31,4 +32,10 @@ QUEUE_INIT_TIMEOUT = 5.0 TILE_SEND_TIMEOUT = 60.0 # Memory operations -MEMORY_CLEAR_DELAY = 0.5 \ No newline at end of file +MEMORY_CLEAR_DELAY = 0.5 + +# Batch processing +MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs) + +# Heartbeat monitoring +HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds \ No newline at end of file diff --git a/utils/usdu_managment.py b/utils/usdu_managment.py new file mode 100644 index 0000000..28fe055 --- /dev/null +++ b/utils/usdu_managment.py @@ -0,0 +1,821 @@ +import asyncio +import time +import json +import copy +import os +import io +from aiohttp import web +import server +from PIL import Image +import numpy as np +import torch + +# Import from other utilities +from .logging import debug_log, log +from .network import handle_api_error, get_client_session +from .image import tensor_to_pil + +# Configure maximum payload size (50MB default, configurable via environment variable) +MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) + +# Import HEARTBEAT_TIMEOUT from constants +from .constants import HEARTBEAT_TIMEOUT + + +# Unified Job Data Structure Keys +JOB_QUEUE = 'queue' +JOB_MODE = 'mode' +JOB_COMPLETED_TASKS = 'completed_tasks' +JOB_WORKER_STATUS = 'worker_status' +JOB_ASSIGNED_TO_WORKERS = 'assigned_to_workers' +JOB_PENDING_TASKS = 'pending_tasks' +JOB_BATCH_SIZE = 'batch_size' # For dynamic +JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image' # For static + +# Task Types +TASK_TYPE_TILE = 'tile' +TASK_TYPE_IMAGE = 'image' + +async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None): + """Unified initialization for job queues in static and dynamic modes.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + debug_log(f"Queue already exists for {multi_job_id}") + return + + job_data = { + JOB_QUEUE: asyncio.Queue(), + JOB_MODE: mode, + JOB_COMPLETED_TASKS: {}, + JOB_WORKER_STATUS: {w: time.time() for w in enabled_workers or []}, + JOB_ASSIGNED_TO_WORKERS: {w: [] for w in enabled_workers or []}, + JOB_PENDING_TASKS: asyncio.Queue(), + } + + if mode == 'dynamic': + job_data[JOB_BATCH_SIZE] = batch_size + pending_queue = job_data[JOB_PENDING_TASKS] + for i in all_indices or range(batch_size): + await pending_queue.put(i) + debug_log(f"Initialized dynamic queue with {batch_size} pending images") + elif mode == 'static': + job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image + # For static with dynamic distribution, populate all tile indices in pending queue + pending_queue = job_data[JOB_PENDING_TASKS] + total_tiles = batch_size * num_tiles_per_image + for i in range(total_tiles): + await pending_queue.put(i) + debug_log(f"Initialized static queue with {total_tiles} pending tiles for dynamic distribution") + + # Keep backward compatibility - if task assignments provided, still track them + if task_assignments and enabled_workers: + # task_assignments[0] is for master, 1+ are for workers + for i, worker_id in enumerate(enabled_workers): + if i + 1 < len(task_assignments): + job_data[JOB_ASSIGNED_TO_WORKERS][worker_id] = task_assignments[i + 1] + debug_log(f"Worker {worker_id} pre-assigned {len(task_assignments[i + 1])} tasks (legacy mode)") + else: + raise ValueError(f"Unknown mode: {mode}") + + prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data + +def _distribute_tasks(items: list, num_participants: int) -> list[list[any]]: + """Distribute a list of items among N participants (master + workers).""" + if num_participants <= 1: + return [items] + + items_per_participant = len(items) // num_participants + remainder = len(items) % num_participants + assignments = [] + start_idx = 0 + for i in range(num_participants): + count = items_per_participant + (1 if i < remainder else 0) + end_idx = start_idx + count + assignments.append(items[start_idx:end_idx]) + start_idx = end_idx + return assignments + +async def _get_next_task(multi_job_id): + """Get next task from pending queue (generalized for tiles/images).""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data or JOB_PENDING_TASKS not in job_data: + return None + try: + task_id = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=1.0) + return task_id + except asyncio.TimeoutError: + return None + +async def _drain_results_queue(multi_job_id): + """Drain pending results from queue and update completed_tasks. Returns count drained.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data or JOB_QUEUE not in job_data or JOB_COMPLETED_TASKS not in job_data: + return 0 + q = job_data[JOB_QUEUE] + completed_tasks = job_data[JOB_COMPLETED_TASKS] + + collected = 0 + while not q.empty(): + try: + result = await asyncio.wait_for(q.get(), timeout=0.1) + worker_id = result['worker_id'] + is_last = result.get('is_last', False) + + if 'image_idx' in result and 'image' in result: + task_id = result['image_idx'] + if task_id not in completed_tasks: + completed_tasks[task_id] = result['image'] + collected += 1 + elif 'tiles' in result: + for tile_data in result['tiles']: + task_id = tile_data.get('global_idx', tile_data['tile_idx']) + if task_id not in completed_tasks: + completed_tasks[task_id] = tile_data + collected += 1 + elif 'tensor' in result and 'tile_idx' in result: # Single tile backward compat + task_id = result.get('global_idx', result['tile_idx']) + if task_id not in completed_tasks: + completed_tasks[task_id] = { + 'tensor': result['tensor'], + 'tile_idx': result['tile_idx'], + 'x': result['x'], + 'y': result['y'], + 'extracted_width': result['extracted_width'], + 'extracted_height': result['extracted_height'], + 'padding': result['padding'], + 'batch_idx': result.get('batch_idx', 0), + 'global_idx': task_id + } + collected += 1 + + if is_last: + # Track worker completion + if worker_id in job_data[JOB_WORKER_STATUS]: + del job_data[JOB_WORKER_STATUS][worker_id] + except asyncio.TimeoutError: + break + + return collected + +async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks): + """Check timed out workers and requeue their tasks. Returns requeued count.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data: + return 0 + + current_time = time.time() + requeued_count = 0 + completed_tasks = job_data.get(JOB_COMPLETED_TASKS, {}) + + for worker, last_heartbeat in list(job_data.get(JOB_WORKER_STATUS, {}).items()): + if current_time - last_heartbeat > HEARTBEAT_TIMEOUT: + log(f"Worker {worker} timed out") + for task_id in job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, []): + if task_id not in completed_tasks: + await job_data[JOB_PENDING_TASKS].put(task_id) + requeued_count += 1 + if JOB_WORKER_STATUS in job_data: + del job_data[JOB_WORKER_STATUS][worker] + if JOB_ASSIGNED_TO_WORKERS in job_data: + job_data[JOB_ASSIGNED_TO_WORKERS][worker] = [] + + return requeued_count + +async def _get_completed_count(multi_job_id): + """Get count of completed tasks.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and JOB_COMPLETED_TASKS in job_data: + return len(job_data[JOB_COMPLETED_TASKS]) + return 0 + +async def _mark_task_completed(multi_job_id, task_id, result): + """Mark a task as completed.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and JOB_COMPLETED_TASKS in job_data: + job_data[JOB_COMPLETED_TASKS][task_id] = result + +async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id): + """Send heartbeat to master.""" + try: + data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)} + session = await get_client_session() + url = f"{master_url}/distributed/heartbeat" + async with session.post(url, json=data) as response: + response.raise_for_status() + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + +async def _cleanup_job(multi_job_id): + """Cleanup the job data.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + del prompt_server.distributed_pending_tile_jobs[multi_job_id] + debug_log(f"Cleaned up job {multi_job_id}") + +# API Endpoints (generalized) + +@server.PromptServer.instance.routes.post("/distributed/request_task") +async def request_task_endpoint(request): + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + mode = job_data.get(JOB_MODE) + pending_queue = job_data.get(JOB_PENDING_TASKS) + if pending_queue: + try: + task_id = await asyncio.wait_for(pending_queue.get(), timeout=0.1) + if JOB_ASSIGNED_TO_WORKERS in job_data and worker_id in job_data[JOB_ASSIGNED_TO_WORKERS]: + job_data[JOB_ASSIGNED_TO_WORKERS][worker_id].append(task_id) + if JOB_WORKER_STATUS in job_data: + job_data[JOB_WORKER_STATUS][worker_id] = time.time() + remaining = pending_queue.qsize() + debug_log(f"Assigned task {task_id} to worker {worker_id} in {mode} mode") + return web.json_response({"task_id": task_id, "estimated_remaining": remaining, "mode": mode}) + except asyncio.TimeoutError: + return web.json_response({"task_id": None}) + else: + return await handle_api_error(request, "No pending tasks", 400) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + + +@server.PromptServer.instance.routes.post("/distributed/heartbeat") +async def heartbeat_endpoint(request): + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_WORKER_STATUS in job_data: + job_data[JOB_WORKER_STATUS][worker_id] = time.time() + debug_log(f"Heartbeat from worker {worker_id}") + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Worker status tracking not available", 400) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + +@server.PromptServer.instance.routes.post("/distributed/submit_tiles") +async def submit_tiles_endpoint(request): + """Endpoint for workers to submit processed tiles in static mode.""" + try: + content_length = request.headers.get('content-length') + if content_length and int(content_length) > MAX_PAYLOAD_SIZE: + return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413) + + data = await request.post() + multi_job_id = data.get('multi_job_id') + worker_id = data.get('worker_id') + is_last = data.get('is_last', 'False').lower() == 'true' + + if multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + + batch_size = int(data.get('batch_size', 0)) + tiles = [] + + # Handle completion signal + if batch_size == 0 and is_last: + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'static': + return await handle_api_error(request, "Mode mismatch: expected static mode", 400) + if JOB_QUEUE in job_data: + await job_data[JOB_QUEUE].put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) + + # Handle batch tiles with metadata + if batch_size > 0: + padding = int(data.get('padding', 32)) + metadata_field = data.get('tiles_metadata') + if metadata_field: + if hasattr(metadata_field, 'file'): + metadata_str = metadata_field.file.read().decode('utf-8') + elif isinstance(metadata_field, (bytes, bytearray)): + metadata_str = metadata_field.decode('utf-8') + else: + metadata_str = str(metadata_field) + + metadata = json.loads(metadata_str) + if len(metadata) != batch_size: + return await handle_api_error(request, "Metadata length mismatch", 400) + + tile_data_list = [] + + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + continue + + img_data = tile_field.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,] + + if i < len(metadata): + tile_meta = metadata[i] + tile_idx = tile_meta.get('tile_idx', i) + tile_info = { + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': tile_meta['x'], + 'y': tile_meta['y'], + 'extracted_width': tile_meta['extracted_width'], + 'extracted_height': tile_meta['extracted_height'], + 'padding': padding, + 'batch_idx': tile_meta.get('batch_idx', 0), + 'global_idx': tile_meta.get('global_idx', tile_idx) + } + tile_data_list.append(tile_info) + + tile_data_list.sort(key=lambda x: x['tile_idx']) + tiles.extend(tile_data_list) + else: + # Legacy format + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + continue + + tile_idx = int(data.get(f'tile_{i}_idx', i)) + x = int(data.get(f'tile_{i}_x', 0)) + y = int(data.get(f'tile_{i}_y', 0)) + extracted_width = int(data.get(f'tile_{i}_width', 512)) + extracted_height = int(data.get(f'tile_{i}_height', 512)) + + img_data = tile_field.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,] + + tiles.append({ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding + }) + else: + # Single tile legacy + image_file = data.get('image') + if not image_file: + return await handle_api_error(request, "Missing image data", 400) + + tile_idx = int(data.get('tile_idx', 0)) + x = int(data.get('x', 0)) + y = int(data.get('y', 0)) + extracted_width = int(data.get('extracted_width', 512)) + extracted_height = int(data.get('extracted_height', 512)) + padding = int(data.get('padding', 32)) + + 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,] + + tiles = [{ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding, + 'batch_idx': 0, + 'global_idx': tile_idx + }] + + # Submit tiles to queue + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'static': + return await handle_api_error(request, "Mode mismatch: expected static mode", 400) + + q = job_data[JOB_QUEUE] + if batch_size > 0 or len(tiles) > 0: + await q.put({ + 'worker_id': worker_id, + 'tiles': tiles, + 'is_last': is_last + }) + else: + await q.put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] + }) + + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + +@server.PromptServer.instance.routes.post("/distributed/submit_image") +async def submit_image_endpoint(request): + """Endpoint for workers to submit processed images in dynamic mode.""" + try: + content_length = request.headers.get('content-length') + if content_length and int(content_length) > MAX_PAYLOAD_SIZE: + return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413) + + data = await request.post() + multi_job_id = data.get('multi_job_id') + worker_id = data.get('worker_id') + is_last = data.get('is_last', 'False').lower() == 'true' + + if multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + + # Handle image submission + if 'full_image' in data and 'image_idx' in data: + image_idx = int(data.get('image_idx')) + img_data = data['full_image'].file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + + debug_log(f"Received full image {image_idx} from worker {worker_id}") + + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic': + return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400) + if JOB_QUEUE in job_data: + await job_data[JOB_QUEUE].put({ + 'worker_id': worker_id, + 'image_idx': image_idx, + 'image': img, + 'is_last': is_last + }) + return web.json_response({"status": "success"}) + + # Handle completion signal (no image data) + elif is_last: + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic': + return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400) + if JOB_QUEUE in job_data: + await job_data[JOB_QUEUE].put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] # For compatibility + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Missing image data or invalid request", 400) + + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + +# Keep legacy endpoint for backward compatibility +@server.PromptServer.instance.routes.post("/distributed/tile_complete") +async def tile_complete_endpoint(request): + try: + content_length = request.headers.get('content-length') + if content_length and int(content_length) > MAX_PAYLOAD_SIZE: + return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413) + + data = await request.post() + multi_job_id = data.get('multi_job_id') + worker_id = data.get('worker_id') + is_last = data.get('is_last', 'False').lower() == 'true' + + if multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + + if 'full_image' in data and 'image_idx' in data: + image_idx = int(data.get('image_idx')) + img_data = data['full_image'].file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + + debug_log(f"Received full image {image_idx} from worker {worker_id}") + + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic': + return await handle_api_error(request, "Mode mismatch for image submission", 400) + if JOB_QUEUE in job_data: + await job_data[JOB_QUEUE].put({ + 'worker_id': worker_id, + 'image_idx': image_idx, + 'image': img, + 'is_last': is_last + }) + return web.json_response({"status": "success"}) + + batch_size = int(data.get('batch_size', 0)) + tiles = [] + + if batch_size == 0 and is_last: + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_QUEUE in job_data: + await job_data[JOB_QUEUE].put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) + + if batch_size > 0: + padding = int(data.get('padding', 32)) + metadata_field = data.get('tiles_metadata') + if metadata_field: + if hasattr(metadata_field, 'file'): + metadata_str = metadata_field.file.read().decode('utf-8') + elif isinstance(metadata_field, (bytes, bytearray)): + metadata_str = metadata_field.decode('utf-8') + else: + metadata_str = str(metadata_field) + + metadata = json.loads(metadata_str) + if len(metadata) != batch_size: + return await handle_api_error(request, "Metadata length mismatch", 400) + + tile_data_list = [] + + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + continue + + img_data = tile_field.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,] + + if i < len(metadata): + tile_meta = metadata[i] + tile_idx = tile_meta.get('tile_idx', i) + tile_info = { + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': tile_meta['x'], + 'y': tile_meta['y'], + 'extracted_width': tile_meta['extracted_width'], + 'extracted_height': tile_meta['extracted_height'], + 'padding': padding, + 'batch_idx': tile_meta.get('batch_idx', 0), + 'global_idx': tile_meta.get('global_idx', tile_idx) + } + tile_data_list.append(tile_info) + + tile_data_list.sort(key=lambda x: x['tile_idx']) + tiles.extend(tile_data_list) + else: + # Legacy + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + continue + + tile_idx = int(data.get(f'tile_{i}_idx', i)) + x = int(data.get(f'tile_{i}_x', 0)) + y = int(data.get(f'tile_{i}_y', 0)) + extracted_width = int(data.get(f'tile_{i}_width', 512)) + extracted_height = int(data.get(f'tile_{i}_height', 512)) + + img_data = tile_field.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,] + + tiles.append({ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding + }) + else: + # Single tile legacy + image_file = data.get('image') + if not image_file: + return await handle_api_error(request, "Missing image data", 400) + + tile_idx = int(data.get('tile_idx', 0)) + x = int(data.get('x', 0)) + y = int(data.get('y', 0)) + extracted_width = int(data.get('extracted_width', 512)) + extracted_height = int(data.get('extracted_height', 512)) + padding = int(data.get('padding', 32)) + + 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,] + + tiles = [{ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding, + 'batch_idx': 0, + 'global_idx': tile_idx + }] + + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'static': + return await handle_api_error(request, "Mode mismatch for tile submission", 400) + + q = job_data[JOB_QUEUE] + if batch_size > 0 or len(tiles) > 0: + await q.put({ + 'worker_id': worker_id, + 'tiles': tiles, + 'is_last': is_last + }) + else: + await q.put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] + }) + + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + + + +# Helper functions for shallow copying conditioning without duplicating models +def clone_control_chain(control, clone_hint=True): + """Shallow copy the ControlNet chain, optionally cloning hints but sharing models.""" + if control is None: + return None + new_control = copy.copy(control) # Shallow copy (shares model) + if clone_hint and hasattr(control, 'cond_hint_original'): + new_control.cond_hint_original = control.cond_hint_original.clone() + if hasattr(control, 'previous_controlnet'): + new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint) + return new_control + +def clone_conditioning(cond_list, clone_hints=True): + """Clone conditioning without duplicating ControlNet models.""" + new_cond = [] + for emb, cond_dict in cond_list: + new_emb = emb.clone() if emb is not None else None + new_dict = cond_dict.copy() + if 'control' in new_dict: + new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints) + if 'mask' in new_dict: + new_dict['mask'] = new_dict['mask'].clone() + # Handle other potential fields if needed + if 'pooled_output' in new_dict: + new_dict['pooled_output'] = new_dict['pooled_output'].clone() + if 'area' in new_dict: + new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple + new_cond.append([new_emb, new_dict]) + return new_cond + +# Helper function to ensure persistent state is initialized +def ensure_tile_jobs_initialized(): + """Ensure tile job storage is initialized on the server instance.""" + prompt_server = server.PromptServer.instance + if not hasattr(prompt_server, 'distributed_pending_tile_jobs'): + debug_log("Initializing persistent tile job queue on server instance.") + prompt_server.distributed_pending_tile_jobs = {} + prompt_server.distributed_tile_jobs_lock = asyncio.Lock() + else: + # Clean up any legacy queue structures that don't have the 'mode' field + # (Should be rare after fixes, but keep for safety) + to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items() + if not isinstance(job_data, dict) or 'mode' not in job_data] + for job_id in to_remove: + debug_log(f"Removing legacy queue structure for job {job_id}") + del prompt_server.distributed_pending_tile_jobs[job_id] + return prompt_server + +# API Endpoint for tile completion +@server.PromptServer.instance.routes.post("/distributed/request_image") +async def request_image_endpoint(request): + """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode).""" + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + return await handle_api_error(request, "Invalid job data structure", 500) + + mode = job_data['mode'] + + # Handle both dynamic and static modes + if mode == 'dynamic' and 'pending_images' in job_data: + pending_queue = job_data['pending_images'] + elif mode == 'static' and JOB_PENDING_TASKS in job_data: + pending_queue = job_data[JOB_PENDING_TASKS] + else: + return await handle_api_error(request, f"Invalid {mode} mode configuration", 400) + + try: + task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) + # Track assigned task + if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']: + job_data['assigned_to_workers'][worker_id].append(task_idx) + # Update worker heartbeat + if 'worker_status' in job_data: + job_data['worker_status'][worker_id] = time.time() + # Get estimated remaining count + remaining = pending_queue.qsize() # Approximate + + # Return appropriate response based on mode + if mode == 'dynamic': + debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}") + return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining}) + else: # static + debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}") + return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining}) + except asyncio.TimeoutError: + if mode == 'dynamic': + return web.json_response({"image_idx": None}) + else: + return web.json_response({"tile_idx": None}) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + +@server.PromptServer.instance.routes.get("/distributed/job_status") +async def job_status_endpoint(request): + """Endpoint to check if a job is ready.""" + multi_job_id = request.query.get('multi_job_id') + if not multi_job_id: + return web.json_response({"ready": False}) + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data) + return web.json_response({"ready": ready}) + + diff --git a/utils/usdu_utils.py b/utils/usdu_utils.py new file mode 100644 index 0000000..21214a7 --- /dev/null +++ b/utils/usdu_utils.py @@ -0,0 +1,456 @@ +import numpy as np +from PIL import Image, ImageFilter +import torch +import torch.nn.functional as F +from torchvision.transforms import GaussianBlur +import math +import os + +if (not hasattr(Image, 'Resampling')): # For older versions of Pillow + Image.Resampling = Image + +BLUR_KERNEL_SIZE = 15 + + +def tensor_to_pil(img_tensor, batch_index=0): + # Takes a batch of images in the form of a tensor of shape [batch_size, height, width, channels] + # and returns an RGB PIL Image. Assumes channels=3 + return Image.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8)) + + +def pil_to_tensor(image): + # Takes a PIL image and returns a tensor of shape [1, height, width, channels] + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image).unsqueeze(0) + if len(image.shape) == 3: # If the image is grayscale, add a channel dimension + image = image.unsqueeze(-1) + return image + + +def controlnet_hint_to_pil(tensor, batch_index=0): + return tensor_to_pil(tensor.movedim(1, -1), batch_index) + + +def pil_to_controlnet_hint(img): + return pil_to_tensor(img).movedim(-1, 1) + + +def crop_tensor(tensor, region): + # Takes a tensor of shape [batch_size, height, width, channels] and crops it to the given region + x1, y1, x2, y2 = region + return tensor[:, y1:y2, x1:x2, :] + + +def resize_tensor(tensor, size, mode="nearest-exact"): + # Takes a tensor of shape [B, C, H, W] and resizes + # it to a shape of [B, C, size[0], size[1]] using the given mode + return torch.nn.functional.interpolate(tensor, size=size, mode=mode) + + +def get_crop_region(mask, pad=0): + # Takes a black and white PIL image in 'L' mode and returns the coordinates of the white rectangular mask region + # Should be equivalent to the get_crop_region function from https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/masking.py + coordinates = mask.getbbox() + if coordinates is not None: + x1, y1, x2, y2 = coordinates + else: + x1, y1, x2, y2 = mask.width, mask.height, 0, 0 + # Apply padding + x1 = max(x1 - pad, 0) + y1 = max(y1 - pad, 0) + x2 = min(x2 + pad, mask.width) + y2 = min(y2 + pad, mask.height) + return fix_crop_region((x1, y1, x2, y2), (mask.width, mask.height)) + + +def fix_crop_region(region, image_size): + # Remove the extra pixel added by the get_crop_region function + image_width, image_height = image_size + x1, y1, x2, y2 = region + if x2 < image_width: + x2 -= 1 + if y2 < image_height: + y2 -= 1 + return x1, y1, x2, y2 + + +def expand_crop(region, width, height, target_width, target_height): + ''' + Expands a crop region to a specified target size. + :param region: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points + of the rectangular region. Expected to have x2 > x1 and y2 > y1. + :param width: The width of the image the crop region is from. + :param height: The height of the image the crop region is from. + :param target_width: The desired width of the crop region. + :param target_height: The desired height of the crop region. + ''' + x1, y1, x2, y2 = region + actual_width = x2 - x1 + actual_height = y2 - y1 + # target_width = math.ceil(actual_width / 8) * 8 + # target_height = math.ceil(actual_height / 8) * 8 + + # Try to expand region to the right of half the difference + width_diff = target_width - actual_width + x2 = min(x2 + width_diff // 2, width) + # Expand region to the left of the difference including the pixels that could not be expanded to the right + width_diff = target_width - (x2 - x1) + x1 = max(x1 - width_diff, 0) + # Try the right again + width_diff = target_width - (x2 - x1) + x2 = min(x2 + width_diff, width) + + # Try to expand region to the bottom of half the difference + height_diff = target_height - actual_height + y2 = min(y2 + height_diff // 2, height) + # Expand region to the top of the difference including the pixels that could not be expanded to the bottom + height_diff = target_height - (y2 - y1) + y1 = max(y1 - height_diff, 0) + # Try the bottom again + height_diff = target_height - (y2 - y1) + y2 = min(y2 + height_diff, height) + + return (x1, y1, x2, y2), (target_width, target_height) + + +def resize_region(region, init_size, resize_size): + # Resize a crop so that it fits an image that was resized to the given width and height + x1, y1, x2, y2 = region + init_width, init_height = init_size + resize_width, resize_height = resize_size + x1 = math.floor(x1 * resize_width / init_width) + x2 = math.ceil(x2 * resize_width / init_width) + y1 = math.floor(y1 * resize_height / init_height) + y2 = math.ceil(y2 * resize_height / init_height) + return (x1, y1, x2, y2) + + +def pad_image(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False): + ''' + Pads an image with the given number of pixels on each side and fills the padding with data from the edges. + :param image: A PIL image + :param left_pad: The number of pixels to pad on the left side + :param right_pad: The number of pixels to pad on the right side + :param top_pad: The number of pixels to pad on the top side + :param bottom_pad: The number of pixels to pad on the bottom side + :param blur: Whether to blur the padded edges + :return: A PIL image with size (image.width + left_pad + right_pad, image.height + top_pad + bottom_pad) + ''' + left_edge = image.crop((0, 1, 1, image.height - 1)) + right_edge = image.crop((image.width - 1, 1, image.width, image.height - 1)) + top_edge = image.crop((1, 0, image.width - 1, 1)) + bottom_edge = image.crop((1, image.height - 1, image.width - 1, image.height)) + new_width = image.width + left_pad + right_pad + new_height = image.height + top_pad + bottom_pad + padded_image = Image.new(image.mode, (new_width, new_height)) + padded_image.paste(image, (left_pad, top_pad)) + if fill: + for i in range(left_pad): + edge = left_edge.resize( + (1, new_height - i * (top_pad + bottom_pad) // left_pad), resample=Image.Resampling.NEAREST) + padded_image.paste(edge, (i, i * top_pad // left_pad)) + for i in range(right_pad): + edge = right_edge.resize( + (1, new_height - i * (top_pad + bottom_pad) // right_pad), resample=Image.Resampling.NEAREST) + padded_image.paste(edge, (new_width - 1 - i, i * top_pad // right_pad)) + for i in range(top_pad): + edge = top_edge.resize( + (new_width - i * (left_pad + right_pad) // top_pad, 1), resample=Image.Resampling.NEAREST) + padded_image.paste(edge, (i * left_pad // top_pad, i)) + for i in range(bottom_pad): + edge = bottom_edge.resize( + (new_width - i * (left_pad + right_pad) // bottom_pad, 1), resample=Image.Resampling.NEAREST) + padded_image.paste(edge, (i * left_pad // bottom_pad, new_height - 1 - i)) + if blur and not (left_pad == right_pad == top_pad == bottom_pad == 0): + padded_image = padded_image.filter(ImageFilter.GaussianBlur(BLUR_KERNEL_SIZE)) + padded_image.paste(image, (left_pad, top_pad)) + return padded_image + + +def pad_image2(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False): + ''' + Pads an image with the given number of pixels on each side and fills the padding with data from the edges. + Faster than pad_image, but only pads with edge data in straight lines. + :param image: A PIL image + :param left_pad: The number of pixels to pad on the left side + :param right_pad: The number of pixels to pad on the right side + :param top_pad: The number of pixels to pad on the top side + :param bottom_pad: The number of pixels to pad on the bottom side + :param blur: Whether to blur the padded edges + :return: A PIL image with size (image.width + left_pad + right_pad, image.height + top_pad + bottom_pad) + ''' + left_edge = image.crop((0, 1, 1, image.height - 1)) + right_edge = image.crop((image.width - 1, 1, image.width, image.height - 1)) + top_edge = image.crop((1, 0, image.width - 1, 1)) + bottom_edge = image.crop((1, image.height - 1, image.width - 1, image.height)) + new_width = image.width + left_pad + right_pad + new_height = image.height + top_pad + bottom_pad + padded_image = Image.new(image.mode, (new_width, new_height)) + padded_image.paste(image, (left_pad, top_pad)) + if fill: + if left_pad > 0: + padded_image.paste(left_edge.resize((left_pad, new_height), resample=Image.Resampling.NEAREST), (0, 0)) + if right_pad > 0: + padded_image.paste(right_edge.resize((right_pad, new_height), + resample=Image.Resampling.NEAREST), (new_width - right_pad, 0)) + if top_pad > 0: + padded_image.paste(top_edge.resize((new_width, top_pad), resample=Image.Resampling.NEAREST), (0, 0)) + if bottom_pad > 0: + padded_image.paste(bottom_edge.resize((new_width, bottom_pad), + resample=Image.Resampling.NEAREST), (0, new_height - bottom_pad)) + if blur and not (left_pad == right_pad == top_pad == bottom_pad == 0): + padded_image = padded_image.filter(ImageFilter.GaussianBlur(BLUR_KERNEL_SIZE)) + padded_image.paste(image, (left_pad, top_pad)) + return padded_image + + +def pad_tensor(tensor, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False): + ''' + Pads an image tensor with the given number of pixels on each side and fills the padding with data from the edges. + :param tensor: A tensor of shape [B, H, W, C] + :param left_pad: The number of pixels to pad on the left side + :param right_pad: The number of pixels to pad on the right side + :param top_pad: The number of pixels to pad on the top side + :param bottom_pad: The number of pixels to pad on the bottom side + :param blur: Whether to blur the padded edges + :return: A tensor of shape [B, H + top_pad + bottom_pad, W + left_pad + right_pad, C] + ''' + batch_size, channels, height, width = tensor.shape + h_pad = left_pad + right_pad + v_pad = top_pad + bottom_pad + new_width = width + h_pad + new_height = height + v_pad + + # Create empty image + padded = torch.zeros((batch_size, channels, new_height, new_width), dtype=tensor.dtype) + + # Copy the original image into the centor of the padded tensor + padded[:, :, top_pad:top_pad + height, left_pad:left_pad + width] = tensor + + # Duplicate the edges of the original image into the padding + if top_pad > 0: + padded[:, :, :top_pad, :] = padded[:, :, top_pad:top_pad + 1, :] # Top edge + if bottom_pad > 0: + padded[:, :, -bottom_pad:, :] = padded[:, :, -bottom_pad - 1:-bottom_pad, :] # Bottom edge + if left_pad > 0: + padded[:, :, :, :left_pad] = padded[:, :, :, left_pad:left_pad + 1] # Left edge + if right_pad > 0: + padded[:, :, :, -right_pad:] = padded[:, :, :, -right_pad - 1:-right_pad] # Right edge + + return padded + + +def resize_and_pad_image(image, width, height, fill=False, blur=False): + ''' + Resizes an image to the given width and height and pads it to the given width and height. + :param image: A PIL image + :param width: The width of the resized image + :param height: The height of the resized image + :param fill: Whether to fill the padding with data from the edges + :param blur: Whether to blur the padded edges + :return: A PIL image of size (width, height) + ''' + width_ratio = width / image.width + height_ratio = height / image.height + if height_ratio > width_ratio: + resize_ratio = width_ratio + else: + resize_ratio = height_ratio + resize_width = round(image.width * resize_ratio) + resize_height = round(image.height * resize_ratio) + resized = image.resize((resize_width, resize_height), resample=Image.Resampling.LANCZOS) + # Pad the sides of the image to get the image to the desired size that wasn't covered by the resize + horizontal_pad = (width - resize_width) // 2 + vertical_pad = (height - resize_height) // 2 + result = pad_image2(resized, horizontal_pad, horizontal_pad, vertical_pad, vertical_pad, fill, blur) + result = result.resize((width, height), resample=Image.Resampling.LANCZOS) + return result, (horizontal_pad, vertical_pad) + + +def resize_and_pad_tensor(tensor, width, height, fill=False, blur=False): + ''' + Resizes an image tensor to the given width and height and pads it to the given width and height. + :param tensor: A tensor of shape [B, H, W, C] + :param width: The width of the resized image + :param height: The height of the resized image + :param fill: Whether to fill the padding with data from the edges + :param blur: Whether to blur the padded edges + :return: A tensor of shape [B, height, width, C] + ''' + # Resize the image to the closest size that maintains the aspect ratio + width_ratio = width / tensor.shape[3] + height_ratio = height / tensor.shape[2] + if height_ratio > width_ratio: + resize_ratio = width_ratio + else: + resize_ratio = height_ratio + resize_width = round(tensor.shape[3] * resize_ratio) + resize_height = round(tensor.shape[2] * resize_ratio) + resized = F.interpolate(tensor, size=(resize_height, resize_width), mode='nearest-exact') + # Pad the sides of the image to get the image to the desired size that wasn't covered by the resize + horizontal_pad = (width - resize_width) // 2 + vertical_pad = (height - resize_height) // 2 + result = pad_tensor(resized, horizontal_pad, horizontal_pad, vertical_pad, vertical_pad, fill, blur) + result = F.interpolate(result, size=(height, width), mode='nearest-exact') + return result + + +def crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): + if "control" not in cond_dict: + return + c = cond_dict["control"] + controlnet = c.copy() + cond_dict["control"] = controlnet + while c is not None: + # hint is shape (B, C, H, W) + hint = controlnet.cond_hint_original + resized_crop = resize_region(region, canvas_size, hint.shape[:-3:-1]) + hint = crop_tensor(hint.movedim(1, -1), resized_crop).movedim(-1, 1) + hint = resize_tensor(hint, tile_size[::-1]) + controlnet.cond_hint_original = hint + c = c.previous_controlnet + controlnet.set_previous_controlnet(c.copy() if c is not None else None) + controlnet = controlnet.previous_controlnet + + +def region_intersection(region1, region2): + """ + Returns the coordinates of the intersection of two rectangular regions. + :param region1: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points + of the first rectangular region. Expected to have x2 > x1 and y2 > y1. + :param region2: The second rectangular region with the same format as the first. + :return: A tuple of the form (x1, y1, x2, y2) denoting the rectangular intersection. + None if there is no intersection. + """ + x1, y1, x2, y2 = region1 + x1_, y1_, x2_, y2_ = region2 + x1 = max(x1, x1_) + y1 = max(y1, y1_) + x2 = min(x2, x2_) + y2 = min(y2, y2_) + if x1 >= x2 or y1 >= y2: + return None + return (x1, y1, x2, y2) + + +def crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): + if "gligen" not in cond_dict: + return + type, model, cond = cond_dict["gligen"] + if type != "position": + from warnings import warn + warn(f"Unknown gligen type {type}") + return + cropped = [] + for c in cond: + emb, h, w, y, x = c + # Get the coordinates of the box in the upscaled image + x1 = x * 8 + y1 = y * 8 + x2 = x1 + w * 8 + y2 = y1 + h * 8 + gligen_upscaled_box = resize_region((x1, y1, x2, y2), init_size, canvas_size) + + # Calculate the intersection of the gligen box and the region + intersection = region_intersection(gligen_upscaled_box, region) + if intersection is None: + continue + x1, y1, x2, y2 = intersection + + # Offset the gligen box so that the origin is at the top left of the tile region + x1 -= region[0] + y1 -= region[1] + x2 -= region[0] + y2 -= region[1] + + # Add the padding + x1 += w_pad + y1 += h_pad + x2 += w_pad + y2 += h_pad + + # Set the new position params + h = (y2 - y1) // 8 + w = (x2 - x1) // 8 + x = x1 // 8 + y = y1 // 8 + cropped.append((emb, h, w, y, x)) + + cond_dict["gligen"] = (type, model, cropped) + + +def crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): + if "area" not in cond_dict: + return + + # Resize the area conditioning to the canvas size and confine it to the tile region + h, w, y, x = cond_dict["area"] + w, h, x, y = 8 * w, 8 * h, 8 * x, 8 * y + x1, y1, x2, y2 = resize_region((x, y, x + w, y + h), init_size, canvas_size) + intersection = region_intersection((x1, y1, x2, y2), region) + if intersection is None: + del cond_dict["area"] + del cond_dict["strength"] + return + x1, y1, x2, y2 = intersection + + # Offset origin to the top left of the tile + x1 -= region[0] + y1 -= region[1] + x2 -= region[0] + y2 -= region[1] + + # Add the padding + x1 += w_pad + y1 += h_pad + x2 += w_pad + y2 += h_pad + + # Set the params for tile + w, h = (x2 - x1) // 8, (y2 - y1) // 8 + x, y = x1 // 8, y1 // 8 + + cond_dict["area"] = (h, w, y, x) + + +def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): + if "mask" not in cond_dict: + return + mask_tensor = cond_dict["mask"] # (B, H, W) + masks = [] + for i in range(mask_tensor.shape[0]): + # Convert to PIL image + mask = tensor_to_pil(mask_tensor, i) # W x H + + # Resize the mask to the canvas size + mask = mask.resize(canvas_size, Image.Resampling.BICUBIC) + + # Crop the mask to the region + mask = mask.crop(region) + + # Add padding + mask, _ = resize_and_pad_image(mask, tile_size[0], tile_size[1], fill=True) + + # Resize the mask to the tile size + if tile_size != mask.size: + mask = mask.resize(tile_size, Image.Resampling.BICUBIC) + + # Convert back to tensor + mask = pil_to_tensor(mask) # (1, H, W, 1) + mask = mask.squeeze(-1) # (1, H, W) + masks.append(mask) + + cond_dict["mask"] = torch.cat(masks, dim=0) # (B, H, W) + + +def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0): + cropped = [] + for emb, x in cond: + cond_dict = x.copy() + n = [emb, cond_dict] + crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) + crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) + crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) + crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) + cropped.append(n) + return cropped \ No newline at end of file diff --git a/web/apiClient.js b/web/apiClient.js new file mode 100644 index 0000000..b4fceaa --- /dev/null +++ b/web/apiClient.js @@ -0,0 +1,146 @@ +import { TIMEOUTS } from './constants.js'; + +export function createApiClient(baseUrl) { + const request = async (endpoint, options = {}, retries = TIMEOUTS.MAX_RETRIES) => { + let lastError; + let delay = TIMEOUTS.RETRY_DELAY; // Initial delay for exponential backoff + + for (let attempt = 0; attempt < retries; attempt++) { + try { + const response = await fetch(`${baseUrl}${endpoint}`, { + headers: { 'Content-Type': 'application/json' }, + ...options + }); + + if (!response.ok) { + const error = await response.json().catch(() => ({ message: 'Request failed' })); + throw new Error(error.message || `HTTP ${response.status}`); + } + + return await response.json(); + } catch (error) { + lastError = error; + console.log(`API Error (attempt ${attempt + 1}/${retries}): ${endpoint} - ${error.message}`); + if (attempt < retries - 1) { + await new Promise(resolve => setTimeout(resolve, delay)); + delay *= 2; // Exponential backoff + } + } + } + throw lastError; + }; + + return { + // Config endpoints + async getConfig() { + return request('/distributed/config'); + }, + + async updateWorker(workerId, data) { + return request('/distributed/config/update_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId, ...data }) + }); + }, + + async deleteWorker(workerId) { + return request('/distributed/config/delete_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + async updateSetting(key, value) { + return request('/distributed/config/update_setting', { + method: 'POST', + body: JSON.stringify({ key, value }) + }); + }, + + async updateMaster(data) { + return request('/distributed/config/update_master', { + method: 'POST', + body: JSON.stringify(data) + }); + }, + + // Worker management endpoints + async launchWorker(workerId) { + return request('/distributed/launch_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + async stopWorker(workerId) { + return request('/distributed/stop_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + async getManagedWorkers() { + return request('/distributed/managed_workers'); + }, + + async getWorkerLog(workerId, lines = 1000) { + return request(`/distributed/worker_log/${workerId}?lines=${lines}`); + }, + + async clearLaunchingFlag(workerId) { + return request('/distributed/worker/clear_launching', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + // Job preparation + async prepareJob(multiJobId) { + return request('/distributed/prepare_job', { + method: 'POST', + body: JSON.stringify({ multi_job_id: multiJobId }) + }); + }, + + // Image loading + async loadImage(imagePath) { + return request('/distributed/load_image', { + method: 'POST', + body: JSON.stringify({ image_path: imagePath }) + }); + }, + + // Network info + async getNetworkInfo() { + return request('/distributed/network_info'); + }, + + // Status checking (with timeout) + async checkStatus(url, timeout = TIMEOUTS.DEFAULT_FETCH) { + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), timeout); + + try { + const response = await fetch(url, { + method: 'GET', + mode: 'cors', + signal: controller.signal + }); + clearTimeout(timeoutId); + + if (!response.ok) throw new Error(`HTTP ${response.status}`); + return await response.json(); + } catch (error) { + clearTimeout(timeoutId); + throw error; + } + }, + + // Batch status checking + async checkMultipleStatuses(urls) { + return Promise.allSettled( + urls.map(url => this.checkStatus(url)) + ); + } + }; +} \ No newline at end of file diff --git a/web/constants.js b/web/constants.js new file mode 100644 index 0000000..df0d251 --- /dev/null +++ b/web/constants.js @@ -0,0 +1,153 @@ +export const BUTTON_STYLES = { + // Base styles with unified padding + base: "width: 100%; padding: 4px 14px; color: white; border: none; border-radius: 4px; cursor: pointer; transition: all 0.2s; font-size: 12px; font-weight: 500;", + + // Context-specific combined styles + workerControl: "flex: 1; font-size: 11px;", + + // Layout modifiers + hidden: "display: none;", + marginLeftAuto: "margin-left: auto;", + + // Color variants + cancel: "background-color: #555;", + info: "background-color: #333;", + success: "background-color: #4a7c4a;", + error: "background-color: #7c4a4a;", + launch: "background-color: #4a7c4a;", + stop: "background-color: #7c4a4a;", + log: "background-color: #685434;", + clearMemory: "background-color: #555; padding: 6px 14px;", + interrupt: "background-color: #555; padding: 6px 14px;", +}; + +export const STATUS_COLORS = { + DISABLED_GRAY: "#666", + OFFLINE_RED: "#c04c4c", + ONLINE_GREEN: "#3ca03c", + PROCESSING_YELLOW: "#f0ad4e" +}; + +export const UI_COLORS = { + MUTED_TEXT: "#888", + SECONDARY_TEXT: "#ccc", + BORDER_LIGHT: "#555", + BORDER_DARK: "#444", + BORDER_DARKER: "#3a3a3a", + BACKGROUND_DARK: "#2a2a2a", + BACKGROUND_DARKER: "#1e1e1e", + ICON_COLOR: "#666", + ACCENT_COLOR: "#777" +}; + +export const PULSE_ANIMATION_CSS = ` + @keyframes pulse { + 0% { + opacity: 1; + transform: scale(0.8); + box-shadow: 0 0 0 0 rgba(240, 173, 78, 0.7); + } + 50% { + opacity: 0.3; + transform: scale(1.1); + box-shadow: 0 0 0 6px rgba(240, 173, 78, 0); + } + 100% { + opacity: 1; + transform: scale(0.8); + box-shadow: 0 0 0 0 rgba(240, 173, 78, 0); + } + } + .status-pulsing { + animation: pulse 1.2s ease-in-out infinite; + transform-origin: center; + } + + /* Button hover effects */ + .distributed-button:hover:not(:disabled) { + filter: brightness(1.2); + transition: filter 0.2s ease; + } + .distributed-button:disabled { + opacity: 0.6; + cursor: not-allowed; + } + + /* Settings button animation */ + .settings-btn { + transition: transform 0.2s ease; + } + + + /* Expanded settings panel */ + .worker-settings { + max-height: 0; + overflow: hidden; + opacity: 0; + transition: max-height 0.3s ease, opacity 0.3s ease, padding 0.3s ease, margin 0.3s ease; + } + .worker-settings.expanded { + max-height: 500px; + opacity: 1; + padding: 12px 0; + } +`; + +export const UI_STYLES = { + statusDot: "display: inline-block; width: 10px; height: 10px; border-radius: 50%; margin-right: 10px;", + controlsDiv: "padding: 0 12px 12px 12px; display: flex; gap: 6px;", + formGroup: "display: flex; flex-direction: column; gap: 5px;", + formLabel: "font-size: 12px; color: #ccc; font-weight: 500;", + formInput: "padding: 6px 10px; background: #2a2a2a; border: 1px solid #444; color: white; font-size: 12px; border-radius: 4px; transition: border-color 0.2s;", + + // Card styles + cardBase: "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex;", + workerCard: "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex; background: #2a2a2a;", + cardBlueprint: "border: 2px dashed #555; cursor: pointer; transition: all 0.2s ease; background: rgba(255, 255, 255, 0.02);", + cardAdd: "border: 1px dashed #444; cursor: pointer; transition: all 0.2s ease; background: transparent;", + + // Column styles + columnBase: "display: flex; align-items: center; justify-content: center;", + checkboxColumn: "flex: 0 0 44px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #3a3a3a; cursor: default; background: rgba(0,0,0,0.1);", + contentColumn: "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;", + iconColumn: "width: 44px; flex-shrink: 0; font-size: 20px; color: #666;", + + // Row and content styles + infoRow: "display: flex; align-items: center; padding: 12px; cursor: pointer; min-height: 64px;", + workerContent: "display: flex; align-items: center; gap: 10px; flex: 1;", + + // Form and controls styles + buttonGroup: "display: flex; gap: 4px; margin-top: 10px;", + settingsForm: "display: flex; flex-direction: column; gap: 10px;", + checkboxGroup: "display: flex; align-items: center; gap: 8px; margin: 5px 0;", + formLabelClickable: "font-size: 12px; color: #ccc; cursor: pointer;", + settingsToggle: "display: flex; align-items: center; gap: 6px; padding: 4px 0; cursor: pointer; user-select: none;", + controlsWrapper: "display: flex; gap: 6px; align-items: stretch; width: 100%;", + + // Existing styles + settingsArrow: "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto; padding: 4px;", + infoBox: "background-color: #333; color: #999; padding: 5px 14px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1; font-weight: 500;", + workerSettings: "margin: 0 12px; padding: 0 12px; background: #1e1e1e; border-radius: 4px; border: 1px solid #2a2a2a;" +}; + +export const TIMEOUTS = { + DEFAULT_FETCH: 5000, // ms for general API calls + STATUS_CHECK: 1200, // ms for status checks + LAUNCH: 90000, // ms for worker launch (longer for model loading) + RETRY_DELAY: 1000, // initial delay for exponential backoff + MAX_RETRIES: 3, // max retry attempts + + // UI feedback delays + BUTTON_RESET: 3000, // button text/state reset after actions + FLASH_SHORT: 1000, // brief success feedback + FLASH_MEDIUM: 1500, // medium error feedback + FLASH_LONG: 2000, // longer error feedback + + // Operational delays + POST_ACTION_DELAY: 500, // delay after operations before status checks + STATUS_CHECK_DELAY: 100, // brief delay before status checks + + // Background tasks + LOG_REFRESH: 2000, // log auto-refresh interval + IMAGE_CACHE_CLEAR: 30000 // delay before clearing image cache +}; \ No newline at end of file diff --git a/web/executionUtils.js b/web/executionUtils.js new file mode 100644 index 0000000..6e2c583 --- /dev/null +++ b/web/executionUtils.js @@ -0,0 +1,590 @@ +import { api } from "../../scripts/api.js"; +import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo } from './workerUtils.js'; +import { TIMEOUTS } from './constants.js'; + +/** + * Convert paths in the API prompt to match the target platform's separator + * @param {Object} apiPrompt - The workflow API prompt + * @param {string} targetSeparator - The target path separator ('\\' or '/') + * @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); + }; + + function convert(obj) { + if (typeof obj === 'string') { + // Only convert strings that look like file paths + if ((obj.includes('\\') || obj.includes('/')) && isLikelyFilename(obj)) { + // Replace any path separator with the target one + return obj.replace(/[\\\/]/g, targetSeparator); + } + return obj; + } else if (Array.isArray(obj)) { + return obj.map(convert); + } else if (typeof obj === 'object' && obj !== null) { + const newObj = {}; + for (const [key, value] of Object.entries(obj)) { + newObj[key] = convert(value); + } + return newObj; + } + return obj; + } + + return convert(apiPrompt); +} + +export function setupInterceptor(extension) { + api.queuePrompt = async (number, prompt) => { + if (extension.isEnabled) { + const hasCollector = findNodesByClass(prompt.output, "DistributedCollector").length > 0; + const hasDistUpscale = findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0; + + if (hasCollector || hasDistUpscale) { + const result = await executeParallelDistributed(extension, prompt); + // Immediate status check for instant feedback + extension.checkAllWorkerStatuses(); + // Another check after a short delay to catch state changes + setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY); + return result; + } + } + return extension.originalQueuePrompt(number, prompt); + }; +} + +export async function executeParallelDistributed(extension, promptWrapper) { + try { + const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution + const enabledWorkers = extension.enabledWorkers; + + // Pre-flight health check on all enabled workers + const activeWorkers = await performPreflightCheck(extension, enabledWorkers); + + // Case: Enabled workers but all offline + if (activeWorkers.length === 0 && enabledWorkers.length > 0) { + extension.log("No active workers found. All enabled workers are offline."); + if (extension.ui?.showToast) { + extension.ui.showToast(extension.app, "error", "All Workers Offline", + `${enabledWorkers.length} worker(s) enabled but all are offline or unreachable. Check worker connections and try again.`, 5000); + } + // Fall back to master-only execution + return extension.originalQueuePrompt(0, promptWrapper); + } + + extension.log(`Pre-flight check: ${activeWorkers.length} of ${enabledWorkers.length} workers are active`, "debug"); + + // Check if master host might be unreachable by workers (cloudflare tunnel down) + const masterHost = extension.config?.master?.host || ''; + const isCloudflareHost = /\.(trycloudflare\.com|cloudflare\.dev)$/i.test(masterHost); + + if (isCloudflareHost && activeWorkers.length > 0) { + // Try to verify if the cloudflare tunnel is actually up + try { + const testUrl = `${window.location.protocol}//${masterHost}/prompt`; + const response = await fetch(testUrl, { + method: 'GET', + mode: 'cors', + cache: 'no-cache', + signal: AbortSignal.timeout(3000) // 3 second timeout + }); + + if (!response.ok) { + throw new Error('Master not reachable'); + } + } catch (error) { + // Cloudflare tunnel appears to be down + extension.log(`Master host ${masterHost} is not reachable - cloudflare tunnel may be down`, "error"); + + if (extension.ui?.showCloudflareWarning) { + extension.ui.showCloudflareWarning(extension, masterHost); + } + + // Stop execution - workers won't be able to send results back + extension.log("Blocking execution - workers cannot reach master at cloudflare domain", "error"); + return null; // This will prevent the workflow from running + } + } + + // Find all distributed nodes in the workflow + const collectorNodes = findNodesByClass(promptWrapper.output, "DistributedCollector"); + const upscaleNodes = findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed"); + const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; + + // Map original node IDs to truly unique job IDs for this specific run + const job_id_map = new Map(allDistributedNodes.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 => prepareDistributedJob(extension, uniqueId)); + await Promise.all(preparePromises); + + const jobs = []; + // Use only active workers + 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 = await prepareApiPromptForParticipant( + extension, 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) { + 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); + } + } + } + + const result = await executeJobs(extension, jobs); + return result; + } catch (error) { + extension.log("Parallel execution failed: " + error.message, "error"); + throw error; + } +} + +export async function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) { + let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt)); + const isMaster = participantId === 'master'; + + // Find all distributed nodes once (before pruning) + const collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector"); + const upscaleNodes = findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed"); + const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; + + // For workers, handle platform-specific path conversion + if (!isMaster) { + const workerInfo = extension.config.workers.find(w => w.id === participantId); + + if (workerInfo && workerInfo.host) { + // Remote or cloud worker - needs path translation + try { + const workerUrl = extension.getWorkerUrl(workerInfo); + const systemInfo = await getCachedWorkerSystemInfo(workerUrl); + const targetSeparator = systemInfo?.platform?.path_separator; + + 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 + } + } + + // Prune the workflow to only include distributed node dependencies + if (allDistributedNodes.length > 0) { + jobApiPrompt = pruneWorkflowForWorker(extension, 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 = extension.config.workers.find(w => w.id === workerId); + const isRemote = workerInfo && workerInfo.host; + + if (isRemote) { + // Find all image/video references in the pruned workflow + const imageReferences = findImageReferences(extension, jobApiPrompt); + if (imageReferences.size > 0) { + extension.log(`Found ${imageReferences.size} media references (images/videos) for remote worker ${workerId}`, "debug"); + // Store image references for later processing + options._imageReferences = imageReferences; + } + } + } + + // Handle Distributed seed nodes + const distributorNodes = findNodesByClass(jobApiPrompt, "DistributedSeed"); + if (distributorNodes.length > 0) { + extension.log(`Found ${distributorNodes.length} seed node(s)`, "debug"); + } + + for (const seedNode of distributorNodes) { + const { inputs } = jobApiPrompt[seedNode.id]; + inputs.is_worker = !isMaster; + if (!isMaster) { + const workerIndex = options.enabled_worker_ids.indexOf(participantId); + inputs.worker_id = `worker_${workerIndex}`; + extension.log(`Set seed node ${seedNode.id} for worker ${workerIndex}`, "debug"); + } + } + + // Handle Distributed collector nodes (already found above) + for (const collector of collectorNodes) { + const { inputs } = jobApiPrompt[collector.id]; + + // Check if this collector is downstream from a distributed upscaler + const hasUpstreamDistributedUpscaler = hasUpstreamNode( + jobApiPrompt, + collector.id, + 'UltimateSDUpscaleDistributed' + ); + + if (hasUpstreamDistributedUpscaler) { + // Set pass_through mode for this collector + inputs.pass_through = true; + extension.log(`Collector ${collector.id} set to pass-through mode (downstream from distributed upscaler)`, "debug"); + } else { + // Normal collector behavior + // 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 = extension.getMasterUrl(); + // Also make the worker_job_id unique to prevent potential caching issues + inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`; + inputs.worker_id = participantId; + } + } + } + + // Handle Ultimate SD Upscale Distributed nodes + for (const upscaleNode of upscaleNodes) { + const { inputs } = jobApiPrompt[upscaleNode.id]; + + // Get the unique job ID from the map + const uniqueJobId = options.job_id_map ? options.job_id_map.get(upscaleNode.id) : upscaleNode.id; + + 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 = extension.getMasterUrl(); + inputs.worker_id = participantId; + // Workers also need the enabled_worker_ids to calculate tile distribution + inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); + } + } + + return jobApiPrompt; +} + +export async function prepareDistributedJob(extension, multi_job_id) { + try { + await extension.api.prepareJob(multi_job_id); + } catch (error) { + extension.log("Error preparing job: " + error.message, "error"); + throw error; + } +} + +export async function executeJobs(extension, 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) { + extension.log(`Pre-loading ${allImageReferences.size} unique media file(s) for all workers`, "debug"); + await loadImagesForWorker(extension, allImageReferences); + } + + // Now dispatch jobs in parallel + const promises = jobs.map(job => { + if (job.type === 'master') { + return extension.originalQueuePrompt(0, job.promptWrapper).then(result => { + masterPromptId = result; + return result; + }); + } else { + return dispatchToWorker(extension, job.worker, job.prompt, job.workflow, job.imageReferences); + } + }); + await Promise.all(promises); + + // Trigger immediate status check for instant feedback + extension.checkAllWorkerStatuses(); + + return masterPromptId || { "prompt_id": "distributed-job-dispatched" }; +} + +async function dispatchToWorker(extension, worker, prompt, workflow, imageReferences) { + const workerUrl = extension.getWorkerUrl(worker); + + // Debug logging - always log to console for debugging + extension.log(`[Distributed] === Dispatching to ${worker.name} (${worker.id}) ===`, "debug"); + extension.log('[Distributed] Worker URL: ' + workerUrl, "debug"); + + // Handle image uploads for remote workers + if (imageReferences && imageReferences.size > 0) { + // Check if this is a local worker (same host as master) + const isLocalWorker = workerUrl.includes('127.0.0.1') || workerUrl.includes('localhost'); + + 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"); + + 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 + } + } + } + + const promptToSend = { + prompt, + extra_data: { extra_pnginfo: { workflow } }, + client_id: api.clientId + }; + + extension.log('[Distributed] Prompt data: ' + JSON.stringify(promptToSend), "debug"); + + try { + await fetch(`${workerUrl}/prompt`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + mode: 'cors', + body: JSON.stringify(promptToSend) + }); + } catch (e) { + extension.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`, "error"); + } +} + +export async function loadImagesForWorker(extension, imageReferences) { + const images = []; + + // Use a cache to avoid loading the same image multiple times + if (!extension._imageCache) { + extension._imageCache = new Map(); + } + + for (const [filename, info] of imageReferences) { + try { + // Check cache first + if (extension._imageCache.has(filename)) { + images.push(extension._imageCache.get(filename)); + extension.log(`Using cached image: ${filename}`, "debug"); + continue; + } + + // Limit cache size + if (extension._imageCache.size >= 10) { + const oldestKey = extension._imageCache.keys().next().value; + extension._imageCache.delete(oldestKey); + extension.log(`Evicted oldest cache entry: ${oldestKey} (cache limit reached)`, "debug"); + } + + // Load image from master's filesystem via API + try { + const data = await extension.api.loadImage(filename); + const imageData = { + name: filename, + image: data.image_data, + hash: data.hash // Include hash from the response + }; + images.push(imageData); + + // Cache the image for future use + extension._imageCache.set(filename, imageData); + extension.log(`Loaded and cached image: ${filename}`, "debug"); + } catch (loadError) { + extension.log(`Failed to load image ${filename}: ${loadError.message}`, "error"); + throw loadError; + } + } catch (error) { + extension.log(`Error loading image ${filename}: ${error.message}`, "error"); + } + } + + // Clear cache after a reasonable time to avoid memory issues + setTimeout(() => { + if (extension._imageCache && extension._imageCache.size > 0) { + extension.log(`Clearing image cache (${extension._imageCache.size} images)`, "debug"); + extension._imageCache.clear(); + } + }, TIMEOUTS.IMAGE_CACHE_CLEAR); // Clear after 30 seconds + + return images; +} + +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) + 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|video)\/\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: mimeType }); + + // Use original filename without heavy cleaning + let cleanName = imageData.name; + let subfolder = ''; + + // 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}`); + } + + extension.log(`Uploaded image to worker: ${imageData.name} -> ${subfolder}/${cleanName}`, "debug"); + } catch (error) { + extension.log(`Failed to upload ${imageData.name}: ${error.message}`, "error"); + // Continue with other images + } + } +} + +export async function performPreflightCheck(extension, workers) { + if (workers.length === 0) return []; + + extension.log(`Performing pre-flight health check on ${workers.length} workers...`, "debug"); + const startTime = Date.now(); + + const checkPromises = workers.map(async (worker) => { + const url = extension.getWorkerUrl(worker, '/prompt'); + + extension.log(`Pre-flight checking ${worker.name} at: ${url}`, "debug"); + + try { + const response = await fetch(url, { + method: 'GET', + mode: 'cors', + signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK) + }); + + if (response.ok) { + extension.log(`Worker ${worker.name} is active`, "debug"); + return { worker, active: true }; + } else { + extension.log(`Worker ${worker.name} returned ${response.status}`, "debug"); + return { worker, active: false }; + } + } catch (error) { + extension.log(`Worker ${worker.name} is offline or unreachable: ${error.message}`, "debug"); + return { worker, active: false }; + } + }); + + const results = await Promise.all(checkPromises); + const activeWorkers = results.filter(r => r.active).map(r => r.worker); + + const elapsed = Date.now() - startTime; + extension.log(`Pre-flight check completed in ${elapsed}ms. Active workers: ${activeWorkers.length}/${workers.length}`, "debug"); + + // Update UI status indicators for inactive workers + results.filter(r => !r.active).forEach(r => { + const statusDot = document.getElementById(`status-${r.worker.id}`); + if (statusDot) { + // Remove pulsing animation once status is determined + statusDot.classList.remove('status-pulsing'); + statusDot.style.backgroundColor = "#c04c4c"; // Red for offline + statusDot.title = "Offline - Cannot connect"; + } + }); + + return activeWorkers; +} \ No newline at end of file diff --git a/web/gpupanel.js b/web/gpupanel.js deleted file mode 100644 index a98fc3e..0000000 --- a/web/gpupanel.js +++ /dev/null @@ -1,2728 +0,0 @@ -import { app } from "../../scripts/app.js"; -import { api } from "../../scripts/api.js"; - -class DistributedExtension { - // 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;", - }; - - // CSS for pulsing animation and button hover effects - static PULSE_ANIMATION_CSS = ` - @keyframes pulse { - 0% { - opacity: 1; - transform: scale(1); - box-shadow: 0 0 0 0 rgba(240, 173, 78, 0.7); - } - 50% { - opacity: 0.3; - transform: scale(1.3); - box-shadow: 0 0 0 6px rgba(240, 173, 78, 0); - } - 100% { - opacity: 1; - transform: scale(1); - box-shadow: 0 0 0 0 rgba(240, 173, 78, 0); - } - } - .status-pulsing { - animation: pulse 1.2s ease-in-out infinite; - } - - /* Button hover effects */ - .distributed-button:hover:not(:disabled) { - filter: brightness(1.2); - transition: filter 0.2s ease; - } - .distributed-button:disabled { - opacity: 0.6; - cursor: not-allowed; - } - - /* Settings button animation */ - .settings-btn { - transition: transform 0.2s ease; - } - - - /* Expanded settings panel */ - .worker-settings { - max-height: 0; - overflow: hidden; - opacity: 0; - transition: max-height 0.3s ease, opacity 0.3s ease, padding 0.3s ease, margin 0.3s ease; - } - .worker-settings.expanded { - max-height: 300px; - opacity: 1; - } - `; - - constructor() { - this.config = null; - this.originalQueuePrompt = api.queuePrompt.bind(api); - this.workerStatuses = new Map(); // Track worker statuses - this.managedWorkers = new Map(); // Track UI-managed workers - this.statusCheckInterval = null; - this.logAutoRefreshInterval = null; // Track log auto-refresh - this.launchingWorkers = new Set(); // Track workers currently launching - this.expandedWorkers = new Set(); // Track which workers have expanded settings - - // Inject CSS for pulsing animation - this.injectStyles(); - - this.loadConfig().then(() => { - this.registerSidebarTab(); - this.setupInterceptor(); - this.startStatusChecking(); - this.loadManagedWorkers(); - this.detectMasterIP(); // Auto-detect master IP on startup - }); - } - - // Debug logging helpers - debugLog(message) { - if (this.config?.settings?.debug) { - console.log(`[Distributed] ${message}`); - } - } - - log(message) { - console.log(`[Distributed] ${message}`); - } - - injectStyles() { - const styleId = 'distributed-styles'; - if (!document.getElementById(styleId)) { - const style = document.createElement('style'); - style.id = styleId; - style.textContent = DistributedExtension.PULSE_ANIMATION_CSS; - document.head.appendChild(style); - } - } - - // --- 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}/distributed/config`); - if (response.ok) { - this.config = await response.json(); - this.debugLog("Loaded config: " + JSON.stringify(this.config)); - } else { - this.log("Failed to load config"); - this.config = { workers: [], settings: {} }; - } - } catch (error) { - this.log("Error loading config: " + error.message); - 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}/distributed/config/update_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ worker_id: workerId, enabled: enabled }) - }); - } catch (error) { - this.log("Error updating worker: " + error.message); - } - } - - async _updateSetting(key, value) { - // Update local config - if (!this.config.settings) { - this.config.settings = {}; - } - this.config.settings[key] = value; - - try { - await fetch(`${window.location.origin}/distributed/config/update_setting`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ key, value }) - }); - - app.extensionManager.toast.add({ - severity: "success", - summary: "Setting Updated", - detail: `${key.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase())} ${value ? 'enabled' : 'disabled'}`, - life: 2000 - }); - } catch (error) { - this.log(`Error updating setting '${key}': ${error.message}`); - app.extensionManager.toast.add({ - severity: "error", - summary: "Setting Update Failed", - detail: error.message, - life: 3000 - }); - } - } - - // --- UI Rendering --- - - registerSidebarTab() { - app.extensionManager.registerSidebarTab({ - id: "distributed", - icon: "pi pi-server", - title: "Distributed", - tooltip: "Distributed Control Panel", - type: "custom", - render: (el) => this.renderSidebarContent(el) - }); - } - - _createButton(text, onClick, style) { - const button = document.createElement("button"); - button.textContent = text; - button.style.cssText = DistributedExtension.BUTTON_STYLES.base + style; - button.onclick = onClick; - button.className = "distributed-button"; - return button; - } - - async _handleInterruptWorkers(button) { - const originalText = button.textContent; - const originalStyle = button.style.cssText; - button.textContent = "Interrupting..."; - button.disabled = true; - button.style.backgroundColor = "#c04c4c"; - - try { - const targetsToInterrupt = this.enabledWorkers.map(w => ({ - name: w.name, - url: `http://${w.host || window.location.hostname}:${w.port}` - })); - - if (targetsToInterrupt.length === 0) { - button.textContent = "No Workers"; - button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; - }, 3000); - return; - } - - const promises = targetsToInterrupt.map(target => - fetch(`${target.url}/interrupt`, { - 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 = "Interrupted!"; - button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.success.split(':')[1].trim().replace(';', ''); - // Update status after interrupt - setTimeout(() => this.checkAllWorkerStatuses(), 500); - } else { - button.textContent = "Error! See Console"; - button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); - this.log("Failed to interrupt: " + failures.map(f => f.name).join(", ")); - } - - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; - }, 3000); - } finally { - button.disabled = false; - } - } - - async _handleClearMemory(button) { - const originalText = button.textContent; - const originalStyle = button.style.cssText; // Store complete original styling - - 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"; - // Preserve original cssText but update background color - button.style.cssText = originalStyle.replace(/background-color:[^;]*;?/, '') + ' background-color: #c04c4c;'; - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; // Restore complete original styling - button.disabled = false; - }, 3000); - return; - } - - const promises = urlsToClear.map(target => - fetch(`${target.url}/distributed/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!"; - // Preserve original cssText but update background color - button.style.cssText = originalStyle.replace(/background-color:[^;]*;?/, '') + ' background-color: #3ca03c;'; - } else { - button.textContent = "Error! See Console"; - // Preserve original cssText but update background color - button.style.cssText = originalStyle.replace(/background-color:[^;]*;?/, '') + ' background-color: #c04c4c;'; - this.log("Failed to clear memory on: " + failures.map(f => f.name).join(", ")); - } - - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; // Restore complete original styling - button.disabled = false; - }, 3000); - } catch (error) { - // Handle any unexpected errors - button.textContent = "Error! See Console"; - button.style.cssText = originalStyle.replace(/background-color:[^;]*;?/, '') + ' background-color: #c04c4c;'; - this.log("Unexpected error in clear memory: " + error.message); - - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; - button.disabled = false; - }, 3000); - } - } - - - renderSidebarContent(el) { - // Panel is being opened/rendered - this.debugLog("Panel opened"); - - if (!el) { - this.debugLog("No element provided to renderSidebarContent"); - return; - } - - // Store reference to the panel element - this.panelElement = el; - - // Reload managed workers when panel opens - this.loadManagedWorkers().then(() => { - // Update all worker controls after loading - if (this.config?.workers) { - this.config.workers.forEach(w => this.updateWorkerControls(w.id)); - } - }); - - // Start checking worker statuses immediately in parallel - // Results will update the UI as they come in - setTimeout(() => this.checkAllWorkerStatuses(), 0); - - 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 = "ComfyUI Distributed"; - 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"); - const workers = this.config?.workers || []; - - // If no workers exist, show a full blueprint placeholder first - if (workers.length === 0) { - const blueprintDiv = document.createElement("div"); - blueprintDiv.style.cssText = "margin-bottom: 12px; border: 2px solid #666; border-radius: 4px; overflow: hidden; display: flex; cursor: pointer; transition: all 0.2s ease; background: rgba(255, 255, 255, 0.02);"; - blueprintDiv.title = "Click to add your first worker"; - - // Left column - plus icon - const blueprintCheckbox = document.createElement("div"); - blueprintCheckbox.style.cssText = "flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: 2px solid #666; color: #888; font-size: 24px; font-weight: bold;"; - blueprintCheckbox.innerHTML = "+"; - - // Right column - const blueprintContent = document.createElement("div"); - blueprintContent.style.cssText = "flex: 1; display: flex; flex-direction: column;"; - - // Info row with placeholder content - const blueprintInfo = document.createElement("div"); - blueprintInfo.style.cssText = "display: flex; align-items: center; padding: 8px; min-height: 60px;"; - - const blueprintWorkerContent = document.createElement("div"); - blueprintWorkerContent.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; - - const blueprintDot = document.createElement("span"); - blueprintDot.style.cssText = "display: inline-block; width: 8px; height: 8px; border-radius: 50%; border: 1px solid #555;"; - - const blueprintText = document.createElement("span"); - blueprintText.innerHTML = `Add New Worker
[CUDA] • [Port]`; - - blueprintWorkerContent.appendChild(blueprintDot); - blueprintWorkerContent.appendChild(blueprintText); - blueprintInfo.appendChild(blueprintWorkerContent); - - // Placeholder button row - single Launch button spanning width - const blueprintButtons = document.createElement("div"); - blueprintButtons.style.cssText = "padding: 0 8px 6px 8px;"; - - const ghostLaunch = document.createElement("div"); - ghostLaunch.style.cssText = "width: 100%; padding: 3px 8px; font-size: 11px; border: 1px solid #444; background: transparent; color: #555; border-radius: 4px; text-align: center;"; - ghostLaunch.textContent = "Launch"; - - blueprintButtons.appendChild(ghostLaunch); - - blueprintContent.appendChild(blueprintInfo); - blueprintContent.appendChild(blueprintButtons); - - blueprintDiv.appendChild(blueprintCheckbox); - blueprintDiv.appendChild(blueprintContent); - - // Click handler - blueprintDiv.onclick = () => this.addNewWorker(); - - // Hover effect - blueprintDiv.onmouseover = () => { - blueprintDiv.style.borderColor = "#888"; - blueprintDiv.style.backgroundColor = "rgba(255, 255, 255, 0.05)"; - blueprintCheckbox.style.borderColor = "#888"; - blueprintCheckbox.style.color = "#aaa"; - }; - blueprintDiv.onmouseout = () => { - blueprintDiv.style.borderColor = "#666"; - blueprintDiv.style.backgroundColor = "rgba(255, 255, 255, 0.02)"; - blueprintCheckbox.style.borderColor = "#666"; - blueprintCheckbox.style.color = "#888"; - }; - - gpuList.appendChild(blueprintDiv); - } - - // Show existing workers - workers.forEach(worker => { - const gpuDiv = document.createElement("div"); - gpuDiv.style.cssText = "margin-bottom: 12px; background: #2a2a2a; border-radius: 4px; overflow: hidden; display: flex;"; - - // Left column - checkbox area (extends full height) - const checkboxColumn = document.createElement("div"); - checkboxColumn.style.cssText = "flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #444; cursor: pointer;"; - checkboxColumn.title = "Enable/disable this worker"; - - // Standard checkbox - const checkbox = document.createElement("input"); - checkbox.type = "checkbox"; - checkbox.id = `gpu-${worker.id}`; - checkbox.checked = worker.enabled; - checkbox.style.cssText = "cursor: pointer; width: 16px; height: 16px; pointer-events: none;"; // pointer-events: none so the column click handles it - - checkboxColumn.appendChild(checkbox); - - // Click handler for checkbox column - checkboxColumn.onclick = async (e) => { - checkbox.checked = !checkbox.checked; - await this.updateWorkerEnabled(worker.id, checkbox.checked); - this.updateSummary(); - }; - - // Right column - content area - const contentColumn = document.createElement("div"); - contentColumn.style.cssText = "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;"; - - // First row: worker info - const infoRow = document.createElement("div"); - infoRow.style.cssText = "display: flex; align-items: center; padding: 8px; cursor: pointer; min-height: 60px;"; - infoRow.title = "Click to expand settings"; - - // Worker content - const workerContent = document.createElement("div"); - workerContent.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; - - const statusDot = document.createElement("span"); - statusDot.id = `status-${worker.id}`; - statusDot.style.cssText = "display: inline-block; width: 8px; height: 8px; border-radius: 50%; background-color: #666;"; - statusDot.title = "Checking status..."; - - const gpuInfo = document.createElement("span"); - const isRemote = this.isRemoteWorker(worker); - - if (isRemote) { - // For remote workers: show host:port - gpuInfo.innerHTML = `${worker.name}
${worker.host}:${worker.port}`; - } else { - // For local workers: show CUDA and port - const cudaInfo = worker.cuda_device !== undefined ? `CUDA ${worker.cuda_device} • ` : ''; - gpuInfo.innerHTML = `${worker.name}
${cudaInfo}Port ${worker.port}`; - } - - workerContent.appendChild(statusDot); - workerContent.appendChild(gpuInfo); - - // Settings arrow - const settingsArrow = document.createElement("span"); - settingsArrow.className = "settings-arrow"; - settingsArrow.innerHTML = "▶"; - settingsArrow.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto;"; - if (this.expandedWorkers.has(worker.id)) { - settingsArrow.style.transform = "rotate(90deg)"; - } - - infoRow.appendChild(workerContent); - infoRow.appendChild(settingsArrow); - - - // Add content and arrow back to info row - infoRow.appendChild(workerContent); - infoRow.appendChild(settingsArrow); - - // Click handler for info row - infoRow.onclick = () => { - this.toggleWorkerExpanded(worker.id); - }; - - // Hover effect for entire content column - contentColumn.onmouseover = () => { - contentColumn.style.backgroundColor = "#333"; - settingsArrow.style.color = "#fff"; - }; - contentColumn.onmouseout = () => { - contentColumn.style.backgroundColor = "transparent"; - settingsArrow.style.color = "#888"; - }; - - contentColumn.appendChild(infoRow); - - // Second row: launch controls - const controlsDiv = document.createElement("div"); - controlsDiv.id = `controls-${worker.id}`; - controlsDiv.style.cssText = "padding: 0 8px 6px 8px; display: flex; gap: 4px;"; - - if (this.isRemoteWorker(worker)) { - // For remote workers, show styled info box - const remoteInfo = document.createElement("div"); - remoteInfo.style.cssText = "background-color: #333; color: #888; padding: 4px 12px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1;"; - remoteInfo.textContent = "Remote worker"; - controlsDiv.appendChild(remoteInfo); - } else { - // Local workers get launch/stop/log buttons - const launchBtn = this._createButton('Launch', - () => this.launchWorker(worker.id), - 'background-color: #555;'); // Neutral color initially - launchBtn.style.cssText += ' padding: 3px 8px; font-size: 11px;'; - launchBtn.title = "Launch worker (runs in background with logging)"; - - const stopBtn = this._createButton('Stop', - () => this.stopWorker(worker.id), - 'background-color: #555;'); // Neutral color initially - stopBtn.style.cssText += ' padding: 3px 8px; font-size: 11px;'; - stopBtn.title = "Stop worker"; - - const logBtn = this._createButton('View Log', - () => this.viewWorkerLog(worker.id), - 'background-color: #685434;'); // ComfyUI darker yellow - logBtn.style.cssText += ' padding: 3px 8px; font-size: 11px; display: none;'; - logBtn.id = `log-${worker.id}`; - logBtn.title = "View worker log file"; - - controlsDiv.appendChild(launchBtn); - controlsDiv.appendChild(stopBtn); - controlsDiv.appendChild(logBtn); - } - - // Third row: expandable settings panel - const settingsDiv = document.createElement("div"); - settingsDiv.id = `settings-${worker.id}`; - settingsDiv.className = "worker-settings"; - if (this.expandedWorkers.has(worker.id)) { - settingsDiv.classList.add("expanded"); - } - settingsDiv.style.cssText = "margin: 0 8px; padding: 0 12px; background: #1e1e1e; border-radius: 4px;"; - if (this.expandedWorkers.has(worker.id)) { - settingsDiv.style.padding = "12px"; - settingsDiv.style.marginTop = "8px"; - settingsDiv.style.marginBottom = "8px"; - } - - // Create settings form - const settingsForm = this.createWorkerSettingsForm(worker); - settingsDiv.appendChild(settingsForm); - - contentColumn.appendChild(controlsDiv); - contentColumn.appendChild(settingsDiv); - - gpuDiv.appendChild(checkboxColumn); - gpuDiv.appendChild(contentColumn); - gpuList.appendChild(gpuDiv); - - // Update controls based on managed worker status - this.updateWorkerControls(worker.id); - }); - gpuSection.appendChild(gpuList); - - // Only show the minimal "Add Worker" box if there are existing workers - if (workers.length > 0) { - // Add Worker placeholder box - const addWorkerDiv = document.createElement("div"); - addWorkerDiv.style.cssText = "margin-top: 10px; margin-bottom: 12px; border: 1px solid #444; border-radius: 4px; overflow: hidden; display: flex; cursor: pointer; transition: all 0.2s ease;"; - addWorkerDiv.title = "Click to add a new worker"; - - // Left column - plus icon area - const plusColumn = document.createElement("div"); - plusColumn.style.cssText = "flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #444; color: #555; font-size: 18px;"; - plusColumn.innerHTML = "+"; - - // Right column - content area - const addContentColumn = document.createElement("div"); - addContentColumn.style.cssText = "flex: 1; display: flex; flex-direction: column;"; - - // Info row - const addInfoRow = document.createElement("div"); - addInfoRow.style.cssText = "display: flex; align-items: center; padding: 8px; min-height: 42px;"; // Reduced from 60px to 42px (30% reduction) - - // Placeholder content - const addContent = document.createElement("div"); - addContent.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; - - // Placeholder dot - const placeholderDot = document.createElement("span"); - placeholderDot.style.cssText = "display: inline-block; width: 8px; height: 8px; border-radius: 50%; border: 1px solid #555;"; - - // Placeholder text - const placeholderInfo = document.createElement("span"); - placeholderInfo.innerHTML = `Add New Worker`; - - addContent.appendChild(placeholderDot); - addContent.appendChild(placeholderInfo); - addInfoRow.appendChild(addContent); - addContentColumn.appendChild(addInfoRow); - - addWorkerDiv.appendChild(plusColumn); - addWorkerDiv.appendChild(addContentColumn); - - // Hover effects - addWorkerDiv.onmouseover = () => { - addWorkerDiv.style.borderColor = "#666"; - addWorkerDiv.style.backgroundColor = "rgba(255, 255, 255, 0.02)"; - plusColumn.style.color = "#888"; - plusColumn.style.borderColor = "#666"; - }; - - addWorkerDiv.onmouseout = () => { - addWorkerDiv.style.borderColor = "#444"; - addWorkerDiv.style.backgroundColor = "transparent"; - plusColumn.style.color = "#555"; - plusColumn.style.borderColor = "#444"; - }; - - // Click handler - addWorkerDiv.onclick = () => this.addNewWorker(); - - gpuSection.appendChild(addWorkerDiv); - } - - container.appendChild(gpuSection); - - const actionsSection = document.createElement("div"); - actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;"; - - // Create a row for both buttons - const buttonRow = document.createElement("div"); - buttonRow.style.cssText = "display: flex; gap: 8px;"; - - const clearMemButton = this._createButton("Clear Worker VRAM", (e) => this._handleClearMemory(e.target), DistributedExtension.BUTTON_STYLES.clearMemory); - clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)"; - clearMemButton.style.cssText += " flex: 1;"; - - const interruptButton = this._createButton("Interrupt Workers", (e) => this._handleInterruptWorkers(e.target), DistributedExtension.BUTTON_STYLES.clearMemory); - interruptButton.title = "Cancel/interrupt execution on all enabled worker GPUs"; - interruptButton.style.cssText += " flex: 1;"; - - buttonRow.appendChild(clearMemButton); - buttonRow.appendChild(interruptButton); - actionsSection.appendChild(buttonRow); - - container.appendChild(actionsSection); - - // Settings section - const settingsSection = document.createElement("div"); - settingsSection.style.cssText = "border-top: 1px solid #444; padding-top: 10px; margin-bottom: 10px;"; - - // Settings header with toggle - const settingsHeader = document.createElement("div"); - settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between; cursor: pointer; user-select: none;"; - - const settingsTitle = document.createElement("h4"); - settingsTitle.textContent = "Settings"; - settingsTitle.style.cssText = "margin: 0; font-size: 14px;"; - - const settingsToggle = document.createElement("span"); - settingsToggle.textContent = "▶"; // Right arrow when collapsed - settingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; - - settingsHeader.appendChild(settingsTitle); - settingsHeader.appendChild(settingsToggle); - - // Hover effect for header - settingsHeader.onmouseover = () => { - settingsToggle.style.color = "#fff"; - }; - settingsHeader.onmouseout = () => { - settingsToggle.style.color = "#888"; - }; - - // Collapsible settings content - const settingsContent = document.createElement("div"); - settingsContent.style.cssText = "max-height: 0; overflow: hidden; transition: max-height 0.3s ease;"; - - const settingsDiv = document.createElement("div"); - settingsDiv.style.cssText = "display: flex; flex-direction: column; gap: 8px; padding-top: 10px;"; - - // Toggle functionality - let settingsExpanded = false; - settingsHeader.onclick = () => { - settingsExpanded = !settingsExpanded; - if (settingsExpanded) { - settingsContent.style.maxHeight = "200px"; - settingsToggle.style.transform = "rotate(90deg)"; - } else { - settingsContent.style.maxHeight = "0"; - settingsToggle.style.transform = "rotate(0deg)"; - } - }; - - // Debug mode setting - const debugGroup = document.createElement("div"); - debugGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; - - const debugCheckbox = document.createElement("input"); - debugCheckbox.type = "checkbox"; - debugCheckbox.id = "setting-debug"; - debugCheckbox.checked = this.config?.settings?.debug || false; - debugCheckbox.onchange = (e) => this._updateSetting('debug', e.target.checked); - - const debugLabel = document.createElement("label"); - debugLabel.htmlFor = "setting-debug"; - debugLabel.textContent = "Debug Mode"; - debugLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; - - debugGroup.appendChild(debugCheckbox); - debugGroup.appendChild(debugLabel); - - // Auto-launch workers setting - const autoLaunchGroup = document.createElement("div"); - autoLaunchGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; - - const autoLaunchCheckbox = document.createElement("input"); - autoLaunchCheckbox.type = "checkbox"; - autoLaunchCheckbox.id = "setting-auto-launch"; - autoLaunchCheckbox.checked = this.config?.settings?.auto_launch_workers || false; - autoLaunchCheckbox.onchange = (e) => this._updateSetting('auto_launch_workers', e.target.checked); - - const autoLaunchLabel = document.createElement("label"); - autoLaunchLabel.htmlFor = "setting-auto-launch"; - autoLaunchLabel.textContent = "Auto-launch Local Workers on Startup"; - autoLaunchLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; - - autoLaunchGroup.appendChild(autoLaunchCheckbox); - autoLaunchGroup.appendChild(autoLaunchLabel); - - // Stop workers on exit setting - const stopOnExitGroup = document.createElement("div"); - stopOnExitGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; - - const stopOnExitCheckbox = document.createElement("input"); - stopOnExitCheckbox.type = "checkbox"; - stopOnExitCheckbox.id = "setting-stop-on-exit"; - stopOnExitCheckbox.checked = this.config?.settings?.stop_workers_on_master_exit !== false; // Default true - stopOnExitCheckbox.onchange = (e) => this._updateSetting('stop_workers_on_master_exit', e.target.checked); - - const stopOnExitLabel = document.createElement("label"); - stopOnExitLabel.htmlFor = "setting-stop-on-exit"; - stopOnExitLabel.textContent = "Stop Local Workers on Master Exit"; - stopOnExitLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; - - stopOnExitGroup.appendChild(stopOnExitCheckbox); - stopOnExitGroup.appendChild(stopOnExitLabel); - - settingsDiv.appendChild(debugGroup); - settingsDiv.appendChild(autoLaunchGroup); - settingsDiv.appendChild(stopOnExitGroup); - settingsContent.appendChild(settingsDiv); - - settingsSection.appendChild(settingsHeader); - settingsSection.appendChild(settingsContent); - container.appendChild(settingsSection); - - const summarySection = document.createElement("div"); - summarySection.style.cssText = "border-top: 1px solid #444; padding-top: 10px;"; - const summary = document.createElement("div"); - summary.id = "distributed-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('distributed-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) { - const hasCollector = this.findNodesByClass(prompt.output, "DistributedCollector").length > 0; - const hasDistUpscale = this.findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0; - - if (hasCollector || hasDistUpscale) { - const result = await this.executeParallelDistributed(prompt); - // Check status after dispatching jobs - setTimeout(() => this.checkAllWorkerStatuses(), 500); - return result; - } - } - return this.originalQueuePrompt(number, prompt); - }; - } - - async executeParallelDistributed(promptWrapper) { - try { - const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution - const enabledWorkers = this.enabledWorkers; - - // Pre-flight health check on all enabled workers - const activeWorkers = await this.performPreflightCheck(enabledWorkers); - - if (activeWorkers.length === 0 && enabledWorkers.length > 0) { - this.log("No active workers found. Running on master only."); - // Fall back to master-only execution - return this.originalQueuePrompt(0, promptWrapper); - } - - this.debugLog(`Pre-flight check: ${activeWorkers.length} of ${enabledWorkers.length} workers are active`); - - // Find all distributed nodes in the workflow - const collectorNodes = this.findNodesByClass(promptWrapper.output, "DistributedCollector"); - const upscaleNodes = this.findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed"); - const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; - - // Map original node IDs to truly unique job IDs for this specific run - const job_id_map = new Map(allDistributedNodes.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._prepareDistributedJob(uniqueId)); - await Promise.all(preparePromises); - - const jobs = []; - // Use only active workers - 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, options - ); - - if (participantId === 'master') { - jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } }); - } else { - const worker = activeWorkers.find(w => w.id === participantId); - 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); - } - } - } - - const result = await this._executeJobs(jobs); - return result; - } catch (error) { - this.log("Parallel execution failed: " + error.message); - throw error; - } - } - - - 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 => { - masterPromptId = result; - return result; - }); - } else { - return this._dispatchToWorker(job.worker, job.prompt, job.workflow, job.imageReferences); - } - }); - await Promise.all(promises); - return masterPromptId || { "prompt_id": "distributed-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 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) - * @param {Object} apiPrompt - The workflow API prompt - * @param {Array} collectorIds - Array of collector node IDs - * @returns {Set} 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; - } - - /** - * Find only upstream nodes (inputs) for distributed collector nodes - * This is used for workers to avoid executing downstream nodes like SaveImage - * @param {Object} apiPrompt - The API prompt containing the workflow - * @param {Array} collectorIds - Array of collector node IDs - * @returns {Set} Set of node IDs that feed into collectors - */ - findCollectorUpstreamNodes(apiPrompt, collectorIds) { - const connected = new Set(collectorIds); // Include all collectors - const toProcess = [...collectorIds]; - - // Only traverse upstream (inputs) - while (toProcess.length > 0) { - const nodeId = toProcess.pop(); - const node = apiPrompt[nodeId]; - - // Traverse upstream (inputs) only - 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); - } - } - } - } - } - - return connected; - } - - /** - * Prune workflow to only include nodes connected to distributed nodes - * @param {Object} apiPrompt - The full workflow API prompt - * @param {Array} distributedNodes - Array of distributed nodes (optional, will find if not provided) - * @returns {Object} Pruned API prompt with only required nodes - */ - pruneWorkflowForWorker(apiPrompt, distributedNodes = null) { - // Find all distributed nodes if not provided - if (!distributedNodes) { - const collectorNodes = this.findNodesByClass(apiPrompt, "DistributedCollector"); - const upscaleNodes = this.findNodesByClass(apiPrompt, "UltimateSDUpscaleDistributed"); - distributedNodes = [...collectorNodes, ...upscaleNodes]; - } - - if (distributedNodes.length === 0) { - // No distributed nodes, return full workflow - return apiPrompt; - } - - // Get all nodes connected to distributed nodes - const distributedIds = distributedNodes.map(node => node.id); - - // For workers, only include upstream nodes (this removes ALL downstream nodes after collectors) - const connectedNodes = this.findCollectorUpstreamNodes(apiPrompt, distributedIds); - - this.debugLog(`Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes (removed all downstream nodes)`); - - // Create pruned prompt with only required nodes - const prunedPrompt = {}; - for (const nodeId of connectedNodes) { - prunedPrompt[nodeId] = JSON.parse(JSON.stringify(apiPrompt[nodeId])); - } - - // Check if any distributed node has downstream SaveImage nodes that were removed - // If so, add a PreviewImage node after the collector - for (const distNode of distributedNodes) { - const distNodeId = distNode.id; - - // Check if this distributed node had any downstream nodes in the original workflow - const originalOutputMap = new Map(); - 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 && String(inputValue[0]) === distNodeId) { - if (!originalOutputMap.has(distNodeId)) { - originalOutputMap.set(distNodeId, []); - } - originalOutputMap.get(distNodeId).push({nodeId, inputName}); - } - } - } - } - - // If this distributed node had downstream nodes that were removed, add a PreviewImage - if (originalOutputMap.has(distNodeId) && originalOutputMap.get(distNodeId).length > 0) { - // Generate a unique ID for the preview node - const previewNodeId = `preview_${distNodeId}`; - - // Add PreviewImage node connected to the distributed node - prunedPrompt[previewNodeId] = { - inputs: { - images: [distNodeId, 0] // Connect to first output of distributed node - }, - class_type: "PreviewImage", - _meta: { - title: "Preview Image (auto-added)" - } - }; - - this.debugLog(`Added PreviewImage node ${previewNodeId} after distributed node ${distNodeId} for worker`); - } - } - - return prunedPrompt; - } - - - _prepareApiPromptForParticipant(baseApiPrompt, participantId, options = {}) { - let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt)); - const isMaster = participantId === 'master'; - - // Find all distributed nodes once (before pruning) - const collectorNodes = this.findNodesByClass(jobApiPrompt, "DistributedCollector"); - const upscaleNodes = this.findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed"); - const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; - - // For workers, prune the workflow to only include distributed node dependencies - if (!isMaster && allDistributedNodes.length > 0) { - 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) { - this.debugLog(`Found ${distributorNodes.length} seed node(s)`); - } - - for (const seedNode of distributorNodes) { - const { inputs } = jobApiPrompt[seedNode.id]; - inputs.is_worker = !isMaster; - if (!isMaster) { - const workerIndex = options.enabled_worker_ids.indexOf(participantId); - inputs.worker_id = `worker_${workerIndex}`; - this.debugLog(`Set seed node ${seedNode.id} for worker ${workerIndex}`); - } - } - - // Handle Distributed collector nodes (already found above) - for (const collector of collectorNodes) { - const { inputs } = jobApiPrompt[collector.id]; - - // Check if this collector is downstream from a distributed upscaler - const hasUpstreamDistributedUpscaler = this._hasUpstreamNode( - jobApiPrompt, - collector.id, - 'UltimateSDUpscaleDistributed' - ); - - if (hasUpstreamDistributedUpscaler) { - // Set pass_through mode for this collector - inputs.pass_through = true; - this.debugLog(`Collector ${collector.id} set to pass-through mode (downstream from distributed upscaler)`); - } else { - // Normal collector behavior - // 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 = this.getMasterUrl(); - // Also make the worker_job_id unique to prevent potential caching issues - inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`; - inputs.worker_id = participantId; - } - } - } - - // Handle Ultimate SD Upscale Distributed nodes - for (const upscaleNode of upscaleNodes) { - const { inputs } = jobApiPrompt[upscaleNode.id]; - - // Get the unique job ID from the map - const uniqueJobId = options.job_id_map ? options.job_id_map.get(upscaleNode.id) : upscaleNode.id; - - 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 = this.getMasterUrl(); - inputs.worker_id = participantId; - // Workers also need the enabled_worker_ids to calculate tile distribution - inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); - } - } - - return jobApiPrompt; - } - - async _prepareDistributedJob(multi_job_id) { - try { - await fetch(`${window.location.origin}/distributed/prepare_job`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ multi_job_id }) - }); - } catch (error) { - this.log("Error preparing job: " + error.message); - throw error; - } - } - - /** - * Check if a node has an upstream node of a specific type - * @param {Object} apiPrompt - The workflow API prompt - * @param {string} nodeId - The node to check - * @param {string} upstreamType - The class_type to look for upstream - * @returns {boolean} True if an upstream node of the specified type exists - */ - _hasUpstreamNode(apiPrompt, nodeId, upstreamType) { - const visited = new Set(); - const toProcess = [nodeId]; - - while (toProcess.length > 0) { - const currentId = toProcess.pop(); - if (visited.has(currentId)) continue; - visited.add(currentId); - - const node = apiPrompt[currentId]; - if (!node) continue; - - // Check inputs for upstream connections - if (node.inputs) { - for (const [inputName, inputValue] of Object.entries(node.inputs)) { - if (Array.isArray(inputValue) && inputValue.length === 2) { - const sourceNodeId = String(inputValue[0]); - const sourceNode = apiPrompt[sourceNodeId]; - - if (sourceNode && sourceNode.class_type === upstreamType) { - return true; - } - - if (!visited.has(sourceNodeId)) { - toProcess.push(sourceNodeId); - } - } - } - } - } - - return false; - } - - startStatusChecking() { - // Start checking every 2 seconds for more responsive updates - // Don't check immediately since panel might not be open yet - this.statusCheckInterval = setInterval(() => { - this.checkAllWorkerStatuses(); - }, 2000); - } - - async checkAllWorkerStatuses() { - if (!this.config || !this.config.workers) return; - - for (const worker of this.config.workers) { - // Only check status for enabled workers - if (worker.enabled) { - this.checkWorkerStatus(worker); - } - } - } - - async checkWorkerStatus(worker) { - const host = worker.host || window.location.hostname; - const url = `http://${host}:${worker.port}/prompt`; - const statusDot = document.getElementById(`status-${worker.id}`); - - // Status dot will be updated when response comes back - - try { - const response = await fetch(url, { - method: 'GET', - mode: 'cors', - signal: AbortSignal.timeout(300) // 300ms for faster response - }); - - if (response.ok) { - const data = await response.json(); - const queueRemaining = data.exec_info?.queue_remaining || 0; - const isProcessing = queueRemaining > 0; - - // Update status - this.workerStatuses.set(worker.id, { - online: true, - processing: isProcessing, - queueCount: queueRemaining - }); - - // Update status dot based on processing state - if (isProcessing) { - this.updateStatusDot(worker.id, "#f0ad4e", `Online - Processing (${queueRemaining} in queue)`, false); - } else { - this.updateStatusDot(worker.id, "#3ca03c", "Online - Idle", false); - } - - // Clear launching state since worker is now online - if (this.launchingWorkers.has(worker.id)) { - this.launchingWorkers.delete(worker.id); - - // Also clear launching flag in backend - this.clearLaunchingFlag(worker.id); - } - } else { - throw new Error(`HTTP ${response.status}`); - } - } catch (error) { - // Worker is offline or unreachable - this.workerStatuses.set(worker.id, { - online: false, - processing: false, - queueCount: 0 - }); - - // Check if worker is launching - if (this.launchingWorkers.has(worker.id)) { - this.updateStatusDot(worker.id, "#f0ad4e", "Launching...", true); - } else { - // Only update to red if not currently launching - this.updateStatusDot(worker.id, "#c04c4c", "Offline - Cannot connect", false); - } - } - - // Update control buttons based on new status - this.updateWorkerControls(worker.id); - } - - async _dispatchToWorker(worker, prompt, workflow, imageReferences) { - const host = worker.host || window.location.hostname; - 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: { }, // Removed workflow from extra_pnginfo to reduce payload size - client_id: api.clientId - }; - - console.log('[Distributed] Prompt data:', promptToSend); - - try { - await fetch(`${workerUrl}/prompt`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - mode: 'cors', - body: JSON.stringify(promptToSend) - }); - } catch (e) { - this.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`); - } - } - - 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 []; - - this.debugLog(`Performing pre-flight health check on ${workers.length} workers...`); - const startTime = Date.now(); - - const checkPromises = workers.map(async (worker) => { - const host = worker.host || window.location.hostname; - const url = `http://${host}:${worker.port}/prompt`; - try { - const response = await fetch(url, { - method: 'GET', - mode: 'cors', - signal: AbortSignal.timeout(300) // 300ms timeout for pre-flight - }); - - if (response.ok) { - this.debugLog(`Worker ${worker.name} is active`); - return { worker, active: true }; - } else { - this.debugLog(`Worker ${worker.name} returned ${response.status}`); - return { worker, active: false }; - } - } catch (error) { - this.debugLog(`Worker ${worker.name} is offline or unreachable`); - return { worker, active: false }; - } - }); - - const results = await Promise.all(checkPromises); - const activeWorkers = results.filter(r => r.active).map(r => r.worker); - - const elapsed = Date.now() - startTime; - this.debugLog(`Pre-flight check completed in ${elapsed}ms. Active workers: ${activeWorkers.length}/${workers.length}`); - - // Update UI status indicators for inactive workers - results.filter(r => !r.active).forEach(r => { - const statusDot = document.getElementById(`status-${r.worker.id}`); - if (statusDot) { - // Remove pulsing animation once status is determined - statusDot.classList.remove('status-pulsing'); - statusDot.style.backgroundColor = "#c04c4c"; // Red for offline - statusDot.title = "Offline - Cannot connect"; - } - }); - - return activeWorkers; - } - - async launchWorker(workerId) { - const worker = this.config.workers.find(w => w.id === workerId); - const launchBtn = document.querySelector(`#controls-${workerId} button`); - - this.updateStatusDot(workerId, "#f0ad4e", "Launching...", true); - this.launchingWorkers.add(workerId); - - // Allow 90 seconds for worker to launch (model loading can take time) - setTimeout(() => { - this.launchingWorkers.delete(workerId); - }, 90000); - - if (!launchBtn) return; - - try { - // Disable button immediately - launchBtn.disabled = true; - - const response = await fetch(`${window.location.origin}/distributed/launch_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ worker_id: workerId }) - }); - - const result = await response.json(); - if (response.ok) { - this.log(`Launched ${worker.name} (PID: ${result.pid})`); - if (result.log_file) { - this.debugLog(`Log file: ${result.log_file}`); - } - - this.managedWorkers.set(String(workerId), { - pid: result.pid, - log_file: result.log_file, - started_at: Date.now() - }); - - // Update controls immediately to hide launch button and show stop/log buttons - this.updateWorkerControls(workerId); - setTimeout(() => this.checkWorkerStatus(worker), 2000); - } else { - throw new Error(result.message || 'Failed to launch worker'); - } - - // Check if worker was already running - if (!response.ok && result && result.message && result.message.includes("already running")) { - if (result.log_file) { - const existing = this.managedWorkers.get(String(workerId)); - if (!existing) { - this.managedWorkers.set(String(workerId), { - log_file: result.log_file - }); - } else { - existing.log_file = result.log_file; - } - } - this.updateWorkerControls(workerId); - setTimeout(() => this.checkWorkerStatus(worker), 100); - } - } catch (error) { - this.log(`Error launching worker: ${error}`); - - // Re-enable button on error - if (launchBtn) { - launchBtn.disabled = false; - } - } - } - - async stopWorker(workerId) { - const worker = this.config.workers.find(w => w.id === workerId); - const stopBtn = document.querySelectorAll(`#controls-${workerId} button`)[1]; - - // Provide immediate feedback - if (stopBtn) { - stopBtn.disabled = true; - stopBtn.textContent = "Stopping..."; - stopBtn.style.backgroundColor = "#666"; - } - - try { - const response = await fetch(`${window.location.origin}/distributed/stop_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ worker_id: workerId }) - }); - - const result = await response.json(); - if (response.ok) { - this.log(`Stopped worker: ${result.message}`); - this.managedWorkers.delete(String(workerId)); - - // Flash success feedback - if (stopBtn) { - stopBtn.style.backgroundColor = DistributedExtension.BUTTON_STYLES.success; - stopBtn.textContent = "Stopped!"; - setTimeout(() => { - this.updateWorkerControls(workerId); - }, 1500); - } - - // Update status - setTimeout(() => this.checkWorkerStatus(worker), 500); - } else { - this.log(`Failed to stop worker: ${result.message}`); - - // Flash error feedback - if (stopBtn) { - stopBtn.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error; - stopBtn.textContent = result.message.includes("already stopped") ? "Not Running" : "Failed"; - setTimeout(() => { - this.updateWorkerControls(workerId); - }, 2000); - } - } - } catch (error) { - this.log(`Error stopping worker: ${error}`); - - // Flash error feedback - if (stopBtn) { - stopBtn.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error; - stopBtn.textContent = "Error"; - setTimeout(() => { - this.updateWorkerControls(workerId); - }, 2000); - } - } - } - - async clearLaunchingFlag(workerId) { - try { - await fetch(`${window.location.origin}/distributed/worker/clear_launching`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ worker_id: workerId }) - }); - this.debugLog(`Cleared launching flag for worker ${workerId}`); - } catch (error) { - this.debugLog(`Error clearing launching flag: ${error}`); - } - } - - // Helper Functions to reduce redundancy - - /** - * Generic async button action handler - */ - async handleAsyncButtonAction(button, action, successText, errorText, resetDelay = 3000) { - const originalText = button.textContent; - const originalStyle = button.style.cssText; - button.disabled = true; - - try { - await action(); - button.textContent = successText; - button.style.cssText = originalStyle; - button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.success; - return true; - } catch (error) { - button.textContent = errorText || `Error: ${error.message}`; - button.style.cssText = originalStyle; - button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error; - throw error; - } finally { - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; - button.disabled = false; - }, resetDelay); - } - } - - /** - * Show toast notification helper - */ - showToast(severity, summary, detail, life = 3000) { - if (app.extensionManager?.toast?.add) { - app.extensionManager.toast.add({ severity, summary, detail, life }); - } - } - - /** - * Update status dot helper - */ - updateStatusDot(workerId, color, title, pulsing = false) { - const statusDot = document.getElementById(`status-${workerId}`); - if (!statusDot) return; - - statusDot.style.backgroundColor = color; - statusDot.title = title; - statusDot.classList.toggle('status-pulsing', pulsing); - } - - /** - * Cleanup method to stop intervals and listeners - */ - cleanup() { - if (this.statusCheckInterval) { - clearInterval(this.statusCheckInterval); - this.statusCheckInterval = null; - } - - if (this.logAutoRefreshInterval) { - clearInterval(this.logAutoRefreshInterval); - this.logAutoRefreshInterval = null; - } - - this.debugLog("Cleaned up intervals"); - } - - async loadManagedWorkers() { - try { - const response = await fetch(`${window.location.origin}/distributed/managed_workers`); - if (response.ok) { - const result = await response.json(); - this.managedWorkers.clear(); - - // Check for launching workers - for (const [workerId, info] of Object.entries(result.managed_workers)) { - this.managedWorkers.set(workerId, info); - - // If worker is marked as launching, add to launchingWorkers set - if (info.launching) { - this.launchingWorkers.add(parseInt(workerId)); - this.debugLog(`Worker ${workerId} is in launching state`); - } - } - - // Update UI for all workers - if (this.config?.workers) { - this.config.workers.forEach(w => this.updateWorkerControls(w.id)); - } - } - } catch (error) { - this.debugLog(`Error loading managed workers: ${error}`); - } - } - - updateWorkerControls(workerId) { - const controlsDiv = document.getElementById(`controls-${workerId}`); - - if (!controlsDiv) return; - - const worker = this.config.workers.find(w => w.id === workerId); - if (!worker) return; - - // Skip button updates for remote workers - if (this.isRemoteWorker(worker)) { - return; - } - - // Ensure we check for string ID - const managedInfo = this.managedWorkers.get(String(workerId)); - const status = this.workerStatuses.get(workerId); - - // Update button states - const buttons = controlsDiv.querySelectorAll('button'); - const launchBtn = buttons[0]; - const stopBtn = buttons[1]; - const logBtn = document.getElementById(`log-${workerId}`); - - // Show log button immediately if we have log file info (even if worker is still starting) - if (managedInfo?.log_file && logBtn) { - logBtn.style.display = 'inline-block'; - } else if (logBtn && !managedInfo) { - logBtn.style.display = 'none'; - } - - if (status?.online || managedInfo) { - // Worker is running or we just launched it - launchBtn.style.display = 'none'; // Hide launch button when running - - if (managedInfo) { - // Only show stop button if we manage this worker - stopBtn.style.display = 'inline-block'; - stopBtn.disabled = false; - stopBtn.textContent = "Stop"; - stopBtn.style.backgroundColor = "#7c4a4a"; // Red when enabled - } else { - // Hide stop button for workers launched outside UI - stopBtn.style.display = 'none'; - } - } else { - // Worker is not running - launchBtn.style.display = 'inline-block'; // Show launch button - launchBtn.disabled = false; - launchBtn.textContent = "Launch"; - launchBtn.style.backgroundColor = "#4a7c4a"; // Green when enabled - - stopBtn.style.display = 'none'; // Hide stop button when not running - } - } - - async viewWorkerLog(workerId) { - const managedInfo = this.managedWorkers.get(String(workerId)); - if (!managedInfo?.log_file) return; - - const logBtn = document.getElementById(`log-${workerId}`); - - // Provide immediate feedback - if (logBtn) { - logBtn.disabled = true; - logBtn.textContent = "Loading..."; - logBtn.style.backgroundColor = "#666"; - } - - try { - // Fetch log content - const response = await fetch(`${window.location.origin}/distributed/worker_log/${workerId}?lines=1000`); - if (!response.ok) { - throw new Error(`Failed to fetch log: ${response.statusText}`); - } - - const data = await response.json(); - - // Create modal dialog - this.showLogModal(workerId, data); - - // Restore button - if (logBtn) { - logBtn.disabled = false; - logBtn.textContent = "View Log"; - logBtn.style.backgroundColor = "#685434"; // Keep the yellow color - } - - } catch (error) { - this.log('Error viewing log: ' + error.message); - app.extensionManager.toast.add({ - severity: "error", - summary: "Error", - detail: `Failed to load log: ${error.message}`, - life: 5000 - }); - - // Flash error and restore button - if (logBtn) { - logBtn.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error; - logBtn.textContent = "Error"; - setTimeout(() => { - logBtn.disabled = false; - logBtn.textContent = "View Log"; - logBtn.style.backgroundColor = "#685434"; // Keep the yellow color - }, 2000); - } - } - } - - showLogModal(workerId, logData) { - // Remove any existing modal - const existingModal = document.getElementById('distributed-log-modal'); - if (existingModal) { - existingModal.remove(); - } - - const worker = this.config.workers.find(w => w.id === workerId); - const workerName = worker?.name || `Worker ${workerId}`; - - // Create modal container - const modal = document.createElement('div'); - modal.id = 'distributed-log-modal'; - modal.style.cssText = ` - position: fixed; - top: 0; - left: 0; - width: 100%; - height: 100%; - background: rgba(0, 0, 0, 0.8); - display: flex; - align-items: center; - justify-content: center; - z-index: 10000; - `; - - // Create modal content - const content = document.createElement('div'); - content.style.cssText = ` - background: #1e1e1e; - border-radius: 8px; - width: 90%; - max-width: 1200px; - height: 80%; - display: flex; - flex-direction: column; - box-shadow: 0 4px 20px rgba(0, 0, 0, 0.5); - `; - - // Header - const header = document.createElement('div'); - header.style.cssText = ` - padding: 15px 20px; - border-bottom: 1px solid #444; - display: flex; - justify-content: space-between; - align-items: center; - `; - - const title = document.createElement('h3'); - title.textContent = `${workerName} - Log Viewer`; - title.style.cssText = 'margin: 0; color: #fff;'; - - const headerButtons = document.createElement('div'); - headerButtons.style.cssText = 'display: flex; gap: 20px; align-items: center;'; - - // Auto-refresh container - const refreshContainer = document.createElement('div'); - refreshContainer.style.cssText = 'display: flex; align-items: center; gap: 4px;'; - - // Auto-refresh checkbox - const refreshCheckbox = document.createElement('input'); - refreshCheckbox.type = 'checkbox'; - refreshCheckbox.id = 'log-auto-refresh'; - refreshCheckbox.checked = true; // Enabled by default - refreshCheckbox.style.cssText = 'cursor: pointer;'; - refreshCheckbox.onchange = (e) => { - if (e.target.checked) { - this.startLogAutoRefresh(workerId); - } else { - this.stopLogAutoRefresh(); - } - }; - - const refreshLabel = document.createElement('label'); - refreshLabel.htmlFor = 'log-auto-refresh'; - refreshLabel.style.cssText = 'font-size: 12px; color: #ccc; cursor: pointer; white-space: nowrap;'; - refreshLabel.textContent = 'Auto-refresh'; - - // Add checkbox and label to container - refreshContainer.appendChild(refreshCheckbox); - refreshContainer.appendChild(refreshLabel); - - // Close button - const closeBtn = this._createButton('✕', - () => { - this.stopLogAutoRefresh(); - modal.remove(); - }, - 'background-color: #c04c4c;'); - closeBtn.style.cssText += ' padding: 5px 10px; font-size: 14px; font-weight: bold;'; - - headerButtons.appendChild(refreshContainer); - headerButtons.appendChild(closeBtn); - - header.appendChild(title); - header.appendChild(headerButtons); - - // Log content area - const logContainer = document.createElement('div'); - logContainer.style.cssText = ` - flex: 1; - overflow: auto; - padding: 15px; - font-family: 'Consolas', 'Monaco', 'Courier New', monospace; - font-size: 12px; - line-height: 1.4; - color: #ddd; - background: #0d0d0d; - white-space: pre-wrap; - word-wrap: break-word; - `; - logContainer.id = 'distributed-log-content'; - logContainer.textContent = logData.content; - - // Auto-scroll to bottom - setTimeout(() => { - logContainer.scrollTop = logContainer.scrollHeight; - }, 0); - - // Status bar - const statusBar = document.createElement('div'); - statusBar.style.cssText = ` - padding: 10px 20px; - border-top: 1px solid #444; - font-size: 11px; - color: #888; - `; - statusBar.textContent = `Log file: ${logData.log_file}`; - if (logData.truncated) { - statusBar.textContent += ` (showing last ${logData.lines_shown} lines of ${this.formatFileSize(logData.file_size)})`; - } - - // Assemble modal - content.appendChild(header); - content.appendChild(logContainer); - content.appendChild(statusBar); - modal.appendChild(content); - - // Close on background click - modal.onclick = (e) => { - if (e.target === modal) { - this.stopLogAutoRefresh(); - modal.remove(); - } - }; - - // Close on Escape key - const handleEscape = (e) => { - if (e.key === 'Escape') { - this.stopLogAutoRefresh(); - modal.remove(); - document.removeEventListener('keydown', handleEscape); - } - }; - document.addEventListener('keydown', handleEscape); - - document.body.appendChild(modal); - - // Start auto-refresh - this.startLogAutoRefresh(workerId); - } - - async refreshLog(workerId, silent = false) { - const logContent = document.getElementById('distributed-log-content'); - if (!logContent) return; - - try { - const response = await fetch(`${window.location.origin}/distributed/worker_log/${workerId}?lines=1000`); - if (!response.ok) throw new Error('Failed to fetch log'); - - const data = await response.json(); - - // Update content - const shouldAutoScroll = logContent.scrollTop + logContent.clientHeight >= logContent.scrollHeight - 50; - logContent.textContent = data.content; - - // Auto-scroll if was at bottom - if (shouldAutoScroll) { - logContent.scrollTop = logContent.scrollHeight; - } - - // Only show toast if not in silent mode (manual refresh) - if (!silent) { - app.extensionManager.toast.add({ - severity: "success", - summary: "Log Refreshed", - detail: "Log content updated", - life: 2000 - }); - } - - } catch (error) { - // Only show error toast if not in silent mode - if (!silent) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Refresh Failed", - detail: error.message, - life: 3000 - }); - } - } - } - - copyLogPath(logPath) { - if (navigator.clipboard) { - navigator.clipboard.writeText(logPath).then(() => { - app.extensionManager.toast.add({ - severity: "success", - summary: "Path Copied", - detail: "Log file path copied to clipboard", - life: 2000 - }); - }).catch(err => { - this.debugLog('Failed to copy: ' + err.message); - }); - } - } - - formatFileSize(bytes) { - if (bytes < 1024) return bytes + ' B'; - if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB'; - return (bytes / (1024 * 1024)).toFixed(1) + ' MB'; - } - - isRemoteWorker(worker) { - const host = worker.host || window.location.hostname; - return host !== "localhost" && host !== "127.0.0.1" && host !== window.location.hostname; - } - - getMasterUrl() { - // Always use the detected/configured master IP for consistency - if (this.config?.master?.host) { - return `http://${this.config.master.host}:${window.location.port}`; - } - - // If no master IP is set but we're on a network address, use it - const hostname = window.location.hostname; - if (hostname !== 'localhost' && hostname !== '127.0.0.1') { - return window.location.origin; - } - - // Fallback warning - this won't work for remote workers - this.debugLog("No master host configured - remote workers won't be able to connect. " + - "Master host should be auto-detected on startup."); - return window.location.origin; - } - - async detectMasterIP() { - try { - const response = await fetch(`${window.location.origin}/distributed/network_info`); - if (!response.ok) { - this.debugLog("Failed to get network info"); - return; - } - - const data = await response.json(); - this.debugLog("Network info: " + JSON.stringify(data)); - - // Check if we already have a master host configured - if (this.config?.master?.host) { - this.debugLog(`Master host already configured: ${this.config.master.host}`); - return; - } - - // Use the recommended IP from the backend - if (data.recommended_ip && data.recommended_ip !== '127.0.0.1') { - this.log(`Auto-detected master IP: ${data.recommended_ip}`); - - // Save the detected IP (pass true to suppress notification) - await this.saveMasterIp(data.recommended_ip, true); - - // Show a single combined notification - if (app.extensionManager?.toast) { - app.extensionManager.toast.add({ - severity: "info", - summary: "Master IP Auto-Detected", - detail: `Master IP set to ${data.recommended_ip} for remote worker communication`, - life: 5000 - }); - } - } - } catch (error) { - this.debugLog("Error detecting master IP: " + error.message); - } - } - - async saveMasterIp(masterIp, suppressNotification = false) { - try { - const response = await fetch(`${window.location.origin}/distributed/config/update_master`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - host: masterIp.trim() - }) - }); - - if (response.ok) { - // Update local config - if (!this.config.master) { - this.config.master = {}; - } - this.config.master.host = masterIp.trim(); - - // Only show notification if not suppressed - if (!suppressNotification) { - app.extensionManager.toast.add({ - severity: "success", - summary: "Master Host Saved", - detail: masterIp.trim() ? `Master host set to ${masterIp.trim()}` : "Master host cleared (auto-detect)", - life: 3000 - }); - } - } else { - throw new Error("Failed to save master host"); - } - } catch (error) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Save Failed", - detail: error.message, - life: 5000 - }); - } - } - - startLogAutoRefresh(workerId) { - // Stop any existing auto-refresh - this.stopLogAutoRefresh(); - - // Refresh every 2 seconds - this.logAutoRefreshInterval = setInterval(() => { - this.refreshLog(workerId, true); // silent mode - }, 2000); - } - - stopLogAutoRefresh() { - if (this.logAutoRefreshInterval) { - clearInterval(this.logAutoRefreshInterval); - this.logAutoRefreshInterval = null; - } - } - - toggleWorkerExpanded(workerId) { - const settingsDiv = document.getElementById(`settings-${workerId}`); - const gpuDiv = settingsDiv.closest('[style*="margin-bottom: 12px"]'); - const settingsArrow = gpuDiv.querySelector('.settings-arrow'); - - if (!settingsDiv) return; - - if (this.expandedWorkers.has(workerId)) { - this.expandedWorkers.delete(workerId); - settingsDiv.classList.remove("expanded"); - if (settingsArrow) { - settingsArrow.style.transform = "rotate(0deg)"; - } - // Animate padding to 0 - settingsDiv.style.padding = "0 12px"; - settingsDiv.style.marginTop = "0"; - settingsDiv.style.marginBottom = "0"; - } else { - this.expandedWorkers.add(workerId); - settingsDiv.classList.add("expanded"); - if (settingsArrow) { - settingsArrow.style.transform = "rotate(90deg)"; - } - // Animate padding to full - settingsDiv.style.padding = "12px"; - settingsDiv.style.marginTop = "8px"; - settingsDiv.style.marginBottom = "8px"; - } - } - - createWorkerSettingsForm(worker) { - const form = document.createElement("div"); - form.style.cssText = "display: flex; flex-direction: column; gap: 10px;"; - - // Name field - const nameGroup = this.createFormGroup("Name:", worker.name, `name-${worker.id}`); - form.appendChild(nameGroup); - - // Remote checkbox - const remoteGroup = document.createElement("div"); - remoteGroup.style.cssText = "display: flex; align-items: center; gap: 8px; margin: 5px 0;"; - - const remoteCheckbox = document.createElement("input"); - remoteCheckbox.type = "checkbox"; - remoteCheckbox.id = `remote-${worker.id}`; - remoteCheckbox.checked = this.isRemoteWorker(worker); - remoteCheckbox.onchange = (e) => { - const isRemote = e.target.checked; - // Show/hide relevant fields - const hostGroup = document.getElementById(`host-group-${worker.id}`); - const cudaGroup = document.getElementById(`cuda-group-${worker.id}`); - const argsGroup = document.getElementById(`args-group-${worker.id}`); - - if (isRemote) { - hostGroup.style.display = "flex"; - cudaGroup.style.display = "none"; - argsGroup.style.display = "none"; - // If switching to remote and host is localhost, clear it - const hostInput = document.getElementById(`host-${worker.id}`); - if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { - hostInput.value = ""; - } - } else { - hostGroup.style.display = "none"; - cudaGroup.style.display = "flex"; - argsGroup.style.display = "flex"; - } - }; - - const remoteLabel = document.createElement("label"); - remoteLabel.htmlFor = `remote-${worker.id}`; - remoteLabel.textContent = "Remote Worker"; - remoteLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; - - remoteGroup.appendChild(remoteCheckbox); - remoteGroup.appendChild(remoteLabel); - form.appendChild(remoteGroup); - - // Host field (only for remote workers) - const hostGroup = this.createFormGroup("Host:", worker.host || "", `host-${worker.id}`, "text", "e.g., 192.168.1.100"); - hostGroup.id = `host-group-${worker.id}`; - hostGroup.style.display = this.isRemoteWorker(worker) ? "flex" : "none"; - form.appendChild(hostGroup); - - // Port field - const portGroup = this.createFormGroup("Port:", worker.port, `port-${worker.id}`, "number"); - form.appendChild(portGroup); - - // CUDA Device field (only for local workers) - const cudaGroup = this.createFormGroup("CUDA Device:", worker.cuda_device || 0, `cuda-${worker.id}`, "number"); - cudaGroup.id = `cuda-group-${worker.id}`; - cudaGroup.style.display = this.isRemoteWorker(worker) ? "none" : "flex"; - form.appendChild(cudaGroup); - - // Extra Args field (only for local workers) - const argsGroup = this.createFormGroup("Extra Args:", worker.extra_args || "", `args-${worker.id}`); - argsGroup.id = `args-group-${worker.id}`; - argsGroup.style.display = this.isRemoteWorker(worker) ? "none" : "flex"; - form.appendChild(argsGroup); - - // Buttons - const buttonGroup = document.createElement("div"); - buttonGroup.style.cssText = "display: flex; gap: 8px; margin-top: 8px;"; - - const saveBtn = this._createButton("Save", - () => this.saveWorkerSettings(worker.id), - "background-color: #4a7c4a;"); - saveBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; - - const cancelBtn = this._createButton("Cancel", - () => this.cancelWorkerSettings(worker.id), - "background-color: #555;"); - cancelBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; - - const deleteBtn = this._createButton("Delete", - () => this.deleteWorker(worker.id), - "background-color: #7c4a4a;"); - deleteBtn.style.cssText += " padding: 6px 12px; font-size: 12px; margin-left: auto;"; - - buttonGroup.appendChild(saveBtn); - buttonGroup.appendChild(cancelBtn); - buttonGroup.appendChild(deleteBtn); - - form.appendChild(buttonGroup); - - return form; - } - - createFormGroup(label, value, id, type = "text", placeholder = "") { - const group = document.createElement("div"); - group.style.cssText = "display: flex; align-items: center; gap: 8px;"; - - const labelEl = document.createElement("label"); - labelEl.textContent = label; - labelEl.style.cssText = "font-size: 12px; color: #ccc; width: 80px;"; - - const input = document.createElement("input"); - input.type = type; - input.id = id; - input.value = value; - if (placeholder) { - input.placeholder = placeholder; - } - input.style.cssText = "flex: 1; padding: 4px 8px; background: #2a2a2a; border: 1px solid #444; border-radius: 3px; color: #fff; font-size: 12px;"; - - group.appendChild(labelEl); - group.appendChild(input); - - return group; - } - - async saveWorkerSettings(workerId) { - const worker = this.config.workers.find(w => w.id === workerId); - if (!worker) return; - - // Get form values - const name = document.getElementById(`name-${workerId}`).value; - const isRemote = document.getElementById(`remote-${workerId}`).checked; - const host = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname; - const port = parseInt(document.getElementById(`port-${workerId}`).value); - const cudaDevice = isRemote ? undefined : parseInt(document.getElementById(`cuda-${workerId}`).value); - const extraArgs = isRemote ? undefined : document.getElementById(`args-${workerId}`).value; - - // Validate - if (!name.trim()) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Validation Error", - detail: "Worker name is required", - life: 3000 - }); - return; - } - - if (isRemote && !host.trim()) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Validation Error", - detail: "Host is required for remote workers", - life: 3000 - }); - return; - } - - if (isNaN(port) || port < 1024 || port > 65535) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Validation Error", - detail: "Port must be between 1024 and 65535", - life: 3000 - }); - return; - } - - // Check for port conflicts - // Remote workers can reuse ports, but local workers cannot share ports with each other or master - if (!isRemote) { - // Check if port conflicts with master - const masterPort = parseInt(window.location.port) || (window.location.protocol === 'https:' ? 443 : 80); - if (port === masterPort) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Port Conflict", - detail: `Port ${port} is already in use by the master server`, - life: 3000 - }); - return; - } - - // Check if port conflicts with other local workers - const localPortConflict = this.config.workers.some(w => - w.id !== workerId && - w.port === port && - !w.host // local workers have no host or host is null - ); - - if (localPortConflict) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Port Conflict", - detail: `Port ${port} is already in use by another local worker`, - life: 3000 - }); - return; - } - } else { - // For remote workers, only check conflicts with other workers on the same host - const sameHostConflict = this.config.workers.some(w => - w.id !== workerId && - w.port === port && - w.host === host.trim() - ); - - if (sameHostConflict) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Port Conflict", - detail: `Port ${port} is already in use by another worker on ${host}`, - life: 3000 - }); - return; - } - } - - try { - const response = await fetch(`${window.location.origin}/distributed/config/update_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - worker_id: workerId, - name: name.trim(), - host: isRemote ? host.trim() : null, - port: port, - cuda_device: isRemote ? null : cudaDevice, - extra_args: isRemote ? null : (extraArgs ? extraArgs.trim() : "") - }) - }); - - if (response.ok) { - // Update local config - worker.name = name.trim(); - if (isRemote) { - worker.host = host.trim(); - delete worker.cuda_device; - delete worker.extra_args; - } else { - delete worker.host; - worker.cuda_device = cudaDevice; - worker.extra_args = extraArgs ? extraArgs.trim() : ""; - } - worker.port = port; - - app.extensionManager.toast.add({ - severity: "success", - summary: "Settings Saved", - detail: `Worker ${name} settings updated`, - life: 3000 - }); - - // Refresh the UI - if (this.panelElement) { - this.renderSidebarContent(this.panelElement); - } - } else { - const error = await response.json(); - throw new Error(error.message || "Failed to save settings"); - } - } catch (error) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Save Failed", - detail: error.message, - life: 5000 - }); - } - } - - cancelWorkerSettings(workerId) { - // Collapse the settings panel - this.toggleWorkerExpanded(workerId); - - // Reset form values to original - const worker = this.config.workers.find(w => w.id === workerId); - if (worker) { - document.getElementById(`name-${workerId}`).value = worker.name; - document.getElementById(`host-${workerId}`).value = worker.host || ""; - document.getElementById(`port-${workerId}`).value = worker.port; - document.getElementById(`cuda-${workerId}`).value = worker.cuda_device || 0; - document.getElementById(`args-${workerId}`).value = worker.extra_args || ""; - - // Reset remote checkbox - const remoteCheckbox = document.getElementById(`remote-${workerId}`); - if (remoteCheckbox) { - remoteCheckbox.checked = this.isRemoteWorker(worker); - } - } - } - - async deleteWorker(workerId) { - const worker = this.config.workers.find(w => w.id === workerId); - if (!worker) return; - - // Confirm deletion - if (!confirm(`Are you sure you want to delete worker "${worker.name}"?`)) { - return; - } - - try { - const response = await fetch(`${window.location.origin}/distributed/config/delete_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ worker_id: workerId }) - }); - - if (response.ok) { - // Remove from local config - const index = this.config.workers.findIndex(w => w.id === workerId); - if (index !== -1) { - this.config.workers.splice(index, 1); - } - - app.extensionManager.toast.add({ - severity: "success", - summary: "Worker Deleted", - detail: `Worker ${worker.name} has been removed`, - life: 3000 - }); - - // Refresh the UI - if (this.panelElement) { - this.renderSidebarContent(this.panelElement); - } - } else { - const error = await response.json(); - throw new Error(error.message || "Failed to delete worker"); - } - } catch (error) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Delete Failed", - detail: error.message, - life: 5000 - }); - } - } - - async addNewWorker() { - // Generate new worker ID - const newId = `worker_${Date.now()}`; - - // Find next available port - const usedPorts = this.config.workers.map(w => w.port); - let nextPort = 8189; - while (usedPorts.includes(nextPort)) { - nextPort++; - } - - // Create new worker object - const newWorker = { - id: newId, - name: `Worker ${this.config.workers.length + 1}`, - port: nextPort, - cuda_device: this.config.workers.length, - enabled: true, // Default to enabled for convenience - extra_args: "" - }; - - // Add to config - this.config.workers.push(newWorker); - - // Save immediately - try { - const response = await fetch(`${window.location.origin}/distributed/config/update_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - worker_id: newId, - name: newWorker.name, - port: newWorker.port, - cuda_device: newWorker.cuda_device, - extra_args: newWorker.extra_args, - enabled: newWorker.enabled - }) - }); - - if (!response.ok) { - // Remove from local config if save failed - this.config.workers.pop(); - throw new Error("Failed to save new worker"); - } - - app.extensionManager.toast.add({ - severity: "success", - summary: "Worker Added", - detail: `New worker created on port ${nextPort}`, - life: 3000 - }); - - // Refresh UI and expand the new worker - this.expandedWorkers.add(newId); - if (this.panelElement) { - this.renderSidebarContent(this.panelElement); - } - - } catch (error) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Failed to Add Worker", - detail: error.message, - life: 5000 - }); - } - } -} - -app.registerExtension({ - name: "Distributed.Panel", - async setup() { - new DistributedExtension(); - } -}); \ No newline at end of file diff --git a/web/image_batch_divider.js b/web/image_batch_divider.js new file mode 100644 index 0000000..522530e --- /dev/null +++ b/web/image_batch_divider.js @@ -0,0 +1,86 @@ +import { app } from "/scripts/app.js"; + +app.registerExtension({ + name: "Distributed.ImageBatchDivider", + async nodeCreated(node) { + if (node.comfyClass === "ImageBatchDivider") { + try { + const updateOutputs = () => { + if (!node.widgets) return; + + const divideByWidget = node.widgets.find(w => w.name === "divide_by"); + if (!divideByWidget) return; + + const divideBy = parseInt(divideByWidget.value, 10) || 1; + const totalOutputs = divideBy; // Direct divide by value + + // Ensure outputs array exists + if (!node.outputs) node.outputs = []; + + // Remove excess outputs + while (node.outputs.length > totalOutputs) { + node.removeOutput(node.outputs.length - 1); + } + + // Add missing outputs + while (node.outputs.length < totalOutputs) { + const outputIndex = node.outputs.length + 1; + node.addOutput(`batch_${outputIndex}`, "IMAGE"); + } + + if (node.setDirty) node.setDirty(true); // Refresh canvas + }; + + // Initial update with delay to allow workflow loading + setTimeout(updateOutputs, 200); + + // Find the widget and set up responsive handlers + const divideByWidget = node.widgets.find(w => w.name === "divide_by"); + if (divideByWidget) { + // Override callback for immediate trigger on value set + const originalCallback = divideByWidget.callback; + divideByWidget.callback = (value) => { + updateOutputs(); + if (originalCallback) originalCallback.call(divideByWidget, value); // Preserve 'this' context + }; + + // Add event listener for real-time input changes (e.g., typing/dragging) + if (divideByWidget.inputEl) { + divideByWidget.inputEl.addEventListener('input', updateOutputs); + } + + // Lightweight MutationObserver as fallback (observe attributes on widget element if available) + const observer = new MutationObserver(updateOutputs); + if (divideByWidget.element) { + observer.observe(divideByWidget.element, { attributes: true, childList: true, subtree: true }); + } + + // Store cleanup function + node._batchDividerCleanup = () => { + observer.disconnect(); + if (divideByWidget.inputEl) { + divideByWidget.inputEl.removeEventListener('input', updateOutputs); + } + divideByWidget.callback = originalCallback; // Restore original + }; + } + + // Add post-configure hook for reliable workflow loading + const originalConfigure = node.configure; + node.configure = function(data) { + const result = originalConfigure ? originalConfigure.call(this, data) : undefined; + updateOutputs(); // Re-run after config load + return result; + }; + } catch (error) { + console.error("Error in ImageBatchDivider extension:", error); + } + } + }, + + nodeBeforeRemove(node) { + if (node.comfyClass === "ImageBatchDivider" && node._batchDividerCleanup) { + node._batchDividerCleanup(); + } + } +}); \ No newline at end of file diff --git a/web/main.js b/web/main.js new file mode 100644 index 0000000..51c2f07 --- /dev/null +++ b/web/main.js @@ -0,0 +1,1336 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; +import { DistributedUI } from './ui.js'; + +import { createStateManager } from './stateManager.js'; +import { createApiClient } from './apiClient.js'; +import { renderSidebarContent } from './sidebarRenderer.js'; +import { handleWorkerOperation, handleInterruptWorkers, handleClearMemory } from './workerUtils.js'; +import { setupInterceptor, executeParallelDistributed } from './executionUtils.js'; +import { BUTTON_STYLES, PULSE_ANIMATION_CSS, TIMEOUTS, STATUS_COLORS } from './constants.js'; + +class DistributedExtension { + constructor() { + this.config = null; + this.originalQueuePrompt = api.queuePrompt.bind(api); + this.statusCheckInterval = null; + this.logAutoRefreshInterval = null; + this.masterSettingsExpanded = false; + this.app = app; // Store app reference for toast notifications + + // Initialize centralized state + this.state = createStateManager(); + + // Initialize UI component factory + this.ui = new DistributedUI(); + + // Initialize API client + this.api = createApiClient(window.location.origin); + + // Initialize status check timeout reference + this.statusCheckTimeout = null; + + // Initialize abort controller for status checks + this.statusCheckAbortController = null; + + // Inject CSS for pulsing animation + this.injectStyles(); + + this.loadConfig().then(async () => { + this.registerSidebarTab(); + this.setupInterceptor(); + // Don't start polling until panel opens + // this.startStatusChecking(); + this.loadManagedWorkers(); + // Detect master IP after everything is set up + this.detectMasterIP(); + }); + } + + // Debug logging helpers + log(message, level = "info") { + if (level === "debug" && !this.config?.settings?.debug) return; + if (level === "error") { + console.error(`[Distributed] ${message}`); + } else { + console.log(`[Distributed] ${message}`); + } + } + + // Generate UUID with fallback for non-secure contexts + generateUUID() { + if (crypto.randomUUID) { + return crypto.randomUUID(); + } + // Fallback for non-secure contexts + return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, function(c) { + const r = Math.random() * 16 | 0; + const v = c == 'x' ? r : (r & 0x3 | 0x8); + return v.toString(16); + }); + } + + injectStyles() { + const styleId = 'distributed-styles'; + if (!document.getElementById(styleId)) { + const style = document.createElement('style'); + style.id = styleId; + style.textContent = PULSE_ANIMATION_CSS; + document.head.appendChild(style); + } + } + + // --- 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 { + this.config = await this.api.getConfig(); + this.log("Loaded config: " + JSON.stringify(this.config), "debug"); + + // Ensure default flag values + if (!this.config.settings) { + this.config.settings = {}; + } + if (this.config.settings.has_auto_populated_workers === undefined) { + this.config.settings.has_auto_populated_workers = false; + } + + // Load stored master CUDA device + this.masterCudaDevice = this.config?.master?.cuda_device ?? undefined; + + // Sync to state + if (this.config.workers) { + this.config.workers.forEach(w => { + this.state.updateWorker(w.id, { enabled: w.enabled }); + }); + } + } catch (error) { + this.log("Failed to load config: " + error.message, "error"); + this.config = { workers: [], settings: { has_auto_populated_workers: false } }; + } + } + + async updateWorkerEnabled(workerId, enabled) { + const worker = this.config.workers.find(w => w.id === workerId); + if (worker) { + worker.enabled = enabled; + this.state.updateWorker(workerId, { enabled }); + + // Immediately update status dot based on enabled state + const statusDot = document.getElementById(`status-${workerId}`); + if (statusDot) { + if (enabled) { + // Enabled: Start with checking state and trigger check + this.ui.updateStatusDot(workerId, STATUS_COLORS.OFFLINE_RED, "Checking status...", true); + setTimeout(() => this.checkWorkerStatus(worker), TIMEOUTS.STATUS_CHECK_DELAY); + } else { + // Disabled: Set to gray + this.ui.updateStatusDot(workerId, STATUS_COLORS.DISABLED_GRAY, "Disabled", false); + } + } + } + + try { + await this.api.updateWorker(workerId, { enabled }); + } catch (error) { + this.log("Error updating worker: " + error.message, "error"); + } + } + + async _updateSetting(key, value) { + // Update local config + if (!this.config.settings) { + this.config.settings = {}; + } + this.config.settings[key] = value; + + try { + await this.api.updateSetting(key, value); + + app.extensionManager.toast.add({ + severity: "success", + summary: "Setting Updated", + detail: `${key.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase())} ${value ? 'enabled' : 'disabled'}`, + life: 2000 + }); + } catch (error) { + this.log(`Error updating setting '${key}': ${error.message}`, "error"); + app.extensionManager.toast.add({ + severity: "error", + summary: "Setting Update Failed", + detail: error.message, + life: 3000 + }); + } + } + + // --- UI Rendering --- + + registerSidebarTab() { + app.extensionManager.registerSidebarTab({ + id: "distributed", + icon: "pi pi-server", + title: "Distributed", + tooltip: "Distributed Control Panel", + type: "custom", + render: (el) => { + this.panelElement = el; + this.onPanelOpen(); + return renderSidebarContent(this, el); + }, + destroy: () => { + this.onPanelClose(); + } + }); + } + + onPanelOpen() { + this.log("Panel opened - starting status polling", "debug"); + if (!this.statusCheckTimeout) { + this.checkAllWorkerStatuses(); + } + } + + onPanelClose() { + this.log("Panel closed - stopping status polling", "debug"); + + // Cancel any pending status checks + if (this.statusCheckAbortController) { + this.statusCheckAbortController.abort(); + this.statusCheckAbortController = null; + } + + // Clear the timeout + if (this.statusCheckTimeout) { + clearTimeout(this.statusCheckTimeout); + this.statusCheckTimeout = null; + } + + this.panelElement = null; + } + + updateSummary() { + const summaryEl = document.getElementById('distributed-summary'); + if (summaryEl) { + // Count active workers from state + const activeWorkers = this.config.workers.filter(w => { + const status = this.state.getWorkerStatus(w.id); + return status.online; + }).length; + const totalGPUs = activeWorkers + 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."; + } + } + } + + // --- Core Logic & Execution --- + + setupInterceptor() { + setupInterceptor(this); + } + + async executeParallelDistributed(promptWrapper) { + return executeParallelDistributed(this, promptWrapper); + } + + startStatusChecking() { + this.checkAllWorkerStatuses(); + } + + async checkAllWorkerStatuses() { + // Don't continue if panel is closed + if (!this.panelElement) return; + + // Create new abort controller for this round of checks + this.statusCheckAbortController = new AbortController(); + + + // Check master status + this.checkMasterStatus(); + + if (!this.config || !this.config.workers) return; + + for (const worker of this.config.workers) { + // Check status for enabled workers OR workers that are launching + if (worker.enabled || this.state.isWorkerLaunching(worker.id)) { + this.checkWorkerStatus(worker); + } + } + + // Determine next interval based on current state + let isActive = this.state.getMasterStatus() === 'processing'; // Master is busy + + // Check workers for activity + this.config.workers.forEach(worker => { + const ws = this.state.getWorker(worker.id); // Get worker state + if (ws.launching || ws.status?.processing) { // Launching or processing + isActive = true; + } + }); + + // Set next delay: 1s if active, 5s if idle + const nextInterval = isActive ? 1000 : 5000; + + // Schedule the next check + this.statusCheckTimeout = setTimeout(() => this.checkAllWorkerStatuses(), nextInterval); + } + + async checkMasterStatus() { + try { + const response = await fetch(`${window.location.origin}/prompt`, { + method: 'GET', + signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK) + }); + + if (response.ok) { + const data = await response.json(); + const queueRemaining = data.exec_info?.queue_remaining || 0; + const isProcessing = queueRemaining > 0; + + // Update master status in state + this.state.setMasterStatus(isProcessing ? 'processing' : 'online'); + + // Update master status dot + const statusDot = document.getElementById('master-status'); + if (statusDot) { + if (isProcessing) { + statusDot.style.backgroundColor = "#f0ad4e"; + statusDot.title = `Processing (${queueRemaining} in queue)`; + } else { + statusDot.style.backgroundColor = "#4CAF50"; + statusDot.title = "Online"; + } + } + } + } catch (error) { + // Master is always online (we're running on it), so keep it green + const statusDot = document.getElementById('master-status'); + if (statusDot) { + statusDot.style.backgroundColor = "#4CAF50"; + statusDot.title = "Online"; + } + } + } + + // Helper to build worker URL + getWorkerUrl(worker, endpoint = '') { + const host = worker.host || window.location.hostname; + + // Cloud workers always use HTTPS + const isCloud = worker.type === 'cloud'; + + // Detect if we're running on Runpod (for local workers on Runpod infrastructure) + const isRunpodProxy = host.endsWith('.proxy.runpod.net'); + + // For local workers on Runpod, construct the port-specific proxy URL + let finalHost = host; + if (!worker.host && isRunpodProxy) { + const match = host.match(/^(.*)\.proxy\.runpod\.net$/); + if (match) { + const podId = match[1]; + const domain = 'proxy.runpod.net'; + finalHost = `${podId}-${worker.port}.${domain}`; + } else { + // Fallback or log error if no match (shouldn't happen) + console.error(`[Distributed] Failed to parse Runpod proxy host: ${host}`); + } + } + + // Determine protocol: HTTPS for cloud, Runpod proxies, or port 443 + const useHttps = isCloud || isRunpodProxy || worker.port === 443; + const protocol = useHttps ? 'https' : 'http'; + + // Only add port if non-standard + const defaultPort = useHttps ? 443 : 80; + const needsPort = !isRunpodProxy && worker.port !== defaultPort; + const portStr = needsPort ? `:${worker.port}` : ''; + + return `${protocol}://${finalHost}${portStr}${endpoint}`; + } + + async checkWorkerStatus(worker) { + // Assume caller ensured enabled; proceed with check + const url = this.getWorkerUrl(worker, '/prompt'); + const statusDot = document.getElementById(`status-${worker.id}`); + + try { + // Combine timeout with abort controller signal + const timeoutSignal = AbortSignal.timeout(TIMEOUTS.STATUS_CHECK); + const signal = this.statusCheckAbortController + ? AbortSignal.any([timeoutSignal, this.statusCheckAbortController.signal]) + : timeoutSignal; + + const response = await fetch(url, { + method: 'GET', + mode: 'cors', + signal: signal + }); + + if (response.ok) { + const data = await response.json(); + const queueRemaining = data.exec_info?.queue_remaining || 0; + const isProcessing = queueRemaining > 0; + + // Update status + this.state.setWorkerStatus(worker.id, { + online: true, + processing: isProcessing, + queueCount: queueRemaining + }); + + // Update status dot based on processing state + if (isProcessing) { + this.ui.updateStatusDot( + worker.id, + "#f0ad4e", + `Online - Processing (${queueRemaining} in queue)`, + false + ); + } else { + this.ui.updateStatusDot(worker.id, "#3ca03c", "Online - Idle", false); + } + + // Clear launching state since worker is now online + if (this.state.isWorkerLaunching(worker.id)) { + this.state.setWorkerLaunching(worker.id, false); + this.clearLaunchingFlag(worker.id); + } + } else { + throw new Error(`HTTP ${response.status}`); + } + } catch (error) { + // Don't process aborted requests + if (error.name === 'AbortError') { + return; + } + + // Worker is offline or unreachable + this.state.setWorkerStatus(worker.id, { + online: false, + processing: false, + queueCount: 0 + }); + + // Check if worker is launching + if (this.state.isWorkerLaunching(worker.id)) { + this.ui.updateStatusDot(worker.id, "#f0ad4e", "Launching...", true); + } else if (worker.enabled) { + // Only update to red if not currently launching AND still enabled + this.ui.updateStatusDot(worker.id, "#c04c4c", "Offline - Cannot connect", false); + } + // If disabled, don't update the dot (leave it gray) + + this.log(`Worker ${worker.id} status check failed: ${error.message}`, "debug"); + } + + // Update control buttons based on new status + this.updateWorkerControls(worker.id); + } + + async launchWorker(workerId) { + const worker = this.config.workers.find(w => w.id === workerId); + const launchBtn = document.querySelector(`#controls-${workerId} button`); + + // If worker is disabled, enable it first + if (!worker.enabled) { + await this.updateWorkerEnabled(workerId, true); + + // Update the checkbox UI + const checkbox = document.getElementById(`gpu-${workerId}`); + if (checkbox) { + checkbox.checked = true; + } + + this.updateSummary(); + } + + this.ui.updateStatusDot(workerId, "#f0ad4e", "Launching...", true); + this.state.setWorkerLaunching(workerId, true); + + // Allow 90 seconds for worker to launch (model loading can take time) + setTimeout(() => { + this.state.setWorkerLaunching(workerId, false); + }, TIMEOUTS.LAUNCH); + + if (!launchBtn) return; + + try { + // Disable button immediately + launchBtn.disabled = true; + + const result = await this.api.launchWorker(workerId); + if (result) { + this.log(`Launched ${worker.name} (PID: ${result.pid})`, "info"); + if (result.log_file) { + this.log(`Log file: ${result.log_file}`, "debug"); + } + + this.state.setWorkerManaged(workerId, { + pid: result.pid, + log_file: result.log_file, + started_at: Date.now() + }); + + // Update controls immediately to hide launch button and show stop/log buttons + this.updateWorkerControls(workerId); + setTimeout(() => this.checkWorkerStatus(worker), TIMEOUTS.STATUS_CHECK); + } + } catch (error) { + // Check if worker was already running + if (error.message && error.message.includes("already running")) { + this.log(`Worker ${worker.name} is already running`, "info"); + this.updateWorkerControls(workerId); + setTimeout(() => this.checkWorkerStatus(worker), TIMEOUTS.STATUS_CHECK_DELAY); + } else { + this.log(`Error launching worker: ${error.message || error}`, "error"); + + // Re-enable button on error + if (launchBtn) { + launchBtn.disabled = false; + } + } + } + } + + async stopWorker(workerId) { + const worker = this.config.workers.find(w => w.id === workerId); + const stopBtn = document.querySelectorAll(`#controls-${workerId} button`)[1]; + + // Provide immediate feedback + if (stopBtn) { + stopBtn.disabled = true; + stopBtn.textContent = "Stopping..."; + stopBtn.style.backgroundColor = "#666"; + } + + try { + const result = await this.api.stopWorker(workerId); + if (result) { + this.log(`Stopped worker: ${result.message}`, "info"); + this.state.setWorkerManaged(workerId, null); + + // Immediately update status to offline + this.ui.updateStatusDot(workerId, "#c04c4c", "Offline"); + this.state.setWorkerStatus(workerId, { online: false }); + + // Flash success feedback + if (stopBtn) { + stopBtn.style.backgroundColor = BUTTON_STYLES.success; + stopBtn.textContent = "Stopped!"; + setTimeout(() => { + this.updateWorkerControls(workerId); + }, TIMEOUTS.FLASH_SHORT); + } + + // Verify status after a short delay + setTimeout(() => this.checkWorkerStatus(worker), TIMEOUTS.STATUS_CHECK); + } else { + this.log(`Failed to stop worker: ${result.message}`, "error"); + + // Flash error feedback + if (stopBtn) { + stopBtn.style.backgroundColor = BUTTON_STYLES.error; + stopBtn.textContent = result.message.includes("already stopped") ? "Not Running" : "Failed"; + + // If already stopped, update status immediately + if (result.message.includes("already stopped")) { + this.ui.updateStatusDot(workerId, "#c04c4c", "Offline"); + this.state.setWorkerStatus(workerId, { online: false }); + } + + setTimeout(() => { + this.updateWorkerControls(workerId); + }, TIMEOUTS.FLASH_MEDIUM); + } + } + } catch (error) { + this.log(`Error stopping worker: ${error}`, "error"); + + // Reset button on error + if (stopBtn) { + stopBtn.style.backgroundColor = BUTTON_STYLES.error; + stopBtn.textContent = "Error"; + setTimeout(() => { + this.updateWorkerControls(workerId); + }, TIMEOUTS.FLASH_MEDIUM); + } + } + } + + async clearLaunchingFlag(workerId) { + try { + await this.api.clearLaunchingFlag(workerId); + this.log(`Cleared launching flag for worker ${workerId}`, "debug"); + } catch (error) { + this.log(`Error clearing launching flag: ${error.message || error}`, "error"); + } + } + + // Generic async button action handler + async handleAsyncButtonAction(button, action, successText, errorText, resetDelay = TIMEOUTS.BUTTON_RESET) { + const originalText = button.textContent; + const originalStyle = button.style.cssText; + button.disabled = true; + + try { + await action(); + button.textContent = successText; + button.style.cssText = originalStyle; + button.style.backgroundColor = BUTTON_STYLES.success; + return true; + } catch (error) { + button.textContent = errorText || `Error: ${error.message}`; + button.style.cssText = originalStyle; + button.style.backgroundColor = BUTTON_STYLES.error; + throw error; + } finally { + setTimeout(() => { + button.textContent = originalText; + button.style.cssText = originalStyle; + button.disabled = false; + }, resetDelay); + } + } + + /** + * Cleanup method to stop intervals and listeners + */ + cleanup() { + if (this.statusCheckInterval) { + clearInterval(this.statusCheckInterval); + this.statusCheckInterval = null; + } + + if (this.logAutoRefreshInterval) { + clearInterval(this.logAutoRefreshInterval); + this.logAutoRefreshInterval = null; + } + + if (this.statusCheckTimeout) { + clearTimeout(this.statusCheckTimeout); + this.statusCheckTimeout = null; + } + + this.log("Cleaned up intervals", "debug"); + } + + async loadManagedWorkers() { + try { + const result = await this.api.getManagedWorkers(); + + // Check for launching workers + for (const [workerId, info] of Object.entries(result.managed_workers)) { + this.state.setWorkerManaged(workerId, info); + + // If worker is marked as launching, add to launchingWorkers set + if (info.launching) { + this.state.setWorkerLaunching(workerId, true); + this.log(`Worker ${workerId} is in launching state`, "debug"); + } + } + + // Update UI for all workers + if (this.config?.workers) { + this.config.workers.forEach(w => this.updateWorkerControls(w.id)); + } + } catch (error) { + this.log(`Error loading managed workers: ${error}`, "error"); + } + } + + updateWorkerControls(workerId) { + const controlsDiv = document.getElementById(`controls-${workerId}`); + + if (!controlsDiv) return; + + const worker = this.config.workers.find(w => w.id === workerId); + if (!worker) return; + + // Skip button updates for remote workers + if (this.isRemoteWorker(worker)) { + return; + } + + // Ensure we check for string ID + const managedInfo = this.state.getWorker(workerId).managed; + const status = this.state.getWorkerStatus(workerId); + + // Update button states - buttons are now inside a wrapper div + const buttons = controlsDiv.querySelectorAll('button'); + const launchBtn = document.getElementById(`launch-${workerId}`); + const stopBtn = document.getElementById(`stop-${workerId}`); + const logBtn = document.getElementById(`log-${workerId}`); + + // Show log button immediately if we have log file info (even if worker is still starting) + if (managedInfo?.log_file && logBtn) { + logBtn.style.display = ''; + } else if (logBtn && !managedInfo) { + logBtn.style.display = 'none'; + } + + if (status?.online || managedInfo) { + // Worker is running or we just launched it + launchBtn.style.display = 'none'; // Hide launch button when running + + if (managedInfo) { + // Only show stop button if we manage this worker + stopBtn.style.display = ''; + stopBtn.disabled = false; + stopBtn.textContent = "Stop"; + stopBtn.style.backgroundColor = "#7c4a4a"; // Red when enabled + } else { + // Hide stop button for workers launched outside UI + stopBtn.style.display = 'none'; + } + } else { + // Worker is not running + launchBtn.style.display = ''; // Show launch button + launchBtn.disabled = false; + launchBtn.textContent = "Launch"; + launchBtn.style.backgroundColor = "#4a7c4a"; // Always green + + stopBtn.style.display = 'none'; // Hide stop button when not running + } + } + + async viewWorkerLog(workerId) { + const managedInfo = this.state.getWorker(workerId).managed; + if (!managedInfo?.log_file) return; + + const logBtn = document.getElementById(`log-${workerId}`); + + // Provide immediate feedback + if (logBtn) { + logBtn.disabled = true; + logBtn.textContent = "Loading..."; + logBtn.style.backgroundColor = "#666"; + } + + try { + // Fetch log content + const data = await this.api.getWorkerLog(workerId, 1000); + + // Create modal dialog + this.ui.showLogModal(this, workerId, data); + + // Restore button + if (logBtn) { + logBtn.disabled = false; + logBtn.textContent = "View Log"; + logBtn.style.backgroundColor = "#685434"; // Keep the yellow color + } + + } catch (error) { + this.log('Error viewing log: ' + error.message, "error"); + app.extensionManager.toast.add({ + severity: "error", + summary: "Error", + detail: `Failed to load log: ${error.message}`, + life: 5000 + }); + + // Flash error and restore button + if (logBtn) { + logBtn.style.backgroundColor = BUTTON_STYLES.error; + logBtn.textContent = "Error"; + setTimeout(() => { + logBtn.disabled = false; + logBtn.textContent = "View Log"; + logBtn.style.backgroundColor = "#685434"; // Keep the yellow color + }, TIMEOUTS.FLASH_LONG); + } + } + } + + async refreshLog(workerId, silent = false) { + const logContent = document.getElementById('distributed-log-content'); + if (!logContent) return; + + try { + const data = await this.api.getWorkerLog(workerId, 1000); + + // Update content + const shouldAutoScroll = logContent.scrollTop + logContent.clientHeight >= logContent.scrollHeight - 50; + logContent.textContent = data.content; + + // Auto-scroll if was at bottom + if (shouldAutoScroll) { + logContent.scrollTop = logContent.scrollHeight; + } + + // Only show toast if not in silent mode (manual refresh) + if (!silent) { + app.extensionManager.toast.add({ + severity: "success", + summary: "Log Refreshed", + detail: "Log content updated", + life: 2000 + }); + } + + } catch (error) { + // Only show error toast if not in silent mode + if (!silent) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Refresh Failed", + detail: error.message, + life: 3000 + }); + } + } + } + + isRemoteWorker(worker) { + // Check if explicitly marked as cloud worker + if (worker.type === "cloud") { + return true; + } + // Otherwise check by host (backward compatibility) + const host = worker.host || window.location.hostname; + return host !== "localhost" && host !== "127.0.0.1" && host !== window.location.hostname; + } + + isCloudWorker(worker) { + return worker.type === "cloud"; + } + + getMasterUrl() { + // Always use the detected/configured master IP for consistency + if (this.config?.master?.host) { + const configuredHost = this.config.master.host; + + // If the configured host already includes protocol, use as-is + if (configuredHost.startsWith('http://') || configuredHost.startsWith('https://')) { + return configuredHost; + } + + // For domain names (not IPs), default to HTTPS + const isIP = /^(\d{1,3}\.){3}\d{1,3}$/.test(configuredHost); + const isLocalhost = configuredHost === 'localhost' || configuredHost === '127.0.0.1'; + + if (!isIP && !isLocalhost && configuredHost.includes('.')) { + // It's a domain name, use HTTPS + return `https://${configuredHost}`; + } else { + // For IPs and localhost, use current access method + const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); + if ((window.location.protocol === 'https:' && port === '443') || + (window.location.protocol === 'http:' && port === '80')) { + return `${window.location.protocol}//${configuredHost}`; + } + return `${window.location.protocol}//${configuredHost}:${port}`; + } + } + + // If no master IP is set but we're on a network address, use it + const hostname = window.location.hostname; + if (hostname !== 'localhost' && hostname !== '127.0.0.1') { + return window.location.origin; + } + + // Fallback warning - this won't work for remote workers + this.log("No master host configured - remote workers won't be able to connect. " + + "Master host should be auto-detected on startup.", "debug"); + return window.location.origin; + } + + async detectMasterIP() { + try { + // Detect if we're running on Runpod + const isRunpod = window.location.hostname.endsWith('.proxy.runpod.net'); + if (isRunpod) { + this.log("Detected Runpod environment", "info"); + } + + const data = await this.api.getNetworkInfo(); + this.log("Network info: " + JSON.stringify(data), "debug"); + + // Store CUDA device info + if (data.cuda_device !== null && data.cuda_device !== undefined) { + this.masterCudaDevice = data.cuda_device; + + // Store persistently in config if not already set or changed + if (!this.config.master) this.config.master = {}; + if (this.config.master.cuda_device === undefined || this.config.master.cuda_device !== data.cuda_device) { + this.config.master.cuda_device = data.cuda_device; + try { + await this.api.updateMaster({ cuda_device: data.cuda_device }); + this.log(`Stored master CUDA device: ${data.cuda_device}`, "debug"); + } catch (error) { + this.log(`Error storing master CUDA device: ${error.message}`, "error"); + } + } + + // Update the master display with CUDA info + this.ui.updateMasterDisplay(this); + } + + // Store CUDA device count for auto-population + if (data.cuda_device_count > 0) { + this.cudaDeviceCount = data.cuda_device_count; + this.log(`Detected ${this.cudaDeviceCount} CUDA devices`, "info"); + + // Auto-populate workers if conditions are met + const shouldAutoPopulate = + !this.config.settings.has_auto_populated_workers && // Never populated before + (!this.config.workers || this.config.workers.length === 0); // No workers exist + + this.log(`Auto-population check: has_populated=${this.config.settings.has_auto_populated_workers}, workers=${this.config.workers ? this.config.workers.length : 'null'}, should_populate=${shouldAutoPopulate}`, "debug"); + + if (shouldAutoPopulate) { + this.log(`Auto-populating workers based on ${this.cudaDeviceCount} CUDA devices (excluding master on CUDA ${this.masterCudaDevice})`, "info"); + + const newWorkers = []; + let workerNum = 1; + let portOffset = 0; + + for (let i = 0; i < this.cudaDeviceCount; i++) { + // Skip the CUDA device used by master + if (i === this.masterCudaDevice) { + this.log(`Skipping CUDA ${i} (used by master)`, "debug"); + continue; + } + + const worker = { + id: crypto.randomUUID(), + name: `Worker ${workerNum}`, + host: isRunpod ? null : "localhost", + port: 8189 + portOffset, + cuda_device: i, + enabled: true, + extra_args: isRunpod ? "--listen" : "" + }; + newWorkers.push(worker); + workerNum++; + portOffset++; + } + + // Only proceed if we have workers to add + if (newWorkers.length > 0) { + this.log(`Auto-populating ${newWorkers.length} workers`, "info"); + + // Add workers to config + this.config.workers = newWorkers; + + // Set the flag to prevent future auto-population + this.config.settings.has_auto_populated_workers = true; + + // Save each worker using the update endpoint + for (const worker of newWorkers) { + try { + await this.api.updateWorker(worker.id, worker); + } catch (error) { + this.log(`Error saving worker ${worker.name}: ${error.message}`, "error"); + } + } + + // Save the updated settings + try { + await this.api.updateSetting('has_auto_populated_workers', true); + } catch (error) { + this.log(`Error saving auto-population flag: ${error.message}`, "error"); + } + + this.log(`Auto-populated ${newWorkers.length} workers and saved config`, "info"); + + // Show success notification + if (app.extensionManager?.toast) { + app.extensionManager.toast.add({ + severity: "success", + summary: "Workers Auto-populated", + detail: `Automatically created ${newWorkers.length} workers based on detected CUDA devices`, + life: 5000 + }); + } + + // Reload the config to include the new workers + await this.loadConfig(); + } else { + this.log("No additional CUDA devices available for workers (all used by master)", "debug"); + } + } + } + + // Check if we already have a master host configured + if (this.config?.master?.host) { + this.log(`Master host already configured: ${this.config.master.host}`, "debug"); + return; + } + + // For Runpod, use the proxy hostname as master host + if (isRunpod) { + const runpodHost = window.location.hostname; + this.log(`Setting Runpod master host: ${runpodHost}`, "info"); + + // Save the Runpod host + await this.api.updateMaster({ host: runpodHost }); + + // Update local config + if (!this.config.master) this.config.master = {}; + this.config.master.host = runpodHost; + + // Show notification + if (app.extensionManager?.toast) { + app.extensionManager.toast.add({ + severity: "info", + summary: "Runpod Auto-Configuration", + detail: `Master host set to ${runpodHost} with --listen flag for workers`, + life: 5000 + }); + } + return; // Skip regular IP detection for Runpod + } + + // Use the recommended IP from the backend + if (data.recommended_ip && data.recommended_ip !== '127.0.0.1') { + this.log(`Auto-detected master IP: ${data.recommended_ip}`, "info"); + + // Save the detected IP (pass true to suppress notification) + await this.api.updateMaster({ host: data.recommended_ip }); + + // Update local config immediately + if (!this.config.master) this.config.master = {}; + this.config.master.host = data.recommended_ip; + } + } catch (error) { + this.log("Error detecting master IP: " + error.message, "error"); + } + } + + async saveWorkerSettings(workerId) { + const worker = this.config.workers.find(w => w.id === workerId); + if (!worker) return; + + // Get form values + const name = document.getElementById(`name-${workerId}`).value; + const workerType = document.getElementById(`worker-type-${workerId}`).value; + const isRemote = workerType === 'remote' || workerType === 'cloud'; + const isCloud = workerType === 'cloud'; + const host = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname; + const port = parseInt(document.getElementById(`port-${workerId}`).value); + const cudaDevice = isRemote ? undefined : parseInt(document.getElementById(`cuda-${workerId}`).value); + const extraArgs = isRemote ? undefined : document.getElementById(`args-${workerId}`).value; + + // Validate + if (!name.trim()) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Validation Error", + detail: "Worker name is required", + life: 3000 + }); + return; + } + + if ((workerType === 'remote' || workerType === 'cloud') && !host.trim()) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Validation Error", + detail: "Host is required for remote workers", + life: 3000 + }); + return; + } + + if (!isCloud && (isNaN(port) || port < 1 || port > 65535)) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Validation Error", + detail: "Port must be between 1 and 65535", + life: 3000 + }); + return; + } + + // Check for port conflicts + // Remote workers can reuse ports, but local workers cannot share ports with each other or master + if (!isRemote) { + // Check if port conflicts with master + const masterPort = parseInt(window.location.port) || (window.location.protocol === 'https:' ? 443 : 80); + if (port === masterPort) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Port Conflict", + detail: `Port ${port} is already in use by the master server`, + life: 3000 + }); + return; + } + + // Check if port conflicts with other local workers + const localPortConflict = this.config.workers.some(w => + w.id !== workerId && + w.port === port && + !w.host // local workers have no host or host is null + ); + + if (localPortConflict) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Port Conflict", + detail: `Port ${port} is already in use by another local worker`, + life: 3000 + }); + return; + } + } else { + // For remote workers, only check conflicts with other workers on the same host + const sameHostConflict = this.config.workers.some(w => + w.id !== workerId && + w.port === port && + w.host === host.trim() + ); + + if (sameHostConflict) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Port Conflict", + detail: `Port ${port} is already in use by another worker on ${host}`, + life: 3000 + }); + return; + } + } + + try { + await this.api.updateWorker(workerId, { + name: name.trim(), + type: workerType, + host: isRemote ? host.trim() : null, + port: port, + cuda_device: isRemote ? null : cudaDevice, + extra_args: isRemote ? null : (extraArgs ? extraArgs.trim() : "") + }); + + // Update local config + worker.name = name.trim(); + worker.type = workerType; + if (isRemote) { + worker.host = host.trim(); + delete worker.cuda_device; + delete worker.extra_args; + } else { + delete worker.host; + worker.cuda_device = cudaDevice; + worker.extra_args = extraArgs ? extraArgs.trim() : ""; + } + worker.port = port; + + // Sync to state + this.state.updateWorker(workerId, { enabled: worker.enabled }); + + app.extensionManager.toast.add({ + severity: "success", + summary: "Settings Saved", + detail: `Worker ${name} settings updated`, + life: 3000 + }); + + // Refresh the UI + if (this.panelElement) { + renderSidebarContent(this, this.panelElement); + } + } catch (error) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Save Failed", + detail: error.message, + life: 5000 + }); + } + } + + cancelWorkerSettings(workerId) { + // Collapse the settings panel + this.toggleWorkerExpanded(workerId); + + // Reset form values to original + const worker = this.config.workers.find(w => w.id === workerId); + if (worker) { + document.getElementById(`name-${workerId}`).value = worker.name; + document.getElementById(`host-${workerId}`).value = worker.host || ""; + document.getElementById(`port-${workerId}`).value = worker.port; + document.getElementById(`cuda-${workerId}`).value = worker.cuda_device || 0; + document.getElementById(`args-${workerId}`).value = worker.extra_args || ""; + + // Reset remote checkbox + const remoteCheckbox = document.getElementById(`remote-${workerId}`); + if (remoteCheckbox) { + remoteCheckbox.checked = this.isRemoteWorker(worker); + } + } + } + + async deleteWorker(workerId) { + const worker = this.config.workers.find(w => w.id === workerId); + if (!worker) return; + + // Confirm deletion + if (!confirm(`Are you sure you want to delete worker "${worker.name}"?`)) { + return; + } + + try { + await this.api.deleteWorker(workerId); + + // Remove from local config + const index = this.config.workers.findIndex(w => w.id === workerId); + if (index !== -1) { + this.config.workers.splice(index, 1); + } + + app.extensionManager.toast.add({ + severity: "success", + summary: "Worker Deleted", + detail: `Worker ${worker.name} has been removed`, + life: 3000 + }); + + // Refresh the UI + if (this.panelElement) { + renderSidebarContent(this, this.panelElement); + } + } catch (error) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Delete Failed", + detail: error.message, + life: 5000 + }); + } + } + + async addNewWorker() { + // Generate new worker ID using UUID (fallback for non-secure contexts) + const newId = this.generateUUID(); + + // Find next available port + const usedPorts = this.config.workers.map(w => w.port); + let nextPort = 8189; + while (usedPorts.includes(nextPort)) { + nextPort++; + } + + // Create new worker object + const newWorker = { + id: newId, + name: `Worker ${this.config.workers.length + 1}`, + port: nextPort, + cuda_device: this.config.workers.length, + enabled: true, // Default to enabled for convenience + extra_args: "" + }; + + // Add to config + this.config.workers.push(newWorker); + + // Save immediately + try { + await this.api.updateWorker(newId, { + name: newWorker.name, + port: newWorker.port, + cuda_device: newWorker.cuda_device, + extra_args: newWorker.extra_args, + enabled: newWorker.enabled + }); + + // Sync to state + this.state.updateWorker(newId, { enabled: true }); + + app.extensionManager.toast.add({ + severity: "success", + summary: "Worker Added", + detail: `New worker created on port ${nextPort}`, + life: 3000 + }); + + // Refresh UI and expand the new worker + this.state.setWorkerExpanded(newId, true); + if (this.panelElement) { + renderSidebarContent(this, this.panelElement); + } + + } catch (error) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Failed to Add Worker", + detail: error.message, + life: 5000 + }); + } + } + + startLogAutoRefresh(workerId) { + // Stop any existing auto-refresh + this.stopLogAutoRefresh(); + + // Refresh every 2 seconds + this.logAutoRefreshInterval = setInterval(() => { + this.refreshLog(workerId, true); // silent mode + }, TIMEOUTS.LOG_REFRESH); + } + + stopLogAutoRefresh() { + if (this.logAutoRefreshInterval) { + clearInterval(this.logAutoRefreshInterval); + this.logAutoRefreshInterval = null; + } + } + + toggleWorkerExpanded(workerId) { + const settingsDiv = document.getElementById(`settings-${workerId}`); + const gpuDiv = settingsDiv.closest('[style*="margin-bottom: 12px"]'); + const settingsArrow = gpuDiv.querySelector('.settings-arrow'); + + if (!settingsDiv) return; + + if (this.state.isWorkerExpanded(workerId)) { + this.state.setWorkerExpanded(workerId, false); + settingsDiv.classList.remove("expanded"); + if (settingsArrow) { + settingsArrow.style.transform = "rotate(0deg)"; + } + // Animate padding to 0 + settingsDiv.style.padding = "0 12px"; + settingsDiv.style.marginTop = "0"; + settingsDiv.style.marginBottom = "0"; + } else { + this.state.setWorkerExpanded(workerId, true); + settingsDiv.classList.add("expanded"); + if (settingsArrow) { + settingsArrow.style.transform = "rotate(90deg)"; + } + // Animate padding to full + settingsDiv.style.padding = "12px"; + settingsDiv.style.marginTop = "8px"; + settingsDiv.style.marginBottom = "8px"; + } + } + + _handleInterruptWorkers(button) { + return handleInterruptWorkers(this, button); + } + + _handleClearMemory(button) { + return handleClearMemory(this, button); + } +} + +app.registerExtension({ + name: "Distributed.Panel", + async setup() { + new DistributedExtension(); + } +}); \ No newline at end of file diff --git a/web/sidebarRenderer.js b/web/sidebarRenderer.js new file mode 100644 index 0000000..9fc95a0 --- /dev/null +++ b/web/sidebarRenderer.js @@ -0,0 +1,277 @@ +import { BUTTON_STYLES } from './constants.js'; + +export async function renderSidebarContent(extension, el) { + // Panel is being opened/rendered + extension.log("Panel opened", "debug"); + + if (!el) { + extension.log("No element provided to renderSidebarContent", "debug"); + return; + } + + // Prevent infinite recursion + if (extension._isRendering) { + extension.log("Already rendering, skipping", "debug"); + return; + } + extension._isRendering = true; + + try { + // Store reference to the panel element + extension.panelElement = el; + + // Show loading indicator + el.innerHTML = ''; + const loadingDiv = document.createElement("div"); + loadingDiv.style.cssText = "display: flex; align-items: center; justify-content: center; height: calc(100vh - 100px); color: #888;"; + loadingDiv.innerHTML = ` + + `; + el.appendChild(loadingDiv); + + // Add rotation animation + const style = document.createElement('style'); + style.textContent = ` + @keyframes rotate { + from { transform: rotate(0deg); } + to { transform: rotate(360deg); } + } + `; + document.head.appendChild(style); + loadingDiv.querySelector('svg').style.animation = 'rotate 1s linear infinite'; + + // Preload data outside render + await extension.loadConfig(); + await extension.loadManagedWorkers(); + + el.innerHTML = ''; + + // Create toolbar header to match ComfyUI style + const toolbar = document.createElement("div"); + toolbar.className = "p-toolbar p-component border-x-0 border-t-0 rounded-none px-2 py-1 min-h-8"; + toolbar.style.cssText = "border-bottom: 1px solid #444; background: transparent; display: flex; align-items: center;"; + + const toolbarStart = document.createElement("div"); + toolbarStart.className = "p-toolbar-start"; + toolbarStart.style.cssText = "display: flex; align-items: center;"; + + const titleSpan = document.createElement("span"); + titleSpan.className = "text-xs 2xl:text-sm truncate"; + titleSpan.textContent = "COMFYUI DISTRIBUTED"; + titleSpan.title = "ComfyUI Distributed"; + + toolbarStart.appendChild(titleSpan); + toolbar.appendChild(toolbarStart); + + const toolbarCenter = document.createElement("div"); + toolbarCenter.className = "p-toolbar-center"; + toolbar.appendChild(toolbarCenter); + + const toolbarEnd = document.createElement("div"); + toolbarEnd.className = "p-toolbar-end"; + toolbar.appendChild(toolbarEnd); + + el.appendChild(toolbar); + + // Main container with adjusted padding + const container = document.createElement("div"); + container.style.cssText = "padding: 15px; display: flex; flex-direction: column; height: calc(100% - 32px);"; + + // Detect master info on panel open (in case CUDA info wasn't available at startup) + extension.log(`Panel opened. CUDA device count: ${extension.cudaDeviceCount}, Workers: ${extension.config?.workers?.length || 0}`, "debug"); + if (!extension.cudaDeviceCount) { + await extension.detectMasterIP(); + } + + + // Now render with guaranteed up-to-date config + // Master Node Section + const masterDiv = extension.ui.renderEntityCard('master', extension.config?.master, extension); + container.appendChild(masterDiv); + + // Workers Section (no heading) + const gpuSection = document.createElement("div"); + gpuSection.style.cssText = "flex: 1; overflow-y: auto; margin-bottom: 15px;"; + + const gpuList = document.createElement("div"); + const workers = extension.config?.workers || []; + + // If no workers exist, show a full blueprint placeholder first + if (workers.length === 0) { + const blueprintDiv = extension.ui.renderEntityCard('blueprint', { onClick: () => extension.addNewWorker() }, extension); + gpuList.appendChild(blueprintDiv); + } + + // Show existing workers + workers.forEach(worker => { + const gpuDiv = extension.ui.renderEntityCard('worker', worker, extension); + gpuList.appendChild(gpuDiv); + }); + gpuSection.appendChild(gpuList); + + // Only show the minimal "Add Worker" box if there are existing workers + if (workers.length > 0) { + const addWorkerDiv = extension.ui.renderEntityCard('add', { onClick: () => extension.addNewWorker() }, extension); + gpuSection.appendChild(addWorkerDiv); + } + + container.appendChild(gpuSection); + + const actionsSection = document.createElement("div"); + actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;"; + + // Create a row for both buttons + const buttonRow = document.createElement("div"); + buttonRow.style.cssText = "display: flex; gap: 8px;"; + + const clearMemButton = extension.ui.createButtonHelper( + "Clear Worker VRAM", + (e) => extension._handleClearMemory(e.target), + BUTTON_STYLES.clearMemory + ); + clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)"; + clearMemButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.clearMemory; + + const interruptButton = extension.ui.createButtonHelper( + "Interrupt Workers", + (e) => extension._handleInterruptWorkers(e.target), + BUTTON_STYLES.interrupt + ); + interruptButton.title = "Cancel/interrupt execution on all enabled worker GPUs"; + interruptButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.interrupt; + + buttonRow.appendChild(clearMemButton); + buttonRow.appendChild(interruptButton); + actionsSection.appendChild(buttonRow); + + container.appendChild(actionsSection); + + // Settings section + const settingsSection = document.createElement("div"); + settingsSection.style.cssText = "border-top: 1px solid #444; padding-top: 10px; margin-bottom: 10px;"; + + // Settings header with toggle + const settingsHeader = document.createElement("div"); + settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between; cursor: pointer; user-select: none;"; + + const workerSettingsTitle = document.createElement("h4"); + workerSettingsTitle.textContent = "Settings"; + workerSettingsTitle.style.cssText = "margin: 0; font-size: 14px;"; + + const workerSettingsToggle = document.createElement("span"); + workerSettingsToggle.textContent = "▶"; // Right arrow when collapsed + workerSettingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; + + settingsHeader.appendChild(workerSettingsTitle); + settingsHeader.appendChild(workerSettingsToggle); + + // Hover effect for header + settingsHeader.onmouseover = () => { + workerSettingsToggle.style.color = "#fff"; + }; + settingsHeader.onmouseout = () => { + workerSettingsToggle.style.color = "#888"; + }; + + // Collapsible settings content + const settingsContent = document.createElement("div"); + settingsContent.style.cssText = "max-height: 0; overflow: hidden; opacity: 0; transition: max-height 0.3s ease, opacity 0.3s ease;"; + + const settingsDiv = document.createElement("div"); + settingsDiv.style.cssText = "display: flex; flex-direction: column; gap: 8px; padding-top: 10px;"; + + // Toggle functionality + let settingsExpanded = false; + settingsHeader.onclick = () => { + settingsExpanded = !settingsExpanded; + if (settingsExpanded) { + settingsContent.style.maxHeight = "200px"; + settingsContent.style.opacity = "1"; + workerSettingsToggle.style.transform = "rotate(90deg)"; + } else { + settingsContent.style.maxHeight = "0"; + settingsContent.style.opacity = "0"; + workerSettingsToggle.style.transform = "rotate(0deg)"; + } + }; + + // Debug mode setting + const debugGroup = document.createElement("div"); + debugGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; + + const debugCheckbox = document.createElement("input"); + debugCheckbox.type = "checkbox"; + debugCheckbox.id = "setting-debug"; + debugCheckbox.checked = extension.config?.settings?.debug || false; + debugCheckbox.onchange = (e) => extension._updateSetting('debug', e.target.checked); + + const debugLabel = document.createElement("label"); + debugLabel.htmlFor = "setting-debug"; + debugLabel.textContent = "Debug Mode"; + debugLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + + debugGroup.appendChild(debugCheckbox); + debugGroup.appendChild(debugLabel); + + // Auto-launch workers setting + const autoLaunchGroup = document.createElement("div"); + autoLaunchGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; + + const autoLaunchCheckbox = document.createElement("input"); + autoLaunchCheckbox.type = "checkbox"; + autoLaunchCheckbox.id = "setting-auto-launch"; + autoLaunchCheckbox.checked = extension.config?.settings?.auto_launch_workers || false; + autoLaunchCheckbox.onchange = (e) => extension._updateSetting('auto_launch_workers', e.target.checked); + + const autoLaunchLabel = document.createElement("label"); + autoLaunchLabel.htmlFor = "setting-auto-launch"; + autoLaunchLabel.textContent = "Auto-launch Local Workers on Startup"; + autoLaunchLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + + autoLaunchGroup.appendChild(autoLaunchCheckbox); + autoLaunchGroup.appendChild(autoLaunchLabel); + + // Stop workers on exit setting + const stopOnExitGroup = document.createElement("div"); + stopOnExitGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; + + const stopOnExitCheckbox = document.createElement("input"); + stopOnExitCheckbox.type = "checkbox"; + stopOnExitCheckbox.id = "setting-stop-on-exit"; + stopOnExitCheckbox.checked = extension.config?.settings?.stop_workers_on_master_exit !== false; // Default true + stopOnExitCheckbox.onchange = (e) => extension._updateSetting('stop_workers_on_master_exit', e.target.checked); + + const stopOnExitLabel = document.createElement("label"); + stopOnExitLabel.htmlFor = "setting-stop-on-exit"; + stopOnExitLabel.textContent = "Stop Local Workers on Master Exit"; + stopOnExitLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + + stopOnExitGroup.appendChild(stopOnExitCheckbox); + stopOnExitGroup.appendChild(stopOnExitLabel); + + settingsDiv.appendChild(debugGroup); + settingsDiv.appendChild(autoLaunchGroup); + settingsDiv.appendChild(stopOnExitGroup); + settingsContent.appendChild(settingsDiv); + + settingsSection.appendChild(settingsHeader); + settingsSection.appendChild(settingsContent); + container.appendChild(settingsSection); + + const summarySection = document.createElement("div"); + summarySection.style.cssText = "border-top: 1px solid #444; padding-top: 10px;"; + const summary = document.createElement("div"); + summary.id = "distributed-summary"; + summary.style.cssText = "font-size: 11px; color: #888;"; + summarySection.appendChild(summary); + container.appendChild(summarySection); + el.appendChild(container); + extension.updateSummary(); + + // Start checking worker statuses immediately in parallel + setTimeout(() => extension.checkAllWorkerStatuses(), 0); + } finally { + // Always reset the rendering flag + extension._isRendering = false; + } +} \ No newline at end of file diff --git a/web/stateManager.js b/web/stateManager.js new file mode 100644 index 0000000..f400f4a --- /dev/null +++ b/web/stateManager.js @@ -0,0 +1,61 @@ +export function createStateManager() { + const state = { + workers: new Map(), // Unified worker state: { status, managed, launching, expanded, ... } + masterStatus: 'online', + }; + + return { + // Worker state management + getWorker(workerId) { + return state.workers.get(String(workerId)) || {}; + }, + + updateWorker(workerId, updates) { + const id = String(workerId); + const current = state.workers.get(id) || {}; + state.workers.set(id, { ...current, ...updates }); + return state.workers.get(id); + }, + + setWorkerStatus(workerId, status) { + return this.updateWorker(workerId, { status }); + }, + + setWorkerManaged(workerId, info) { + return this.updateWorker(workerId, { managed: info }); + }, + + setWorkerLaunching(workerId, launching) { + return this.updateWorker(workerId, { launching }); + }, + + setWorkerExpanded(workerId, expanded) { + return this.updateWorker(workerId, { expanded }); + }, + + isWorkerLaunching(workerId) { + return this.getWorker(workerId).launching || false; + }, + + isWorkerExpanded(workerId) { + return this.getWorker(workerId).expanded || false; + }, + + isWorkerManaged(workerId) { + return !!this.getWorker(workerId).managed; + }, + + getWorkerStatus(workerId) { + return this.getWorker(workerId).status || {}; + }, + + // Master state + setMasterStatus(status) { + state.masterStatus = status; + }, + + getMasterStatus() { + return state.masterStatus; + } + }; +} \ No newline at end of file diff --git a/web/ui.js b/web/ui.js new file mode 100644 index 0000000..e3681ea --- /dev/null +++ b/web/ui.js @@ -0,0 +1,1236 @@ +import { BUTTON_STYLES, UI_STYLES, STATUS_COLORS, UI_COLORS, TIMEOUTS } from './constants.js'; + +const cardConfigs = { + master: { + checkbox: { + enabled: false, + checked: true, + disabled: true, + opacity: 0.6, + title: "Master node is always enabled" + }, + statusDot: { + color: STATUS_COLORS.ONLINE_GREEN, + title: 'Online', + id: 'master-status', + dynamic: true + }, + infoText: (data, extension) => { + const cudaDevice = extension.config?.master?.cuda_device ?? extension.masterCudaDevice; + const cudaInfo = cudaDevice !== undefined ? `CUDA ${cudaDevice} • ` : ''; + const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); + return `${data?.name || extension.config?.master?.name || "Master"}
${cudaInfo}Port ${port}`; + }, + controls: { + type: 'info', + text: 'Master', + style: "background-color: #333; color: #999;" + }, + settings: { + formType: 'master', + id: 'master-settings', + expandedTracker: 'masterSettingsExpanded' + }, + hover: true, + expand: true, + border: 'solid' + }, + worker: { + checkbox: { + enabled: true, + title: "Enable/disable this worker" + }, + statusDot: { + dynamic: true, + initialColor: (data) => data.enabled ? STATUS_COLORS.OFFLINE_RED : STATUS_COLORS.DISABLED_GRAY, + initialTitle: (data) => data.enabled ? "Checking status..." : "Disabled", + pulsing: (data) => data.enabled, + id: (data) => `status-${data.id}` + }, + infoText: (data, extension) => { + const isRemote = extension.isRemoteWorker(data); + const isCloud = data.type === 'cloud'; + + if (isCloud) { + // For cloud workers, don't show port (it's always 443) + return `${data.name}
${data.host}`; + } else if (isRemote) { + return `${data.name}
${data.host}:${data.port}`; + } else { + const cudaInfo = data.cuda_device !== undefined ? `CUDA ${data.cuda_device} • ` : ''; + return `${data.name}
${cudaInfo}Port ${data.port}`; + } + }, + controls: { + dynamic: true + }, + settings: { + formType: 'worker', + id: (data) => `settings-${data.id}`, + expandedId: (data) => data?.id + }, + hover: true, + expand: true, + border: 'solid' + }, + blueprint: { + checkbox: { + type: 'icon', + content: '+', + width: 42, + style: `border-right: 2px dashed ${UI_COLORS.BORDER_LIGHT}; color: ${UI_COLORS.ACCENT_COLOR}; font-size: 24px; font-weight: 500;` + }, + statusDot: { + color: 'transparent', + border: `1px solid ${UI_COLORS.BORDER_LIGHT}` + }, + infoText: () => `Add New Worker
[CUDA] • [Port]`, + controls: { + type: 'ghost', + text: 'Configure', + style: `border: 1px solid ${UI_COLORS.BORDER_DARK}; background: transparent; color: ${UI_COLORS.BORDER_LIGHT};` + }, + hover: 'placeholder', + expand: false, + border: 'dashed' + }, + add: { + checkbox: { + type: 'icon', + content: '+', + width: 43, + style: `border-right: 1px dashed ${UI_COLORS.BORDER_DARK}; color: ${UI_COLORS.BORDER_LIGHT}; font-size: 18px;` + }, + statusDot: { + color: 'transparent', + border: `1px solid ${UI_COLORS.BORDER_LIGHT}` + }, + infoText: () => `Add New Worker`, + controls: null, + hover: 'placeholder', + expand: false, + border: 'dashed', + minHeight: '48px' + } +}; + +export class DistributedUI { + constructor() { + // UI element styles + this.styles = UI_STYLES; + } + + createStatusDot(id, color = "#666", title = "Status") { + const dot = document.createElement("span"); + if (id) dot.id = id; + dot.style.cssText = this.styles.statusDot + ` background-color: ${color};`; + dot.title = title; + return dot; + } + + createButton(text, onClick, customStyle = "") { + const button = document.createElement("button"); + button.textContent = text; + button.className = "distributed-button"; + button.style.cssText = BUTTON_STYLES.base + customStyle; + if (onClick) button.onclick = onClick; + return button; + } + + createButtonGroup(buttons, style = "") { + const group = document.createElement("div"); + group.style.cssText = this.styles.buttonGroup + style; + buttons.forEach(button => group.appendChild(button)); + return group; + } + + createWorkerControls(workerId, handlers = {}) { + const controlsDiv = document.createElement("div"); + controlsDiv.id = `controls-${workerId}`; + controlsDiv.style.cssText = this.styles.controlsDiv; + + const buttons = []; + + if (handlers.launch) { + const launchBtn = this.createButton('Launch', handlers.launch); + launchBtn.id = `launch-${workerId}`; + launchBtn.title = "Launch this worker instance"; + buttons.push(launchBtn); + } + + if (handlers.stop) { + const stopBtn = this.createButton('Stop', handlers.stop); + stopBtn.id = `stop-${workerId}`; + stopBtn.title = "Stop this worker instance"; + buttons.push(stopBtn); + } + + if (handlers.viewLog) { + const logBtn = this.createButton('View Log', handlers.viewLog); + logBtn.id = `log-${workerId}`; + logBtn.title = "View worker log file"; + buttons.push(logBtn); + } + + buttons.forEach(btn => controlsDiv.appendChild(btn)); + return controlsDiv; + } + + createFormGroup(label, value, id, type = "text", placeholder = "") { + const group = document.createElement("div"); + group.style.cssText = this.styles.formGroup; + + const labelEl = document.createElement("label"); + labelEl.textContent = label; + labelEl.htmlFor = id; + labelEl.style.cssText = this.styles.formLabel; + + const input = document.createElement("input"); + input.type = type; + input.id = id; + input.value = value; + input.placeholder = placeholder; + input.style.cssText = this.styles.formInput; + + group.appendChild(labelEl); + group.appendChild(input); + return { group, input }; + } + + + createInfoBox(text) { + const box = document.createElement("div"); + box.style.cssText = this.styles.infoBox; + box.textContent = text; + return box; + } + + addHoverEffect(element, onHover, onLeave) { + element.onmouseover = onHover; + element.onmouseout = onLeave; + } + + createCard(type = 'worker', options = {}) { + const card = document.createElement("div"); + + switch(type) { + case 'master': + case 'worker': + card.style.cssText = this.styles.workerCard; + break; + case 'blueprint': + card.style.cssText = this.styles.cardBase + this.styles.cardBlueprint; + if (options.onClick) card.onclick = options.onClick; + if (options.title) card.title = options.title; + break; + case 'add': + card.style.cssText = this.styles.cardBase + this.styles.cardAdd; + if (options.onClick) card.onclick = options.onClick; + if (options.title) card.title = options.title; + break; + } + + if (options.onMouseEnter) { + card.addEventListener('mouseenter', options.onMouseEnter); + } + if (options.onMouseLeave) { + card.addEventListener('mouseleave', options.onMouseLeave); + } + + return card; + } + + createCardColumn(type = 'checkbox', options = {}) { + const column = document.createElement("div"); + + switch(type) { + case 'checkbox': + column.style.cssText = this.styles.checkboxColumn; + if (options.title) column.title = options.title; + break; + case 'icon': + column.style.cssText = this.styles.columnBase + this.styles.iconColumn; + break; + case 'content': + column.style.cssText = this.styles.contentColumn; + break; + } + + return column; + } + + createInfoRow(options = {}) { + const row = document.createElement("div"); + row.style.cssText = this.styles.infoRow; + if (options.onClick) row.onclick = options.onClick; + return row; + } + + createWorkerContent() { + const content = document.createElement("div"); + content.style.cssText = this.styles.workerContent; + return content; + } + + createSettingsForm(fields = [], options = {}) { + const form = document.createElement("div"); + form.style.cssText = this.styles.settingsForm; + + fields.forEach(field => { + if (field.type === 'checkbox') { + const group = document.createElement("div"); + group.style.cssText = this.styles.checkboxGroup; + + const checkbox = document.createElement("input"); + checkbox.type = "checkbox"; + checkbox.id = field.id; + checkbox.checked = field.checked || false; + if (field.onChange) checkbox.onchange = field.onChange; + + const label = document.createElement("label"); + label.htmlFor = field.id; + label.textContent = field.label; + label.style.cssText = this.styles.formLabelClickable; + + group.appendChild(checkbox); + group.appendChild(label); + form.appendChild(group); + } else { + const result = this.createFormGroup(field.label, field.value, field.id, field.type, field.placeholder); + if (field.groupId) result.group.id = field.groupId; + if (field.display) result.group.style.display = field.display; + form.appendChild(result.group); + } + }); + + if (options.buttons) { + const buttonGroup = this.createButtonGroup(options.buttons, options.buttonStyle || " margin-top: 8px;"); + form.appendChild(buttonGroup); + } + + return form; + } + + + createButtonHelper(text, onClick, style) { + return this.createButton(text, onClick, style); + } + + updateMasterDisplay(extension) { + // Use persistent config value as fallback + const cudaDevice = extension?.config?.master?.cuda_device ?? extension?.masterCudaDevice; + + // Update CUDA info if element exists + const cudaInfo = document.getElementById('master-cuda-info'); + if (cudaInfo) { + const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); + if (cudaDevice !== undefined && cudaDevice !== null) { + cudaInfo.textContent = `CUDA ${cudaDevice} • Port ${port}`; + } else { + cudaInfo.textContent = `Port ${port}`; + } + } + + // Update name if changed + const nameDisplay = document.getElementById('master-name-display'); + if (nameDisplay && extension?.config?.master?.name) { + nameDisplay.textContent = extension.config.master.name; + } + } + + showToast(app, severity, summary, detail, life = 3000) { + if (app.extensionManager?.toast?.add) { + app.extensionManager.toast.add({ severity, summary, detail, life }); + } + } + + showCloudflareWarning(extension, masterHost) { + // Remove any existing banner first + const existingBanner = document.getElementById('cloudflare-warning-banner'); + if (existingBanner) { + existingBanner.remove(); + } + + // Create warning banner + const banner = document.createElement('div'); + banner.id = 'cloudflare-warning-banner'; + banner.style.cssText = ` + position: fixed; + top: 0; + left: 0; + right: 0; + background: #ff9800; + color: #333; + padding: 8px 16px; + text-align: center; + z-index: 10000; + display: flex; + align-items: center; + justify-content: center; + gap: 16px; + box-shadow: 0 2px 5px rgba(0,0,0,0.2); + font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; + `; + + const messageSpan = document.createElement('span'); + messageSpan.innerHTML = `Connection issue: Master address ${masterHost} is not reachable. The cloudflare tunnel may be offline.`; + messageSpan.style.fontSize = '13px'; + + const resetButton = document.createElement('button'); + resetButton.textContent = 'Reset Master Address'; + resetButton.style.cssText = ` + background: #333; + color: white; + border: none; + padding: 6px 14px; + border-radius: 4px; + cursor: pointer; + font-weight: 500; + font-size: 13px; + transition: background 0.2s; + `; + resetButton.onmouseover = () => resetButton.style.background = '#555'; + resetButton.onmouseout = () => resetButton.style.background = '#333'; + + const dismissButton = document.createElement('button'); + dismissButton.textContent = 'Dismiss'; + dismissButton.style.cssText = ` + background: transparent; + color: #333; + border: 1px solid #333; + padding: 6px 14px; + border-radius: 4px; + cursor: pointer; + font-size: 13px; + transition: opacity 0.2s; + `; + dismissButton.onmouseover = () => dismissButton.style.opacity = '0.7'; + dismissButton.onmouseout = () => dismissButton.style.opacity = '1'; + + // Add click handlers + resetButton.onclick = async () => { + resetButton.disabled = true; + resetButton.textContent = 'Resetting...'; + + try { + // Save with empty host - this will trigger auto-detection + await extension.api.updateMaster({ + name: extension.config?.master?.name || "Master", + host: "" + }); + + // Clear the local config host so detectMasterIP() doesn't skip + if (extension.config?.master) { + extension.config.master.host = ""; + } + + // The API call above doesn't trigger auto-detection, so we need to do it + await extension.detectMasterIP(); + + // Reload config to get the new detected IP + await extension.loadConfig(); + + // Log the new master URL for debugging + const newMasterUrl = extension.getMasterUrl(); + extension.log(`Master host reset. New URL: ${newMasterUrl}`, "info"); + + // Update UI if sidebar is open + if (extension.panelElement) { + const hostInput = document.getElementById('master-host'); + if (hostInput) { + hostInput.value = extension.config?.master?.host || ""; + } + } + + // Show success message with the actual URL that will be used + extension.app.extensionManager.toast.add({ + severity: "success", + summary: "Master Host Reset", + detail: `New address: ${newMasterUrl}`, + life: 4000 + }); + + banner.remove(); + } catch (error) { + resetButton.disabled = false; + resetButton.textContent = 'Reset Master Host'; + extension.log(`Failed to reset master host: ${error.message}`, "error"); + } + }; + + dismissButton.onclick = () => banner.remove(); + + // Assemble banner + banner.appendChild(messageSpan); + banner.appendChild(resetButton); + banner.appendChild(dismissButton); + + // Add to page + document.body.prepend(banner); + + // Auto-dismiss after 30 seconds + setTimeout(() => { + if (document.getElementById('cloudflare-warning-banner')) { + banner.style.transition = 'opacity 0.5s'; + banner.style.opacity = '0'; + setTimeout(() => banner.remove(), 500); + } + }, 30000); + } + + updateStatusDot(workerId, color, title, pulsing = false) { + const statusDot = document.getElementById(`status-${workerId}`); + if (!statusDot) return; + + statusDot.style.backgroundColor = color; + statusDot.title = title; + statusDot.classList.toggle('status-pulsing', pulsing); + } + + showLogModal(extension, workerId, logData) { + // Remove any existing modal + const existingModal = document.getElementById('distributed-log-modal'); + if (existingModal) { + existingModal.remove(); + } + + const worker = extension.config.workers.find(w => w.id === workerId); + const workerName = worker?.name || `Worker ${workerId}`; + + // Create modal container + const modal = document.createElement('div'); + modal.id = 'distributed-log-modal'; + modal.style.cssText = ` + position: fixed; + top: 0; + left: 0; + width: 100%; + height: 100%; + background: rgba(0, 0, 0, 0.8); + display: flex; + align-items: center; + justify-content: center; + z-index: 10000; + `; + + // Create modal content + const content = document.createElement('div'); + content.style.cssText = ` + background: #1e1e1e; + border-radius: 8px; + width: 90%; + max-width: 1200px; + height: 80%; + display: flex; + flex-direction: column; + border: 1px solid #444; + `; + + // Header + const header = document.createElement('div'); + header.style.cssText = ` + padding: 15px 20px; + border-bottom: 1px solid #444; + display: flex; + justify-content: space-between; + align-items: center; + `; + + const title = document.createElement('h3'); + title.textContent = `${workerName} - Log Viewer`; + title.style.cssText = 'margin: 0; color: #fff;'; + + const headerButtons = document.createElement('div'); + headerButtons.style.cssText = 'display: flex; gap: 20px; align-items: center;'; + + // Auto-refresh container + const refreshContainer = document.createElement('div'); + refreshContainer.style.cssText = 'display: flex; align-items: center; gap: 4px;'; + + // Auto-refresh checkbox + const refreshCheckbox = document.createElement('input'); + refreshCheckbox.type = 'checkbox'; + refreshCheckbox.id = 'log-auto-refresh'; + refreshCheckbox.checked = true; // Enabled by default + refreshCheckbox.style.cssText = 'cursor: pointer;'; + refreshCheckbox.onchange = (e) => { + if (e.target.checked) { + extension.startLogAutoRefresh(workerId); + } else { + extension.stopLogAutoRefresh(); + } + }; + + const refreshLabel = document.createElement('label'); + refreshLabel.htmlFor = 'log-auto-refresh'; + refreshLabel.style.cssText = 'font-size: 12px; color: #ccc; cursor: pointer; white-space: nowrap;'; + refreshLabel.textContent = 'Auto-refresh'; + + // Add checkbox and label to container + refreshContainer.appendChild(refreshCheckbox); + refreshContainer.appendChild(refreshLabel); + + // Close button + const closeBtn = this.createButton('✕', + () => { + extension.stopLogAutoRefresh(); + modal.remove(); + }, + 'background-color: #c04c4c;'); + closeBtn.style.cssText += ' padding: 5px 10px; font-size: 14px; font-weight: bold;'; + + headerButtons.appendChild(refreshContainer); + headerButtons.appendChild(closeBtn); + + header.appendChild(title); + header.appendChild(headerButtons); + + // Log content area + const logContainer = document.createElement('div'); + logContainer.style.cssText = ` + flex: 1; + overflow: auto; + padding: 15px; + font-family: 'Consolas', 'Monaco', 'Courier New', monospace; + font-size: 12px; + line-height: 1.4; + color: #ddd; + background: #0d0d0d; + white-space: pre-wrap; + word-wrap: break-word; + `; + logContainer.id = 'distributed-log-content'; + logContainer.textContent = logData.content; + + // Auto-scroll to bottom + setTimeout(() => { + logContainer.scrollTop = logContainer.scrollHeight; + }, 0); + + // Status bar + const statusBar = document.createElement('div'); + statusBar.style.cssText = ` + padding: 10px 20px; + border-top: 1px solid #444; + font-size: 11px; + color: #888; + `; + statusBar.textContent = `Log file: ${logData.log_file}`; + if (logData.truncated) { + statusBar.textContent += ` (showing last ${logData.lines_shown} lines of ${this.formatFileSize(logData.file_size)})`; + } + + // Assemble modal + content.appendChild(header); + content.appendChild(logContainer); + content.appendChild(statusBar); + modal.appendChild(content); + + // Close on background click + modal.onclick = (e) => { + if (e.target === modal) { + extension.stopLogAutoRefresh(); + modal.remove(); + } + }; + + // Close on Escape key + const handleEscape = (e) => { + if (e.key === 'Escape') { + extension.stopLogAutoRefresh(); + modal.remove(); + document.removeEventListener('keydown', handleEscape); + } + }; + document.addEventListener('keydown', handleEscape); + + document.body.appendChild(modal); + + // Start auto-refresh + extension.startLogAutoRefresh(workerId); + } + + formatFileSize(bytes) { + if (bytes < 1024) return bytes + ' B'; + if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB'; + return (bytes / (1024 * 1024)).toFixed(1) + ' MB'; + } + + createWorkerSettingsForm(extension, worker) { + const form = document.createElement("div"); + form.style.cssText = "display: flex; flex-direction: column; gap: 8px;"; + + // Name field + const nameGroup = this.createFormGroup("Name:", worker.name, `name-${worker.id}`); + form.appendChild(nameGroup.group); + + // Worker type dropdown + const typeGroup = document.createElement("div"); + typeGroup.style.cssText = "display: flex; flex-direction: column; gap: 4px; margin: 5px 0;"; + + const typeLabel = document.createElement("label"); + typeLabel.htmlFor = `worker-type-${worker.id}`; + typeLabel.textContent = "Worker Type:"; + typeLabel.style.cssText = "font-size: 12px; color: #ccc;"; + + const typeSelect = document.createElement("select"); + typeSelect.id = `worker-type-${worker.id}`; + typeSelect.style.cssText = "padding: 4px 8px; background: #333; color: #fff; border: 1px solid #555; border-radius: 4px; font-size: 12px;"; + + // Create options + const localOption = document.createElement("option"); + localOption.value = "local"; + localOption.textContent = "Local"; + + const remoteOption = document.createElement("option"); + remoteOption.value = "remote"; + remoteOption.textContent = "Remote"; + + const cloudOption = document.createElement("option"); + cloudOption.value = "cloud"; + cloudOption.textContent = "Cloud"; + + typeSelect.appendChild(localOption); + typeSelect.appendChild(remoteOption); + typeSelect.appendChild(cloudOption); + + // Create powered by Runpod text (initially hidden) + const runpodText = document.createElement("a"); + runpodText.id = `runpod-text-${worker.id}`; + runpodText.href = "https://get.runpod.io/0bw29uf3ug0p"; + runpodText.target = "_blank"; + runpodText.textContent = "Deploy Cloud Worker with Runpod"; + runpodText.style.cssText = "font-size: 12px; color: #4a90e2; text-decoration: none; margin-top: 4px; display: none; cursor: pointer;"; + + // Store the onchange function to be assigned later + const createOnChangeHandler = () => { + return (e) => { + const workerType = e.target.value; + // Show/hide relevant fields + const hostGroup = document.getElementById(`host-group-${worker.id}`); + const hostInput = document.getElementById(`host-${worker.id}`); + const portGroup = document.getElementById(`port-group-${worker.id}`); + const portInput = document.getElementById(`port-${worker.id}`); + const cudaGroup = document.getElementById(`cuda-group-${worker.id}`); + const argsGroup = document.getElementById(`args-group-${worker.id}`); + const runpodTextElem = document.getElementById(`runpod-text-${worker.id}`); + + // Check if elements exist before accessing them + if (!hostGroup || !portGroup || !cudaGroup || !argsGroup || !runpodTextElem || !hostInput || !portInput) { + return; // Elements not ready yet + } + + if (workerType === "local") { + hostGroup.style.display = "none"; + portGroup.style.display = "flex"; + cudaGroup.style.display = "flex"; + argsGroup.style.display = "flex"; + runpodTextElem.style.display = "none"; + } else if (workerType === "remote") { + hostGroup.style.display = "flex"; + portGroup.style.display = "flex"; + cudaGroup.style.display = "none"; + argsGroup.style.display = "none"; + runpodTextElem.style.display = "none"; + // Update placeholder for remote workers + hostInput.placeholder = "e.g., 192.168.1.100"; + // If switching to remote and host is localhost, clear it + if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { + hostInput.value = ""; + } + } else if (workerType === "cloud") { + hostGroup.style.display = "flex"; + portGroup.style.display = "flex"; // Keep port visible for cloud workers + cudaGroup.style.display = "none"; + argsGroup.style.display = "none"; + runpodTextElem.style.display = "block"; + // Update placeholder for cloud workers + hostInput.placeholder = "e.g., your-cloud-worker.trycloudflare.com"; + // Set port to 443 for cloud workers + portInput.value = "443"; + // If switching to cloud and host is localhost, clear it + if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { + hostInput.value = ""; + } + } + }; + }; + + typeGroup.appendChild(typeLabel); + typeGroup.appendChild(typeSelect); + typeGroup.appendChild(runpodText); + form.appendChild(typeGroup); + + // Host field (only for remote workers) + const hostGroup = this.createFormGroup("Host:", worker.host || "", `host-${worker.id}`, "text", "e.g., 192.168.1.100"); + hostGroup.group.id = `host-group-${worker.id}`; + hostGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "flex" : "none"; + form.appendChild(hostGroup.group); + + // Port field + const portGroup = this.createFormGroup("Port:", worker.port, `port-${worker.id}`, "number"); + portGroup.group.id = `port-group-${worker.id}`; + form.appendChild(portGroup.group); + + // CUDA Device field (only for local workers) + const cudaGroup = this.createFormGroup("CUDA Device:", worker.cuda_device || 0, `cuda-${worker.id}`, "number"); + cudaGroup.group.id = `cuda-group-${worker.id}`; + cudaGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "none" : "flex"; + form.appendChild(cudaGroup.group); + + // Extra Args field (only for local workers) + const argsGroup = this.createFormGroup("Extra Args:", worker.extra_args || "", `args-${worker.id}`); + argsGroup.group.id = `args-group-${worker.id}`; + argsGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "none" : "flex"; + form.appendChild(argsGroup.group); + + // Buttons + const saveBtn = this.createButton("Save", + () => extension.saveWorkerSettings(worker.id), + "background-color: #4a7c4a;"); + saveBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.success; + + const cancelBtn = this.createButton("Cancel", + () => extension.cancelWorkerSettings(worker.id), + "background-color: #555;"); + cancelBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.cancel; + + const deleteBtn = this.createButton("Delete", + () => extension.deleteWorker(worker.id), + "background-color: #7c4a4a;"); + deleteBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.error + BUTTON_STYLES.marginLeftAuto; + + const buttonGroup = this.createButtonGroup([saveBtn, cancelBtn, deleteBtn], " margin-top: 8px;"); + form.appendChild(buttonGroup); + + // Assign the onchange handler now that all elements are created + typeSelect.onchange = createOnChangeHandler(); + + // Set initial value and trigger state after all DOM elements are created + if (worker.type === "cloud") { + typeSelect.value = "cloud"; + // Show Runpod text immediately for cloud workers + runpodText.style.display = "block"; + } else if (extension.isRemoteWorker(worker)) { + typeSelect.value = "remote"; + } else { + typeSelect.value = "local"; + } + + // Trigger initial state now that all elements exist + typeSelect.dispatchEvent(new Event('change')); + + return form; + } + + createSettingsToggle() { + const settingsRow = document.createElement("div"); + settingsRow.style.cssText = this.styles.settingsToggle; + + const settingsTitle = document.createElement("h4"); + settingsTitle.textContent = "Settings"; + settingsTitle.style.cssText = "margin: 0; font-size: 14px;"; + + const settingsToggle = document.createElement("span"); + settingsToggle.textContent = "▶"; // Right arrow when collapsed + settingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; + + settingsRow.appendChild(settingsToggle); + settingsRow.appendChild(settingsTitle); + + return { settingsRow, settingsToggle }; + } + + + createCheckboxOrIconColumn(config, data, extension) { + const column = this.createCardColumn('checkbox'); + + if (config?.type === 'icon') { + column.style.flex = `0 0 ${config.width || 44}px`; + column.innerHTML = config.content || '+'; + if (config.style) { + const styles = config.style.split(';').filter(s => s.trim()); + styles.forEach(style => { + const [prop, value] = style.split(':').map(s => s.trim()); + if (prop && value) { + column.style[prop.replace(/-([a-z])/g, (g) => g[1].toUpperCase())] = value; + } + }); + } + } else { + const checkbox = document.createElement("input"); + checkbox.type = "checkbox"; + checkbox.id = `gpu-${data?.id || 'master'}`; + checkbox.checked = config?.checked !== undefined ? config.checked : data?.enabled; + checkbox.disabled = config?.disabled || false; + checkbox.style.cssText = `cursor: ${config?.disabled ? 'default' : 'pointer'}; width: 16px; height: 16px;`; + + if (config?.opacity) checkbox.style.opacity = config.opacity; + if (config?.title) column.title = config.title; + + if (config?.enabled && !config?.disabled && data?.id) { + checkbox.style.pointerEvents = "none"; + column.style.cursor = "pointer"; + column.onclick = async () => { + checkbox.checked = !checkbox.checked; + await extension.updateWorkerEnabled(data.id, checkbox.checked); + extension.updateSummary(); + }; + } + + column.appendChild(checkbox); + } + + return column; + } + + createStatusDotHelper(config, data, extension) { + let color = config.color || "#666"; + let title = config.title || "Status"; + let id = config.id; + + if (typeof config.initialColor === 'function') { + color = config.initialColor(data); + } + if (typeof config.initialTitle === 'function') { + title = config.initialTitle(data); + } + if (typeof config.id === 'function') { + id = config.id(data); + } + + const dot = this.createStatusDot(id, color, title); + + if (config.border) { + dot.style.border = config.border; + } + + if (config.pulsing && (typeof config.pulsing !== 'function' || config.pulsing(data))) { + dot.classList.add('status-pulsing'); + } + + return dot; + } + + createSettingsToggleHelper(expandedId, extension) { + const arrow = document.createElement("span"); + arrow.className = "settings-arrow"; + arrow.innerHTML = "▶"; + arrow.style.cssText = this.styles.settingsArrow; + + const isExpanded = typeof expandedId === 'function' ? + extension.state.isWorkerExpanded(expandedId(extension)) : + (expandedId === 'master' ? false : extension.state.isWorkerExpanded(expandedId)); + + if (isExpanded) { + arrow.style.transform = "rotate(90deg)"; + } + + return arrow; + } + + createControlsSection(config, data, extension, isRemote) { + if (!config) return null; + + const controlsDiv = document.createElement("div"); + controlsDiv.id = `controls-${data?.id || 'master'}`; + controlsDiv.style.cssText = this.styles.controlsDiv; + + // Always create a wrapper div for consistent layout + const controlsWrapper = document.createElement("div"); + controlsWrapper.style.cssText = this.styles.controlsWrapper; + + if (config.dynamic && data) { + if (isRemote) { + const isCloud = data.type === 'cloud'; + const workerTypeText = isCloud ? "Cloud worker" : "Remote worker"; + const remoteInfo = this.createButton(workerTypeText, null, BUTTON_STYLES.info); + remoteInfo.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.info + " color: #999; cursor: default;"; + remoteInfo.disabled = true; + controlsWrapper.appendChild(remoteInfo); + } else { + const controls = this.createWorkerControls(data.id, { + launch: () => extension.launchWorker(data.id), + stop: () => extension.stopWorker(data.id), + viewLog: () => extension.viewWorkerLog(data.id) + }); + + const launchBtn = controls.querySelector(`#launch-${data.id}`); + const stopBtn = controls.querySelector(`#stop-${data.id}`); + const logBtn = controls.querySelector(`#log-${data.id}`); + + launchBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.launch; + launchBtn.title = "Launch worker (runs in background with logging)"; + + stopBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.stop + BUTTON_STYLES.hidden; + stopBtn.title = "Stop worker"; + + logBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.log + BUTTON_STYLES.hidden; + + while (controls.firstChild) { + controlsWrapper.appendChild(controls.firstChild); + } + } + } else if (config.type === 'info') { + const infoBtn = this.createButton(config.text, null, config.style || ""); + infoBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + (config.style || BUTTON_STYLES.info) + " cursor: default;"; + infoBtn.disabled = true; + controlsWrapper.appendChild(infoBtn); + } else if (config.type === 'ghost') { + const ghostBtn = document.createElement("button"); + ghostBtn.style.cssText = `flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; border-radius: 4px; cursor: default; ${config.style || ""}`; + ghostBtn.textContent = config.text; + ghostBtn.disabled = true; + controlsWrapper.appendChild(ghostBtn); + } + + controlsDiv.appendChild(controlsWrapper); + return controlsDiv; + } + + createSettingsSection(config, data, extension) { + const settingsDiv = document.createElement("div"); + const settingsId = typeof config.id === 'function' ? config.id(data) : config.id; + settingsDiv.id = settingsId; + settingsDiv.className = "worker-settings"; + + const expandedId = typeof config.expandedId === 'function' ? config.expandedId(data) : config.expandedId; + const isExpanded = expandedId === 'master-settings' ? false : extension.state.isWorkerExpanded(expandedId); + + settingsDiv.style.cssText = this.styles.workerSettings; + + if (isExpanded) { + settingsDiv.classList.add("expanded"); + settingsDiv.style.padding = "12px"; + settingsDiv.style.marginTop = "8px"; + settingsDiv.style.marginBottom = "8px"; + } + + let settingsForm; + if (config.formType === 'master') { + settingsForm = this.createMasterSettingsForm(extension, data); + } else if (config.formType === 'worker') { + settingsForm = this.createWorkerSettingsForm(extension, data); + } + + if (settingsForm) { + settingsDiv.appendChild(settingsForm); + } + + return settingsDiv; + } + + createMasterSettingsForm(extension, data) { + const settingsForm = document.createElement("div"); + settingsForm.style.cssText = "display: flex; flex-direction: column; gap: 8px;"; + + const nameResult = this.createFormGroup("Name:", extension.config?.master?.name || "Master", "master-name"); + settingsForm.appendChild(nameResult.group); + + const hostResult = this.createFormGroup("Host:", extension.config?.master?.host || "", "master-host", "text", "Auto-detect if empty"); + settingsForm.appendChild(hostResult.group); + + const saveBtn = this.createButton("Save", async () => { + const nameInput = document.getElementById('master-name'); + const hostInput = document.getElementById('master-host'); + + if (!extension.config.master) extension.config.master = {}; + extension.config.master.name = nameInput.value.trim() || "Master"; + + const hostValue = hostInput.value.trim(); + + await extension.api.updateMaster({ + host: hostValue, + name: extension.config.master.name + }); + + // Reload config to refresh any updated values + await extension.loadConfig(); + + // If host was emptied, trigger auto-detection + if (!hostValue) { + extension.log("Host field cleared, triggering IP auto-detection", "debug"); + await extension.detectMasterIP(); + // Reload config again to get the auto-detected IP + await extension.loadConfig(); + // Update the input field with the detected IP + document.getElementById('master-host').value = extension.config?.master?.host || ""; + } + + document.getElementById('master-name-display').textContent = extension.config.master.name; + this.updateMasterDisplay(extension); + + // Show toast notification + if (extension.app?.extensionManager?.toast) { + const message = !hostValue ? + "Master settings saved and IP auto-detected" : + "Master settings saved successfully"; + extension.app.extensionManager.toast.add({ + severity: "success", + summary: "Master Updated", + detail: message, + life: 3000 + }); + } + + saveBtn.textContent = "Saved!"; + setTimeout(() => { saveBtn.textContent = "Save"; }, TIMEOUTS.FLASH_LONG); + }, "background-color: #4a7c4a;"); + saveBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.success; + + const cancelBtn = this.createButton("Cancel", () => { + document.getElementById('master-name').value = extension.config?.master?.name || "Master"; + document.getElementById('master-host').value = extension.config?.master?.host || ""; + }, "background-color: #555;"); + cancelBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.cancel; + + const buttonGroup = this.createButtonGroup([saveBtn, cancelBtn], " margin-top: 8px;"); + settingsForm.appendChild(buttonGroup); + + return settingsForm; + } + + addPlaceholderHover(card, leftColumn, entityType) { + card.onmouseover = () => { + if (entityType === 'blueprint') { + card.style.borderColor = "#777"; + card.style.backgroundColor = "rgba(255, 255, 255, 0.05)"; + leftColumn.style.color = "#999"; + } else { + card.style.borderColor = "#666"; + card.style.backgroundColor = "rgba(255, 255, 255, 0.02)"; + leftColumn.style.color = "#888"; + leftColumn.style.borderColor = "#666"; + } + }; + + card.onmouseout = () => { + if (entityType === 'blueprint') { + card.style.borderColor = "#555"; + card.style.backgroundColor = "rgba(255, 255, 255, 0.02)"; + leftColumn.style.color = "#777"; + } else { + card.style.borderColor = "#444"; + card.style.backgroundColor = "transparent"; + leftColumn.style.color = "#555"; + leftColumn.style.borderColor = "#444"; + } + }; + } + + renderEntityCard(entityType, data, extension) { + const config = cardConfigs[entityType] || {}; + const isPlaceholder = entityType === 'blueprint' || entityType === 'add'; + const isWorker = entityType === 'worker'; + const isMaster = entityType === 'master'; + const isRemote = isWorker && extension.isRemoteWorker(data); + + const cardOptions = { + onClick: isPlaceholder ? data?.onClick : null + }; + if (isPlaceholder) { + cardOptions.title = entityType === 'blueprint' ? "Click to add your first worker" : "Click to add a new worker"; + } + const card = this.createCard(entityType, cardOptions); + + const leftColumn = this.createCheckboxOrIconColumn(config.checkbox, data, extension); + card.appendChild(leftColumn); + + const rightColumn = this.createCardColumn('content'); + + const infoRow = this.createInfoRow(); + if (config.infoRowPadding) { + infoRow.style.padding = config.infoRowPadding; + } + if (config.minHeight === 'auto') { + infoRow.style.minHeight = 'auto'; + } else if (config.minHeight) { + infoRow.style.minHeight = config.minHeight; + } + if (config.expand) { + infoRow.title = "Click to expand settings"; + infoRow.onclick = () => { + if (isMaster) { + const masterSettingsExpanded = !extension.masterSettingsExpanded; + extension.masterSettingsExpanded = masterSettingsExpanded; + const masterSettingsDiv = document.getElementById("master-settings"); + const arrow = infoRow.querySelector('.settings-arrow'); + if (masterSettingsExpanded) { + masterSettingsDiv.classList.add("expanded"); + masterSettingsDiv.style.padding = "12px"; + masterSettingsDiv.style.marginTop = "8px"; + masterSettingsDiv.style.marginBottom = "8px"; + arrow.style.transform = "rotate(90deg)"; + } else { + masterSettingsDiv.classList.remove("expanded"); + masterSettingsDiv.style.padding = "0 12px"; + masterSettingsDiv.style.marginTop = "0"; + masterSettingsDiv.style.marginBottom = "0"; + arrow.style.transform = "rotate(0deg)"; + } + } else { + extension.toggleWorkerExpanded(data.id); + } + }; + } + + const workerContent = this.createWorkerContent(); + if (entityType === 'add') { + workerContent.style.alignItems = "center"; + } + + const statusDot = this.createStatusDotHelper(config.statusDot, data, extension); + workerContent.appendChild(statusDot); + + const infoSpan = document.createElement("span"); + infoSpan.innerHTML = config.infoText(data, extension); + workerContent.appendChild(infoSpan); + + infoRow.appendChild(workerContent); + + let settingsArrow; + if (config.expand) { + const expandedId = config.settings?.expandedId || (isMaster ? 'master' : data?.id); + settingsArrow = this.createSettingsToggleHelper(expandedId, extension); + if (isMaster && !extension.masterSettingsExpanded) { + settingsArrow.style.transform = "rotate(0deg)"; + } + infoRow.appendChild(settingsArrow); + } + + rightColumn.appendChild(infoRow); + + if (config.hover === true) { + rightColumn.onmouseover = () => { + rightColumn.style.backgroundColor = "#333"; + if (settingsArrow) settingsArrow.style.color = "#fff"; + }; + rightColumn.onmouseout = () => { + rightColumn.style.backgroundColor = "transparent"; + if (settingsArrow) settingsArrow.style.color = "#888"; + }; + } + + const controlsDiv = this.createControlsSection(config.controls, data, extension, isRemote); + if (controlsDiv) { + rightColumn.appendChild(controlsDiv); + } + + if (config.settings) { + const settingsDiv = this.createSettingsSection(config.settings, data, extension); + rightColumn.appendChild(settingsDiv); + } + + card.appendChild(rightColumn); + + if (config.hover === 'placeholder') { + this.addPlaceholderHover(card, leftColumn, entityType); + } + + if (isWorker && !isRemote) { + extension.updateWorkerControls(data.id); + } + + return card; + } +} \ No newline at end of file diff --git a/web/workerUtils.js b/web/workerUtils.js new file mode 100644 index 0000000..ef87542 --- /dev/null +++ b/web/workerUtils.js @@ -0,0 +1,325 @@ +import { BUTTON_STYLES, TIMEOUTS } from './constants.js'; + +export async function handleWorkerOperation(extension, button, operation, successText, errorText) { + const originalText = button.textContent; + const originalStyle = button.style.cssText; + + button.textContent = operation.loadingText; + button.disabled = true; + + try { + const urlsToProcess = extension.enabledWorkers.map(w => ({ + name: w.name, + url: extension.getWorkerUrl(w) + })); + + if (urlsToProcess.length === 0) { + button.textContent = "No Workers"; + button.style.backgroundColor = "#c04c4c"; + setTimeout(() => { + button.textContent = originalText; + button.style.cssText = originalStyle; + button.disabled = false; + }, TIMEOUTS.BUTTON_RESET); + return; + } + + const promises = urlsToProcess.map(target => + fetch(`${target.url}${operation.endpoint}`, { + 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 = successText; + button.style.backgroundColor = BUTTON_STYLES.success.split(':')[1].trim().replace(';', ''); + if (operation.onSuccess) operation.onSuccess(); + } else { + button.textContent = errorText; + button.style.backgroundColor = BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); + extension.log(`${operation.name} failed on: ${failures.map(f => f.name).join(", ")}`, "error"); + } + + setTimeout(() => { + button.textContent = originalText; + button.style.cssText = originalStyle; + }, TIMEOUTS.BUTTON_RESET); + } finally { + button.disabled = false; + } +} + +export async function handleInterruptWorkers(extension, button) { + return handleWorkerOperation(extension, button, { + name: "Interrupt", + endpoint: "/interrupt", + loadingText: "Interrupting...", + onSuccess: () => setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY) + }, "Interrupted!", "Error! See Console"); +} + +export async function handleClearMemory(extension, button) { + return handleWorkerOperation(extension, button, { + name: "Clear memory", + endpoint: "/distributed/clear_memory", + loadingText: "Clearing..." + }, "Success!", "Error! See Console"); +} + +export function findNodesByClass(apiPrompt, className) { + return Object.entries(apiPrompt) + .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 + */ +export function findImageReferences(extension, 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]" + // - 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) { + 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 = mediaValue.replace(/\s*\[\w+\]$/, '').trim(); + if (imageExtensions.test(cleanValue)) { + images.set(cleanValue, { + nodeId, + nodeType: node.class_type, + inputName: 'image' // Keep as 'image' for compatibility + }); + extension.log(`Found media reference: ${cleanValue} in node ${nodeId} (${node.class_type})`, "debug"); + } + } + } + } + + return images; +} + +/** + * Find only upstream nodes (inputs) for distributed collector nodes + * This is used for workers to avoid executing downstream nodes like SaveImage + * @param {Object} apiPrompt - The API prompt containing the workflow + * @param {Array} collectorIds - Array of collector node IDs + * @returns {Set} Set of node IDs that feed into collectors + */ +export function findCollectorUpstreamNodes(apiPrompt, collectorIds) { + const connected = new Set(collectorIds); // Include all collectors + const toProcess = [...collectorIds]; + + // Only traverse upstream (inputs) + while (toProcess.length > 0) { + const nodeId = toProcess.pop(); + const node = apiPrompt[nodeId]; + + // Traverse upstream (inputs) only + 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); + } + } + } + } + } + + return connected; +} + +/** + * Prune workflow to only include nodes connected to distributed nodes + * @param {Object} apiPrompt - The full workflow API prompt + * @param {Array} distributedNodes - Array of distributed nodes (optional, will find if not provided) + * @returns {Object} Pruned API prompt with only required nodes + */ +export function pruneWorkflowForWorker(extension, apiPrompt, distributedNodes = null) { + // Find all distributed nodes if not provided + if (!distributedNodes) { + const collectorNodes = findNodesByClass(apiPrompt, "DistributedCollector"); + const upscaleNodes = findNodesByClass(apiPrompt, "UltimateSDUpscaleDistributed"); + distributedNodes = [...collectorNodes, ...upscaleNodes]; + } + + if (distributedNodes.length === 0) { + // No distributed nodes, return full workflow + return apiPrompt; + } + + // Get all nodes connected to distributed nodes + const distributedIds = distributedNodes.map(node => node.id); + + // For workers, only include upstream nodes (this removes ALL downstream nodes after collectors) + const connectedNodes = findCollectorUpstreamNodes(apiPrompt, distributedIds); + + extension.log(`Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes (removed all downstream nodes)`, "debug"); + + // Create pruned prompt with only required nodes + const prunedPrompt = {}; + for (const nodeId of connectedNodes) { + prunedPrompt[nodeId] = JSON.parse(JSON.stringify(apiPrompt[nodeId])); + } + + // Check if any distributed node has downstream SaveImage nodes that were removed + // If so, add a PreviewImage node after the collector + for (const distNode of distributedNodes) { + const distNodeId = distNode.id; + + // Check if this distributed node had any downstream nodes in the original workflow + const originalOutputMap = new Map(); + 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 && String(inputValue[0]) === distNodeId) { + if (!originalOutputMap.has(distNodeId)) { + originalOutputMap.set(distNodeId, []); + } + originalOutputMap.get(distNodeId).push({nodeId, inputName}); + } + } + } + } + + // If this distributed node had downstream nodes that were removed, add a PreviewImage + if (originalOutputMap.has(distNodeId) && originalOutputMap.get(distNodeId).length > 0) { + // Generate unique numeric ID: max existing numeric key +1 + const existingIds = Object.keys(prunedPrompt) + .filter(k => !isNaN(parseInt(k))) + .map(k => parseInt(k)); + const maxId = existingIds.length > 0 ? Math.max(...existingIds) : 0; + const previewNodeId = String(maxId + 1); + + // Add PreviewImage node connected to the distributed node + prunedPrompt[previewNodeId] = { + inputs: { + images: [distNodeId, 0] // Connect to first output of distributed node + }, + class_type: "PreviewImage", + _meta: { + title: "Preview Image (auto-added)" + } + }; + + extension.log(`Added PreviewImage node ${previewNodeId} after distributed node ${distNodeId} for worker`, "debug"); + } + } + + return prunedPrompt; +} + +/** + * Check if a node has an upstream node of a specific type + * @param {Object} apiPrompt - The workflow API prompt + * @param {string} nodeId - The node to check + * @param {string} upstreamType - The class_type to look for upstream + * @returns {boolean} True if an upstream node of the specified type exists + */ +export function hasUpstreamNode(apiPrompt, nodeId, upstreamType) { + const visited = new Set(); + const toProcess = [nodeId]; + + while (toProcess.length > 0) { + const currentId = toProcess.pop(); + if (visited.has(currentId)) continue; + visited.add(currentId); + + const node = apiPrompt[currentId]; + if (!node) continue; + + // Check inputs for upstream connections + if (node.inputs) { + for (const [inputName, inputValue] of Object.entries(node.inputs)) { + if (Array.isArray(inputValue) && inputValue.length === 2) { + const sourceNodeId = String(inputValue[0]); + const sourceNode = apiPrompt[sourceNodeId]; + + if (sourceNode && sourceNode.class_type === upstreamType) { + return true; + } + + if (!visited.has(sourceNodeId)) { + toProcess.push(sourceNodeId); + } + } + } + } + } + + return false; +} + +/** + * Get system information from a worker + * @param {string} workerUrl - The worker URL + * @returns {Promise} System information including platform details + */ +export async function getWorkerSystemInfo(workerUrl) { + try { + const response = await fetch(`${workerUrl}/distributed/system_info`); + if (!response.ok) { + throw new Error(`HTTP ${response.status}`); + } + return await response.json(); + } catch (error) { + console.warn(`Failed to get system info from ${workerUrl}:`, error); + // Return sensible defaults + return { + platform: { + os_name: 'posix', // Assume Linux + path_separator: '/', + system: 'Linux' + } + }; + } +} + +// Cache system info to avoid repeated calls +const systemInfoCache = new Map(); + +/** + * Get cached system information from a worker + * @param {string} workerUrl - The worker URL + * @returns {Promise} Cached or fresh system information + */ +export async function getCachedWorkerSystemInfo(workerUrl) { + if (systemInfoCache.has(workerUrl)) { + return systemInfoCache.get(workerUrl); + } + + const info = await getWorkerSystemInfo(workerUrl); + systemInfoCache.set(workerUrl, info); + return info; +} + +/** + * Clear the system info cache + */ +export function clearSystemInfoCache() { + systemInfoCache.clear(); +} \ No newline at end of file diff --git a/workflows/distributed-upscale-batch.json b/workflows/distributed-upscale-batch.json new file mode 100644 index 0000000..89f6871 --- /dev/null +++ b/workflows/distributed-upscale-batch.json @@ -0,0 +1,1277 @@ +{ + "id": "6b169846-6204-4336-bda1-f9dc43cd52e7", + "revision": 0, + "last_node_id": 26, + "last_link_id": 26, + "nodes": [ + { + "id": 1, + "type": "Reroute", + "pos": [ + 3355.76220703125, + 1302.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 1 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 6 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 2, + "type": "Reroute", + "pos": [ + 3355.76220703125, + 1252.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 2 + } + ], + "outputs": [ + { + "name": "", + "type": "MODEL", + "links": [ + 4 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 3, + "type": "Reroute", + "pos": [ + 4035.76220703125, + 1302.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 3 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 23 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 4, + "type": "Reroute", + "pos": [ + 4035.76220703125, + 1252.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 4 + } + ], + "outputs": [ + { + "name": "", + "type": "MODEL", + "links": [ + 20 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 5, + "type": "Reroute", + "pos": [ + 4035.76220703125, + 1222.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 5 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 19 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 6, + "type": "CheckpointLoaderSimple", + "pos": [ + 2925.76220703125, + 1522.14697265625 + ], + "size": [ + 370, + 98 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "slot_index": 0, + "links": [ + 2 + ] + }, + { + "name": "CLIP", + "type": "CLIP", + "slot_index": 1, + "links": [ + 7, + 18 + ] + }, + { + "name": "VAE", + "type": "VAE", + "slot_index": 2, + "links": [ + 1 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CheckpointLoaderSimple", + "models": [ + { + "name": "v1-5-pruned-emaonly-fp16.safetensors", + "url": "https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/v1-5-pruned-emaonly-fp16.safetensors?download=true", + "directory": "checkpoints" + } + ] + }, + "widgets_values": [ + "SDXL\\juggernautXL_ragnarokBy.safetensors" + ] + }, + { + "id": 7, + "type": "Reroute", + "pos": [ + 3685.76220703125, + 1302.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 6 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 3, + 15 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 8, + "type": "ControlNetLoader", + "pos": [ + 3135.76220703125, + 1762.14697265625 + ], + "size": [ + 390, + 58 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 8 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetLoader" + }, + "widgets_values": [ + "SDXL\\xinsir-controlnet-union-sdxl-1.0-promax.safetensors" + ] + }, + { + "id": 9, + "type": "CLIPTextEncode", + "pos": [ + 3355.76220703125, + 1622.14697265625 + ], + "size": [ + 420, + 88 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 7 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 12 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "text, watermark" + ] + }, + { + "id": 10, + "type": "SetUnionControlNetType", + "pos": [ + 3565.76220703125, + 1762.14697265625 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "control_net", + "type": "CONTROL_NET", + "link": 8 + } + ], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 13 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "SetUnionControlNetType" + }, + "widgets_values": [ + "tile" + ] + }, + { + "id": 11, + "type": "Reroute", + "pos": [ + 3675.76220703125, + 1132.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 9 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 14 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 12, + "type": "ImageResize+", + "pos": [ + 2925.76220703125, + 1222.14697265625 + ], + "size": [ + 270, + 218 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 10 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 5 + ] + }, + { + "name": "width", + "type": "INT", + "links": null + }, + { + "name": "height", + "type": "INT", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui_essentials", + "ver": "9d9f4bedfc9f0321c19faf71855e228c93bd0dc9", + "Node name for S&R": "ImageResize+" + }, + "widgets_values": [ + 2048, + 2048, + "lanczos", + "keep proportion", + "always", + 8 + ] + }, + { + "id": 13, + "type": "ControlNetApplyAdvanced", + "pos": [ + 3865.76220703125, + 1402.14697265625 + ], + "size": [ + 270, + 186 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "link": 11 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 12 + }, + { + "name": "control_net", + "type": "CONTROL_NET", + "link": 13 + }, + { + "name": "image", + "type": "IMAGE", + "link": 14 + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": 15 + } + ], + "outputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 21 + ] + }, + { + "name": "negative", + "type": "CONDITIONING", + "links": [ + 22 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetApplyAdvanced" + }, + "widgets_values": [ + 1.0000000000000002, + 0, + 0.9000000000000002 + ] + }, + { + "id": 14, + "type": "ImageUpscaleWithModel", + "pos": [ + 2635.76220703125, + 1222.14697265625 + ], + "size": [ + 222.75416564941406, + 46 + ], + "flags": {}, + "order": 17, + "mode": 4, + "inputs": [ + { + "name": "upscale_model", + "type": "UPSCALE_MODEL", + "link": 16 + }, + { + "name": "image", + "type": "IMAGE", + "link": 17 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 10 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ImageUpscaleWithModel" + }, + "widgets_values": [] + }, + { + "id": 15, + "type": "MarkdownNote", + "pos": [ + 3145.76220703125, + 1862.14697265625 + ], + "size": [ + 370, + 88 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "[https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main](https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main)" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 17, + "type": "Note", + "pos": [ + 2685.76220703125, + 1532.14697265625 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Choose an SDXL model" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 18, + "type": "CLIPTextEncode", + "pos": [ + 3355.76220703125, + 1412.14697265625 + ], + "size": [ + 420, + 160 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 18 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 11 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "beautiful scenery nature glass bottle landscape, , purple galaxy bottle," + ] + }, + { + "id": 19, + "type": "UltimateSDUpscaleDistributed", + "pos": [ + 4225.76220703125, + 1362.14697265625 + ], + "size": [ + 326.691650390625, + 450 + ], + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "upscaled_image", + "type": "IMAGE", + "link": 19 + }, + { + "name": "model", + "type": "MODEL", + "link": 20 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 21 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 22 + }, + { + "name": "vae", + "type": "VAE", + "link": 23 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 25 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "dd23503883fdf319e8beb6e7a190445ecf89973c", + "Node name for S&R": "UltimateSDUpscaleDistributed" + }, + "widgets_values": [ + 320425997971252, + "randomize", + 20, + 7, + "dpmpp_2m_sde", + "karras", + 0.7000000000000002, + 1024, + 1024, + 32, + 16, + true, + false, + true + ] + }, + { + "id": 20, + "type": "Reroute", + "pos": [ + 2455.76220703125, + 1132.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 26 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 9, + 17 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 21, + "type": "UpscaleModelLoader", + "pos": [ + 2355.76220703125, + 1322.14697265625 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "UPSCALE_MODEL", + "type": "UPSCALE_MODEL", + "links": [ + 16 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "UpscaleModelLoader" + }, + "widgets_values": [ + "4xNomos8kDAT.pth" + ] + }, + { + "id": 23, + "type": "PreviewImage", + "pos": [ + 4605.76220703125, + 1372.14697265625 + ], + "size": [ + 490, + 550 + ], + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 25 + } + ], + "outputs": [], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 25, + "type": "VHS_LoadVideoPath", + "pos": [ + 1950, + 1130 + ], + "size": [ + 330, + 286 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 26 + ] + }, + { + "name": "frame_count", + "type": "INT", + "links": null + }, + { + "name": "audio", + "type": "AUDIO", + "links": null + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_LoadVideoPath" + }, + "widgets_values": { + "video": "input/test/", + "force_rate": 0, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "format": "None", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "input/test/", + "type": "path", + "format": "video/input/test/", + "force_rate": 0, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1 + } + } + } + }, + { + "id": 22, + "type": "Note", + "pos": [ + 2645.76220703125, + 1322.14697265625 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Optional" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 26, + "type": "Note", + "pos": [ + 1950, + 980 + ], + "size": [ + 320, + 88 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Use this node to load a batch of images from a folder.\n\nNOTE: images have to be the same dimensions." + ], + "color": "#432", + "bgcolor": "#653" + } + ], + "links": [ + [ + 1, + 6, + 2, + 1, + 0, + "*" + ], + [ + 2, + 6, + 0, + 2, + 0, + "*" + ], + [ + 3, + 7, + 0, + 3, + 0, + "*" + ], + [ + 4, + 2, + 0, + 4, + 0, + "*" + ], + [ + 5, + 12, + 0, + 5, + 0, + "*" + ], + [ + 6, + 1, + 0, + 7, + 0, + "*" + ], + [ + 7, + 6, + 1, + 9, + 0, + "CLIP" + ], + [ + 8, + 8, + 0, + 10, + 0, + "CONTROL_NET" + ], + [ + 9, + 20, + 0, + 11, + 0, + "*" + ], + [ + 10, + 14, + 0, + 12, + 0, + "IMAGE" + ], + [ + 11, + 18, + 0, + 13, + 0, + "CONDITIONING" + ], + [ + 12, + 9, + 0, + 13, + 1, + "CONDITIONING" + ], + [ + 13, + 10, + 0, + 13, + 2, + "CONTROL_NET" + ], + [ + 14, + 11, + 0, + 13, + 3, + "IMAGE" + ], + [ + 15, + 7, + 0, + 13, + 4, + "VAE" + ], + [ + 16, + 21, + 0, + 14, + 0, + "UPSCALE_MODEL" + ], + [ + 17, + 20, + 0, + 14, + 1, + "IMAGE" + ], + [ + 18, + 6, + 1, + 18, + 0, + "CLIP" + ], + [ + 19, + 5, + 0, + 19, + 0, + "IMAGE" + ], + [ + 20, + 4, + 0, + 19, + 1, + "MODEL" + ], + [ + 21, + 13, + 0, + 19, + 2, + "CONDITIONING" + ], + [ + 22, + 13, + 1, + 19, + 3, + "CONDITIONING" + ], + [ + 23, + 3, + 0, + 19, + 4, + "VAE" + ], + [ + 25, + 19, + 0, + 23, + 0, + "IMAGE" + ], + [ + 26, + 25, + 0, + 20, + 0, + "*" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.9229599817706583, + "offset": [ + -1960.8728618221426, + -538.777843033621 + ] + }, + "frontendVersion": "1.23.4", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/workflows/distributed-upscale-video.json b/workflows/distributed-upscale-video.json new file mode 100644 index 0000000..fc82bd5 --- /dev/null +++ b/workflows/distributed-upscale-video.json @@ -0,0 +1,1414 @@ +{ + "id": "a778fe2e-2eb6-417f-8b97-86177243bce8", + "revision": 0, + "last_node_id": 140, + "last_link_id": 210, + "nodes": [ + { + "id": 86, + "type": "Reroute", + "pos": [ + 2130, + 1050 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 121 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 127 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 107, + "type": "Reroute", + "pos": [ + 2810, + 1050 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 171 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 172 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 110, + "type": "Reroute", + "pos": [ + 2810, + 970 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 210 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 176 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 7, + "type": "CheckpointLoaderSimple", + "pos": [ + 1700, + 1270 + ], + "size": [ + 370, + 98 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "slot_index": 0, + "links": [ + 124 + ] + }, + { + "name": "CLIP", + "type": "CLIP", + "slot_index": 1, + "links": [ + 5, + 6 + ] + }, + { + "name": "VAE", + "type": "VAE", + "slot_index": 2, + "links": [ + 121 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CheckpointLoaderSimple", + "models": [ + { + "name": "v1-5-pruned-emaonly-fp16.safetensors", + "url": "https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/v1-5-pruned-emaonly-fp16.safetensors?download=true", + "directory": "checkpoints" + } + ] + }, + "widgets_values": [ + "SDXL\\juggernautXL_ragnarokBy.safetensors" + ] + }, + { + "id": 90, + "type": "Reroute", + "pos": [ + 2460, + 1050 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 127 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 129, + 171 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 44, + "type": "ControlNetLoader", + "pos": [ + 1910, + 1510 + ], + "size": [ + 390, + 58 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 81 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetLoader" + }, + "widgets_values": [ + "SDXL\\xinsir-controlnet-union-sdxl-1.0-promax.safetensors" + ] + }, + { + "id": 52, + "type": "SetUnionControlNetType", + "pos": [ + 2340, + 1510 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "control_net", + "type": "CONTROL_NET", + "link": 81 + } + ], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 192 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "SetUnionControlNetType" + }, + "widgets_values": [ + "tile" + ] + }, + { + "id": 125, + "type": "Reroute", + "pos": [ + 2450, + 880 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 200 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 206 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 71, + "type": "RIFE VFI", + "pos": [ + 3367.979248046875, + 1110 + ], + "size": [ + 290.73333740234375, + 198 + ], + "flags": {}, + "order": 24, + "mode": 0, + "inputs": [ + { + "name": "frames", + "type": "IMAGE", + "link": 108 + }, + { + "name": "optional_interpolation_states", + "shape": 7, + "type": "INTERPOLATION_STATES", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 109 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-frame-interpolation", + "ver": "a969c01dbccd9e5510641be04eb51fe93f6bfc3d", + "Node name for S&R": "RIFE VFI" + }, + "widgets_values": [ + "rife47.pth", + 10, + 2, + true, + true, + 1 + ] + }, + { + "id": 34, + "type": "ImageResize+", + "pos": [ + 1700, + 970 + ], + "size": [ + 270, + 218 + ], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 197 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 210 + ] + }, + { + "name": "width", + "type": "INT", + "links": null + }, + { + "name": "height", + "type": "INT", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui_essentials", + "ver": "9d9f4bedfc9f0321c19faf71855e228c93bd0dc9", + "Node name for S&R": "ImageResize+" + }, + "widgets_values": [ + 2048, + 2048, + "lanczos", + "keep proportion", + "always", + 8 + ] + }, + { + "id": 43, + "type": "ControlNetApplyAdvanced", + "pos": [ + 2640, + 1150 + ], + "size": [ + 270, + 186 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "link": 63 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 64 + }, + { + "name": "control_net", + "type": "CONTROL_NET", + "link": 192 + }, + { + "name": "image", + "type": "IMAGE", + "link": 206 + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": 129 + } + ], + "outputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 190 + ] + }, + { + "name": "negative", + "type": "CONDITIONING", + "links": [ + 191 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetApplyAdvanced" + }, + "widgets_values": [ + 1.0000000000000002, + 0, + 0.9000000000000002 + ] + }, + { + "id": 123, + "type": "ImageUpscaleWithModel", + "pos": [ + 1410, + 970 + ], + "size": [ + 222.75416564941406, + 46 + ], + "flags": {}, + "order": 18, + "mode": 4, + "inputs": [ + { + "name": "upscale_model", + "type": "UPSCALE_MODEL", + "link": 195 + }, + { + "name": "image", + "type": "IMAGE", + "link": 209 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 197 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ImageUpscaleWithModel" + }, + "widgets_values": [] + }, + { + "id": 131, + "type": "Note", + "pos": [ + 1430, + 1070 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Optional" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 88, + "type": "Reroute", + "pos": [ + 2130, + 1020 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 124 + } + ], + "outputs": [ + { + "name": "", + "type": "MODEL", + "links": [ + 137 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 89, + "type": "Reroute", + "pos": [ + 2810, + 1020 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 137 + } + ], + "outputs": [ + { + "name": "", + "type": "MODEL", + "links": [ + 170 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 124, + "type": "Reroute", + "pos": [ + 1230, + 880 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 208 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 200, + 209 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 132, + "type": "Note", + "pos": [ + 1460, + 1270 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Choose an SDXL model" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": [ + 2140, + 1360 + ], + "size": [ + 420, + 88 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 6 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 64 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "text, watermark" + ] + }, + { + "id": 41, + "type": "VHS_VideoCombine", + "pos": [ + 3700, + 1110 + ], + "size": [ + 450, + 594.9140625 + ], + "flags": {}, + "order": 25, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 109 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 30, + "loop_count": 0, + "filename_prefix": "AnimateDiff", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": false, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00231.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 30, + "workflow": "AnimateDiff_00231.png", + "fullpath": "C:\\venvs\\ComfyUI\\ComfyUI\\output\\AnimateDiff_00231.mp4" + } + } + } + }, + { + "id": 134, + "type": "MarkdownNote", + "pos": [ + 1910, + 1610 + ], + "size": [ + 380, + 88 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "[https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main](https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main)" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 122, + "type": "UpscaleModelLoader", + "pos": [ + 1120, + 1040 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "UPSCALE_MODEL", + "type": "UPSCALE_MODEL", + "links": [ + 195 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "UpscaleModelLoader" + }, + "widgets_values": [ + "4xNomos8kDAT.pth" + ] + }, + { + "id": 5, + "type": "CLIPTextEncode", + "pos": [ + 2140, + 1150 + ], + "size": [ + 420, + 160 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 5 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 63 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "beautiful scenery nature glass bottle landscape, , purple galaxy bottle," + ] + }, + { + "id": 30, + "type": "UltimateSDUpscaleDistributed", + "pos": [ + 3000, + 1110 + ], + "size": [ + 326.691650390625, + 450 + ], + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "upscaled_image", + "type": "IMAGE", + "link": 176 + }, + { + "name": "model", + "type": "MODEL", + "link": 170 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 190 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 191 + }, + { + "name": "vae", + "type": "VAE", + "link": 172 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 108 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "dd23503883fdf319e8beb6e7a190445ecf89973c", + "Node name for S&R": "UltimateSDUpscaleDistributed" + }, + "widgets_values": [ + 320425997971252, + "randomize", + 20, + 7, + "euler", + "normal", + 0.7000000000000002, + 1024, + 1024, + 32, + 16, + true, + false, + true + ] + }, + { + "id": 139, + "type": "Note", + "pos": [ + 3000, + 1610 + ], + "size": [ + 330, + 88 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "If all your GPUs are the same/similar, set static_distribution to true\n" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 138, + "type": "VHS_LoadVideo", + "pos": [ + 670, + 880 + ], + "size": [ + 340, + 509 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 208 + ] + }, + { + "name": "frame_count", + "type": "INT", + "links": null + }, + { + "name": "audio", + "type": "AUDIO", + "links": null + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_LoadVideo" + }, + "widgets_values": { + "video": "AnimateDiff_00155.mp4", + "force_rate": 16, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "format": "Wan", + "choose video to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00155.mp4", + "type": "input", + "format": "video/mp4", + "force_rate": 16, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1 + } + } + } + }, + { + "id": 140, + "type": "MarkdownNote", + "pos": [ + 680, + 680 + ], + "size": [ + 390, + 140 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "# WARNING: Experimental workflow\n\nThe tiled upscaling works with video, but it is experimental.\n\nThe results can be flickery." + ], + "color": "#432", + "bgcolor": "#653" + } + ], + "links": [ + [ + 5, + 7, + 1, + 5, + 0, + "CLIP" + ], + [ + 6, + 7, + 1, + 6, + 0, + "CLIP" + ], + [ + 63, + 5, + 0, + 43, + 0, + "CONDITIONING" + ], + [ + 64, + 6, + 0, + 43, + 1, + "CONDITIONING" + ], + [ + 81, + 44, + 0, + 52, + 0, + "CONTROL_NET" + ], + [ + 108, + 30, + 0, + 71, + 0, + "IMAGE" + ], + [ + 109, + 71, + 0, + 41, + 0, + "IMAGE" + ], + [ + 121, + 7, + 2, + 86, + 0, + "*" + ], + [ + 124, + 7, + 0, + 88, + 0, + "*" + ], + [ + 127, + 86, + 0, + 90, + 0, + "*" + ], + [ + 129, + 90, + 0, + 43, + 4, + "VAE" + ], + [ + 137, + 88, + 0, + 89, + 0, + "*" + ], + [ + 170, + 89, + 0, + 30, + 1, + "MODEL" + ], + [ + 171, + 90, + 0, + 107, + 0, + "*" + ], + [ + 172, + 107, + 0, + 30, + 4, + "VAE" + ], + [ + 176, + 110, + 0, + 30, + 0, + "IMAGE" + ], + [ + 190, + 43, + 0, + 30, + 2, + "CONDITIONING" + ], + [ + 191, + 43, + 1, + 30, + 3, + "CONDITIONING" + ], + [ + 192, + 52, + 0, + 43, + 2, + "CONTROL_NET" + ], + [ + 195, + 122, + 0, + 123, + 0, + "UPSCALE_MODEL" + ], + [ + 197, + 123, + 0, + 34, + 0, + "IMAGE" + ], + [ + 200, + 124, + 0, + 125, + 0, + "*" + ], + [ + 206, + 125, + 0, + 43, + 3, + "IMAGE" + ], + [ + 208, + 138, + 0, + 124, + 0, + "*" + ], + [ + 209, + 124, + 0, + 123, + 1, + "IMAGE" + ], + [ + 210, + 34, + 0, + 110, + 0, + "*" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.6934334949441466, + "offset": [ + -745.9669584283469, + -262.3877185095017 + ] + }, + "frontendVersion": "1.23.4", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/workflows/distributed-upscale.json b/workflows/distributed-upscale.json index 339a7df..044a9f9 100644 --- a/workflows/distributed-upscale.json +++ b/workflows/distributed-upscale.json @@ -1,317 +1,197 @@ { - "id": "0a68c65f-df29-492e-8525-7621fada0ddd", + "id": "817bbfe2-06b8-44c8-8c14-b82b63b335d5", "revision": 0, - "last_node_id": 17, - "last_link_id": 23, + "last_node_id": 137, + "last_link_id": 211, "nodes": [ { - "id": 17, - "type": "ImageScale", + "id": 86, + "type": "Reroute", "pos": [ - 1540, + 2130, 1050 ], "size": [ - 270, - 130 + 75, + 26 ], "flags": {}, - "order": 6, + "order": 11, "mode": 0, "inputs": [ { - "name": "image", - "type": "IMAGE", - "link": 22 + "name": "", + "type": "*", + "link": 121 } ], "outputs": [ { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 23 - ] - } - ], - "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "cnr_id": "comfy-core", - "ver": "0.3.43", - "Node name for S&R": "ImageScale" - }, - "widgets_values": [ - "lanczos", - 2048, - 2048, - "disabled" - ] - }, - { - "id": 13, - "type": "UltimateSDUpscaleDistributed", - "pos": [ - 1890, - 660 - ], - "size": [ - 326.691650390625, - 402 - ], - "flags": {}, - "order": 7, - "mode": 0, - "inputs": [ - { - "name": "upscaled_image", - "type": "IMAGE", - "link": 23 - }, - { - "name": "model", - "type": "MODEL", - "link": 15 - }, - { - "name": "positive", - "type": "CONDITIONING", - "link": 16 - }, - { - "name": "negative", - "type": "CONDITIONING", - "link": 17 - }, - { - "name": "vae", + "name": "", "type": "VAE", - "link": 18 - } - ], - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", "links": [ - 12 + 127 ] } ], "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "aux_id": "robertvoy/ComfyUI-Distributed", - "ver": "99021363d65cc2b2f0f3a0f12a76a358f0fb330f", - "Node name for S&R": "UltimateSDUpscaleDistributed" - }, - "widgets_values": [ - 402634549992051, - "randomize", - 10, - 8, - "euler", - "normal", - 0.24000000000000005, - 512, - 512, - 32, - 16, - true - ] + "showOutputText": false, + "horizontal": false + } }, { - "id": 10, - "type": "PreviewImage", + "id": 88, + "type": "Reroute", "pos": [ - 2260, - 680 + 2130, + 1000 ], "size": [ - 410, - 410 + 75, + 26 ], "flags": {}, "order": 8, "mode": 0, "inputs": [ { - "name": "images", - "type": "IMAGE", - "link": 12 + "name": "", + "type": "*", + "link": 124 + } + ], + "outputs": [ + { + "name": "", + "type": "MODEL", + "links": [ + 137 + ] } ], - "outputs": [], "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "cnr_id": "comfy-core", - "ver": "0.3.43", - "Node name for S&R": "PreviewImage" - }, - "widgets_values": [] + "showOutputText": false, + "horizontal": false + } }, { - "id": 5, - "type": "CLIPTextEncode", + "id": 107, + "type": "Reroute", "pos": [ - 1390, - 540 + 2810, + 1050 ], "size": [ - 422.84503173828125, - 164.31304931640625 + 75, + 26 ], "flags": {}, - "order": 4, + "order": 18, "mode": 0, "inputs": [ { - "name": "clip", - "type": "CLIP", - "link": 5 + "name": "", + "type": "*", + "link": 171 } ], "outputs": [ { - "name": "CONDITIONING", - "type": "CONDITIONING", - "slot_index": 0, + "name": "", + "type": "VAE", "links": [ - 16 + 172 ] } ], "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "cnr_id": "comfy-core", - "ver": "0.3.41", - "Node name for S&R": "CLIPTextEncode" - }, - "widgets_values": [ - "beautiful scenery nature glass bottle landscape, , purple galaxy bottle," - ] + "showOutputText": false, + "horizontal": false + } }, { - "id": 14, - "type": "UpscaleModelLoader", + "id": 89, + "type": "Reroute", "pos": [ - 970, - 950 + 2810, + 1000 ], "size": [ - 270, - 58 + 75, + 26 ], "flags": {}, - "order": 0, + "order": 14, "mode": 0, - "inputs": [], + "inputs": [ + { + "name": "", + "type": "*", + "link": 137 + } + ], "outputs": [ { - "name": "UPSCALE_MODEL", - "type": "UPSCALE_MODEL", + "name": "", + "type": "MODEL", "links": [ - 19 + 170 ] } ], "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "cnr_id": "comfy-core", - "ver": "0.3.43", - "Node name for S&R": "UpscaleModelLoader" - }, - "widgets_values": [ - "4xNomos8kDAT.pth" - ] + "showOutputText": false, + "horizontal": false + } }, { - "id": 12, - "type": "LoadImage", + "id": 110, + "type": "Reroute", "pos": [ - 960, - 1080 + 2810, + 970 ], "size": [ - 274.375, - 314.00006103515625 + 75, + 26 ], "flags": {}, - "order": 1, + "order": 21, "mode": 0, - "inputs": [], + "inputs": [ + { + "name": "", + "type": "*", + "link": 211 + } + ], "outputs": [ { - "name": "IMAGE", + "name": "", "type": "IMAGE", "links": [ - 20 + 176 ] - }, - { - "name": "MASK", - "type": "MASK", - "links": null } ], "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "cnr_id": "comfy-core", - "ver": "0.3.43", - "Node name for S&R": "LoadImage" - }, - "widgets_values": [ - "example.png", - "image" - ] + "showOutputText": false, + "horizontal": false + } }, { "id": 7, "type": "CheckpointLoaderSimple", "pos": [ - 960, - 700 + 1700, + 1270 ], "size": [ - 315, + 370, 98 ], "flags": {}, - "order": 2, + "order": 0, "mode": 0, "inputs": [], "outputs": [ @@ -320,7 +200,7 @@ "type": "MODEL", "slot_index": 0, "links": [ - 15 + 124 ] }, { @@ -337,7 +217,7 @@ "type": "VAE", "slot_index": 2, "links": [ - 18 + 121 ] } ], @@ -364,19 +244,95 @@ "SDXL\\juggernautXL_ragnarokBy.safetensors" ] }, + { + "id": 90, + "type": "Reroute", + "pos": [ + 2460, + 1050 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 127 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 129, + 171 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 44, + "type": "ControlNetLoader", + "pos": [ + 1910, + 1510 + ], + "size": [ + 390, + 58 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 81 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetLoader" + }, + "widgets_values": [ + "SDXL\\xinsir-controlnet-union-sdxl-1.0-promax.safetensors" + ] + }, { "id": 6, "type": "CLIPTextEncode", "pos": [ - 1390, - 760 + 2130, + 1370 ], "size": [ - 425.27801513671875, - 180.6060791015625 + 420, + 88 ], "flags": {}, - "order": 5, + "order": 10, "mode": 0, "inputs": [ { @@ -391,7 +347,7 @@ "type": "CONDITIONING", "slot_index": 0, "links": [ - 17 + 64 ] } ], @@ -412,29 +368,105 @@ ] }, { - "id": 15, - "type": "ImageUpscaleWithModel", + "id": 52, + "type": "SetUnionControlNetType", "pos": [ - 1280, - 1050 + 2340, + 1510 ], "size": [ - 221.96665954589844, - 46 + 210, + 58 ], "flags": {}, - "order": 3, + "order": 12, "mode": 0, "inputs": [ { - "name": "upscale_model", - "type": "UPSCALE_MODEL", - "link": 19 - }, + "name": "control_net", + "type": "CONTROL_NET", + "link": 81 + } + ], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 192 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "SetUnionControlNetType" + }, + "widgets_values": [ + "tile" + ] + }, + { + "id": 125, + "type": "Reroute", + "pos": [ + 2450, + 880 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 200 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 206 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 34, + "type": "ImageResize+", + "pos": [ + 1700, + 970 + ], + "size": [ + 270, + 218 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ { "name": "image", "type": "IMAGE", - "link": 20 + "link": 197 } ], "outputs": [ @@ -442,7 +474,73 @@ "name": "IMAGE", "type": "IMAGE", "links": [ - 22 + 211 + ] + }, + { + "name": "width", + "type": "INT", + "links": null + }, + { + "name": "height", + "type": "INT", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui_essentials", + "ver": "9d9f4bedfc9f0321c19faf71855e228c93bd0dc9", + "Node name for S&R": "ImageResize+" + }, + "widgets_values": [ + 2048, + 2048, + "lanczos", + "keep proportion", + "always", + 8 + ] + }, + { + "id": 123, + "type": "ImageUpscaleWithModel", + "pos": [ + 1410, + 970 + ], + "size": [ + 222.75416564941406, + 46 + ], + "flags": {}, + "order": 17, + "mode": 4, + "inputs": [ + { + "name": "upscale_model", + "type": "UPSCALE_MODEL", + "link": 195 + }, + { + "name": "image", + "type": "IMAGE", + "link": 210 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 197 ] } ], @@ -457,7 +555,460 @@ "cnr_id": "comfy-core", "ver": "0.3.43", "Node name for S&R": "ImageUpscaleWithModel" + }, + "widgets_values": [] + }, + { + "id": 134, + "type": "MarkdownNote", + "pos": [ + 1920, + 1610 + ], + "size": [ + 370, + 88 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "[https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main](https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main)" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 135, + "type": "LoadImage", + "pos": [ + 780, + 880 + ], + "size": [ + 274.375, + 314.00006103515625 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 207 + ] + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "example.png", + "image" + ] + }, + { + "id": 132, + "type": "Note", + "pos": [ + 1460, + 1280 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Choose an SDXL model" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 5, + "type": "CLIPTextEncode", + "pos": [ + 2130, + 1160 + ], + "size": [ + 420, + 160 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 5 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 63 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "beautiful scenery nature glass bottle landscape, , purple galaxy bottle," + ] + }, + { + "id": 124, + "type": "Reroute", + "pos": [ + 1230, + 880 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 207 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 200, + 210 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false } + }, + { + "id": 122, + "type": "UpscaleModelLoader", + "pos": [ + 1130, + 1070 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "UPSCALE_MODEL", + "type": "UPSCALE_MODEL", + "links": [ + 195 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "UpscaleModelLoader" + }, + "widgets_values": [ + "4xNomos8kDAT.pth" + ] + }, + { + "id": 131, + "type": "Note", + "pos": [ + 1420, + 1070 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Optional" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 137, + "type": "Note", + "pos": [ + 3000, + 1600 + ], + "size": [ + 330, + 90 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "If all your GPUs are the same/similar, set static_distribution to true\n" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 136, + "type": "PreviewImage", + "pos": [ + 3380, + 1100 + ], + "size": [ + 490, + 550 + ], + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 209 + } + ], + "outputs": [], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 43, + "type": "ControlNetApplyAdvanced", + "pos": [ + 2640, + 1150 + ], + "size": [ + 270, + 186 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "link": 63 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 64 + }, + { + "name": "control_net", + "type": "CONTROL_NET", + "link": 192 + }, + { + "name": "image", + "type": "IMAGE", + "link": 206 + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": 129 + } + ], + "outputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 190 + ] + }, + { + "name": "negative", + "type": "CONDITIONING", + "links": [ + 191 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetApplyAdvanced" + }, + "widgets_values": [ + 1.0000000000000002, + 0, + 0.8000000000000002 + ] + }, + { + "id": 30, + "type": "UltimateSDUpscaleDistributed", + "pos": [ + 3000, + 1110 + ], + "size": [ + 326.691650390625, + 450 + ], + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "upscaled_image", + "type": "IMAGE", + "link": 176 + }, + { + "name": "model", + "type": "MODEL", + "link": 170 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 190 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 191 + }, + { + "name": "vae", + "type": "VAE", + "link": 172 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 209 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "dd23503883fdf319e8beb6e7a190445ecf89973c", + "Node name for S&R": "UltimateSDUpscaleDistributed" + }, + "widgets_values": [ + 269777990474642, + "randomize", + 20, + 7, + "dpmpp_2m_sde", + "karras", + 0.6000000000000001, + 1024, + 1024, + 32, + 16, + true, + false, + true + ] } ], "links": [ @@ -478,89 +1029,205 @@ "CLIP" ], [ - 12, - 13, - 0, - 10, - 0, - "IMAGE" - ], - [ - 15, - 7, - 0, - 13, - 1, - "MODEL" - ], - [ - 16, + 63, 5, 0, - 13, - 2, + 43, + 0, "CONDITIONING" ], [ - 17, + 64, 6, 0, - 13, - 3, + 43, + 1, "CONDITIONING" ], [ - 18, + 81, + 44, + 0, + 52, + 0, + "CONTROL_NET" + ], + [ + 121, 7, 2, - 13, + 86, + 0, + "*" + ], + [ + 124, + 7, + 0, + 88, + 0, + "*" + ], + [ + 127, + 86, + 0, + 90, + 0, + "*" + ], + [ + 129, + 90, + 0, + 43, 4, "VAE" ], [ - 19, - 14, + 137, + 88, 0, - 15, + 89, + 0, + "*" + ], + [ + 170, + 89, + 0, + 30, + 1, + "MODEL" + ], + [ + 171, + 90, + 0, + 107, + 0, + "*" + ], + [ + 172, + 107, + 0, + 30, + 4, + "VAE" + ], + [ + 176, + 110, + 0, + 30, + 0, + "IMAGE" + ], + [ + 190, + 43, + 0, + 30, + 2, + "CONDITIONING" + ], + [ + 191, + 43, + 1, + 30, + 3, + "CONDITIONING" + ], + [ + 192, + 52, + 0, + 43, + 2, + "CONTROL_NET" + ], + [ + 195, + 122, + 0, + 123, 0, "UPSCALE_MODEL" ], [ - 20, - 12, + 197, + 123, 0, - 15, + 34, + 0, + "IMAGE" + ], + [ + 200, + 124, + 0, + 125, + 0, + "*" + ], + [ + 206, + 125, + 0, + 43, + 3, + "IMAGE" + ], + [ + 207, + 135, + 0, + 124, + 0, + "*" + ], + [ + 209, + 30, + 0, + 136, + 0, + "IMAGE" + ], + [ + 210, + 124, + 0, + 123, 1, "IMAGE" ], [ - 22, - 15, + 211, + 34, 0, - 17, + 110, 0, - "IMAGE" - ], - [ - 23, - 17, - 0, - 13, - 0, - "IMAGE" + "*" ] ], "groups": [], "config": {}, "extra": { "ds": { - "scale": 0.6649272177973091, + "scale": 1.0152559799477252, "offset": [ - -510.50184649377775, - -337.73126769987704 + -2260.53316345765, + -499.7179536588252 ] }, - "frontendVersion": "1.23.4" + "frontendVersion": "1.23.4", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true }, "version": 0.4 } \ No newline at end of file diff --git a/workflows/distributed-wan.json b/workflows/distributed-wan.json new file mode 100644 index 0000000..3e2faf0 --- /dev/null +++ b/workflows/distributed-wan.json @@ -0,0 +1,1688 @@ +{ + "id": "00000000-0000-0000-0000-000000000000", + "revision": 0, + "last_node_id": 234, + "last_link_id": 79, + "nodes": [ + { + "id": 67, + "type": "ModelSamplingSD3", + "pos": [ + 1700, + 550 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 23 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 17, + 26 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ModelSamplingSD3" + }, + "widgets_values": [ + 8.000000000000002 + ] + }, + { + "id": 50, + "type": "WanImageToVideo", + "pos": [ + 1590, + 230 + ], + "size": [ + 270, + 210 + ], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "link": 7 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 8 + }, + { + "name": "vae", + "type": "VAE", + "link": 9 + }, + { + "name": "clip_vision_output", + "shape": 7, + "type": "CLIP_VISION_OUTPUT", + "link": 10 + }, + { + "name": "start_image", + "shape": 7, + "type": "IMAGE", + "link": 11 + }, + { + "name": "width", + "type": "INT", + "widget": { + "name": "width" + }, + "link": 5 + }, + { + "name": "height", + "type": "INT", + "widget": { + "name": "height" + }, + "link": 6 + } + ], + "outputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 27 + ] + }, + { + "name": "negative", + "type": "CONDITIONING", + "links": null + }, + { + "name": "latent", + "type": "LATENT", + "links": [ + 22 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "WanImageToVideo" + }, + "widgets_values": [ + 832, + 480, + 17, + 1 + ] + }, + { + "id": 60, + "type": "RandomNoise", + "pos": [ + 1850, + 10 + ], + "size": [ + 270, + 82 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "noise_seed", + "type": "INT", + "widget": { + "name": "noise_seed" + }, + "link": 16 + } + ], + "outputs": [ + { + "name": "NOISE", + "type": "NOISE", + "links": [ + 18 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "RandomNoise" + }, + "widgets_values": [ + 200526686850175, + "randomize" + ] + }, + { + "id": 150, + "type": "DistributedSeed", + "pos": [ + 1540, + 0 + ], + "size": [ + 270, + 82 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "seed", + "type": "INT", + "links": [ + 16 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "9650280c7f50898720bf9773e8874d75e7c20846", + "Node name for S&R": "DistributedSeed" + }, + "widgets_values": [ + 177896884005433, + "randomize" + ] + }, + { + "id": 61, + "type": "KSamplerSelect", + "pos": [ + 2140, + 350 + ], + "size": [ + 260, + 60 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "SAMPLER", + "type": "SAMPLER", + "links": [ + 20 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "KSamplerSelect" + }, + "widgets_values": [ + "lcm" + ] + }, + { + "id": 221, + "type": "DistributedCollector", + "pos": [ + 2660, + 200 + ], + "size": [ + 166.50416564941406, + 26 + ], + "flags": { + "collapsed": false + }, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 67 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 78 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "b4ed92a363fab14d944b87c610aa9a0b4c87c085", + "Node name for S&R": "DistributedCollector" + }, + "widgets_values": [] + }, + { + "id": 231, + "type": "ImageBatchDivider", + "pos": [ + 2850, + 200 + ], + "size": [ + 210, + 118 + ], + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 78 + } + ], + "outputs": [ + { + "name": "batch_1", + "type": "IMAGE", + "links": [ + 74 + ] + }, + { + "name": "batch_2", + "type": "IMAGE", + "links": [ + 75 + ] + }, + { + "name": "batch_3", + "type": "IMAGE", + "links": [ + 76 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "fa09a2da10dfc63cecfbec1f8d4a0516e923b911", + "Node name for S&R": "ImageBatchDivider" + }, + "widgets_values": [ + 3 + ] + }, + { + "id": 139, + "type": "CLIPLoaderGGUF", + "pos": [ + 490, + 60 + ], + "size": [ + 270, + 82 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 1, + 2 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-gguf", + "ver": "1.1.1", + "Node name for S&R": "CLIPLoaderGGUF" + }, + "widgets_values": [ + "umt5-xxl-encoder-Q8_0.gguf", + "wan" + ] + }, + { + "id": 142, + "type": "UnetLoaderGGUF", + "pos": [ + 930, + 850 + ], + "size": [ + 340, + 60 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 24 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-gguf", + "ver": "1.1.1", + "Node name for S&R": "UnetLoaderGGUF" + }, + "widgets_values": [ + "wan2.1-i2v-14b-720p-Q8_0.gguf" + ] + }, + { + "id": 118, + "type": "LoraLoaderModelOnly", + "pos": [ + 1290, + 840 + ], + "size": [ + 310, + 82 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 24 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 23 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "LoraLoaderModelOnly" + }, + "widgets_values": [ + "Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors", + 0.8100000000000002 + ] + }, + { + "id": 135, + "type": "ImageNoiseAugmentation", + "pos": [ + 1250, + 310 + ], + "size": [ + 270, + 106 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 25 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 11 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-kjnodes", + "ver": "1.1.2", + "Node name for S&R": "ImageNoiseAugmentation" + }, + "widgets_values": [ + 0.10000000000000002, + 1100164865582526, + "randomize" + ] + }, + { + "id": 49, + "type": "CLIPVisionLoader", + "pos": [ + 1260, + 610 + ], + "size": [ + 290, + 60 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP_VISION", + "type": "CLIP_VISION", + "links": [ + 12 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "CLIPVisionLoader" + }, + "widgets_values": [ + "clip_vision_h.safetensors" + ] + }, + { + "id": 7, + "type": "CLIPTextEncode", + "pos": [ + 790, + 120 + ], + "size": [ + 380, + 110 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 2 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 8 + ] + } + ], + "title": "CLIP Text Encode (Negative Prompt)", + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" + ] + }, + { + "id": 64, + "type": "SamplerCustomAdvanced", + "pos": [ + 2140, + 190 + ], + "size": [ + 260, + 110 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "noise", + "type": "NOISE", + "link": 18 + }, + { + "name": "guider", + "type": "GUIDER", + "link": 19 + }, + { + "name": "sampler", + "type": "SAMPLER", + "link": 20 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 21 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 22 + } + ], + "outputs": [ + { + "name": "output", + "type": "LATENT", + "links": [ + 3 + ] + }, + { + "name": "denoised_output", + "type": "LATENT", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "SamplerCustomAdvanced" + }, + "widgets_values": [] + }, + { + "id": 233, + "type": "Note", + "pos": [ + 2840, + 40 + ], + "size": [ + 220, + 90 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Use this node to set how many videos to output.\n\nExample: if you have 1x master and 2x workers, set it to 3." + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 234, + "type": "Note", + "pos": [ + 1540, + -140 + ], + "size": [ + 270, + 90 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "This node will give you a different variation from each worker. Delete this node if you want them to be all the same." + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 227, + "type": "VHS_VideoCombine", + "pos": [ + 3090, + 200 + ], + "size": [ + 380, + 555.6923217773438 + ], + "flags": {}, + "order": 24, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 74 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 16, + "loop_count": 0, + "filename_prefix": "WAN", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00163.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 8, + "workflow": "AnimateDiff_00163.png", + "fullpath": "C:\\venvs\\ComfyUI\\ComfyUI\\output\\AnimateDiff_00163.mp4" + } + } + } + }, + { + "id": 228, + "type": "VHS_VideoCombine", + "pos": [ + 3490, + 190 + ], + "size": [ + 380, + 334 + ], + "flags": {}, + "order": 25, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 75 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 16, + "loop_count": 0, + "filename_prefix": "WAN", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00165.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 8, + "workflow": "AnimateDiff_00165.png", + "fullpath": "C:\\venvs\\ComfyUI\\ComfyUI\\output\\AnimateDiff_00165.mp4" + } + } + } + }, + { + "id": 229, + "type": "VHS_VideoCombine", + "pos": [ + 3880, + 190 + ], + "size": [ + 390, + 334 + ], + "flags": {}, + "order": 26, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 76 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 16, + "loop_count": 0, + "filename_prefix": "WAN", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00164.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 8, + "workflow": "AnimateDiff_00164.png", + "fullpath": "C:\\venvs\\ComfyUI\\ComfyUI\\output\\AnimateDiff_00164.mp4" + } + } + } + }, + { + "id": 8, + "type": "VAEDecode", + "pos": [ + 2450, + 190 + ], + "size": [ + 140, + 46 + ], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 3 + }, + { + "name": "vae", + "type": "VAE", + "link": 4 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 67 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "VAEDecode" + }, + "widgets_values": [] + }, + { + "id": 51, + "type": "CLIPVisionEncode", + "pos": [ + 1260, + 490 + ], + "size": [ + 290.97918701171875, + 78 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "clip_vision", + "type": "CLIP_VISION", + "link": 12 + }, + { + "name": "image", + "type": "IMAGE", + "link": 13 + } + ], + "outputs": [ + { + "name": "CLIP_VISION_OUTPUT", + "type": "CLIP_VISION_OUTPUT", + "links": [ + 10 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "CLIPVisionEncode" + }, + "widgets_values": [ + "none" + ] + }, + { + "id": 192, + "type": "LoadImage", + "pos": [ + 530, + 310 + ], + "size": [ + 310, + 370 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 79 + ] + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "example.png", + "image" + ] + }, + { + "id": 55, + "type": "ImageResize+", + "pos": [ + 890, + 320 + ], + "size": [ + 270, + 218 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 79 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 13, + 25 + ] + }, + { + "name": "width", + "type": "INT", + "links": [ + 5 + ] + }, + { + "name": "height", + "type": "INT", + "links": [ + 6 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui_essentials", + "ver": "9d9f4bedfc9f0321c19faf71855e228c93bd0dc9", + "Node name for S&R": "ImageResize+" + }, + "widgets_values": [ + 832, + 480, + "lanczos", + "fill / crop", + "always", + 0 + ] + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": [ + 790, + -10 + ], + "size": [ + 380, + 88 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 1 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 7 + ] + } + ], + "title": "CLIP Text Encode (Positive Prompt)", + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "beautiful scenery nature glass bottle landscape, , purple galaxy bottle," + ] + }, + { + "id": 39, + "type": "VAELoader", + "pos": [ + 1250, + 200 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 4, + 9 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "wan_2.1_vae.safetensors" + ] + }, + { + "id": 148, + "type": "BasicGuider", + "pos": [ + 1980, + 260 + ], + "size": [ + 156.0208282470703, + 46 + ], + "flags": { + "collapsed": true + }, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 26 + }, + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 27 + } + ], + "outputs": [ + { + "name": "GUIDER", + "type": "GUIDER", + "links": [ + 19 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "BasicGuider" + }, + "widgets_values": [] + }, + { + "id": 62, + "type": "BasicScheduler", + "pos": [ + 2130, + 460 + ], + "size": [ + 270, + 106 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 17 + } + ], + "outputs": [ + { + "name": "SIGMAS", + "type": "SIGMAS", + "links": [ + 21 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "BasicScheduler" + }, + "widgets_values": [ + "simple", + 4, + 1 + ] + } + ], + "links": [ + [ + 1, + 139, + 0, + 6, + 0, + "CLIP" + ], + [ + 2, + 139, + 0, + 7, + 0, + "CLIP" + ], + [ + 3, + 64, + 0, + 8, + 0, + "LATENT" + ], + [ + 4, + 39, + 0, + 8, + 1, + "VAE" + ], + [ + 5, + 55, + 1, + 50, + 5, + "INT" + ], + [ + 6, + 55, + 2, + 50, + 6, + "INT" + ], + [ + 7, + 6, + 0, + 50, + 0, + "CONDITIONING" + ], + [ + 8, + 7, + 0, + 50, + 1, + "CONDITIONING" + ], + [ + 9, + 39, + 0, + 50, + 2, + "VAE" + ], + [ + 10, + 51, + 0, + 50, + 3, + "CLIP_VISION_OUTPUT" + ], + [ + 11, + 135, + 0, + 50, + 4, + "IMAGE" + ], + [ + 12, + 49, + 0, + 51, + 0, + "CLIP_VISION" + ], + [ + 13, + 55, + 0, + 51, + 1, + "IMAGE" + ], + [ + 16, + 150, + 0, + 60, + 0, + "INT" + ], + [ + 17, + 67, + 0, + 62, + 0, + "MODEL" + ], + [ + 18, + 60, + 0, + 64, + 0, + "NOISE" + ], + [ + 19, + 148, + 0, + 64, + 1, + "GUIDER" + ], + [ + 20, + 61, + 0, + 64, + 2, + "SAMPLER" + ], + [ + 21, + 62, + 0, + 64, + 3, + "SIGMAS" + ], + [ + 22, + 50, + 2, + 64, + 4, + "LATENT" + ], + [ + 23, + 118, + 0, + 67, + 0, + "MODEL" + ], + [ + 24, + 142, + 0, + 118, + 0, + "MODEL" + ], + [ + 25, + 55, + 0, + 135, + 0, + "IMAGE" + ], + [ + 26, + 67, + 0, + 148, + 0, + "MODEL" + ], + [ + 27, + 50, + 0, + 148, + 1, + "CONDITIONING" + ], + [ + 67, + 8, + 0, + 221, + 0, + "IMAGE" + ], + [ + 74, + 231, + 0, + 227, + 0, + "IMAGE" + ], + [ + 75, + 231, + 1, + 228, + 0, + "IMAGE" + ], + [ + 76, + 231, + 2, + 229, + 0, + "IMAGE" + ], + [ + 78, + 221, + 0, + 231, + 0, + "IMAGE" + ], + [ + 79, + 192, + 0, + 55, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.9090909090909091, + "offset": [ + -1587.6131275248988, + 280.3225523752761 + ] + }, + "frontendVersion": "1.23.4", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file