diff --git a/api/orchestration/serverless_media_prompt.py b/api/orchestration/serverless_media_prompt.py index 17087e5..6cabf4e 100644 --- a/api/orchestration/serverless_media_prompt.py +++ b/api/orchestration/serverless_media_prompt.py @@ -13,6 +13,7 @@ import json from graphlib import CycleError, TopologicalSorter from ...serverless_bridge.contract import SOCKET_KINDS, identifier, validate_bindings +from ...serverless_bridge.image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options from .serverless_prompt import ( _PRIMITIVE_NODES, _REMOVED_SCALAR_NODES, @@ -87,6 +88,8 @@ def _validate_prompt(prepared_prompt): for collector_id in collectors: inputs = prepared_prompt[collector_id].get("inputs", {}) + validate_options(inputs.get("transfer_format", DEFAULT_FORMAT), + inputs.get("transfer_quality", DEFAULT_QUALITY)) for field in ("multi_job_id", "worker_id"): _validate_identifier(inputs.get(field), f"collector {collector_id!r} {field}") connected = [socket for socket in SOCKET_KINDS if socket in inputs] @@ -159,6 +162,11 @@ def _build_exports(source_nodes, collector_ids, request_id): socket: copy.deepcopy(inputs[socket]), }, } + if kind == "IMAGE": + fmt, quality = validate_options(inputs.get("transfer_format", DEFAULT_FORMAT), + inputs.get("transfer_quality", DEFAULT_QUALITY)) + if fmt != DEFAULT_FORMAT: + exports[node_id]["inputs"].update(transfer_format=fmt, transfer_quality=quality) validate_bindings(bindings) return bindings, exports diff --git a/docs/serverless-media.md b/docs/serverless-media.md index b55bcc4..a7c16b4 100644 --- a/docs/serverless-media.md +++ b/docs/serverless-media.md @@ -1,6 +1,6 @@ # Typed serverless media (PR2) -`comfy-media-v1` adds typed, lossless list results for `IMAGE`, `AUDIO`, `MASK`, +`comfy-media-v1` adds typed list results, exact by default, for `IMAGE`, `AUDIO`, `MASK`, `STRING`, and native ComfyUI `VIDEO`. It is an **explicit endpoint profile**, not an automatically detected capability. Omitting `media_profile` still selects `legacy-image-v1`; existing image collectors and handlers keep their old contract. @@ -44,6 +44,32 @@ sockets. Different items may have different valid shapes or audio sample rates. The entire remote response is validated and decoded before any of its results are delivered; missing, duplicate, mismatched, or corrupt results fail collection. +## Image return format and quality + +`transfer_format` defaults to **`exact`** and `transfer_quality` to **90**. +These controls affect **remote IMAGE returns only**, not master-local images, +uploaded inputs, MASK, AUDIO, STRING or native VIDEO values. + +- `exact`: preserve floating-point dtype and values, with no image quantization. +- `legacy_png`, `png`, `ffv1`: lossless **8-bit image pixels**, not exact float + tensors. Image values are clipped to [0,1] and quantized as in the existing + Collector. `legacy_png` uses PNG without compression; `png` compresses it. +- `jpeg`, `webp`: lossy still-image compression. +- `h264`, `h265`, `av1`: lossy video encoding of each IMAGE tensor's frame batch. + This does not transcode the native `video` socket or combine separate list items. +- Quality 1–100 affects lossy formats only, using the existing Collector's + quality-to-CRF mapping for video. Dimensions, frame order and channel count + are preserved; alpha is carried separately where needed, not silently dropped. + +Both master and endpoint must have this bridge revision for compressed image +returns. Pillow handles still formats and CPU/software PyAV encoders handle video; +unavailable codecs fail clearly rather than silently changing the selection. +Compression can reduce result egress and inline response size, but does not relax +encoded/decoded byte budgets or the input request limit. Compressed IMAGE batches +are capped at 240 frames and 8192 pixels per dimension, with a conservative +256 MiB aggregate decode budget (including codec scratch and alpha). The original Collector's +settings, default and fallback behavior are unchanged. + ## Inputs and fidelity Automatic local-file transfer supports stock `LoadImage`, plus explicit stock diff --git a/nodes/media_collector.py b/nodes/media_collector.py index b994c9c..981a6ba 100644 --- a/nodes/media_collector.py +++ b/nodes/media_collector.py @@ -13,6 +13,8 @@ from comfy_api.v0_0_2 import io import server from ..serverless_bridge.contract import MAX_ITEMS, SOCKET_KINDS, identifier +from ..serverless_bridge.image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY, + FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options) from ..utils.async_helpers import run_async_in_server_loop POLL_SECONDS = 0.1 @@ -193,7 +195,11 @@ class DistributedMediaCollector(io.ComfyNode): *[kind.Input(socket, optional=True, **({'force_input': True} if socket == 'texts' else {'extra_dict': {'forceInput': True}})) - for socket, kind in zip(SOCKET_KINDS, types)]], + for socket, kind in zip(SOCKET_KINDS, types)], + io.Combo.Input("transfer_format", options=list(FORMATS), default=DEFAULT_FORMAT, + optional=True, tooltip=FORMAT_TOOLTIP), + io.Int.Input("transfer_quality", default=DEFAULT_QUALITY, min=1, max=100, + optional=True, tooltip=QUALITY_TOOLTIP)], outputs=[kind.Output(display_name=socket, is_output_list=True) for socket, kind in zip(SOCKET_KINDS, types)], is_input_list=True, accept_all_inputs=True, @@ -203,7 +209,9 @@ class DistributedMediaCollector(io.ComfyNode): def execute(cls, images=None, audio=None, masks=None, texts=None, video=None, load_balance=False, multi_job_id='', is_worker=False, enabled_worker_ids='[]', delegate_only=False, pass_through=False, - expected_media='[]'): + expected_media='[]', transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY): + validate_options(_control(transfer_format, 'transfer_format', str), + _control(transfer_quality, 'transfer_quality', int)) media = dict(images=images, audio=audio, masks=masks, texts=texts, video=video) controls = dict(load_balance=load_balance, is_worker=is_worker, enabled_worker_ids=enabled_worker_ids, delegate_only=delegate_only, diff --git a/serverless_bridge/README.md b/serverless_bridge/README.md index c2bc06a..fbc547d 100644 --- a/serverless_bridge/README.md +++ b/serverless_bridge/README.md @@ -155,6 +155,28 @@ only and must not be returned to the master. The executor reads validated `media_exports` history records and bounded `/view` payloads before constructing the external envelope. [contract.py](contract.py) is the schema authority. +## Optional image return encoding + +IMAGE `MediaExportV1` nodes may also contain literal `transfer_format` and +`transfer_quality` controls. Defaults are `exact` and `90`; old exact requests +omit the controls and retain their original representation. Other media exports +must not carry image controls. The master compiler supplies them from the typed +Collector's widgets; callers should not alter binding identity metadata. + +Formats: `exact`, `legacy_png`, `png`, `ffv1`, `jpeg`, `webp`, `h264`, `h265`, `av1`. +`exact` preserves float tensor bits. The image/video codecs quantize RGB/grayscale +to 8-bit; PNG/FFV1 are lossless at that precision, while the other codecs are +lossy. Quality is an integer 1–100 and affects only lossy formats. IMAGE batch +order, shape and channel count are retained, and alpha is not discarded. +Native VIDEO/MASK/AUDIO/STRING still use their original exact codecs. + +Copy/update the complete bridge directory on both sides, including +`image_transfer.py` and `image_codec.py`. Pillow and the selected CPU/software +PyAV codec must be available; missing encoders fail without automatic fallback. +Compressed IMAGE artifacts use `image-transfer-v1`; decode through the shared +`codec.decode_payload`, which enforces payload and decoded-shape limits. +See [image transfer choices](../docs/serverless-media.md#image-return-format-and-quality). + ## Transport budgets and ownership - Contract maxima: 256 MiB per artifact, 512 MiB aggregate artifact bytes, diff --git a/serverless_bridge/codec.py b/serverless_bridge/codec.py index cac8c13..f72faf3 100644 --- a/serverless_bridge/codec.py +++ b/serverless_bridge/codec.py @@ -146,6 +146,10 @@ def decode_payload(kind: str, descriptor: dict, payload: bytes, *, max_bytes: in raise ValueError("Payload must be bytes") if len(payload) > limit: raise ValueError("Payload exceeds byte limit") + if kind == "IMAGE" and descriptor.get("encoding") == "image-transfer-v1": + # Lazy and IMAGE-only: exact tensors and other kinds stay unchanged. + from .image_codec import decode_image + return decode_image(descriptor, payload, max_bytes=limit) if kind == "STRING": if set(descriptor) != {"encoding"} or descriptor["encoding"] != "utf8-v1": raise ValueError("Invalid STRING descriptor encoding or fields") diff --git a/serverless_bridge/handler.py b/serverless_bridge/handler.py index 3c28caf..c6a7f9b 100644 --- a/serverless_bridge/handler.py +++ b/serverless_bridge/handler.py @@ -119,8 +119,13 @@ def _validate_request(input_data): if node is None or node["class_type"] != "MediaExportV1": raise ValueError("Missing declared MediaExportV1 node") controls = node["inputs"] - if set(controls) != {"request_id", "binding_json", binding["socket"]}: + from .image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options + required = {"request_id", "binding_json", binding["socket"]} + optional = {"transfer_format", "transfer_quality"} if binding["kind"] == "IMAGE" else set() + if not required <= set(controls) or set(controls) - required - optional: raise ValueError("Media exporter must have exactly its declared socket and controls") + validate_options(controls.get("transfer_format", DEFAULT_FORMAT), + controls.get("transfer_quality", DEFAULT_QUALITY)) if controls["request_id"] != media["request_id"]: raise ValueError("Media exporter request identity mismatch") text = controls["binding_json"] diff --git a/serverless_bridge/image_codec.py b/serverless_bridge/image_codec.py new file mode 100644 index 0000000..7523191 --- /dev/null +++ b/serverless_bridge/image_codec.py @@ -0,0 +1,634 @@ +"""Bounded optional IMAGE returns, independent of ComfyUI and Distributed. + +Exact mode is the existing tensor wire format, unchanged. All other modes clip +and truncate RGB/luma to uint8 (as the Collector does), normalize in float32, +then cast to the source dtype. RGBA always carries an exact tensor alpha +sidecar, including PNG/WebP. ``legacy_png`` means uncompressed binary PNG here, +not the old Collector's JSON/base64 transport. + +The compressed wire format concatenates length-delimited stills or one Matroska +video, followed by optional alpha bytes. No archive, path, URL or pickle is +accepted. Limits cover wire bytes, descriptor structure, output tensor bytes, +and aggregate padded pixel buffers, not process RSS/native codec overhead. +Video preflight reads EBML track dimensions and block counts before creating a +decoder. We extract packets ourselves, avoiding native container stream probing +(which can decode before caller checks). FFmpeg max_pixels also bounds +bitstream-driven resolution changes, including native buffer alignment. +Software codecs only, one thread, no fallback to another encoder/format. +""" + +import io +import struct +import zlib +from fractions import Fraction + +import torch + +from . import codec +from .codec import MAX_PAYLOAD_BYTES +from .image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options + +MAX_FRAMES = 240 +MAX_DIMENSION = 8192 +MAX_DESCRIPTOR_BYTES = 64 * 1024 +MAX_CONTAINER_ELEMENTS = 16384 +_VIDEO_CODECS = { + "ffv1": ("ffv1", None, "bgr0"), + "h264": ("libx264", 51, "yuv420p"), + "h265": ("libx265", 51, "yuv420p"), + "av1": ("libsvtav1", 63, "yuv420p"), +} +_CODEC_IDS = {"ffv1": b"V_FFV1", "h264": b"V_MPEG4/ISO/AVC", + "h265": b"V_MPEGH/ISO/HEVC", "av1": b"V_AV1"} + + +class _BoundedIO(io.BytesIO): + """Enforce the final file size on every write, including muxer flush/seek.""" + def __init__(self, limit): + super().__init__() + self.limit = limit + + def write(self, data): + if self.tell() + len(data) > self.limit: + raise ValueError("Encoded IMAGE exceeds byte limit") + return super().write(data) + + def seek(self, offset, whence=0): + if whence == 0: + position = offset + elif whence == 1: + position = self.tell() + offset + elif whence == 2: + position = self.getbuffer().nbytes + offset + else: + raise ValueError("Invalid image buffer seek") + if not 0 <= position <= self.limit: + raise ValueError("Encoded IMAGE exceeds byte limit") + return super().seek(offset, whence) + + +def _quality_to_crf(quality, maximum): + return int(round(maximum * (1.0 - quality / 100.0) ** 0.6)) + + +def _padded_size(fmt, height, width): + if fmt in ("h264", "h265", "av1"): + height, width = height + height % 2, width + width % 2 + # libx265 / SVT-AV1 cannot safely encode arbitrarily small pictures. + if fmt in ("h265", "av1"): + height, width = max(64, height), max(64, width) + return height, width + + +def _shape_budget(shape, fmt, limit, itemsize=4, alpha_size=0, coded_frames=None): + size = codec._tensor_size("IMAGE", shape, itemsize, limit) + count, height, width, _ = shape + if count > MAX_FRAMES or max(height, width) > MAX_DIMENSION: + raise ValueError("IMAGE frame count or dimensions exceed limit") + ph, pw = _padded_size(fmt, height, width) + # Reserve output and padded RGBA-sized decoded buffers BEFORE allocation. + bh, bw = ((ph + 63) // 64 * 64, (pw + 63) // 64 * 64) if fmt in _VIDEO_CODECS else (ph, pw) + # Slice/OBU counts conservatively bound native frames, including AV1 hidden + # frames and multiple outputs from one packet. Reserve 16 bytes per pixel + # for hostile high-bit-depth streams, not only our usual 8-bit RGB frames. + native_count = max(count, coded_frames or count) + if native_count > MAX_FRAMES * 4: + raise ValueError("IMAGE coded frame count exceeds limit") + pixels = codec._tensor_size("IMAGE", [native_count, bh, bw, 4], 4, limit) + if size + pixels + alpha_size > limit: + raise ValueError("IMAGE decoded aggregate exceeds byte limit") + return ph, pw + + +def _descriptor_budget(value): + """Bound the JSON-shaped tree without serializing attacker-sized values.""" + remaining = MAX_DESCRIPTOR_BYTES + def walk(obj, depth): + nonlocal remaining + if depth > 5: + raise ValueError("IMAGE descriptor nesting exceeds limit") + remaining -= 8 + if remaining < 0: + raise ValueError("IMAGE descriptor exceeds byte limit") + if obj is None: + return + if type(obj) is str: + remaining -= len(obj) * 6 + elif type(obj) is int: + if obj.bit_length() > 63: + raise ValueError("IMAGE descriptor integer exceeds limit") + remaining -= 20 + elif type(obj) is list: + if len(obj) > MAX_FRAMES: + raise ValueError("IMAGE descriptor list exceeds limit") + for item in obj: + walk(item, depth + 1) + elif type(obj) is dict: + if len(obj) > 10: + raise ValueError("IMAGE descriptor fields exceed limit") + for key, item in obj.items(): + if type(key) is not str: + raise ValueError("Invalid IMAGE descriptor key") + walk(key, depth + 1) + walk(item, depth + 1) + else: + raise ValueError("Invalid IMAGE descriptor value") + if remaining < 0: + raise ValueError("IMAGE descriptor exceeds byte limit") + walk(value, 0) + + +def _alpha_metadata(alpha, shape, limit): + if shape[-1] != 4: + if alpha is not None: + raise ValueError("Unexpected IMAGE alpha sidecar") + return 0, torch.float32 + if type(alpha) is not dict or set(alpha) != {"descriptor", "size"}: + raise ValueError("Missing or invalid IMAGE alpha sidecar") + d = alpha["descriptor"] + if type(d) is not dict or set(d) != {"encoding", "dtype", "shape", "layout"}: + raise ValueError("Invalid IMAGE alpha descriptor") + if d["encoding"] != "tensor-le-v1" or d["layout"] != "BHWC" or d["shape"] != shape[:3] + [1]: + raise ValueError("Invalid IMAGE alpha shape or encoding") + if type(d["dtype"]) is not str or d["dtype"] not in codec._DTYPES: + raise ValueError("Invalid IMAGE alpha dtype") + dtype, itemsize = codec._DTYPES[d["dtype"]] + size = codec._tensor_size("IMAGE", d["shape"], itemsize, limit) + if type(alpha["size"]) is not int or alpha["size"] != size: + raise ValueError("Invalid IMAGE alpha size") + return size, torch.float64 if dtype == torch.float64 else torch.float32 + + +def _import_av(): + try: + import av + except ImportError: + raise ValueError("PyAV is required for IMAGE video transfer") from None + return av + + +def _encode_still(frame, fmt, quality, limit): + from PIL import Image + sink = _BoundedIO(limit) + if frame.shape[-1] == 1: + frame = frame[..., 0] + image = Image.fromarray(frame) + if fmt in ("legacy_png", "png"): + image.save(sink, format="PNG", compress_level=0 if fmt == "legacy_png" else 4) + elif fmt == "jpeg": + image.save(sink, format="JPEG", quality=quality, subsampling=0 if quality >= 90 else 2) + else: + image.save(sink, format="WEBP", quality=quality, method=4) + return sink.getvalue() + + +def _encode_video(frames, fmt, quality, height, width, limit): + import numpy as np + av = _import_av() + encoder, maximum, pix_fmt = _VIDEO_CODECS[fmt] + try: + av.codec.Codec(encoder, "w") + except Exception: + raise ValueError("IMAGE software encoder unavailable: " + encoder) from None + sink = _BoundedIO(limit) + options = {} + if maximum is not None: + options["crf"] = str(_quality_to_crf(quality, maximum)) + if fmt == "h264": + options["preset"] = "medium" + elif fmt == "h265": + options.update(preset="fast", **{"x265-params": "log-level=error:pools=none:frame-threads=1"}) + elif fmt == "av1": + options.update(preset="8", **{"svtav1-params": "lp=1"}) + else: + options.update(level="3", slicecrc="1") + try: + with av.open(sink, mode="w", format="matroska") as container: + stream = container.add_stream(encoder, rate=24) + stream.width, stream.height = width, height + stream.pix_fmt = pix_fmt + stream.codec_context.thread_count = 1 + stream.options = options + for index, frame in enumerate(frames): + if frame.shape[-1] == 1: + frame = np.repeat(frame, 3, axis=-1) + h, w = frame.shape[:2] + frame = np.pad(frame, ((0, height - h), (0, width - w), (0, 0)), mode="edge") + vf = av.VideoFrame.from_ndarray(frame, format="rgb24") + vf.pts = index + vf.time_base = Fraction(1, 24) + for packet in stream.encode(vf): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + data = sink.getvalue() + if len(data) > limit: + raise ValueError("Encoded IMAGE exceeds byte limit") + return data + except av.error.FFmpegError: + raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None + except ValueError: + raise + except Exception: + raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None + + +def encode_image(value, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY, + max_bytes=MAX_PAYLOAD_BYTES): + """Return (descriptor, bytes); reject unsupported inputs, never fall back.""" + fmt, quality = validate_options(transfer_format, transfer_quality) + if fmt == "exact": + return codec.encode_payload("IMAGE", value, max_bytes=max_bytes) + limit = codec._byte_limit(max_bytes) + if not isinstance(value, torch.Tensor) or value.device.type != "cpu": + raise ValueError("IMAGE must be a CPU tensor") + if value.layout != torch.strided or value.is_nested: + raise ValueError("IMAGE must have dense strided layout") + if value.dtype not in codec._DTYPE_NAMES: + raise ValueError("Unsupported IMAGE dtype") + shape = list(value.shape) + codec._tensor_size("IMAGE", shape, value.element_size(), limit) + alpha_size = shape[0] * shape[1] * shape[2] * value.element_size() if shape[-1] == 4 else 0 + itemsize = value.element_size() + ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size) + if not torch.isfinite(value).all().item(): + raise ValueError("IMAGE values must be finite") + value = value.detach().resolve_neg() + alpha, alpha_payload = None, b"" + if shape[-1] == 4: + ad, alpha_payload = codec.encode_payload("IMAGE", value[..., 3:4], max_bytes=limit) + alpha = {"descriptor": ad, "size": len(alpha_payload)} + import numpy as np + def frames(): + for frame in value: + arr = frame[..., :3].float().numpy() + with np.errstate(over="ignore", invalid="ignore"): + yield np.ascontiguousarray(np.clip(arr * 255.0, 0, 255).astype(np.uint8)) + budget = limit - len(alpha_payload) + if fmt in _VIDEO_CODECS: + parts = [_encode_video(frames(), fmt, quality, ph, pw, budget)] + header = _video_header(parts[0], fmt, ph, pw, shape[0]) + coded_frames = _coded_frame_bound(parts[0], fmt, header) + _shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames) + else: + parts = [] + for frame in frames(): + part = _encode_still(frame, fmt, quality, budget) + budget -= len(part) + parts.append(part) + descriptor = {"encoding": "image-transfer-v1", "format": fmt, "quality": quality, + "shape": shape, "layout": "BHWC", "dtype": codec._DTYPE_NAMES[value.dtype], + "parts": [len(p) for p in parts], "alpha": alpha} + _descriptor_budget(descriptor) + payload = b"".join(parts) + alpha_payload + if len(payload) > limit: + raise ValueError("Encoded IMAGE exceeds byte limit") + return descriptor, payload + + +def _still_header(data, fmt): + """Read dimensions before Pillow (WebP may allocate pixels during open).""" + if fmt in ("png", "legacy_png"): + if data[:8] != b"\x89PNG\r\n\x1a\n": + raise ValueError("Invalid IMAGE PNG header") + offset, dimensions, saw_data, chunks = 8, None, False, 0 + while offset < len(data): + chunks += 1 + if chunks > MAX_CONTAINER_ELEMENTS: + raise ValueError("IMAGE PNG chunk count exceeds limit") + if offset + 12 > len(data): + raise ValueError("Truncated IMAGE PNG chunk") + length = int.from_bytes(data[offset:offset + 4], "big") + end = offset + 12 + length + if end > len(data): + raise ValueError("Truncated IMAGE PNG chunk") + tag = data[offset + 4:offset + 8] + body = data[offset + 8:end - 4] + if zlib.crc32(data[offset + 4:end - 4]) != int.from_bytes(data[end - 4:end], "big"): + raise ValueError("Invalid IMAGE PNG checksum") + if tag == b"IHDR" and offset == 8 and length == 13: + width, height, depth, color, compression, filtering, interlace = struct.unpack(">IIBBBBB", body) + if depth != 8 or color not in (0, 2) or (compression, filtering, interlace) != (0, 0, 0): + raise ValueError("Unsupported IMAGE PNG layout") + dimensions = (height, width) + elif tag == b"IDAT" and dimensions: + saw_data = True + elif tag == b"IEND" and not length and end == len(data) and saw_data: + return dimensions + else: + raise ValueError("Unsupported IMAGE PNG chunk") + offset = end + raise ValueError("Truncated IMAGE PNG") + if fmt == "webp": + # Our encoder writes a single lossy VP8 still. Reject animation/alpha/ + # extended WebP before WebPAnimDecoder can allocate its canvas. + if len(data) < 30 or data[:4] != b"RIFF" or data[8:16] != b"WEBPVP8 ": + raise ValueError("Unsupported IMAGE WebP header") + size = int.from_bytes(data[16:20], "little") + if int.from_bytes(data[4:8], "little") + 8 != len(data) or 20 + size + size % 2 != len(data): + raise ValueError("Invalid IMAGE WebP length") + if data[23:26] != b"\x9d\x01\x2a": + raise ValueError("Invalid IMAGE WebP frame") + width, height = struct.unpack(" len(data): + raise ValueError("Truncated IMAGE JPEG marker") + if marker in (0xc0, 0xc2): + if dimensions or size < 8 or data[offset + 4] != 8 or data[offset + 9] not in (1, 3): + raise ValueError("Unsupported IMAGE JPEG layout") + dimensions = struct.unpack(">HH", data[offset + 5:offset + 9]) + elif marker == 0xda: + if dimensions is None: + raise ValueError("Missing IMAGE JPEG dimensions") + return dimensions + offset = end + raise ValueError("Missing IMAGE JPEG pixels") + + +def _video_header(data, fmt, height, width, count): + """Strict bounded EBML preflight: one track, no lacing/attachments/encodings. + + Native stream probing is not the first dimension check: it may instantiate a + decoder. Validate the Matroska header and count blocks without decompressing. + """ + nodes = 0 + def vint(pos, end, identifier=False): + if pos >= end or data[pos] == 0: + raise ValueError("Invalid IMAGE EBML integer") + length = 9 - data[pos].bit_length() + if length > (4 if identifier else 8) or pos + length > end: + raise ValueError("Truncated IMAGE EBML integer") + value = int.from_bytes(data[pos:pos + length], "big") + if not identifier: + value &= (1 << (7 * length)) - 1 + if value == (1 << (7 * length)) - 1: + raise ValueError("Unknown-sized IMAGE EBML elements are forbidden") + return value, pos + length + def elements(start, end): + nonlocal nodes + while start < end: + nodes += 1 + if nodes > MAX_CONTAINER_ELEMENTS: + raise ValueError("IMAGE container element count exceeds limit") + tag, pos = vint(start, end, True) + length, body = vint(pos, end) + finish = body + length + if finish > end: + raise ValueError("Truncated IMAGE EBML element") + yield tag, body, finish + start = finish + def integer(start, end): + if not 1 <= end - start <= 8: + raise ValueError("Invalid IMAGE EBML numeric field") + return int.from_bytes(data[start:end], "big") + roots = list(elements(0, len(data))) + if len(roots) != 2 or [r[0] for r in roots] != [0x1a45dfa3, 0x18538067]: + raise ValueError("Invalid IMAGE Matroska structure") + tracks, blocks, packets = [], 0, [] + for tag, start, end in elements(roots[1][1], roots[1][2]): + if tag == 0x1654ae6b: + for entry, es, ee in elements(start, end): + if entry in (0xbf, 0xec): + continue + if entry != 0xae: + raise ValueError("Unexpected IMAGE track element") + fields, dims = {}, {} + for field, fs, fe in elements(es, ee): + if field in fields: + raise ValueError("Duplicate IMAGE track field") + fields[field] = (fs, fe) + if field == 0x6d80: + raise ValueError("IMAGE track encodings are forbidden") + if field == 0xe0: + for dim, ds, de in elements(fs, fe): + if dim in dims: + raise ValueError("Duplicate IMAGE video field") + dims[dim] = (ds, de) + if any(key not in fields for key in (0xd7, 0x83, 0x86)) or any(key not in dims for key in (0xb0, 0xba)): + raise ValueError("Incomplete IMAGE video track") + if integer(*fields[0xd7]) != 1 or integer(*fields[0x83]) != 1: + raise ValueError("Unsupported IMAGE video track") + if fields[0x86][1] - fields[0x86][0] > 32: + raise ValueError("IMAGE codec identifier exceeds limit") + codec_id = bytes(data[slice(*fields[0x86])]) + if codec_id != _CODEC_IDS[fmt] and not (fmt == "ffv1" and codec_id == b"V_MS/VFW/FOURCC"): + raise ValueError("IMAGE video codec mismatch") + if integer(*dims[0xb0]) != width or integer(*dims[0xba]) != height: + raise ValueError("IMAGE video header dimensions mismatch") + if codec_id == b"V_MS/VFW/FOURCC": + cs, ce = fields.get(0x63a2, (0, 0)) + if ce - cs < 40 or data[cs + 16:cs + 20] != b"FFV1": + raise ValueError("Invalid IMAGE FFV1 codec header") + if struct.unpack("= be or data[pos + 2] & 0x06: + raise ValueError("Invalid or laced IMAGE video block") + blocks += 1 + if blocks > count: + raise ValueError("IMAGE video frame count exceeds descriptor") + packets.append((pos + 3, be)) + for field, fs, fe in elements(start, end): + if field == 0xa3: + block(fs, fe) + elif field == 0xa0: + group = [e for e in elements(fs, fe) if e[0] == 0xa1] + if len(group) != 1: + raise ValueError("Invalid IMAGE video block group") + block(group[0][1], group[0][2]) + elif field not in (0xe7, 0xec, 0xbf): + raise ValueError("Unsupported IMAGE cluster field") + elif tag not in (0x114d9b74, 0x1549a966, 0x1c53bb6b, 0x1254c367, 0xec, 0xbf): + raise ValueError("Unsupported IMAGE Matroska element") + if len(tracks) != 1 or blocks != count: + raise ValueError("IMAGE video track or frame count mismatch") + cs, ce = tracks[0].get(0x63a2, (0, 0)) + if not 0 < ce - cs <= MAX_DESCRIPTOR_BYTES: + raise ValueError("IMAGE codec configuration exceeds limit or is missing") + if bytes(data[slice(*tracks[0][0x86])]) == b"V_MS/VFW/FOURCC": + cs += 40 # Skip BITMAPINFOHEADER; FFV1 receives only its codec config. + return bytes(data[cs:ce]), packets + + +def _coded_frame_bound(data, fmt, header): + """Conservative per-packet allocation bound before native decode. + + FFV1 has one picture per packet. Every AVC/HEVC picture needs at least one + VCL NAL; every AV1 output (including show-existing) needs a frame/header OBU. + Counting slices overestimates pictures, which is safe. Do not trust a + Matroska block count alone: a malicious packet may contain several pictures. + """ + config, packets = header + if fmt == "ffv1": + return len(packets) + total, units = 0, 0 + if fmt in ("h264", "h265"): + index = 4 if fmt == "h264" else 21 + if len(config) <= index or config[0] != 1: + raise ValueError("Invalid IMAGE video configuration") + length_size = (config[index] & 3) + 1 + for start, end in packets: + while start < end: + units += 1 + if units > MAX_CONTAINER_ELEMENTS: + raise ValueError("IMAGE coded unit count exceeds limit") + if fmt in ("h264", "h265"): + if start + length_size > end: + raise ValueError("Truncated IMAGE NAL length") + length = int.from_bytes(data[start:start + length_size], "big") + start += length_size + if length < (1 if fmt == "h264" else 2) or length > end - start: + raise ValueError("Truncated IMAGE NAL unit") + kind = (data[start] & 31) if fmt == "h264" else (data[start] >> 1) & 63 + total += (1 <= kind <= 5) if fmt == "h264" else (kind <= 31) + else: + tag = data[start] + start += 1 + if tag & 0x81 or not tag & 2: + raise ValueError("Unsupported IMAGE AV1 OBU header") + kind = (tag >> 3) & 15 + if tag & 4: + start += 1 + length = 0 + for shift in range(0, 56, 7): + if start >= end: + raise ValueError("Truncated IMAGE AV1 OBU length") + value = data[start] + start += 1 + length |= (value & 127) << shift + if not value & 128: + break + else: + raise ValueError("IMAGE AV1 OBU length exceeds limit") + if length > end - start: + raise ValueError("Truncated IMAGE AV1 OBU") + total += kind in (3, 6) + start += length + if total > MAX_FRAMES * 4: + raise ValueError("IMAGE coded frame count exceeds limit") + if total < len(packets): + raise ValueError("IMAGE video has missing coded frames") + return total + + +def _decode_video(data, fmt, height, width, count, header): + av = _import_av() + maximum = ((height + 63) // 64 * 64) * ((width + 63) // 64 * 64) + ctx = av.CodecContext.create("libdav1d" if fmt == "av1" else "hevc" if fmt == "h265" else fmt, "r") + ctx.width, ctx.height = width, height + ctx.extradata, packets = header + ctx.thread_count = 1 + ctx.options = {"max_pixels": str(maximum), "threads": "1", "err_detect": "explode"} + ctx.open() + if "max_pixels" in ctx.options: + raise ValueError("IMAGE decoder does not support max_pixels safety limit") + if (ctx.width, ctx.height) != (width, height): + raise ValueError("IMAGE bitstream dimensions mismatch") + decoded = 0 + for bounds in packets + [None]: + packet = av.Packet(data[slice(*bounds)]) if bounds is not None else None + for frame in ctx.decode(packet): + if decoded >= count or (frame.width, frame.height) != (width, height): + raise ValueError("IMAGE decoded dimensions or frame count mismatch") + decoded += 1 + yield frame.to_ndarray(format="rgb24") + if decoded != count: + raise ValueError("IMAGE decoded frame count mismatch") + + +def decode_image(descriptor, payload, *, max_bytes=MAX_PAYLOAD_BYTES): + """Decode to an independent contiguous CPU tensor after bounded preflight.""" + limit = codec._byte_limit(max_bytes) + if not isinstance(descriptor, dict): + raise ValueError("IMAGE descriptor must be a dictionary") + if descriptor.get("encoding") == "tensor-le-v1": + return codec.decode_payload("IMAGE", descriptor, payload, max_bytes=limit) + _descriptor_budget(descriptor) + if set(descriptor) != {"encoding", "format", "quality", "shape", "layout", "dtype", "parts", "alpha"}: + raise ValueError("Invalid IMAGE descriptor fields") + if descriptor["encoding"] != "image-transfer-v1" or descriptor["layout"] != "BHWC": + raise ValueError("Invalid IMAGE descriptor encoding or layout") + fmt, _ = validate_options(descriptor["format"], descriptor["quality"]) + if fmt == "exact": + raise ValueError("Exact IMAGE requires the tensor encoding") + if not isinstance(payload, bytes): + raise ValueError("IMAGE payload must be bytes") + if len(payload) > limit: + raise ValueError("IMAGE payload exceeds byte limit") + shape = descriptor["shape"] + if type(descriptor["dtype"]) is not str or descriptor["dtype"] not in codec._DTYPES: + raise ValueError("Unsupported IMAGE descriptor dtype") + dtype, itemsize = codec._DTYPES[descriptor["dtype"]] + codec._tensor_size("IMAGE", shape, itemsize, limit) + alpha_size, _ = _alpha_metadata(descriptor["alpha"], shape, limit) + if alpha_size and descriptor["alpha"]["descriptor"]["dtype"] != descriptor["dtype"]: + raise ValueError("IMAGE alpha dtype mismatch") + ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size) + count, height, width, channels = shape + lengths = descriptor["parts"] + if type(lengths) is not list or len(lengths) != (1 if fmt in _VIDEO_CODECS else count): + raise ValueError("IMAGE part count mismatch") + total = alpha_size + for length in lengths: + if type(length) is not int or not 0 < length <= limit - total: + raise ValueError("IMAGE part length exceeds byte limit") + total += length + if total != len(payload): + raise ValueError("IMAGE payload length does not match parts") + parts, offset = [], 0 + for length in lengths: + parts.append(memoryview(payload)[offset:offset + length]) + offset += length + try: + if fmt in _VIDEO_CODECS: + video_header = _video_header(parts[0], fmt, ph, pw, count) + coded_frames = _coded_frame_bound(parts[0], fmt, video_header) + _shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames) + else: + for part in parts: + if _still_header(part, fmt) != (height, width): + raise ValueError("IMAGE header dimensions mismatch") + # All headers and aggregate budgets are checked before output allocation. + alpha = None + if alpha_size: + alpha = codec.decode_payload("IMAGE", descriptor["alpha"]["descriptor"], payload[offset:], max_bytes=limit) + result = torch.empty(shape, dtype=dtype, device="cpu") + if fmt in _VIDEO_CODECS: + frames = _decode_video(parts[0], fmt, ph, pw, count, video_header) + else: + def stills(): + import numpy as np + from PIL import Image + for part in parts: + with Image.open(io.BytesIO(part)) as image: + if image.size != (width, height) or getattr(image, "n_frames", 1) != 1: + raise ValueError("IMAGE still dimensions or frame count mismatch") + if image.mode not in ("L", "RGB"): + raise ValueError("Unsupported IMAGE still pixel mode") + yield np.array(image.convert("RGB"), dtype=np.uint8) + frames = stills() + color = 3 if channels == 4 else channels + for index, frame in enumerate(frames): + # Float32 division followed by the source dtype cast reproduces + # Collector normalized pixels without narrowing exact alpha. + result[index, ..., :color] = torch.from_numpy(frame[:height, :width, :color].copy()).float().div_(255) + if alpha is not None: + result[..., 3:4] = alpha + return result + except ValueError: + raise + except Exception: + raise ValueError("Malformed or unsupported IMAGE compressed payload") from None diff --git a/serverless_bridge/image_transfer.py b/serverless_bridge/image_transfer.py new file mode 100644 index 0000000..0d0ce2e --- /dev/null +++ b/serverless_bridge/image_transfer.py @@ -0,0 +1,30 @@ +"""Image-return policy shared by the master and standalone endpoint bridge. + +This module is deliberately dependency-free so graph compilation does not import +Torch, Pillow, PyAV or ComfyUI. Existing Collector settings stay unchanged. +""" + +FORMATS = ("exact", "legacy_png", "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1") +DEFAULT_FORMAT = "exact" +DEFAULT_QUALITY = 90 +FORMAT_TOOLTIP = ( + "Image returns only; master-local images and other media stay unchanged. " + "exact: original floating-point tensor, no quantization. " + "legacy_png / png / ffv1: lossless 8-bit pixels, not exact float tensors. " + "jpeg / webp: lossy stills; h264 / h265 / av1: lossy frame-batch video. " + "Alpha is preserved separately when the codec cannot carry it. " + "Requires matching transfer support on the worker/endpoint." +) +QUALITY_TOOLTIP = ( + "Quality 1–100 for lossy image returns; ignored by exact and lossless modes. " + "JPEG/WebP use it directly; video formats use the existing Collector CRF mapping." +) + + +def validate_options(transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY): + """Reject malformed controls rather than repairing or silently substituting.""" + if type(transfer_format) is not str or transfer_format not in FORMATS: + raise ValueError("Unsupported image transfer_format") + if type(transfer_quality) is not int or not 1 <= transfer_quality <= 100: + raise ValueError("Image transfer_quality must be an integer from 1 to 100") + return transfer_format, transfer_quality diff --git a/serverless_bridge/nodes.py b/serverless_bridge/nodes.py index 5c61328..0ef6b0f 100644 --- a/serverless_bridge/nodes.py +++ b/serverless_bridge/nodes.py @@ -8,6 +8,8 @@ import folder_paths from comfy_api.v0_0_2 import io from .codec import decode_payload, encode_payload +from .image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY, + FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options) from .contract import ( MAX_ARTIFACT_BYTES, MAX_ITEMS, MAX_TOTAL_BYTES, SOCKET_KINDS, bounded_int, identifier, validate_binding, validate_json, @@ -43,7 +45,11 @@ def _cpu(kind, value): return value -def encode_value(kind, value): +def encode_value(kind, value, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY): + fmt, quality = validate_options(transfer_format, transfer_quality) + if kind == "IMAGE" and fmt != DEFAULT_FORMAT: + from .image_codec import encode_image + return encode_image(_cpu(kind, value), transfer_format=fmt, transfer_quality=quality) if kind == "VIDEO": from .video_codec import encode_video_cpu return encode_video_cpu(value) @@ -57,10 +63,13 @@ def decode_value(kind, descriptor, payload): return decode_payload(kind, descriptor, payload) -def export_values(request_id, binding, values): +def export_values(request_id, binding, values, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY): """Write bounded owned payloads; history contains metadata, never base64.""" identifier(request_id) validate_binding(binding) + fmt, quality = validate_options(transfer_format, transfer_quality) + if binding["kind"] != "IMAGE" and fmt != DEFAULT_FORMAT: + raise ValueError("Image transfer controls apply only to IMAGE exports") if not isinstance(values, list) or not 1 <= len(values) <= MAX_ITEMS: raise ValueError("Connected media output has no values or exceeds item limit") request_key = hashlib.sha256(request_id.encode()).hexdigest() @@ -74,7 +83,8 @@ def export_values(request_id, binding, values): total = 0 try: for index, value in enumerate(values): - descriptor, payload = encode_value(binding["kind"], value) + descriptor, payload = encode_value(binding["kind"], value, + transfer_format=fmt, transfer_quality=quality) validate_json(descriptor) total += len(payload) if total > MAX_TOTAL_BYTES: @@ -144,14 +154,22 @@ class MediaExportV1(io.ComfyNode): inputs=[io.String.Input("request_id"), io.String.Input("binding_json"), io.Image.Input("images", optional=True), io.Audio.Input("audio", optional=True), io.Mask.Input("masks", optional=True), io.String.Input("texts", optional=True, force_input=True), - io.Video.Input("video", optional=True)], outputs=[], is_output_node=True, is_input_list=True) + io.Video.Input("video", optional=True), + io.Combo.Input("transfer_format", options=list(FORMATS), default=DEFAULT_FORMAT, + optional=True, tooltip=FORMAT_TOOLTIP), + io.Int.Input("transfer_quality", default=DEFAULT_QUALITY, min=1, max=100, + optional=True, tooltip=QUALITY_TOOLTIP)], + outputs=[], is_output_node=True, is_input_list=True) @classmethod - def execute(cls, request_id, binding_json, images=None, audio=None, masks=None, texts=None, video=None): + def execute(cls, request_id, binding_json, images=None, audio=None, masks=None, texts=None, video=None, + transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY): binding = _json(_scalar(binding_json)) validate_binding(binding) provided = dict(images=images, audio=audio, masks=masks, texts=texts, video=video) if {name for name, value in provided.items() if value is not None} != {binding["socket"]}: raise ValueError("Media export must receive exactly its declared typed socket") - exported = export_values(_scalar(request_id), binding, provided[binding["socket"]]) + exported = export_values(_scalar(request_id), binding, provided[binding["socket"]], + transfer_format=_scalar(transfer_format), + transfer_quality=_scalar(transfer_quality)) return io.NodeOutput(ui={"media_exports": [exported]}) diff --git a/tests/helpers/media_end_to_end_check.py b/tests/helpers/media_end_to_end_check.py index e70b160..a40158f 100644 --- a/tests/helpers/media_end_to_end_check.py +++ b/tests/helpers/media_end_to_end_check.py @@ -361,6 +361,43 @@ async def check(base): else: raise AssertionError("Result decoder accepted a corrupt final artifact") assert ps.distributed_pending_jobs == {} + # Exercise user-selected image return encoding through the actual + # compiler, HTTP handler, V3 exporter and atomic response decoder. + for fmt in ("png", "jpeg", "h264"): + key = "encoded-" + fmt + fixtures[key] = native_fixtures(3) + # Keep full alpha precision and use nontrivial image dimensions. + torch.manual_seed(41) + rgba = torch.rand((2, 64, 64, 4), dtype=torch.float64) + fixtures[key]["images"] = [rgba] + request_id = "request-" + key + graph = {"source": node("AcceptanceMediaSource", key=key), + "collector": node("DistributedMediaCollector", transfer_format=fmt, + transfer_quality=73, multi_job_id="job-" + key, + worker_id="serverless:" + key, + **{socket: ["source", i] for i, socket in enumerate(sockets)})} + original = copy.deepcopy(graph) + workflow, manifest = compiler.compile_media_prompt(graph, request_id) + media = {"version": 1, "profile": contract.PROFILE, "request_id": request_id, + "artifacts": [], "inputs": [], "outputs": manifest} + body = json.loads(serializer(workflow, [], 30, media=media)) + result = await handler.execute_media_request(body["input"], comfy_url=url, + input_dir=base / "input", temp_dir=base / "temp", timeout_seconds=30) + assert graph == original + grouped = await decoder.collect_typed_results(result, media, manifest) + actual = grouped["job-" + key] + image = actual["images"][0] + assert image.shape == rgba.shape and image.dtype == rgba.dtype + assert torch.equal(image[..., 3], rgba[..., 3]), "alpha must not be discarded/narrowed" + assert not torch.equal(image[..., :3], rgba[..., :3]), "selection must actually encode" + if fmt == "png": + pixels = (rgba[..., :3].float().clamp(0, 1) * 255).to(torch.uint8) + assert torch.equal(image[..., :3], (pixels.float() / 255).to(rgba.dtype)) + for socket, kind in zip(sockets[1:], kinds[1:]): + for got, wanted in zip(actual[socket], fixtures[key][socket]): + exact(kind, got, video_codec.encode_video(wanted) if kind == "VIDEO" else wanted) + assert not list((base / "temp").rglob("*.bin")) + print("IMAGE_TRANSFER_HTTP_OK png jpeg h264; alpha and other media exact") print("REAL_HTTP_BRIDGE_RESULTS_OK all five kinds; 6 bindings/request; native VIDEO remains readable after cleanup") print("ATOMIC_DECODE_REJECTION_OK corrupt final artifact rejected before delivery") @@ -404,7 +441,7 @@ async def check(base): assert executor.history_result["outputs"]["assert"]["accepted"] == [seen_sinks[tag]] assert executor.history_result["outputs"]["text-assert"]["accepted"] == [seen_sinks[text_tag]] assert job not in ps.distributed_pending_jobs and text_job not in ps.distributed_pending_jobs - assert source_calls == ["worker-a", "worker-b", "local"], source_calls + assert source_calls == ["worker-a", "worker-b", "encoded-png", "encoded-jpeg", "encoded-h264", "local"], source_calls assert set(seen_sinks) == {"master", "delegate", "master-text", "delegate-text"} assert not list((base / "output").rglob("*")) assert not torch.cuda.is_initialized() diff --git a/tests/test_media_image_codec.py b/tests/test_media_image_codec.py new file mode 100644 index 0000000..96d6230 --- /dev/null +++ b/tests/test_media_image_codec.py @@ -0,0 +1,330 @@ +"""Real CPU-only IMAGE transfers; no ComfyUI/bootstrap/provider imports.""" +import importlib +import sys +import types +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image + +# Load a private package, not Distributed's root __init__. +ROOT = Path(__file__).resolve().parents[1] / "serverless_bridge" +PACKAGE = "image_codec_tests_bridge" +pkg = types.ModuleType(PACKAGE) +pkg.__path__ = [str(ROOT)] +sys.modules[PACKAGE] = pkg +codec = importlib.import_module(PACKAGE + ".codec") +images = importlib.import_module(PACKAGE + ".image_codec") +FORMATS = ("legacy_png", "png", "jpeg", "webp", "ffv1", "h264", "h265", "av1") +LOSSLESS = ("legacy_png", "png", "ffv1") + + +@pytest.fixture(autouse=True) +def no_gpu(monkeypatch): + def forbidden(*a, **kw): + pytest.fail("GPU query or initialization is forbidden") + for name in ("is_available", "init", "_lazy_init", "device_count"): + monkeypatch.setattr(torch.cuda, name, forbidden) + + +def available(fmt): + if fmt in images._VIDEO_CODECS: + av = pytest.importorskip("av") + name = images._VIDEO_CODECS[fmt][0] + if name not in av.codecs_available: + pytest.skip("Software encoder unavailable: " + name) + + +def sample(channels=3, dtype=torch.float32, shape=(3, 65, 67)): + g = torch.Generator(device="cpu").manual_seed(42) + return torch.rand((*shape, channels), generator=g, device="cpu", dtype=dtype) + + +def quantized(value): + return torch.from_numpy(np.clip(value.float().numpy() * 255, 0, 255).astype(np.uint8)).float() / 255 + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +def test_exact_unchanged(dtype): + value = sample(4, dtype, (2, 3, 5)) + value.flatten()[0] = -0.0 + expected = codec.encode_payload("IMAGE", value) + assert images.encode_image(value) == expected + result = images.decode_image(*expected) + assert torch.equal(result.view(torch.uint8), value.view(torch.uint8)) + + +@pytest.mark.parametrize("fmt", FORMATS) +@pytest.mark.parametrize("channels", [1, 3, 4]) +def test_real_roundtrips_shape_order_quantization_alpha(fmt, channels): + available(fmt) + value = sample(channels) + # Large inter-frame difference also verifies frame ordering. + value[0] *= 0.2 + value[2] = value[2] * 0.2 + 0.8 + if channels == 4: + value[..., 3] = torch.linspace(-0.1, 1.1, value[..., 3].numel()).reshape(value[..., 3].shape) + descriptor, payload = images.encode_image(value, transfer_format=fmt) + assert descriptor["encoding"] == "image-transfer-v1" + result = codec.decode_payload("IMAGE", descriptor, payload) + assert result.shape == value.shape + assert result.device.type == "cpu" and result.is_contiguous() + color = 3 if channels == 4 else channels + expected = quantized(value[..., :color]) + if fmt in LOSSLESS: + assert torch.equal(result[..., :color], expected) + else: + # Random RGB noise loses chroma detail in 4:2:0 even at high quality. + assert (result[..., :color] - expected).abs().mean() < 0.12 + assert result[0, ..., :color].mean() < result[2, ..., :color].mean() + if channels == 4: + assert torch.equal(result[..., 3], value[..., 3]) + + +@pytest.mark.parametrize("fmt", ["png", "ffv1", "jpeg", "webp", "h264", "h265", "av1"]) +def test_float64_alpha_is_exact(fmt): + available(fmt) + value = sample(4, torch.float64, (1, 65, 67)) + descriptor, payload = images.encode_image(value, transfer_format=fmt) + out = images.decode_image(descriptor, payload) + assert out.dtype == torch.float64 + assert torch.equal(out[..., 3].contiguous().view(torch.uint8), value[..., 3].contiguous().view(torch.uint8)) + + +@pytest.mark.parametrize("fmt", ["jpeg", "webp", "h264", "h265", "av1"]) +def test_lossy_quality_changes_final_payload_size(fmt): + available(fmt) + value = sample(shape=(2, 96, 96)) + low = images.encode_image(value, transfer_format=fmt, transfer_quality=15)[1] + high = images.encode_image(value, transfer_format=fmt, transfer_quality=95)[1] + assert len(high) > len(low) + + +def test_truncation_clipping_grayscale_noncontiguous(): + value = torch.tensor([-0.1, 0.5, 1.1, 0.12345], device="cpu").reshape(1, 2, 2, 1).transpose(1, 2) + out = images.decode_image(*images.encode_image(value, transfer_format="png")) + assert torch.equal(out, quantized(value)) + assert out[0, 0, 1, 0] == 1 + + +@pytest.mark.parametrize("value", [None, [], torch.ones(2, 3, 3), torch.ones(1, 2, 3, 2), torch.ones(1, 2, 3, 0), torch.ones(1, 2, 3, 3, dtype=torch.uint8), torch.ones(1, 2, 3, 3, device="meta"), torch.ones(1, 2, 3, 3).to_sparse()]) +def test_invalid_values(value): + with pytest.raises(ValueError): + images.encode_image(value, transfer_format="png") + + +@pytest.mark.parametrize("bad", [float("nan"), float("inf"), -float("inf")]) +def test_nonfinite_rejected(bad): + with pytest.raises(ValueError, match="finite"): + images.encode_image(torch.full((1, 1, 1, 3), bad), transfer_format="png") + + +@pytest.mark.parametrize("options", [{"transfer_format": "PNG"}, {"transfer_format": []}, {"transfer_quality": True}, {"transfer_quality": 0}, {"transfer_quality": 101}, {"transfer_quality": 1.5}, {"max_bytes": True}, {"max_bytes": -1}]) +def test_strict_options(options): + with pytest.raises(ValueError): + images.encode_image(sample(shape=(1, 1, 1)), **options) + + +@pytest.mark.parametrize("change", [{"extra": 1}, {"format": "exact"}, {"quality": True}, {"layout": "BCHW"}, {"shape": [1, 2**60, 1, 3]}, {"shape": [True, 1, 1, 3]}, {"parts": [0]}, {"parts": [True]}, {"parts": [2**60]}, {"alpha": {"extra": 0}}]) +def test_bad_descriptors(change): + descriptor, payload = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png") + descriptor.update(change) + with pytest.raises(ValueError): + images.decode_image(descriptor, payload) + + +def test_truncated_appended_and_mismatched_parts(): + d, p = images.encode_image(sample(shape=(2, 3, 5)), transfer_format="png") + for payload in (p[:-1], p + b"x"): + with pytest.raises(ValueError): + images.decode_image(d, payload) + for change in ({"shape": [1, 3, 5, 3]}, {"parts": [len(p)]}, {"shape": [2, 2, 5, 3]}): + bad = dict(d, **change) + with pytest.raises(ValueError): + images.decode_image(bad, p) + + +@pytest.mark.parametrize("fmt", ["png", "jpeg", "webp", "ffv1", "h264", "h265", "av1"]) +def test_payload_and_decoded_budget(fmt): + available(fmt) + value = sample(shape=(1, 65, 67)) + d, p = images.encode_image(value, transfer_format=fmt) + with pytest.raises(ValueError, match="limit|budget"): + images.decode_image(d, p, max_bytes=100) + with pytest.raises(ValueError, match="limit|budget"): + images.encode_image(value, transfer_format=fmt, max_bytes=100) + + +def test_oversized_shape_before_materialization(monkeypatch): + value = torch.ones(1).expand(1, 8192, 8192, 4) + def forbidden(*a, **kw): + pytest.fail("Materialized an oversized tensor") + monkeypatch.setattr(torch.Tensor, "contiguous", forbidden) + monkeypatch.setattr(torch, "isfinite", forbidden) + with pytest.raises(ValueError): + images.encode_image(value, transfer_format="png") + + +def test_header_dimensions_checked_before_pixels(monkeypatch): + d, p = images.encode_image(sample(shape=(1, 9, 11)), transfer_format="png") + d["shape"] = [1, 1, 1, 3] + def forbidden(*a, **kw): + pytest.fail("Pixels decoded before header dimensions were checked") + monkeypatch.setattr(Image.Image, "load", forbidden) + with pytest.raises(ValueError): + images.decode_image(d, p) + + +def test_other_kinds_do_not_dispatch(): + d, p = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png") + for kind in ("MASK", "AUDIO", "STRING"): + with pytest.raises(ValueError): + codec.decode_payload(kind, d, p) + + +def test_missing_video_encoder_is_explicit(monkeypatch): + monkeypatch.setitem(images._VIDEO_CODECS, "h264", ("missing_software_encoder", 51, "yuv420p")) + with pytest.raises(ValueError, match="encoder|unavailable"): + images.encode_image(sample(shape=(1, 64, 64)), transfer_format="h264") + + +def test_bounded_sink_catches_final_flush_bytes(): + sink = images._BoundedIO(4) + sink.write(b"1234") + with pytest.raises(ValueError, match="limit"): + sink.write(b"5") + assert sink.getvalue() == b"1234" + + +def test_crf_mapping(): + assert images._quality_to_crf(90, 51) == 13 + assert images._quality_to_crf(90, 63) == 16 + assert images._quality_to_crf(100, 63) == 0 + + +@pytest.mark.parametrize("fmt", FORMATS) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +def test_compressed_preserves_source_dtype_and_exact_alpha(fmt, dtype): + available(fmt) + value = sample(4, dtype, (1, 17, 19)) + value[..., 3].flatten()[0] = -0.0 + d, p = images.encode_image(value, transfer_format=fmt) + with torch.device("meta"): + result = images.decode_image(d, p) + assert result.device.type == "cpu" and result.dtype == dtype + assert d["dtype"] == str(dtype).removeprefix("torch.") + assert torch.equal(result[..., 3].contiguous().view(torch.uint8), value[..., 3].contiguous().view(torch.uint8)) + if fmt in LOSSLESS: + assert torch.equal(result[..., :3], quantized(value[..., :3]).to(dtype)) + + +@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"]) +def test_video_count_and_header_fail_before_native_decoder(fmt, monkeypatch): + available(fmt) + d, p = images.encode_image(sample(shape=(2, 65, 67)), transfer_format=fmt) + def forbidden(*a, **kw): + pytest.fail("Native decoder reached before header validation") + monkeypatch.setattr(images, "_decode_video", forbidden) + for shape in ([1, 65, 67, 3], [3, 65, 67, 3], [2, 32, 32, 3]): + bad = dict(d, shape=shape) + with pytest.raises(ValueError): + images.decode_image(bad, p) + for payload in (p[:-1], p[:40], p + b"extra"): + bad = dict(d, parts=[len(payload)]) + with pytest.raises(ValueError): + images.decode_image(bad, payload) + + +@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"]) +def test_video_does_not_use_native_container_probe(fmt, monkeypatch): + available(fmt) + d, p = images.encode_image(sample(shape=(2, 65, 67)), transfer_format=fmt) + av = images._import_av() + def forbidden(*a, **kw): + pytest.fail("Unbounded native stream probing is forbidden") + monkeypatch.setattr(av, "open", forbidden) + assert images.decode_image(d, p).shape == (2, 65, 67, 3) + + +@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"]) +def test_real_video_flush_bound(fmt): + available(fmt) + frames = [np.zeros((64, 64, 3), dtype=np.uint8)] + good = images._encode_video(iter(frames), fmt, 90, 64, 64, 100000) + with pytest.raises(ValueError, match="limit"): + images._encode_video(iter(frames), fmt, 90, 64, 64, len(good) - 1) + + +def test_alpha_invalid_dtype_nonfinite_and_extra_fields(): + import copy + d, p = images.encode_image(sample(4, shape=(1, 3, 5)), transfer_format="png") + for field, value in (("dtype", "float64"), ("shape", [1, 3, 6, 1]), ("extra", 1)): + bad = copy.deepcopy(d) + bad["alpha"]["descriptor"][field] = value + with pytest.raises(ValueError): + images.decode_image(bad, p) + import struct + bad = p[:-4] + struct.pack(" compiler -> endpoint validation, no network or GPU.""" +import copy +import json + +import pytest + +from test_serverless_media_prompt import compiler, _prompt +from test_media_handler import handler +from test_media_collector import module, collect, native_media + +FORMATS = ('exact', 'legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1') + + +def compile_request(fmt='exact', quality=90): + prompt = _prompt() + prompt['source:1'] = {'class_type': 'Source', 'inputs': {}} + prompt['collector:1']['inputs'].update(transfer_format=fmt, transfer_quality=quality, + masks=['source:1', 1], texts=['source:1', 2]) + original = copy.deepcopy(prompt) + workflow, outputs = compiler.compile_media_prompt(prompt, 'request:1') + assert prompt == original + return {'workflow': workflow, 'images': [], 'media': { + 'version': 1, 'profile': 'comfy-media-v1', 'request_id': 'request:1', + 'artifacts': [], 'inputs': [], 'outputs': outputs}} + + +@pytest.mark.parametrize('fmt', FORMATS) +def test_compiler_and_handler_preserve_image_policy_only(fmt): + data = compile_request(fmt, 73) + for binding in data['media']['outputs']: + inputs = data['workflow'][binding['node_id']]['inputs'] + if binding['kind'] == 'IMAGE' and fmt != 'exact': + assert inputs['transfer_format'] == fmt + assert inputs['transfer_quality'] == 73 + else: + assert 'transfer_format' not in inputs and 'transfer_quality' not in inputs + handler._validate_request(data) + + +@pytest.mark.parametrize('fmt,quality', [('JPEG',90),('',90),('jpeg',True),('jpeg',0), + ('jpeg',101),('jpeg','90'),(['source:1',0],90)]) +def test_compiler_rejects_invalid_or_linked_controls(fmt, quality): + with pytest.raises(ValueError): + compile_request(fmt, quality) + + +@pytest.mark.parametrize('field,value', [('transfer_format','unknown'),('transfer_quality',0), + ('transfer_quality',True),('transfer_quality','90')]) +def test_handler_rejects_tampered_policy(field, value): + data = compile_request('jpeg') + binding = next(b for b in data['media']['outputs'] if b['kind'] == 'IMAGE') + data['workflow'][binding['node_id']]['inputs'][field] = value + with pytest.raises(ValueError): + handler._validate_request(data) + + +def test_handler_rejects_image_controls_on_mask_export(): + data = compile_request() + binding = next(b for b in data['media']['outputs'] if b['kind'] == 'MASK') + data['workflow'][binding['node_id']]['inputs']['transfer_format'] = 'jpeg' + with pytest.raises(ValueError): + handler._validate_request(data) + + +@pytest.mark.parametrize('fmt', FORMATS) +def test_master_passthrough_stays_native_for_every_format(module, fmt): + media = native_media() + actual = collect(module, **media, transfer_format=[fmt], transfer_quality=[73]) + for values, original in zip(actual, media.values()): + assert values is original + + +@pytest.mark.parametrize('quality', [True, 0, 101, '90', [90,91]]) +def test_node_rejects_bad_quality_even_without_remote_participants(module, quality): + with pytest.raises(ValueError): + collect(module, texts=['hello'], transfer_quality=quality)