Main release

utils/__init__.py

@@ -0,0 +1,5 @@
+"""
+Utility modules for ComfyUI-Distributed extension.
+"""
+
+# Make utils importable as a package

utils/async_helpers.py

@@ -0,0 +1,50 @@
+"""
+Async helper utilities for ComfyUI-Distributed.
+"""
+import asyncio
+import threading
+from typing import Optional, Any, Coroutine
+from .network import get_server_loop
+
+def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) -> Any:
+    """
+    Run async coroutine in server's event loop and wait for result.
+    
+    This is useful when you need to run async code from a synchronous context
+    but want to use the server's existing event loop instead of creating a new one.
+    
+    Args:
+        coro: The coroutine to run
+        timeout: Optional timeout in seconds
+        
+    Returns:
+        The result of the coroutine
+        
+    Raises:
+        TimeoutError: If the operation times out
+        Exception: Any exception raised by the coroutine
+    """
+    event = threading.Event()
+    result = None
+    error = None
+    
+    async def wrapper():
+        nonlocal result, error
+        try:
+            result = await coro
+        except Exception as e:
+            error = e
+        finally:
+            event.set()
+    
+    # Schedule on server's event loop
+    loop = get_server_loop()
+    asyncio.run_coroutine_threadsafe(wrapper(), loop)
+    
+    # Wait for completion
+    if not event.wait(timeout):
+        raise TimeoutError(f"Async operation timed out after {timeout} seconds")
+    
+    if error:
+        raise error
+    return result

utils/config.py

@@ -0,0 +1,50 @@
+"""
+Configuration management for ComfyUI-Distributed.
+"""
+import os
+import json
+from .logging import log
+
+CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json")
+
+def get_default_config():
+    """Returns the default configuration dictionary. Single source of truth."""
+    return {
+        "master": {"host": ""},
+        "workers": [],
+        "settings": {
+            "debug": False,
+            "auto_launch_workers": False,
+            "stop_workers_on_master_exit": True
+        }
+    }
+
+def load_config():
+    """Loads the config, falling back to defaults if the file is missing or invalid."""
+    if os.path.exists(CONFIG_FILE):
+        try:
+            with open(CONFIG_FILE, 'r') as f:
+                return json.load(f)
+        except Exception as e:
+            log(f"Error loading config, using defaults: {e}")
+    return get_default_config()
+
+def save_config(config):
+    """Saves the configuration to file."""
+    try:
+        with open(CONFIG_FILE, 'w') as f:
+            json.dump(config, f, indent=2)
+        return True
+    except Exception as e:
+        log(f"Error saving config: {e}")
+        return False
+
+def ensure_config_exists():
+    """Creates default config file if it doesn't exist. Used by __init__.py"""
+    if not os.path.exists(CONFIG_FILE):
+        default_config = get_default_config()
+        if save_config(default_config):
+            from .logging import debug_log
+            debug_log("Created default config file")
+        else:
+            log("Could not create default config file")

utils/constants.py

@@ -0,0 +1,34 @@
+"""
+Shared constants for ComfyUI-Distributed.
+"""
+
+# Timeouts (in seconds)
+WORKER_JOB_TIMEOUT = 10.0
+TILE_COLLECTION_TIMEOUT = 60.0
+TILE_WAIT_TIMEOUT = 30.0
+PROCESS_TERMINATION_TIMEOUT = 5.0
+
+# Process monitoring
+WORKER_CHECK_INTERVAL = 2.0
+STATUS_CHECK_INTERVAL = 5.0
+
+# Network
+CHUNK_SIZE = 8192
+LOG_TAIL_BYTES = 65536  # 64KB
+
+# File paths
+WORKER_LOG_PATTERN = "distributed_worker_*.log"
+
+# Worker management
+WORKER_STARTUP_DELAY = 2.0
+
+# Tile transfer
+TILE_TRANSFER_TIMEOUT = 30.0
+
+# Process cleanup
+PROCESS_WAIT_TIMEOUT = 3.0
+QUEUE_INIT_TIMEOUT = 5.0
+TILE_SEND_TIMEOUT = 60.0
+
+# Memory operations  
+MEMORY_CLEAR_DELAY = 0.5

utils/image.py

@@ -0,0 +1,25 @@
+"""
+Image and tensor conversion utilities for ComfyUI-Distributed.
+"""
+import torch
+import numpy as np
+from PIL import Image
+
+def tensor_to_pil(img_tensor, batch_index=0):
+    """Takes a batch of images in tensor form [B, H, W, C] and returns an RGB PIL Image."""
+    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, H, W, C]."""
+    image = np.array(image).astype(np.float32) / 255.0
+    image = torch.from_numpy(image).unsqueeze(0)
+    if len(image.shape) == 3:  # If grayscale, add channel dimension
+        image = image.unsqueeze(-1)
+    return image
+
+def ensure_contiguous(tensor):
+    """Ensure tensor is contiguous in memory."""
+    if not tensor.is_contiguous():
+        return tensor.contiguous()
+    return tensor
+

utils/logging.py

@@ -0,0 +1,28 @@
+"""
+Shared logging utilities for ComfyUI-Distributed.
+"""
+import os
+import json
+
+# Config file is in parent directory
+CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json")
+
+def is_debug_enabled():
+    """Check if debug is enabled."""
+    if os.path.exists(CONFIG_FILE):
+        try:
+            with open(CONFIG_FILE, 'r') as f:
+                config = json.load(f)
+                return config.get("settings", {}).get("debug", False)
+        except:
+            pass
+    return False
+
+def debug_log(message):
+    """Log debug messages only if debug is enabled in config."""
+    if is_debug_enabled():
+        print(f"[Distributed] {message}")
+
+def log(message):
+    """Always log important messages."""
+    print(f"[Distributed] {message}")

utils/network.py

@@ -0,0 +1,41 @@
+"""
+Network and API utilities for ComfyUI-Distributed.
+"""
+import aiohttp
+from aiohttp import web
+from .logging import debug_log
+
+# Shared session for connection pooling
+_client_session = None
+
+async def get_client_session():
+    """Get or create a shared aiohttp client session."""
+    global _client_session
+    if _client_session is None or _client_session.closed:
+        connector = aiohttp.TCPConnector(limit=100, limit_per_host=30)
+        # Don't set timeout here - set it per request
+        _client_session = aiohttp.ClientSession(connector=connector)
+    return _client_session
+
+async def cleanup_client_session():
+    """Clean up the shared client session."""
+    global _client_session
+    if _client_session and not _client_session.closed:
+        await _client_session.close()
+        _client_session = None
+
+async def handle_api_error(request, error, status=500):
+    """Standardized error response handler."""
+    debug_log(f"API Error: {error}")
+    return web.json_response({"status": "error", "message": str(error)}, status=status)
+
+def get_server_port():
+    """Get the ComfyUI server port."""
+    import server
+    return server.PromptServer.instance.port
+
+def get_server_loop():
+    """Get the ComfyUI server event loop."""
+    import server
+    return server.PromptServer.instance.loop
+

utils/process.py

@@ -0,0 +1,38 @@
+"""
+Process management utilities for ComfyUI-Distributed.
+"""
+import os
+import subprocess
+import platform
+import signal
+
+def is_process_alive(pid):
+    """Check if a process with given PID is still alive."""
+    try:
+        if platform.system() == "Windows":
+            # Windows: use tasklist
+            result = subprocess.run(['tasklist', '/FI', f'PID eq {pid}'], 
+                                  capture_output=True, text=True)
+            return str(pid) in result.stdout
+        else:
+            # Unix: send signal 0
+            os.kill(pid, 0)
+            return True
+    except (OSError, subprocess.SubprocessError):
+        return False
+
+def terminate_process(process, timeout=5):
+    """Gracefully terminate a process with timeout."""
+    if process.poll() is None:  # Still running
+        process.terminate()
+        try:
+            process.wait(timeout=timeout)
+        except subprocess.TimeoutExpired:
+            process.kill()
+            process.wait()
+
+def get_python_executable():
+    """Get the Python executable path."""
+    import sys
+    return sys.executable
+