feat: add explicit image return codecs and quality controls
This commit is contained in:
Robert Wojciechowski
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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("<HH", data[26:30])
|
||||
return height & 0x3fff, width & 0x3fff
|
||||
if data[:2] != b"\xff\xd8" or data[-2:] != b"\xff\xd9":
|
||||
raise ValueError("Invalid IMAGE JPEG header or trailer")
|
||||
offset, dimensions = 2, None
|
||||
while offset + 4 <= len(data):
|
||||
if data[offset] != 0xff:
|
||||
raise ValueError("Invalid IMAGE JPEG marker")
|
||||
marker = data[offset + 1]
|
||||
size = int.from_bytes(data[offset + 2:offset + 4], "big")
|
||||
end = offset + 2 + size
|
||||
if size < 2 or end > 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("<ii", data[cs + 4:cs + 12]) != (width, height):
|
||||
raise ValueError("IMAGE FFV1 header dimensions mismatch")
|
||||
tracks.append(fields)
|
||||
elif tag == 0x1f43b675:
|
||||
def block(bs, be):
|
||||
nonlocal blocks
|
||||
track, pos = vint(bs, be)
|
||||
if track != 1 or pos + 3 >= 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
|
||||
@@ -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
|
||||
@@ -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]})
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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("<f", float("nan"))
|
||||
with pytest.raises(ValueError, match="finite"):
|
||||
images.decode_image(d, bad)
|
||||
|
||||
|
||||
def test_descriptor_structure_bound_and_frame_limit():
|
||||
d, p = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
|
||||
for bad in (dict(d, quality="x" * 70000), dict(d, parts=[1] * 241), dict(d, shape=[241, 1, 1, 3]), dict(d, dtype="x")):
|
||||
with pytest.raises(ValueError):
|
||||
images.decode_image(bad, p)
|
||||
with pytest.raises(ValueError):
|
||||
images.encode_image(torch.zeros(241, 1, 1, 3), transfer_format="png")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["png", "jpeg", "webp"])
|
||||
def test_all_still_headers_checked_before_pillow_open(fmt, monkeypatch):
|
||||
d, p = images.encode_image(sample(shape=(1, 31, 33)), transfer_format=fmt)
|
||||
d["shape"] = [1, 2, 2, 3]
|
||||
def forbidden(*a, **kw):
|
||||
pytest.fail("Pillow opened a mismatched header")
|
||||
monkeypatch.setattr(Image, "open", forbidden)
|
||||
with pytest.raises(ValueError, match="dimensions"):
|
||||
images.decode_image(d, p)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
|
||||
def test_video_longer_batch_order_and_coded_allocation_bound(fmt):
|
||||
available(fmt)
|
||||
value = torch.linspace(0, 1, 16).reshape(16, 1, 1, 1).expand(16, 64, 64, 3)
|
||||
d, p = images.encode_image(value, transfer_format=fmt)
|
||||
header = images._video_header(p, fmt, 64, 64, 16)
|
||||
coded = images._coded_frame_bound(p, fmt, header)
|
||||
assert 16 <= coded <= images.MAX_FRAMES * 4
|
||||
result = images.decode_image(d, p)
|
||||
assert torch.all((result.mean(dim=(1, 2, 3)) - value.mean(dim=(1, 2, 3))).abs() < 0.03)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["h264", "h265", "av1"])
|
||||
def test_packet_unit_bounds_before_native_decode(fmt):
|
||||
available(fmt)
|
||||
d, p = images.encode_image(sample(shape=(1, 64, 64)), transfer_format=fmt)
|
||||
config, bounds = images._video_header(p, fmt, 64, 64, 1)
|
||||
start, end = bounds[0]
|
||||
packet = p[start:end]
|
||||
with pytest.raises(ValueError):
|
||||
images._coded_frame_bound(packet[:-1], fmt, (config, [(0, len(packet) - 1)]))
|
||||
# An adversarial single block can contain many coded frames: bound it too.
|
||||
enormous = packet * (images.MAX_FRAMES * 4 + 1)
|
||||
with pytest.raises(ValueError, match="limit"):
|
||||
images._coded_frame_bound(enormous, fmt, (config, [(0, len(enormous))]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["h264", "h265", "av1"])
|
||||
def test_bitstream_resolution_cannot_bypass_native_pixel_limit(fmt, monkeypatch):
|
||||
available(fmt)
|
||||
d, p = images.encode_image(sample(shape=(1, 256, 256)), transfer_format=fmt)
|
||||
header = images._video_header(p, fmt, 256, 256, 1)
|
||||
av = images._import_av()
|
||||
# A lying container is not enough to disable max_pixels in the decoder.
|
||||
# The bitstream still contains 256x256 pixels but the budget permits 64x64.
|
||||
with pytest.raises((ValueError, av.error.FFmpegError)):
|
||||
list(images._decode_video(p, fmt, 64, 64, 1, header))
|
||||
@@ -0,0 +1,76 @@
|
||||
"""Collector controls -> 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)
|
||||
Reference in New Issue
Block a user