Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9345a8e71b | ||
|
|
a630f55bd9 | ||
|
|
6edd3fda38 |
@@ -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
@@ -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}, "
|
||||
|
||||
@@ -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)`.
|
||||
@@ -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
@@ -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
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
@@ -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])
|
||||
|
||||
@@ -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"}
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user