feat: add explicit image return codecs and quality controls
This commit is contained in:
Robert Wojciechowski
@@ -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
|
||||
Reference in New Issue
Block a user