Compare commits

...
Author SHA1 Message Date
Robert Wojciechowski 9345a8e71b docs: link community add-ons and forks 2026-10-12 01:20:46 +00:00
Robert Wojciechowski a630f55bd9 docs: move compressed result transfer guide out of README 2026-10-11 20:55:37 +00:00
6edd3fda38 feat(collector): optional compressed worker → master result transfer (#97)
* feat(collector): compressed worker -> master result transfer

Workers sent every frame as an uncompressed PNG (compress_level=0),
base64-encoded inside a JSON request, one request per frame. For cloud
workers this makes result download (egress) the dominant transfer cost:
a 121-frame 720p batch is ~446 MB per worker.

Add `transfer_format` / `transfer_quality` optional widgets to
Distributed Collector:

- legacy_png (default): unchanged behaviour
- png, ffv1: lossless, bit-identical master tensors to legacy_png
- jpeg, webp: lossy stills
- h264, h265, av1: lossy video of the whole batch

Compressed results are posted as multipart/form-data chunks (binary, no
base64, ~32 MB soft cap per request, configurable with
COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES). The master decodes each chunk
off the event loop and enqueues one item per frame, so collection,
ordering and audio handling are unchanged. Encoding also runs in an
executor on the worker.

Compatibility:
- /distributed/job_complete still accepts the legacy JSON envelope.
- A worker talking to an older master gets a 400 on the first chunk and
  resends everything as legacy_png.
- Missing video encoders fall back to h264 (lossy) or png (ffv1).

Measured on 121 frames at 1280x720: png 163 MB (2.7x), ffv1 128 MB
(3.5x), h264 q90 2.4 MB (186x, 40.2 dB), av1 q90 0.4 MB.

Also clips IMAGE values before 8-bit quantization on the new path so
out-of-range values no longer wrap around.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0166t8T3CVaxZkJbP2RtSpo6

* fix(collector): wire transfer options through V3 node; bound finalized payload size

Addresses review on #97.

V3 integration
- Add optional `transfer_format` / `transfer_quality` inputs to the
  registered nodes/v3.py DistributedCollector schema (legacy defaults) and
  accept/forward them in execute() to the runtime helper.
- tests/test_v3_collector_transfer.py checks the registered node's public
  schema and argument forwarding without a ComfyUI checkout; the opt-in
  real-runtime acceptance (tests/helpers/v3_runtime_check.py) now allows
  exactly these appended optional inputs, verifies their contract, and
  checks forwarding through the real executor.

Payload size bound
- encode_chunks() enforces limits on finalized bytes (after encoder flush
  and muxer finalization). Encoder lookahead kept the pre-flush buffer
  near zero, so a chunk could hold the whole batch and finalize far above
  the limit. Oversized video chunks are re-encoded with proportionally
  fewer frames; the next chunk is sized from observed bytes/frame.
- Multi-frame chunks stay within TRANSFER_MAX_CHUNK_BYTES (32 MiB); a
  single frame may use up to TRANSFER_MAX_REQUEST_BYTES (96 MiB, new,
  COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES) minus envelope reserve.
  A frame that cannot fit raises FrameTooLargeError.
- The worker encodes all chunks before sending, so an unfittable frame
  fails before anything is delivered (no partial batch on the master).
- Audio rides on the last chunk only when the request stays in bounds;
  otherwise it is sent as a final audio-only completion.
- Legacy fallback triggers only on an older master's rejection
  ("Invalid JSON payload"); other errors, e.g. size limits, raise with a
  pointer to the limit settings instead of resending as legacy_png.
- The master returns 413 when a request exceeds --max-upload-size.

Tests: codec flush-overshoot and single-frame cases, worker fallback /
nothing-sent / audio placement, and real aiohttp server tests with
client_max_size covering bounded high-entropy video, 413 reporting, and
unfittable frames.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0166t8T3CVaxZkJbP2RtSpo6

---------

Co-authored-by: ymiyad <ymiyad@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-11 20:50:23 +00:00
14 changed files with 1858 additions and 46 deletions
+20
View File
@@ -220,6 +220,26 @@ Yes, it is compatible with Docker setups, but you will need to configure your Do
---
## Community projects
Community-maintained add-ons and forks that extend ComfyUI-Distributed.
### [ComfyUI-H3MC-Distributed](https://github.com/ymiyad/ComfyUI-H3MC-Distributed) by ymiyad
Connects H3 Motion Context with ComfyUI-Distributed to generate multiple candidate takes of the next video clip across your workers. Compare the results, pick your preferred take, and use it as context for the next clip.
Each saved candidate video includes its continuation latent, so you don't need a shared folder or separate latent files. This add-on speeds up iteration through parallel candidates; it does not split a single clip across GPUs.
### [ComfyUI-Deadline-Distributed](https://github.com/doubletwisted/ComfyUI-Deadline-Distributed) by doubletwisted
A fork of ComfyUI-Distributed with Thinkbox/AWS Deadline 10 integration for render farms. Claim Deadline workers from the Distributed panel, register them automatically, and collect their results back into your master workflow.
Requires Deadline 10 and the author's ComfyUI-Deadline-Plugin.
See each repository for installation instructions, requirements, and support.
---
## Disclaimer
This software is provided "as is" without any warranties, express or implied, including merchantability, fitness for a particular purpose, or non-infringement. The developers and copyright holders are not liable for any claims, damages, or liabilities arising from the use, modification, or distribution of the software. Users are solely responsible for ensuring compliance with applicable laws and regulations and for securing their networks against unauthorized access, hacking, data breaches, or loss. The developers assume no liability for any damages or incidents resulting from misuse, improper configuration, or external threats.
+153 -24
View File
@@ -15,6 +15,7 @@ from ..utils.logging import debug_log
from ..utils.image import pil_to_tensor, ensure_contiguous
from ..utils.network import handle_api_error
from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY
from ..utils import transfer_codec
try:
from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution
except ImportError:
@@ -26,8 +27,16 @@ from .queue_request import parse_queue_request_payload
prompt_server = server.PromptServer.instance
# Canonical worker result envelope accepted by POST /distributed/job_complete:
# { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool }
# Worker result envelopes accepted by POST /distributed/job_complete:
#
# application/json (legacy, one frame per request):
# { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool,
# "audio"?: {...} }
#
# multipart/form-data (compressed, many frames per request):
# envelope: JSON { "job_id", "worker_id", "format", "start_index", "count", "width", "height",
# "is_last", "audio"?: {...} }
# media: one part per frame for still formats, or one video file for video formats
def _decode_image_sync(image_path):
@@ -270,8 +279,136 @@ async def check_file_endpoint(request):
return await handle_api_error(request, e, 500)
async def _enqueue_worker_items(multi_job_id, items):
"""Put decoded items on the job queue, waiting briefly for the master to create it.
Returns the queue size after enqueueing, or None if the job never appeared.
"""
deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
while True:
async with prompt_server.distributed_jobs_lock:
pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
if pending is not None:
for item in items:
await pending.put(item)
return pending.qsize()
if time.monotonic() > deadline:
return None
await asyncio.sleep(0.05)
def _decode_compressed_chunk_sync(envelope, parts):
"""Decode a compressed chunk into per-frame IMAGE tensors (threadpool)."""
frames = transfer_codec.decode_chunk(envelope, parts)
return [
ensure_contiguous(torch.from_numpy(transfer_codec.uint8_frame_to_tensor_array(frame)))
for frame in frames
]
def _validate_compressed_envelope(envelope):
errors = []
if not isinstance(envelope, dict):
return ["envelope: expected JSON object"]
for key in ("job_id", "worker_id"):
value = envelope.get(key)
if not isinstance(value, str) or not value.strip():
errors.append(f"{key}: expected non-empty string")
start_index = envelope.get("start_index")
if not isinstance(start_index, int) or isinstance(start_index, bool) or start_index < 0:
errors.append("start_index: expected non-negative integer")
if not isinstance(envelope.get("is_last"), bool):
errors.append("is_last: expected boolean")
try:
if transfer_codec.normalize_format(envelope.get("format")) == transfer_codec.LEGACY_FORMAT:
errors.append("format: legacy_png must be sent as JSON")
except transfer_codec.TransferCodecError as exc:
errors.append(f"format: {exc}")
audio_payload = envelope.get("audio")
if audio_payload is not None and not isinstance(audio_payload, dict):
errors.append("audio: expected object when provided")
return errors
async def _job_complete_compressed(request):
try:
form = await request.post()
except Exception as exc:
too_large = getattr(web, "HTTPRequestEntityTooLarge", None)
if too_large is not None and isinstance(exc, too_large):
return await handle_api_error(
request, f"Request body exceeds the master's upload limit (--max-upload-size): {exc}", 413)
return await handle_api_error(request, f"Invalid multipart payload: {exc}", 400)
raw_envelope = form.get("envelope")
if isinstance(raw_envelope, (bytes, bytearray)):
raw_envelope = raw_envelope.decode("utf-8")
elif hasattr(raw_envelope, "file"):
raw_envelope = raw_envelope.file.read().decode("utf-8")
try:
envelope = json.loads(raw_envelope) if isinstance(raw_envelope, str) else None
except ValueError as exc:
return await handle_api_error(request, f"envelope: invalid JSON ({exc})", 400)
if envelope is None:
return await handle_api_error(request, "envelope: missing field", 400)
errors = _validate_compressed_envelope(envelope)
if errors:
return await handle_api_error(request, errors, 400)
parts = []
for field in form.getall("media", []):
if hasattr(field, "file"):
parts.append(field.file.read())
elif isinstance(field, (bytes, bytearray)):
parts.append(bytes(field))
else:
return await handle_api_error(request, "media: expected binary file parts", 400)
loop = asyncio.get_running_loop()
try:
tensors = await loop.run_in_executor(None, _decode_compressed_chunk_sync, envelope, parts)
decoded_audio = _decode_audio_payload(envelope.get("audio")) if envelope.get("audio") is not None else None
except ValueError as exc:
return await handle_api_error(request, exc, 400)
multi_job_id = envelope["job_id"].strip()
worker_id = envelope["worker_id"].strip()
start_index = envelope["start_index"]
is_last = envelope["is_last"]
last_pos = len(tensors) - 1
items = [
{
"tensor": tensor,
"worker_id": worker_id,
"image_index": start_index + pos,
"is_last": bool(is_last and pos == last_pos),
"audio": decoded_audio if pos == last_pos else None,
}
for pos, tensor in enumerate(tensors)
]
queue_size = await _enqueue_worker_items(multi_job_id, items)
if queue_size is None:
return await handle_api_error(request, "job not initialized", 404)
debug_log(
f"job_complete received {envelope['format']} chunk - job_id: {multi_job_id}, worker: {worker_id}, "
f"frames: {start_index}..{start_index + len(items) - 1}, "
f"bytes: {sum(len(p) for p in parts)}, is_last: {is_last}, queue_size: {queue_size}"
)
return web.json_response({"status": "success", "frames": len(items)})
@server.PromptServer.instance.routes.post("/distributed/job_complete")
async def job_complete_endpoint(request):
content_type = str(getattr(request, "content_type", "") or "").lower()
if content_type.startswith("multipart/"):
try:
return await _job_complete_compressed(request)
except Exception as e:
return await handle_api_error(request, e)
try:
data = await request.json()
except Exception as exc:
@@ -316,28 +453,20 @@ async def job_complete_endpoint(request):
multi_job_id = job_id.strip()
worker_id = worker_id.strip()
pending = None
queue_size = 0
deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
while pending is None:
async with prompt_server.distributed_jobs_lock:
pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
if pending is not None:
await pending.put(
{
"tensor": tensor,
"worker_id": worker_id,
"image_index": int(batch_idx),
"is_last": is_last,
"audio": decoded_audio,
}
)
queue_size = pending.qsize()
break
if time.monotonic() > deadline:
return await handle_api_error(request, "job not initialized", 404)
await asyncio.sleep(0.05)
queue_size = await _enqueue_worker_items(
multi_job_id,
[
{
"tensor": tensor,
"worker_id": worker_id,
"image_index": int(batch_idx),
"is_last": is_last,
"audio": decoded_audio,
}
],
)
if queue_size is None:
return await handle_api_error(request, "job not initialized", 404)
debug_log(
f"job_complete received canonical envelope - job_id: {multi_job_id}, "
+38
View File
@@ -0,0 +1,38 @@
# Compressed Result Transfer
**Distributed Collector** has two optional widgets that control how workers send their images back to the master. This matters most when workers run in the cloud and you pay for data leaving it (egress).
| `transfer_format` | Type | Notes |
|---|---|---|
| `legacy_png` (default) | lossless | Original behaviour: one JSON request per frame, uncompressed PNG in base64. Works with any master version. |
| `png` | lossless | Compressed PNG, binary, several frames per request. Same tensors as `legacy_png`. |
| `ffv1` | lossless | Whole batch as one FFV1 video. Same tensors as `legacy_png`; best lossless choice for frame sequences. |
| `jpeg`, `webp` | lossy | Per-frame stills. Good for image batches. |
| `h264`, `h265`, `av1` | lossy | Whole batch as one video. By far the smallest for video frame batches. |
`transfer_quality` (1–100, default 90) applies to the lossy formats. JPEG/WebP use it directly; video codecs map it to CRF (90 → CRF 13 for h264/h265, CRF 16 for av1; 100 → CRF 0).
Measured on a 121-frame 1280×720 batch (one worker's result):
| Format | Sent | vs legacy | PSNR |
|---|---|---|---|
| `legacy_png` | 446 MB | 1× | lossless |
| `png` | 163 MB | 2.7× | lossless |
| `ffv1` | 128 MB | 3.5× | lossless |
| `jpeg` q90 | 27 MB | 17× | 42.7 dB |
| `webp` q90 | 13 MB | 33× | 40.6 dB |
| `h264` q90 | 2.4 MB | 186× | 40.2 dB |
| `h265` q90 | 1.0 MB | 434× | 40.1 dB |
| `av1` q90 | 0.4 MB | 1031× | 40.5 dB |
The test clip is a slow pan, which video codecs compress especially well; expect smaller (still large) gains on content with more motion.
Notes:
- Choose the format on the master; the value is sent to workers with the workflow.
- The master must run this version to accept compressed results. A worker sending to an older master gets a 400 on its first request and automatically resends with `legacy_png`.
- If a worker's PyAV build lacks the requested video encoder, it falls back to `h264` (lossy formats) or `png` (`ffv1`), and logs it.
- Lossy formats only affect frames produced by workers; the master's own frames are never re-encoded.
- Request size is bounded on the final encoded bytes. Frames are batched into requests of up to 32 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES`); a single large frame may go alone in a request of up to 96 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES`). Keep the latter below the master's `--max-upload-size` (default 100 MB) if you change either.
- All chunks are encoded before anything is sent. If a single frame cannot fit within the request limit, the worker fails with an error and sends nothing; use a lossy format or raise both limits.
- A rejection other than the older-master case above (for example the master's size limit, HTTP 413) is reported as an error rather than retried as `legacy_png`.
- Each worker logs what it sent, e.g. `Sent 121 image(s) as h264 in 1 request(s): 2.40 MB (raw 334.49 MB, 139.4x smaller)`.
+2
View File
@@ -74,6 +74,8 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
> ComfyUI instances running on a cloud service like Runpod.
For faster result transfers and lower bandwidth usage, see [Compressed Result Transfer](compressed-result-transfer.md).
### Deploy Cloud Worker on Runpod
📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
+193 -18
View File
@@ -12,15 +12,21 @@ from comfy.utils import ProgressBar
from ..utils.logging import debug_log, log
from ..utils.config import get_worker_timeout_seconds, load_config, is_master_delegate_only
from ..utils.constants import HEARTBEAT_INTERVAL
from ..utils.constants import HEARTBEAT_INTERVAL, TRANSFER_MAX_CHUNK_BYTES, TRANSFER_MAX_REQUEST_BYTES
from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
from ..utils.network import build_worker_url, get_client_session, probe_worker
from ..utils.audio_payload import encode_audio_payload
from ..utils.async_helpers import run_async_in_server_loop
from ..utils import transfer_codec
prompt_server = _server.PromptServer.instance
def _encode_all_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes):
"""Materialize every transfer chunk (runs in an executor)."""
return list(transfer_codec.encode_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes))
class DistributedCollectorNode:
INPUT_IS_LIST = True
EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
@@ -40,6 +46,23 @@ class DistributedCollectorNode:
"optional": {
"images": ("IMAGE",),
"audio": ("AUDIO",),
"transfer_format": (
list(transfer_codec.TRANSFER_FORMATS),
{
"default": transfer_codec.DEFAULT_TRANSFER_FORMAT,
"tooltip": transfer_codec.TRANSFER_FORMAT_TOOLTIP,
},
),
"transfer_quality": (
"INT",
{
"default": transfer_codec.DEFAULT_TRANSFER_QUALITY,
"min": 1,
"max": 100,
"step": 1,
"tooltip": transfer_codec.TRANSFER_QUALITY_TOOLTIP,
},
),
},
"hidden": {
"multi_job_id": ("STRING", {"default": ""}),
@@ -104,7 +127,7 @@ class DistributedCollectorNode:
return None
return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
def run(self, images=None, load_balance=False, 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):
def run(self, images=None, load_balance=False, 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, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
if images is not None:
images = self._normalize_images_input(images)
audio = self._normalize_audio_input(audio)
@@ -117,6 +140,8 @@ class DistributedCollectorNode:
worker_id = self._unwrap_list_input(worker_id)
pass_through = self._unwrap_list_input(pass_through)
delegate_only = self._unwrap_list_input(delegate_only)
transfer_format = self._unwrap_list_input(transfer_format)
transfer_quality = self._unwrap_list_input(transfer_quality)
remote_only_master = (
bool(multi_job_id)
@@ -147,12 +172,172 @@ class DistributedCollectorNode:
worker_batch_size,
worker_id,
delegate_only,
transfer_format,
transfer_quality,
)
)
return result
async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
async def send_batch_to_master(
self,
image_batch,
audio,
multi_job_id,
master_url,
worker_id,
transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT,
transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY,
):
"""Send an image batch, optionally with audio, or an audio-only completion."""
batch_size = 0 if image_batch is None else image_batch.shape[0]
fmt = self._resolve_transfer_format(transfer_format) if batch_size else transfer_codec.LEGACY_FORMAT
if fmt != transfer_codec.LEGACY_FORMAT:
sent = await self._send_compressed_batch(
image_batch, audio, multi_job_id, master_url, worker_id, fmt, transfer_quality
)
if sent:
return
log(
"Worker - Master did not accept compressed transfer "
"(older ComfyUI-Distributed on the master?). Falling back to legacy_png."
)
await self._send_legacy_batch(image_batch, audio, multi_job_id, master_url, worker_id)
@staticmethod
def _resolve_transfer_format(transfer_format):
try:
fmt = transfer_codec.normalize_format(transfer_format)
except transfer_codec.TransferCodecError as exc:
log(f"Worker - {exc}; using legacy_png.")
return transfer_codec.LEGACY_FORMAT
if fmt not in transfer_codec.VIDEO_FORMATS or transfer_codec.video_codec_available(fmt):
return fmt
fallback = "h264" if fmt != "ffv1" and transfer_codec.video_codec_available("h264") else "png"
log(f"Worker - Encoder for '{fmt}' is not available in this PyAV build; using '{fallback}'.")
return fallback
@staticmethod
def _build_multipart(envelope, parts):
form = aiohttp.FormData()
form.add_field("envelope", json.dumps(envelope), content_type="application/json")
for i, data in enumerate(parts):
form.add_field("media", data, filename=f"{i}.bin", content_type="application/octet-stream")
return form
# Room left in each request for the multipart envelope (ids, frame
# metadata, part headers). Audio is budgeted separately.
TRANSFER_ENVELOPE_RESERVE_BYTES = 256 * 1024
# Body of a 400 from a master without compressed-transfer support: it
# tries request.json() on the multipart body and fails.
_OLD_MASTER_REJECTION = "Invalid JSON payload"
async def _send_compressed_batch(self, image_batch, audio, multi_job_id, master_url, worker_id, fmt, quality):
"""Send the batch as compressed multipart chunks.
All chunks are encoded before anything is sent, so a frame that cannot
fit in one request fails the transfer without delivering a partial
batch. Every request stays within TRANSFER_MAX_REQUEST_BYTES.
Returns False (without having delivered anything) when the master
rejects the first chunk because it predates compressed transfer, so
the caller can retry with legacy_png. Any other rejection raises.
"""
encoded_audio = encode_audio_payload(audio)
session = await get_client_session()
url = f"{master_url}/distributed/job_complete"
loop = asyncio.get_running_loop()
frames = await loop.run_in_executor(None, transfer_codec.tensor_batch_to_uint8, image_batch)
raw_bytes = int(frames.nbytes)
total = int(frames.shape[0])
media_limit = max(1, TRANSFER_MAX_REQUEST_BYTES - self.TRANSFER_ENVELOPE_RESERVE_BYTES)
started = time.time()
# Encoding is CPU-bound (seconds for video); keep the server loop responsive.
chunks = await loop.run_in_executor(
None, _encode_all_chunks, frames, fmt, quality, min(TRANSFER_MAX_CHUNK_BYTES, media_limit), media_limit,
)
del frames
# Attach audio to the final chunk only if the request stays in bounds;
# otherwise finish with a separate audio-only completion.
audio_bytes = len(json.dumps(encoded_audio)) if encoded_audio is not None else 0
last_media = sum(len(p) for p in chunks[-1]["parts"])
audio_inline = encoded_audio is not None and last_media + audio_bytes <= media_limit
audio_separate = encoded_audio is not None and not audio_inline
sent_bytes = 0
for index, chunk in enumerate(chunks):
parts = chunk.pop("parts")
is_final_chunk = index == len(chunks) - 1
envelope = {
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"is_last": bool(is_final_chunk and not audio_separate),
**chunk,
}
if is_final_chunk and audio_inline:
envelope["audio"] = encoded_audio
chunk_bytes = sum(len(p) for p in parts)
timeout_seconds = min(600, 60 + chunk_bytes // (256 * 1024))
try:
async with session.post(
url,
data=self._build_multipart(envelope, parts),
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
) as response:
if response.status >= 400:
body = await response.text()
if index == 0 and response.status == 400 and self._OLD_MASTER_REJECTION in body:
debug_log(f"Worker - Master rejected compressed chunk: {body[:300]}")
return False
raise RuntimeError(
f"Master rejected {fmt} chunk {index + 1}/{len(chunks)} "
f"({chunk_bytes} bytes) with HTTP {response.status}: {body[:300]}. "
"If this is a size limit, make COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES "
"on the worker smaller than the master's --max-upload-size."
)
except Exception as e:
log(f"Worker - Failed to send {fmt} chunk to master: {e}")
debug_log(f"Worker - Full error details: URL={url}")
raise
sent_bytes += chunk_bytes
requests = len(chunks)
if audio_separate:
await self._post_json_completion(session, url, {
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": 0,
"audio": encoded_audio,
"is_last": True,
})
requests += 1
log(
f"Worker - Sent {total} image(s) as {fmt} in {requests} request(s): "
f"{sent_bytes / 1e6:.2f} MB (raw {raw_bytes / 1e6:.2f} MB, "
f"{transfer_codec.describe_ratio(sent_bytes, raw_bytes)} smaller) in {time.time() - started:.1f}s"
)
return True
@staticmethod
async def _post_json_completion(session, url, payload):
timeout_seconds = 60 if "image" in payload else 600
try:
async with session.post(
url,
json=payload,
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
) as response:
response.raise_for_status()
except Exception as e:
media_type = "image/audio" if "image" in payload else "audio-only"
log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
debug_log(f"Worker - Full error details: URL={url}")
raise
async def _send_legacy_batch(self, image_batch, audio, multi_job_id, master_url, worker_id):
"""Original transfer: one JSON request per frame with a base64 PNG."""
encoded_audio = encode_audio_payload(audio)
session = await get_client_session()
url = f"{master_url}/distributed/job_complete"
@@ -189,19 +374,7 @@ class DistributedCollectorNode:
payloads.append(payload)
for payload in payloads:
timeout_seconds = 60 if "image" in payload else 600
try:
async with session.post(
url,
json=payload,
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
) as response:
response.raise_for_status()
except Exception as e:
media_type = "image/audio" if "image" in payload else "audio-only"
log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
debug_log(f"Worker - Full error details: URL={url}")
raise
await self._post_json_completion(session, url, payload)
def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
"""Combine audio from master and workers into a single audio output.
@@ -319,12 +492,14 @@ class DistributedCollectorNode:
return ensure_contiguous(fallback_images)
return None
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
if is_worker:
# Worker mode: send images and audio to master in a single batch
image_count = 0 if images is None else images.shape[0]
debug_log(f"Worker - Job {multi_job_id} complete. Sending {image_count} image(s) to master")
await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
await self.send_batch_to_master(
images, audio, multi_job_id, master_url, worker_id, transfer_format, transfer_quality
)
return (images, audio if audio is not None else self.EMPTY_AUDIO)
else:
delegate_mode = delegate_only or is_master_delegate_only()
+13 -3
View File
@@ -8,6 +8,7 @@ from comfy_api.v0_0_2 import io
from . import utilities as _utilities
from .collector import DistributedCollectorNode as _CollectorRuntime
from ..utils import transfer_codec as _transfer_codec
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
@@ -121,7 +122,13 @@ class DistributedCollector(io.ComfyNode):
node_id='DistributedCollector', display_name='Distributed Collector', category='image',
inputs=[io.Boolean.Input('load_balance', default=False,
tooltip='Run this workflow on one least-busy participant (master included when participating).'),
io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True)],
io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True),
io.Combo.Input('transfer_format', options=list(_transfer_codec.TRANSFER_FORMATS),
default=_transfer_codec.DEFAULT_TRANSFER_FORMAT, optional=True,
tooltip=_transfer_codec.TRANSFER_FORMAT_TOOLTIP),
io.Int.Input('transfer_quality', default=_transfer_codec.DEFAULT_TRANSFER_QUALITY,
min=1, max=100, step=1, optional=True,
tooltip=_transfer_codec.TRANSFER_QUALITY_TOOLTIP)],
outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
is_input_list=True, accept_all_inputs=True,
)
@@ -129,12 +136,15 @@ class DistributedCollector(io.ComfyNode):
@classmethod
def execute(cls, images=None, load_balance=False, 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):
worker_id='', pass_through=False, delegate_only=False,
transfer_format=_transfer_codec.DEFAULT_TRANSFER_FORMAT,
transfer_quality=_transfer_codec.DEFAULT_TRANSFER_QUALITY):
return io.NodeOutput(*_CollectorRuntime().run(
images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
worker_batch_size=worker_batch_size, worker_id=worker_id,
pass_through=pass_through, delegate_only=delegate_only))
pass_through=pass_through, delegate_only=delegate_only,
transfer_format=transfer_format, transfer_quality=transfer_quality))
class UltimateSDUpscaleDistributed(io.ComfyNode):
+153
View File
@@ -4,6 +4,8 @@ import types
import unittest
import asyncio
import base64
import io
import json
from dataclasses import dataclass
from pathlib import Path
from unittest.mock import AsyncMock, patch
@@ -26,6 +28,17 @@ class _FakeRequest:
return self._payload
def _register_real_transfer_codec(package_name, utils_pkg):
"""Load the real utils/transfer_codec.py into the synthetic test package."""
codec_path = Path(__file__).resolve().parents[2] / "utils" / "transfer_codec.py"
codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
codec_module = importlib.util.module_from_spec(codec_spec)
sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
codec_spec.loader.exec_module(codec_module)
utils_pkg.transfer_codec = codec_module
return codec_module
def _load_job_routes_module():
module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py"
package_name = "dist_api_queue_testpkg"
@@ -46,6 +59,7 @@ def _load_job_routes_module():
utils_pkg = types.ModuleType(f"{package_name}.utils")
utils_pkg.__path__ = []
sys.modules[f"{package_name}.utils"] = utils_pkg
_register_real_transfer_codec(package_name, utils_pkg)
# aiohttp.web stub
created_aiohttp_stub = False
@@ -394,5 +408,144 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
job_routes._decode_audio_payload(payload)
class _FakeFileField:
def __init__(self, data):
self.file = io.BytesIO(data)
class _FakeMultipartRequest:
content_type = "multipart/form-data"
def __init__(self, envelope, parts):
from multidict import MultiDict
fields = [("envelope", envelope if isinstance(envelope, str) else json.dumps(envelope))]
fields += [("media", _FakeFileField(part)) for part in parts]
self._form = MultiDict(fields)
async def post(self):
return self._form
async def json(self):
raise ValueError("not JSON")
class JobCompleteCompressedPayloadTests(unittest.IsolatedAsyncioTestCase):
codec = job_routes.transfer_codec
def setUp(self):
self.queue = asyncio.Queue()
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
job_routes.prompt_server.distributed_pending_jobs = {"job-z": self.queue}
def _frames(self, count=4, height=20, width=30):
rng = np.random.default_rng(1)
return rng.integers(0, 256, (count, height, width, 3), dtype=np.uint8)
def _request_for(self, frames, fmt, *, start_offset=0, is_last=True, audio=None, **overrides):
chunk = next(self.codec.encode_chunks(frames, fmt, 90, 10**9))
parts = chunk.pop("parts")
envelope = {
"job_id": "job-z",
"worker_id": "worker-1",
"is_last": is_last,
**chunk,
"start_index": chunk["start_index"] + start_offset,
}
if audio is not None:
envelope["audio"] = audio
envelope.update(overrides)
return _FakeMultipartRequest(envelope, parts)
async def _drain(self):
items = []
while not self.queue.empty():
items.append(await self.queue.get())
return items
async def test_lossless_chunk_enqueues_one_item_per_frame_bit_exact(self):
frames = self._frames()
response = await job_routes.job_complete_endpoint(self._request_for(frames, "ffv1", start_offset=10))
self.assertEqual(response.status, 200)
items = await self._drain()
self.assertEqual([item["image_index"] for item in items], [10, 11, 12, 13])
self.assertEqual([item["is_last"] for item in items], [False, False, False, True])
self.assertEqual({item["worker_id"] for item in items}, {"worker-1"})
for item, frame in zip(items, frames):
self.assertEqual(tuple(item["tensor"].shape), (1, 20, 30, 3))
self.assertEqual(item["tensor"].dtype, torch.float32)
expected = frame.astype(np.float32) / 255.0
self.assertTrue(np.array_equal(item["tensor"][0].numpy(), expected))
async def test_non_final_chunk_never_marks_worker_done(self):
response = await job_routes.job_complete_endpoint(
self._request_for(self._frames(count=2), "png", is_last=False)
)
self.assertEqual(response.status, 200)
self.assertFalse(any(item["is_last"] for item in await self._drain()))
async def test_audio_is_attached_only_to_final_frame(self):
audio = JobCompleteAudioPayloadTests._encoded_audio_payload(None)
response = await job_routes.job_complete_endpoint(
self._request_for(self._frames(count=3), "jpeg", audio=audio)
)
self.assertEqual(response.status, 200)
items = await self._drain()
self.assertEqual([item["audio"] is not None for item in items], [False, False, True])
self.assertEqual(tuple(items[-1]["audio"]["waveform"].shape), (1, 2, 4))
async def test_rejects_frame_count_mismatch(self):
response = await job_routes.job_complete_endpoint(
self._request_for(self._frames(count=2), "ffv1", count=5)
)
self.assertEqual(response.status, 400)
self.assertIn("frames", response.payload["message"])
self.assertTrue(self.queue.empty())
async def test_rejects_bad_envelope(self):
response = await job_routes.job_complete_endpoint(
self._request_for(self._frames(count=1), "png", format="gif", is_last="yes")
)
self.assertEqual(response.status, 400)
self.assertIn("format", response.payload["message"])
self.assertIn("is_last", response.payload["message"])
async def test_rejects_invalid_envelope_json(self):
response = await job_routes.job_complete_endpoint(_FakeMultipartRequest("{not json", [b"x"]))
self.assertEqual(response.status, 400)
self.assertIn("envelope", response.payload["message"])
async def test_unknown_job_returns_404(self):
job_routes.prompt_server.distributed_pending_jobs = {}
with patch.object(job_routes, "JOB_INIT_GRACE_PERIOD", 0.0):
response = await job_routes.job_complete_endpoint(self._request_for(self._frames(count=1), "png"))
self.assertEqual(response.status, 404)
async def test_legacy_json_payload_still_accepted(self):
request = _FakeRequest(
{
"job_id": "job-z",
"worker_id": "worker-1",
"batch_idx": 3,
"image": "data:image/png;base64,AAAA",
"is_last": True,
}
)
with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
response = await job_routes.job_complete_endpoint(request)
self.assertEqual(response.status, 200)
items = await self._drain()
self.assertEqual(len(items), 1)
self.assertEqual((items[0]["image_index"], items[0]["tensor"]), (3, "tensor-data"))
if __name__ == "__main__":
unittest.main()
+216
View File
@@ -0,0 +1,216 @@
"""End-to-end: real worker send path -> real aiohttp server -> job_complete endpoint."""
import asyncio
import importlib.util
import sys
from pathlib import Path
import numpy as np
import pytest
import torch
_TESTS_DIR = Path(__file__).resolve().parents[1]
def _load_from(path, name):
spec = importlib.util.spec_from_file_location(name, path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def _load_modules():
real_aiohttp = sys.modules["aiohttp"] # imported inside the isolated test
queue_tests = _load_from(_TESTS_DIR / "api" / "test_distributed_queue.py", "dist_transfer_it_queue")
collector_tests = _load_from(_TESTS_DIR / "test_collector_list_inputs.py", "dist_transfer_it_collector")
job_routes = queue_tests._load_job_routes_module()
collector = collector_tests._load_collector_module()
# Error responses must be real HTTP responses here (the unit-test harness
# returns plain objects); mirror utils.network.handle_api_error.
from aiohttp import web
async def handle_api_error(_request, error, status=500):
if isinstance(error, list):
return web.json_response({"errors": [str(item) for item in error]}, status=status)
return web.json_response({"error": str(error)}, status=status)
job_routes.handle_api_error = handle_api_error
# The collector harness stubs aiohttp; this test needs the real client.
sys.modules["aiohttp"] = real_aiohttp
collector.aiohttp = real_aiohttp
return job_routes, collector
@pytest.fixture(autouse=True)
def _isolate_sys_modules():
"""The loaders below install stub modules (server, comfy, aiohttp...); undo that."""
saved = dict(sys.modules)
# Earlier test modules may have left a stub `aiohttp` behind; drop it so the
# real package is imported for this test (restored afterwards).
for name in [n for n in sys.modules if n == "aiohttp" or n.startswith("aiohttp.")]:
if getattr(sys.modules[name], "__file__", None) is None:
del sys.modules[name]
yield
sys.modules.clear()
sys.modules.update(saved)
async def _run_transfer(job_routes, collector_module, images, fmt, *, client_max_size,
request_limit, chunk_target, audio=None):
"""Send `images` from a real worker node to the real completion route.
Returns (queued items, sizes of every request body the server received,
exception raised by the worker or None).
"""
import aiohttp
from aiohttp import web
queue = asyncio.Queue()
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
body_sizes = []
@web.middleware
async def record_size(request, handler):
body_sizes.append(request.content_length)
return await handler(request)
app = web.Application(client_max_size=client_max_size, middlewares=[record_size])
app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, "127.0.0.1", 0)
await site.start()
port = site._server.sockets[0].getsockname()[1]
error = None
try:
async with aiohttp.ClientSession() as session:
async def _get_session():
return session
collector_module.get_client_session = _get_session
collector_module.TRANSFER_MAX_REQUEST_BYTES = request_limit
collector_module.TRANSFER_MAX_CHUNK_BYTES = chunk_target
node = collector_module.DistributedCollectorNode()
try:
await node.send_batch_to_master(
images, audio, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
)
except Exception as exc: # surfaced to the caller for assertions
error = exc
finally:
await runner.cleanup()
items = []
while not queue.empty():
items.append(queue.get_nowait())
return items, body_sizes, error
def _noise_images(count, height=360, width=640, seed=97):
return torch.from_numpy(np.random.default_rng(seed).random((count, height, width, 3), dtype=np.float32))
@pytest.mark.parametrize("fmt", ["h264", "ffv1"])
def test_high_entropy_video_stays_under_server_upload_limit(fmt):
"""Encoder flush used to push finalized chunks past the limit (HTTP 400)."""
pytest.importorskip("aiohttp")
job_routes, collector_module = _load_modules()
images = _noise_images(24)
limit = 2 * 1024 * 1024
reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
items, body_sizes, error = asyncio.run(_run_transfer(
job_routes, collector_module, images, fmt,
client_max_size=limit, request_limit=limit, chunk_target=limit - reserve))
assert error is None, error
assert len(body_sizes) > 1 and max(body_sizes) <= limit
assert [item["image_index"] for item in items] == list(range(24))
assert [item["is_last"] for item in items] == [False] * 23 + [True]
def test_server_rejection_for_size_is_reported_not_retried_as_legacy():
"""A worker configured above the master's limit fails loudly instead of
silently resending the whole batch as uncompressed PNG."""
pytest.importorskip("aiohttp")
job_routes, collector_module = _load_modules()
images = _noise_images(8)
items, body_sizes, error = asyncio.run(_run_transfer(
job_routes, collector_module, images, "ffv1",
client_max_size=1024 * 1024, request_limit=8 * 1024 * 1024, chunk_target=4 * 1024 * 1024))
assert isinstance(error, RuntimeError)
assert "HTTP 413" in str(error) and "--max-upload-size" in str(error)
assert len(body_sizes) == 1
assert items == []
def test_frame_that_cannot_fit_sends_nothing():
pytest.importorskip("aiohttp")
job_routes, collector_module = _load_modules()
images = _noise_images(3)
reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
items, body_sizes, error = asyncio.run(_run_transfer(
job_routes, collector_module, images, "png",
client_max_size=100 * 1024 * 1024, request_limit=reserve + 64 * 1024, chunk_target=64 * 1024))
assert isinstance(error, collector_module.transfer_codec.FrameTooLargeError)
assert body_sizes == []
assert items == []
@pytest.mark.parametrize("fmt", ["png", "ffv1", "h264", "jpeg"])
def test_worker_to_master_over_real_http(fmt):
# Import here, not at module level: other test harnesses only install their
# aiohttp stubs when the real package has not been imported yet.
aiohttp = pytest.importorskip("aiohttp")
from aiohttp import web
job_routes, collector_module = _load_modules()
images = torch.rand(5, 24, 36, 3)
async def scenario():
queue = asyncio.Queue()
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
app = web.Application()
app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, "127.0.0.1", 0)
await site.start()
port = site._server.sockets[0].getsockname()[1]
async with aiohttp.ClientSession() as session:
async def _get_session():
return session
collector_module.get_client_session = _get_session
collector_module.TRANSFER_MAX_CHUNK_BYTES = 2000 # several requests for still formats
node = collector_module.DistributedCollectorNode()
await node.send_batch_to_master(
images, None, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
)
await runner.cleanup()
items = []
while not queue.empty():
items.append(queue.get_nowait())
return items
items = asyncio.run(scenario())
assert [item["image_index"] for item in items] == [0, 1, 2, 3, 4]
assert [item["is_last"] for item in items] == [False, False, False, False, True]
expected = np.clip(255.0 * images.numpy(), 0, 255).astype(np.uint8).astype(np.float32) / 255.0
got = torch.cat([item["tensor"] for item in items]).numpy()
assert got.shape == expected.shape
if fmt in ("png", "ffv1"):
assert np.array_equal(got, expected)
else:
assert np.abs(got - expected).mean() < 0.25 # random noise: only sanity-check lossy codecs
+38 -1
View File
@@ -30,6 +30,19 @@ import torch
BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text())
# Inputs intentionally added after the V1 baseline was captured. They must be
# optional, appended after the baseline inputs, and keep legacy behaviour as
# their default so saved workflows are unaffected.
ADDED_OPTIONAL_INPUTS = {
'DistributedCollector': {
'transfer_format': [
['legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1'],
{'default': 'legacy_png'},
],
'transfer_quality': ['INT', {'default': 90, 'min': 1, 'max': 100, 'step': 1}],
},
}
def normalized_input(value):
kind = value[0]
@@ -134,6 +147,22 @@ async def check_execution(mapping, module, prompt_server, asset_manager):
assert torch.equal(result[0].result[0], images[:3])
assert 'v3-aggregate' not in prompt_server.distributed_pending_jobs
# The registered node must forward the transfer options to the runtime
# helper, and default to the legacy transfer when they are not set.
collector_runtime = sys.modules[module.__name__ + '.nodes.collector'].DistributedCollectorNode
forwarded = []
def fake_run(self, **kwargs):
forwarded.append(kwargs)
return (kwargs['images'], kwargs['audio'])
with patch.object(collector_runtime, 'run', fake_run):
await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True,
'transfer_format': 'av1', 'transfer_quality': 70})
await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True})
# Set values arrive list-wrapped (is_input_list); unset ones use the
# execute() defaults. The runtime helper unwraps both forms.
assert [(kw['transfer_format'], kw['transfer_quality']) for kw in forwarded] == [
(['av1'], [70]), ('legacy_png', 90)], forwarded
# Prove V3's per-execution class clones do not leak helper instance state;
# replace only the GPU/model boundary, not the input parser or executor.
upscale_runtime = sys.modules[module.__name__ + '.nodes.distributed_upscale'].UltimateSDUpscaleDistributed
@@ -221,7 +250,15 @@ async def main():
new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
for group in ('required', 'optional'):
old_inputs = old['input'].get(group, {})
new_inputs = new['input'].get(group, {})
new_inputs = dict(new['input'].get(group, {}))
if group == 'optional':
for name, expected in ADDED_OPTIONAL_INPUTS.get(node_id, {}).items():
assert name in new_inputs, (node_id, name, 'added input missing')
actual = normalized_input(new_inputs.pop(name))
kind, opts = normalized_input(expected)
assert actual[0] == kind, (node_id, name, actual)
for key, value in opts.items():
assert actual[1].get(key) == value, (node_id, name, key, actual)
assert list(old_inputs) == list(new_inputs), (node_id, group, 'input order')
for name, original in old_inputs.items():
assert normalized_input(original) == normalized_input(new_inputs[name]), (node_id, name, original, new_inputs[name])
+266
View File
@@ -7,6 +7,17 @@ from pathlib import Path
import torch
def _register_real_transfer_codec(package_name, utils_pkg):
"""Load the real utils/transfer_codec.py into the synthetic test package."""
codec_path = Path(__file__).resolve().parents[1] / "utils" / "transfer_codec.py"
codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
codec_module = importlib.util.module_from_spec(codec_spec)
sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
codec_spec.loader.exec_module(codec_module)
utils_pkg.transfer_codec = codec_module
return codec_module
def _load_collector_module():
module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
package_name = "dist_collector_list_testpkg"
@@ -26,6 +37,7 @@ def _load_collector_module():
utils_pkg = types.ModuleType(f"{package_name}.utils")
utils_pkg.__path__ = []
sys.modules[f"{package_name}.utils"] = utils_pkg
_register_real_transfer_codec(package_name, utils_pkg)
class _Routes:
def post(self, _path):
@@ -86,6 +98,8 @@ def _load_collector_module():
constants_module = types.ModuleType(f"{package_name}.utils.constants")
constants_module.HEARTBEAT_INTERVAL = 1.0
constants_module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024 * 1024
constants_module.TRANSFER_MAX_REQUEST_BYTES = 96 * 1024 * 1024
sys.modules[f"{package_name}.utils.constants"] = constants_module
image_module = types.ModuleType(f"{package_name}.utils.image")
@@ -362,3 +376,255 @@ def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
assert images is None
assert combined_audio["sample_rate"] == 48000
assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])
# --------------------------------------------------------------------------- #
# Compressed worker -> master transfer
# --------------------------------------------------------------------------- #
class _RecordingResponse:
def __init__(self, status=200, text=""):
self.status = status
self._text = text
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
async def text(self):
return self._text
def raise_for_status(self):
if self.status >= 400:
raise RuntimeError(f"HTTP {self.status}")
class _RecordingSession:
def __init__(self, statuses=None, text="Invalid JSON payload"):
self.calls = []
self._statuses = list(statuses or [])
self._text = text
def post(self, url, timeout, json=None, data=None):
self.calls.append({"url": url, "json": json, "data": data, "timeout": timeout.total})
status = self._statuses.pop(0) if self._statuses else 200
return _RecordingResponse(status=status, text=self._text)
def _install_fake_transport(module, session):
async def _fake_get_client_session():
return session
module.get_client_session = _fake_get_client_session
# Record the multipart content instead of building a real aiohttp.FormData.
module.DistributedCollectorNode._build_multipart = staticmethod(
lambda envelope, parts: {"envelope": envelope, "parts": list(parts)}
)
def _run_worker(collector, images, audio=None, **kwargs):
return collector.run(
images=[images],
audio=[audio],
multi_job_id=["job-c"],
is_worker=[True],
master_url=["http://master"],
worker_id=["worker-a"],
**{key: [value] for key, value in kwargs.items()},
)
def test_collector_exposes_transfer_widgets_with_legacy_default():
module = _load_collector_module()
optional = module.DistributedCollectorNode.INPUT_TYPES()["optional"]
formats, options = optional["transfer_format"]
assert "legacy_png" in formats and "h264" in formats and "ffv1" in formats
assert options["default"] == "legacy_png"
assert optional["transfer_quality"][0] == "INT"
def test_worker_sends_video_batch_as_one_multipart_chunk_with_audio_on_last():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
module.encode_audio_payload = lambda value: None if value is None else {"encoded": True}
images = torch.rand(6, 32, 48, 3)
_run_worker(
collector,
images,
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
transfer_format="ffv1",
transfer_quality=90,
)
assert len(session.calls) == 1
call = session.calls[0]
assert call["url"] == "http://master/distributed/job_complete"
assert call["json"] is None
envelope = call["data"]["envelope"]
assert envelope["format"] == "ffv1"
assert (envelope["start_index"], envelope["count"]) == (0, 6)
assert (envelope["width"], envelope["height"]) == (48, 32)
assert envelope["is_last"] is True
assert envelope["audio"] == {"encoded": True}
assert envelope["job_id"] == "job-c" and envelope["worker_id"] == "worker-a"
assert len(call["data"]["parts"]) == 1
def test_worker_marks_only_final_chunk_as_last():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
module.TRANSFER_MAX_CHUNK_BYTES = 1 # force one frame per request
_run_worker(collector, torch.rand(3, 8, 8, 3), transfer_format="png")
envelopes = [call["data"]["envelope"] for call in session.calls]
assert [(e["start_index"], e["count"], e["is_last"]) for e in envelopes] == [
(0, 1, False),
(1, 1, False),
(2, 1, True),
]
assert "audio" not in envelopes[-1]
def test_worker_falls_back_to_legacy_png_when_master_rejects_multipart():
from PIL import Image
import numpy as np
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession(statuses=[400])
_install_fake_transport(module, session)
module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
(255 * tensor[idx].cpu().numpy()).astype(np.uint8)
)
_run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="h264")
assert session.calls[0]["data"] is not None # compressed attempt
legacy = [call["json"] for call in session.calls[1:]]
assert [payload["batch_idx"] for payload in legacy] == [0, 1]
assert all(payload["image"].startswith("data:image/png;base64,") for payload in legacy)
assert [payload["is_last"] for payload in legacy] == [False, True]
def test_worker_uses_legacy_path_for_unknown_transfer_format():
from PIL import Image
import numpy as np
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
(255 * tensor[idx].cpu().numpy()).astype(np.uint8)
)
_run_worker(collector, torch.rand(1, 8, 8, 3), transfer_format="gif")
assert len(session.calls) == 1
assert session.calls[0]["json"]["image"].startswith("data:image/png;base64,")
def test_worker_falls_back_when_video_encoder_missing():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
codec = module.transfer_codec
original = codec.video_codec_available
codec.video_codec_available = lambda fmt: fmt == "h264" and original("h264")
try:
_run_worker(collector, torch.rand(2, 16, 16, 3), transfer_format="av1")
finally:
codec.video_codec_available = original
assert session.calls[0]["data"]["envelope"]["format"] == "h264"
def test_worker_does_not_fall_back_on_non_legacy_rejection():
import pytest
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession(statuses=[400], text="Invalid multipart payload: Request Entity Too Large")
_install_fake_transport(module, session)
with pytest.raises(RuntimeError, match="MAX_REQUEST_BYTES"):
_run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="png")
# One compressed attempt, no legacy resend of the batch.
assert len(session.calls) == 1
assert session.calls[0]["json"] is None
def test_worker_sends_nothing_when_a_frame_cannot_fit():
import numpy as np
import pytest
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
module.TRANSFER_MAX_REQUEST_BYTES = reserve + 32 * 1024
module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024
# Frame 0 is tiny; frame 2 is noise that cannot fit in 32 KiB as PNG.
images = torch.zeros(3, 128, 128, 3)
images[2] = torch.from_numpy(np.random.default_rng(1).random((128, 128, 3), dtype=np.float32))
with pytest.raises(module.transfer_codec.FrameTooLargeError):
_run_worker(collector, images, transfer_format="png")
assert session.calls == []
def test_worker_sends_audio_separately_when_it_would_overflow_last_request():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
module.TRANSFER_MAX_REQUEST_BYTES = reserve + 64 * 1024
module.TRANSFER_MAX_CHUNK_BYTES = 64 * 1024
module.encode_audio_payload = lambda value: None if value is None else {"data": "x" * (128 * 1024)}
_run_worker(
collector,
torch.zeros(2, 16, 16, 3),
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
transfer_format="png",
)
multipart = [call["data"]["envelope"] for call in session.calls if call["data"] is not None]
assert [e["is_last"] for e in multipart] == [False]
assert all("audio" not in e for e in multipart)
final = session.calls[-1]["json"]
assert final["is_last"] is True and final["audio"]["data"].startswith("x")
assert "image" not in final
def test_worker_attaches_audio_to_last_chunk_when_it_fits():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
module.encode_audio_payload = lambda value: None if value is None else {"data": "small"}
_run_worker(
collector,
torch.zeros(2, 16, 16, 3),
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
transfer_format="png",
)
assert len(session.calls) == 1
envelope = session.calls[0]["data"]["envelope"]
assert envelope["is_last"] is True and envelope["audio"] == {"data": "small"}
+219
View File
@@ -0,0 +1,219 @@
import importlib.util
import io
import sys
from pathlib import Path
import numpy as np
import pytest
import torch
from PIL import Image
def _load_codec():
path = Path(__file__).resolve().parents[1] / "utils" / "transfer_codec.py"
spec = importlib.util.spec_from_file_location("dist_transfer_codec_test", path)
module = importlib.util.module_from_spec(spec)
sys.modules["dist_transfer_codec_test"] = module
spec.loader.exec_module(module)
return module
tc = _load_codec()
def _smooth_frames(count=4, height=67, width=91):
"""Smooth moving gradient, closer to generated content than random noise."""
yy, xx = np.mgrid[0:height, 0:width].astype(np.float32)
frames = []
for i in range(count):
r = 0.5 + 0.5 * np.sin((xx + 3 * i) / 13.0)
g = 0.5 + 0.5 * np.cos((yy - 2 * i) / 11.0)
b = (xx + yy + 5 * i) / (width + height + 5 * count)
frames.append(np.stack([r, g, b], axis=-1))
return torch.from_numpy(np.stack(frames)).float()
def _legacy_master_tensor(frame_tensor):
"""What the unmodified pipeline delivered: tensor_to_pil -> PNG -> pil_to_tensor."""
img = Image.fromarray((255 * frame_tensor.numpy()).astype(np.uint8))
buf = io.BytesIO()
img.save(buf, format="PNG", compress_level=0)
with Image.open(io.BytesIO(buf.getvalue())) as decoded:
return np.array(decoded.convert("RGB")).astype(np.float32) / 255.0
def _round_trip(frames_u8, fmt, quality=90, max_chunk_bytes=10**9):
out = []
chunks = list(tc.encode_chunks(frames_u8, fmt, quality, max_chunk_bytes))
for chunk in chunks:
envelope = {k: v for k, v in chunk.items() if k != "parts"}
out.extend(tc.decode_chunk(envelope, chunk["parts"]))
return chunks, out
def _psnr(a, b):
mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
return float("inf") if mse == 0 else 10 * np.log10(255.0**2 / mse)
@pytest.mark.parametrize("fmt", ["png", "ffv1"])
def test_lossless_formats_reproduce_legacy_master_tensors_exactly(fmt):
images = _smooth_frames()
frames = tc.tensor_batch_to_uint8(images)
_, decoded = _round_trip(frames, fmt)
assert len(decoded) == images.shape[0]
for i, frame in enumerate(decoded):
got = tc.uint8_frame_to_tensor_array(frame)[0]
assert np.array_equal(got, _legacy_master_tensor(images[i]))
@pytest.mark.parametrize("fmt", ["jpeg", "webp", "h264", "h265", "av1"])
def test_lossy_formats_round_trip_odd_sizes_with_reasonable_quality(fmt):
if fmt in tc.VIDEO_FORMATS and not tc.video_codec_available(fmt):
pytest.skip(f"{fmt} encoder not available in this PyAV build")
frames = tc.tensor_batch_to_uint8(_smooth_frames())
_, decoded = _round_trip(frames, fmt, quality=90)
assert len(decoded) == frames.shape[0]
for i, frame in enumerate(decoded):
assert frame.shape == frames[i].shape
assert frame.dtype == np.uint8
assert _psnr(frame, frames[i]) > 30
def test_h265_handles_frames_smaller_than_ctu():
if not tc.video_codec_available("h265"):
pytest.skip("h265 encoder not available")
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=3, height=17, width=23))
_, decoded = _round_trip(frames, "h265")
assert [f.shape for f in decoded] == [(17, 23, 3)] * 3
def test_image_formats_split_chunks_by_byte_cap_and_keep_order():
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=5))
one = len(tc.encode_image_frame(frames[0], "png", 90))
chunks, decoded = _round_trip(frames, "png", max_chunk_bytes=one * 2 + 1)
assert len(chunks) >= 3
assert [c["start_index"] for c in chunks] == sorted(c["start_index"] for c in chunks)
assert sum(c["count"] for c in chunks) == 5
assert all(len(c["parts"]) == c["count"] for c in chunks)
assert np.array_equal(np.stack(decoded), frames)
def test_video_formats_split_chunks_by_frame_cap(monkeypatch):
monkeypatch.setattr(tc, "MAX_VIDEO_CHUNK_FRAMES", 2)
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=5))
chunks, decoded = _round_trip(frames, "ffv1")
assert [(c["start_index"], c["count"]) for c in chunks] == [(0, 2), (2, 2), (4, 1)]
assert all(len(c["parts"]) == 1 for c in chunks)
assert np.array_equal(np.stack(decoded), frames)
def test_tensor_batch_to_uint8_clips_and_normalizes_channels():
rgba = torch.tensor([[[[1.5, -0.2, 0.5, 0.3]]]])
assert tc.tensor_batch_to_uint8(rgba).tolist() == [[[[255, 0, 127]]]]
gray = torch.full((1, 2, 2, 1), 0.5)
out = tc.tensor_batch_to_uint8(gray)
assert out.shape == (1, 2, 2, 3)
assert (out == 127).all()
def test_quality_to_crf_mapping():
assert tc.quality_to_crf(100, 51) == 0
assert tc.quality_to_crf(90, 51) == 13
assert tc.quality_to_crf(80, 51) == 19
assert tc.quality_to_crf(90, 63) == 16
assert tc.quality_to_crf(1, 51) == 51
assert tc.quality_to_crf("bogus", 51) == tc.quality_to_crf(tc.DEFAULT_TRANSFER_QUALITY, 51)
def test_normalize_format_rejects_unknown_values():
assert tc.normalize_format(" H264 ") == "h264"
assert tc.normalize_format(None) == tc.DEFAULT_TRANSFER_FORMAT
with pytest.raises(tc.TransferCodecError):
tc.normalize_format("gif")
def test_decode_chunk_validates_envelope_against_payload():
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=2))
chunk = next(tc.encode_chunks(frames, "png", 90, 10**9))
envelope = {k: v for k, v in chunk.items() if k != "parts"}
with pytest.raises(tc.TransferCodecError, match="image parts"):
tc.decode_chunk({**envelope, "count": 3}, chunk["parts"])
with pytest.raises(tc.TransferCodecError, match="expected"):
tc.decode_chunk({**envelope, "width": 10}, chunk["parts"])
with pytest.raises(tc.TransferCodecError, match="count"):
tc.decode_chunk({**envelope, "count": 0}, chunk["parts"])
with pytest.raises(tc.TransferCodecError):
tc.decode_chunk(envelope, [b"not an image", b"nope"])
video = next(tc.encode_chunks(frames, "ffv1", 90, 10**9))
video_env = {k: v for k, v in video.items() if k != "parts"}
with pytest.raises(tc.TransferCodecError, match="frames"):
tc.decode_chunk({**video_env, "count": 5}, video["parts"])
def test_legacy_format_is_not_chunk_encoded():
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=1))
with pytest.raises(tc.TransferCodecError):
list(tc.encode_chunks(frames, "legacy_png", 90, 10**9))
# --------------------------------------------------------------------------- #
# Size bounds on finalized payloads
# --------------------------------------------------------------------------- #
def _noise_frames(count=30, height=360, width=640, seed=97):
"""High-entropy frames: worst case for every codec, and enough encoder
lookahead that pre-flush buffer sizes badly underestimate final sizes."""
return np.random.default_rng(seed).integers(0, 256, size=(count, height, width, 3), dtype=np.uint8)
@pytest.mark.parametrize("fmt", ["h264", "ffv1"])
def test_video_chunks_are_bounded_after_encoder_flush(fmt):
frames = _noise_frames()
limit = 1024 * 1024
# Sanity: one unconstrained encode of this batch is far above the limit,
# so the bound below is actually being enforced.
unbounded = next(tc.encode_chunks(frames, fmt, 90, 10**12))
assert len(unbounded["parts"][0]) > 3 * limit
chunks = list(tc.encode_chunks(frames, fmt, 90, limit))
assert all(len(c["parts"]) == 1 and len(c["parts"][0]) <= limit for c in chunks)
assert [c["start_index"] for c in chunks] == list(np.cumsum([0] + [c["count"] for c in chunks[:-1]]))
assert sum(c["count"] for c in chunks) == len(frames)
decoded = [f for c in chunks for f in tc.decode_chunk(c, c["parts"])]
assert len(decoded) == len(frames)
if fmt == "ffv1":
assert np.array_equal(np.stack(decoded), frames)
def test_single_video_frame_may_exceed_chunk_target_up_to_frame_limit():
frames = _noise_frames(count=3)
one_frame = len(next(tc.encode_chunks(frames[:1], "ffv1", 90, 10**12))["parts"][0])
chunks = list(tc.encode_chunks(frames, "ffv1", 90, one_frame // 2, one_frame * 2))
assert [c["count"] for c in chunks] == [1, 1, 1]
assert all(one_frame // 2 < len(c["parts"][0]) <= one_frame * 2 for c in chunks)
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "png"])
def test_frame_that_cannot_fit_alone_raises(fmt):
frames = _noise_frames(count=2)
with pytest.raises(tc.FrameTooLargeError, match="request limit"):
list(tc.encode_chunks(frames, fmt, 90, 64 * 1024, 64 * 1024))
def test_still_chunks_are_bounded():
frames = _noise_frames(count=10)
limit = 1024 * 1024
chunks = list(tc.encode_chunks(frames, "jpeg", 90, limit))
assert all(sum(len(p) for p in c["parts"]) <= limit for c in chunks)
assert sum(c["count"] for c in chunks) == 10
+125
View File
@@ -0,0 +1,125 @@
"""Registered V3 DistributedCollector: transfer options in the public schema and forwarded to the runtime.
Runs without a ComfyUI checkout by loading nodes/v3.py against a recording
stand-in for comfy_api's io module. tests/test_v3_migration.py covers the same
boundary against the real framework when COMFYUI_SOURCE_ROOT is set.
"""
import importlib.util
import sys
import types
from pathlib import Path
import pytest
ROOT = Path(__file__).resolve().parents[1]
PKG = "dist_v3_transfer_testpkg"
class _Field:
def __init__(self, kind, id=None, **kwargs):
self.kind = kind
self.id = id
self.kwargs = kwargs
def _io_type(kind):
return types.SimpleNamespace(
Input=lambda id, **kw: _Field(kind, id, **kw),
Output=lambda id=None, **kw: _Field(kind, id, **kw),
)
def _fake_io():
io = types.ModuleType("io")
class ComfyNode:
pass
io.ComfyNode = ComfyNode
io.Schema = lambda **kw: types.SimpleNamespace(**kw)
io.NodeOutput = lambda *values, **kw: ("NodeOutput", values, kw)
io.NumberDisplay = types.SimpleNamespace(number="number")
io.Hidden = types.SimpleNamespace(unique_id="unique_id", extra_pnginfo="extra_pnginfo")
for name in ("Int", "Float", "String", "Boolean", "Combo", "Image", "Audio",
"AnyType", "Model", "Conditioning", "Vae"):
setattr(io, name, _io_type(name))
return io
@pytest.fixture
def v3(monkeypatch):
calls = []
class CollectorRuntime:
def run(self, **kwargs):
calls.append(kwargs)
return (kwargs["images"], kwargs["audio"])
def module(name, **attrs):
mod = types.ModuleType(name)
mod.__dict__.update(attrs)
monkeypatch.setitem(sys.modules, name, mod)
return mod
for name in [n for n in sys.modules if n == PKG or n.startswith(PKG + ".")]:
monkeypatch.delitem(sys.modules, name)
module(PKG, __path__=[])
module(f"{PKG}.nodes", __path__=[])
module(f"{PKG}.utils", __path__=[])
module(f"{PKG}.nodes.utilities")
module(f"{PKG}.nodes.distributed_upscale", UltimateSDUpscaleDistributed=object)
module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime)
module("comfy", __path__=[])
module("comfy.samplers", KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["normal"]))
module("comfy_api", __path__=[])
module("comfy_api.v0_0_2", io=_fake_io())
codec_spec = importlib.util.spec_from_file_location(f"{PKG}.utils.transfer_codec", ROOT / "utils" / "transfer_codec.py")
codec = importlib.util.module_from_spec(codec_spec)
monkeypatch.setitem(sys.modules, f"{PKG}.utils.transfer_codec", codec)
codec_spec.loader.exec_module(codec)
spec = importlib.util.spec_from_file_location(f"{PKG}.nodes.v3", ROOT / "nodes" / "v3.py")
mod = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, f"{PKG}.nodes.v3", mod)
spec.loader.exec_module(mod)
return mod, codec, calls
def test_registered_collector_schema_exposes_optional_transfer_inputs(v3):
mod, codec, _ = v3
schema = mod.DistributedCollector.define_schema()
inputs = {field.id: field for field in schema.inputs}
assert [field.id for field in schema.inputs] == [
"load_balance", "images", "audio", "transfer_format", "transfer_quality"]
fmt = inputs["transfer_format"]
assert fmt.kind == "Combo"
assert fmt.kwargs["options"] == list(codec.TRANSFER_FORMATS)
assert fmt.kwargs["default"] == "legacy_png"
assert fmt.kwargs["optional"] is True
quality = inputs["transfer_quality"]
assert quality.kind == "Int"
assert (quality.kwargs["default"], quality.kwargs["min"], quality.kwargs["max"]) == (90, 1, 100)
assert quality.kwargs["optional"] is True
assert schema.is_input_list is True
def test_registered_collector_forwards_transfer_options(v3):
mod, _, calls = v3
mod.DistributedCollector.execute(images=["img"], audio=[None], transfer_format=["av1"], transfer_quality=[70],
multi_job_id=["job"], is_worker=[True])
assert calls[-1]["transfer_format"] == ["av1"]
assert calls[-1]["transfer_quality"] == [70]
assert calls[-1]["multi_job_id"] == ["job"]
def test_registered_collector_defaults_to_legacy_transfer(v3):
mod, _, calls = v3
mod.DistributedCollector.execute(images=["img"])
assert calls[-1]["transfer_format"] == "legacy_png"
assert calls[-1]["transfer_quality"] == 90
+11
View File
@@ -42,6 +42,17 @@ 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)
# Compressed worker -> master result transfer (DistributedCollector)
# Hard cap on one multipart request body. Must not exceed the master's
# --max-upload-size (ComfyUI default: 100 MiB), which rejects larger requests.
TRANSFER_MAX_REQUEST_BYTES = int(os.environ.get('COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES', str(96 * 1024 * 1024)))
# Target media bytes per request when batching several frames into one
# request; a single frame may exceed this up to the hard cap above.
TRANSFER_MAX_CHUNK_BYTES = min(
int(os.environ.get('COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES', str(32 * 1024 * 1024))),
TRANSFER_MAX_REQUEST_BYTES,
)
# Heartbeat monitoring
HEARTBEAT_INTERVAL = float(os.environ.get('COMFYUI_HEARTBEAT_INTERVAL', '10')) # Heartbeat/check interval in seconds
HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds (default 60s)
+411
View File
@@ -0,0 +1,411 @@
"""
Compressed worker -> master image transfer.
Workers encode their IMAGE batch with one of the formats below and POST it to
the master as multipart/form-data chunks; the master decodes each chunk back
into per-frame tensors. The goal is to cut network egress (e.g. cloud workers
returning results to a local master), so formats range from bit-exact
lossless to strongly compressed video.
Formats
-------
legacy_png Original behaviour: one JSON request per frame, uncompressed PNG
in base64. Kept for compatibility with unmodified masters.
png Lossless PNG per frame, sent as binary.
ffv1 Lossless FFV1 video for the whole batch (good for frame sequences).
jpeg, webp Lossy per-frame stills.
h264, h265, av1
Lossy video for the whole batch. Smallest by far for video
frame batches, where consecutive frames are similar.
All formats quantize to 8-bit exactly like the legacy path did, so the
lossless formats reproduce the legacy master tensors bit-for-bit.
"""
import io
import math
import os
import numpy as np
from PIL import Image
LEGACY_FORMAT = "legacy_png"
IMAGE_FORMATS = ("png", "jpeg", "webp")
VIDEO_FORMATS = ("ffv1", "h264", "h265", "av1")
LOSSLESS_FORMATS = (LEGACY_FORMAT, "png", "ffv1")
TRANSFER_FORMATS = (LEGACY_FORMAT, "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1")
DEFAULT_TRANSFER_FORMAT = LEGACY_FORMAT
DEFAULT_TRANSFER_QUALITY = 90
TRANSFER_FORMAT_TOOLTIP = (
"How workers send images back to the master. "
"legacy_png: original uncompressed PNG (largest). "
"png / ffv1: lossless, identical result. "
"jpeg / webp: lossy stills. "
"h264 / h265 / av1: lossy video of the whole batch, smallest for frame sequences."
)
TRANSFER_QUALITY_TOOLTIP = (
"Quality for lossy transfer formats (ignored by lossless ones). "
"JPEG/WebP use it directly; video codecs map it to CRF "
"(90 -> h264/h265 CRF 13, av1 CRF 16)."
)
# encoder name, max CRF, container-friendly pixel format
_VIDEO_CODECS = {
"ffv1": ("ffv1", None, "bgr0"),
"h264": ("libx264", 51, "yuv420p"),
"h265": ("libx265", 51, "yuv420p"),
"av1": ("libsvtav1", 63, "yuv420p"),
}
_VIDEO_RATE = 24
# Upper bound on frames per video chunk, so one request never holds an
# unbounded batch even when frames compress extremely well.
MAX_VIDEO_CHUNK_FRAMES = 240
# When sizing the next video chunk from observed bytes/frame, aim this far
# below the limit so ordinary variation does not force a re-encode.
VIDEO_CHUNK_SIZE_HEADROOM = 0.85
class TransferCodecError(ValueError):
"""Raised when a transfer payload cannot be encoded or decoded."""
class FrameTooLargeError(TransferCodecError):
"""A single encoded frame does not fit in one request.
Raised by encode_chunks before anything has been yielded for that frame;
callers that materialize all chunks first can therefore fail without
having delivered a partial batch.
"""
def normalize_format(value):
fmt = str(value or DEFAULT_TRANSFER_FORMAT).strip().lower()
if fmt not in TRANSFER_FORMATS:
raise TransferCodecError(
f"Unknown transfer format '{value}'. Expected one of: {', '.join(TRANSFER_FORMATS)}"
)
return fmt
def normalize_quality(value):
try:
quality = int(value)
except (TypeError, ValueError):
quality = DEFAULT_TRANSFER_QUALITY
return max(1, min(100, quality))
def quality_to_crf(quality, max_crf):
"""Map 1-100 quality onto a codec CRF scale (0 = lossless/best).
The curve is chosen so the same quality number looks roughly the same
across JPEG/WebP and video codecs: 90 -> CRF 13 (h264/h265), 16 (av1);
80 -> 19 / 24; 100 -> 0.
"""
quality = normalize_quality(quality)
return int(round(max_crf * (1.0 - quality / 100.0) ** 0.6))
def tensor_batch_to_uint8(images):
"""IMAGE tensor [B,H,W,C] float 0..1 -> uint8 numpy [B,H,W,3].
Uses the same truncating quantization as utils.image.tensor_to_pil so that
lossless formats reproduce the legacy transfer exactly, but clips first so
out-of-range values no longer wrap around.
"""
arr = images.detach().cpu().float().numpy()
if arr.ndim == 3:
arr = arr[None, ...]
if arr.ndim != 4:
raise TransferCodecError(f"Expected IMAGE tensor [B,H,W,C], got shape {tuple(arr.shape)}")
channels = arr.shape[-1]
if channels == 1:
arr = np.repeat(arr, 3, axis=-1)
elif channels >= 3:
arr = arr[..., :3]
else:
raise TransferCodecError(f"Unsupported channel count {channels}")
return np.ascontiguousarray(np.clip(255.0 * arr, 0, 255).astype(np.uint8))
def uint8_frame_to_tensor_array(frame):
"""uint8 [H,W,3] -> float32 [1,H,W,3] numpy, matching utils.image.pil_to_tensor."""
return (frame.astype(np.float32) / 255.0)[None, ...]
# --------------------------------------------------------------------------- #
# Still images
# --------------------------------------------------------------------------- #
def encode_image_frame(frame, fmt, quality):
img = Image.fromarray(frame)
quality = normalize_quality(quality)
buf = io.BytesIO()
if fmt == "png":
img.save(buf, format="PNG", compress_level=4) # ~same size as 6, 2x faster
elif fmt == "jpeg":
img.save(buf, format="JPEG", quality=quality, subsampling=0 if quality >= 90 else 2)
elif fmt == "webp":
img.save(buf, format="WEBP", quality=quality, method=4)
else:
raise TransferCodecError(f"'{fmt}' is not a still-image transfer format")
return buf.getvalue()
def decode_image_frame(data):
try:
with Image.open(io.BytesIO(data)) as img:
return np.asarray(img.convert("RGB"), dtype=np.uint8)
except Exception as exc:
raise TransferCodecError(f"Failed to decode image frame: {exc}") from exc
# --------------------------------------------------------------------------- #
# Video
# --------------------------------------------------------------------------- #
def _import_av():
# SVT-AV1 prints a banner per encoder instance unless told otherwise.
os.environ.setdefault("SVT_LOG", "1")
try:
import av # PyAV ships with ComfyUI
except ImportError as exc:
raise TransferCodecError("PyAV ('av') is required for video transfer formats") from exc
return av
def video_codec_available(fmt):
if fmt not in _VIDEO_CODECS:
return False
try:
av = _import_av()
av.codec.Codec(_VIDEO_CODECS[fmt][0], "w")
return True
except Exception:
return False
def _padded_size(fmt, pix_fmt, width, height):
"""Encoder-friendly frame size; the decoder crops back to width x height.
4:2:0 needs even dimensions, and libx265 crashes on frames smaller than
its 64px CTU, so pad up to whichever is larger.
"""
if pix_fmt == "yuv420p":
width, height = width + width % 2, height + height % 2
if fmt == "h265":
width, height = max(width, 64), max(height, 64)
return width, height
def _pad_to(frame, width, height):
h, w = frame.shape[:2]
if h == height and w == width:
return frame
return np.pad(frame, ((0, height - h), (0, width - w), (0, 0)), mode="edge")
class _VideoChunkWriter:
def __init__(self, fmt, quality, width, height):
av = _import_av()
self._av = av
codec_name, max_crf, pix_fmt = _VIDEO_CODECS[fmt]
self.fmt = fmt
self.pix_fmt = pix_fmt
self.width, self.height = _padded_size(fmt, pix_fmt, width, height)
self.buffer = io.BytesIO()
self.container = av.open(self.buffer, mode="w", format="matroska")
self.stream = self.container.add_stream(codec_name, rate=_VIDEO_RATE)
self.stream.width = self.width
self.stream.height = self.height
self.stream.pix_fmt = pix_fmt
options = {}
if max_crf is not None:
options["crf"] = str(quality_to_crf(quality, max_crf))
if fmt == "h264":
options["preset"] = "medium"
elif fmt == "h265":
options["preset"] = "fast"
options["x265-params"] = "log-level=error"
elif fmt == "av1":
options["preset"] = "8"
elif fmt == "ffv1":
options["level"] = "3"
options["slicecrc"] = "0"
self.stream.options = options
self.count = 0
def add(self, frame):
frame = _pad_to(frame, self.width, self.height)
vf = self._av.VideoFrame.from_ndarray(frame, format="rgb24")
for packet in self.stream.encode(vf):
self.container.mux(packet)
self.count += 1
def bytes_so_far(self):
return self.buffer.tell()
def close(self):
for packet in self.stream.encode():
self.container.mux(packet)
self.container.close()
return self.buffer.getvalue()
def decode_video(data, width, height, expected_frames=None):
av = _import_av()
frames = []
try:
with av.open(io.BytesIO(data), mode="r") as container:
for vf in container.decode(video=0):
arr = vf.to_ndarray(format="rgb24")
frames.append(np.ascontiguousarray(arr[:height, :width]))
except TransferCodecError:
raise
except Exception as exc:
raise TransferCodecError(f"Failed to decode video chunk: {exc}") from exc
if expected_frames is not None and len(frames) != expected_frames:
raise TransferCodecError(
f"Video chunk decoded to {len(frames)} frames, expected {expected_frames}"
)
return frames
# --------------------------------------------------------------------------- #
# Chunking
# --------------------------------------------------------------------------- #
def encode_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes=None):
"""Yield transfer chunks for a uint8 frame batch [B,H,W,3].
Each chunk is a dict:
{"format", "start_index", "count", "width", "height", "parts": [bytes, ...]}
Image formats carry one part per frame; video formats carry one part
(a Matroska file) holding `count` frames.
Size guarantees, measured on the finalized bytes (after encoder flush and
muxer finalization, not on encoder buffers):
* a chunk holding more than one frame never exceeds `max_chunk_bytes`;
* a chunk holding a single frame may exceed `max_chunk_bytes` but never
`max_frame_bytes` (defaults to `max_chunk_bytes`).
A frame that cannot fit in `max_frame_bytes` even on its own raises
FrameTooLargeError. Video chunks that come out too large are re-encoded
with fewer frames.
"""
fmt = normalize_format(fmt)
if fmt == LEGACY_FORMAT:
raise TransferCodecError("legacy_png is sent through the JSON path, not encode_chunks")
total = int(frames.shape[0])
if total == 0:
return
height, width = int(frames.shape[1]), int(frames.shape[2])
max_chunk_bytes = max(1, int(max_chunk_bytes))
max_frame_bytes = max_chunk_bytes if max_frame_bytes is None else max(max_chunk_bytes, int(max_frame_bytes))
def _chunk(start, parts, count):
return {
"format": fmt,
"start_index": start,
"count": count,
"width": width,
"height": height,
"parts": parts,
}
def _too_large(index, size):
return FrameTooLargeError(
f"Frame {index} encodes to {size} bytes as {fmt}, above the {max_frame_bytes}-byte request limit. "
"Use a lossy format or lower quality, or raise the master's --max-upload-size together with "
"COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES."
)
if fmt in IMAGE_FORMATS:
start, parts, size = 0, [], 0
for i in range(total):
data = encode_image_frame(frames[i], fmt, quality)
if len(data) > max_frame_bytes:
raise _too_large(i, len(data))
if parts and size + len(data) > max_chunk_bytes:
yield _chunk(start, parts, len(parts))
start, parts, size = i, [], 0
parts.append(data)
size += len(data)
if parts:
yield _chunk(start, parts, len(parts))
return
start = 0
budget = MAX_VIDEO_CHUNK_FRAMES
while start < total:
want = min(budget, total - start)
data, count = _encode_video_range(frames, start, want, fmt, quality, width, height, max_chunk_bytes)
if len(data) <= max_chunk_bytes or (count == 1 and len(data) <= max_frame_bytes):
yield _chunk(start, [data], count)
start += count
per_frame = len(data) / count
budget = max(1, min(MAX_VIDEO_CHUNK_FRAMES, int(max_chunk_bytes * VIDEO_CHUNK_SIZE_HEADROOM / per_frame)))
continue
if count == 1:
raise _too_large(start, len(data))
# Finalized chunk overshot: retry this range with proportionally fewer
# frames. Strictly fewer than `count`, so this always terminates.
estimate = int(count * max_chunk_bytes * VIDEO_CHUNK_SIZE_HEADROOM / len(data))
budget = max(1, min(count - 1, estimate))
def _encode_video_range(frames, start, count, fmt, quality, width, height, stop_after_bytes):
"""Encode frames[start:start+count] as one finalized video.
Stops feeding frames early once the muxed output already exceeds
`stop_after_bytes`, since the chunk is then known to be too large and will
be re-encoded smaller; this bounds wasted work. Encoder lookahead means the
pre-flush size underestimates the final size, so the caller must still
check the returned (finalized) bytes. Returns (bytes, frames_encoded).
"""
writer = _VideoChunkWriter(fmt, quality, width, height)
for i in range(start, start + count):
writer.add(frames[i])
if writer.bytes_so_far() > stop_after_bytes:
break
return writer.close(), writer.count
def decode_chunk(envelope, parts):
"""Decode one received chunk into a list of uint8 [H,W,3] frames."""
fmt = normalize_format(envelope.get("format"))
if fmt == LEGACY_FORMAT:
raise TransferCodecError("legacy_png chunks are not sent as multipart")
count = envelope.get("count")
width = envelope.get("width")
height = envelope.get("height")
for name, value in (("count", count), ("width", width), ("height", height)):
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
raise TransferCodecError(f"envelope.{name}: expected positive integer")
if not parts:
raise TransferCodecError("chunk has no media parts")
if fmt in IMAGE_FORMATS:
if len(parts) != count:
raise TransferCodecError(f"expected {count} image parts, got {len(parts)}")
frames = [decode_image_frame(p) for p in parts]
else:
if len(parts) != 1:
raise TransferCodecError(f"expected 1 video part, got {len(parts)}")
frames = decode_video(parts[0], width, height, expected_frames=count)
for frame in frames:
if frame.shape[0] != height or frame.shape[1] != width:
raise TransferCodecError(
f"decoded frame is {frame.shape[1]}x{frame.shape[0]}, expected {width}x{height}"
)
return frames
def describe_ratio(sent_bytes, raw_bytes):
if sent_bytes <= 0:
return "n/a"
ratio = raw_bytes / sent_bytes
return f"{ratio:.1f}x" if math.isfinite(ratio) else "n/a"