Merge feat/runpod-serverless-workers (now including main) into feat/dynamic-image-list-queue
Resolve api/queue_orchestration.py, nodes/v3.py and tests/helpers/v3_runtime_check.py to keep both the worker sink (#99) and the image-list nodes. Stub nodes.image_list in the V3 collector transfer test.
This commit is contained in:
Robert Wojciechowski
@@ -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