nodes/robertvoy-ComfyUI-Distributed
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Feat: Audio Batch Divider

Browse Source
This commit is contained in:
robertvoy
2026-01-25 17:36:55 +11:00
parent 64ba4b75db
commit d691ba841f
2 changed files with 658 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
distributed.py
+584 -46
View File
@@ -4,11 +4,13 @@ from PIL import Image
import folder_paths
import os
import json
import re
import asyncio
import aiohttp
from aiohttp import web
import io
import server
import execution
import comfy.model_management
import subprocess
import platform
@@ -36,6 +38,7 @@ from .utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
from .utils.process import is_process_alive, terminate_process, get_python_executable
from .utils.network import handle_api_error, get_server_port, get_server_loop, get_client_session, cleanup_client_session
from .utils.async_helpers import run_async_in_server_loop
from .utils.cloudflare import cloudflare_tunnel_manager
from .utils.constants import (
WORKER_JOB_TIMEOUT, PROCESS_TERMINATION_TIMEOUT, WORKER_CHECK_INTERVAL,
STATUS_CHECK_INTERVAL, CHUNK_SIZE, LOG_TAIL_BYTES, WORKER_LOG_PATTERN,
@@ -59,6 +62,17 @@ def cleanup():
atexit.register(cleanup)
def normalize_host(value):
if value is None:
return None
if not isinstance(value, str):
return value
host = value.strip()
if not host:
return host
host = re.sub(r"^https?://", "", host, flags=re.IGNORECASE)
return host.split("/")[0]
# --- API Endpoints ---
@server.PromptServer.instance.routes.get("/distributed/config")
async def get_config_endpoint(request):
@@ -83,6 +97,102 @@ async def queue_status_endpoint(request):
except Exception as e:
return await handle_api_error(request, e, 500)
prompt_server = server.PromptServer.instance
async def _queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
"""Validate and queue a prompt via ComfyUI's prompt queue."""
payload = {"prompt": prompt_obj}
payload = prompt_server.trigger_on_prompt(payload)
prompt = payload["prompt"]
prompt_id = str(uuid.uuid4())
valid = await execution.validate_prompt(prompt_id, prompt, None)
if not valid[0]:
raise RuntimeError(f"Invalid prompt: {valid[1]}")
extra_data = {}
if workflow_meta:
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
if client_id:
extra_data["client_id"] = client_id
sensitive = {}
for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
if key in extra_data:
sensitive[key] = extra_data.pop(key)
number = getattr(prompt_server, "number", 0)
prompt_server.number = number + 1
prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
prompt_server.prompt_queue.put(prompt_queue_item)
return prompt_id
@server.PromptServer.instance.routes.get("/distributed/worker_ws")
async def worker_ws_endpoint(request):
"""WebSocket endpoint for worker prompt dispatch."""
ws = web.WebSocketResponse(heartbeat=30)
await ws.prepare(request)
async for msg in ws:
if msg.type == aiohttp.WSMsgType.TEXT:
try:
data = json.loads(msg.data or "{}")
except json.JSONDecodeError:
await ws.send_json({
"type": "dispatch_ack",
"request_id": None,
"ok": False,
"error": "Invalid JSON payload.",
})
continue
if data.get("type") != "dispatch_prompt":
await ws.send_json({
"type": "dispatch_ack",
"request_id": data.get("request_id"),
"ok": False,
"error": "Unsupported websocket message type.",
})
continue
prompt = data.get("prompt")
if not isinstance(prompt, dict):
await ws.send_json({
"type": "dispatch_ack",
"request_id": data.get("request_id"),
"ok": False,
"error": "Field 'prompt' must be an object.",
})
continue
try:
prompt_id = await _queue_prompt_payload(
prompt,
workflow_meta=data.get("workflow"),
client_id=data.get("client_id"),
)
await ws.send_json({
"type": "dispatch_ack",
"request_id": data.get("request_id"),
"ok": True,
"prompt_id": prompt_id,
})
except Exception as exc:
await ws.send_json({
"type": "dispatch_ack",
"request_id": data.get("request_id"),
"ok": False,
"error": str(exc),
})
elif msg.type == aiohttp.WSMsgType.ERROR:
log(f"[Distributed] Worker websocket error: {ws.exception()}")
return ws
@server.PromptServer.instance.routes.post("/distributed/worker/clear_launching")
async def clear_launching_state(request):
"""Clear the launching flag when worker is confirmed running."""
@@ -269,6 +379,49 @@ async def get_network_info_endpoint(request):
"recommended_ip": None
})
@server.PromptServer.instance.routes.get("/distributed/tunnel/status")
async def tunnel_status_endpoint(request):
"""Return Cloudflare tunnel status and last known details."""
try:
status = cloudflare_tunnel_manager.get_status()
config = load_config()
master_host = (config.get("master") or {}).get("host")
return web.json_response({
"status": "success",
"tunnel": status,
"master_host": master_host
})
except Exception as e:
return await handle_api_error(request, e, 500)
@server.PromptServer.instance.routes.post("/distributed/tunnel/start")
async def tunnel_start_endpoint(request):
"""Start a Cloudflare tunnel pointing at the current ComfyUI server."""
try:
result = await cloudflare_tunnel_manager.start_tunnel()
config = load_config()
return web.json_response({
"status": "success",
"tunnel": result,
"master_host": (config.get("master") or {}).get("host")
})
except Exception as e:
return await handle_api_error(request, e, 500)
@server.PromptServer.instance.routes.post("/distributed/tunnel/stop")
async def tunnel_stop_endpoint(request):
"""Stop the managed Cloudflare tunnel if running."""
try:
result = await cloudflare_tunnel_manager.stop_tunnel()
config = load_config()
return web.json_response({
"status": "success",
"tunnel": result,
"master_host": (config.get("master") or {}).get("host")
})
except Exception as e:
return await handle_api_error(request, e, 500)
@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."""
@@ -323,7 +476,7 @@ async def update_worker_endpoint(request):
if data["host"] is None:
worker.pop("host", None)
else:
worker["host"] = data["host"]
worker["host"] = normalize_host(data["host"])
# Handle cuda_device field - remove it if None
if "cuda_device" in data:
@@ -352,7 +505,7 @@ async def update_worker_endpoint(request):
new_worker = {
"id": worker_id,
"name": data["name"],
"host": data.get("host", "localhost"),
"host": normalize_host(data.get("host", "localhost")),
"port": data["port"],
"cuda_device": data["cuda_device"],
"enabled": data.get("enabled", False),
@@ -629,7 +782,8 @@ async def get_local_worker_status_endpoint(request):
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"]:
host = normalize_host(worker.get("host")) or ""
if not host or host in ["localhost", "127.0.0.1"]:
worker_id = worker["id"]
port = worker["port"]
@@ -1280,7 +1434,7 @@ async def get_client_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')
host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
if host in ['localhost', '127.0.0.1', '0.0.0.0', ''] or worker_config.get('type') == 'local':
return True
@@ -1297,7 +1451,7 @@ async def is_same_physical_host(worker_config):
master_machine_id = get_machine_id()
# Fetch worker's machine ID via API
host = worker_config.get('host', 'localhost')
host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
port = worker_config.get('port', 8188)
session = await get_client_session()
@@ -1360,11 +1514,11 @@ async def get_comms_channel(worker_id, worker_config):
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')
host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
return f"http://{host}:{worker_config['port']}"
else:
# Remote worker: use configured endpoint
host = worker_config.get('host', 'localhost')
host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
return f"http://{host}:{worker_config['port']}"
# Auto-launch workers if enabled
@@ -1391,7 +1545,7 @@ def auto_launch_workers():
worker_name = worker.get('name', f'Worker {worker_id}')
# Skip remote workers
host = worker.get('host', 'localhost').lower()
host = (normalize_host(worker.get('host', 'localhost')) or 'localhost').lower()
if host not in ['localhost', '127.0.0.1', '', None]:
debug_log(f"Skipping remote worker {worker_name} (host: {host})")
continue
@@ -1445,6 +1599,10 @@ async def async_cleanup_and_exit(signum=None):
else:
print("\n[Distributed] Master shutting down, workers will continue running")
worker_manager.save_processes()
try:
await cloudflare_tunnel_manager.stop_tunnel()
except Exception as tunnel_error:
debug_log(f"Error stopping Cloudflare tunnel during shutdown: {tunnel_error}")
except Exception as e:
print(f"[Distributed] Error during cleanup: {e}")
@@ -1503,6 +1661,17 @@ def sync_cleanup():
else:
print("\n[Distributed] Master shutting down, workers will continue running")
worker_manager.save_processes()
try:
loop = asyncio.get_event_loop()
if loop.is_running():
loop.create_task(cloudflare_tunnel_manager.stop_tunnel())
else:
loop.run_until_complete(cloudflare_tunnel_manager.stop_tunnel())
except RuntimeError:
# No running loop; create a temporary one
asyncio.run(cloudflare_tunnel_manager.stop_tunnel())
except Exception as tunnel_error:
debug_log(f"Error stopping Cloudflare tunnel during sync cleanup: {tunnel_error}")
except Exception as e:
print(f"[Distributed] Error during cleanup: {e}")
@@ -1525,6 +1694,10 @@ if not os.environ.get('COMFYUI_MASTER_PID'):
worker_manager.cleanup_all()
else:
worker_manager.save_processes()
try:
asyncio.run(cloudflare_tunnel_manager.stop_tunnel())
except Exception as tunnel_error:
print(f"[Distributed] Error stopping Cloudflare tunnel: {tunnel_error}")
except Exception as cleanup_error:
print(f"[Distributed] Error during cleanup: {cleanup_error}")
sys.exit(0)
@@ -1814,6 +1987,21 @@ async def job_complete_endpoint(request):
log(f"Error processing image from worker {worker_id}: {e}")
return await handle_api_error(request, f"Image processing error: {e}", 400)
# Parse audio data if present
audio_data = None
audio_waveform_field = data.get('audio_waveform')
audio_sample_rate = data.get('audio_sample_rate')
if audio_waveform_field is not None:
try:
waveform_bytes = audio_waveform_field.file.read()
waveform_buffer = io.BytesIO(waveform_bytes)
waveform_tensor = torch.load(waveform_buffer, weights_only=True)
sample_rate = int(audio_sample_rate) if audio_sample_rate else 44100
audio_data = {"waveform": waveform_tensor, "sample_rate": sample_rate}
debug_log(f"Received audio from worker {worker_id}: shape={waveform_tensor.shape}, sample_rate={sample_rate}")
except Exception as e:
log(f"Error parsing audio from worker {worker_id}: {e}")
# Put batch into queue
async with prompt_server.distributed_jobs_lock:
debug_log(f"Current pending jobs: {list(prompt_server.distributed_pending_jobs.keys())}")
@@ -1824,7 +2012,8 @@ async def job_complete_endpoint(request):
'worker_id': worker_id,
'tensors': tensors,
'indices': indices,
'is_last': is_last
'is_last': is_last,
'audio': audio_data
})
debug_log(f"Received batch result for job {multi_job_id} from worker {worker_id}, size={len(tensors)}")
else:
@@ -1833,7 +2022,8 @@ async def job_complete_endpoint(request):
'tensor': tensors[0],
'worker_id': worker_id,
'image_index': int(image_index) if image_index else 0,
'is_last': is_last
'is_last': is_last,
'audio': audio_data
})
debug_log(f"Received single result for job {multi_job_id} from worker {worker_id}")
@@ -1852,6 +2042,7 @@ class DistributedCollectorNode:
def INPUT_TYPES(s):
return {
"required": { "images": ("IMAGE",) },
"optional": { "audio": ("AUDIO",) },
"hidden": {
"multi_job_id": ("STRING", {"default": ""}),
"is_worker": ("BOOLEAN", {"default": False}),
@@ -1864,45 +2055,47 @@ class DistributedCollectorNode:
},
}
RETURN_TYPES = ("IMAGE",)
RETURN_TYPES = ("IMAGE", "AUDIO")
RETURN_NAMES = ("images", "audio")
FUNCTION = "run"
CATEGORY = "image"
def run(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
def run(self, images, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
# Create empty audio if not provided
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
if not multi_job_id or pass_through:
if pass_through:
print(f"[Distributed Collector] Pass-through mode enabled, returning images unchanged")
return (images,)
return (images, audio if audio is not None else empty_audio)
# Use async helper to run in server loop
result = run_async_in_server_loop(
self.execute(images, multi_job_id, is_worker, master_url, enabled_worker_ids, worker_batch_size, worker_id, delegate_only)
self.execute(images, audio, multi_job_id, is_worker, master_url, enabled_worker_ids, worker_batch_size, worker_id, delegate_only)
)
return result
async def send_batch_to_master(self, image_batch, multi_job_id, master_url, worker_id):
"""Send image batch to master, chunked if large."""
async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
"""Send image batch and optional audio to master, chunked if large."""
batch_size = image_batch.shape[0]
if batch_size == 0:
if batch_size == 0 and audio is None:
return
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
@@ -1911,9 +2104,21 @@ class DistributedCollectorNode:
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')
# Add audio data only on the final chunk to avoid duplication
if is_chunk_last and audio is not None:
waveform = audio.get("waveform")
sample_rate = audio.get("sample_rate", 44100)
if waveform is not None and waveform.numel() > 0:
# Serialize waveform tensor to bytes
audio_bytes = io.BytesIO()
torch.save(waveform, audio_bytes)
audio_bytes.seek(0)
data.add_field('audio_waveform', audio_bytes, filename='audio.pt', content_type='application/octet-stream')
data.add_field('audio_sample_rate', str(sample_rate))
debug_log(f"Worker - Including audio: shape={waveform.shape}, sample_rate={sample_rate}")
try:
session = await get_client_session()
url = f"{master_url}/distributed/job_complete"
async with session.post(url, data=data) as response:
@@ -1923,36 +2128,76 @@ class DistributedCollectorNode:
debug_log(f"Worker - Full error details: URL={url}")
raise # Re-raise to handle at caller level
def _combine_audio(self, master_audio, worker_audio, empty_audio):
"""Combine audio from master and workers into a single audio output."""
audio_pieces = []
sample_rate = 44100
# Add master audio first if present
if master_audio is not None:
waveform = master_audio.get("waveform")
if waveform is not None and waveform.numel() > 0:
audio_pieces.append(waveform)
sample_rate = master_audio.get("sample_rate", 44100)
# Add worker audio in sorted order
for worker_id_str in sorted(worker_audio.keys()):
w_audio = worker_audio[worker_id_str]
if w_audio is not None:
waveform = w_audio.get("waveform")
if waveform is not None and waveform.numel() > 0:
audio_pieces.append(waveform)
# Use first available sample rate
if sample_rate == 44100:
sample_rate = w_audio.get("sample_rate", 44100)
if not audio_pieces:
return empty_audio
try:
# Concatenate along the samples dimension (dim=-1)
# Ensure all pieces have same batch and channel dimensions
combined_waveform = torch.cat(audio_pieces, dim=-1)
debug_log(f"Master - Combined audio: {len(audio_pieces)} pieces, final shape={combined_waveform.shape}")
return {"waveform": combined_waveform, "sample_rate": sample_rate}
except Exception as e:
log(f"Master - Error combining audio: {e}")
return empty_audio
async def execute(self, images, audio, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
async def execute(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
if is_worker:
# Worker mode: send images to master in a single batch
# Worker mode: send images and audio to master in a single batch
debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master")
await self.send_batch_to_master(images, multi_job_id, master_url, worker_id)
return (images,)
await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
return (images, audio if audio is not None else empty_audio)
else:
delegate_mode = delegate_only or is_master_delegate_only()
# Master mode: collect images from workers
# Master mode: collect images and audio from workers
enabled_workers = json.loads(enabled_worker_ids)
num_workers = len(enabled_workers)
if num_workers == 0:
return (images,)
return (images, audio if audio is not None else empty_audio)
if delegate_mode:
master_batch_size = 0
images_on_cpu = None
master_audio = None
debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
else:
images_on_cpu = images.cpu()
master_batch_size = images.shape[0]
master_audio = audio # Keep master's audio for later
debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
# Ensure master images are contiguous
images_on_cpu = ensure_contiguous(images_on_cpu)
# Initialize storage for collected images
# Initialize storage for collected images and audio
worker_images = {} # Dict to store images by worker_id and index
worker_audio = {} # Dict to store audio by worker_id
# Get the existing queue - it should already exist from prepare_job
async with prompt_server.distributed_jobs_lock:
@@ -2020,19 +2265,25 @@ class DistributedCollectorNode:
# 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
# Collect audio data if present
result_audio = result.get('audio')
if result_audio is not None:
worker_audio[worker_id] = result_audio
debug_log(f"Master - Got audio from worker {worker_id}")
# Record activity and refresh timeout baseline
last_activity = time.time()
base_timeout = float(get_worker_timeout_seconds())
if is_last:
workers_done.add(worker_id)
p.update(1) # +1 per completed worker
@@ -2058,7 +2309,7 @@ class DistributedCollectorNode:
if not wrec:
debug_log(f"Collector probe: worker {wid} not found in config")
continue
host = wrec.get('host') or 'localhost'
host = normalize_host(wrec.get('host') or 'localhost') or 'localhost'
port = int(wrec.get('port', 8188))
url = f"http://{host}:{port}/prompt"
try:
@@ -2188,11 +2439,15 @@ class DistributedCollectorNode:
# Ensure the combined tensor is contiguous and properly formatted
combined = ensure_contiguous(combined)
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total (master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
return (combined,)
# Combine audio from master and workers
combined_audio = self._combine_audio(master_audio, worker_audio, empty_audio)
return (combined, combined_audio)
except Exception as e:
log(f"Master - Error combining images: {e}")
# Return just the master images as fallback
return (images,)
return (images, audio if audio is not None else empty_audio)
# --- Distributor Node ---
class DistributedSeed:
@@ -2256,6 +2511,68 @@ class AnyType(str):
any_type = AnyType("*")
class DistributedModelName:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"text": ("STRING", {"default": ""}),
},
"hidden": {
"unique_id": "UNIQUE_ID",
"extra_pnginfo": "EXTRA_PNGINFO",
},
}
RETURN_TYPES = (any_type,)
RETURN_NAMES = ("output",)
FUNCTION = "log_input"
OUTPUT_NODE = True
CATEGORY = "utils"
def _stringify(self, value):
if isinstance(value, str):
return value
if isinstance(value, (int, float, bool)):
return str(value)
try:
return json.dumps(value, indent=4)
except Exception:
return str(value)
def _update_workflow(self, extra_pnginfo, unique_id, values):
if not extra_pnginfo:
return
info = extra_pnginfo[0] if isinstance(extra_pnginfo, list) else extra_pnginfo
if not isinstance(info, dict) or "workflow" not in info:
return
node_id = None
if isinstance(unique_id, list) and unique_id:
node_id = str(unique_id[0])
elif unique_id is not None:
node_id = str(unique_id)
if not node_id:
return
workflow = info["workflow"]
node = next((x for x in workflow["nodes"] if str(x.get("id")) == node_id), None)
if node:
node["widgets_values"] = [values]
def log_input(self, text, unique_id=None, extra_pnginfo=None):
values = []
if isinstance(text, list):
for val in text:
values.append(self._stringify(val))
else:
values.append(self._stringify(text))
# Keep widget display in workflow metadata if available.
self._update_workflow(extra_pnginfo, unique_id, values)
if isinstance(values, list) and len(values) == 1:
return {"ui": {"text": values}, "result": (values[0],)}
return {"ui": {"text": values}, "result": (values,)}
class ByPassTypeTuple(tuple):
def __getitem__(self, index):
if index > 0:
@@ -2323,6 +2640,75 @@ class ImageBatchDivider:
return tuple(outputs)
class AudioBatchDivider:
"""Divides an audio waveform into multiple parts along the time/samples dimension."""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"audio": ("AUDIO",),
"divide_by": ("INT", {
"default": 2,
"min": 1,
"max": 10,
"step": 1,
"display": "number",
"tooltip": "Number of parts to divide the audio into"
}),
}
}
RETURN_TYPES = ByPassTypeTuple(("AUDIO",)) # Flexible for variable outputs
RETURN_NAMES = ByPassTypeTuple(tuple([f"audio_{i+1}" for i in range(10)]))
FUNCTION = "divide_audio"
OUTPUT_NODE = True
CATEGORY = "audio"
def divide_audio(self, audio, divide_by):
import torch
waveform = audio.get("waveform")
sample_rate = audio.get("sample_rate", 44100)
if waveform is None or waveform.numel() == 0:
# Return empty audio for all outputs
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": sample_rate}
return tuple([empty_audio] * 10)
total_splits = min(divide_by, 10) # Cap to max 10
# Waveform shape: [batch, channels, samples]
total_samples = waveform.shape[-1]
samples_per_split = total_samples // total_splits
remainder = total_samples % total_splits
outputs = []
start_idx = 0
for i in range(total_splits):
current_samples = samples_per_split + (1 if i < remainder else 0)
end_idx = start_idx + current_samples
split_waveform = waveform[..., start_idx:end_idx]
outputs.append({
"waveform": split_waveform,
"sample_rate": sample_rate
})
start_idx = end_idx
# Pad with empty audio up to max (10) to match RETURN_TYPES length
empty_audio = {
"waveform": torch.zeros(waveform.shape[0], waveform.shape[1], 1,
dtype=waveform.dtype, device=waveform.device),
"sample_rate": sample_rate
}
while len(outputs) < 10:
outputs.append(empty_audio)
return tuple(outputs)
class DistributedEmptyImage:
"""Produces an empty IMAGE batch used when the master delegates all work."""
@@ -2348,15 +2734,167 @@ class DistributedEmptyImage:
return (tensor,)
NODE_CLASS_MAPPINGS = {
# --- Distributed Queue Node ---
class DistributedQueueNode:
"""Routes entire workflows to the least-busy worker for job scheduling/load balancing."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {},
"hidden": {
"prompt": "PROMPT",
"extra_pnginfo": "EXTRA_PNGINFO",
"client_id": "CLIENT_ID",
"skip_dispatch": ("BOOLEAN", {"default": False}),
},
}
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("worker_id", "prompt_id")
FUNCTION = "queue"
OUTPUT_NODE = True
CATEGORY = "utils"
def queue(self, prompt=None, extra_pnginfo=None, client_id=None, skip_dispatch=False):
if skip_dispatch:
debug_log("DistributedQueue: skip_dispatch enabled; not dispatching.")
return ("", "")
if not prompt:
log("DistributedQueue: No prompt supplied; skipping dispatch.")
return ("", "")
return run_async_in_server_loop(
self._queue_on_worker(prompt, extra_pnginfo, client_id)
)
async def _queue_on_worker(self, prompt_obj, workflow_meta, client_id):
config = load_config()
enabled_workers = [
worker for worker in config.get("workers", [])
if worker.get("enabled", False)
]
if not enabled_workers:
log("DistributedQueue: No enabled workers found.")
return ("", "")
statuses = await self._fetch_worker_statuses(enabled_workers)
if not statuses:
log("DistributedQueue: No reachable workers found.")
return ("", "")
selected = self._select_worker(statuses)
worker = selected["worker"]
queue_remaining = selected["queue_remaining"]
prompt_copy = json.loads(json.dumps(prompt_obj))
self._mark_skip_dispatch(prompt_copy)
payload = {"prompt": prompt_copy}
extra_data = {}
if workflow_meta:
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
if client_id:
extra_data["client_id"] = client_id
if extra_data:
payload["extra_data"] = extra_data
url = self._build_worker_url(worker, "/prompt")
session = await get_client_session()
async with session.post(
url,
json=payload,
timeout=aiohttp.ClientTimeout(total=60),
) as resp:
resp.raise_for_status()
response_payload = await resp.json()
prompt_id = str(response_payload.get("prompt_id", ""))
worker_id = str(worker.get("id", ""))
log(
f"DistributedQueue: queued prompt {prompt_id} on worker {worker_id} (queue_remaining={queue_remaining})."
)
return (worker_id, prompt_id)
async def _fetch_worker_statuses(self, workers):
session = await get_client_session()
statuses = []
for worker in workers:
url = self._build_worker_url(worker, "/prompt")
try:
async with session.get(
url,
timeout=aiohttp.ClientTimeout(total=3),
) as resp:
if resp.status != 200:
continue
data = await resp.json()
queue_remaining = int(data.get("exec_info", {}).get("queue_remaining", 0))
statuses.append({
"worker": worker,
"queue_remaining": queue_remaining,
})
except Exception as e:
debug_log(f"DistributedQueue: Worker status check failed ({url}): {e}")
return statuses
def _select_worker(self, statuses):
idle_workers = [status for status in statuses if status["queue_remaining"] == 0]
if idle_workers:
return self._select_round_robin(idle_workers)
return min(statuses, key=lambda status: status["queue_remaining"])
def _select_round_robin(self, statuses):
prompt_server = server.PromptServer.instance
if not hasattr(prompt_server, "distributed_queue_rr_index"):
prompt_server.distributed_queue_rr_index = 0
index = prompt_server.distributed_queue_rr_index % len(statuses)
prompt_server.distributed_queue_rr_index += 1
return statuses[index]
def _mark_skip_dispatch(self, prompt_obj):
for node in prompt_obj.values():
if isinstance(node, dict) and node.get("class_type") == "DistributedQueue":
inputs = node.setdefault("inputs", {})
inputs["skip_dispatch"] = True
def _build_worker_url(self, worker, endpoint=""):
host = (worker.get("host") or "").strip()
port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
if not host:
host = getattr(server.PromptServer.instance, "address", "127.0.0.1") or "127.0.0.1"
if host.startswith(("http://", "https://")):
base = host.rstrip("/")
else:
is_cloud = worker.get("type") == "cloud" or host.endswith(".proxy.runpod.net") or port == 443
scheme = "https" if is_cloud else "http"
default_port = 443 if scheme == "https" else 80
port_part = "" if port == default_port else f":{port}"
base = f"{scheme}://{host}{port_part}"
return f"{base}{endpoint}"
NODE_CLASS_MAPPINGS = {
"DistributedQueue": DistributedQueueNode,
"DistributedCollector": DistributedCollectorNode,
"DistributedSeed": DistributedSeed,
"DistributedModelName": DistributedModelName,
"ImageBatchDivider": ImageBatchDivider,
"AudioBatchDivider": AudioBatchDivider,
"DistributedEmptyImage": DistributedEmptyImage,
}
NODE_DISPLAY_NAME_MAPPINGS = {
NODE_DISPLAY_NAME_MAPPINGS = {
"DistributedQueue": "Distributed Queue",
"DistributedCollector": "Distributed Collector",
"DistributedSeed": "Distributed Seed",
"DistributedModelName": "Distributed Model Name",
"ImageBatchDivider": "Image Batch Divider",
"AudioBatchDivider": "Audio Batch Divider",
"DistributedEmptyImage": "Distributed Empty Image",
}
web/image_batch_divider.js
+74 -72
View File
@@ -1,86 +1,88 @@
import { app } from "/scripts/app.js";
// Configuration for each batch divider node type
const BATCH_DIVIDER_NODES = {
"ImageBatchDivider": { outputPrefix: "batch_", outputType: "IMAGE" },
"AudioBatchDivider": { outputPrefix: "audio_", outputType: "AUDIO" }
};
app.registerExtension({
name: "Distributed.ImageBatchDivider",
name: "Distributed.BatchDividers",
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 config = BATCH_DIVIDER_NODES[node.comfyClass];
if (!config) return;
try {
const updateOutputs = () => {
if (!node.widgets) return;
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
};
if (!divideByWidget) return;
const divideBy = parseInt(divideByWidget.value, 10) || 1;
const totalOutputs = divideBy;
// Ensure outputs array exists
if (!node.outputs) node.outputs = [];
// Remove excess outputs
while (node.outputs.length > totalOutputs) {
node.removeOutput(node.outputs.length - 1);
}
// 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;
// Add missing outputs
while (node.outputs.length < totalOutputs) {
const outputIndex = node.outputs.length + 1;
node.addOutput(`${config.outputPrefix}${outputIndex}`, config.outputType);
}
if (node.setDirty) node.setDirty(true);
};
// 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) {
const originalCallback = divideByWidget.callback;
divideByWidget.callback = (value) => {
updateOutputs();
if (originalCallback) originalCallback.call(divideByWidget, value);
};
if (divideByWidget.inputEl) {
divideByWidget.inputEl.addEventListener('input', updateOutputs);
}
const observer = new MutationObserver(updateOutputs);
if (divideByWidget.element) {
observer.observe(divideByWidget.element, { attributes: true, childList: true, subtree: true });
}
node._batchDividerCleanup = () => {
observer.disconnect();
if (divideByWidget.inputEl) {
divideByWidget.inputEl.removeEventListener('input', updateOutputs);
}
divideByWidget.callback = originalCallback;
};
} catch (error) {
console.error("Error in ImageBatchDivider extension:", error);
}
const originalConfigure = node.configure;
node.configure = function(data) {
const result = originalConfigure ? originalConfigure.call(this, data) : undefined;
updateOutputs();
return result;
};
} catch (error) {
console.error(`Error in ${node.comfyClass} extension:`, error);
}
},
nodeBeforeRemove(node) {
if (node.comfyClass === "ImageBatchDivider" && node._batchDividerCleanup) {
if (BATCH_DIVIDER_NODES[node.comfyClass] && node._batchDividerCleanup) {
node._batchDividerCleanup();
}
}
});
});