Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
593047badf | ||
|
|
a9ef36f951 | ||
|
|
b636fd5202 | ||
|
|
f0927f1166 |
+19
-5
@@ -97,10 +97,11 @@ _SECRET_FIELDS = {"api_key", "clear_api_key"}
|
||||
_SERVERLESS_LOCAL_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
|
||||
_SERVERLESS_ONLY_FIELDS = {
|
||||
"provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "api_key_configured",
|
||||
"media_profile",
|
||||
}
|
||||
_SERVERLESS_CONFIG_FIELDS = {
|
||||
"worker_id", "type", "provider", "endpoint_id", "jobs_per_run", "concurrency",
|
||||
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key",
|
||||
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key", "media_profile",
|
||||
}
|
||||
|
||||
|
||||
@@ -199,7 +200,13 @@ def _clean_bulk_type_conversions(old_workers, new_workers):
|
||||
if not isinstance(worker, dict) or worker.get("id") is None:
|
||||
continue
|
||||
old = old_by_id.get(str(worker["id"]))
|
||||
if old is None or old.get("type") != "serverless" or worker.get("type") == "serverless":
|
||||
if old is None or old.get("type") != "serverless":
|
||||
continue
|
||||
if worker.get("type") == "serverless":
|
||||
# Older clients do not know this field. Only an explicit selection
|
||||
# may change the endpoint contract, including on bulk saves.
|
||||
if "media_profile" not in worker and "media_profile" in old:
|
||||
worker["media_profile"] = old["media_profile"]
|
||||
continue
|
||||
converted = dict(worker)
|
||||
for field in _SERVERLESS_ONLY_FIELDS:
|
||||
@@ -218,7 +225,10 @@ def _bulk_serverless_changes(old_workers, new_workers):
|
||||
if new is None or new.get("type") != "serverless":
|
||||
_ensure_pool_idle(worker_id)
|
||||
secret_ids.add(worker_id)
|
||||
elif old.get("endpoint_id") != new.get("endpoint_id"):
|
||||
elif (
|
||||
old.get("endpoint_id") != new.get("endpoint_id")
|
||||
or old.get("media_profile", "legacy-image-v1") != new.get("media_profile", "legacy-image-v1")
|
||||
):
|
||||
_ensure_pool_idle(worker_id)
|
||||
elif new is not None and new.get("type") == "serverless":
|
||||
secret_ids.add(worker_id)
|
||||
@@ -401,7 +411,7 @@ async def _update_worker_endpoint(request, secret_changes):
|
||||
candidate.setdefault("provider", "runpod")
|
||||
candidate.setdefault("enabled", False)
|
||||
# Legacy "concurrency" is accepted (see _SERVERLESS_CONFIG_FIELDS) but ignored.
|
||||
for field in ("provider", "endpoint_id", "jobs_per_run", "timeout_seconds", "enabled", "name"):
|
||||
for field in ("provider", "endpoint_id", "jobs_per_run", "timeout_seconds", "enabled", "name", "media_profile"):
|
||||
if field in data:
|
||||
candidate[field] = data[field]
|
||||
candidate = _clean_serverless_worker(candidate)
|
||||
@@ -424,8 +434,12 @@ async def _update_worker_endpoint(request, secret_changes):
|
||||
current_type == "serverless"
|
||||
and existing.get("endpoint_id") != validated.get("endpoint_id")
|
||||
)
|
||||
media_profile_changed = (
|
||||
current_type == "serverless"
|
||||
and existing.get("media_profile", "legacy-image-v1") != validated.get("media_profile", "legacy-image-v1")
|
||||
)
|
||||
if current_type == "serverless" and (
|
||||
endpoint_changed or next_secret != prior_secret
|
||||
endpoint_changed or media_profile_changed or next_secret != prior_secret
|
||||
):
|
||||
_ensure_pool_idle(worker_id)
|
||||
|
||||
|
||||
@@ -37,6 +37,8 @@ prompt_server = server.PromptServer.instance
|
||||
# envelope: JSON { "job_id", "worker_id", "format", "start_index", "count", "width", "height",
|
||||
# "is_last", "audio"?: {...} }
|
||||
# media: one part per frame for still formats, or one video file for video formats
|
||||
# format "exact" adds envelope "descriptors" (one tensor-le-v1 descriptor per
|
||||
# frame) and sends one raw tensor part per frame; values are not quantized.
|
||||
|
||||
|
||||
def _decode_image_sync(image_path):
|
||||
@@ -299,6 +301,8 @@ async def _enqueue_worker_items(multi_job_id, items):
|
||||
|
||||
def _decode_compressed_chunk_sync(envelope, parts):
|
||||
"""Decode a compressed chunk into per-frame IMAGE tensors (threadpool)."""
|
||||
if envelope.get("format") == transfer_codec.EXACT_FORMAT:
|
||||
return _decode_exact_chunk_sync(envelope, parts)
|
||||
frames = transfer_codec.decode_chunk(envelope, parts)
|
||||
return [
|
||||
ensure_contiguous(torch.from_numpy(transfer_codec.uint8_frame_to_tensor_array(frame)))
|
||||
@@ -306,6 +310,35 @@ def _decode_compressed_chunk_sync(envelope, parts):
|
||||
]
|
||||
|
||||
|
||||
def _decode_exact_chunk_sync(envelope, parts):
|
||||
"""Decode raw exact IMAGE frames, checking each against the envelope geometry."""
|
||||
from ..serverless_bridge import codec as bridge_codec
|
||||
count, width, height = envelope.get("count"), envelope.get("width"), envelope.get("height")
|
||||
for name, value in (("count", count), ("width", width), ("height", height)):
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
|
||||
raise transfer_codec.TransferCodecError(f"envelope.{name}: expected positive integer")
|
||||
descriptors = envelope.get("descriptors")
|
||||
if not isinstance(descriptors, list) or len(descriptors) != count or len(parts) != count:
|
||||
raise transfer_codec.TransferCodecError(
|
||||
f"exact chunk expected {count} descriptors and frames, got "
|
||||
f"{len(descriptors) if isinstance(descriptors, list) else 'no'} descriptors and {len(parts)} frames"
|
||||
)
|
||||
tensors = []
|
||||
for descriptor, payload in zip(descriptors, parts):
|
||||
if not isinstance(descriptor, dict) or descriptor.get("encoding") != "tensor-le-v1":
|
||||
raise transfer_codec.TransferCodecError("exact frames require tensor-le-v1 encoding")
|
||||
shape = descriptor.get("shape") if isinstance(descriptor, dict) else None
|
||||
if not isinstance(shape, list) or len(shape) != 4 or shape[:3] != [1, height, width]:
|
||||
raise transfer_codec.TransferCodecError(
|
||||
f"exact frame shape {shape} does not match the envelope's 1x{height}x{width} frames"
|
||||
)
|
||||
try:
|
||||
tensors.append(bridge_codec.decode_payload("IMAGE", descriptor, payload))
|
||||
except ValueError as exc:
|
||||
raise transfer_codec.TransferCodecError(f"exact frame: {exc}") from None
|
||||
return tensors
|
||||
|
||||
|
||||
def _validate_compressed_envelope(envelope):
|
||||
errors = []
|
||||
if not isinstance(envelope, dict):
|
||||
@@ -384,6 +417,7 @@ async def _job_complete_compressed(request):
|
||||
"image_index": start_index + pos,
|
||||
"is_last": bool(is_last and pos == last_pos),
|
||||
"audio": decoded_audio if pos == last_pos else None,
|
||||
"format": envelope["format"],
|
||||
}
|
||||
for pos, tensor in enumerate(tensors)
|
||||
]
|
||||
@@ -462,6 +496,7 @@ async def job_complete_endpoint(request):
|
||||
"image_index": int(batch_idx),
|
||||
"is_last": is_last,
|
||||
"audio": decoded_audio,
|
||||
"format": transfer_codec.LEGACY_FORMAT,
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
@@ -18,7 +18,7 @@ class AmbiguousSubmission(RunpodError):
|
||||
"""Submission may have been accepted: never retry it automatically."""
|
||||
|
||||
|
||||
def serialize_request(workflow, images, timeout_seconds):
|
||||
def serialize_request(workflow, images, timeout_seconds, *, media=None):
|
||||
"""Serialize and size-check the exact JSON body used by Runpod ``/run``."""
|
||||
payload = {
|
||||
"input": {"workflow": workflow, "images": images},
|
||||
@@ -27,6 +27,10 @@ def serialize_request(workflow, images, timeout_seconds):
|
||||
"ttl": timeout_seconds * 1000,
|
||||
},
|
||||
}
|
||||
if media is not None:
|
||||
from ...serverless_bridge.handler import _validate_request
|
||||
payload["input"]["media"] = media
|
||||
_validate_request(payload["input"])
|
||||
body = json.dumps(
|
||||
payload,
|
||||
allow_nan=False,
|
||||
@@ -108,8 +112,8 @@ class RunpodClient:
|
||||
async def health(self):
|
||||
return await self._request("GET", "health")
|
||||
|
||||
async def submit(self, workflow, images, timeout_seconds):
|
||||
body = serialize_request(workflow, images, timeout_seconds)
|
||||
async def submit(self, workflow, images, timeout_seconds, *, media=None):
|
||||
body = serialize_request(workflow, images, timeout_seconds, media=media)
|
||||
result = await self._request("POST", "run", request_body=body)
|
||||
try:
|
||||
return checked_id(result.get("id"))
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
"""Pure compilation of prepared Distributed Collectors for comfy-media-v1 pools.
|
||||
|
||||
The legacy image compiler remains separate. This compiler emits one MediaExportV1
|
||||
execution target per connected Collector socket (``images`` and/or ``audio``),
|
||||
retaining only its upstream computation. It never imports ComfyUI or
|
||||
inspects/executes runtime node classes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import hashlib
|
||||
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, validate_options
|
||||
from .serverless_prompt import (
|
||||
_PRIMITIVE_NODES,
|
||||
_REMOVED_SCALAR_NODES,
|
||||
_is_link,
|
||||
_is_unsupported_distributed,
|
||||
_make_scalar_resolver,
|
||||
_prune_graph,
|
||||
)
|
||||
|
||||
|
||||
_COLLECTOR = "DistributedCollector"
|
||||
# Collector inputs, in output-slot order, that a typed pool returns.
|
||||
COLLECTOR_SOCKETS = ("images", "audio")
|
||||
# Collector widget defaults (the saved-workflow contract is legacy_png / 90).
|
||||
COLLECTOR_DEFAULT_FORMAT = "legacy_png"
|
||||
COLLECTOR_DEFAULT_QUALITY = 90
|
||||
|
||||
|
||||
def collector_image_policy(inputs):
|
||||
"""Return the endpoint IMAGE (format, quality) for one Collector.
|
||||
|
||||
legacy_png is the Collector's saved default and names a JSON/base64 framing
|
||||
that only applies to direct workers; serverless returns the same 8-bit
|
||||
pixels as compressed png. exact stays opt-in. No other format is altered.
|
||||
"""
|
||||
fmt, quality = validate_options(inputs.get("transfer_format", COLLECTOR_DEFAULT_FORMAT),
|
||||
inputs.get("transfer_quality", COLLECTOR_DEFAULT_QUALITY))
|
||||
return ("png" if fmt == "legacy_png" else fmt), quality
|
||||
|
||||
|
||||
def _validate_identifier(value, context):
|
||||
# Validate the original token before any lookup. In particular, neither
|
||||
# integer node IDs nor whitespace in prepared identities is normalized.
|
||||
try:
|
||||
identifier(value)
|
||||
except ValueError:
|
||||
raise ValueError(f"{context} must be a valid nonempty media identifier") from None
|
||||
|
||||
|
||||
def _validate_prompt(prepared_prompt):
|
||||
"""Validate the entire source graph, including branches later discarded."""
|
||||
if not isinstance(prepared_prompt, dict) or not prepared_prompt:
|
||||
raise ValueError("prepared prompt must be a nonempty API graph mapping")
|
||||
|
||||
collectors = []
|
||||
for node_id, node in prepared_prompt.items():
|
||||
_validate_identifier(node_id, "prompt node ID")
|
||||
if not isinstance(node, dict):
|
||||
raise ValueError(f"prompt node {node_id!r} must be an object")
|
||||
class_type = node.get("class_type")
|
||||
if not isinstance(class_type, str) or not class_type.strip():
|
||||
raise ValueError(f"prompt node {node_id!r} must have a nonempty class_type")
|
||||
if class_type == _COLLECTOR:
|
||||
collectors.append(node_id)
|
||||
elif _is_unsupported_distributed(class_type):
|
||||
raise ValueError(f"unsupported distributed or batch node: {class_type}")
|
||||
inputs = node.get("inputs", {})
|
||||
if not isinstance(inputs, dict):
|
||||
raise ValueError(f"prompt node {node_id!r} inputs must be an object")
|
||||
if any(not isinstance(name, str) or not name for name in inputs):
|
||||
raise ValueError(f"prompt node {node_id!r} input names must be nonempty strings")
|
||||
|
||||
if not collectors:
|
||||
raise ValueError("prepared prompt contains no DistributedCollector nodes")
|
||||
|
||||
dependencies = {}
|
||||
for node_id, node in prepared_prompt.items():
|
||||
parents = set()
|
||||
for input_name, value in node.get("inputs", {}).items():
|
||||
# Ordinary inputs may contain literal lists. As in the legacy
|
||||
# compiler, exactly two elements denotes an API link, not media.
|
||||
if not _is_link(value):
|
||||
continue
|
||||
source_id, slot = value
|
||||
context = f"link at node {node_id!r} input {input_name!r}"
|
||||
_validate_identifier(source_id, context + " source ID")
|
||||
if source_id not in prepared_prompt:
|
||||
raise ValueError(f"invalid {context}: source {source_id!r} does not exist")
|
||||
if type(slot) is not int or slot < 0:
|
||||
raise ValueError(f"invalid {context}: output slot must be a nonnegative integer")
|
||||
source_class = prepared_prompt[source_id]["class_type"]
|
||||
if source_class in _REMOVED_SCALAR_NODES or source_class in _PRIMITIVE_NODES:
|
||||
if slot != 0:
|
||||
raise ValueError(f"invalid output slot {slot} for scalar node {source_id!r}")
|
||||
elif source_class == _COLLECTOR and slot >= len(COLLECTOR_SOCKETS):
|
||||
raise ValueError(f"invalid output slot {slot} for collector {source_id!r}")
|
||||
parents.add(source_id)
|
||||
dependencies[node_id] = parents
|
||||
|
||||
for collector_id in collectors:
|
||||
inputs = prepared_prompt[collector_id].get("inputs", {})
|
||||
collector_image_policy(inputs)
|
||||
for field in ("multi_job_id", "worker_id"):
|
||||
_validate_identifier(inputs.get(field), f"collector {collector_id!r} {field}")
|
||||
connected = [socket for socket in COLLECTOR_SOCKETS if socket in inputs]
|
||||
if not connected:
|
||||
raise ValueError(f"collector {collector_id!r} requires at least one connected media socket")
|
||||
for socket in connected:
|
||||
value = inputs[socket]
|
||||
if not _is_link(value):
|
||||
raise ValueError(f"collector {collector_id!r} {socket} must be linked; use PrimitiveString for literal text")
|
||||
if prepared_prompt[value[0]]["class_type"] in _REMOVED_SCALAR_NODES:
|
||||
# Export sockets must remain actual links after scalar folding.
|
||||
# A PrimitiveString wrapper supports distributed text values.
|
||||
raise ValueError(f"collector {collector_id!r} {socket} cannot link directly to a distributed scalar; use a typed node or PrimitiveString")
|
||||
|
||||
# Iterative topological sorting checks ordinary and unused cycles too, not
|
||||
# only the primitive chains that the shared scalar resolver traverses.
|
||||
try:
|
||||
ordered = list(TopologicalSorter(dependencies).static_order())
|
||||
except CycleError:
|
||||
raise ValueError("cycle in prepared media prompt") from None
|
||||
collector_set = set(collectors)
|
||||
has_collector_ancestor = {}
|
||||
for node_id in ordered:
|
||||
has_collector_ancestor[node_id] = any(
|
||||
parent in collector_set or has_collector_ancestor[parent]
|
||||
for parent in dependencies[node_id]
|
||||
)
|
||||
if node_id in collector_set and has_collector_ancestor[node_id]:
|
||||
raise ValueError("collector-to-collector dependence is unsupported")
|
||||
|
||||
return sorted(collectors)
|
||||
|
||||
|
||||
def _build_exports(source_nodes, collector_ids, request_id):
|
||||
bindings = []
|
||||
exports = {}
|
||||
reserved = set(source_nodes)
|
||||
for collector_id in collector_ids:
|
||||
inputs = source_nodes[collector_id].get("inputs", {})
|
||||
for socket in COLLECTOR_SOCKETS:
|
||||
kind = SOCKET_KINDS[socket]
|
||||
if socket not in inputs:
|
||||
continue
|
||||
# Hashing avoids separator ambiguity and keeps long valid source
|
||||
# identities within the bridge's identifier length bound. Identity
|
||||
# is stable across request IDs and dictionary insertion orders.
|
||||
identity = json.dumps([collector_id, socket], separators=(",", ":"))
|
||||
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()
|
||||
base_id = f"serverless:media:{digest}"
|
||||
node_id = base_id
|
||||
suffix = 0
|
||||
while node_id in reserved:
|
||||
suffix += 1
|
||||
node_id = f"{base_id}:{suffix}"
|
||||
reserved.add(node_id)
|
||||
binding = {
|
||||
"binding_id": f"b:{digest}",
|
||||
"collector_id": collector_id,
|
||||
"job_id": inputs["multi_job_id"],
|
||||
"participant_id": inputs["worker_id"],
|
||||
"node_id": node_id,
|
||||
"socket": socket,
|
||||
"kind": kind,
|
||||
}
|
||||
bindings.append(binding)
|
||||
exports[node_id] = {
|
||||
"class_type": "MediaExportV1",
|
||||
"inputs": {
|
||||
"request_id": request_id,
|
||||
"binding_json": json.dumps(binding),
|
||||
socket: copy.deepcopy(inputs[socket]),
|
||||
},
|
||||
}
|
||||
if kind == "IMAGE":
|
||||
fmt, quality = collector_image_policy(inputs)
|
||||
if fmt != DEFAULT_FORMAT:
|
||||
exports[node_id]["inputs"].update(transfer_format=fmt, transfer_quality=quality)
|
||||
validate_bindings(bindings)
|
||||
return bindings, exports
|
||||
|
||||
|
||||
def compile_media_prompt(prepared_prompt: dict, request_id: str) -> tuple[dict, list[dict]]:
|
||||
"""Return a worker graph and canonical typed-media export bindings.
|
||||
|
||||
Each DistributedCollector must have prepared literal ``multi_job_id`` and
|
||||
``worker_id`` identities and at least one linked ``images``/``audio`` input.
|
||||
Present media inputs must be API links. Collector transport controls become
|
||||
the IMAGE export policy (``legacy_png`` is sent as ``png``); other controls
|
||||
are not exported.
|
||||
|
||||
Only MediaExportV1 targets and their ancestors survive. Required ordinary
|
||||
output/utility nodes are preserved regardless of their runtime OUTPUT_NODE
|
||||
flag; callers must dispatch exactly the returned binding ``node_id`` values
|
||||
as partial_execution_targets. Distributed scalar dependencies use the same
|
||||
resolver as the legacy compiler. Direct scalar-to-media links need a typed
|
||||
wrapper so that each export input remains a link.
|
||||
|
||||
Validation and compilation never modify the caller's source graph. This
|
||||
structural compiler does not infer socket types or unknown nodes' output
|
||||
counts; ComfyUI performs runtime type validation at the endpoint.
|
||||
"""
|
||||
_validate_identifier(request_id, "request_id")
|
||||
collector_ids = _validate_prompt(prepared_prompt)
|
||||
bindings, exports = _build_exports(prepared_prompt, collector_ids, request_id)
|
||||
compiled = _prune_graph(
|
||||
prepared_prompt, bindings, exports, _make_scalar_resolver(prepared_prompt)
|
||||
)
|
||||
return {node_id: compiled[node_id] for node_id in sorted(compiled)}, bindings
|
||||
@@ -146,6 +146,10 @@ class ServerlessManager:
|
||||
|
||||
async def start(self, execution_id, prompt_id, prepared_jobs):
|
||||
self.ensure_capacity(len(prepared_jobs))
|
||||
from .serverless_typed_results import validate_prepared_job
|
||||
# Validate the entire batch before creating even the first local slot.
|
||||
for prepared in prepared_jobs:
|
||||
validate_prepared_job(prepared)
|
||||
new_runs = []
|
||||
# Journal every local slot before starting any billed submission.
|
||||
for prepared in prepared_jobs:
|
||||
@@ -177,14 +181,30 @@ class ServerlessManager:
|
||||
return any(run["worker"]["id"] == participant_id and
|
||||
any(m["job_id"] == collector_job_id for m in run["manifest"]) for run in self.runs.values())
|
||||
|
||||
async def _deliver(self, run, images=None, error=None):
|
||||
async with self.server.distributed_jobs_lock:
|
||||
async def _deliver(self, run, images=None, error=None, *, media=None):
|
||||
if media is not None and not error:
|
||||
# Validate all destinations before exposing any partial success.
|
||||
for binding in run["manifest"]:
|
||||
if binding["participant_id"] != run["worker"]["id"]:
|
||||
raise ValueError("Serverless media participant identity mismatch")
|
||||
values = media[binding["job_id"]][binding["socket"]]
|
||||
if not isinstance(values, list) or not values:
|
||||
raise ValueError("Missing serverless media values")
|
||||
async with self.server.distributed_jobs_lock:
|
||||
delivered = set()
|
||||
for binding in run["manifest"]:
|
||||
job_id = binding["job_id"]
|
||||
if error or media is not None:
|
||||
if job_id in delivered:
|
||||
continue
|
||||
delivered.add(job_id)
|
||||
queue = self.server.distributed_pending_jobs.get(binding["job_id"])
|
||||
if queue is None:
|
||||
continue # A cancelled/finished collector must never be resurrected.
|
||||
if error:
|
||||
queue.put_nowait({"worker_id": run["worker"]["id"], "serverless_error": error, "is_last": True})
|
||||
elif media is not None:
|
||||
queue.put_nowait({"worker_id": run["worker"]["id"], "media": media[job_id], "is_last": True})
|
||||
else:
|
||||
batch = images[binding["job_id"]]
|
||||
for index, tensor in enumerate(batch):
|
||||
@@ -218,13 +238,18 @@ class ServerlessManager:
|
||||
submitted = False
|
||||
provider_terminal = False
|
||||
try:
|
||||
from .serverless_typed_results import validate_prepared_job, collect_typed_results
|
||||
typed = validate_prepared_job(run)
|
||||
async with asyncio.timeout(max(0, run["deadline"] - time.monotonic())):
|
||||
async with semaphore:
|
||||
await self._reserve_provider_capacity(run)
|
||||
client = await self._client(run["worker"])
|
||||
run["client"] = client
|
||||
submitted = True
|
||||
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
|
||||
if typed:
|
||||
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"], media=run["media"])
|
||||
else:
|
||||
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
|
||||
run["provider_job_id"] = job_id
|
||||
# Persist immediately: all subsequent actions use this exact ID.
|
||||
self.ledger.update(rid, job_id=job_id, status="IN_QUEUE", updated_at=time.time())
|
||||
@@ -240,8 +265,12 @@ class ServerlessManager:
|
||||
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
|
||||
raise RunpodError("Serverless workflow reported an output error; inspect the provider job")
|
||||
self.ledger.update(rid, status="DOWNLOADING", updated_at=time.time())
|
||||
images = await decode_results(output, run["manifest"])
|
||||
await self._deliver(run, images=images)
|
||||
if typed:
|
||||
media = await collect_typed_results(output, run["media"], run["manifest"])
|
||||
await self._deliver(run, media=media)
|
||||
else:
|
||||
images = await decode_results(output, run["manifest"])
|
||||
await self._deliver(run, images=images)
|
||||
self.ledger.update(rid, status="COMPLETED", updated_at=time.time())
|
||||
return
|
||||
self.ledger.update(rid, status=status, updated_at=time.time())
|
||||
@@ -276,6 +305,7 @@ class ServerlessManager:
|
||||
# base64 inputs, workflow metadata, or a client holding credentials.
|
||||
run.pop("workflow", None)
|
||||
run.pop("images", None)
|
||||
run.pop("media", None)
|
||||
run["client"] = None
|
||||
# A released local semaphore alone cannot release a durable reservation.
|
||||
await self._notify_capacity()
|
||||
@@ -359,6 +389,7 @@ class ServerlessManager:
|
||||
self.ledger.update(rid, status="CANCELLED", updated_at=time.time())
|
||||
run.pop("workflow", None)
|
||||
run.pop("images", None)
|
||||
run.pop("media", None)
|
||||
await self._deliver(run, error="Serverless job cancelled before submission")
|
||||
continue
|
||||
# Do not send two cancellation requests for the same live job in
|
||||
|
||||
@@ -0,0 +1,225 @@
|
||||
"""Bounded local native inputs for comfy-media-v1 requests.
|
||||
|
||||
Only stock LoadAudio/LoadVideo schemas are adapted. Original nodes and links
|
||||
remain in the private request graph for endpoint validation; only the existing
|
||||
image preparer rewrites filenames. No upload, transcode, resample or provider
|
||||
operation belongs here. The caller still enforces the complete request-body cap.
|
||||
"""
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from ...serverless_bridge import codec, contract, video_codec
|
||||
from . import serverless_media
|
||||
|
||||
|
||||
_LOADERS = {"LoadAudio": ("audio", "AUDIO"), "LoadVideo": ("file", "VIDEO")}
|
||||
Error = serverless_media.ServerlessMediaError
|
||||
|
||||
|
||||
def _loader_references(workflow):
|
||||
if not isinstance(workflow, dict):
|
||||
raise Error("Serverless workflow must be an API graph")
|
||||
loaders = {}
|
||||
for node_id, node in workflow.items():
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
class_type = node.get("class_type")
|
||||
if not isinstance(class_type, str) or class_type not in _LOADERS:
|
||||
continue
|
||||
contract.identifier(node_id)
|
||||
key, kind = _LOADERS[class_type]
|
||||
inputs = node.get("inputs")
|
||||
if not isinstance(inputs, dict) or set(inputs) != {key}:
|
||||
raise Error(f"{class_type} inputs must contain only its native filename; extra controls are unsupported")
|
||||
if not isinstance(inputs[key], str):
|
||||
raise Error(f"{class_type} requires a literal filename; linked inputs are unsupported")
|
||||
normalized, _suffix = serverless_media._normalize_media_reference(inputs[key])
|
||||
loaders[node_id] = (kind, normalized)
|
||||
|
||||
references = []
|
||||
used = set()
|
||||
for node_id, node in workflow.items():
|
||||
if not isinstance(node, dict) or not isinstance(node.get("inputs"), dict):
|
||||
continue
|
||||
for key, value in node["inputs"].items():
|
||||
if not isinstance(value, (list, tuple)) or not value or not isinstance(value[0], str):
|
||||
continue
|
||||
if value[0] not in loaders:
|
||||
continue
|
||||
if type(value) is not list or len(value) != 2 or type(value[1]) is not int or value[1] != 0:
|
||||
raise Error("Native media loader links require output slot 0")
|
||||
contract.identifier(node_id)
|
||||
contract.identifier(key)
|
||||
references.append((node_id, key, value[0]))
|
||||
used.add(value[0])
|
||||
if used != set(loaders):
|
||||
raise Error("Native media loader has no consumer")
|
||||
return loaders, references
|
||||
|
||||
|
||||
def _source_path(value, folder_paths):
|
||||
"""Resolve native annotations; authorization/opening is owned by _read_path."""
|
||||
annotation = serverless_media._MEDIA_ANNOTATION.search(value)
|
||||
unannotated = value[:annotation.start()] if annotation else value
|
||||
if Path(unannotated).is_absolute():
|
||||
return unannotated
|
||||
try:
|
||||
return folder_paths.get_annotated_filepath(value)
|
||||
except Exception:
|
||||
raise Error("Native media input path is missing or invalid") from None
|
||||
|
||||
|
||||
def _decode_audio(data, limit):
|
||||
"""Match core nodes_audio.load, bounding each frame *before* accumulation.
|
||||
|
||||
Float32 output bytes are charged even for narrow integer PCM. Native float64
|
||||
remains float64, as in f32_pcm (there is no silent dtype conversion). Raw
|
||||
retained arrays and each live AV plane allocation have a separate <= limit
|
||||
guard, so wider/padded raw frames cannot bypass the output budget. PyAV owns
|
||||
the one frame being decoded; it is checked before to_ndarray or torch.cat.
|
||||
Concatenation/conversion use bounded temporary copies, not an RSS guarantee.
|
||||
"""
|
||||
if limit <= 0:
|
||||
raise Error("AUDIO decoded input exceeds byte limit")
|
||||
try:
|
||||
import av
|
||||
import torch
|
||||
from comfy_extras.nodes_audio import f32_pcm
|
||||
except ImportError:
|
||||
raise Error("Native AUDIO decoding is unavailable") from None
|
||||
|
||||
try:
|
||||
with av.open(io.BytesIO(data)) as container:
|
||||
if not container.streams.audio:
|
||||
raise Error("Native AUDIO input has no audio stream")
|
||||
stream = container.streams.audio[0]
|
||||
rate = stream.codec_context.sample_rate
|
||||
channels = stream.channels
|
||||
if type(rate) is not int or rate <= 0 or type(channels) is not int or channels <= 0:
|
||||
raise Error("Native AUDIO stream shape or sample rate is invalid")
|
||||
frames = []
|
||||
output_bytes = raw_bytes = 0
|
||||
dtype = None
|
||||
for frame in container.decode(streams=stream.index):
|
||||
samples = frame.samples
|
||||
if (type(samples) is not int or samples <= 0 or frame.sample_rate != rate
|
||||
or len(frame.layout.channels) != channels):
|
||||
raise Error("Native AUDIO frame shape or sample rate changed")
|
||||
if frame.format.name not in {"s16", "s16p", "s32", "s32p", "flt", "fltp", "dbl", "dblp"}:
|
||||
raise Error("Unsupported native AUDIO PCM format")
|
||||
itemsize = frame.format.bytes
|
||||
count = samples * channels
|
||||
# Native f32_pcm promotes integer PCM to float32, not float64.
|
||||
decoded_bytes = count * max(4, itemsize)
|
||||
frame_bytes = count * itemsize
|
||||
plane_bytes = sum(plane.buffer_size for plane in frame.planes)
|
||||
if (decoded_bytes > limit - output_bytes or frame_bytes > limit - raw_bytes
|
||||
or plane_bytes > limit - raw_bytes):
|
||||
raise Error("AUDIO decoded input or raw frame exceeds byte limit")
|
||||
output_bytes += decoded_bytes
|
||||
raw_bytes += frame_bytes
|
||||
array = frame.to_ndarray()
|
||||
if array.ndim != 2 or array.size != count or array.nbytes != frame_bytes:
|
||||
raise Error("Native AUDIO frame shape is invalid")
|
||||
buf = torch.from_numpy(array)
|
||||
if buf.dtype not in (torch.int16, torch.int32, torch.float32, torch.float64):
|
||||
raise Error("Unsupported native AUDIO PCM dtype")
|
||||
if dtype is not None and buf.dtype != dtype:
|
||||
raise Error("Native AUDIO frame dtype changed")
|
||||
dtype = buf.dtype
|
||||
if buf.shape[0] != channels:
|
||||
buf = buf.view(-1, channels).t()
|
||||
if tuple(buf.shape) != (channels, samples):
|
||||
raise Error("Native AUDIO frame shape is invalid")
|
||||
frames.append(buf)
|
||||
if not frames:
|
||||
raise Error("Native AUDIO input has no decoded frames")
|
||||
waveform = f32_pcm(torch.cat(frames, dim=1)).unsqueeze(0)
|
||||
return {"waveform": waveform, "sample_rate": rate}
|
||||
except Error:
|
||||
raise
|
||||
except Exception:
|
||||
raise Error("Native AUDIO input could not be decoded safely") from None
|
||||
|
||||
|
||||
def _artifact_id(kind, descriptor, payload):
|
||||
digest = hashlib.sha256()
|
||||
digest.update(kind.encode("ascii") + b"\0")
|
||||
# Include interpretation (not just waveform bytes), notably sample_rate.
|
||||
digest.update(json.dumps(descriptor, sort_keys=True, separators=(",", ":"), allow_nan=False).encode())
|
||||
digest.update(b"\0")
|
||||
digest.update(payload)
|
||||
return "input-" + kind.lower() + "-" + digest.hexdigest()
|
||||
|
||||
|
||||
async def prepare_media_inputs(workflow: dict) -> tuple[list, list, list]:
|
||||
"""Return legacy images, inline typed artifacts, and every consumer reference.
|
||||
|
||||
Paths are read once into bounded bytes through the video adapter's no-follow
|
||||
directory-fd reader. Decoders receive BytesIO, never a second path open.
|
||||
Source files and unique decoded artifacts have independent aggregate caps.
|
||||
Repeated content of the same kind is decoded once and charged once.
|
||||
"""
|
||||
loaders, references = _loader_references(workflow)
|
||||
image_view = {node_id: node for node_id, node in workflow.items() if node_id not in loaders}
|
||||
# Retain the strict scan for every other loader before expensive decoding.
|
||||
serverless_media._iter_image_references(image_view)
|
||||
if not loaders:
|
||||
return await serverless_media.prepare_images(image_view), [], []
|
||||
try:
|
||||
import folder_paths
|
||||
except ImportError:
|
||||
raise Error("ComfyUI media directories are unavailable") from None
|
||||
roots = list(serverless_media._folder_roots(folder_paths).values())
|
||||
artifacts = {}
|
||||
source_artifacts = {}
|
||||
source_digests = set()
|
||||
filenames = {}
|
||||
loader_artifacts = {}
|
||||
file_total = decoded_total = 0
|
||||
for node_id, (kind, filename) in loaders.items():
|
||||
# Avoid retaining file buffers merely to deduplicate repeat references.
|
||||
if (kind, filename) in filenames:
|
||||
loader_artifacts[node_id] = filenames[kind, filename]
|
||||
continue
|
||||
path = _source_path(filename, folder_paths)
|
||||
data = video_codec._read_path(path, contract.MAX_ARTIFACT_BYTES, roots)
|
||||
source_digest = hashlib.sha256(data).hexdigest()
|
||||
if source_digest not in source_digests:
|
||||
if len(data) > contract.MAX_TOTAL_BYTES - file_total:
|
||||
raise Error("Native media source files exceed total byte limit")
|
||||
file_total += len(data)
|
||||
source_digests.add(source_digest)
|
||||
source_key = (kind, source_digest)
|
||||
if source_key not in source_artifacts:
|
||||
if len(artifacts) >= contract.MAX_ARTIFACTS:
|
||||
raise Error("Native media input artifact count exceeds limit")
|
||||
limit = min(contract.MAX_ARTIFACT_BYTES, contract.MAX_TOTAL_BYTES - decoded_total)
|
||||
if kind == "AUDIO":
|
||||
value = _decode_audio(data, limit)
|
||||
descriptor, payload = codec.encode_payload(kind, value, max_bytes=limit)
|
||||
del value
|
||||
else:
|
||||
from comfy_api.latest import InputImpl
|
||||
value = InputImpl.VideoFromFile(io.BytesIO(data))
|
||||
descriptor, payload = video_codec.encode_video(value, max_bytes=limit, allowed_roots=roots)
|
||||
del value
|
||||
aid = _artifact_id(kind, descriptor, payload)
|
||||
if aid not in artifacts:
|
||||
if len(payload) > limit:
|
||||
raise Error("Native media decoded inputs exceed total byte limit")
|
||||
artifacts[aid] = contract.artifact_from_payload(
|
||||
aid, kind, descriptor, payload, max_bytes=contract.MAX_ARTIFACT_BYTES)
|
||||
decoded_total += len(payload)
|
||||
source_artifacts[source_key] = aid
|
||||
del payload
|
||||
aid = source_artifacts[source_key]
|
||||
loader_artifacts[node_id] = filenames[kind, filename] = aid
|
||||
del data
|
||||
refs = [{"node_id": node_id, "input_name": key, "artifact_id": loader_artifacts[source]}
|
||||
for node_id, key, source in references]
|
||||
# Shared node dictionaries deliberately propagate only legacy image rewrites.
|
||||
images = await serverless_media.prepare_images(image_view)
|
||||
return images, list(artifacts.values()), refs
|
||||
@@ -0,0 +1,79 @@
|
||||
"""Explicit media-v1 preflight and atomic, lossless result materialization."""
|
||||
from urllib.parse import urljoin
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ...serverless_bridge import contract
|
||||
from ...serverless_bridge.contract import MAX_ARTIFACT_BYTES, PROFILE, validate_bindings
|
||||
from .serverless_media import (
|
||||
MAX_REDIRECTS, _CHUNK_SIZE, _DOWNLOAD_TIMEOUT_SECONDS,
|
||||
_PublicResolver, _validate_artifact_url,
|
||||
)
|
||||
|
||||
|
||||
def validate_prepared_job(prepared):
|
||||
"""Fail before journalling/submission; legacy jobs keep their original path."""
|
||||
profile = prepared["worker"].get("media_profile", "legacy-image-v1")
|
||||
media = prepared.get("media")
|
||||
if profile == "legacy-image-v1" and media is None:
|
||||
return False
|
||||
if profile != PROFILE or media is None:
|
||||
raise ValueError("Serverless media profile/request mismatch")
|
||||
from .runpod_client import serialize_request
|
||||
# The shared endpoint validator also checks exporter controls and input refs.
|
||||
serialize_request(prepared["workflow"], prepared["images"],
|
||||
prepared["worker"]["timeout_seconds"], media=media)
|
||||
expected = validate_bindings(prepared["manifest"])
|
||||
if expected != validate_bindings(media["outputs"]):
|
||||
raise ValueError("Serverless media manifest identity mismatch")
|
||||
if any(binding["participant_id"] != prepared["worker"]["id"] for binding in expected.values()):
|
||||
raise ValueError("Serverless media participant identity mismatch")
|
||||
return True
|
||||
|
||||
|
||||
async def download_public_artifact(url, *, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
"""Contract-sized public HTTPS download, never using a provider session/key."""
|
||||
if type(max_bytes) is not int or not 0 <= max_bytes <= MAX_ARTIFACT_BYTES:
|
||||
raise ValueError("Media artifact exceeds the permitted size")
|
||||
current_url = _validate_artifact_url(url)
|
||||
connector = aiohttp.TCPConnector(resolver=_PublicResolver(), use_dns_cache=False, limit=4)
|
||||
timeout = aiohttp.ClientTimeout(total=_DOWNLOAD_TIMEOUT_SECONDS, connect=10, sock_read=15)
|
||||
try:
|
||||
async with aiohttp.ClientSession(connector=connector, timeout=timeout, trust_env=False,
|
||||
cookie_jar=aiohttp.DummyCookieJar()) as session:
|
||||
for redirect_count in range(MAX_REDIRECTS + 1):
|
||||
async with session.get(current_url, allow_redirects=False) as response:
|
||||
if response.status in {301, 302, 303, 307, 308}:
|
||||
location = response.headers.get("Location")
|
||||
if not location or redirect_count >= MAX_REDIRECTS:
|
||||
raise ValueError("Media artifact redirect is invalid")
|
||||
current_url = _validate_artifact_url(urljoin(current_url, location))
|
||||
continue
|
||||
if response.status != 200:
|
||||
raise ValueError("Media artifact download failed")
|
||||
payload = bytearray()
|
||||
async for chunk in response.content.iter_chunked(_CHUNK_SIZE):
|
||||
if not isinstance(chunk, bytes) or len(payload) + len(chunk) > max_bytes:
|
||||
raise ValueError("Media artifact exceeds the permitted size")
|
||||
payload.extend(chunk)
|
||||
return bytes(payload)
|
||||
except Exception:
|
||||
# Signed URLs and transport exception details must not reach the ledger.
|
||||
raise ValueError("Media artifact download failed or exceeded its byte limit") from None
|
||||
raise ValueError("Media artifact download failed")
|
||||
|
||||
|
||||
async def collect_typed_results(output, media, manifest, *, download=None):
|
||||
"""Decode every binding before returning any collector's native value lists."""
|
||||
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
|
||||
raise ValueError("Invalid serverless media output")
|
||||
decoded = await contract.decode_response(output.get("media"), request_id=media["request_id"],
|
||||
bindings=manifest,
|
||||
download=download if download is not None else download_public_artifact)
|
||||
grouped = {}
|
||||
for binding in manifest:
|
||||
values = decoded[binding["binding_id"]]
|
||||
if not isinstance(values, list) or not values:
|
||||
raise ValueError("Missing serverless media values")
|
||||
grouped.setdefault(binding["job_id"], {}).setdefault(binding["socket"], []).extend(values)
|
||||
return grouped
|
||||
+83
-12
@@ -1,4 +1,5 @@
|
||||
import asyncio
|
||||
import copy
|
||||
import time
|
||||
import uuid
|
||||
|
||||
@@ -86,6 +87,10 @@ def _resolve_enabled_workers(config, requested_ids=None):
|
||||
if worker.get("type") == "serverless":
|
||||
from ..utils.serverless_config import validate_serverless_worker
|
||||
pool = validate_serverless_worker(worker)
|
||||
if pool.get("media_profile") == "comfy-media-v1":
|
||||
from ..serverless_bridge.contract import identifier
|
||||
# Do not repair malformed media identities by stripping them.
|
||||
identifier(pool["id"])
|
||||
for slot in range(pool["jobs_per_run"]):
|
||||
workers.append(dict(pool, pool_id=worker_id, id=f"serverless:{worker_id}:{slot + 1}",
|
||||
name=f"{pool.get('name', worker_id)} / job {slot + 1}"))
|
||||
@@ -155,6 +160,32 @@ def _prompt_requests_load_balance(prompt_index):
|
||||
return False
|
||||
|
||||
|
||||
def _validate_media_workers(workers, prompt_index=None):
|
||||
"""Reject unknown serverless media profiles before probing or dispatching.
|
||||
|
||||
A Distributed Collector accepts direct workers, legacy-image-v1 pools and
|
||||
comfy-media-v1 pools in the same run; each pool's profile selects its
|
||||
endpoint contract. Unknown profiles fail rather than guessing a contract.
|
||||
legacy-image-v1 endpoints return 8-bit files, so a Collector set to exact
|
||||
cannot use them.
|
||||
"""
|
||||
exact = prompt_index is not None and any(
|
||||
str(prompt_index.inputs_by_node.get(node_id, {}).get("transfer_format", "")).strip().lower() == "exact"
|
||||
for node_id in prompt_index.nodes_for_class("DistributedCollector")
|
||||
)
|
||||
for worker in workers:
|
||||
if worker.get("type") != "serverless":
|
||||
continue
|
||||
profile = worker.get("media_profile", "legacy-image-v1")
|
||||
if profile not in ("legacy-image-v1", "comfy-media-v1"):
|
||||
raise ValueError("Unsupported serverless media_profile")
|
||||
if exact and profile == "legacy-image-v1":
|
||||
raise ValueError(
|
||||
"transfer_format exact requires comfy-media-v1 serverless pools; "
|
||||
"legacy-image-v1 endpoints return 8-bit images"
|
||||
)
|
||||
|
||||
|
||||
async def _prepare_worker_payload(
|
||||
worker,
|
||||
prompt_index,
|
||||
@@ -170,11 +201,33 @@ async def _prepare_worker_payload(
|
||||
):
|
||||
"""Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
|
||||
async with worker_prep_semaphore:
|
||||
_validate_media_workers([worker], prompt_index)
|
||||
worker_prompt = prompt_index.copy_prompt()
|
||||
if worker.get("type") == "serverless":
|
||||
from .orchestration.runpod_client import serialize_request
|
||||
if worker.get("media_profile", "legacy-image-v1") == "comfy-media-v1":
|
||||
from ..serverless_bridge.contract import PROFILE, VERSION
|
||||
from .orchestration.serverless_media_prompt import compile_media_prompt
|
||||
from .orchestration.serverless_typed_inputs import prepare_media_inputs
|
||||
# The typed compiler validates the complete private graph and
|
||||
# prunes it itself, retaining every connected Collector branch.
|
||||
worker_prompt = apply_participant_overrides(
|
||||
worker_prompt, worker["id"], enabled_ids, job_id_map,
|
||||
"", delegate_master, prompt_index,
|
||||
)
|
||||
request_id = str(uuid.uuid4())
|
||||
workflow, manifest = compile_media_prompt(worker_prompt, request_id)
|
||||
images, artifacts, refs = await prepare_media_inputs(workflow)
|
||||
media = {
|
||||
"version": VERSION, "profile": PROFILE, "request_id": request_id,
|
||||
"artifacts": artifacts, "inputs": refs, "outputs": copy.deepcopy(manifest),
|
||||
}
|
||||
serialize_request(workflow, images, worker["timeout_seconds"], media=media)
|
||||
return worker, {"worker": worker, "workflow": workflow, "manifest": manifest,
|
||||
"images": images, "media": media}
|
||||
|
||||
from .orchestration.serverless_prompt import compile_serverless_prompt
|
||||
from .orchestration.serverless_media import prepare_images
|
||||
from .orchestration.runpod_client import serialize_request
|
||||
worker_prompt = apply_participant_overrides(
|
||||
prune_prompt_for_worker(worker_prompt), worker["id"], enabled_ids,
|
||||
job_id_map, "", delegate_master, prompt_index,
|
||||
@@ -270,22 +323,22 @@ async def orchestrate_distributed_execution(
|
||||
for node in prompt_obj.values())
|
||||
target_kwargs = {}
|
||||
if image_list_mode:
|
||||
# Check the source graph before execution-root pruning can remove an
|
||||
# incompatible mode (or mask it with an unrelated boundary error).
|
||||
if any(isinstance(node, dict) and node.get("class_type") in (
|
||||
"DistributedCollector", "UltimateSDUpscaleDistributed") for node in prompt_obj.values()):
|
||||
raise ValueError(
|
||||
"Distributed Image List v1 cannot be mixed with Distributed Collector or "
|
||||
"Ultimate SD Upscale nodes in one prompt"
|
||||
)
|
||||
prompt_obj, outer_targets = prepare_image_list_execution(prompt_obj)
|
||||
target_kwargs["partial_execution_targets"] = outer_targets
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
if image_list_mode:
|
||||
# Validate before probing or queueing anything so malformed boundaries and
|
||||
# old worker releases fail before any upstream image is materialized.
|
||||
|
||||
validate_image_list_regions(prompt_obj, explicit_targets=True)
|
||||
other_distributed = (
|
||||
prompt_index.nodes_for_class("DistributedCollector")
|
||||
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
|
||||
)
|
||||
if other_distributed:
|
||||
raise ValueError(
|
||||
"Distributed Image List v1 cannot be mixed with Distributed Collector or "
|
||||
"Ultimate SD Upscale nodes in one prompt"
|
||||
)
|
||||
configured_serverless = [
|
||||
str(worker.get("id")) for worker in config.get("workers", [])
|
||||
if worker.get("enabled", False)
|
||||
@@ -306,6 +359,13 @@ async def orchestrate_distributed_execution(
|
||||
)
|
||||
else:
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
_validate_media_workers(workers, prompt_index)
|
||||
if any(worker.get("media_profile") == "comfy-media-v1" for worker in workers):
|
||||
from ..serverless_bridge.contract import identifier
|
||||
# Validate the source, not PromptIndex's JSON copy (which coerces
|
||||
# integer keys to strings), before probing or dispatching workers.
|
||||
for node_id in prompt_obj:
|
||||
identifier(node_id)
|
||||
load_balance_requested = _prompt_requests_load_balance(prompt_index)
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
@@ -488,6 +548,15 @@ async def orchestrate_distributed_execution(
|
||||
manager = None
|
||||
if serverless_jobs:
|
||||
from .orchestration.serverless_runtime import get_manager
|
||||
# Repeat the manager's whole-batch identity/profile preflight *before*
|
||||
# any direct dispatch, master queueing or collector queue allocation.
|
||||
# Pure legacy jobs are a no-op in that validator; keep their path free
|
||||
# of typed imports (and compatible with existing legacy media adapters).
|
||||
for prepared in serverless_jobs:
|
||||
if (prepared["worker"].get("media_profile", "legacy-image-v1") != "legacy-image-v1"
|
||||
or prepared.get("media") is not None):
|
||||
from .orchestration.serverless_typed_results import validate_prepared_job
|
||||
validate_prepared_job(prepared)
|
||||
manager = get_manager()
|
||||
manager.ensure_capacity(len(serverless_jobs))
|
||||
|
||||
@@ -534,8 +603,10 @@ async def orchestrate_distributed_execution(
|
||||
await manager.cancel(execution_id=execution_trace_id)
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
for prepared in serverless_jobs:
|
||||
for binding in prepared["manifest"]:
|
||||
q = prompt_server.distributed_pending_jobs.get(binding["job_id"])
|
||||
# Typed manifests have several sockets per collector, but
|
||||
# each participant gets only one terminal collector event.
|
||||
for job_id in {binding["job_id"] for binding in prepared["manifest"]}:
|
||||
q = prompt_server.distributed_pending_jobs.get(job_id)
|
||||
if q is not None:
|
||||
q.put_nowait({"worker_id": prepared["worker"]["id"], "is_last": True,
|
||||
"serverless_error": "Serverless job setup failed; no automatic retry was attempted"})
|
||||
|
||||
@@ -141,13 +141,16 @@ Also, for collector-based flows:
|
||||
|
||||
## Endpoint: `POST /distributed/job_complete`
|
||||
|
||||
Submit one completed worker image back to the master collector queue.
|
||||
Submit direct-worker image/audio results to the master Collector queue. The route
|
||||
accepts canonical JSON completions or multipart image chunks, including opt-in
|
||||
exact floating-point IMAGE transport. Serverless typed results use their separate
|
||||
bridge/provider contract.
|
||||
|
||||
### URL
|
||||
|
||||
- `http://<master-host>:<master-port>/distributed/job_complete`
|
||||
|
||||
### Request Body
|
||||
### JSON request body
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -164,13 +167,35 @@ Submit one completed worker image back to the master collector queue.
|
||||
- `job_id` (string, required)
|
||||
- `worker_id` (string, required)
|
||||
- `batch_idx` (integer >= 0, required)
|
||||
- `image` (string, required)
|
||||
- `image` (string, optional if `audio` is present)
|
||||
- PNG payload as either:
|
||||
- data URL: `data:image/png;base64,...`
|
||||
- raw base64 PNG bytes
|
||||
- `is_last` (boolean, required)
|
||||
|
||||
Legacy multipart/tensor payload formats are no longer accepted on this endpoint.
|
||||
At least one of `image` or `audio` must be present. Optional `audio` is an object
|
||||
with `sample_rate` (positive integer), `shape` (`[batch, channels, samples]`),
|
||||
`dtype: "float32"`, and `data` (base64 waveform bytes). Audio-only completions
|
||||
still include `batch_idx`; the Collector uses `0` and `is_last: true`.
|
||||
|
||||
### Multipart image chunks
|
||||
|
||||
Send `multipart/form-data` with a JSON `envelope` field and repeated binary
|
||||
`media` file fields. The envelope contains `job_id`, `worker_id`, `format`,
|
||||
`start_index` (nonnegative integer), `count`, `width`, `height`, `is_last`, and
|
||||
optional canonical `audio`.
|
||||
|
||||
- `png`, `jpeg`, `webp`: one `media` part per frame.
|
||||
- `ffv1`, `h264`, `h265` (experimental), `av1`: one video part per chunk.
|
||||
- `exact`: one raw tensor part per frame and an envelope `descriptors` list with
|
||||
one `tensor-le-v1` IMAGE descriptor per part. Each descriptor shape must be
|
||||
`[1, height, width, channels]`; decoding preserves supported float dtype and
|
||||
values instead of quantizing to 8-bit.
|
||||
|
||||
`legacy_png` uses JSON, not multipart. Exact transport requires support on both
|
||||
sides and never retries as an 8-bit format when an older master rejects it.
|
||||
See [result transfer](compressed-result-transfer.md) for format controls,
|
||||
non-exact compatibility fallbacks and request limits.
|
||||
|
||||
### CORS note
|
||||
|
||||
|
||||
@@ -1,38 +1,49 @@
|
||||
# Compressed Result Transfer
|
||||
# Compressed and Exact Result Transfer
|
||||
|
||||
**Distributed Collector** has two optional widgets that control how workers send their images back to the master. This matters most when workers run in the cloud and you pay for data leaving it (egress).
|
||||
**Distributed Collector** (`DistributedCollector`) collects images and/or audio. Its `transfer_format` and `transfer_quality` widgets control **remote image returns only**; they do not re-encode master-local images or change audio transport. The saved defaults remain `legacy_png` and quality `90`.
|
||||
|
||||
| `transfer_format` | Type | Notes |
|
||||
| `transfer_format` | Fidelity | Direct-worker transport |
|
||||
|---|---|---|
|
||||
| `legacy_png` (default) | lossless | Original behaviour: one JSON request per frame, uncompressed PNG in base64. Works with any master version. |
|
||||
| `png` | lossless | Compressed PNG, binary, several frames per request. Same tensors as `legacy_png`. |
|
||||
| `ffv1` | lossless | Whole batch as one FFV1 video. Same tensors as `legacy_png`; best lossless choice for frame sequences. |
|
||||
| `jpeg`, `webp` | lossy | Per-frame stills. Good for image batches. |
|
||||
| `h264`, `h265`, `av1` | lossy | Whole batch as one video. By far the smallest for video frame batches. |
|
||||
| `legacy_png` (default) | 8-bit pixels | Original behaviour: one JSON request per frame, uncompressed PNG in base64. Compatible with legacy masters. |
|
||||
| `png` | Lossless 8-bit pixels | Compressed PNG, binary, several frames per request. |
|
||||
| `ffv1` | Lossless 8-bit pixels | FFV1 video chunks; useful for frame sequences. |
|
||||
| `jpeg`, `webp` | Lossy 8-bit images | Per-frame stills. |
|
||||
| `h264`, `h265` (**experimental**), `av1` | Lossy 8-bit images | Video chunks; useful for similar consecutive frames. |
|
||||
| `exact` (opt-in) | Original floating-point tensors | Binary multipart chunks with one raw tensor payload and `tensor-le-v1` descriptor per frame. No 8-bit quantization. |
|
||||
|
||||
`transfer_quality` (1–100, default 90) applies to the lossy formats. JPEG/WebP use it directly; video codecs map it to CRF (90 → CRF 13 for h264/h265, CRF 16 for av1; 100 → CRF 0).
|
||||
“Lossless” PNG/FFV1 means lossless **after 8-bit quantization**, not preservation of the source float tensor. Direct non-exact returns decode to RGB float32. `exact` preserves supported floating-point dtype, shape, channels and values: float16, bfloat16, float32 or float64, finite dense BHWC tensors with 1, 3 or 4 channels. It generally costs substantially more bandwidth.
|
||||
|
||||
Measured on a 121-frame 1280×720 batch (one worker's result):
|
||||
`transfer_quality` (1–100, default 90) applies to lossy formats, not `exact`, PNG or FFV1. JPEG/WebP use it directly; video codecs map it to CRF (90 → CRF 13 for h264/h265, CRF 16 for av1; 100 → CRF 0). A quality of 100 is not a substitute for exact float transport.
|
||||
|
||||
## Serverless profiles
|
||||
|
||||
- **`comfy-media-v1`:** the same Collector supplies the endpoint's image-return policy. The saved `legacy_png` default is mapped to compressed **`png`** with the same 8-bit pixel precision; normal serverless returns are not exact by default. All other formats retain their selection, including opt-in `exact` and experimental `h265`. Missing codecs or rejected formats fail instead of selecting another codec. Typed image returns preserve channel count and carry alpha separately where needed. See [serverless images and audio](serverless-media.md).
|
||||
- **`legacy-image-v1`:** retains its image-file/`SaveImage` contract, not the typed return-format controls. Selecting a legacy pool when any Collector requests `exact` is rejected before dispatch; there is no 8-bit fallback. Audio requires direct workers or `comfy-media-v1`.
|
||||
|
||||
## Direct-worker compatibility and limits
|
||||
|
||||
- Choose the format on the master; the value is sent to workers with the workflow.
|
||||
- **`exact` requires support on both master and worker.** If an older master rejects an exact chunk, the worker raises an error and never resends it as `legacy_png` or another 8-bit format. An exact Collector also rejects direct-worker image returns labelled with a non-exact format. Update both sides or explicitly select a non-exact format.
|
||||
- For **non-exact compressed** formats only, the specific first-request older-master rejection (`400` with `Invalid JSON payload`) causes a retry using `legacy_png`. Other rejections, including HTTP 413, remain errors.
|
||||
- For direct video returns only, an unavailable encoder falls back to `h264` for lossy video if available, otherwise `png`; unavailable `ffv1` falls back to `png`. The worker logs this. This compatibility fallback does **not** apply to `exact` or typed serverless returns.
|
||||
- Multipart requests use a 32 MiB chunk target (`COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES`) and a 96 MiB request budget (`COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES`), reserving space for the envelope and accounting separately for audio. A single large frame may exceed the chunk target but must fit the request budget. Keep the request budget below the master's `--max-upload-size` (default 100 MB).
|
||||
- Non-exact compressed chunks are all encoded before sending. Exact transfer instead plans frame sizes from shape/dtype before sending and encodes one chunk at a time. An oversized individual frame is rejected before any chunks are sent; choose another format or raise the worker and master limits together.
|
||||
- Audio is attached to the final image chunk when it fits, otherwise sent separately. Audio-only completions use the canonical JSON envelope. Direct audio uses float32 waveform bytes regardless of the image format; `exact` is an IMAGE option, not a new audio codec.
|
||||
- Each worker logs the format and bytes sent. These request limits are for direct multipart transfers; serverless request/artifact limits are documented separately.
|
||||
|
||||
## Previously documented size example
|
||||
|
||||
The following measurements are retained from a 121-frame 1280×720 slow-pan batch (one worker's result); they are not new validation of the current implementation or exact transport.
|
||||
|
||||
| Format | Sent | vs legacy | PSNR |
|
||||
|---|---|---|---|
|
||||
| `legacy_png` | 446 MB | 1× | lossless |
|
||||
| `png` | 163 MB | 2.7× | lossless |
|
||||
| `ffv1` | 128 MB | 3.5× | lossless |
|
||||
| `legacy_png` | 446 MB | 1× | lossless at 8-bit precision |
|
||||
| `png` | 163 MB | 2.7× | lossless at 8-bit precision |
|
||||
| `ffv1` | 128 MB | 3.5× | lossless at 8-bit precision |
|
||||
| `jpeg` q90 | 27 MB | 17× | 42.7 dB |
|
||||
| `webp` q90 | 13 MB | 33× | 40.6 dB |
|
||||
| `h264` q90 | 2.4 MB | 186× | 40.2 dB |
|
||||
| `h265` q90 | 1.0 MB | 434× | 40.1 dB |
|
||||
| `h265` q90 (experimental) | 1.0 MB | 434× | 40.1 dB |
|
||||
| `av1` q90 | 0.4 MB | 1031× | 40.5 dB |
|
||||
|
||||
The test clip is a slow pan, which video codecs compress especially well; expect smaller (still large) gains on content with more motion.
|
||||
|
||||
Notes:
|
||||
- Choose the format on the master; the value is sent to workers with the workflow.
|
||||
- The master must run this version to accept compressed results. A worker sending to an older master gets a 400 on its first request and automatically resends with `legacy_png`.
|
||||
- If a worker's PyAV build lacks the requested video encoder, it falls back to `h264` (lossy formats) or `png` (`ffv1`), and logs it.
|
||||
- Lossy formats only affect frames produced by workers; the master's own frames are never re-encoded.
|
||||
- Request size is bounded on the final encoded bytes. Frames are batched into requests of up to 32 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES`); a single large frame may go alone in a request of up to 96 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES`). Keep the latter below the master's `--max-upload-size` (default 100 MB) if you change either.
|
||||
- All chunks are encoded before anything is sent. If a single frame cannot fit within the request limit, the worker fails with an error and sends nothing; use a lossy format or raise both limits.
|
||||
- A rejection other than the older-master case above (for example the master's size limit, HTTP 413) is reported as an error rather than retried as `legacy_png`.
|
||||
- Each worker logs what it sent, e.g. `Sent 121 image(s) as h264 in 1 request(s): 2.40 MB (raw 334.49 MB, 139.4x smaller)`.
|
||||
Video codecs compress this slow pan especially well; gains depend on content and motion.
|
||||
|
||||
@@ -0,0 +1,187 @@
|
||||
# Serverless images and audio
|
||||
|
||||
**Distributed Collector** (`DistributedCollector`) collects `IMAGE` and/or `AUDIO`
|
||||
from direct workers and serverless pools. There is no separate
|
||||
`DistributedMediaCollector`; replace that removed node with `DistributedCollector`
|
||||
for image/audio workflows. The Collector has no MASK, STRING, or native VIDEO
|
||||
output sockets.
|
||||
|
||||
`comfy-media-v1` is an **explicit endpoint profile**, not an automatically detected
|
||||
capability. It returns typed image/audio items to the Collector. The standalone
|
||||
bridge's generic contract also supports `MASK`, `STRING`, and native ComfyUI
|
||||
`VIDEO`, but that does not add those outputs to the Collector. Omitting
|
||||
`media_profile` still selects the image-only `legacy-image-v1` contract.
|
||||
See [legacy serverless setup](serverless-workers.md) for credentials and recovery.
|
||||
|
||||
## Use the Collector with a typed endpoint
|
||||
|
||||
1. Prepare an endpoint with the [standalone media bridge](../serverless_bridge/README.md)
|
||||
and the workflow's models/custom nodes. The endpoint needs only the bridge,
|
||||
**not the full ComfyUI-Distributed extension**.
|
||||
2. In the serverless worker settings, set **Endpoint media profile** to
|
||||
**Typed media v1 (requires endpoint bridge)** (`media_profile: "comfy-media-v1"`).
|
||||
Start with **Images per run = 1**. This existing setting counts remote jobs,
|
||||
not the number of typed items each job returns.
|
||||
3. Use **Distributed Collector** (`DistributedCollector`), connect `images`,
|
||||
`audio`, or both, and connect the corresponding outputs downstream. Image-only
|
||||
and audio-only workflows use the same node.
|
||||
4. Queue through Distributed. **Delegate Only** excludes local upstream
|
||||
generation; normal collection includes the master's local values.
|
||||
|
||||
**Test Saved Endpoint** checks provider authentication/health, not bridge
|
||||
installation, model availability, or typed execution. Selecting a profile does
|
||||
not install or upgrade anything on the endpoint.
|
||||
|
||||
The Collector accepts optional `images` and `audio` inputs, including ComfyUI list
|
||||
inputs, and returns two ordinary outputs—not list outputs:
|
||||
|
||||
| Socket / output index | Native type | Aggregation |
|
||||
| --- | --- | --- |
|
||||
| `images` / 0 | `IMAGE` | Concatenate BHWC tensors along the batch dimension |
|
||||
| `audio` / 1 | `AUDIO` | Concatenate waveforms along the sample dimension |
|
||||
|
||||
For each connected socket, order is **master first, then prepared participant
|
||||
order, then that participant's item order**—never completion order. Delegate-only
|
||||
mode omits the master. Image dimensions/channels and audio batch/channel dimensions
|
||||
must be compatible for concatenation. Use a common audio sample rate; the Collector
|
||||
does not resample or align streams. Audio-only collection returns `None` for images;
|
||||
when no audio is supplied it returns a silent placeholder, not an empty list.
|
||||
|
||||
The typed wire contract carries 1–64 items per connected socket. The entire remote
|
||||
response is validated and decoded before any of its results are delivered;
|
||||
missing, duplicate, mismatched, or corrupt typed results fail collection. Validated
|
||||
wire lists are then combined into the Collector's ordinary outputs.
|
||||
|
||||
Direct workers, `legacy-image-v1` pools and `comfy-media-v1` pools can share an
|
||||
image-only Collector run. Audio requires direct workers or `comfy-media-v1` pools.
|
||||
If any Collector requests `exact`, a selected `legacy-image-v1` pool is rejected
|
||||
before dispatch instead of silently returning 8-bit images.
|
||||
|
||||
## Image return format and quality
|
||||
|
||||
The Collector's saved `transfer_format` default remains **`legacy_png`**, with
|
||||
`transfer_quality` **90**. For `comfy-media-v1`, the compiler maps `legacy_png` to
|
||||
**`png`**, so normal serverless returns use compressed 8-bit PNG. `exact` is opt-in.
|
||||
These controls affect **remote IMAGE returns only**, not master-local images,
|
||||
uploaded inputs or AUDIO. The bridge's generic MASK, STRING and native VIDEO
|
||||
codecs are also unaffected.
|
||||
|
||||
- `exact`: preserve supported floating-point dtype and values, with no image
|
||||
quantization. Available for direct workers and `comfy-media-v1` pools; it is
|
||||
never downgraded to an 8-bit format. An older master that rejects exact transfer
|
||||
causes an error, not a legacy retry. Update both sides before selecting it.
|
||||
- `legacy_png`, `png`, `ffv1`: lossless **8-bit image pixels**, not exact float
|
||||
tensors. For typed serverless returns, image values are clipped to [0,1] and
|
||||
quantized; `legacy_png` is sent as `png`. Direct workers retain the original
|
||||
JSON/base64 framing for `legacy_png`.
|
||||
- `jpeg`, `webp`: lossy still-image compression.
|
||||
- `h264`, `h265` (**experimental**), `av1`: lossy video encoding of each IMAGE
|
||||
tensor's frame batch. This is an image-return encoding, not a native VIDEO output.
|
||||
- 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 for typed serverless returns; 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). Direct
|
||||
workers have different limits and non-exact compatibility fallbacks; see
|
||||
[result transfer](compressed-result-transfer.md). The bridge's internal
|
||||
`MediaExportV1` default is still `exact` for generic callers that omit controls;
|
||||
the Collector compiler explicitly supplies `png` for the normal saved default.
|
||||
|
||||
## Inputs and fidelity
|
||||
|
||||
Automatic local-file transfer supports stock `LoadImage`, plus explicit stock
|
||||
`LoadAudio` (`audio` filename) and `LoadVideo` (`file` filename) adapters. Audio and
|
||||
video loaders require a literal filename, no extra controls, and output slot 0.
|
||||
Files must be under ComfyUI's input/output/temp roots. Custom file loaders are not
|
||||
supported; detected unsupported file references fail rather than silently
|
||||
becoming endpoint-local paths. Models and custom nodes are never transferred.
|
||||
|
||||
- `LoadImage` uses the existing bounded file/base64 image staging path.
|
||||
- `LoadAudio` is decoded on the master using native PCM semantics with bounded
|
||||
accumulation, then transferred as a waveform and sample rate. It does not
|
||||
resample; integer PCM conversion follows ComfyUI's `f32_pcm` behavior.
|
||||
- `LoadVideo` transfers opaque file bytes as native `VideoFromFile`, without
|
||||
decoding frames or transcoding the container.
|
||||
- Typed tensor transport is `tensor-le-v1`: contiguous little-endian float16,
|
||||
bfloat16, float32, or float64 bytes with exact dtype/shape/layout metadata.
|
||||
Only dense, finite tensors with positive dimensions are supported. IMAGE has
|
||||
1, 3, or 4 channels. Bridge export explicitly detaches/stages tensors to CPU;
|
||||
the pure codec itself rejects non-CPU storage. Remote decoding returns
|
||||
independently owned CPU values. STRING transport is exact UTF-8 (`utf8-v1`).
|
||||
|
||||
### Bridge-level native VIDEO support
|
||||
|
||||
The generic bridge can import/export native VIDEO, and the typed input adapter
|
||||
can stage `LoadVideo` for upstream processing. This is not a Collector VIDEO
|
||||
output or a generic video converter.
|
||||
|
||||
The adapter supports the pinned native classes **exactly**, not arbitrary
|
||||
subclasses or objects exposing `save_to()` / `get_components()`:
|
||||
|
||||
- `VideoFromComponents`: images, optional matching BHW or BHW1 alpha, optional
|
||||
audio, positive rational frame rate, bounded JSON metadata, positive integer
|
||||
bit depth, and `sRGB`, `HDR`, or `HDR PQ` color space are preserved.
|
||||
- `VideoFromFile`: regular local files within allowed media roots or an exact
|
||||
`BytesIO`; opaque bytes, finite signed start/duration, and optional four-integer
|
||||
crop are preserved without normalization.
|
||||
- `VideoFromList`: 1–64 flat exact `VideoFromFile` segments, plus optional
|
||||
`complete_audio`, in order.
|
||||
|
||||
Unknown classes, incompatible native state, nested/component segment lists,
|
||||
unsupported tensors, invalid metadata, and oversized payloads are rejected.
|
||||
There is no fallback transcode, frame extraction, color conversion, or container
|
||||
repair/validation. A downstream native consumer still needs to read the original
|
||||
container. Decoded file-backed VIDEO values own independent `BytesIO` storage;
|
||||
component tensors also own their storage, so staging-file cleanup cannot
|
||||
invalidate a returned VIDEO.
|
||||
|
||||
## Size and operational limits
|
||||
|
||||
- The typed contract caps each artifact at **256 MiB**, aggregate artifact bytes
|
||||
at **512 MiB**, and each artifact list at **128** entries. These are validation
|
||||
ceilings, not a promise that inline requests can carry that much data.
|
||||
- The built-in master packs inputs inline and enforces **10,000,000 bytes for the
|
||||
complete serialized Runpod `/run` request**, including workflow, policy and
|
||||
base64. Legacy image reads also retain their 10,000,000-byte aggregate guard.
|
||||
There is **no built-in master-side large-input uploader** in PR2.
|
||||
- The endpoint defaults to inline results and caps its result JSON at
|
||||
**10,000,000 bytes**. Larger results require an operator-supplied publisher
|
||||
returning expiring public HTTPS artifact URLs. See the bridge guide for
|
||||
bounded download/upload responsibilities and retention.
|
||||
- The generic five-kind typed envelope is a serverless bridge contract. Direct
|
||||
workers support Collector images/audio, including opt-in exact IMAGE transport,
|
||||
not generic MASK/STRING/VIDEO Collector outputs. The typed compiler rejects
|
||||
unsupported distributed/batch nodes and collector-to-collector dependencies.
|
||||
Downstream processing stays on the master. [Image-list scheduling](image-list-queue.md)
|
||||
is a separate execution mode.
|
||||
|
||||
## Timeout is not proof that Comfy stopped
|
||||
|
||||
Provider cancellation, a handler timeout, or cancellation of its Python task
|
||||
**does not establish that the local Comfy prompt stopped**. The bridge never
|
||||
posts `/interrupt`: the pinned core route checks a queue snapshot and then
|
||||
interrupts globally, so even a prompt ID does not make it race-free.
|
||||
|
||||
If submission/termination or callback teardown cannot be confirmed, the bridge
|
||||
raises `MediaExecutionUncertain` and retains owned input/exporter files—even
|
||||
after terminal Comfy history if a downloader/publisher ignores cancellation.
|
||||
Callback teardown has a separate 10-second grace; late callback exit does not
|
||||
clean retained files. Callbacks must be cooperative async code; synchronous
|
||||
event-loop blocking cannot be hard-timed-out.
|
||||
|
||||
The endpoint operator must **quarantine the owning worker/environment, reconcile
|
||||
the local prompt and readers or terminate that environment, and establish that
|
||||
they have stopped before cleanup or worker reuse**. Do not automatically retry
|
||||
the request or turn every handler exception into a safely reusable terminal
|
||||
provider job. Master-side provider status/reconciliation alone is not proof
|
||||
that a surviving local Comfy process or callback has stopped.
|
||||
|
||||
These docs describe the implementation, not a test result or live Runpod/provider
|
||||
or GPU end-to-end validation.
|
||||
@@ -1,9 +1,12 @@
|
||||
# Runpod serverless workers (preview)
|
||||
|
||||
A **Serverless (Runpod)** worker is a queue-based endpoint, not a permanent
|
||||
ComfyUI machine. Runpod chooses the actual GPU workers. Distributed submits one
|
||||
normal ComfyUI API workflow per requested image and collects the returned images.
|
||||
The endpoint workers do not need ComfyUI-Distributed installed.
|
||||
ComfyUI machine. Runpod chooses the actual GPU workers. This page describes the
|
||||
default **`legacy-image-v1`** profile: Distributed submits normal ComfyUI API
|
||||
workflows and collects returned image files. For image/audio collection and
|
||||
opt-in exact IMAGE returns, use [the `comfy-media-v1` profile](serverless-media.md)
|
||||
with the same **Distributed Collector** (`DistributedCollector`). The endpoint
|
||||
workers do not need the full ComfyUI-Distributed extension installed.
|
||||
|
||||
## Setup
|
||||
|
||||
@@ -14,7 +17,8 @@ The endpoint workers do not need ComfyUI-Distributed installed.
|
||||
4. Save, then choose **Test Saved Endpoint**. This checks authentication and
|
||||
endpoint health only; it does not start GPU inference or prove model/node compatibility.
|
||||
5. Enable the worker when you want queued Distributed workflows to use it.
|
||||
6. Use an image **Distributed Collector** and queue through Distributed as usual.
|
||||
6. Connect the **Distributed Collector**'s `images` input and queue through
|
||||
Distributed as usual. Leave `audio` unconnected for this legacy profile.
|
||||
Enable **Delegate Only** if the master should collect without generating locally.
|
||||
|
||||
Keys are write-only. Leave the key field blank to retain a saved key, or select
|
||||
@@ -152,9 +156,16 @@ Recovery uses the same trusted-master API as the panel; it never resubmits a job
|
||||
Worker setting changes become available once all its local tasks and unresolved records
|
||||
are finished. No credentials or ledger files need to be edited by hand.
|
||||
|
||||
## First-version limits
|
||||
## Legacy profile limits
|
||||
|
||||
- Image collectors only. Audio/video transfer and audio collectors are rejected.
|
||||
- Only the Collector's image branch is supported by `legacy-image-v1`; connected
|
||||
audio is rejected. `comfy-media-v1` adds audio through the same Collector and
|
||||
supports the configured image-return formats.
|
||||
- `exact` is opt-in and requires direct workers or `comfy-media-v1` pools. If any
|
||||
Collector requests it, a selected legacy pool is rejected before dispatch,
|
||||
not used as an 8-bit fallback. Other Collector transfer-format controls do not
|
||||
change this legacy profile's `SaveImage`/image-file contract. With
|
||||
`comfy-media-v1`, the saved `legacy_png` default is mapped to compressed `png`.
|
||||
- Coordinated `UltimateSDUpscaleDistributed` workloads and unsupported Distributed
|
||||
node types are rejected before submission.
|
||||
- Downstream post-collector processing remains local; the remote graph contains
|
||||
@@ -168,6 +179,6 @@ are finished. No credentials or ledger files need to be edited by hand.
|
||||
- No remote GPU discovery, endpoint provisioning, model synchronization, automatic
|
||||
scaling changes, or inbound master callback/tunnel is needed.
|
||||
|
||||
This feature is verified with offline provider/transport tests. A successful
|
||||
health test is not a paid end-to-end generation test; test your deployed handler
|
||||
These docs describe the implementation, not a test result. A successful health
|
||||
check is not a paid end-to-end generation test; validate your deployed handler
|
||||
and model inventory with one image per run before increasing it.
|
||||
|
||||
+188
-3
@@ -27,6 +27,39 @@ def _encode_all_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes):
|
||||
return list(transfer_codec.encode_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes))
|
||||
|
||||
|
||||
def _exact_frame_bytes(image_batch):
|
||||
"""Wire size of one exact frame, computed from shape/dtype before encoding."""
|
||||
from ..serverless_bridge import codec as bridge_codec
|
||||
if not isinstance(image_batch, torch.Tensor) or image_batch.layout != torch.strided or image_batch.is_nested:
|
||||
raise transfer_codec.TransferCodecError("exact transfer requires a dense IMAGE tensor")
|
||||
if image_batch.dtype not in bridge_codec._DTYPE_NAMES:
|
||||
raise transfer_codec.TransferCodecError(f"exact transfer does not support dtype {image_batch.dtype}")
|
||||
if image_batch.ndim != 4:
|
||||
raise transfer_codec.TransferCodecError(
|
||||
f"Expected IMAGE tensor [B,H,W,C], got shape {tuple(image_batch.shape)}")
|
||||
batch, height, width, channels = (int(value) for value in image_batch.shape)
|
||||
if min(batch, height, width) <= 0 or channels not in (1, 3, 4):
|
||||
raise transfer_codec.TransferCodecError("exact IMAGE dimensions must be positive with 1, 3 or 4 channels")
|
||||
size = height * width * channels * image_batch.element_size()
|
||||
if size > bridge_codec.MAX_PAYLOAD_BYTES:
|
||||
raise transfer_codec.FrameTooLargeError("exact frame exceeds the tensor codec byte limit")
|
||||
return size
|
||||
|
||||
|
||||
def _encode_exact_frames(frames):
|
||||
"""Encode [1,H,W,C] CPU slices losslessly (runs in an executor)."""
|
||||
from ..serverless_bridge import codec as bridge_codec
|
||||
descriptors, parts = [], []
|
||||
for frame in frames:
|
||||
try:
|
||||
descriptor, payload = bridge_codec.encode_payload("IMAGE", frame)
|
||||
except ValueError as exc:
|
||||
raise transfer_codec.TransferCodecError(f"exact transfer cannot encode this IMAGE: {exc}") from None
|
||||
descriptors.append(descriptor)
|
||||
parts.append(payload)
|
||||
return descriptors, parts
|
||||
|
||||
|
||||
class DistributedCollectorNode:
|
||||
INPUT_IS_LIST = True
|
||||
EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
|
||||
@@ -200,6 +233,10 @@ class DistributedCollectorNode:
|
||||
"""Send an image batch, optionally with audio, or an audio-only completion."""
|
||||
batch_size = 0 if image_batch is None else image_batch.shape[0]
|
||||
fmt = self._resolve_transfer_format(transfer_format) if batch_size else transfer_codec.LEGACY_FORMAT
|
||||
if fmt == transfer_codec.EXACT_FORMAT:
|
||||
# Never downgrade exact to 8-bit: an older master fails the job clearly.
|
||||
await self._send_exact_batch(image_batch, audio, multi_job_id, master_url, worker_id)
|
||||
return
|
||||
if fmt != transfer_codec.LEGACY_FORMAT:
|
||||
sent = await self._send_compressed_batch(
|
||||
image_batch, audio, multi_job_id, master_url, worker_id, fmt, transfer_quality
|
||||
@@ -329,6 +366,100 @@ class DistributedCollectorNode:
|
||||
)
|
||||
return True
|
||||
|
||||
async def _send_exact_batch(self, image_batch, audio, multi_job_id, master_url, worker_id):
|
||||
"""Send original floating-point frames as binary multipart chunks.
|
||||
|
||||
Frame sizes are known from shape and dtype, so every request is planned
|
||||
(and an oversized frame rejected) before anything is sent. Frames are
|
||||
encoded one chunk at a time to avoid holding a second copy of the batch.
|
||||
Any master rejection raises: there is no lossy or 8-bit fallback.
|
||||
"""
|
||||
encoded_audio = encode_audio_payload(audio)
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/job_complete"
|
||||
loop = asyncio.get_running_loop()
|
||||
media_limit = max(1, TRANSFER_MAX_REQUEST_BYTES - self.TRANSFER_ENVELOPE_RESERVE_BYTES)
|
||||
chunk_limit = min(TRANSFER_MAX_CHUNK_BYTES, media_limit)
|
||||
|
||||
frame_bytes = _exact_frame_bytes(image_batch)
|
||||
total = int(image_batch.shape[0])
|
||||
if frame_bytes > media_limit:
|
||||
raise transfer_codec.FrameTooLargeError(
|
||||
f"Frame 0 is {frame_bytes} bytes as exact, above the {media_limit}-byte request limit. "
|
||||
"Use png or a lossy format, or raise the master's --max-upload-size together with "
|
||||
"COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES."
|
||||
)
|
||||
# Bound descriptor/part-header overhead even for tiny frames.
|
||||
per_chunk = min(128, max(1, chunk_limit // frame_bytes))
|
||||
ranges = [(start, min(per_chunk, total - start)) for start in range(0, total, per_chunk)]
|
||||
audio_bytes = len(json.dumps(encoded_audio)) if encoded_audio is not None else 0
|
||||
last_media = ranges[-1][1] * frame_bytes
|
||||
audio_inline = encoded_audio is not None and last_media + audio_bytes <= media_limit
|
||||
audio_separate = encoded_audio is not None and not audio_inline
|
||||
height, width = int(image_batch.shape[1]), int(image_batch.shape[2])
|
||||
if audio_separate and audio_bytes > media_limit:
|
||||
raise transfer_codec.FrameTooLargeError("exact completion audio exceeds the request limit")
|
||||
cpu_batch = ensure_contiguous(image_batch.detach().cpu())
|
||||
|
||||
started = time.time()
|
||||
sent_bytes = 0
|
||||
for index, (start, count) in enumerate(ranges):
|
||||
frames = [cpu_batch[i:i + 1] for i in range(start, start + count)]
|
||||
descriptors, parts = await loop.run_in_executor(None, _encode_exact_frames, frames)
|
||||
is_final_chunk = index == len(ranges) - 1
|
||||
envelope = {
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"format": transfer_codec.EXACT_FORMAT,
|
||||
"start_index": start,
|
||||
"count": count,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"descriptors": descriptors,
|
||||
"is_last": bool(is_final_chunk and not audio_separate),
|
||||
}
|
||||
if is_final_chunk and audio_inline:
|
||||
envelope["audio"] = encoded_audio
|
||||
chunk_bytes = sum(len(p) for p in parts)
|
||||
body = self._build_multipart(envelope, parts)()
|
||||
if body.size is None or body.size > TRANSFER_MAX_REQUEST_BYTES:
|
||||
raise transfer_codec.FrameTooLargeError("exact multipart body exceeds the request limit")
|
||||
timeout_seconds = min(600, 60 + chunk_bytes // (256 * 1024))
|
||||
try:
|
||||
async with session.post(
|
||||
url,
|
||||
data=body,
|
||||
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
|
||||
) as response:
|
||||
if response.status >= 400:
|
||||
body = await response.text()
|
||||
raise RuntimeError(
|
||||
f"Master rejected exact chunk {index + 1}/{len(ranges)} ({chunk_bytes} bytes) "
|
||||
f"with HTTP {response.status}: {body[:300]}. exact needs a master with exact "
|
||||
"transfer support and requests within its --max-upload-size; it is never "
|
||||
"resent as an 8-bit format. Update the master or choose png."
|
||||
)
|
||||
except Exception as e:
|
||||
log(f"Worker - Failed to send exact chunk to master: {e}")
|
||||
debug_log(f"Worker - Full error details: URL={url}")
|
||||
raise
|
||||
sent_bytes += chunk_bytes
|
||||
|
||||
requests = len(ranges)
|
||||
if audio_separate:
|
||||
await self._post_json_completion(session, url, {
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"batch_idx": 0,
|
||||
"audio": encoded_audio,
|
||||
"is_last": True,
|
||||
})
|
||||
requests += 1
|
||||
log(
|
||||
f"Worker - Sent {total} image(s) as exact in {requests} request(s): "
|
||||
f"{sent_bytes / 1e6:.2f} MB in {time.time() - started:.1f}s"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _post_json_completion(session, url, payload):
|
||||
timeout_seconds = 60 if "image" in payload else 600
|
||||
@@ -457,6 +588,31 @@ class DistributedCollectorNode:
|
||||
worker_images[worker_id][image_index] = tensor
|
||||
return 1
|
||||
|
||||
# Collector sockets a comfy-media-v1 serverless pool may return.
|
||||
TYPED_MEDIA_SOCKETS = ("images", "audio")
|
||||
|
||||
def _store_typed_media(self, worker_images: dict, worker_audio: dict, worker_id: str, media) -> int:
|
||||
"""Store one atomic typed serverless delivery: ordered IMAGE/AUDIO item lists.
|
||||
|
||||
Each IMAGE item is a BHWC batch kept at its list position; AUDIO items are
|
||||
joined in order, exactly as list inputs are on the master.
|
||||
"""
|
||||
if (not isinstance(media, dict) or not media
|
||||
or any(socket not in self.TYPED_MEDIA_SOCKETS for socket in media)):
|
||||
raise RuntimeError(f"Invalid serverless media delivery from {worker_id}")
|
||||
for socket, items in media.items():
|
||||
if not isinstance(items, list) or not items or any(item is None for item in items):
|
||||
raise RuntimeError(f"Serverless participant {worker_id} returned no {socket} items")
|
||||
stored = 0
|
||||
for index, tensor in enumerate(media.get("images", [])):
|
||||
if not isinstance(tensor, torch.Tensor) or tensor.ndim != 4:
|
||||
raise RuntimeError(f"Serverless participant {worker_id} returned an invalid IMAGE item")
|
||||
worker_images.setdefault(worker_id, {})[index] = tensor
|
||||
stored += 1
|
||||
if "audio" in media:
|
||||
worker_audio[worker_id] = self._normalize_audio_input(media["audio"])
|
||||
return stored
|
||||
|
||||
def _reorder_and_combine_tensors(
|
||||
self,
|
||||
worker_images: dict,
|
||||
@@ -534,6 +690,10 @@ class DistributedCollectorNode:
|
||||
serverless_manager = self._get_serverless_manager()
|
||||
return serverless_manager
|
||||
|
||||
try:
|
||||
exact_requested = transfer_codec.normalize_format(transfer_format) == transfer_codec.EXACT_FORMAT
|
||||
except transfer_codec.TransferCodecError:
|
||||
exact_requested = False
|
||||
num_workers = len(expected_workers)
|
||||
if num_workers == 0:
|
||||
return (images, audio if audio is not None else self.EMPTY_AUDIO)
|
||||
@@ -609,6 +769,27 @@ class DistributedCollectorNode:
|
||||
raise RuntimeError(f"Serverless participant {participant} failed: {message}")
|
||||
|
||||
result_worker_id = str(result["worker_id"])
|
||||
if "media" in result:
|
||||
if result_worker_id not in serverless_workers:
|
||||
raise RuntimeError(f"Unexpected serverless media participant {result_worker_id}")
|
||||
if result_worker_id in workers_done:
|
||||
raise RuntimeError(f"Duplicate serverless media delivery from {result_worker_id}")
|
||||
if result.get("is_last") is not True:
|
||||
raise RuntimeError("Serverless media delivery must be atomic and final")
|
||||
collected_count += self._store_typed_media(
|
||||
worker_images, worker_audio, result_worker_id, result["media"]
|
||||
)
|
||||
debug_log(f"Master - Got typed media from serverless participant {result_worker_id}")
|
||||
mark_worker_done(result_worker_id)
|
||||
return
|
||||
|
||||
if (exact_requested and result.get("tensor") is not None
|
||||
and result.get("format") != transfer_codec.EXACT_FORMAT):
|
||||
raise RuntimeError(
|
||||
f"Worker {result_worker_id} sent {result.get('format')} images, but this Collector "
|
||||
"requested exact. Update ComfyUI-Distributed on that worker; exact is never "
|
||||
"replaced by an 8-bit format."
|
||||
)
|
||||
is_last = result.get("is_last", False)
|
||||
collected_count += self._store_worker_result(worker_images, result)
|
||||
debug_log(
|
||||
@@ -733,6 +914,10 @@ class DistributedCollectorNode:
|
||||
consume_result(item)
|
||||
break
|
||||
|
||||
if exact_requested and workers_done != expected_workers:
|
||||
raise RuntimeError(
|
||||
f"Exact transfer incomplete: missing completion from {sorted(expected_workers - workers_done)}"
|
||||
)
|
||||
total_collected = sum(len(imgs) for imgs in worker_images.values())
|
||||
|
||||
# Remove the queue before combining so a late provider callback cannot
|
||||
@@ -753,10 +938,10 @@ class DistributedCollectorNode:
|
||||
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
|
||||
return (combined, combined_audio)
|
||||
except Exception as e:
|
||||
if serverless_results:
|
||||
if serverless_results or exact_requested:
|
||||
raise RuntimeError(
|
||||
"Serverless result batch shape mismatch while combining participants "
|
||||
f"{sorted(serverless_results)}: {e}"
|
||||
"Serverless or exact result batch shape mismatch while combining participants "
|
||||
f"{sorted(worker_images)}: {e}"
|
||||
) from e
|
||||
log(f"Master - Error combining images: {e}")
|
||||
# Preserve collected audio for regular workers when image assembly fails.
|
||||
|
||||
@@ -0,0 +1,266 @@
|
||||
# Standalone `comfy-media-v1` endpoint bridge
|
||||
|
||||
This directory is independently installable. It registers two internal Comfy V3
|
||||
nodes, `MediaImportV1` and `MediaExportV1`, and provides the asynchronous HTTP
|
||||
executor `execute_media_request`. It does **not** start ComfyUI, provision an
|
||||
endpoint, install models, or provide a provider worker/supervisor. The endpoint
|
||||
needs this bridge, not the full ComfyUI-Distributed extension.
|
||||
|
||||
The bridge's generic contract supports `IMAGE`, `AUDIO`, `MASK`, `STRING`, and
|
||||
native ComfyUI `VIDEO`. **Distributed Collector** (`DistributedCollector`) uses
|
||||
only the image/audio outputs and combines their returned lists into ordinary
|
||||
IMAGE/AUDIO outputs. It has no MASK/STRING/VIDEO outputs, and the separate
|
||||
`DistributedMediaCollector` node has been removed. For Collector setup, ordering,
|
||||
and limits, see [serverless images and audio](../docs/serverless-media.md).
|
||||
|
||||
## Install in the endpoint image
|
||||
|
||||
1. Copy this complete directory to `<ComfyUI>/custom_nodes/serverless_bridge`.
|
||||
Use a compatible ComfyUI with `comfy_api.v0_0_2` and the native VIDEO classes
|
||||
used by [video_codec.py](video_codec.py). These exact-class/private-state
|
||||
adapters are tied to the supported core implementation, not arbitrary VIDEO
|
||||
plugins or future core layouts.
|
||||
2. Restart ComfyUI so `comfy_entrypoint()` registers the two nodes. Keep the
|
||||
workflow's own models and custom nodes installed separately.
|
||||
3. Make `<ComfyUI>/custom_nodes` importable by the handler process (for example,
|
||||
via its `PYTHONPATH`). The handler uses `aiohttp`; the nodes/codecs use
|
||||
ComfyUI's PyTorch/native media dependencies. The handler module itself does
|
||||
not import ComfyUI, PyTorch, Runpod, or Distributed bootstrap code.
|
||||
4. Give the handler filesystem access to **the same input directory** used by
|
||||
that ComfyUI instance, plus its trusted HTTP URL. These are operator settings,
|
||||
never request-controlled paths/URLs. Staging/cleanup requires POSIX directory
|
||||
descriptor and no-follow file operations.
|
||||
|
||||
Installing the nodes alone is insufficient: the endpoint must route media-v1
|
||||
inputs through this executor and implement the uncertainty policy below. A
|
||||
legacy image handler cannot interpret the typed envelope.
|
||||
|
||||
## Callable integration
|
||||
|
||||
```python
|
||||
from serverless_bridge.handler import execute_media_request
|
||||
|
||||
|
||||
async def execute_endpoint_input(
|
||||
input_data, *, comfy_url, input_dir, temp_dir=None,
|
||||
download=None, publish=None,
|
||||
):
|
||||
# Invoke only under a supervisor implementing the uncertainty policy below.
|
||||
# Exceptions intentionally propagate; this is NOT a provider worker loop.
|
||||
return await execute_media_request(
|
||||
input_data,
|
||||
comfy_url=comfy_url,
|
||||
input_dir=input_dir,
|
||||
timeout_seconds=1800,
|
||||
session=None,
|
||||
download=download,
|
||||
publish=publish,
|
||||
temp_dir=temp_dir,
|
||||
)
|
||||
```
|
||||
|
||||
Pass the provider job's **`input` object**, not the whole job. Return the resulting
|
||||
`{"media": ...}` object as provider output; Runpod exposes it as `output.media`.
|
||||
Do not add another `output` wrapper inside it. `timeout_seconds` must be finite
|
||||
and positive. If supplied, `session` is a caller-owned `aiohttp.ClientSession`
|
||||
for local Comfy HTTP and remains open; otherwise the executor creates/closes its
|
||||
own session. Input dictionaries are copied before graph rewriting.
|
||||
|
||||
**This callable is not a turnkey Runpod handler.** An outer SDK that catches every
|
||||
exception and immediately reuses the worker is unsafe. The supervisor must
|
||||
recognize `serverless_bridge.handler.MediaExecutionUncertain`, stop intake, and
|
||||
quarantine/reconcile or terminate the owning environment before reuse. Merely
|
||||
returning `{"error": ...}` or marking a provider job `FAILED` is not that policy.
|
||||
|
||||
## Request and response
|
||||
|
||||
A minimal generic bridge text request builder (no files or model required by this
|
||||
graph). This illustrates the bridge contract, not a Collector text output:
|
||||
|
||||
```python
|
||||
import json
|
||||
import uuid
|
||||
|
||||
|
||||
def make_text_request(text="hello"):
|
||||
request_id = uuid.uuid4().hex
|
||||
binding = {
|
||||
"binding_id": "text-binding",
|
||||
"collector_id": "collector-1",
|
||||
"job_id": "collector-job-1",
|
||||
"participant_id": "serverless:example:0",
|
||||
"node_id": "export-1",
|
||||
"socket": "texts",
|
||||
"kind": "STRING",
|
||||
}
|
||||
return {
|
||||
"workflow": {
|
||||
"text-1": {"class_type": "PrimitiveString", "inputs": {"value": text}},
|
||||
"export-1": {
|
||||
"class_type": "MediaExportV1",
|
||||
"inputs": {
|
||||
"request_id": request_id,
|
||||
"binding_json": json.dumps(binding),
|
||||
"texts": ["text-1", 0],
|
||||
},
|
||||
},
|
||||
},
|
||||
"images": [],
|
||||
"media": {
|
||||
"version": 1,
|
||||
"profile": "comfy-media-v1",
|
||||
"request_id": request_id,
|
||||
"artifacts": [],
|
||||
"inputs": [],
|
||||
"outputs": [binding],
|
||||
},
|
||||
}
|
||||
```
|
||||
|
||||
`workflow` is an API graph, not an editor workflow. Top-level input fields are
|
||||
exactly `workflow`, `media`, and optional `images`. `media` has exactly the six
|
||||
fields shown above. `media.outputs` declares one exporter per connected socket;
|
||||
its binding must match that node's `binding_json` and request ID exactly. Socket
|
||||
names/types are `images/IMAGE`, `audio/AUDIO`, `masks/MASK`, `texts/STRING`, and
|
||||
`video/VIDEO`. Use unique request IDs. Identifiers are 1–200 characters from
|
||||
`[A-Za-z0-9_.:/-]`; request scopes additionally reject empty, `.` or `..` slash
|
||||
components. The executor validates the graph and submits **once**, with exactly
|
||||
the declared exporter IDs as `partial_execution_targets`. Unrelated graph roots,
|
||||
cycles, undeclared exporters, or caller-supplied `MediaImportV1` nodes fail.
|
||||
|
||||
For typed inputs, `media.artifacts` contains records with exactly:
|
||||
|
||||
| Field | Value |
|
||||
| --- | --- |
|
||||
| `id`, `kind` | Unique artifact ID and one of the five native types |
|
||||
| `descriptor` | Kind-specific codec metadata; use the shared codec, not pickle |
|
||||
| `byte_length`, `sha256` | Exact payload byte count and lowercase SHA-256 hex |
|
||||
| `transport` | `{"type": "inline", "data_base64": "..."}` or the URL transport below |
|
||||
|
||||
Each `media.inputs` entry is
|
||||
`{"node_id": "consumer", "input_name": "socket", "artifact_id": "artifact-id"}`.
|
||||
It replaces an existing consumer input with a generated import link; every
|
||||
artifact must be used, and destinations cannot repeat. The executor verifies
|
||||
bytes/checksums, stages owned files, inserts imports, then prunes disconnected
|
||||
loader ancestors. Optional legacy `images` entries are `{ "name": "input.png",
|
||||
"image": "<base64>" }`, must match stock `LoadImage` references, and are rewritten
|
||||
to owned staging paths. They are not an arbitrary filename-upload interface.
|
||||
|
||||
The returned `media` object has exactly `version`, `profile`, `request_id`,
|
||||
`bindings`, `artifacts`, and `items`:
|
||||
|
||||
- `bindings`: every declared binding, unchanged, plus `count` (1–64).
|
||||
- `artifacts`: the same artifact schema above with inline or URL transport.
|
||||
- `items`: `{ "binding_id": "...", "item_index": 0, "artifact_id": "..." }`
|
||||
records covering every index exactly once, ordered by `item_index` when decoded.
|
||||
|
||||
All identity, completeness, type, size, and checksum checks must pass; do not
|
||||
consume a partial result. The internal history transport `comfy-file` is local
|
||||
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. The **bridge-level defaults** are `exact` and `90`;
|
||||
generic callers that omit the controls retain that representation. Other media
|
||||
exports must not carry image controls. Callers should not alter binding identity
|
||||
metadata.
|
||||
|
||||
The **Collector default is different**: its saved widget value remains
|
||||
`legacy_png` with quality `90`. For a `comfy-media-v1` pool, the master compiler
|
||||
maps that value to an explicit `png` export, so normal serverless image returns
|
||||
use compressed 8-bit PNG. `exact` is opt-in at the Collector. Other selected
|
||||
formats are unchanged. Direct workers also support opt-in exact IMAGE transport;
|
||||
an older master rejection is an error, not an 8-bit retry. A Collector requesting
|
||||
`exact` cannot dispatch to a `legacy-image-v1` pool.
|
||||
|
||||
Formats: `exact`, `legacy_png`, `png`, `ffv1`, `jpeg`, `webp`, `h264`,
|
||||
`h265` (**experimental**), `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 bridge values still use their original exact codecs;
|
||||
the image-format controls do not add those sockets to the Collector.
|
||||
|
||||
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,
|
||||
128 artifacts, 64 items per binding. Request input totals include decoded
|
||||
legacy-image bytes. Decoders validate tensor/VIDEO sizes and layouts before
|
||||
allocating native storage; these are payload limits, **not peak-RSS guarantees**.
|
||||
- The built-in master currently emits inline inputs only. Its complete Runpod
|
||||
request JSON is capped at **10,000,000 bytes**, including base64, graph and
|
||||
policy. The larger contract limits do not bypass this cap; there is no
|
||||
built-in large-input/object-storage uploader.
|
||||
- The executor's result JSON is capped at **10,000,000 bytes**. Inline results
|
||||
are the default. When needed, `publish` is called for selected larger
|
||||
artifacts to fit that budget; supplying it does not force all results to URLs.
|
||||
Without a publisher, an oversized response fails instead of being truncated.
|
||||
- Optional **async** `download(url, *, max_bytes)` returns `bytes`. The executor
|
||||
supplies the declared artifact length and checks exact length and SHA-256.
|
||||
The callback must enforce public HTTPS, DNS/IP restrictions, redirect checks,
|
||||
and bounded streaming **before allocation**; syntactic contract validation
|
||||
does not provide SSRF protection. No downloader is installed by default.
|
||||
- Optional **async** `publish(metadata, payload)` uploads the exact `bytes` and
|
||||
returns exactly `{"type": "url", "url": "https://...", "expires_at": "..."}`.
|
||||
`metadata` is the artifact record without `transport`. Expiry must be a future,
|
||||
timezone-aware ISO timestamp; URLs use HTTPS port 443, no userinfo/fragment,
|
||||
and at most 8192 characters. Keep objects available long enough for collection.
|
||||
The callback/operator owns upload credentials, object expiry and cleanup,
|
||||
including orphan uploads after failure. There is no bridge storage backend.
|
||||
- The master uses a separate bounded public-HTTPS downloader without provider
|
||||
credentials. Never send provider credentials to artifact hosts or log signed
|
||||
URLs. Endpoint callbacks must enforce equivalent protections themselves.
|
||||
|
||||
Callbacks must cooperate with async cancellation, finish/abort their own I/O,
|
||||
and not detach readers/uploads that outlive the callback. Synchronous code that
|
||||
blocks the event loop cannot be hard-timed-out by this executor.
|
||||
|
||||
## Cancellation, quarantine, and file retention
|
||||
|
||||
`timeout_seconds` covers staging, submission, polling, downloads and publication.
|
||||
On failure/cancellation, there is a separate **10-second callback-drain grace**
|
||||
and, when needed, a **10-second prompt-reconciliation grace**. These are recovery
|
||||
allowances, not evidence that execution stopped. Ambiguous `/prompt` submission
|
||||
is never retried. An identified pending prompt may be deleted from `/queue`,
|
||||
with exact read-back; running prompts are never sent `/interrupt`. The pinned
|
||||
core route has a snapshot/global-interrupt race even when given a prompt ID.
|
||||
|
||||
`MediaExecutionUncertain` means prompt acceptance/termination **or callback
|
||||
teardown** is unconfirmed. Even successful terminal Comfy history is insufficient
|
||||
if an async downloader/publisher ignores cancellation. The bridge retains input
|
||||
and exporter files in that case; late callback exit does not automatically
|
||||
remove retained files. Before accepting another job on that environment:
|
||||
|
||||
1. Quarantine it and stop new intake; do not automatically resubmit the request.
|
||||
2. Reconcile the local prompt and all readers, or terminate the owning execution
|
||||
environment when their termination cannot be established.
|
||||
3. Only then clean retained files and allow reuse/replacement. A provider terminal
|
||||
job status alone cannot certify a surviving local Comfy process or callback.
|
||||
|
||||
Staged inputs use random exclusive files under
|
||||
`input_dir/distributed-media-v1/<sha256(request_id)>/`. Cleanup removes only files
|
||||
created by that invocation, after confirmed terminal state (or before submission)
|
||||
and reader teardown; it never recursively deletes a shared namespace.
|
||||
|
||||
`temp_dir=None` leaves exporter retention to the Comfy temp-directory/environment
|
||||
lifecycle. Optionally pass the **trusted local Comfy temp root** to enable
|
||||
best-effort cleanup of validated, invocation-owned exporter UUID `.bin` files
|
||||
under the corresponding hashed namespace, only after terminal history and
|
||||
collection/publisher teardown. Symlinked, hard-linked, replaced, missing, or
|
||||
unidentifiable files are not claimed; directories and all unidentified leftovers
|
||||
remain lifecycle-owned. This is not a temp-directory sweep. Independently decoded
|
||||
VIDEO objects do not depend on either staging or exporter files.
|
||||
|
||||
This guide describes the source contract, not a test result. No live provider/GPU
|
||||
validation is claimed, and this directory is not a turnkey provider deployment.
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Standalone typed-media bridge; copy this directory into endpoint custom_nodes.
|
||||
|
||||
Importing codecs does not load ComfyUI, providers or Distributed bootstrap.
|
||||
"""
|
||||
|
||||
async def comfy_entrypoint():
|
||||
from comfy_api.v0_0_2 import ComfyExtension
|
||||
from .nodes import MediaImportV1, MediaExportV1
|
||||
|
||||
class MediaBridgeExtension(ComfyExtension):
|
||||
async def get_node_list(self):
|
||||
return [MediaImportV1, MediaExportV1]
|
||||
|
||||
return MediaBridgeExtension()
|
||||
@@ -0,0 +1,189 @@
|
||||
"""Lossless, bounded CPU payloads shared with the standalone endpoint bridge.
|
||||
|
||||
Descriptors contain only kind-specific metadata. Tensor bytes are contiguous,
|
||||
little-endian floating-point values, never a pickle or an image/audio container.
|
||||
The optional byte budget can lower, but never raise, MAX_PAYLOAD_BYTES.
|
||||
"""
|
||||
|
||||
import ctypes
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
MAX_PAYLOAD_BYTES = 256 * 1024 * 1024
|
||||
|
||||
_DTYPES = {
|
||||
"float16": (torch.float16, 2),
|
||||
"bfloat16": (torch.bfloat16, 2),
|
||||
"float32": (torch.float32, 4),
|
||||
"float64": (torch.float64, 8),
|
||||
}
|
||||
_DTYPE_NAMES = {dtype: name for name, (dtype, _) in _DTYPES.items()}
|
||||
_LAYOUTS = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}
|
||||
|
||||
|
||||
def _validate_kind(kind):
|
||||
if type(kind) is not str or kind not in ("IMAGE", "MASK", "STRING", "AUDIO"):
|
||||
raise ValueError("Unsupported payload kind")
|
||||
|
||||
|
||||
def _byte_limit(max_bytes):
|
||||
if type(max_bytes) is not int or max_bytes < 0:
|
||||
raise ValueError("max_bytes must be a nonnegative integer")
|
||||
return min(max_bytes, MAX_PAYLOAD_BYTES)
|
||||
|
||||
|
||||
def _sample_rate(value):
|
||||
if type(value) is not int or value <= 0:
|
||||
raise ValueError("AUDIO sample_rate must be a positive integer")
|
||||
return value
|
||||
|
||||
|
||||
def _tensor_size(kind, shape, itemsize, limit):
|
||||
if type(shape) is not list or len(shape) != len(_LAYOUTS[kind]):
|
||||
raise ValueError("Tensor shape has an invalid rank")
|
||||
if any(type(dimension) is not int or dimension <= 0 for dimension in shape):
|
||||
raise ValueError("Tensor shape dimensions must be positive integers")
|
||||
if kind == "IMAGE" and shape[-1] not in (1, 3, 4):
|
||||
raise ValueError("IMAGE shape must have 1, 3 or 4 channels")
|
||||
size = itemsize
|
||||
for dimension in shape:
|
||||
# Divide before multiplying, including for arbitrary-size Python ints.
|
||||
if dimension > limit // size:
|
||||
raise ValueError("Tensor payload exceeds byte limit")
|
||||
size *= dimension
|
||||
return size
|
||||
|
||||
|
||||
def _native_little_endian(payload, itemsize):
|
||||
"""Convert native to wire byte order, or vice versa (an involution)."""
|
||||
if sys.byteorder == "little":
|
||||
return payload
|
||||
swapped = bytearray(len(payload))
|
||||
for offset in range(itemsize):
|
||||
swapped[offset::itemsize] = payload[itemsize - 1 - offset::itemsize]
|
||||
return bytes(swapped)
|
||||
|
||||
|
||||
def _encode_string(value, limit):
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("STRING value must be text")
|
||||
# Check the exact UTF-8 size before allocating the encoded payload.
|
||||
if len(value) > limit:
|
||||
raise ValueError("STRING payload exceeds byte limit")
|
||||
size = 0
|
||||
for character in value:
|
||||
codepoint = ord(character)
|
||||
if 0xD800 <= codepoint <= 0xDFFF:
|
||||
raise ValueError("STRING value must be valid UTF-8 text")
|
||||
size += 1 + (codepoint > 0x7F) + (codepoint > 0x7FF) + (codepoint > 0xFFFF)
|
||||
if size > limit:
|
||||
raise ValueError("STRING payload exceeds byte limit")
|
||||
return {"encoding": "utf8-v1"}, value.encode("utf-8")
|
||||
|
||||
|
||||
def encode_payload(kind: str, value, *, max_bytes: int = MAX_PAYLOAD_BYTES) -> tuple[dict, bytes]:
|
||||
"""Encode a typed value without dtype conversion, normalization or GPU use.
|
||||
|
||||
IMAGE/MASK accept dense CPU tensors; AUDIO accepts a waveform CPU tensor
|
||||
plus a positive integer sample_rate. STRING accepts exact Unicode text.
|
||||
Invalid values and exceeded budgets raise nonsensitive ValueError messages.
|
||||
"""
|
||||
_validate_kind(kind)
|
||||
limit = _byte_limit(max_bytes)
|
||||
if kind == "STRING":
|
||||
return _encode_string(value, limit)
|
||||
|
||||
sample_rate = None
|
||||
if kind == "AUDIO":
|
||||
if not isinstance(value, dict) or set(value) != {"waveform", "sample_rate"}:
|
||||
raise ValueError("AUDIO value must contain waveform and sample_rate")
|
||||
sample_rate = _sample_rate(value["sample_rate"])
|
||||
value = value["waveform"]
|
||||
if not isinstance(value, torch.Tensor):
|
||||
raise ValueError("Tensor value must be a torch.Tensor")
|
||||
if value.device.type != "cpu":
|
||||
raise ValueError("Tensor value must be on CPU")
|
||||
if value.layout != torch.strided or value.is_nested:
|
||||
raise ValueError("Tensor value must have a dense strided layout")
|
||||
if value.dtype not in _DTYPE_NAMES:
|
||||
raise ValueError("Unsupported tensor dtype")
|
||||
|
||||
dtype_name = _DTYPE_NAMES[value.dtype]
|
||||
_, itemsize = _DTYPES[dtype_name]
|
||||
shape = list(value.shape)
|
||||
size = _tensor_size(kind, shape, itemsize, limit)
|
||||
tensor = value.detach().resolve_neg().contiguous()
|
||||
if not torch.isfinite(tensor).all().item():
|
||||
raise ValueError("Tensor values must be finite")
|
||||
# A uint16 view preserves BF16 bits without conversion through float32.
|
||||
if tensor.dtype == torch.bfloat16:
|
||||
tensor = tensor.view(torch.uint16)
|
||||
# The live, contiguous CPU tensor owns this pointer for the entire copy.
|
||||
payload = _native_little_endian(ctypes.string_at(tensor.data_ptr(), size), itemsize)
|
||||
descriptor = {
|
||||
"encoding": "tensor-le-v1", "dtype": dtype_name,
|
||||
"shape": shape, "layout": _LAYOUTS[kind],
|
||||
}
|
||||
if kind == "AUDIO":
|
||||
descriptor["sample_rate"] = sample_rate
|
||||
return descriptor, payload
|
||||
|
||||
|
||||
def decode_payload(kind: str, descriptor: dict, payload: bytes, *, max_bytes: int = MAX_PAYLOAD_BYTES):
|
||||
"""Decode validated wire bytes into independent contiguous CPU tensors/text.
|
||||
|
||||
All metadata, shape products and byte lengths are checked before tensor
|
||||
allocation. No GPU device discovery, transfers or object deserialization
|
||||
occur. AUDIO returns {'waveform': tensor, 'sample_rate': int}.
|
||||
"""
|
||||
_validate_kind(kind)
|
||||
limit = _byte_limit(max_bytes)
|
||||
if not isinstance(descriptor, dict):
|
||||
raise ValueError("Payload descriptor must be a dictionary")
|
||||
if not isinstance(payload, bytes):
|
||||
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")
|
||||
try:
|
||||
return payload.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
raise ValueError("STRING payload must be valid UTF-8 bytes") from None
|
||||
|
||||
fields = {"encoding", "dtype", "shape", "layout"}
|
||||
if kind == "AUDIO":
|
||||
fields.add("sample_rate")
|
||||
if set(descriptor) != fields:
|
||||
raise ValueError("Invalid tensor descriptor fields")
|
||||
if descriptor["encoding"] != "tensor-le-v1":
|
||||
raise ValueError("Unsupported tensor descriptor encoding")
|
||||
if descriptor["layout"] != _LAYOUTS[kind]:
|
||||
raise ValueError("Invalid tensor descriptor layout")
|
||||
dtype_name = descriptor["dtype"]
|
||||
if type(dtype_name) is not str or dtype_name not in _DTYPES:
|
||||
raise ValueError("Unsupported tensor descriptor dtype")
|
||||
dtype, itemsize = _DTYPES[dtype_name]
|
||||
size = _tensor_size(kind, descriptor["shape"], itemsize, limit)
|
||||
if len(payload) != size:
|
||||
raise ValueError("Tensor payload length does not match shape and dtype")
|
||||
sample_rate = _sample_rate(descriptor["sample_rate"]) if kind == "AUDIO" else None
|
||||
|
||||
# frombuffer is always CPU. Copy immutable wire bytes into a writable private
|
||||
# buffer, then clone so the returned tensor owns its storage independently.
|
||||
buffer = bytearray(_native_little_endian(payload, itemsize))
|
||||
storage_dtype = torch.uint16 if dtype == torch.bfloat16 else dtype
|
||||
tensor = torch.frombuffer(buffer, dtype=storage_dtype).view(dtype)
|
||||
if not torch.isfinite(tensor).all().item():
|
||||
raise ValueError("Tensor payload values must be finite")
|
||||
tensor = tensor.clone().reshape(descriptor["shape"])
|
||||
if kind == "AUDIO":
|
||||
return {"waveform": tensor, "sample_rate": sample_rate}
|
||||
return tensor
|
||||
@@ -0,0 +1,249 @@
|
||||
"""Versioned, bounded media envelopes shared by master and endpoint bridge.
|
||||
|
||||
No provider or ComfyUI bootstrap imports. Network I/O is supplied by the caller,
|
||||
which must enforce public-HTTPS resolution, redirect and streaming byte limits.
|
||||
"""
|
||||
import base64
|
||||
import binascii
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from datetime import datetime, timezone
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
PROFILE = "comfy-media-v1"
|
||||
VERSION = 1
|
||||
MAX_ARTIFACT_BYTES = 256 * 1024 * 1024
|
||||
MAX_TOTAL_BYTES = 512 * 1024 * 1024
|
||||
MAX_ARTIFACTS = 128
|
||||
MAX_ITEMS = 64
|
||||
MAX_DESCRIPTOR_BYTES = 128 * 1024
|
||||
SOCKET_KINDS = {"images": "IMAGE", "audio": "AUDIO", "masks": "MASK", "texts": "STRING", "video": "VIDEO"}
|
||||
BINDING_FIELDS = {"binding_id", "collector_id", "job_id", "participant_id", "node_id", "socket", "kind"}
|
||||
_ID = re.compile(r"[A-Za-z0-9_.:/-]{1,200}\Z")
|
||||
_HASH = re.compile(r"[a-f0-9]{64}\Z")
|
||||
|
||||
|
||||
def identifier(value):
|
||||
if not isinstance(value, str) or not _ID.fullmatch(value):
|
||||
raise ValueError("Invalid media identifier")
|
||||
return value
|
||||
|
||||
|
||||
def bounded_int(value, maximum, *, minimum=0):
|
||||
if type(value) is not int or not minimum <= value <= maximum:
|
||||
raise ValueError("Invalid media count or size")
|
||||
return value
|
||||
|
||||
|
||||
def validate_json(value, *, max_bytes=MAX_DESCRIPTOR_BYTES):
|
||||
"""Bound metadata depth and size without accepting non-finite JSON numbers."""
|
||||
def walk(item, depth=0):
|
||||
if depth > 12:
|
||||
raise ValueError("Media metadata nesting exceeds limit")
|
||||
if isinstance(item, dict):
|
||||
if len(item) > 4096 or any(not isinstance(k, str) for k in item):
|
||||
raise ValueError("Invalid media metadata object")
|
||||
for child in item.values():
|
||||
walk(child, depth + 1)
|
||||
elif isinstance(item, list):
|
||||
if len(item) > 4096:
|
||||
raise ValueError("Media metadata list exceeds limit")
|
||||
for child in item:
|
||||
walk(child, depth + 1)
|
||||
elif item is not None and type(item) not in (str, int, float, bool):
|
||||
raise ValueError("Invalid media metadata value")
|
||||
walk(value)
|
||||
try:
|
||||
raw = json.dumps(value, allow_nan=False, separators=(",", ":")).encode("utf-8")
|
||||
except (ValueError, TypeError, UnicodeError, RecursionError):
|
||||
raise ValueError("Invalid media JSON") from None
|
||||
if len(raw) > max_bytes:
|
||||
raise ValueError("Media JSON exceeds byte limit")
|
||||
return raw
|
||||
|
||||
|
||||
def validate_binding(binding):
|
||||
if not isinstance(binding, dict) or set(binding) != BINDING_FIELDS:
|
||||
raise ValueError("Invalid media binding")
|
||||
for field in BINDING_FIELDS - {"socket", "kind"}:
|
||||
identifier(binding[field])
|
||||
if SOCKET_KINDS.get(binding["socket"]) != binding["kind"]:
|
||||
raise ValueError("Media socket/type mismatch")
|
||||
return binding
|
||||
|
||||
|
||||
def validate_bindings(bindings):
|
||||
if not isinstance(bindings, list) or not 1 <= len(bindings) <= MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid media bindings count")
|
||||
result = {}
|
||||
destinations = set()
|
||||
for binding in bindings:
|
||||
validate_binding(binding)
|
||||
bid = binding["binding_id"]
|
||||
destination = (binding["job_id"], binding["participant_id"], binding["collector_id"], binding["socket"])
|
||||
if bid in result or destination in destinations:
|
||||
raise ValueError("Duplicate media binding")
|
||||
result[bid] = binding
|
||||
destinations.add(destination)
|
||||
return result
|
||||
|
||||
|
||||
def _validate_url_transport(transport):
|
||||
if set(transport) != {"type", "url", "expires_at"}:
|
||||
raise ValueError("Invalid media URL transport")
|
||||
url = transport["url"]
|
||||
if not isinstance(url, str) or len(url) > 8192:
|
||||
raise ValueError("Invalid media URL")
|
||||
try:
|
||||
parsed = urlsplit(url)
|
||||
if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.fragment:
|
||||
raise ValueError()
|
||||
if parsed.port not in (None, 443):
|
||||
raise ValueError()
|
||||
expiry = datetime.fromisoformat(transport["expires_at"].replace("Z", "+00:00"))
|
||||
if expiry.tzinfo is None or expiry <= datetime.now(timezone.utc):
|
||||
raise ValueError()
|
||||
except (ValueError, TypeError, AttributeError):
|
||||
raise ValueError("Invalid or expired media URL") from None
|
||||
|
||||
|
||||
def validate_artifact(artifact, *, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
fields = {"id", "kind", "descriptor", "byte_length", "sha256", "transport"}
|
||||
if not isinstance(artifact, dict) or set(artifact) != fields:
|
||||
raise ValueError("Invalid media artifact")
|
||||
identifier(artifact["id"])
|
||||
if artifact["kind"] not in SOCKET_KINDS.values():
|
||||
raise ValueError("Unsupported media kind")
|
||||
if not isinstance(artifact["descriptor"], dict):
|
||||
raise ValueError("Invalid media descriptor")
|
||||
validate_json(artifact["descriptor"])
|
||||
size = bounded_int(artifact["byte_length"], max_bytes)
|
||||
digest = artifact["sha256"]
|
||||
if not isinstance(digest, str) or not _HASH.fullmatch(digest):
|
||||
raise ValueError("Invalid media checksum")
|
||||
transport = artifact["transport"]
|
||||
if not isinstance(transport, dict):
|
||||
raise ValueError("Invalid media transport")
|
||||
if transport.get("type") == "inline":
|
||||
if set(transport) != {"type", "data_base64"} or not isinstance(transport["data_base64"], str):
|
||||
raise ValueError("Invalid inline media transport")
|
||||
if len(transport["data_base64"]) != 4 * ((size + 2) // 3):
|
||||
raise ValueError("Inline media length mismatch")
|
||||
elif transport.get("type") == "url":
|
||||
_validate_url_transport(transport)
|
||||
else:
|
||||
raise ValueError("Unsupported media transport")
|
||||
return artifact
|
||||
|
||||
|
||||
def artifact_from_payload(artifact_id, kind, descriptor, payload, *, transport=None, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
if not isinstance(payload, bytes):
|
||||
raise ValueError("Media payload must be bytes")
|
||||
bounded_int(len(payload), max_bytes)
|
||||
result = {"id": artifact_id, "kind": kind, "descriptor": descriptor,
|
||||
"byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(),
|
||||
"transport": transport if transport is not None else {
|
||||
"type": "inline", "data_base64": base64.b64encode(payload).decode("ascii")}}
|
||||
return validate_artifact(result, max_bytes=max_bytes)
|
||||
|
||||
|
||||
async def resolve_artifact(artifact, *, download=None, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
"""download(url, *, max_bytes) must return bounded bytes; never log its URL."""
|
||||
validate_artifact(artifact, max_bytes=max_bytes)
|
||||
transport = artifact["transport"]
|
||||
if transport["type"] == "inline":
|
||||
try:
|
||||
payload = base64.b64decode(transport["data_base64"], validate=True)
|
||||
except (ValueError, binascii.Error):
|
||||
raise ValueError("Invalid media base64") from None
|
||||
else:
|
||||
if download is None:
|
||||
raise ValueError("Media URL download is not configured")
|
||||
payload = await download(transport["url"], max_bytes=artifact["byte_length"])
|
||||
if not isinstance(payload, bytes) or len(payload) != artifact["byte_length"]:
|
||||
raise ValueError("Media byte length mismatch")
|
||||
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
|
||||
raise ValueError("Media checksum mismatch")
|
||||
return payload
|
||||
|
||||
|
||||
def validate_response(response, *, request_id, bindings, max_total_bytes=MAX_TOTAL_BYTES):
|
||||
"""Validate ALL identity/completeness before downloading or delivering anything."""
|
||||
expected = validate_bindings(bindings)
|
||||
identifier(request_id)
|
||||
if not isinstance(response, dict) or set(response) != {"version", "profile", "request_id", "bindings", "artifacts", "items"}:
|
||||
raise ValueError("Invalid media response envelope")
|
||||
if type(response["version"]) is not int or response["version"] != VERSION or response["profile"] != PROFILE:
|
||||
raise ValueError("Unsupported media response version/profile")
|
||||
if response["request_id"] != request_id:
|
||||
raise ValueError("Media request identity mismatch")
|
||||
returned = response["bindings"]
|
||||
if not isinstance(returned, list) or len(returned) != len(expected):
|
||||
raise ValueError("Missing or extra media bindings")
|
||||
counts = {}
|
||||
for row in returned:
|
||||
if not isinstance(row, dict) or set(row) != BINDING_FIELDS | {"count"}:
|
||||
raise ValueError("Invalid result binding")
|
||||
binding = {key: row[key] for key in BINDING_FIELDS}
|
||||
validate_binding(binding)
|
||||
bid = binding["binding_id"]
|
||||
if bid in counts or expected.get(bid) != binding:
|
||||
raise ValueError("Media binding identity mismatch")
|
||||
# Connected outputs must produce at least one value. Empty STRING is a value.
|
||||
counts[bid] = bounded_int(row["count"], MAX_ITEMS, minimum=1)
|
||||
artifacts = response["artifacts"]
|
||||
if not isinstance(artifacts, list) or not 1 <= len(artifacts) <= MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid media artifacts count")
|
||||
indexed = {}
|
||||
total = 0
|
||||
for artifact in artifacts:
|
||||
validate_artifact(artifact)
|
||||
aid = artifact["id"]
|
||||
if aid in indexed:
|
||||
raise ValueError("Duplicate media artifact")
|
||||
total += artifact["byte_length"]
|
||||
if total > max_total_bytes:
|
||||
raise ValueError("Media response exceeds total byte limit")
|
||||
indexed[aid] = artifact
|
||||
items = response["items"]
|
||||
if not isinstance(items, list) or len(items) != sum(counts.values()):
|
||||
raise ValueError("Missing or extra media items")
|
||||
ordered = {bid: [None] * count for bid, count in counts.items()}
|
||||
used = set()
|
||||
for item in items:
|
||||
if not isinstance(item, dict) or set(item) != {"binding_id", "item_index", "artifact_id"}:
|
||||
raise ValueError("Invalid media item")
|
||||
bid, aid = identifier(item["binding_id"]), identifier(item["artifact_id"])
|
||||
if bid not in counts or aid not in indexed:
|
||||
raise ValueError("Unknown media item identity")
|
||||
index = bounded_int(item["item_index"], counts[bid] - 1)
|
||||
if ordered[bid][index] is not None or indexed[aid]["kind"] != expected[bid]["kind"]:
|
||||
raise ValueError("Duplicate or mistyped media item")
|
||||
ordered[bid][index] = indexed[aid]
|
||||
used.add(aid)
|
||||
if used != set(indexed) or any(value is None for values in ordered.values() for value in values):
|
||||
raise ValueError("Incomplete media result manifest")
|
||||
return ordered
|
||||
|
||||
|
||||
async def decode_response(response, *, request_id, bindings, download=None, max_total_bytes=MAX_TOTAL_BYTES):
|
||||
ordered = validate_response(response, request_id=request_id, bindings=bindings, max_total_bytes=max_total_bytes)
|
||||
# Resolve and decode before exposing a single result to the collector queue.
|
||||
from .codec import decode_payload
|
||||
cache = {}
|
||||
decoded = {}
|
||||
for bid, artifacts in ordered.items():
|
||||
values = []
|
||||
for artifact in artifacts:
|
||||
aid = artifact["id"]
|
||||
if aid not in cache:
|
||||
payload = await resolve_artifact(artifact, download=download)
|
||||
if artifact["kind"] == "VIDEO":
|
||||
from .video_codec import decode_video
|
||||
cache[aid] = decode_video(artifact["descriptor"], payload)
|
||||
else:
|
||||
cache[aid] = decode_payload(artifact["kind"], artifact["descriptor"], payload)
|
||||
values.append(cache[aid])
|
||||
decoded[bid] = values
|
||||
return decoded
|
||||
@@ -0,0 +1,794 @@
|
||||
"""Pure HTTP executor for media-v1 requests against a local ComfyUI instance.
|
||||
|
||||
No ComfyUI, tensor, or provider imports. ``download(url, *, max_bytes)`` must
|
||||
provide public-HTTPS/redirect/DNS protection and bounded streaming. Optional
|
||||
``publish(metadata, payload)`` returns a contract-valid, expiring HTTPS transport.
|
||||
|
||||
Only identified invocation-owned files are deleted. Optional trusted local
|
||||
``temp_dir`` enables exporter cleanup after terminal history and collection.
|
||||
Without it, exporter retention belongs to the ComfyUI temp-directory lifecycle.
|
||||
The pinned native /interrupt route is not atomic: running-prompt cancellation
|
||||
fails closed and retains files instead of risking an unrelated successor prompt.
|
||||
"""
|
||||
import asyncio
|
||||
import base64
|
||||
import binascii
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import stat
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from urllib.parse import quote, urlsplit
|
||||
|
||||
import aiohttp
|
||||
|
||||
from .contract import (
|
||||
BINDING_FIELDS, MAX_ARTIFACT_BYTES, MAX_ARTIFACTS, MAX_TOTAL_BYTES,
|
||||
PROFILE, VERSION, artifact_from_payload, identifier, resolve_artifact,
|
||||
validate_artifact, validate_bindings, validate_json, validate_response,
|
||||
)
|
||||
|
||||
MAX_RESPONSE_JSON_BYTES = 10_000_000
|
||||
MAX_HTTP_JSON_BYTES = 32 * 1024 * 1024
|
||||
POLL_INTERVAL_SECONDS = 0.25
|
||||
CANCEL_GRACE_SECONDS = 10.0
|
||||
DRAIN_GRACE_SECONDS = 10.0
|
||||
HTTP_TIMEOUT_SECONDS = 15.0
|
||||
_UNCERTAIN_WORK = set()
|
||||
GET_ATTEMPTS = 3
|
||||
_IMPORT_SLOTS = {"IMAGE": 0, "AUDIO": 1, "MASK": 2, "STRING": 3, "VIDEO": 4}
|
||||
# Used ONLY to run the shared manifest validator before /view. Never fetched.
|
||||
_VALIDATION_TRANSPORT = {"type": "url", "url": "https://media-validation.invalid/artifact",
|
||||
"expires_at": "9999-01-01T00:00:00Z"}
|
||||
|
||||
|
||||
class MediaExecutionUncertain(RuntimeError):
|
||||
"""The local prompt may still execute; staged files were deliberately retained."""
|
||||
|
||||
|
||||
class _HTTPError(RuntimeError):
|
||||
def __init__(self, status):
|
||||
super().__init__(f"Local Comfy HTTP request failed (status {status})")
|
||||
self.status = status
|
||||
|
||||
|
||||
def _scope_id(value):
|
||||
identifier(value)
|
||||
if any(part in ("", ".", "..") for part in value.split("/")):
|
||||
raise ValueError("Unsafe media scope identifier")
|
||||
return value
|
||||
|
||||
|
||||
def _json_bytes(value):
|
||||
return json.dumps(value, allow_nan=False).encode("utf-8")
|
||||
|
||||
|
||||
def _link(value, workflow):
|
||||
if (not isinstance(value, list) or len(value) != 2 or
|
||||
not isinstance(value[0], str) or value[0] not in workflow or
|
||||
type(value[1]) is not int or value[1] < 0):
|
||||
raise ValueError("Invalid workflow socket link")
|
||||
return value[0]
|
||||
|
||||
|
||||
def _validate_request(input_data):
|
||||
if (not isinstance(input_data, dict) or not {"workflow", "media"} <= set(input_data)
|
||||
or set(input_data) - {"workflow", "media", "images"}):
|
||||
raise ValueError("Invalid media request envelope")
|
||||
media = input_data["media"]
|
||||
if (not isinstance(media, dict) or set(media) !=
|
||||
{"version", "profile", "request_id", "artifacts", "inputs", "outputs"}):
|
||||
raise ValueError("Invalid media request controls")
|
||||
if type(media["version"]) is not int or media["version"] != VERSION or media["profile"] != PROFILE:
|
||||
raise ValueError("Unsupported media request version/profile")
|
||||
_scope_id(media["request_id"])
|
||||
bindings = validate_bindings(media["outputs"])
|
||||
targets = []
|
||||
for binding in bindings.values():
|
||||
for field in BINDING_FIELDS - {"socket", "kind"}:
|
||||
_scope_id(binding[field])
|
||||
targets.append(binding["node_id"])
|
||||
if len(set(targets)) != len(targets):
|
||||
raise ValueError("Duplicate media exporter node")
|
||||
|
||||
workflow = input_data["workflow"]
|
||||
if not isinstance(workflow, dict) or not workflow:
|
||||
raise ValueError("Invalid API workflow")
|
||||
validate_json(workflow, max_bytes=16 * 1024 * 1024)
|
||||
dependencies = {}
|
||||
for node_id, node in workflow.items():
|
||||
_scope_id(node_id)
|
||||
if (not isinstance(node, dict) or not isinstance(node.get("class_type"), str)
|
||||
or not node["class_type"] or not isinstance(node.get("inputs"), dict)):
|
||||
raise ValueError("Invalid API workflow node")
|
||||
if node["class_type"] == "MediaImportV1":
|
||||
raise ValueError("Request may not supply its own media import filenames")
|
||||
if node["class_type"] == "MediaExportV1" and node_id not in targets:
|
||||
raise ValueError("Undeclared media exporter")
|
||||
dependencies[node_id] = set()
|
||||
for name, value in node["inputs"].items():
|
||||
_scope_id(name)
|
||||
if isinstance(value, list):
|
||||
dependencies[node_id].add(_link(value, workflow))
|
||||
for binding in bindings.values():
|
||||
node = workflow.get(binding["node_id"])
|
||||
if node is None or node["class_type"] != "MediaExportV1":
|
||||
raise ValueError("Missing declared MediaExportV1 node")
|
||||
controls = node["inputs"]
|
||||
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"]
|
||||
if not isinstance(text, str) or len(text) > 128 * 1024:
|
||||
raise ValueError("Invalid media exporter binding JSON")
|
||||
try:
|
||||
declared = json.loads(text)
|
||||
except (ValueError, RecursionError):
|
||||
raise ValueError("Invalid media exporter binding JSON") from None
|
||||
if declared != binding:
|
||||
raise ValueError("Media exporter binding identity mismatch")
|
||||
_link(controls[binding["socket"]], workflow)
|
||||
# The execution roots are exactly the declared exporters, not every class
|
||||
# with OUTPUT_NODE=True. Required upstream output utilities may remain.
|
||||
visited = set()
|
||||
pending = list(targets)
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id not in visited:
|
||||
visited.add(node_id)
|
||||
pending.extend(dependencies[node_id])
|
||||
if visited != set(workflow):
|
||||
raise ValueError("Workflow contains nodes outside declared output roots")
|
||||
if any(set(targets) & sources for sources in dependencies.values()):
|
||||
raise ValueError("A media exporter cannot be an upstream dependency")
|
||||
# Kahn traversal bounds cycle checking without relying on recursion depth.
|
||||
counts = {node: len(sources) for node, sources in dependencies.items()}
|
||||
consumers = {node: [] for node in workflow}
|
||||
for node, sources in dependencies.items():
|
||||
for source in sources:
|
||||
consumers[source].append(node)
|
||||
pending = [node for node, count in counts.items() if not count]
|
||||
processed = 0
|
||||
while pending:
|
||||
processed += 1
|
||||
for consumer in consumers[pending.pop()]:
|
||||
counts[consumer] -= 1
|
||||
if not counts[consumer]:
|
||||
pending.append(consumer)
|
||||
if processed != len(workflow):
|
||||
raise ValueError("Workflow contains a dependency cycle")
|
||||
|
||||
artifacts = media["artifacts"]
|
||||
if not isinstance(artifacts, list) or len(artifacts) > MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid input artifacts count")
|
||||
indexed, total = {}, 0
|
||||
for artifact in artifacts:
|
||||
validate_artifact(artifact)
|
||||
aid = _scope_id(artifact["id"])
|
||||
if aid in indexed:
|
||||
raise ValueError("Duplicate input artifact")
|
||||
indexed[aid] = artifact
|
||||
total += artifact["byte_length"]
|
||||
inputs = media["inputs"]
|
||||
if not isinstance(inputs, list) or len(inputs) > 4096:
|
||||
raise ValueError("Invalid input references")
|
||||
destinations, used = set(), set()
|
||||
exporters = {binding["node_id"]: binding for binding in bindings.values()}
|
||||
for ref in inputs:
|
||||
if not isinstance(ref, dict) or set(ref) != {"node_id", "input_name", "artifact_id"}:
|
||||
raise ValueError("Invalid media input reference")
|
||||
for value in ref.values():
|
||||
_scope_id(value)
|
||||
node_id, name, aid = ref["node_id"], ref["input_name"], ref["artifact_id"]
|
||||
if (node_id not in workflow or name not in workflow[node_id]["inputs"]
|
||||
or aid not in indexed or (node_id, name) in destinations):
|
||||
raise ValueError("Unknown or duplicate media input reference")
|
||||
binding = exporters.get(node_id)
|
||||
if binding is not None and (name != binding["socket"] or indexed[aid]["kind"] != binding["kind"]):
|
||||
raise ValueError("Exporter input must match its declared media socket and kind")
|
||||
destinations.add((node_id, name))
|
||||
used.add(aid)
|
||||
if used != set(indexed):
|
||||
raise ValueError("Every input artifact must be used")
|
||||
|
||||
images = input_data.get("images", [])
|
||||
if not isinstance(images, list) or len(images) > MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid legacy images count")
|
||||
names = set()
|
||||
legacy = []
|
||||
for image in images:
|
||||
if not isinstance(image, dict) or set(image) != {"name", "image"}:
|
||||
raise ValueError("Invalid legacy image")
|
||||
name, encoded = image["name"], image["image"]
|
||||
if (not isinstance(name, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_. -]{0,254}", name)
|
||||
or name in names or not isinstance(encoded, str)):
|
||||
raise ValueError("Invalid or duplicate legacy image basename")
|
||||
names.add(name)
|
||||
if len(encoded) > 4 * ((MAX_ARTIFACT_BYTES + 2) // 3):
|
||||
raise ValueError("Legacy image exceeds byte limit")
|
||||
# Exact decoded length, including padding, before allocating bytes.
|
||||
size = (len(encoded) // 4) * 3 - (len(encoded) - len(encoded.rstrip("=")))
|
||||
if len(encoded) % 4 or size < 0 or size > MAX_ARTIFACT_BYTES:
|
||||
raise ValueError("Invalid legacy image base64 length")
|
||||
total += size
|
||||
matches = [node_id for node_id, node in workflow.items()
|
||||
if node["class_type"] == "LoadImage" and node["inputs"].get("image") == name]
|
||||
if not matches or any((node_id, "image") in destinations for node_id in matches):
|
||||
raise ValueError("Unused or conflicting legacy image binding")
|
||||
legacy.append((encoded, size, matches))
|
||||
if total > MAX_TOTAL_BYTES:
|
||||
raise ValueError("Media inputs exceed total byte limit")
|
||||
return copy.deepcopy(workflow), copy.deepcopy(media), legacy, targets
|
||||
|
||||
|
||||
class _Staging:
|
||||
"""Exclusive random files under a hashed namespace; never delete siblings."""
|
||||
def __init__(self, input_dir, request_id):
|
||||
self.root = Path(input_dir).resolve()
|
||||
self.key = hashlib.sha256(request_id.encode()).hexdigest()
|
||||
self.relative = f"distributed-media-v1/{self.key}"
|
||||
self.names = []
|
||||
self.fd = None
|
||||
|
||||
def _open(self):
|
||||
self.root.mkdir(parents=True, exist_ok=True)
|
||||
directory = self.root
|
||||
for component in ("distributed-media-v1", self.key):
|
||||
directory = directory / component
|
||||
if directory.is_symlink():
|
||||
raise ValueError("Unsafe media staging directory symlink")
|
||||
directory.mkdir(exist_ok=True, mode=0o700)
|
||||
if not directory.resolve().is_relative_to(self.root):
|
||||
raise ValueError("Unsafe media staging directory")
|
||||
self.fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
|
||||
|
||||
def write(self, payload):
|
||||
if self.fd is None:
|
||||
self._open()
|
||||
name = uuid.uuid4().hex + ".bin"
|
||||
fd = os.open(name, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600, dir_fd=self.fd)
|
||||
self.names.append(name)
|
||||
with os.fdopen(fd, "wb") as stream:
|
||||
stream.write(payload)
|
||||
return f"{self.relative}/{name}"
|
||||
|
||||
def close(self, *, cleanup):
|
||||
if self.fd is not None:
|
||||
try:
|
||||
if cleanup:
|
||||
for name in self.names:
|
||||
try:
|
||||
os.unlink(name, dir_fd=self.fd)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
finally:
|
||||
os.close(self.fd)
|
||||
self.fd = None
|
||||
if cleanup:
|
||||
# Never recursively remove a shared namespace or another run.
|
||||
for directory in (self.root / self.relative, self.root / "distributed-media-v1"):
|
||||
try:
|
||||
directory.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
class _OwnedExports:
|
||||
"""Best-effort retirement of validated terminal-history files, never a sweep.
|
||||
|
||||
Walk every directory component with openat/O_NOFOLLOW, then pin the directory
|
||||
fd and each regular file's identity. Missing, linked, replaced, or unidentifiable
|
||||
files stay with the endpoint lifecycle. Never use history-supplied path strings.
|
||||
"""
|
||||
def __init__(self, temp_dir, request_id):
|
||||
self.root = Path(temp_dir).absolute() if temp_dir is not None else None
|
||||
self.key = hashlib.sha256(request_id.encode()).hexdigest()
|
||||
self.fd = None
|
||||
self.files = {}
|
||||
|
||||
def _open(self):
|
||||
if self.fd is not None:
|
||||
return True
|
||||
if self.root is None or ".." in self.root.parts:
|
||||
return False
|
||||
flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
|
||||
fd = None
|
||||
try:
|
||||
fd = os.open(self.root.anchor, flags)
|
||||
for component in (*self.root.parts[1:], "distributed-media-v1", self.key):
|
||||
child = os.open(component, flags, dir_fd=fd)
|
||||
os.close(fd)
|
||||
fd = child
|
||||
self.fd, fd = fd, None
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
finally:
|
||||
if fd is not None:
|
||||
os.close(fd)
|
||||
|
||||
def _identity(self, name):
|
||||
info = os.stat(name, dir_fd=self.fd, follow_symlinks=False)
|
||||
if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1:
|
||||
return None
|
||||
return (info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns)
|
||||
|
||||
def identify(self, entry, media):
|
||||
if self.root is None or not isinstance(entry.get("outputs"), dict):
|
||||
return
|
||||
# A failed prompt may contain valid completed exporters alongside missing
|
||||
# or malformed ones. Claim only independently validated declared records.
|
||||
for binding in media["outputs"]:
|
||||
node_id = binding["node_id"]
|
||||
try:
|
||||
_, files = _manifest({"outputs": {node_id: entry["outputs"].get(node_id)}},
|
||||
dict(media, outputs=[binding]))
|
||||
except ValueError:
|
||||
continue
|
||||
if not files or not self._open():
|
||||
continue
|
||||
for artifact in files:
|
||||
name = artifact["id"] + ".bin" # _manifest enforces UUID hex.
|
||||
try:
|
||||
identity = self._identity(name)
|
||||
except OSError:
|
||||
continue
|
||||
if identity is not None:
|
||||
self.files.setdefault(name, identity)
|
||||
|
||||
def close(self, *, cleanup):
|
||||
if self.fd is None:
|
||||
return
|
||||
try:
|
||||
if cleanup:
|
||||
for name, identity in self.files.items():
|
||||
try:
|
||||
if self._identity(name) == identity:
|
||||
os.unlink(name, dir_fd=self.fd)
|
||||
except OSError:
|
||||
pass # Endpoint lifecycle owns inaccessible leftovers.
|
||||
finally:
|
||||
os.close(self.fd)
|
||||
self.fd = None
|
||||
# Deliberately retain directories: they may belong to concurrent runs.
|
||||
|
||||
|
||||
async def _body(response, maximum):
|
||||
length = response.content_length
|
||||
if length is not None and length > maximum:
|
||||
raise ValueError("Local Comfy response exceeds declared byte length")
|
||||
body = bytearray()
|
||||
async for chunk in response.content.iter_chunked(64 * 1024):
|
||||
if len(body) + len(chunk) > maximum:
|
||||
raise ValueError("Local Comfy response exceeds declared byte length")
|
||||
body.extend(chunk)
|
||||
return bytes(body)
|
||||
|
||||
|
||||
class _Execution:
|
||||
def __init__(self, session, url, stage, media, targets, deadline, exports):
|
||||
self.session, self.url, self.stage = session, url, stage
|
||||
self.media, self.targets = media, targets
|
||||
self.deadline = deadline
|
||||
self.run_id = uuid.uuid4().hex
|
||||
self.prompt_id = None
|
||||
self.attempted = False
|
||||
self.terminal = False
|
||||
self.exports = exports
|
||||
self.interruption_blocked = False
|
||||
|
||||
def remaining(self):
|
||||
remaining = self.deadline - asyncio.get_running_loop().time()
|
||||
if remaining <= 0:
|
||||
raise asyncio.TimeoutError("Media execution deadline expired")
|
||||
return remaining
|
||||
|
||||
async def json_request(self, method, path, *, body=None, permit_rejection=False):
|
||||
async with self.session.request(method, self.url + path, json=body,
|
||||
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining())),
|
||||
allow_redirects=False) as response:
|
||||
raw = await _body(response, MAX_HTTP_JSON_BYTES)
|
||||
# Native queue-delete and interrupt routes return an empty 200.
|
||||
if response.status == 200 and not raw and method == "POST" and path in {"/queue", "/interrupt"}:
|
||||
return {}
|
||||
try:
|
||||
value = json.loads(raw)
|
||||
except (ValueError, UnicodeError, RecursionError):
|
||||
if response.status != 200:
|
||||
raise _HTTPError(response.status) from None
|
||||
raise ValueError("Invalid local Comfy JSON response") from None
|
||||
if response.status != 200:
|
||||
# Native /prompt rejects the entire graph with a 400 error.
|
||||
# Other failures may hide acceptance and are never retried.
|
||||
if (permit_rejection and response.status == 400 and isinstance(value, dict)
|
||||
and value.get("error") and "prompt_id" not in value):
|
||||
self.terminal = True
|
||||
raise _HTTPError(response.status)
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError("Invalid local Comfy JSON object")
|
||||
return value
|
||||
|
||||
async def get_json(self, path):
|
||||
for attempt in range(GET_ATTEMPTS):
|
||||
try:
|
||||
return await self.json_request("GET", path)
|
||||
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
||||
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
|
||||
raise
|
||||
if attempt == GET_ATTEMPTS - 1:
|
||||
raise
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
|
||||
|
||||
async def history(self):
|
||||
response = await self.get_json("/history/" + quote(self.prompt_id, safe=""))
|
||||
entry = response.get(self.prompt_id)
|
||||
if entry is None:
|
||||
return None
|
||||
if not isinstance(entry, dict) or not isinstance(entry.get("status"), dict):
|
||||
raise ValueError("Malformed Comfy prompt history")
|
||||
status = entry["status"]
|
||||
messages = status.get("messages", [])
|
||||
if not isinstance(messages, list):
|
||||
raise ValueError("Malformed Comfy history messages")
|
||||
failed = status.get("status_str") in {"error", "interrupted"} or any(
|
||||
isinstance(message, list) and message and message[0] in {"execution_error", "execution_interrupted"}
|
||||
for message in messages)
|
||||
if failed or status.get("completed") is True:
|
||||
self.terminal = True
|
||||
self.exports.identify(entry, self.media)
|
||||
if failed or status.get("status_str") != "success":
|
||||
raise RuntimeError("Local Comfy prompt execution failed")
|
||||
return entry
|
||||
return None
|
||||
|
||||
async def submit_and_wait(self, workflow):
|
||||
body = {"prompt": workflow, "client_id": self.run_id,
|
||||
"extra_data": {"distributed_media_run_id": self.run_id},
|
||||
"partial_execution_targets": list(self.targets)}
|
||||
self.remaining() # CPU staging must not spend past the absolute deadline.
|
||||
self.attempted = True
|
||||
response = await self.json_request("POST", "/prompt", body=body, permit_rejection=True)
|
||||
self.prompt_id = _scope_id(response.get("prompt_id"))
|
||||
if response.get("node_errors"):
|
||||
raise RuntimeError("Local Comfy accepted an incomplete set of execution targets")
|
||||
while True:
|
||||
entry = await self.history()
|
||||
if entry is not None:
|
||||
return entry
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
||||
|
||||
async def settle(self):
|
||||
"""Bounded recovery; caller shields this from task cancellation."""
|
||||
deleted = False
|
||||
while True:
|
||||
if self.prompt_id:
|
||||
try:
|
||||
await self.history()
|
||||
except RuntimeError:
|
||||
if not self.terminal:
|
||||
raise
|
||||
if self.terminal:
|
||||
return
|
||||
queue = await self.get_json("/queue")
|
||||
groups = []
|
||||
for key in ("queue_running", "queue_pending"):
|
||||
rows = queue.get(key)
|
||||
if not isinstance(rows, list):
|
||||
raise ValueError("Malformed Comfy queue")
|
||||
group = {}
|
||||
for row in rows:
|
||||
if not isinstance(row, list) or len(row) < 4 or not isinstance(row[3], dict):
|
||||
raise ValueError("Malformed Comfy queue entry")
|
||||
pid = _scope_id(row[1])
|
||||
if pid in group:
|
||||
raise ValueError("Duplicate Comfy queue identity")
|
||||
group[pid] = row[3]
|
||||
groups.append(group)
|
||||
running, pending = groups
|
||||
if self.prompt_id is None:
|
||||
own = [pid for group in groups for pid, extra in group.items()
|
||||
if extra.get("distributed_media_run_id") == self.run_id
|
||||
and extra.get("client_id") == self.run_id]
|
||||
if len(own) > 1:
|
||||
raise ValueError("Ambiguous Comfy queue ownership")
|
||||
if own:
|
||||
self.prompt_id = own[0]
|
||||
if self.prompt_id in pending and not deleted:
|
||||
await self.json_request("POST", "/queue", body={"delete": [self.prompt_id]})
|
||||
deleted = True
|
||||
elif self.prompt_id in running:
|
||||
# Pinned ComfyUI server.py checks a snapshot then globally
|
||||
# interrupts, outside the queue mutex. Supplying prompt_id does
|
||||
# NOT close that race. No safe capability negotiation exists on
|
||||
# this route; do not invoke it, even after our own queue read.
|
||||
self.interruption_blocked = True
|
||||
raise RuntimeError("Native prompt-scoped interruption unavailable")
|
||||
elif deleted and self.prompt_id not in running and self.prompt_id not in pending:
|
||||
self.terminal = True # queue deletion confirmed by exact read-back
|
||||
return
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
||||
|
||||
async def fetch(self, artifact):
|
||||
transport = artifact["transport"]
|
||||
params = {"filename": transport["filename"], "subfolder": transport["subfolder"], "type": "temp"}
|
||||
for attempt in range(GET_ATTEMPTS):
|
||||
try:
|
||||
async with self.session.get(self.url + "/view", params=params, allow_redirects=False,
|
||||
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining()))) as response:
|
||||
if response.status != 200:
|
||||
raise _HTTPError(response.status)
|
||||
payload = await _body(response, artifact["byte_length"])
|
||||
break
|
||||
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
||||
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
|
||||
raise
|
||||
if attempt == GET_ATTEMPTS - 1:
|
||||
raise
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
|
||||
if len(payload) != artifact["byte_length"]:
|
||||
raise ValueError("Media output byte length mismatch")
|
||||
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
|
||||
raise ValueError("Media output checksum mismatch")
|
||||
return payload
|
||||
|
||||
|
||||
def _manifest(entry, media):
|
||||
outputs = entry.get("outputs")
|
||||
if not isinstance(outputs, dict):
|
||||
raise ValueError("Missing media history outputs")
|
||||
expected_ids = {binding["node_id"] for binding in media["outputs"]}
|
||||
if any(node_id not in expected_ids and isinstance(value, dict) and "media_exports" in value
|
||||
for node_id, value in outputs.items()):
|
||||
raise ValueError("Unexpected media export history")
|
||||
response = {"version": VERSION, "profile": PROFILE, "request_id": media["request_id"],
|
||||
"bindings": [], "artifacts": [], "items": []}
|
||||
files = []
|
||||
subfolder = "distributed-media-v1/" + hashlib.sha256(media["request_id"].encode()).hexdigest()
|
||||
for binding in media["outputs"]:
|
||||
node = outputs.get(binding["node_id"], {})
|
||||
records = node.get("media_exports") if isinstance(node, dict) else None
|
||||
if not isinstance(records, list) or len(records) != 1:
|
||||
raise ValueError("Missing or duplicate media export result")
|
||||
record = records[0]
|
||||
if (not isinstance(record, dict) or set(record) != {"request_id", "binding", "artifacts", "items"}
|
||||
or record["request_id"] != media["request_id"]):
|
||||
raise ValueError("Media export request identity mismatch")
|
||||
if not isinstance(record["artifacts"], list) or len(record["artifacts"]) > MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid exported artifacts count")
|
||||
converted = []
|
||||
for artifact in record["artifacts"]:
|
||||
if not isinstance(artifact, dict) or not isinstance(artifact.get("transport"), dict):
|
||||
raise ValueError("Invalid exported artifact")
|
||||
transport = artifact["transport"]
|
||||
aid = artifact.get("id")
|
||||
if (not isinstance(aid, str) or not re.fullmatch(r"[a-f0-9]{32}", aid)
|
||||
or transport != {"type": "comfy-file", "filename": aid + ".bin",
|
||||
"subfolder": subfolder, "folder_type": "temp"}):
|
||||
raise ValueError("Unsafe or unowned exported media file")
|
||||
converted.append(dict(artifact, transport=dict(_VALIDATION_TRANSPORT)))
|
||||
partial = dict(response, bindings=[record["binding"]], artifacts=converted, items=record["items"])
|
||||
validate_response(partial, request_id=media["request_id"], bindings=[binding])
|
||||
response["bindings"].append(copy.deepcopy(record["binding"]))
|
||||
response["artifacts"].extend(converted)
|
||||
response["items"].extend(copy.deepcopy(record["items"]))
|
||||
files.extend(copy.deepcopy(record["artifacts"]))
|
||||
# Validate the whole manifest BEFORE the first /view or publisher callback.
|
||||
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
|
||||
return response, files
|
||||
|
||||
|
||||
def _publication_plan(response, files, publish):
|
||||
skeleton = dict(response, artifacts=[dict(a, transport={"type": "inline", "data_base64": ""}) for a in files])
|
||||
base_size = len(_json_bytes({"media": skeleton}))
|
||||
sizes = {a["id"]: 4 * ((a["byte_length"] + 2) // 3) for a in files}
|
||||
predicted = base_size + sum(sizes.values())
|
||||
published = set()
|
||||
if predicted > MAX_RESPONSE_JSON_BYTES and publish is None:
|
||||
raise ValueError("Media response exceeds inline JSON limit; a publisher is required")
|
||||
# Reserve enough space for a maximum-size escaped URL, including metadata.
|
||||
# Small remaining inline values preserve their exact bytes and ordering.
|
||||
for aid in sorted(sizes, key=sizes.get, reverse=True):
|
||||
if predicted <= MAX_RESPONSE_JSON_BYTES:
|
||||
break
|
||||
if sizes[aid] > 65_536:
|
||||
published.add(aid)
|
||||
predicted -= sizes[aid] - 65_536
|
||||
if predicted > MAX_RESPONSE_JSON_BYTES:
|
||||
raise ValueError("Media response metadata exceeds JSON limit")
|
||||
return published
|
||||
|
||||
|
||||
async def _run(execution, workflow, legacy, download, publish):
|
||||
media = execution.media
|
||||
imports = {}
|
||||
for artifact in media["artifacts"]:
|
||||
execution.remaining()
|
||||
payload = await resolve_artifact(artifact, download=download)
|
||||
execution.remaining()
|
||||
filename = execution.stage.write(payload)
|
||||
node_id = "distributed-media-import:" + uuid.uuid4().hex
|
||||
while node_id in workflow:
|
||||
node_id = "distributed-media-import:" + uuid.uuid4().hex
|
||||
metadata = {key: value for key, value in artifact.items() if key != "transport"}
|
||||
workflow[node_id] = {"class_type": "MediaImportV1", "inputs": {
|
||||
"artifact_json": json.dumps(metadata, allow_nan=False), "filename": filename}}
|
||||
imports[artifact["id"]] = [node_id, _IMPORT_SLOTS[artifact["kind"]]]
|
||||
for ref in media["inputs"]:
|
||||
workflow[ref["node_id"]]["inputs"][ref["input_name"]] = list(imports[ref["artifact_id"]])
|
||||
for encoded, size, matches in legacy:
|
||||
execution.remaining()
|
||||
try:
|
||||
payload = base64.b64decode(encoded, validate=True)
|
||||
except (ValueError, binascii.Error):
|
||||
raise ValueError("Invalid legacy image base64") from None
|
||||
if len(payload) != size:
|
||||
raise ValueError("Legacy image byte length mismatch")
|
||||
execution.remaining()
|
||||
filename = execution.stage.write(payload)
|
||||
for node_id in matches:
|
||||
workflow[node_id]["inputs"]["image"] = filename
|
||||
# Replacements can disconnect file loaders whose original endpoint-local
|
||||
# paths are unavailable. Keep exactly the post-rewrite exporter closure;
|
||||
# shared ancestors still needed by another branch must survive.
|
||||
reachable, pending = set(), list(execution.targets)
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id in reachable:
|
||||
continue
|
||||
reachable.add(node_id)
|
||||
pending.extend(_link(value, workflow) for value in workflow[node_id]["inputs"].values()
|
||||
if isinstance(value, list))
|
||||
workflow = {node_id: node for node_id, node in workflow.items() if node_id in reachable}
|
||||
entry = await execution.submit_and_wait(workflow)
|
||||
response, files = _manifest(entry, media)
|
||||
published = _publication_plan(response, files, publish)
|
||||
artifacts = []
|
||||
for artifact in files:
|
||||
payload = await execution.fetch(artifact)
|
||||
transport = None
|
||||
if artifact["id"] in published:
|
||||
metadata = {key: copy.deepcopy(value) for key, value in artifact.items() if key != "transport"}
|
||||
transport = await publish(metadata, payload)
|
||||
if not isinstance(transport, dict) or transport.get("type") != "url":
|
||||
raise ValueError("Publisher must return an expiring HTTPS URL transport")
|
||||
artifacts.append(artifact_from_payload(artifact["id"], artifact["kind"], artifact["descriptor"],
|
||||
payload, transport=transport))
|
||||
response["artifacts"] = artifacts
|
||||
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
|
||||
result = {"media": response}
|
||||
if len(_json_bytes(result)) > MAX_RESPONSE_JSON_BYTES:
|
||||
raise ValueError("Media response exceeds JSON byte limit")
|
||||
execution.remaining()
|
||||
return result
|
||||
|
||||
|
||||
def _release_uncertain_work(task):
|
||||
_UNCERTAIN_WORK.discard(task)
|
||||
try:
|
||||
task.result()
|
||||
except BaseException:
|
||||
pass
|
||||
|
||||
|
||||
async def _drain_work(work):
|
||||
"""Cancel once; a misbehaving callback cannot make teardown unbounded."""
|
||||
if not work.done():
|
||||
work.cancel()
|
||||
deadline = asyncio.get_running_loop().time() + DRAIN_GRACE_SECONDS
|
||||
while not work.done():
|
||||
remaining = deadline - asyncio.get_running_loop().time()
|
||||
if remaining <= 0:
|
||||
# Keep a strong reference until it exits, but never delete its files
|
||||
# asynchronously: uncertain ownership belongs to endpoint recovery.
|
||||
_UNCERTAIN_WORK.add(work)
|
||||
work.add_done_callback(_release_uncertain_work)
|
||||
return False
|
||||
try:
|
||||
# Unlike wait_for, wait's timeout does not wait for cancellation.
|
||||
await asyncio.wait({work}, timeout=remaining)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
_release_uncertain_work(work)
|
||||
return True
|
||||
|
||||
|
||||
async def execute_media_request(input_data, *, comfy_url, input_dir, temp_dir=None,
|
||||
timeout_seconds=1800, session=None, download=None, publish=None) -> dict:
|
||||
"""Validate, stage, submit once, collect exact media, and retire owned inputs.
|
||||
|
||||
``timeout_seconds`` covers staging, HTTP, polling, and publishing. A separate
|
||||
bounded cancellation grace is allowed ONLY to establish the local prompt is
|
||||
terminal. On uncertain acceptance/termination, raise MediaExecutionUncertain
|
||||
and retain input files rather than racing an executing prompt. Never retry an
|
||||
ambiguous submission, even when the exception is a timeout.
|
||||
Async callback teardown has its own bounded grace; non-cooperative callbacks
|
||||
also raise MediaExecutionUncertain and retain files, even after terminal history.
|
||||
|
||||
``temp_dir`` is a trusted endpoint setting for this ComfyUI's local temp root,
|
||||
never a request-supplied path. Only validated exporter UUID.bin files beneath
|
||||
distributed-media-v1/sha256(request_id) are retired, after all collection and
|
||||
publisher work has unwound. Unidentified files remain lifecycle-owned. With
|
||||
``temp_dir=None`` all exporter files remain lifecycle-owned as before.
|
||||
|
||||
The pinned native /interrupt route has a check/global-interrupt race. Running
|
||||
prompts therefore fail closed with MediaExecutionUncertain, without posting
|
||||
/interrupt. A prompt-scoped request body alone is not a safety guarantee.
|
||||
"""
|
||||
if (type(timeout_seconds) not in (int, float) or not math.isfinite(timeout_seconds)
|
||||
or timeout_seconds <= 0):
|
||||
raise ValueError("Invalid media execution timeout")
|
||||
if not isinstance(comfy_url, str):
|
||||
raise ValueError("Invalid local Comfy URL")
|
||||
try:
|
||||
parsed = urlsplit(comfy_url)
|
||||
if (parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username
|
||||
or parsed.password or parsed.query or parsed.fragment):
|
||||
raise ValueError()
|
||||
parsed.port
|
||||
except ValueError:
|
||||
raise ValueError("Invalid local Comfy URL") from None
|
||||
deadline = asyncio.get_running_loop().time() + timeout_seconds
|
||||
workflow, media, legacy, targets = _validate_request(input_data)
|
||||
stage = _Staging(input_dir, media["request_id"])
|
||||
exports = _OwnedExports(temp_dir, media["request_id"])
|
||||
owned_session = session is None
|
||||
if owned_session:
|
||||
session = aiohttp.ClientSession()
|
||||
execution = _Execution(session, comfy_url.rstrip("/"), stage, media, targets, deadline, exports)
|
||||
work = None
|
||||
try:
|
||||
remaining = execution.remaining()
|
||||
work = asyncio.create_task(_run(execution, workflow, legacy, download, publish))
|
||||
return await asyncio.wait_for(asyncio.shield(work), remaining)
|
||||
except BaseException as exc:
|
||||
# Retire files only after readers unwind. Cooperative async callbacks
|
||||
# normally drain immediately; an uncooperative one requires quarantine.
|
||||
if work is not None and not await _drain_work(work):
|
||||
raise MediaExecutionUncertain(
|
||||
"Media callback teardown unconfirmed; owned files retained; "
|
||||
"quarantine this endpoint before reuse; do not resubmit automatically"
|
||||
) from exc
|
||||
if execution.attempted and not execution.terminal:
|
||||
execution.deadline = asyncio.get_running_loop().time() + CANCEL_GRACE_SECONDS
|
||||
recovery = asyncio.create_task(asyncio.wait_for(execution.settle(), CANCEL_GRACE_SECONDS))
|
||||
# Shield the bounded recovery even from repeated task.cancel() calls.
|
||||
while not recovery.done():
|
||||
try:
|
||||
await asyncio.shield(recovery)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
except Exception:
|
||||
break
|
||||
try:
|
||||
recovery.result()
|
||||
except BaseException:
|
||||
pass
|
||||
if not execution.terminal:
|
||||
raise MediaExecutionUncertain(
|
||||
"Local Comfy prompt terminal state unconfirmed; owned staged files retained "
|
||||
f"under {stage.relative}; prompt_id={execution.prompt_id or 'unknown'}; "
|
||||
"do not resubmit this request automatically"
|
||||
+ ("; native prompt-scoped interruption unavailable (non-atomic /interrupt route)"
|
||||
if execution.interruption_blocked else "")
|
||||
) from exc
|
||||
raise
|
||||
finally:
|
||||
try:
|
||||
try:
|
||||
exports.close(cleanup=(work is None or work.done()) and execution.terminal)
|
||||
finally:
|
||||
stage.close(cleanup=(work is None or work.done()) and
|
||||
(not execution.attempted or execution.terminal))
|
||||
finally:
|
||||
if owned_session:
|
||||
await session.close()
|
||||
@@ -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
|
||||
@@ -0,0 +1,175 @@
|
||||
"""Internal ComfyUI boundary nodes for the standalone media-v1 handler."""
|
||||
import hashlib
|
||||
import json
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
def _scalar(value):
|
||||
if isinstance(value, list):
|
||||
if len(value) != 1:
|
||||
raise ValueError("Media bridge control must have exactly one value")
|
||||
return value[0]
|
||||
return value
|
||||
|
||||
|
||||
def _json(text):
|
||||
if not isinstance(text, str) or len(text) > 256 * 1024:
|
||||
raise ValueError("Invalid media bridge JSON")
|
||||
try:
|
||||
value = json.loads(text)
|
||||
except (ValueError, TypeError, RecursionError):
|
||||
raise ValueError("Invalid media bridge JSON") from None
|
||||
validate_json(value)
|
||||
return value
|
||||
|
||||
|
||||
def _cpu(kind, value):
|
||||
# Runtime transfer only, never sampling or device discovery. Pure codecs
|
||||
# deliberately reject non-CPU storage rather than making hidden transfers.
|
||||
if kind in ("IMAGE", "MASK"):
|
||||
return value.detach().cpu()
|
||||
if kind == "AUDIO":
|
||||
return {"waveform": value["waveform"].detach().cpu(), "sample_rate": value["sample_rate"]}
|
||||
return 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)
|
||||
return encode_payload(kind, _cpu(kind, value))
|
||||
|
||||
|
||||
def decode_value(kind, descriptor, payload):
|
||||
if kind == "VIDEO":
|
||||
from .video_codec import decode_video
|
||||
return decode_video(descriptor, payload)
|
||||
return decode_payload(kind, descriptor, payload)
|
||||
|
||||
|
||||
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()
|
||||
subfolder = f"distributed-media-v1/{request_key}"
|
||||
root = Path(folder_paths.get_temp_directory()).resolve()
|
||||
directory = (root / subfolder).resolve()
|
||||
if not directory.is_relative_to(root):
|
||||
raise ValueError("Unsafe media output directory")
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
written, artifacts, items = [], [], []
|
||||
total = 0
|
||||
try:
|
||||
for index, value in enumerate(values):
|
||||
descriptor, payload = encode_value(binding["kind"], value,
|
||||
transfer_format=fmt, transfer_quality=quality)
|
||||
validate_json(descriptor)
|
||||
total += len(payload)
|
||||
if total > MAX_TOTAL_BYTES:
|
||||
raise ValueError("Media output exceeds total byte limit")
|
||||
aid = uuid.uuid4().hex
|
||||
filename = aid + ".bin"
|
||||
path = directory / filename
|
||||
with path.open("xb") as stream:
|
||||
written.append(path)
|
||||
stream.write(payload)
|
||||
artifacts.append({"id": aid, "kind": binding["kind"], "descriptor": descriptor,
|
||||
"byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(),
|
||||
"transport": {"type": "comfy-file", "filename": filename,
|
||||
"subfolder": subfolder, "folder_type": "temp"}})
|
||||
items.append({"binding_id": binding["binding_id"], "item_index": index, "artifact_id": aid})
|
||||
except BaseException:
|
||||
for path in written:
|
||||
path.unlink(missing_ok=True)
|
||||
raise
|
||||
return {"request_id": request_id, "binding": dict(binding, count=len(values)),
|
||||
"artifacts": artifacts, "items": items}
|
||||
|
||||
|
||||
def import_value(artifact, filename):
|
||||
"""Read a handler-staged, checksummed file within native media roots only."""
|
||||
fields = {"id", "kind", "descriptor", "byte_length", "sha256"}
|
||||
if not isinstance(artifact, dict) or set(artifact) != fields:
|
||||
raise ValueError("Invalid staged media metadata")
|
||||
identifier(artifact["id"])
|
||||
if artifact["kind"] not in SOCKET_KINDS.values():
|
||||
raise ValueError("Invalid staged media kind")
|
||||
size = bounded_int(artifact["byte_length"], MAX_ARTIFACT_BYTES)
|
||||
if not isinstance(filename, str):
|
||||
raise ValueError("Invalid staged media filename")
|
||||
roots = [Path(getter()).resolve() for getter in (
|
||||
folder_paths.get_input_directory, folder_paths.get_output_directory, folder_paths.get_temp_directory)]
|
||||
path = Path(folder_paths.get_annotated_filepath(filename)).resolve()
|
||||
if not any(path.is_relative_to(root) for root in roots) or not path.is_file():
|
||||
raise ValueError("Staged media file is outside allowed roots or unavailable")
|
||||
if path.stat().st_size != size:
|
||||
raise ValueError("Staged media file length mismatch")
|
||||
with path.open("rb") as stream:
|
||||
payload = stream.read(size + 1)
|
||||
if len(payload) != size or hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
|
||||
raise ValueError("Staged media file failed integrity validation")
|
||||
return decode_value(artifact["kind"], artifact["descriptor"], payload)
|
||||
|
||||
|
||||
class MediaImportV1(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(node_id="MediaImportV1", category="distributed/internal", is_deprecated=True,
|
||||
inputs=[io.String.Input("artifact_json"), io.String.Input("filename")],
|
||||
outputs=[io.Image.Output(), io.Audio.Output(), io.Mask.Output(), io.String.Output(), io.Video.Output()])
|
||||
|
||||
@classmethod
|
||||
def execute(cls, artifact_json, filename):
|
||||
artifact = _json(artifact_json)
|
||||
value = import_value(artifact, filename)
|
||||
return io.NodeOutput(*(value if kind == artifact["kind"] else None for kind in SOCKET_KINDS.values()))
|
||||
|
||||
|
||||
class MediaExportV1(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(node_id="MediaExportV1", category="distributed/internal", is_deprecated=True,
|
||||
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),
|
||||
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,
|
||||
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"]],
|
||||
transfer_format=_scalar(transfer_format),
|
||||
transfer_quality=_scalar(transfer_quality))
|
||||
return io.NodeOutput(ui={"media_exports": [exported]})
|
||||
@@ -0,0 +1,482 @@
|
||||
"""Lossless transport of pinned native ComfyUI VIDEO representations.
|
||||
|
||||
No save_to/get_components fallback is allowed: an unknown class must not silently
|
||||
transcode. Tensor components use the shared CPU codec. File bytes are opaque and
|
||||
retain signed trim/crop parameters; decoded files own independent BytesIOs.
|
||||
VideoFromList's native constructor flattens its inputs, so wire segments are a
|
||||
bounded, flat list of exact VideoFromFile objects, never recursive playlists.
|
||||
|
||||
Offsets are relative to the containing VIDEO payload. Canonical part order is
|
||||
images/alpha/audio, or segments/complete_audio. Every byte is covered exactly
|
||||
once. All descriptor structure and tensor sizes are preflighted before decoding
|
||||
any media storage. File containers are not parsed until a native consumer reads
|
||||
them; this adapter does not attempt to validate or repair container codecs.
|
||||
"""
|
||||
from fractions import Fraction
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
import stat
|
||||
|
||||
from .codec import MAX_PAYLOAD_BYTES, decode_payload, encode_payload
|
||||
|
||||
|
||||
MAX_SEGMENTS = 64
|
||||
MAX_METADATA_BYTES = 64 * 1024
|
||||
MAX_METADATA_DEPTH = 16
|
||||
_ENCODING = "native-video-v1"
|
||||
_BASE_FIELDS = {"encoding", "representation", "byte_length"}
|
||||
_COMPONENT_FIELDS = {"images", "alpha", "audio", "frame_rate", "metadata", "bit_depth", "color_space"}
|
||||
_FILE_FIELDS = {"data", "start_time", "duration", "crop"}
|
||||
_SEGMENT_FIELDS = {"segments", "complete_audio"}
|
||||
_TENSOR_BYTES = {"float16": 2, "bfloat16": 2, "float32": 4, "float64": 8}
|
||||
|
||||
|
||||
def _limit(max_bytes):
|
||||
if type(max_bytes) is not int or max_bytes < 0:
|
||||
raise ValueError("VIDEO max_bytes must be a nonnegative integer")
|
||||
return min(max_bytes, MAX_PAYLOAD_BYTES)
|
||||
|
||||
|
||||
def _native_types():
|
||||
# Lazy imports keep non-VIDEO users independent of ComfyUI and PyAV.
|
||||
from comfy_api.latest._input_impl.video_types import (
|
||||
VideoFromComponents, VideoFromFile, VideoFromList,
|
||||
)
|
||||
from comfy_api.latest._util import VideoComponents
|
||||
return VideoFromComponents, VideoFromFile, VideoFromList, VideoComponents
|
||||
|
||||
|
||||
def _fields(value, fields, label):
|
||||
if type(value) is not dict or set(value) != fields:
|
||||
raise ValueError(f"Invalid VIDEO {label} fields")
|
||||
|
||||
|
||||
def _integer(value, *, positive=False):
|
||||
if type(value) is not int or value < (1 if positive else 0):
|
||||
raise ValueError("VIDEO length, offset and dimensions must be integers in range")
|
||||
return value
|
||||
|
||||
|
||||
def _number(value):
|
||||
if type(value) not in (int, float) or (type(value) is float and not math.isfinite(value)):
|
||||
raise ValueError("VIDEO trim values must be finite numbers")
|
||||
return value
|
||||
|
||||
|
||||
def _crop(value, *, wire=False):
|
||||
if value is None:
|
||||
return None
|
||||
if type(value) is not (list if wire else tuple) or len(value) != 4:
|
||||
raise ValueError("VIDEO crop must contain four integers")
|
||||
if any(type(part) is not int for part in value):
|
||||
raise ValueError("VIDEO crop must contain four integers")
|
||||
# Native crop normalization is a consumer concern, not a transport operation.
|
||||
return list(value) if not wire else tuple(value)
|
||||
|
||||
|
||||
def _frame_rate(value):
|
||||
_fields(value, {"numerator", "denominator"}, "frame_rate")
|
||||
_integer(value["numerator"], positive=True)
|
||||
_integer(value["denominator"], positive=True)
|
||||
|
||||
|
||||
def _encoding_options(bit_depth, color_space):
|
||||
_integer(bit_depth, positive=True)
|
||||
if type(color_space) is not str or color_space not in ("sRGB", "HDR", "HDR PQ"):
|
||||
raise ValueError("Unsupported native VIDEO color_space")
|
||||
|
||||
|
||||
def _metadata(value):
|
||||
"""Validate bounded JSON before allocating its serialized form; return a copy."""
|
||||
if value is None:
|
||||
return None
|
||||
if type(value) is not dict:
|
||||
raise ValueError("VIDEO metadata must be a JSON object or null")
|
||||
remaining = MAX_METADATA_BYTES
|
||||
|
||||
def charge(size):
|
||||
nonlocal remaining
|
||||
remaining -= size
|
||||
if remaining < 0:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
|
||||
def visit(item, depth):
|
||||
if depth > MAX_METADATA_DEPTH:
|
||||
raise ValueError("VIDEO metadata exceeds depth limit")
|
||||
kind = type(item)
|
||||
if kind in (dict, list):
|
||||
# Each member costs at least one JSON byte. Also bounds iteration
|
||||
# before walking potentially huge containers or cyclic structures.
|
||||
if len(item) > remaining:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
charge(2 + max(0, len(item) - 1))
|
||||
if kind is dict:
|
||||
for key, child in item.items():
|
||||
if type(key) is not str:
|
||||
raise ValueError("VIDEO metadata keys must be strings")
|
||||
visit(key, depth + 1)
|
||||
charge(1)
|
||||
visit(child, depth + 1)
|
||||
else:
|
||||
for child in item:
|
||||
visit(child, depth + 1)
|
||||
elif kind in (str, int, float, bool) or item is None:
|
||||
if kind is str and len(item) > remaining:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
# Avoid converting an unbounded Python integer to decimal first.
|
||||
if kind is int and item.bit_length() > MAX_METADATA_BYTES * 4:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
if kind is float and not math.isfinite(item):
|
||||
raise ValueError("VIDEO metadata numbers must be finite")
|
||||
try:
|
||||
charge(len(json.dumps(item, ensure_ascii=True, allow_nan=False)))
|
||||
except (ValueError, OverflowError):
|
||||
raise ValueError("Invalid or oversized VIDEO metadata scalar") from None
|
||||
else:
|
||||
raise ValueError("VIDEO metadata must contain only JSON values")
|
||||
|
||||
visit(value, 0)
|
||||
return json.loads(json.dumps(value, ensure_ascii=True, allow_nan=False, separators=(",", ":")))
|
||||
|
||||
|
||||
def _attrs(value, names):
|
||||
attributes = vars(value)
|
||||
if not all(name in attributes for name in names):
|
||||
raise ValueError("Native VIDEO attributes do not match the pinned adapter")
|
||||
return [attributes[name] for name in names]
|
||||
|
||||
|
||||
def _roots(allowed_roots):
|
||||
if allowed_roots is None:
|
||||
import folder_paths
|
||||
allowed_roots = [folder_paths.get_input_directory(), folder_paths.get_output_directory(),
|
||||
folder_paths.get_temp_directory()]
|
||||
if type(allowed_roots) not in (list, tuple):
|
||||
raise ValueError("VIDEO allowed_roots must be a list or tuple")
|
||||
try:
|
||||
roots = [Path(root).resolve() for root in allowed_roots]
|
||||
except (OSError, TypeError, ValueError, RuntimeError):
|
||||
raise ValueError("Invalid VIDEO allowed roots") from None
|
||||
return roots
|
||||
|
||||
|
||||
def _read_path(value, limit, allowed_roots):
|
||||
"""Resolve authorization, reject special files, and bound all reads.
|
||||
|
||||
Open each resolved path component with NOFOLLOW using directory fds, so an
|
||||
intermediate symlink swapped in after resolve() cannot bypass authorization.
|
||||
NONBLOCK prevents an attacker swapping in a FIFO from hanging the worker.
|
||||
"""
|
||||
roots = _roots(allowed_roots)
|
||||
directory_fd = file_fd = None
|
||||
try:
|
||||
path = Path(value).resolve(strict=True)
|
||||
if not any(path.is_relative_to(root) for root in roots):
|
||||
raise ValueError("VIDEO file is outside allowed roots")
|
||||
directory_flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
|
||||
directory_fd = os.open(path.anchor, directory_flags)
|
||||
for part in path.parts[1:-1]:
|
||||
next_fd = os.open(part, directory_flags, dir_fd=directory_fd)
|
||||
os.close(directory_fd)
|
||||
directory_fd = next_fd
|
||||
file_fd = os.open(path.name, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=directory_fd)
|
||||
info = os.fstat(file_fd)
|
||||
if not stat.S_ISREG(info.st_mode):
|
||||
raise ValueError("VIDEO source must be a regular file")
|
||||
if info.st_size > limit:
|
||||
raise ValueError("VIDEO file exceeds byte limit")
|
||||
# Read at most the stat size plus a sentinel; reject growth or truncation.
|
||||
chunks = []
|
||||
count = 0
|
||||
while count <= info.st_size:
|
||||
chunk = os.read(file_fd, min(1024 * 1024, info.st_size + 1 - count))
|
||||
if not chunk:
|
||||
break
|
||||
chunks.append(chunk)
|
||||
count += len(chunk)
|
||||
if count != info.st_size:
|
||||
raise ValueError("VIDEO source changed size during bounded read")
|
||||
return b"".join(chunks)
|
||||
except (OSError, RuntimeError):
|
||||
raise ValueError("VIDEO file cannot be read safely") from None
|
||||
finally:
|
||||
if file_fd is not None:
|
||||
os.close(file_fd)
|
||||
if directory_fd is not None:
|
||||
os.close(directory_fd)
|
||||
|
||||
|
||||
def _file_bytes(value, limit, allowed_roots):
|
||||
if type(value) is io.BytesIO:
|
||||
try:
|
||||
with value.getbuffer() as view:
|
||||
if view.nbytes > limit:
|
||||
raise ValueError("VIDEO file exceeds byte limit")
|
||||
return bytes(view)
|
||||
except ValueError:
|
||||
raise ValueError("VIDEO BytesIO is closed or exceeds byte limit") from None
|
||||
if isinstance(value, (str, os.PathLike)):
|
||||
return _read_path(value, limit, allowed_roots)
|
||||
raise ValueError("VIDEO file must be a local path or exact BytesIO")
|
||||
|
||||
|
||||
def encode_video(value, *, max_bytes=MAX_PAYLOAD_BYTES, allowed_roots=None):
|
||||
"""Encode an exact native VIDEO without opening/transcoding its container.
|
||||
|
||||
Paths are authorized only under allowed_roots (when supplied), otherwise
|
||||
under ComfyUI input/output/temp roots. Budgets may lower, never raise, the
|
||||
shared maximum. Unsupported native state raises ValueError, not a fallback.
|
||||
Non-CPU tensors are rejected; runtime staging belongs to the bridge boundary.
|
||||
"""
|
||||
return _encode_video(value, max_bytes=max_bytes, allowed_roots=allowed_roots, stage_cpu=False)
|
||||
|
||||
|
||||
def encode_video_cpu(value, *, max_bytes=MAX_PAYLOAD_BYTES, allowed_roots=None):
|
||||
"""Bridge-only boundary: detach/stage native tensors, then losslessly encode.
|
||||
|
||||
Share the exact-type adapter with encode_video rather than reconstructing
|
||||
native videos, calling get_components, or duplicating private attributes.
|
||||
Images, alpha and both audio representations move without dtype conversion;
|
||||
metadata/options and opaque file bytes/trim/crop take the unchanged path.
|
||||
"""
|
||||
return _encode_video(value, max_bytes=max_bytes, allowed_roots=allowed_roots, stage_cpu=True)
|
||||
|
||||
|
||||
def _stage_tensor_cpu(value):
|
||||
import torch
|
||||
if not isinstance(value, torch.Tensor):
|
||||
raise ValueError("Tensor value must be a torch.Tensor")
|
||||
return value.detach().cpu()
|
||||
|
||||
|
||||
def _encode_video(value, *, max_bytes, allowed_roots, stage_cpu):
|
||||
limit = _limit(max_bytes)
|
||||
components_class, file_class, list_class, components_type = _native_types()
|
||||
descriptor = {"encoding": _ENCODING}
|
||||
chunks = []
|
||||
size = 0
|
||||
|
||||
def add(payload, nested=None):
|
||||
nonlocal size
|
||||
if len(payload) > limit - size:
|
||||
raise ValueError("VIDEO payload exceeds byte limit")
|
||||
slot = {"offset": size, "length": len(payload)}
|
||||
if nested is not None:
|
||||
slot["descriptor"] = nested
|
||||
chunks.append(payload)
|
||||
size += len(payload)
|
||||
return slot
|
||||
|
||||
def tensor(kind, tensor_value):
|
||||
if stage_cpu:
|
||||
if kind == "AUDIO":
|
||||
if not isinstance(tensor_value, dict) or set(tensor_value) != {"waveform", "sample_rate"}:
|
||||
raise ValueError("AUDIO value must contain waveform and sample_rate")
|
||||
tensor_value = {**tensor_value, "waveform": _stage_tensor_cpu(tensor_value["waveform"])}
|
||||
else:
|
||||
tensor_value = _stage_tensor_cpu(tensor_value)
|
||||
nested, payload = encode_payload(kind, tensor_value, max_bytes=limit - size)
|
||||
return add(payload, nested)
|
||||
|
||||
if type(value) is components_class:
|
||||
component, bit_depth, color_space = _attrs(value, [
|
||||
"_VideoFromComponents__components", "_VideoFromComponents__bit_depth",
|
||||
"_VideoFromComponents__color_space",
|
||||
])
|
||||
if type(component) is not components_type:
|
||||
raise ValueError("Unsupported native VIDEO components class")
|
||||
images, fps, audio, metadata, alpha = _attrs(component, [
|
||||
"images", "frame_rate", "audio", "metadata", "alpha",
|
||||
])
|
||||
if type(fps) is not Fraction or fps <= 0:
|
||||
raise ValueError("VIDEO frame_rate must be a positive Fraction")
|
||||
_encoding_options(bit_depth, color_space)
|
||||
descriptor.update(representation="components", bit_depth=bit_depth, color_space=color_space,
|
||||
frame_rate={"numerator": fps.numerator, "denominator": fps.denominator},
|
||||
metadata=_metadata(metadata))
|
||||
descriptor["images"] = tensor("IMAGE", images)
|
||||
# Native get_components() returns BHW1 alpha; MASK nodes use BHW.
|
||||
# Select an existing lossless tensor layout without squeezing/converting.
|
||||
alpha_kind = "IMAGE" if getattr(alpha, "ndim", None) == 4 else "MASK"
|
||||
descriptor["alpha"] = tensor(alpha_kind, alpha) if alpha is not None else None
|
||||
if alpha is not None:
|
||||
_alpha_header(descriptor["alpha"], descriptor["images"])
|
||||
descriptor["audio"] = tensor("AUDIO", audio) if audio is not None else None
|
||||
elif type(value) is file_class:
|
||||
source, start, duration, crop = _attrs(value, [
|
||||
"_VideoFromFile__file", "_VideoFromFile__start_time", "_VideoFromFile__duration",
|
||||
"_VideoFromFile__crop",
|
||||
])
|
||||
descriptor.update(representation="file", start_time=_number(start), duration=_number(duration),
|
||||
crop=_crop(crop))
|
||||
descriptor["data"] = add(_file_bytes(source, limit, allowed_roots))
|
||||
elif type(value) is list_class:
|
||||
videos, complete_audio = _attrs(value, ["videos", "complete_audio"])
|
||||
if type(videos) is not list or not 1 <= len(videos) <= MAX_SEGMENTS:
|
||||
raise ValueError("VIDEO must have between 1 and 64 segments")
|
||||
if any(type(segment) is not file_class for segment in videos):
|
||||
raise ValueError("VIDEO segments must be exact native VideoFromFile objects")
|
||||
descriptor.update(representation="segments", segments=[])
|
||||
for segment in videos:
|
||||
nested, payload = encode_video(segment, max_bytes=limit - size, allowed_roots=allowed_roots)
|
||||
descriptor["segments"].append(add(payload, nested))
|
||||
descriptor["complete_audio"] = tensor("AUDIO", complete_audio) if complete_audio is not None else None
|
||||
else:
|
||||
raise ValueError("Unsupported native VIDEO class")
|
||||
descriptor["byte_length"] = size
|
||||
return descriptor, b"".join(chunks)
|
||||
|
||||
|
||||
def _tensor_header(kind, descriptor, length):
|
||||
"""Allocation-free preflight of the shared tensor-le-v1 descriptor.
|
||||
|
||||
The shared decoder still validates the bytes/finite values. Doing its small
|
||||
structural check here first prevents allocation for an early part when a
|
||||
later part's descriptor is already known to be invalid.
|
||||
"""
|
||||
fields = {"encoding", "dtype", "shape", "layout"}
|
||||
if kind == "AUDIO":
|
||||
fields.add("sample_rate")
|
||||
_fields(descriptor, fields, "tensor")
|
||||
layout = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}[kind]
|
||||
if type(descriptor["encoding"]) is not str or descriptor["encoding"] != "tensor-le-v1":
|
||||
raise ValueError("Invalid VIDEO tensor encoding")
|
||||
if type(descriptor["layout"]) is not str or descriptor["layout"] != layout:
|
||||
raise ValueError("Invalid VIDEO tensor layout")
|
||||
dtype = descriptor["dtype"]
|
||||
if type(dtype) is not str or dtype not in _TENSOR_BYTES:
|
||||
raise ValueError("Invalid VIDEO tensor dtype")
|
||||
shape = descriptor["shape"]
|
||||
if type(shape) is not list or len(shape) != len(layout):
|
||||
raise ValueError("Invalid VIDEO tensor rank")
|
||||
size = _TENSOR_BYTES[dtype]
|
||||
for dimension in shape:
|
||||
_integer(dimension, positive=True)
|
||||
if dimension > length // size:
|
||||
raise ValueError("VIDEO tensor shape exceeds byte limit")
|
||||
size *= dimension
|
||||
if kind == "IMAGE" and shape[-1] not in (1, 3, 4):
|
||||
raise ValueError("Invalid VIDEO image channels")
|
||||
if size != length:
|
||||
raise ValueError("VIDEO tensor length does not match shape")
|
||||
if kind == "AUDIO":
|
||||
_integer(descriptor["sample_rate"], positive=True)
|
||||
|
||||
|
||||
def _alpha_header(alpha, images):
|
||||
_fields(alpha, {"offset", "length", "descriptor"}, "alpha part")
|
||||
descriptor = alpha["descriptor"]
|
||||
if type(descriptor) is not dict:
|
||||
raise ValueError("Invalid VIDEO alpha descriptor")
|
||||
layout = descriptor.get("layout")
|
||||
if type(layout) is not str or layout not in ("BHW", "BHWC"):
|
||||
raise ValueError("Invalid VIDEO alpha layout")
|
||||
kind = "MASK" if layout == "BHW" else "IMAGE"
|
||||
_tensor_header(kind, descriptor, _integer(alpha["length"]))
|
||||
expected = images["descriptor"]["shape"][:3] + ([1] if kind == "IMAGE" else [])
|
||||
if descriptor["shape"] != expected:
|
||||
raise ValueError("VIDEO alpha shape must match images")
|
||||
return kind
|
||||
|
||||
|
||||
def _validate_video(descriptor, size, *, file_only=False):
|
||||
if type(descriptor) is not dict:
|
||||
raise ValueError("VIDEO descriptor must be a dictionary")
|
||||
representation = descriptor.get("representation")
|
||||
if type(representation) is not str or representation not in ("components", "file", "segments"):
|
||||
raise ValueError("Unsupported VIDEO representation")
|
||||
if file_only and representation != "file":
|
||||
# Depth is at most one: native lists already flatten their segments.
|
||||
raise ValueError("VIDEO segments cannot contain nested lists or components")
|
||||
fields = {"components": _COMPONENT_FIELDS, "file": _FILE_FIELDS, "segments": _SEGMENT_FIELDS}[representation]
|
||||
_fields(descriptor, _BASE_FIELDS | fields, "descriptor")
|
||||
if type(descriptor["encoding"]) is not str or descriptor["encoding"] != _ENCODING:
|
||||
raise ValueError("Unsupported VIDEO encoding")
|
||||
if _integer(descriptor["byte_length"]) != size:
|
||||
raise ValueError("VIDEO payload length does not match descriptor")
|
||||
result = dict(descriptor)
|
||||
cursor = 0
|
||||
|
||||
def part(slot, kind=None, video=False):
|
||||
nonlocal cursor
|
||||
fields = {"offset", "length", "descriptor"} if kind or video else {"offset", "length"}
|
||||
_fields(slot, fields, "part")
|
||||
offset = _integer(slot["offset"])
|
||||
length = _integer(slot["length"])
|
||||
if offset != cursor or length > size - cursor:
|
||||
raise ValueError("VIDEO parts must not overlap, leave gaps or exceed bounds")
|
||||
cursor += length
|
||||
if kind:
|
||||
_tensor_header(kind, slot["descriptor"], length)
|
||||
if video:
|
||||
return {**slot, "descriptor": _validate_video(slot["descriptor"], length, file_only=True)}
|
||||
return slot
|
||||
|
||||
if representation == "components":
|
||||
_frame_rate(descriptor["frame_rate"])
|
||||
_encoding_options(descriptor["bit_depth"], descriptor["color_space"])
|
||||
result["metadata"] = _metadata(descriptor["metadata"])
|
||||
part(descriptor["images"], "IMAGE")
|
||||
if descriptor["alpha"] is not None:
|
||||
part(descriptor["alpha"], _alpha_header(descriptor["alpha"], descriptor["images"]))
|
||||
if descriptor["audio"] is not None:
|
||||
part(descriptor["audio"], "AUDIO")
|
||||
elif representation == "file":
|
||||
_number(descriptor["start_time"])
|
||||
_number(descriptor["duration"])
|
||||
_crop(descriptor["crop"], wire=True)
|
||||
part(descriptor["data"])
|
||||
else:
|
||||
segments = descriptor["segments"]
|
||||
if type(segments) is not list or not 1 <= len(segments) <= MAX_SEGMENTS:
|
||||
raise ValueError("VIDEO must have between 1 and 64 segments")
|
||||
result["segments"] = [part(segment, video=True) for segment in segments]
|
||||
if descriptor["complete_audio"] is not None:
|
||||
part(descriptor["complete_audio"], "AUDIO")
|
||||
if cursor != size:
|
||||
raise ValueError("VIDEO payload contains trailing bytes")
|
||||
return result
|
||||
|
||||
|
||||
def _decode_validated(descriptor, payload, types):
|
||||
components_class, file_class, list_class, components_type = types
|
||||
|
||||
def data(slot):
|
||||
return payload[slot["offset"]:slot["offset"] + slot["length"]]
|
||||
|
||||
def tensor(kind, slot):
|
||||
if slot is None:
|
||||
return None
|
||||
return decode_payload(kind, slot["descriptor"], data(slot), max_bytes=slot["length"])
|
||||
|
||||
if descriptor["representation"] == "file":
|
||||
# VideoFromFile retains the BytesIO, unlike a temporary filesystem path.
|
||||
return file_class(io.BytesIO(data(descriptor["data"])), start_time=descriptor["start_time"],
|
||||
duration=descriptor["duration"], crop=_crop(descriptor["crop"], wire=True))
|
||||
if descriptor["representation"] == "components":
|
||||
fps = descriptor["frame_rate"]
|
||||
return components_class(components_type(
|
||||
images=tensor("IMAGE", descriptor["images"]),
|
||||
frame_rate=Fraction(fps["numerator"], fps["denominator"]),
|
||||
audio=tensor("AUDIO", descriptor["audio"]),
|
||||
metadata=descriptor["metadata"],
|
||||
alpha=tensor("IMAGE" if descriptor["alpha"] is not None and
|
||||
descriptor["alpha"]["descriptor"]["layout"] == "BHWC" else "MASK",
|
||||
descriptor["alpha"]),
|
||||
), bit_depth=descriptor["bit_depth"], color_space=descriptor["color_space"])
|
||||
videos = [_decode_validated(slot["descriptor"], data(slot), types) for slot in descriptor["segments"]]
|
||||
return list_class(videos, complete_audio=tensor("AUDIO", descriptor["complete_audio"]))
|
||||
|
||||
|
||||
def decode_video(descriptor, payload, *, max_bytes=MAX_PAYLOAD_BYTES):
|
||||
"""Decode a bounded VIDEO into independently owned native CPU media objects."""
|
||||
limit = _limit(max_bytes)
|
||||
if type(payload) is not bytes:
|
||||
raise ValueError("VIDEO payload must be bytes")
|
||||
if len(payload) > limit:
|
||||
raise ValueError("VIDEO payload exceeds byte limit")
|
||||
validated = _validate_video(descriptor, len(payload))
|
||||
return _decode_validated(validated, payload, _native_types())
|
||||
@@ -0,0 +1,512 @@
|
||||
"""Real CPU Comfy V3 + HTTP proof; invoked only by the opt-in parent test.
|
||||
|
||||
The only harness pieces are deterministic source/assertion nodes and a bounded
|
||||
queue worker. /prompt, /history and /view are Comfy's real handlers; compilation,
|
||||
packing, staging, codecs, bridge execution, response decoding and Collector are
|
||||
unmodified production code. No normal instance or Runpod client is contacted.
|
||||
"""
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib
|
||||
import io as bytes_io
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import socket
|
||||
import sys
|
||||
import types
|
||||
import wave
|
||||
from contextlib import asynccontextmanager
|
||||
from fractions import Fraction
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def bootstrap():
|
||||
if len(sys.argv) != 3:
|
||||
raise SystemExit("Use test_media_end_to_end.py with explicit COMFYUI_SOURCE_ROOT")
|
||||
core, run_root = (Path(value).resolve() for value in sys.argv[1:])
|
||||
assert run_root.is_dir() and run_root != core and not run_root.is_relative_to(core)
|
||||
for name in ("CUDA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"):
|
||||
assert os.environ.get(name) == "", f"CPU isolation missing: {name}"
|
||||
assert not any(name in sys.modules for name in ("torch", "execution", "folder_paths"))
|
||||
sys.path.insert(0, str(core))
|
||||
import comfy.cli_args
|
||||
args = comfy.cli_args.args
|
||||
# Set CPU and all writable roots BEFORE any framework/tensor import.
|
||||
args.cpu = True
|
||||
args.disable_assets = True
|
||||
args.disable_all_custom_nodes = True
|
||||
args.enable_manager = False
|
||||
args.offline = True
|
||||
args.base_directory = str(run_root)
|
||||
args.front_end_root = str(run_root)
|
||||
args.database_url = "sqlite:///:memory:"
|
||||
for name in ("input", "output", "temp", "user"):
|
||||
(run_root / name).mkdir()
|
||||
setattr(args, name + "_directory", str(run_root / name))
|
||||
return run_root
|
||||
|
||||
|
||||
async def check(base):
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
import execution
|
||||
import folder_paths
|
||||
import nodes
|
||||
import server
|
||||
from aiohttp import web
|
||||
from app.assets.manager import default_asset_manager
|
||||
from comfy_api.v0_0_2 import io
|
||||
from comfy_api.latest._input_impl.video_types import (
|
||||
VideoFromComponents, VideoFromFile, VideoFromList,
|
||||
)
|
||||
from comfy_api.latest._util import VideoComponents
|
||||
from comfy_extras import nodes_audio
|
||||
|
||||
assert not torch.cuda.is_initialized()
|
||||
torch.set_num_threads(1)
|
||||
assert execution.comfy.model_management.get_torch_device().type == "cpu"
|
||||
for name in ("input", "output", "temp", "user"):
|
||||
assert Path(getattr(folder_paths, "get_" + name + "_directory")()).is_relative_to(base)
|
||||
|
||||
# Namespace scaffolding avoids the unrelated pack bootstrap (service/config
|
||||
# side effects). Every imported module below is the real source on disk.
|
||||
package_name = "_media_e2e_acceptance"
|
||||
for suffix in ("", ".api", ".api.orchestration", ".nodes"):
|
||||
package = types.ModuleType(package_name + suffix)
|
||||
package.__path__ = [str(ROOT.joinpath(*suffix.lstrip(".").split(".")))] if suffix else [str(ROOT)]
|
||||
sys.modules[package.__name__] = package
|
||||
def load(name):
|
||||
return importlib.import_module(package_name + "." + name)
|
||||
|
||||
asset_manager = default_asset_manager()
|
||||
ps = server.PromptServer(asyncio.get_running_loop(), asset_manager)
|
||||
ps.distributed_pending_jobs = {}
|
||||
ps.distributed_jobs_lock = asyncio.Lock()
|
||||
compiler = load("api.orchestration.serverless_media_prompt")
|
||||
packer = load("api.orchestration.serverless_typed_inputs")
|
||||
decoder = load("api.orchestration.serverless_typed_results")
|
||||
serializer = load("api.orchestration.runpod_client").serialize_request
|
||||
handler = load("serverless_bridge.handler")
|
||||
contract = load("serverless_bridge.contract")
|
||||
video_codec = load("serverless_bridge.video_codec")
|
||||
bridge = load("serverless_bridge")
|
||||
collector = load("nodes.v3").DistributedCollector
|
||||
extension = await bridge.comfy_entrypoint()
|
||||
bridge_classes = await extension.get_node_list()
|
||||
assert {cls.GET_SCHEMA().node_id for cls in bridge_classes} == {"MediaImportV1", "MediaExportV1"}
|
||||
sockets = list(contract.SOCKET_KINDS)
|
||||
collector_sockets = ["images", "audio"]
|
||||
kinds = list(contract.SOCKET_KINDS.values())
|
||||
assert kinds == ["IMAGE", "AUDIO", "MASK", "STRING", "VIDEO"]
|
||||
io_types = [io.Image, io.Audio, io.Mask, io.String, io.Video]
|
||||
fixtures, expected_sinks, seen_sinks, source_calls = {}, {}, {}, []
|
||||
|
||||
class AcceptanceMediaSource(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(node_id="AcceptanceMediaSource", inputs=[io.String.Input("key")],
|
||||
outputs=[kind.Output(is_output_list=True) for kind in io_types])
|
||||
|
||||
@classmethod
|
||||
def execute(cls, key):
|
||||
source_calls.append(key)
|
||||
return io.NodeOutput(*(fixtures[key][socket] for socket in sockets))
|
||||
|
||||
def tensor_equal(actual, expected):
|
||||
assert actual.device.type == "cpu" and actual.dtype == expected.dtype
|
||||
assert actual.shape == expected.shape
|
||||
# Byte comparison preserves signed zero and catches dtype narrowing.
|
||||
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
|
||||
|
||||
def exact(kind, actual, expected):
|
||||
if kind in ("IMAGE", "MASK"):
|
||||
tensor_equal(actual, expected)
|
||||
elif kind == "AUDIO":
|
||||
assert actual["sample_rate"] == expected["sample_rate"]
|
||||
tensor_equal(actual["waveform"], expected["waveform"])
|
||||
elif kind == "STRING":
|
||||
assert type(actual) is str and actual.encode("utf-8") == expected.encode("utf-8")
|
||||
else:
|
||||
assert video_codec.encode_video(actual) == expected
|
||||
if type(actual) is VideoFromFile:
|
||||
buffer = vars(actual)["_VideoFromFile__file"]
|
||||
assert type(buffer) is bytes_io.BytesIO and not buffer.closed
|
||||
position = buffer.tell()
|
||||
buffer.seek(0)
|
||||
with av.open(buffer) as container:
|
||||
assert len(list(container.decode(video=0))) == 3
|
||||
buffer.seek(position)
|
||||
elif type(actual) is VideoFromList:
|
||||
for item in actual.videos:
|
||||
exact("VIDEO", item, video_codec.encode_video(item))
|
||||
else:
|
||||
assert type(actual) is VideoFromComponents
|
||||
components = actual.get_components()
|
||||
assert type(components.frame_rate) is Fraction
|
||||
assert components.images.dtype == torch.float64
|
||||
assert components.alpha.dtype == torch.bfloat16
|
||||
assert components.audio["waveform"].dtype == torch.float64
|
||||
assert actual.get_bit_depth() == 10 and actual.get_color_space() == "HDR PQ"
|
||||
|
||||
class AcceptanceMediaAssert(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(node_id="AcceptanceMediaAssert", inputs=[io.String.Input("tag"),
|
||||
*[kind.Input(socket, optional=True, **({"force_input": True} if socket == "texts" else {}))
|
||||
for socket, kind in zip(sockets, io_types)]], outputs=[],
|
||||
is_input_list=True, is_output_node=True)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, tag, **media):
|
||||
assert len(tag) == 1 and tag[0] not in seen_sinks
|
||||
expected = expected_sinks[tag[0]]
|
||||
assert set(media) == set(expected)
|
||||
for socket in expected:
|
||||
assert len(media[socket]) == len(expected[socket]), (tag, socket)
|
||||
# Existing Collector outputs are a batched IMAGE and one AUDIO,
|
||||
# not lists. Native V3 wraps each output once for this list sink.
|
||||
for actual, wanted in zip(media[socket], expected[socket]):
|
||||
exact(contract.SOCKET_KINDS[socket], actual, wanted)
|
||||
counts = {socket: len(media[socket]) for socket in expected}
|
||||
seen_sinks[tag[0]] = counts
|
||||
return io.NodeOutput(ui={"accepted": [counts]})
|
||||
|
||||
for cls in [*bridge_classes, collector, AcceptanceMediaSource, AcceptanceMediaAssert]:
|
||||
assert issubclass(cls, io.ComfyNode)
|
||||
cls.GET_NODE_INFO_V1()
|
||||
nodes.NODE_CLASS_MAPPINGS[cls.GET_SCHEMA().node_id] = cls
|
||||
|
||||
def node(class_type, **inputs):
|
||||
return {"class_type": class_type, "inputs": inputs}
|
||||
|
||||
def add_bridge_exports(workflow, manifest, request_id, participant, job_id, links):
|
||||
"""Exercise bridge-only kinds without inventing Collector sockets."""
|
||||
for socket, link in links.items():
|
||||
export_id = "bridge-export:" + socket
|
||||
binding = {"binding_id": "bridge-binding:" + socket, "collector_id": "bridge-only",
|
||||
"job_id": job_id, "participant_id": participant, "node_id": export_id,
|
||||
"socket": socket, "kind": contract.SOCKET_KINDS[socket]}
|
||||
assert export_id not in workflow
|
||||
manifest.append(binding)
|
||||
workflow[export_id] = node("MediaExportV1", request_id=request_id,
|
||||
binding_json=json.dumps(binding), **{socket: link})
|
||||
contract.validate_bindings(manifest)
|
||||
|
||||
def tiny_video(variant):
|
||||
buffer = bytes_io.BytesIO()
|
||||
with av.open(buffer, "w", format="matroska") as container:
|
||||
stream = container.add_stream("ffv1", rate=Fraction(24, 1))
|
||||
stream.width = stream.height = 16
|
||||
stream.pix_fmt = "yuv420p"
|
||||
for level in (16 + variant, 64 + variant, 128 + variant):
|
||||
frame = av.VideoFrame.from_ndarray(np.full((16, 16, 3), level, dtype=np.uint8), format="rgb24")
|
||||
for packet in stream.encode(frame):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode():
|
||||
container.mux(packet)
|
||||
return buffer.getvalue()
|
||||
|
||||
def native_fixtures(variant):
|
||||
image = torch.arange(48, dtype=torch.float64).reshape(2, 2, 3, 4) / 997 + variant / 100
|
||||
image[0, 0, 0, 0] = -0.0
|
||||
mask = torch.tensor([[[-0.0, 0.125, 0.875]]], dtype=torch.bfloat16)
|
||||
audio = {"waveform": torch.arange(38, dtype=torch.float64).reshape(1, 2, 19) / 997,
|
||||
"sample_rate": 44100}
|
||||
audio["waveform"][0, 0, 0] = -0.0
|
||||
component = VideoFromComponents(VideoComponents(
|
||||
images=image[..., :3], frame_rate=Fraction(30000, 1001), alpha=mask.reshape(1, 1, 3, 1).repeat(2, 2, 1, 1),
|
||||
audio=audio, metadata={"title": "雪", "nested": [True, None, -4, 0.125]}),
|
||||
bit_depth=10, color_space="HDR PQ")
|
||||
payload = tiny_video(variant)
|
||||
file = VideoFromFile(bytes_io.BytesIO(payload), start_time=-0.0, duration=0, crop=(0, 0, 8, 8))
|
||||
sequence = VideoFromList([VideoFromFile(bytes_io.BytesIO(payload)), file], complete_audio=audio)
|
||||
return {"images": [image, image[:1].clone()], "audio": [audio], "masks": [mask],
|
||||
"texts": [f" café 🎥 {variant}\r\n雪 ", ""], "video": [component, sequence]}
|
||||
|
||||
for key, variant in (("local", 0), ("worker-a", 1), ("worker-b", 2)):
|
||||
fixtures[key] = native_fixtures(variant)
|
||||
native_root = base / "master-input"
|
||||
native_root.mkdir()
|
||||
wav = native_root / "stereo.wav"
|
||||
pcm = np.array([-32768, -12001, -1, 0, 1, 8193, 32767, 17] * 2, dtype="<i2")
|
||||
with wave.open(str(wav), "wb") as out:
|
||||
out.setnchannels(2)
|
||||
out.setsampwidth(2)
|
||||
out.setframerate(44100)
|
||||
out.writeframes(pcm.tobytes())
|
||||
waveform, sample_rate = nodes_audio.load(str(wav))
|
||||
clip = native_root / "native.mkv"
|
||||
clip.write_bytes(tiny_video(3))
|
||||
native_expected = copy.copy(fixtures["worker-a"])
|
||||
native_expected["audio"] = [{"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}]
|
||||
native_expected["video"] = [VideoFromFile(str(clip))]
|
||||
expectations = {}
|
||||
for key, values in (("worker-a", native_expected), ("worker-b", fixtures["worker-b"])):
|
||||
expectations[key] = {socket: [video_codec.encode_video(v, allowed_roots=[native_root]) if socket == "video" else v
|
||||
for v in values[socket]] for socket in sockets}
|
||||
|
||||
# Native routes only, bound to a new ephemeral loopback listener. No fake
|
||||
# history, permissive validator, bridge/codec mock, or existing service.
|
||||
allowed_paths = {"/prompt", "/history/{prompt_id}", "/view", "/queue"}
|
||||
ps.app.add_routes([route for route in ps.routes if route.path in allowed_paths])
|
||||
completed = []
|
||||
stop = asyncio.Event()
|
||||
|
||||
async def queue_worker():
|
||||
while not stop.is_set():
|
||||
item = await asyncio.to_thread(ps.prompt_queue.get, 0.1)
|
||||
if item is None:
|
||||
continue
|
||||
queued, item_id = item
|
||||
assert len(queued) in (5, 6), "Unsupported core PromptQueue tuple; adapt acceptance worker explicitly"
|
||||
executor = execution.PromptExecutor(ps, cache_args={"ram": 0, "ram_inactive": 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, queued[2], queued[1], queued[3], queued[4])
|
||||
status = execution.PromptQueue.ExecutionStatus(
|
||||
status_str="success" if executor.success else "error", completed=executor.success,
|
||||
messages=executor.status_messages)
|
||||
ps.prompt_queue.task_done(item_id, executor.history_result, status=status)
|
||||
completed.append((queued, executor.success, executor.status_messages))
|
||||
|
||||
@asynccontextmanager
|
||||
async def serve():
|
||||
runner = web.AppRunner(ps.app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
ALLOWED_CONNECTIONS.add(("127.0.0.1", port))
|
||||
worker = asyncio.create_task(queue_worker())
|
||||
try:
|
||||
yield f"http://127.0.0.1:{port}"
|
||||
finally:
|
||||
stop.set()
|
||||
await asyncio.wait_for(worker, 10)
|
||||
await runner.cleanup()
|
||||
ALLOWED_CONNECTIONS.clear()
|
||||
|
||||
grouped_results = {}
|
||||
async with serve() as url:
|
||||
for key in ("worker-a", "worker-b"):
|
||||
participant = "serverless:" + key
|
||||
request_id = "acceptance-" + key
|
||||
graph = {"source": node("AcceptanceMediaSource", key=key),
|
||||
"collector": node("DistributedCollector", load_balance=False, transfer_format="exact",
|
||||
multi_job_id="job-main", worker_id=participant,
|
||||
**{socket: ["source", index] for index, socket in enumerate(collector_sockets)}),
|
||||
"audio-collector": node("DistributedCollector", load_balance=False,
|
||||
multi_job_id="job-audio", worker_id=participant, audio=["source", 1])}
|
||||
if key == "worker-a":
|
||||
graph.update({"wav": node("LoadAudio", audio=wav.name), "clip": node("LoadVideo", file=clip.name)})
|
||||
graph["collector"]["inputs"]["audio"] = ["wav", 0]
|
||||
before = copy.deepcopy(graph)
|
||||
workflow, manifest = compiler.compile_media_prompt(graph, request_id)
|
||||
assert graph == before
|
||||
bridge_links = {socket: ["source", sockets.index(socket)] for socket in sockets[2:]}
|
||||
if key == "worker-a":
|
||||
workflow["clip"] = copy.deepcopy(graph["clip"])
|
||||
bridge_links["video"] = ["clip", 0]
|
||||
add_bridge_exports(workflow, manifest, request_id, participant, "job-bridge", bridge_links)
|
||||
assert len(manifest) == 6 and {b["kind"] for b in manifest} == set(kinds)
|
||||
folder_paths.set_input_directory(str(native_root))
|
||||
workflow_before = copy.deepcopy(workflow)
|
||||
images, artifacts, refs = await packer.prepare_media_inputs(workflow)
|
||||
assert images == [] and workflow == workflow_before
|
||||
if key == "worker-a":
|
||||
assert {a["kind"] for a in artifacts} == {"AUDIO", "VIDEO"} and len(artifacts) == 2
|
||||
assert len(refs) == 2
|
||||
assert {ref["input_name"] for ref in refs} == {"audio", "video"}
|
||||
assert all(workflow[ref["node_id"]]["class_type"] == "MediaExportV1" for ref in refs)
|
||||
# Endpoint cannot fall back to the master's local files/loaders.
|
||||
wav.unlink()
|
||||
clip.unlink()
|
||||
print("NATIVE_INPUTS_AND_DIRECT_EXPORT_REFS_OK WAV + native file VIDEO")
|
||||
else:
|
||||
assert not artifacts and not refs
|
||||
folder_paths.set_input_directory(str(base / "input"))
|
||||
media = {"version": 1, "profile": contract.PROFILE, "request_id": request_id,
|
||||
"artifacts": artifacts, "inputs": refs, "outputs": manifest}
|
||||
body = json.loads(serializer(workflow, images, 30, media=media))
|
||||
assert decoder.validate_prepared_job({"worker": {"id": participant, "media_profile": contract.PROFILE,
|
||||
"timeout_seconds": 30}, "workflow": workflow, "images": images,
|
||||
"media": media, "manifest": manifest}) is True
|
||||
before = copy.deepcopy(body["input"])
|
||||
result = await handler.execute_media_request(body["input"], comfy_url=url,
|
||||
input_dir=base / "input", temp_dir=base / "temp", timeout_seconds=30)
|
||||
assert body["input"] == before
|
||||
assert completed and completed[-1][1], completed
|
||||
queued = completed[-1][0]
|
||||
assert set(queued[4]) == {b["node_id"] for b in manifest}
|
||||
submitted = queued[2]
|
||||
if key == "worker-a":
|
||||
assert "wav" not in submitted and "clip" not in submitted
|
||||
imports = [n for n in submitted.values() if n["class_type"] == "MediaImportV1"]
|
||||
assert len(imports) == 2
|
||||
assert {json.loads(n["inputs"]["artifact_json"])["kind"] for n in imports} == {"AUDIO", "VIDEO"}
|
||||
assert not list((base / "input").rglob("*.bin"))
|
||||
assert not list((base / "temp").rglob("*.bin"))
|
||||
if key == "worker-b":
|
||||
# Retire the originals as well as handler-owned staging/export
|
||||
# files before constructing the result's lazy native videos.
|
||||
for original in fixtures[key]["video"][1].videos:
|
||||
vars(original)["_VideoFromFile__file"].close()
|
||||
# Exercise the wire format, not an in-memory handler-only object.
|
||||
result = json.loads(json.dumps(result, allow_nan=False))
|
||||
grouped = await decoder.collect_typed_results(result, media, manifest)
|
||||
assert set(grouped) == {"job-main", "job-audio", "job-bridge"}
|
||||
assert set(grouped["job-main"]) == set(collector_sockets)
|
||||
assert set(grouped["job-bridge"]) == set(sockets[2:])
|
||||
exact("AUDIO", grouped["job-audio"]["audio"][0], fixtures[key]["audio"][0])
|
||||
all_media = {**grouped["job-main"], **grouped["job-bridge"]}
|
||||
for socket, kind in zip(sockets, kinds):
|
||||
values, expected = all_media[socket], expectations[key][socket]
|
||||
assert len(values) == len(expected), (key, socket)
|
||||
for actual, wanted in zip(values, expected):
|
||||
exact(kind, actual, wanted)
|
||||
if key == "worker-b":
|
||||
original = fixtures[key]["video"][0].get_components()
|
||||
decoded = grouped["job-bridge"]["video"][0].get_components()
|
||||
for field in ("images", "alpha"):
|
||||
assert getattr(decoded, field).data_ptr() != getattr(original, field).data_ptr()
|
||||
assert decoded.audio["waveform"].data_ptr() != original.audio["waveform"].data_ptr()
|
||||
assert decoded.metadata == original.metadata and decoded.metadata is not original.metadata
|
||||
grouped_results[participant] = grouped
|
||||
# Late invalid artifact must reject the entire grouped response.
|
||||
corrupt = copy.deepcopy(result)
|
||||
corrupt["media"]["artifacts"][-1]["sha256"] = "0" * 64
|
||||
try:
|
||||
await decoder.collect_typed_results(corrupt, media, manifest)
|
||||
except ValueError:
|
||||
pass
|
||||
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 ("legacy_png", "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("DistributedCollector", transfer_format=fmt,
|
||||
transfer_quality=73, multi_job_id="job-" + key,
|
||||
worker_id="serverless:" + key,
|
||||
**{socket: ["source", i] for i, socket in enumerate(collector_sockets)})}
|
||||
if fmt == "legacy_png":
|
||||
# A saved Collector with no newly selected controls still uses
|
||||
# legacy_png / 90, mapped to PNG by the serverless compiler.
|
||||
graph["collector"]["inputs"].pop("transfer_format")
|
||||
graph["collector"]["inputs"].pop("transfer_quality")
|
||||
original = copy.deepcopy(graph)
|
||||
workflow, manifest = compiler.compile_media_prompt(graph, request_id)
|
||||
if fmt == "legacy_png":
|
||||
binding = next(b for b in manifest if b["kind"] == "IMAGE")
|
||||
controls = workflow[binding["node_id"]]["inputs"]
|
||||
assert (controls["transfer_format"], controls["transfer_quality"]) == ("png", 90)
|
||||
add_bridge_exports(workflow, manifest, request_id, "serverless:" + key, "job-bridge-" + key,
|
||||
{socket: ["source", sockets.index(socket)] for socket in sockets[2:]})
|
||||
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], **grouped["job-bridge-" + 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 in ("legacy_png", "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 default legacy_png -> png/90; png jpeg h264; alpha and other media exact")
|
||||
print("REAL_HTTP_BRIDGE_RESULTS_OK Collector IMAGE/AUDIO + independent MASK/STRING/VIDEO bridge exports; native VIDEO survives cleanup")
|
||||
print("ATOMIC_DECODE_REJECTION_OK corrupt final artifact rejected before delivery")
|
||||
|
||||
participants = list(grouped_results)
|
||||
|
||||
def expected_batches(ordered, include_images=True):
|
||||
audio_items = [item for media in ordered for item in media["audio"]]
|
||||
expected = {"audio": [{"waveform": torch.cat([item["waveform"] for item in audio_items], dim=-1),
|
||||
"sample_rate": audio_items[0]["sample_rate"]}]}
|
||||
if include_images:
|
||||
expected["images"] = [torch.cat([image for media in ordered for image in media["images"]])]
|
||||
return expected
|
||||
|
||||
for delegate in (False, True):
|
||||
tag = "delegate" if delegate else "master"
|
||||
job = "collect-" + tag
|
||||
queue = asyncio.Queue()
|
||||
for participant in reversed(participants):
|
||||
await queue.put({"worker_id": participant, "media": grouped_results[participant]["job-main"], "is_last": True})
|
||||
ps.distributed_pending_jobs[job] = queue
|
||||
ordered = ([] if delegate else [fixtures["local"]]) + [grouped_results[p]["job-main"] for p in participants]
|
||||
expected_sinks[tag] = expected_batches(ordered)
|
||||
graph = {"collector": node("DistributedCollector", load_balance=False, multi_job_id=job,
|
||||
enabled_worker_ids=json.dumps(participants), delegate_only=delegate, transfer_format="exact"),
|
||||
"assert": node("AcceptanceMediaAssert", tag=tag,
|
||||
**{socket: ["collector", index] for index, socket in enumerate(collector_sockets)})}
|
||||
if not delegate:
|
||||
graph["source"] = node("AcceptanceMediaSource", key="local")
|
||||
graph["collector"]["inputs"].update({socket: ["source", index] for index, socket in enumerate(collector_sockets)})
|
||||
# A second audio-only Collector consumes its own job group from the same
|
||||
# atomic response. It must return None for images and batch audio in order.
|
||||
audio_job, audio_tag = job + "-audio", tag + "-audio"
|
||||
audio_queue = asyncio.Queue()
|
||||
for participant in reversed(participants):
|
||||
await audio_queue.put({"worker_id": participant,
|
||||
"media": grouped_results[participant]["job-audio"], "is_last": True})
|
||||
ps.distributed_pending_jobs[audio_job] = audio_queue
|
||||
audio_ordered = ([] if delegate else [fixtures["local"]]) + [grouped_results[p]["job-audio"] for p in participants]
|
||||
expected_sinks[audio_tag] = expected_batches(audio_ordered, include_images=False)
|
||||
graph["audio-collector"] = node("DistributedCollector", load_balance=False, multi_job_id=audio_job,
|
||||
enabled_worker_ids=json.dumps(participants), delegate_only=delegate)
|
||||
graph["audio-assert"] = node("AcceptanceMediaAssert", tag=audio_tag, audio=["audio-collector", 1])
|
||||
if not delegate:
|
||||
graph["audio-collector"]["inputs"]["audio"] = ["source", 1]
|
||||
valid = await execution.validate_prompt(job, graph, None)
|
||||
assert valid[0] and not valid[3], valid
|
||||
executor = execution.PromptExecutor(ps, cache_args={"ram": 0, "ram_inactive": 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, graph, job, {}, valid[2])
|
||||
assert executor.success, executor.status_messages
|
||||
assert executor.history_result["outputs"]["assert"]["accepted"] == [seen_sinks[tag]]
|
||||
assert executor.history_result["outputs"]["audio-assert"]["accepted"] == [seen_sinks[audio_tag]]
|
||||
assert job not in ps.distributed_pending_jobs and audio_job not in ps.distributed_pending_jobs
|
||||
assert source_calls == ["worker-a", "worker-b", "encoded-legacy_png", "encoded-png", "encoded-jpeg", "encoded-h264", "local"], source_calls
|
||||
assert set(seen_sinks) == {"master", "delegate", "master-audio", "delegate-audio"}
|
||||
|
||||
assert not list((base / "output").rglob("*"))
|
||||
assert not torch.cuda.is_initialized()
|
||||
print("MASTER_AND_DELEGATE_ORDER_OK", json.dumps(seen_sinks, sort_keys=True))
|
||||
print("MEDIA_END_TO_END_CPU_OK real PromptExecutor/HTTP; no queue-orchestration or provider acceptance claimed")
|
||||
|
||||
|
||||
ALLOWED_CONNECTIONS = set()
|
||||
|
||||
|
||||
def network_guard(event, args):
|
||||
"""Fail closed on provider/normal-instance traffic, including during imports."""
|
||||
if event == "socket.connect":
|
||||
sock, address = args
|
||||
if sock.family in (socket.AF_INET, socket.AF_INET6):
|
||||
assert address in ALLOWED_CONNECTIONS, "Acceptance forbids non-owned network connections"
|
||||
elif event == "socket.bind":
|
||||
sock, address = args
|
||||
if sock.family in (socket.AF_INET, socket.AF_INET6):
|
||||
assert address == ("127.0.0.1", 0), "Acceptance binds only ephemeral loopback"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.addaudithook(network_guard)
|
||||
asyncio.run(asyncio.wait_for(check(bootstrap()), timeout=120))
|
||||
@@ -36,7 +36,7 @@ BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text
|
||||
ADDED_OPTIONAL_INPUTS = {
|
||||
'DistributedCollector': {
|
||||
'transfer_format': [
|
||||
['legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1'],
|
||||
['legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1', 'exact'],
|
||||
{'default': 'legacy_png'},
|
||||
],
|
||||
'transfer_quality': ['INT', {'default': 90, 'min': 1, 'max': 100, 'step': 1}],
|
||||
@@ -372,6 +372,12 @@ async def main():
|
||||
assert set(mapping) == set(BASELINE) | ADDED_NODES
|
||||
for node_id in ADDED_NODES:
|
||||
assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id]
|
||||
assert 'DistributedMediaCollector' not in mapping
|
||||
collector_info = mapping['DistributedCollector'].GET_NODE_INFO_V1()
|
||||
assert collector_info['output'] == ['IMAGE', 'AUDIO'], collector_info
|
||||
assert not any(collector_info['output_is_list']), collector_info
|
||||
assert list(collector_info['input']['optional']) == [
|
||||
'images', 'audio', 'transfer_format', 'transfer_quality']
|
||||
sink_info = mapping['DistributedWorkerSink'].GET_NODE_INFO_V1()
|
||||
assert sink_info['output_node'] and sink_info['dev_only'], sink_info
|
||||
for node_id, cls in mapping.items():
|
||||
|
||||
@@ -0,0 +1,524 @@
|
||||
"""Exact Collector transport through real serializers, HTTP, route and queue.
|
||||
|
||||
Only ComfyUI's process-wide server/model-management and unrelated orchestration
|
||||
are replaced. Tensor/audio codecs, multipart bodies and aiohttp are real; no
|
||||
running ComfyUI service, GPU, provider or external network access is required.
|
||||
"""
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
import importlib
|
||||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PKG = "dist_collector_exact_testpkg"
|
||||
JOB_ID = "exact-job"
|
||||
WORKER_ID = "worker-1"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def transport(monkeypatch):
|
||||
# Other legacy test loaders may leave aiohttp stubs behind. Replace only
|
||||
# those entries, never unload/reload real extension-backed dependencies.
|
||||
for name, loaded_module in list(sys.modules.items()):
|
||||
if (name == "aiohttp" or name.startswith("aiohttp.")) and not getattr(loaded_module, "__file__", None):
|
||||
monkeypatch.delitem(sys.modules, name)
|
||||
aiohttp = importlib.import_module("aiohttp")
|
||||
web = importlib.import_module("aiohttp.web")
|
||||
|
||||
def module(name, **attrs):
|
||||
result = types.ModuleType(name)
|
||||
result.__dict__.update(attrs)
|
||||
monkeypatch.setitem(sys.modules, name, result)
|
||||
return result
|
||||
|
||||
def load(relative):
|
||||
name = f"{PKG}.{relative}"
|
||||
spec = importlib.util.spec_from_file_location(name, ROOT.joinpath(*relative.split(".")).with_suffix(".py"))
|
||||
assert spec is not None and spec.loader is not None
|
||||
result = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, name, result)
|
||||
spec.loader.exec_module(result)
|
||||
parent, _, child = name.rpartition(".")
|
||||
setattr(sys.modules[parent], child, result)
|
||||
return result
|
||||
|
||||
module(PKG, __path__=[str(ROOT)])
|
||||
for package in ("nodes", "api", "utils", "serverless_bridge"):
|
||||
module(f"{PKG}.{package}", __path__=[str(ROOT / package)])
|
||||
server = types.SimpleNamespace(routes=web.RouteTableDef(), distributed_pending_jobs={})
|
||||
module("server", PromptServer=types.SimpleNamespace(instance=server))
|
||||
mm = module("comfy.model_management", throw_exception_if_processing_interrupted=lambda: None)
|
||||
progress = module("comfy.utils", ProgressBar=lambda _total: types.SimpleNamespace(update=lambda _n: None))
|
||||
module("comfy", __path__=[], model_management=mm, utils=progress)
|
||||
module(f"{PKG}.utils.logging", debug_log=lambda *_a, **_kw: None, log=lambda *_a, **_kw: None)
|
||||
module(f"{PKG}.utils.config", get_worker_timeout_seconds=lambda: 0.1,
|
||||
load_config=lambda: {"workers": []}, is_master_delegate_only=lambda: False)
|
||||
|
||||
async def handle_api_error(_request, error, status=500):
|
||||
return web.json_response({"error": str(error)}, status=status)
|
||||
|
||||
module(f"{PKG}.utils.network", handle_api_error=handle_api_error,
|
||||
get_client_session=AsyncMock(), build_worker_url=lambda _worker: "http://unused.invalid",
|
||||
probe_worker=AsyncMock(return_value=None))
|
||||
module(f"{PKG}.utils.async_helpers", run_async_in_server_loop=asyncio.run)
|
||||
module(f"{PKG}.api.queue_orchestration", ensure_distributed_state=lambda: None,
|
||||
orchestrate_distributed_execution=AsyncMock())
|
||||
module(f"{PKG}.api.queue_request", parse_queue_request_payload=lambda value: value)
|
||||
for name in ("utils.constants", "utils.image", "utils.audio_payload", "utils.transfer_codec"):
|
||||
load(name)
|
||||
bridge_codec = load("serverless_bridge.codec")
|
||||
load("serverless_bridge.image_transfer")
|
||||
image_codec = load("serverless_bridge.image_codec")
|
||||
collector = load("nodes.collector")
|
||||
routes = load("api.job_routes")
|
||||
return types.SimpleNamespace(aiohttp=aiohttp, web=web, server=server,
|
||||
collector=collector, routes=routes, bridge_codec=bridge_codec,
|
||||
image_codec=image_codec)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def endpoint(transport, monkeypatch, *, request_limit=2 * 1024 * 1024,
|
||||
chunk_target=1024 * 1024, client_max_size=None, handler=None):
|
||||
"""A disposable loopback endpoint; record actual HTTP lengths, not estimates."""
|
||||
queue = asyncio.Queue()
|
||||
transport.server.distributed_jobs_lock = asyncio.Lock()
|
||||
transport.server.distributed_pending_jobs = {JOB_ID: queue}
|
||||
requests = []
|
||||
web = transport.web
|
||||
|
||||
@web.middleware
|
||||
async def record_request(request, handler):
|
||||
record = {"bytes": request.content_length, "content_type": request.content_type}
|
||||
requests.append(record)
|
||||
response = await handler(request)
|
||||
record["status"] = response.status
|
||||
if response.status < 400 and request.content_type.startswith("multipart/"):
|
||||
form = await request.post()
|
||||
record["envelope"] = json.loads(form["envelope"])
|
||||
return response
|
||||
|
||||
app = web.Application(client_max_size=client_max_size or request_limit, middlewares=[record_request])
|
||||
app.router.add_post("/distributed/job_complete", handler or transport.routes.job_complete_endpoint)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
try:
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
async with transport.aiohttp.ClientSession() as session:
|
||||
monkeypatch.setattr(transport.collector, "get_client_session", AsyncMock(return_value=session))
|
||||
monkeypatch.setattr(transport.collector, "TRANSFER_MAX_REQUEST_BYTES", request_limit)
|
||||
monkeypatch.setattr(transport.collector, "TRANSFER_MAX_CHUNK_BYTES", chunk_target)
|
||||
yield types.SimpleNamespace(
|
||||
session=session, url=f"http://127.0.0.1:{port}", requests=requests, queue=queue,
|
||||
node=transport.collector.DistributedCollectorNode(),
|
||||
)
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
def frames(dtype=torch.float32, channels=3, count=5, height=7, width=9):
|
||||
# Non-contiguous, out-of-range and non-8-bit-representable values exercise
|
||||
# detach/contiguous without permitting clipping, alpha loss or quantization.
|
||||
values = torch.arange(count * height * width * channels, dtype=torch.float64)
|
||||
values = (values.remainder(41) - 11) / 17
|
||||
images = values.reshape(count, width, height, channels).transpose(1, 2).to(dtype)
|
||||
images[0, 0, 0, 0] = -0.0
|
||||
return images.requires_grad_(True)
|
||||
|
||||
|
||||
def assert_exact(actual, expected):
|
||||
assert actual.dtype == expected.dtype
|
||||
assert actual.shape == expected.shape
|
||||
assert actual.device.type == "cpu"
|
||||
assert actual.is_contiguous()
|
||||
# Compare raw bits as well as values, including BF16 and signed zero.
|
||||
assert torch.equal(actual, expected)
|
||||
assert torch.equal(actual.view(torch.uint8), expected.detach().contiguous().view(torch.uint8))
|
||||
|
||||
|
||||
def queued_items(queue):
|
||||
return [queue.get_nowait() for _ in range(queue.qsize())]
|
||||
|
||||
|
||||
async def send(ctx, images, *, audio=None, fmt="exact"):
|
||||
await ctx.node.send_batch_to_master(images, audio, JOB_ID, ctx.url, WORKER_ID, fmt, 17)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("delegate_only", [False, True])
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
||||
@pytest.mark.parametrize("channels", [1, 3, 4])
|
||||
def test_exact_sender_route_and_collector_preserve_dtype_values_and_channels(
|
||||
transport, monkeypatch, dtype, channels, delegate_only):
|
||||
images = frames(dtype, channels)
|
||||
frame_bytes = images[0].numel() * images.element_size()
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=reserve + 2 * frame_bytes,
|
||||
chunk_target=2 * frame_bytes) as ctx:
|
||||
await send(ctx, images)
|
||||
items = queued_items(ctx.queue)
|
||||
assert [item["image_index"] for item in items] == list(range(len(images)))
|
||||
assert [item["is_last"] for item in items] == [False] * 4 + [True]
|
||||
assert all(item["worker_id"] == WORKER_ID and item["format"] == "exact" for item in items)
|
||||
for item, expected in zip(items, images.split(1)):
|
||||
assert_exact(item["tensor"], expected)
|
||||
await ctx.queue.put(item)
|
||||
envelopes = [record["envelope"] for record in ctx.requests]
|
||||
assert [(e["start_index"], e["count"]) for e in envelopes] == [(0, 2), (2, 2), (4, 1)]
|
||||
assert all(record["bytes"] <= reserve + 2 * frame_bytes for record in ctx.requests)
|
||||
master_images = None if delegate_only else images[:1].detach().clone().add(0.03125)
|
||||
actual, _audio = await ctx.node.execute(
|
||||
master_images, None, multi_job_id=JOB_ID, enabled_worker_ids=json.dumps([WORKER_ID]),
|
||||
delegate_only=delegate_only, transfer_format="exact",
|
||||
)
|
||||
expected = images if delegate_only else torch.cat([master_images, images])
|
||||
assert_exact(actual, expected)
|
||||
assert JOB_ID not in transport.server.distributed_pending_jobs
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_single_frame_can_exceed_chunk_target_but_not_request_limit(transport, monkeypatch):
|
||||
images = frames(count=3)
|
||||
frame_bytes = images[0].numel() * images.element_size()
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=reserve + frame_bytes,
|
||||
chunk_target=frame_bytes // 2) as ctx:
|
||||
await send(ctx, images)
|
||||
assert [record["envelope"]["count"] for record in ctx.requests] == [1, 1, 1]
|
||||
assert all(record["bytes"] <= reserve + frame_bytes for record in ctx.requests)
|
||||
assert_exact(torch.cat([item["tensor"] for item in queued_items(ctx.queue)]), images)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_oversized_frame_is_rejected_before_encoding_or_post(transport, monkeypatch):
|
||||
images = frames(count=3)
|
||||
frame_bytes = images[0].numel() * images.element_size()
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
encode = AsyncMock(side_effect=AssertionError("oversized frame must not be encoded"))
|
||||
monkeypatch.setattr(transport.collector, "_encode_exact_frames", encode)
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=reserve + frame_bytes - 1) as ctx:
|
||||
with pytest.raises(transport.collector.transfer_codec.FrameTooLargeError, match="request limit"):
|
||||
await send(ctx, images)
|
||||
assert ctx.requests == []
|
||||
assert ctx.queue.empty()
|
||||
encode.assert_not_called()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(0, 2, 3, 3), (1, 0, 3, 3), (1, 2, 0, 3), (1, 2, 3, 0),
|
||||
(1, 2, 3, 2), (1, 2, 3, 5), (1, 2, 3), (2, 3)])
|
||||
def test_exact_sender_rejects_invalid_geometry_before_post(transport, monkeypatch, shape):
|
||||
image = torch.empty(shape)
|
||||
with pytest.raises(transport.collector.transfer_codec.TransferCodecError):
|
||||
transport.collector._exact_frame_bytes(image)
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
# An empty batch is also a no-media completion error (ValueError).
|
||||
with pytest.raises(ValueError):
|
||||
await send(ctx, image)
|
||||
assert ctx.requests == []
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.uint8, torch.int32, torch.bool, torch.complex64])
|
||||
def test_exact_sender_rejects_unsupported_dtype_before_post(transport, monkeypatch, dtype):
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="dtype"):
|
||||
await send(ctx, torch.zeros(1, 2, 3, 3, dtype=dtype))
|
||||
assert ctx.requests == []
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_frame_planning_enforces_bridge_byte_limit(transport, monkeypatch):
|
||||
monkeypatch.setattr(transport.bridge_codec, "MAX_PAYLOAD_BYTES", 32)
|
||||
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="[Bb]yte|limit"):
|
||||
transport.collector._exact_frame_bytes(torch.zeros(1, 2, 3, 3))
|
||||
|
||||
|
||||
def test_exact_frame_planning_rejects_non_dense_layout(transport):
|
||||
image = torch.zeros(1, 2, 3, 3).to_sparse()
|
||||
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="dense|layout|strided"):
|
||||
transport.collector._exact_frame_bytes(image)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")])
|
||||
def test_exact_sender_rejects_nonfinite_payload(transport, monkeypatch, value):
|
||||
images = torch.full((1, 2, 3, 3), value)
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="finite"):
|
||||
await send(ctx, images)
|
||||
assert ctx.requests == []
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def exact_chunk(transport, count=2):
|
||||
images = frames(count=count).detach()
|
||||
encoded = [transport.bridge_codec.encode_payload("IMAGE", frame) for frame in images.split(1)]
|
||||
return {
|
||||
"job_id": JOB_ID, "worker_id": WORKER_ID, "format": "exact",
|
||||
"start_index": 0, "count": count, "width": images.shape[2], "height": images.shape[1],
|
||||
"is_last": True, "descriptors": [descriptor for descriptor, _ in encoded],
|
||||
}, [payload for _, payload in encoded]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", [
|
||||
"zero_count", "bool_width", "negative_height", "negative_index", "bool_index",
|
||||
"missing_descriptor", "extra_descriptor", "missing_part", "extra_part", "not_a_descriptor",
|
||||
"wrong_height", "wrong_batch", "bad_channels", "bad_dtype", "bad_layout", "bad_encoding",
|
||||
"extra_descriptor_field", "truncated", "trailing_bytes", "nonfinite", "compressed_spoof",
|
||||
])
|
||||
def test_exact_route_rejects_bad_geometry_and_payload_atomically(transport, monkeypatch, case):
|
||||
envelope, parts = exact_chunk(transport)
|
||||
descriptor = envelope["descriptors"][-1]
|
||||
if case == "zero_count":
|
||||
envelope["count"] = 0
|
||||
elif case == "bool_width":
|
||||
envelope["width"] = True
|
||||
elif case == "negative_height":
|
||||
envelope["height"] = -1
|
||||
elif case == "negative_index":
|
||||
envelope["start_index"] = -1
|
||||
elif case == "bool_index":
|
||||
envelope["start_index"] = True
|
||||
elif case == "missing_descriptor":
|
||||
envelope["descriptors"].pop()
|
||||
elif case == "extra_descriptor":
|
||||
envelope["descriptors"].append(dict(descriptor))
|
||||
elif case == "missing_part":
|
||||
parts.pop()
|
||||
elif case == "extra_part":
|
||||
parts.append(parts[-1])
|
||||
elif case == "not_a_descriptor":
|
||||
envelope["descriptors"][-1] = None
|
||||
elif case == "wrong_height":
|
||||
descriptor["shape"][1] += 1
|
||||
elif case == "wrong_batch":
|
||||
descriptor["shape"][0] = 2
|
||||
elif case == "bad_channels":
|
||||
descriptor["shape"][-1] = 2
|
||||
elif case == "bad_dtype":
|
||||
descriptor["dtype"] = "uint8"
|
||||
elif case == "bad_layout":
|
||||
descriptor["layout"] = "BCHW"
|
||||
elif case == "bad_encoding":
|
||||
descriptor["encoding"] = "pickle"
|
||||
elif case == "extra_descriptor_field":
|
||||
descriptor["unknown"] = True
|
||||
elif case == "truncated":
|
||||
parts[-1] = parts[-1][:-1]
|
||||
elif case == "trailing_bytes":
|
||||
parts[-1] += b"\0"
|
||||
elif case == "nonfinite":
|
||||
bad = torch.full(tuple(descriptor["shape"]), float("nan"))
|
||||
parts[-1] = bad.numpy().tobytes()
|
||||
elif case == "compressed_spoof":
|
||||
# A valid serverless image-transfer descriptor must NOT be accepted as
|
||||
# exact by the direct route, even though the shared bridge can decode it.
|
||||
envelope["descriptors"][-1], parts[-1] = transport.image_codec.encode_image(
|
||||
frames(count=1).detach(), transfer_format="png", transfer_quality=90,
|
||||
)
|
||||
# Prove the spoof is a valid compressed bridge payload, not simply junk.
|
||||
decoded = transport.bridge_codec.decode_payload("IMAGE", envelope["descriptors"][-1], parts[-1])
|
||||
assert list(decoded.shape) == descriptor["shape"]
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
async with ctx.session.post(
|
||||
ctx.url + "/distributed/job_complete",
|
||||
data=ctx.node._build_multipart(envelope, parts),
|
||||
) as response:
|
||||
body = await response.text()
|
||||
assert response.status == 400, body
|
||||
# A corrupt second frame must not leave the first in the queue.
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status,body", [(400, "Invalid JSON payload: old master"),
|
||||
(413, "upload limit"), (500, "unexpected failure")])
|
||||
def test_exact_master_rejection_never_falls_back(transport, monkeypatch, status, body):
|
||||
async def reject(_request):
|
||||
return transport.web.Response(status=status, text=body)
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, handler=reject, chunk_target=1) as ctx:
|
||||
legacy = AsyncMock(side_effect=AssertionError("exact must never use legacy transport"))
|
||||
compressed = AsyncMock(side_effect=AssertionError("exact must never use compressed transport"))
|
||||
monkeypatch.setattr(ctx.node, "_send_legacy_batch", legacy)
|
||||
monkeypatch.setattr(ctx.node, "_send_compressed_batch", compressed)
|
||||
with pytest.raises(RuntimeError, match=f"HTTP {status}") as raised:
|
||||
await send(ctx, frames(count=3))
|
||||
assert "never" in str(raised.value) and "8-bit" in str(raised.value)
|
||||
assert len(ctx.requests) == 1
|
||||
assert ctx.requests[0]["content_type"].startswith("multipart/")
|
||||
assert ctx.queue.empty()
|
||||
legacy.assert_not_called()
|
||||
compressed.assert_not_called()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["legacy_png", "png"])
|
||||
def test_exact_collector_rejects_old_worker_8bit_result(transport, monkeypatch, fmt):
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
await send(ctx, frames(count=1).detach(), fmt=fmt)
|
||||
assert ctx.queue.qsize() == 1
|
||||
with pytest.raises(RuntimeError, match="requested exact"):
|
||||
await ctx.node.execute(
|
||||
None, None, multi_job_id=JOB_ID, enabled_worker_ids=json.dumps([WORKER_ID]),
|
||||
delegate_only=True, transfer_format="exact",
|
||||
)
|
||||
assert JOB_ID not in transport.server.distributed_pending_jobs
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_collector_shape_mismatch_raises_instead_of_returning_local_images(transport, monkeypatch):
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
await send(ctx, frames(count=1, channels=4))
|
||||
with pytest.raises(RuntimeError, match="shape|combin|mismatch"):
|
||||
await ctx.node.execute(
|
||||
frames(count=1, channels=3), None, multi_job_id=JOB_ID,
|
||||
enabled_worker_ids=json.dumps([WORKER_ID]), transfer_format="exact",
|
||||
)
|
||||
assert JOB_ID not in transport.server.distributed_pending_jobs
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_incomplete_worker_raises_and_cleans_queue(transport, monkeypatch):
|
||||
"""A timed-out worker's first frame is not a successful partial result."""
|
||||
envelope, parts = exact_chunk(transport, count=1)
|
||||
envelope["is_last"] = False
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
async with ctx.session.post(
|
||||
ctx.url + "/distributed/job_complete", data=ctx.node._build_multipart(envelope, parts),
|
||||
) as response:
|
||||
assert response.status == 200, await response.text()
|
||||
assert ctx.queue.qsize() == 1
|
||||
with pytest.raises(RuntimeError, match="Exact transfer incomplete"):
|
||||
await asyncio.wait_for(ctx.node.execute(
|
||||
frames(count=1).detach(), None, multi_job_id=JOB_ID,
|
||||
enabled_worker_ids=json.dumps([WORKER_ID]), transfer_format="exact",
|
||||
), timeout=2)
|
||||
assert JOB_ID not in transport.server.distributed_pending_jobs
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_oversized_multipart_metadata_is_rejected_before_post(transport, monkeypatch):
|
||||
images = frames(count=1)
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
limit = reserve + images.numel() * images.element_size()
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=limit) as ctx:
|
||||
with pytest.raises(transport.collector.transfer_codec.FrameTooLargeError, match="request limit"):
|
||||
await ctx.node.send_batch_to_master(
|
||||
images, None, JOB_ID, ctx.url, "w" * limit, "exact", 90,
|
||||
)
|
||||
assert ctx.requests == []
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_actual_server_upload_rejection_is_413_without_retry(transport, monkeypatch):
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, client_max_size=1024) as ctx:
|
||||
with pytest.raises(RuntimeError, match="HTTP 413"):
|
||||
await send(ctx, frames(count=3, height=64, width=64))
|
||||
assert len(ctx.requests) == 1
|
||||
assert ctx.requests[0]["status"] == 413
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_oversized_audio_is_rejected_before_any_image_post(transport, monkeypatch):
|
||||
images = frames(count=3)
|
||||
audio = {"waveform": torch.zeros(1, 1, 100000), "sample_rate": 48000}
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=reserve + 1024) as ctx:
|
||||
with pytest.raises(transport.collector.transfer_codec.FrameTooLargeError, match="audio.*request limit"):
|
||||
await send(ctx, images, audio=audio)
|
||||
assert ctx.requests == []
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_tiny_frames_cannot_overflow_multipart_metadata_budget(transport, monkeypatch):
|
||||
images = torch.full((2048, 1, 1, 1), 0.125, dtype=torch.float16)
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
limit = reserve + images.numel() * images.element_size()
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=limit, chunk_target=limit) as ctx:
|
||||
await send(ctx, images)
|
||||
assert len(ctx.requests) > 1
|
||||
assert all(record["bytes"] <= limit for record in ctx.requests)
|
||||
assert_exact(torch.cat([item["tensor"] for item in queued_items(ctx.queue)]), images)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("separate", [False, True])
|
||||
def test_exact_audio_is_preserved_and_only_final_completion_is_terminal(transport, monkeypatch, separate):
|
||||
images = frames(count=3)
|
||||
audio = {"waveform": torch.linspace(-0.9, 0.9, 60).reshape(1, 1, -1), "sample_rate": 48000}
|
||||
frame_bytes = images[0].numel() * images.element_size()
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
media_budget = frame_bytes if separate else frame_bytes + 8192
|
||||
limit = reserve + media_budget
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=limit, chunk_target=frame_bytes) as ctx:
|
||||
await send(ctx, images, audio=audio)
|
||||
items = queued_items(ctx.queue)
|
||||
assert len(ctx.requests) == 3 + int(separate)
|
||||
assert [item["is_last"] for item in items] == [False] * (len(items) - 1) + [True]
|
||||
assert all(item["audio"] is None for item in items[:-1])
|
||||
assert items[-1]["audio"]["sample_rate"] == audio["sample_rate"]
|
||||
assert_exact(items[-1]["audio"]["waveform"], audio["waveform"])
|
||||
assert_exact(torch.cat([item["tensor"] for item in items if item["tensor"] is not None]), images)
|
||||
assert all(record["bytes"] <= limit for record in ctx.requests)
|
||||
if separate:
|
||||
assert items[-1]["tensor"] is None
|
||||
assert ctx.requests[-1]["content_type"] == "application/json"
|
||||
|
||||
asyncio.run(scenario())
|
||||
@@ -4,6 +4,7 @@ import types
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
@@ -19,6 +20,13 @@ def _register_real_transfer_codec(package_name, utils_pkg):
|
||||
|
||||
|
||||
def _load_collector_module():
|
||||
# Restore shared dependencies even when a different test imports this loader.
|
||||
# Keep real modules imported during loading (notably extension modules) cached.
|
||||
with pytest.MonkeyPatch.context() as monkeypatch:
|
||||
return _load_collector_module_with_stubs(monkeypatch)
|
||||
|
||||
|
||||
def _load_collector_module_with_stubs(monkeypatch):
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
|
||||
package_name = "dist_collector_list_testpkg"
|
||||
|
||||
@@ -53,7 +61,7 @@ def _load_collector_module():
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
||||
sys.modules["server"] = server_module
|
||||
monkeypatch.setitem(sys.modules, "server", server_module)
|
||||
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
model_management = types.ModuleType("comfy.model_management")
|
||||
@@ -77,13 +85,13 @@ def _load_collector_module():
|
||||
|
||||
comfy_utils.ProgressBar = ProgressBar
|
||||
comfy_module.utils = comfy_utils
|
||||
sys.modules["comfy"] = comfy_module
|
||||
sys.modules["comfy.model_management"] = model_management
|
||||
sys.modules["comfy.utils"] = comfy_utils
|
||||
monkeypatch.setitem(sys.modules, "comfy", comfy_module)
|
||||
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
|
||||
monkeypatch.setitem(sys.modules, "comfy.utils", comfy_utils)
|
||||
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
monkeypatch.setitem(sys.modules, "aiohttp", aiohttp_module)
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Real Collector aggregation of typed bridge deliveries and direct results.
|
||||
|
||||
Only Comfy's host services are stubbed by the shared Collector loader. Tensor
|
||||
normalization, participant ordering, queue consumption and cleanup are real.
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import types
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from test_collector_list_inputs import _load_collector_module
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def collector_module():
|
||||
module = _load_collector_module()
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
module.prompt_server.distributed_serverless = types.SimpleNamespace(cancel=AsyncMock())
|
||||
return module
|
||||
|
||||
|
||||
def native_media(value=0):
|
||||
image = torch.arange(24, dtype=torch.float64).reshape(1, 2, 3, 4) / 997 + value
|
||||
image[0, 0, 0, 0] = -0.0
|
||||
audio = {"waveform": torch.arange(6, dtype=torch.float64).reshape(1, 2, 3) / 997 + value,
|
||||
"sample_rate": 24000}
|
||||
return {"images": image, "audio": audio}
|
||||
|
||||
|
||||
def assert_exact(actual, expected):
|
||||
assert actual.dtype == expected.dtype and actual.shape == expected.shape
|
||||
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
|
||||
|
||||
|
||||
def collect_queue(module, messages, participants, *, local=None, delegate=True):
|
||||
queue = asyncio.Queue()
|
||||
for message in messages:
|
||||
queue.put_nowait(message)
|
||||
module.prompt_server.distributed_pending_jobs["typed-job"] = queue
|
||||
values = local or {"images": None, "audio": None}
|
||||
return asyncio.run(asyncio.wait_for(module.DistributedCollectorNode().execute(
|
||||
**values, multi_job_id="typed-job", enabled_worker_ids=json.dumps(participants),
|
||||
delegate_only=delegate), timeout=2))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("delegate", [False, True])
|
||||
@pytest.mark.parametrize("socket_set", [("images",), ("audio",), ("images", "audio")])
|
||||
def test_typed_and_direct_results_preserve_batch_and_audio_order(collector_module, delegate, socket_set):
|
||||
module = collector_module
|
||||
local, typed_a, typed_b, direct = [native_media(n) for n in range(4)]
|
||||
typed = {socket: [typed_a[socket], typed_b[socket]] for socket in socket_set}
|
||||
direct_message = {"worker_id": "direct", "is_last": True}
|
||||
if "images" in socket_set:
|
||||
direct_message.update(tensor=direct["images"], image_index=0)
|
||||
if "audio" in socket_set:
|
||||
direct_message["audio"] = direct["audio"]
|
||||
messages = [direct_message, {"worker_id": "serverless:pool:1", "media": typed, "is_last": True}]
|
||||
# Delivery order differs from participant order. Duplicate configured IDs
|
||||
# must not duplicate a participant's data or leave collection waiting.
|
||||
participants = ["serverless:pool:1", "direct", "serverless:pool:1"]
|
||||
local_inputs = {socket: local[socket] if socket in socket_set else None for socket in ("images", "audio")}
|
||||
images, audio = collect_queue(module, messages, participants, local=local_inputs, delegate=delegate)
|
||||
expected = ([] if delegate else [local]) + [typed_a, typed_b, direct]
|
||||
if "images" in socket_set:
|
||||
assert_exact(images, torch.cat([row["images"] for row in expected]))
|
||||
else:
|
||||
assert images is None
|
||||
if "audio" in socket_set:
|
||||
assert audio["sample_rate"] == 24000
|
||||
assert_exact(audio["waveform"], torch.cat([row["audio"]["waveform"] for row in expected], dim=-1))
|
||||
else:
|
||||
assert torch.equal(audio["waveform"], module.DistributedCollectorNode.EMPTY_AUDIO["waveform"])
|
||||
assert module.prompt_server.distributed_pending_jobs == {}
|
||||
module.prompt_server.distributed_serverless.cancel.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("media", [{}, {"texts": ["not a Collector socket"]}, {"images": []},
|
||||
{"audio": [None]}, {"images": "not a list"},
|
||||
{"images": [torch.zeros(2, 2, 3)]}])
|
||||
def test_invalid_typed_delivery_cancels_and_removes_queue(collector_module, media):
|
||||
module = collector_module
|
||||
with pytest.raises(RuntimeError, match="media|items|IMAGE"):
|
||||
collect_queue(module, [{"worker_id": "serverless:pool:1", "media": media, "is_last": True}],
|
||||
["serverless:pool:1"])
|
||||
assert module.prompt_server.distributed_pending_jobs == {}
|
||||
module.prompt_server.distributed_serverless.cancel.assert_awaited_once_with(collector_job_id="typed-job")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["nonfinal", "unexpected", "duplicate", "provider-error", "shape"])
|
||||
def test_typed_delivery_protocol_failures_are_not_partial_success(collector_module, failure):
|
||||
media = {"images": [native_media()["images"]]}
|
||||
first = {"worker_id": "serverless:pool:1", "media": media, "is_last": True}
|
||||
messages = [first]
|
||||
participants = ["serverless:pool:1", "serverless:pool:2"]
|
||||
if failure == "nonfinal":
|
||||
first["is_last"] = False
|
||||
elif failure == "unexpected":
|
||||
first["worker_id"] = "other"
|
||||
elif failure == "duplicate":
|
||||
messages.append(dict(first))
|
||||
elif failure == "provider-error":
|
||||
messages.append({"worker_id": "serverless:pool:2", "serverless_error": "offline failure", "is_last": True})
|
||||
else:
|
||||
messages.append({"worker_id": "serverless:pool:2", "media": {"images": [torch.zeros(1, 4, 5, 3)]},
|
||||
"is_last": True})
|
||||
with pytest.raises(RuntimeError, match="atomic|Unexpected|Duplicate|failed|shape mismatch"):
|
||||
collect_queue(collector_module, messages, participants)
|
||||
assert collector_module.prompt_server.distributed_pending_jobs == {}
|
||||
collector_module.prompt_server.distributed_serverless.cancel.assert_awaited_once_with(collector_job_id="typed-job")
|
||||
@@ -0,0 +1,198 @@
|
||||
"""Real ComfyUI V3 bridge mapping, file ownership and exact CPU round trips."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def test_real_media_bridge_nodes():
|
||||
core = os.environ.get("COMFYUI_SOURCE_ROOT")
|
||||
if core is None:
|
||||
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU media bridge acceptance")
|
||||
assert core and (Path(core) / "comfy_api/latest/_input_impl/video_types.py").is_file(), (
|
||||
"Configured COMFYUI_SOURCE_ROOT must provide native video classes"
|
||||
)
|
||||
probe = r'''
|
||||
import asyncio, importlib.util, io, json, os, sys, tempfile
|
||||
from fractions import Fraction
|
||||
import av
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
import torch
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('Bridge tests must not query or initialize a GPU')
|
||||
torch.cuda.is_available = lambda: False
|
||||
torch.cuda.device_count = lambda: 0
|
||||
for name in ('init', '_lazy_init', 'current_device'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
import execution
|
||||
import folder_paths
|
||||
from comfy_api.latest._input_impl.video_types import VideoFromComponents, VideoFromFile, VideoFromList
|
||||
from comfy_api.latest._util import VideoComponents
|
||||
for name in ('is_available', 'device_count'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
spec = importlib.util.spec_from_file_location('typed_bridge_probe', Path(sys.argv[2])/'serverless_bridge/__init__.py', submodule_search_locations=[str(Path(sys.argv[2])/'serverless_bridge')])
|
||||
bridge = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = bridge
|
||||
spec.loader.exec_module(bridge)
|
||||
from typed_bridge_probe.nodes import MediaImportV1, MediaExportV1, export_values, import_value
|
||||
from typed_bridge_probe.contract import SOCKET_KINDS
|
||||
from typed_bridge_probe.video_codec import encode_video
|
||||
|
||||
def tiny_video():
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, 'w', format='matroska') as container:
|
||||
stream = container.add_stream('ffv1', rate=Fraction(24, 1))
|
||||
stream.width = stream.height = 32
|
||||
stream.pix_fmt = 'yuv420p'
|
||||
for level in (16, 64, 128):
|
||||
frame = av.VideoFrame.from_ndarray(np.full((32, 32, 3), level, dtype=np.uint8), format='rgb24')
|
||||
for packet in stream.encode(frame): container.mux(packet)
|
||||
for packet in stream.encode(): container.mux(packet)
|
||||
return buffer.getvalue()
|
||||
|
||||
def exact_tensor(actual, expected):
|
||||
assert actual.device.type == 'cpu' and actual.dtype == expected.dtype
|
||||
assert actual.shape == expected.shape
|
||||
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
|
||||
|
||||
def exact_video(actual, expected, wire):
|
||||
assert type(actual) is type(expected)
|
||||
assert encode_video(actual) == wire
|
||||
if type(actual) is VideoFromComponents:
|
||||
a = vars(actual)['_VideoFromComponents__components']
|
||||
e = vars(expected)['_VideoFromComponents__components']
|
||||
assert type(a) is VideoComponents and type(a.frame_rate) is Fraction
|
||||
assert a.frame_rate == e.frame_rate and a.metadata == e.metadata
|
||||
assert a.metadata is not e.metadata
|
||||
for name in ('images', 'alpha'): exact_tensor(getattr(a, name), getattr(e, name))
|
||||
exact_tensor(a.audio['waveform'], e.audio['waveform'])
|
||||
assert a.audio['sample_rate'] == e.audio['sample_rate']
|
||||
assert actual.get_bit_depth() == 10 and actual.get_color_space() == 'HDR PQ'
|
||||
else:
|
||||
values = actual.videos if type(actual) is VideoFromList else [actual]
|
||||
for value in values:
|
||||
buffer = vars(value)['_VideoFromFile__file']
|
||||
assert type(buffer) is io.BytesIO and not buffer.closed
|
||||
buffer.seek(0)
|
||||
with av.open(buffer) as container:
|
||||
assert len(list(container.decode(video=0))) == 3
|
||||
if type(actual) is VideoFromList:
|
||||
exact_tensor(actual.complete_audio['waveform'], expected.complete_audio['waveform'])
|
||||
|
||||
async def check():
|
||||
extension = await bridge.comfy_entrypoint()
|
||||
assert await extension.get_node_list() == [MediaImportV1, MediaExportV1]
|
||||
imp = MediaImportV1.GET_NODE_INFO_V1()
|
||||
exp = MediaExportV1.GET_NODE_INFO_V1()
|
||||
assert imp['output'] == ['IMAGE','AUDIO','MASK','STRING','VIDEO'], imp
|
||||
assert exp['output_node'] and MediaExportV1.define_schema().is_input_list, exp
|
||||
with tempfile.TemporaryDirectory(prefix='typed-media-', dir=os.environ.get('TMPDIR')) as scratch:
|
||||
base = Path(scratch)
|
||||
roots = {key: base/key for key in ('input','output','temp')}
|
||||
for path in roots.values(): path.mkdir()
|
||||
def annotated(name):
|
||||
for key, root in roots.items():
|
||||
suffix = f' [{key}]'
|
||||
if name.endswith(suffix): return str(root/name[:-len(suffix)])
|
||||
return str(roots['input']/name)
|
||||
with patch.object(folder_paths,'get_input_directory',lambda: str(roots['input'])), patch.object(folder_paths,'get_output_directory',lambda: str(roots['output'])), patch.object(folder_paths,'get_temp_directory',lambda: str(roots['temp'])), patch.object(folder_paths,'get_annotated_filepath',annotated):
|
||||
examples = {'images': torch.rand(2,2,3,4,dtype=torch.float64),
|
||||
'masks': torch.tensor([[[0.123456789,0.987654321]]],dtype=torch.float64),
|
||||
'audio': {'waveform':torch.randn(1,6,31,dtype=torch.float64),'sample_rate':44100},
|
||||
'texts': ' café\n🎥\r\n'}
|
||||
media = tiny_video()
|
||||
sources = [roots['input']/f'source-{n}.mkv' for n in range(3)]
|
||||
for path in sources: path.write_bytes(media)
|
||||
images = torch.arange(36, dtype=torch.float64, device='cpu').reshape(2,2,3,3) / 997
|
||||
images[0,0,0,0] = -0.0
|
||||
component_video = VideoFromComponents(VideoComponents(
|
||||
images=images, frame_rate=Fraction(30000,1001),
|
||||
alpha=torch.full((2,2,3,1),0.123456789,dtype=torch.float64,device='cpu'),
|
||||
audio=examples['audio'], metadata={'title':'雪', 'nested':[True,None,-4,0.125]}),
|
||||
bit_depth=10,color_space='HDR PQ')
|
||||
videos = [component_video,
|
||||
VideoFromFile(str(sources[0]),start_time=-0.0,duration=0,crop=(0,0,2,2)),
|
||||
VideoFromList([VideoFromFile(str(sources[1]),start_time=-1.25,duration=-2.5,crop=(-3,5,7,-9)),
|
||||
VideoFromFile(str(sources[2]),start_time=0,duration=0)],
|
||||
complete_audio=examples['audio'])]
|
||||
for socket,value in [*examples.items(), *[('video',v) for v in videos]]:
|
||||
kind = SOCKET_KINDS[socket]
|
||||
expected_wire = encode_video(value) if kind == 'VIDEO' else None
|
||||
binding = dict(binding_id='b-'+socket,collector_id='collector',job_id='job',participant_id='worker',node_id='export-'+socket,socket=socket,kind=kind)
|
||||
controls = {'request_id':'request','binding_json':json.dumps(binding)}
|
||||
data, missing, v3 = execution.get_input_data(controls,MediaExportV1,'export-'+socket)
|
||||
assert not missing, missing
|
||||
data[socket] = [value]
|
||||
results = await execution._async_map_node_over_list('probe','export-'+socket,MediaExportV1,data,MediaExportV1.FUNCTION,v3_data=v3)
|
||||
results = await execution.resolve_map_node_over_list_results(results)
|
||||
raw = results[0].ui
|
||||
exported = raw['media_exports'][0]
|
||||
assert exported['binding']['count'] == 1
|
||||
artifact = exported['artifacts'][0]
|
||||
transport = artifact.pop('transport')
|
||||
filename = transport['subfolder']+'/'+transport['filename']+' [temp]'
|
||||
decoded = import_value(artifact,filename)
|
||||
if kind in ('IMAGE','MASK'):
|
||||
assert decoded.dtype == value.dtype and torch.equal(decoded,value)
|
||||
elif kind == 'AUDIO':
|
||||
assert decoded['sample_rate'] == value['sample_rate']
|
||||
assert torch.equal(decoded['waveform'],value['waveform'])
|
||||
elif kind == 'VIDEO': exact_video(decoded,value,expected_wire)
|
||||
else: assert decoded == value
|
||||
mapped, missing, hidden = execution.get_input_data({'artifact_json':json.dumps(artifact),'filename':filename},MediaImportV1,'import')
|
||||
got = await execution._async_map_node_over_list('probe','import',MediaImportV1,mapped,MediaImportV1.FUNCTION,v3_data=hidden)
|
||||
got = await execution.resolve_map_node_over_list_results(got)
|
||||
assert len(got[0].result) == 5
|
||||
assert all(x is None for n,x in enumerate(got[0].result) if n != list(SOCKET_KINDS).index(socket))
|
||||
bad = dict(artifact,sha256='0'*64)
|
||||
try: import_value(bad,filename)
|
||||
except ValueError: pass
|
||||
else: raise AssertionError('checksum accepted')
|
||||
try: import_value(artifact,str(base/'outside.bin'))
|
||||
except ValueError: pass
|
||||
else: raise AssertionError('outside path accepted')
|
||||
if kind == 'VIDEO':
|
||||
mapped_video = got[0].result[list(SOCKET_KINDS).index(socket)]
|
||||
# Remove both the exported artifact and the original input
|
||||
# containers before checking independently owned BytesIOs.
|
||||
Path(annotated(filename)).unlink()
|
||||
originals = value.videos if type(value) is VideoFromList else [value]
|
||||
for original in originals:
|
||||
if type(original) is VideoFromFile:
|
||||
Path(vars(original)['_VideoFromFile__file']).unlink()
|
||||
exact_video(decoded,value,expected_wire)
|
||||
exact_video(mapped_video,value,expected_wire)
|
||||
if type(value) is VideoFromFile:
|
||||
assert tuple(decoded.get_components().images.shape) == (3,2,2,3)
|
||||
text_binding = dict(binding_id='empty',collector_id='collector',job_id='job',participant_id='worker',node_id='export-empty',socket='texts',kind='STRING')
|
||||
out = export_values('request',text_binding,[''])
|
||||
assert out['artifacts'][0]['byte_length'] == 0
|
||||
before = set(roots['temp'].rglob('*.bin'))
|
||||
try: export_values('request',text_binding,['first',None])
|
||||
except ValueError: pass
|
||||
else: raise AssertionError('invalid text accepted')
|
||||
assert set(roots['temp'].rglob('*.bin')) == before
|
||||
print('MEDIA_BRIDGE_CPU_OK')
|
||||
asyncio.run(check())
|
||||
'''
|
||||
env = dict(os.environ)
|
||||
env.pop("PYTHONPATH", None)
|
||||
env.pop("COMFYUI_SOURCE_ROOT", None)
|
||||
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="", PYTHONDONTWRITEBYTECODE="1",
|
||||
TMPDIR=str(scratch))
|
||||
result = subprocess.run([sys.executable, "-c", probe, str(Path(core).resolve()), str(ROOT)],
|
||||
capture_output=True, text=True, timeout=90, env=env)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert "MEDIA_BRIDGE_CPU_OK" in result.stdout
|
||||
@@ -0,0 +1,328 @@
|
||||
import importlib.util
|
||||
import math
|
||||
import struct
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
MODULE_PATH = Path(__file__).resolve().parents[1] / "serverless_bridge" / "codec.py"
|
||||
spec = importlib.util.spec_from_file_location("dist_media_codec_test", MODULE_PATH)
|
||||
codec = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(codec)
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
|
||||
SHAPES = [
|
||||
("IMAGE", (2, 3, 5, 1), "BHWC"),
|
||||
("IMAGE", (2, 3, 5, 3), "BHWC"),
|
||||
("IMAGE", (2, 3, 5, 4), "BHWC"),
|
||||
("MASK", (2, 3, 5), "BHW"),
|
||||
("AUDIO", (2, 1, 17), "BCT"),
|
||||
("AUDIO", (2, 6, 17), "BCT"),
|
||||
]
|
||||
|
||||
|
||||
def audio(tensor, sample_rate=48000):
|
||||
return {"waveform": tensor, "sample_rate": sample_rate}
|
||||
|
||||
|
||||
def descriptor(_kind="MASK", **changes):
|
||||
result = {
|
||||
"encoding": "tensor-le-v1", "dtype": "float32",
|
||||
"shape": [1, 1, 1], "layout": "BHW",
|
||||
}
|
||||
if _kind == "IMAGE":
|
||||
result.update(shape=[1, 1, 1, 3], layout="BHWC")
|
||||
elif _kind == "AUDIO":
|
||||
result.update(layout="BCT", sample_rate=48000)
|
||||
result.update(changes)
|
||||
return result
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def forbid_gpu_queries(monkeypatch):
|
||||
def forbidden(*args, **kwargs):
|
||||
pytest.fail("codec must not query or initialize a GPU")
|
||||
for name in ("is_available", "init", "_lazy_init", "device_count"):
|
||||
monkeypatch.setattr(torch.cuda, name, forbidden)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("kind,shape,layout", SHAPES)
|
||||
def test_exact_roundtrip_preserves_dtype_shape_channels_and_bits(dtype, kind, shape, layout):
|
||||
count = math.prod(shape)
|
||||
tensor = (torch.arange(count, device="cpu", dtype=torch.float64) / 997 - 0.125).to(dtype).reshape(shape)
|
||||
tensor.flatten()[0] = -0.0
|
||||
tensor.requires_grad_(True)
|
||||
value = audio(tensor) if kind == "AUDIO" else tensor
|
||||
|
||||
metadata, payload = codec.encode_payload(kind, value)
|
||||
expected = {
|
||||
"encoding": "tensor-le-v1", "dtype": str(dtype).removeprefix("torch."),
|
||||
"shape": list(shape), "layout": layout,
|
||||
}
|
||||
if kind == "AUDIO":
|
||||
expected["sample_rate"] = 48000
|
||||
assert metadata == expected
|
||||
assert type(payload) is bytes
|
||||
assert len(payload) == count * tensor.element_size()
|
||||
result = codec.decode_payload(kind, metadata, payload)
|
||||
if kind == "AUDIO":
|
||||
assert set(result) == {"waveform", "sample_rate"}
|
||||
assert result["sample_rate"] == 48000
|
||||
result = result["waveform"]
|
||||
assert result.device.type == "cpu"
|
||||
assert result.is_contiguous()
|
||||
assert not result.requires_grad
|
||||
assert result.dtype == dtype
|
||||
assert result.shape == shape
|
||||
assert torch.equal(result.view(torch.uint8), tensor.detach().view(torch.uint8))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype,fmt,expected_bits", [
|
||||
(torch.float16, "e", None), (torch.bfloat16, "H", [0x3F80, 0x8000, 0xC020]),
|
||||
(torch.float32, "f", None), (torch.float64, "d", None),
|
||||
])
|
||||
def test_wire_bytes_are_little_endian(dtype, fmt, expected_bits):
|
||||
values = [1.0, -0.0, -2.5]
|
||||
tensor = torch.tensor(values, dtype=dtype, device="cpu").reshape(1, 1, 3)
|
||||
metadata, payload = codec.encode_payload("MASK", tensor)
|
||||
expected = struct.pack("<3" + fmt, *(expected_bits if expected_bits else values))
|
||||
assert payload == expected
|
||||
assert torch.equal(codec.decode_payload("MASK", metadata, expected).view(torch.uint8), tensor.view(torch.uint8))
|
||||
|
||||
|
||||
def test_bfloat16_extremes_and_subnormal_bits_are_lossless():
|
||||
bits = torch.tensor([0, 0x8000, 1, 0x8001, 0x3F81, 0x7F7F, 0xFF7F], dtype=torch.uint16, device="cpu")
|
||||
value = bits.view(torch.bfloat16).reshape(1, 1, -1)
|
||||
metadata, payload = codec.encode_payload("MASK", value)
|
||||
assert payload == struct.pack("<7H", *bits.tolist())
|
||||
assert torch.equal(codec.decode_payload("MASK", metadata, payload).view(torch.uint16).flatten(), bits)
|
||||
|
||||
|
||||
def test_fractional_mask_is_not_quantized_clipped_or_normalized():
|
||||
value = torch.tensor([0.1234567890123, 1.0000000000001, -2.5], dtype=torch.float64, device="cpu").reshape(1, 1, 3)
|
||||
metadata, payload = codec.encode_payload("MASK", value)
|
||||
assert torch.equal(codec.decode_payload("MASK", metadata, payload), value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text", ["", " ", "\t\n ", "café 雪 🎵", "e\u0301", "\x00\ufeff\r\n"])
|
||||
def test_strings_preserve_exact_utf8(text):
|
||||
metadata, payload = codec.encode_payload("STRING", text)
|
||||
assert metadata == {"encoding": "utf8-v1"}
|
||||
assert payload == text.encode("utf-8")
|
||||
assert codec.decode_payload("STRING", metadata, payload) == text
|
||||
|
||||
|
||||
def test_noncontiguous_and_lazy_negative_tensors_roundtrip():
|
||||
value = torch.arange(24, dtype=torch.float32, device="cpu").reshape(2, 3, 4).transpose(1, 2)
|
||||
for tensor in (value, torch._neg_view(value)):
|
||||
metadata, payload = codec.encode_payload("MASK", tensor)
|
||||
result = codec.decode_payload("MASK", metadata, payload)
|
||||
assert result.is_contiguous()
|
||||
assert torch.equal(result, tensor)
|
||||
|
||||
|
||||
def test_decode_is_independent_and_ignores_default_device(monkeypatch):
|
||||
source = torch.tensor([0.25], device="cpu").reshape(1, 1, 1)
|
||||
metadata, payload = codec.encode_payload("MASK", source)
|
||||
buffers = []
|
||||
original = torch.frombuffer
|
||||
def capture(buffer, *args, **kwargs):
|
||||
buffers.append(buffer)
|
||||
return original(buffer, *args, **kwargs)
|
||||
monkeypatch.setattr(torch, "frombuffer", capture)
|
||||
with torch.device("meta"):
|
||||
first = codec.decode_payload("MASK", metadata, payload)
|
||||
second = codec.decode_payload("MASK", metadata, payload)
|
||||
assert first.device.type == second.device.type == "cpu"
|
||||
first.fill_(9)
|
||||
source.fill_(7)
|
||||
for buffer in buffers:
|
||||
if isinstance(buffer, bytearray):
|
||||
buffer[:] = b"\x00" * len(buffer)
|
||||
assert second.item() == 0.25
|
||||
assert codec.decode_payload("MASK", metadata, payload).item() == 0.25
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", ["image", " MASK", "LATENT", "", None, 1, [], {}])
|
||||
def test_unknown_kinds_raise_value_error(kind):
|
||||
with pytest.raises(ValueError, match="kind"):
|
||||
codec.encode_payload(kind, "")
|
||||
with pytest.raises(ValueError, match="kind"):
|
||||
codec.decode_payload(kind, {}, b"")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind,value", [
|
||||
("STRING", None), ("STRING", b"text"), ("STRING", 123),
|
||||
("IMAGE", [[[[0.5]]]]), ("MASK", None),
|
||||
("AUDIO", torch.ones(1, 1, 1, device="cpu")), ("AUDIO", {}),
|
||||
("AUDIO", {"waveform": "invalid", "sample_rate": 48000}),
|
||||
("AUDIO", {"waveform": torch.ones(1, 1, 1, device="cpu")}),
|
||||
])
|
||||
def test_unsupported_values_raise_value_error(kind, value):
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_payload(kind, value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.uint8, torch.int64, torch.bool, torch.complex64])
|
||||
def test_unsupported_tensor_dtype(dtype):
|
||||
with pytest.raises(ValueError, match="dtype"):
|
||||
codec.encode_payload("MASK", torch.ones(1, 1, 1, dtype=dtype, device="cpu"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind,shape", [
|
||||
("IMAGE", (1, 2, 3)), ("IMAGE", (1, 1, 1, 2)), ("IMAGE", (1, 1, 1, 5)),
|
||||
("MASK", (1, 2)), ("MASK", (1, 2, 3, 1)), ("AUDIO", (2, 4)),
|
||||
("IMAGE", (1, 0, 2, 3)), ("MASK", (0, 1, 1)), ("AUDIO", (1, 0, 4)),
|
||||
("AUDIO", (1, 2, 0)),
|
||||
])
|
||||
def test_invalid_tensor_shape(kind, shape):
|
||||
tensor = torch.empty(shape, dtype=torch.float32, device="cpu")
|
||||
with pytest.raises(ValueError, match="shape|rank|channel|dimension"):
|
||||
codec.encode_payload(kind, audio(tensor) if kind == "AUDIO" else tensor)
|
||||
|
||||
|
||||
def test_rejects_non_cpu_and_sparse_without_transferring(monkeypatch):
|
||||
meta = torch.empty((1, 1, 1), device="meta")
|
||||
sparse = torch.ones((1, 1, 1), device="cpu").to_sparse()
|
||||
def forbidden(*args, **kwargs):
|
||||
pytest.fail("codec must not move tensors between devices")
|
||||
monkeypatch.setattr(torch.Tensor, "cpu", forbidden)
|
||||
monkeypatch.setattr(torch.Tensor, "to", forbidden)
|
||||
for tensor in (meta, sparse):
|
||||
with pytest.raises(ValueError, match="CPU|layout|strided"):
|
||||
codec.encode_payload("MASK", tensor)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sample_rate", [True, False, 0, -1, 44100.0, "44100", None])
|
||||
def test_invalid_sample_rate(sample_rate):
|
||||
with pytest.raises(ValueError, match="sample_rate"):
|
||||
codec.encode_payload("AUDIO", audio(torch.ones(1, 1, 1, device="cpu"), sample_rate))
|
||||
with pytest.raises(ValueError, match="sample_rate"):
|
||||
codec.decode_payload("AUDIO", descriptor("AUDIO", sample_rate=sample_rate), b"\x00" * 4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [None, [], "descriptor", {}, {"encoding": "utf8-v1"}])
|
||||
def test_malformed_tensor_descriptor(bad):
|
||||
with pytest.raises(ValueError, match="descriptor"):
|
||||
codec.decode_payload("MASK", bad, b"\x00" * 4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("changes", [
|
||||
{"encoding": "pickle"}, {"encoding": None}, {"encoding": []},
|
||||
{"dtype": "uint8"}, {"dtype": torch.float32}, {"dtype": []},
|
||||
{"layout": "BHWC"}, {"layout": None}, {"layout": []},
|
||||
{"shape": (1, 1, 1)}, {"shape": None}, {"shape": "1,1,1"},
|
||||
{"shape": [1, 1]}, {"shape": [1, 1, 1, 1]}, {"shape": [1, 0, 1]},
|
||||
{"shape": [1, -1, 1]}, {"shape": [1, True, 1]}, {"shape": [1, 1.0, 1]},
|
||||
{"shape": [1, "1", 1]}, {"shape": [1, [], 1]}, {"kind": "MASK"},
|
||||
])
|
||||
def test_invalid_tensor_descriptor_fields(changes):
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_payload("MASK", descriptor(**changes), b"\x00" * 4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metadata", [{}, {"encoding": "tensor-le-v1"}, {"encoding": "utf8-v1", "kind": "STRING"}, {"encoding": []}])
|
||||
def test_invalid_string_descriptor(metadata):
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_payload("STRING", metadata, b"")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [b"", b"\x00" * 3, b"\x00" * 5, b"\x00" * 8])
|
||||
def test_shape_length_mismatch(payload):
|
||||
with pytest.raises(ValueError, match="length"):
|
||||
codec.decode_payload("MASK", descriptor(), payload)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [None, "text", [0, 0, 0, 0], 4])
|
||||
def test_nonbyte_payload_rejected(payload):
|
||||
with pytest.raises(ValueError, match="bytes"):
|
||||
codec.decode_payload("MASK", descriptor(), payload)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")])
|
||||
def test_nonfinite_values_fail_on_encode_and_decode(dtype, value):
|
||||
tensor = torch.tensor([value], dtype=dtype, device="cpu").reshape(1, 1, 1)
|
||||
with pytest.raises(ValueError, match="finite"):
|
||||
codec.encode_payload("MASK", tensor)
|
||||
fmt = {torch.float16: "e", torch.float32: "f", torch.float64: "d"}
|
||||
if dtype == torch.bfloat16:
|
||||
payload = struct.pack("<H", 0x7FC0 if math.isnan(value) else (0x7F80 if value > 0 else 0xFF80))
|
||||
else:
|
||||
payload = struct.pack("<" + fmt[dtype], value)
|
||||
metadata = descriptor(dtype=str(dtype).removeprefix("torch."))
|
||||
with pytest.raises(ValueError, match="finite"):
|
||||
codec.decode_payload("MASK", metadata, payload)
|
||||
|
||||
|
||||
def test_invalid_utf8_errors_do_not_echo_content():
|
||||
with pytest.raises(ValueError) as error:
|
||||
codec.encode_payload("STRING", "private-value\ud800")
|
||||
assert "private-value" not in str(error.value)
|
||||
with pytest.raises(ValueError) as error:
|
||||
codec.decode_payload("STRING", {"encoding": "utf8-v1"}, b"private-value\xff")
|
||||
assert "private-value" not in str(error.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("max_bytes", [True, False, -1, 4.0, "4", None])
|
||||
def test_invalid_byte_budget(max_bytes):
|
||||
with pytest.raises(ValueError, match="max_bytes"):
|
||||
codec.encode_payload("STRING", "", max_bytes=max_bytes)
|
||||
with pytest.raises(ValueError, match="max_bytes"):
|
||||
codec.decode_payload("STRING", {"encoding": "utf8-v1"}, b"", max_bytes=max_bytes)
|
||||
|
||||
|
||||
def test_exact_byte_budget_and_zero_budget_for_empty_string():
|
||||
tensor = torch.ones(1, 1, 1, device="cpu")
|
||||
metadata, payload = codec.encode_payload("MASK", tensor, max_bytes=4)
|
||||
assert codec.decode_payload("MASK", metadata, payload, max_bytes=4).item() == 1
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.encode_payload("MASK", tensor, max_bytes=3)
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.decode_payload("MASK", metadata, payload, max_bytes=3)
|
||||
metadata, payload = codec.encode_payload("STRING", "", max_bytes=0)
|
||||
assert codec.decode_payload("STRING", metadata, payload, max_bytes=0) == ""
|
||||
for text in ("a", "é", "🎵"):
|
||||
byte_count = len(text.encode("utf-8"))
|
||||
metadata, payload = codec.encode_payload("STRING", text, max_bytes=byte_count)
|
||||
assert codec.decode_payload("STRING", metadata, payload, max_bytes=byte_count) == text
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.encode_payload("STRING", text, max_bytes=byte_count - 1)
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.decode_payload("STRING", metadata, payload, max_bytes=byte_count - 1)
|
||||
|
||||
|
||||
def test_oversized_shape_rejected_before_tensor_allocation(monkeypatch):
|
||||
assert codec.MAX_PAYLOAD_BYTES == 256 * 1024 * 1024
|
||||
def forbidden(*args, **kwargs):
|
||||
pytest.fail("oversized shape must fail before tensor allocation")
|
||||
for name in ("frombuffer", "empty", "tensor", "zeros", "isfinite"):
|
||||
monkeypatch.setattr(torch, name, forbidden)
|
||||
for shape in ([1, 1, codec.MAX_PAYLOAD_BYTES], [2**100, 2**100, 1]):
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.decode_payload("MASK", descriptor(shape=shape), b"\x00" * 4)
|
||||
|
||||
|
||||
def test_oversized_encode_rejected_before_materialization(monkeypatch):
|
||||
huge_view = torch.ones(1, device="cpu").expand(1, 1, codec.MAX_PAYLOAD_BYTES)
|
||||
def forbidden(*args, **kwargs):
|
||||
pytest.fail("oversized tensor must fail before materialization or scanning")
|
||||
monkeypatch.setattr(torch.Tensor, "contiguous", forbidden)
|
||||
monkeypatch.setattr(torch, "isfinite", forbidden)
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.encode_payload("MASK", huge_view)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("itemsize", [2, 4, 8])
|
||||
def test_native_little_endian_conversion_handles_big_endian(monkeypatch, itemsize):
|
||||
raw = bytes(range(itemsize * 3))
|
||||
monkeypatch.setattr(codec, "sys", types.SimpleNamespace(byteorder="little"))
|
||||
assert codec._native_little_endian(raw, itemsize) == raw
|
||||
monkeypatch.setattr(codec, "sys", types.SimpleNamespace(byteorder="big"))
|
||||
expected = b"".join(raw[i:i + itemsize][::-1] for i in range(0, len(raw), itemsize))
|
||||
assert codec._native_little_endian(raw, itemsize) == expected
|
||||
assert codec._native_little_endian(expected, itemsize) == raw
|
||||
@@ -0,0 +1,152 @@
|
||||
"""Offline contract tests: malformed results never reach transport/delivery."""
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
spec = importlib.util.spec_from_file_location("media_contract_test", ROOT / "serverless_bridge" / "contract.py")
|
||||
media = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(media)
|
||||
|
||||
|
||||
def binding(**changes):
|
||||
return dict(binding_id="binding-1", collector_id="100", job_id="job-1",
|
||||
participant_id="serverless:worker:0", node_id="export-1", socket="texts", kind="STRING", **changes)
|
||||
|
||||
|
||||
def response():
|
||||
artifact = media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b" \n")
|
||||
return {"version": 1, "profile": media.PROFILE, "request_id": "request-1",
|
||||
"bindings": [dict(binding(), count=1)], "artifacts": [artifact],
|
||||
"items": [{"binding_id": "binding-1", "item_index": 0, "artifact_id": "a"}]}
|
||||
|
||||
|
||||
def validate(value):
|
||||
return media.validate_response(value, request_id="request-1", bindings=[binding()])
|
||||
|
||||
|
||||
def test_valid_manifest_and_exact_inline_payload():
|
||||
value = response()
|
||||
ordered = validate(value)
|
||||
assert asyncio.run(media.resolve_artifact(ordered["binding-1"][0])) == b" \n"
|
||||
|
||||
|
||||
def test_order_uses_indices_not_arrival():
|
||||
value = response()
|
||||
value["bindings"][0]["count"] = 2
|
||||
value["artifacts"].append(media.artifact_from_payload("b", "STRING", {"encoding": "utf8-v1"}, b"second"))
|
||||
value["items"] = [{"binding_id": "binding-1", "item_index": 1, "artifact_id": "b"}, value["items"][0]]
|
||||
assert [item["id"] for item in validate(value)["binding-1"]] == ["a", "b"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field,value", [("version", True), ("version", 2), ("profile", "legacy"), ("request_id", "another")])
|
||||
def test_envelope_identity(field, value):
|
||||
obj = response()
|
||||
obj[field] = value
|
||||
with pytest.raises(ValueError):
|
||||
validate(obj)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["collector_id", "participant_id", "job_id", "node_id", "binding_id"])
|
||||
def test_binding_identity(field):
|
||||
value = response()
|
||||
value["bindings"][0][field] = "another"
|
||||
with pytest.raises(ValueError):
|
||||
validate(value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("count", [0, -1, True, 1.0, "1", 65])
|
||||
def test_invalid_count(count):
|
||||
value = response()
|
||||
value["bindings"][0]["count"] = count
|
||||
with pytest.raises(ValueError):
|
||||
validate(value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", ["missing_binding", "missing_item", "extra_item", "duplicate_binding", "duplicate_artifact", "unused_artifact", "wrong_type", "wrong_index"])
|
||||
def test_incomplete_or_ambiguous_manifest(mutation):
|
||||
value = response()
|
||||
if mutation == "missing_binding":
|
||||
value["bindings"] = []
|
||||
elif mutation == "missing_item":
|
||||
value["items"] = []
|
||||
elif mutation == "extra_item":
|
||||
value["items"] *= 2
|
||||
elif mutation == "duplicate_binding":
|
||||
value["bindings"] *= 2
|
||||
elif mutation == "duplicate_artifact":
|
||||
value["artifacts"] *= 2
|
||||
elif mutation == "unused_artifact":
|
||||
artifact = copy.deepcopy(value["artifacts"][0])
|
||||
artifact["id"] = "unused"
|
||||
value["artifacts"].append(artifact)
|
||||
elif mutation == "wrong_type":
|
||||
value["artifacts"][0]["kind"] = "MASK"
|
||||
elif mutation == "wrong_index":
|
||||
value["items"][0]["item_index"] = True
|
||||
with pytest.raises(ValueError):
|
||||
validate(value)
|
||||
|
||||
|
||||
def test_duplicate_item_indices():
|
||||
value = response()
|
||||
value["bindings"][0]["count"] = 2
|
||||
value["items"] *= 2
|
||||
with pytest.raises(ValueError, match="Duplicate"):
|
||||
validate(value)
|
||||
|
||||
|
||||
def test_total_budget():
|
||||
with pytest.raises(ValueError, match="total byte"):
|
||||
media.validate_response(response(), request_id="request-1", bindings=[binding()], max_total_bytes=1)
|
||||
|
||||
|
||||
def test_checksum_and_base64_validation():
|
||||
artifact = response()["artifacts"][0]
|
||||
artifact["sha256"] = "0" * 64
|
||||
with pytest.raises(ValueError, match="checksum"):
|
||||
asyncio.run(media.resolve_artifact(artifact))
|
||||
artifact["transport"]["data_base64"] = "!!!!"
|
||||
with pytest.raises(ValueError, match="base64"):
|
||||
asyncio.run(media.resolve_artifact(artifact))
|
||||
|
||||
|
||||
def test_empty_text_is_a_value():
|
||||
artifact = media.artifact_from_payload("empty", "STRING", {"encoding": "utf8-v1"}, b"")
|
||||
assert asyncio.run(media.resolve_artifact(artifact)) == b""
|
||||
|
||||
|
||||
@pytest.mark.parametrize("url,expires", [("http://example.com/a", "2099-01-01T00:00:00Z"),
|
||||
("https://user:secret@example.com/a", "2099-01-01T00:00:00Z"),
|
||||
("https://example.com/a", "2000-01-01T00:00:00Z"),
|
||||
("https://example.com/a", "2099-01-01"),
|
||||
("https://example.com:1234/a", "2099-01-01T00:00:00Z")])
|
||||
def test_invalid_url_envelopes(url, expires):
|
||||
with pytest.raises(ValueError):
|
||||
media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b"a",
|
||||
transport={"type": "url", "url": url, "expires_at": expires})
|
||||
|
||||
|
||||
def test_url_resolution_is_bounded_and_hash_verified():
|
||||
artifact = media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b"hello",
|
||||
transport={"type": "url", "url": "https://example.com/a", "expires_at": "2099-01-01T00:00:00Z"})
|
||||
calls = []
|
||||
async def download(url, *, max_bytes):
|
||||
calls.append((url, max_bytes))
|
||||
return b"hello"
|
||||
assert asyncio.run(media.resolve_artifact(artifact, download=download)) == b"hello"
|
||||
assert calls == [("https://example.com/a", 5)]
|
||||
|
||||
|
||||
def test_descriptor_depth_and_nonfinite_rejected():
|
||||
value = {}
|
||||
cursor = value
|
||||
for _ in range(15):
|
||||
cursor["next"] = {}
|
||||
cursor = cursor["next"]
|
||||
for descriptor in [value, {"x": float("nan")}, {"x": object()}]:
|
||||
with pytest.raises(ValueError):
|
||||
media.artifact_from_payload("a", "STRING", descriptor, b"")
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Opt-in CPU acceptance of Collector IMAGE/AUDIO and the five-kind bridge.
|
||||
|
||||
COMFYUI_SOURCE_ROOT must name an existing compatible checkout; run with that
|
||||
checkout's Python environment. No installation, normal instance, or provider is
|
||||
used. The child owns its loopback server and every writable Comfy directory.
|
||||
"""
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def test_typed_media_end_to_end_cpu():
|
||||
configured = os.environ.get("COMFYUI_SOURCE_ROOT")
|
||||
if configured is None:
|
||||
pytest.skip("Set COMFYUI_SOURCE_ROOT for real CPU typed-media end-to-end acceptance")
|
||||
core = Path(configured).expanduser().resolve()
|
||||
assert configured and all((core / name).is_file() for name in (
|
||||
"execution.py", "server.py", "comfy/cli_args.py",
|
||||
"comfy_api/latest/_input_impl/video_types.py",
|
||||
)), "Configured COMFYUI_SOURCE_ROOT must provide PromptExecutor and native V3 VIDEO"
|
||||
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
env = dict(os.environ)
|
||||
for name in ("PYTHONPATH", "COMFYUI_SOURCE_ROOT"):
|
||||
env.pop(name, None)
|
||||
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="", ROCR_VISIBLE_DEVICES="",
|
||||
PYTHONDONTWRITEBYTECODE="1", PYTHONHASHSEED="0", COMFYUI_IS_WORKER="1",
|
||||
OMP_NUM_THREADS="1", MKL_NUM_THREADS="1")
|
||||
with tempfile.TemporaryDirectory(prefix="media-e2e-", dir=scratch) as run_root:
|
||||
env["TMPDIR"] = run_root
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(ROOT / "tests/helpers/media_end_to_end_check.py"),
|
||||
str(core), run_root], cwd=run_root, env=env,
|
||||
capture_output=True, text=True, timeout=150,
|
||||
)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
for marker in ("NATIVE_INPUTS_AND_DIRECT_EXPORT_REFS_OK", "REAL_HTTP_BRIDGE_RESULTS_OK",
|
||||
"MASTER_AND_DELEGATE_ORDER_OK", "ATOMIC_DECODE_REJECTION_OK",
|
||||
"IMAGE_TRANSFER_HTTP_OK",
|
||||
"MEDIA_END_TO_END_CPU_OK"):
|
||||
assert marker in result.stdout, result.stdout + result.stderr
|
||||
print(result.stdout)
|
||||
@@ -0,0 +1,877 @@
|
||||
"""CPU-only endpoint adapter tests against an actual loopback Comfy HTTP fake."""
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import hashlib
|
||||
import importlib
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE = "media_handler_testpkg"
|
||||
package = types.ModuleType(PACKAGE)
|
||||
package.__path__ = [str(ROOT / "serverless_bridge")]
|
||||
sys.modules[PACKAGE] = package
|
||||
handler = importlib.import_module(PACKAGE + ".handler")
|
||||
contract = importlib.import_module(PACKAGE + ".contract")
|
||||
|
||||
|
||||
def request(kind="STRING", *, incoming=True):
|
||||
socket = {value: key for key, value in contract.SOCKET_KINDS.items()}[kind]
|
||||
binding = dict(binding_id="binding:1", collector_id="collector:1", job_id="job/42",
|
||||
participant_id="worker:0", node_id="export:1", socket=socket, kind=kind)
|
||||
artifact = contract.artifact_from_payload("input:1", kind, {"encoding": "test"}, b"input bytes")
|
||||
return {"workflow": {
|
||||
"source:1": {"class_type": "Source", "inputs": {"value": "literal"}},
|
||||
"export:1": {"class_type": "MediaExportV1", "inputs": {
|
||||
"request_id": "request:1", "binding_json": json.dumps(binding), socket: ["source:1", 0]}},
|
||||
}, "media": {"version": 1, "profile": contract.PROFILE, "request_id": "request:1",
|
||||
"artifacts": [artifact] if incoming else [],
|
||||
"inputs": [{"node_id": "source:1", "input_name": "value", "artifact_id": "input:1"}] if incoming else [],
|
||||
"outputs": [binding]}}
|
||||
|
||||
|
||||
class Comfy:
|
||||
def __init__(self, root, data=None, *, payload=b"result bytes"):
|
||||
self.root = root
|
||||
self.data = data or request()
|
||||
self.payload = payload
|
||||
self.submitted = []
|
||||
self.views = []
|
||||
self.deletes = []
|
||||
self.interrupts = []
|
||||
self.polls = 0
|
||||
self.active = False
|
||||
self.mode = "success"
|
||||
self.history_failures = 0
|
||||
self.record_mutation = None
|
||||
self.prompt_seen = asyncio.Event()
|
||||
self.terminal_seen = False
|
||||
|
||||
def staged(self):
|
||||
return list(self.root.glob("distributed-media-v1/*/*.bin"))
|
||||
|
||||
def record(self):
|
||||
binding = self.data["media"]["outputs"][0]
|
||||
aid = "a" * 32
|
||||
artifact = contract.artifact_from_payload(aid, binding["kind"], {"encoding": "test"}, self.payload)
|
||||
artifact["transport"] = {"type": "comfy-file", "filename": aid + ".bin",
|
||||
"subfolder": "distributed-media-v1/" + hashlib.sha256(b"request:1").hexdigest(),
|
||||
"folder_type": "temp"}
|
||||
value = {"request_id": "request:1", "binding": dict(binding, count=1), "artifacts": [artifact],
|
||||
"items": [{"binding_id": binding["binding_id"], "item_index": 0, "artifact_id": aid}]}
|
||||
if self.record_mutation:
|
||||
self.record_mutation(value)
|
||||
return value
|
||||
|
||||
async def submit(self, req):
|
||||
data = await req.json()
|
||||
self.submitted.append(data)
|
||||
self.active = True
|
||||
self.prompt_seen.set()
|
||||
if self.data["media"]["artifacts"]:
|
||||
assert self.staged()
|
||||
link = data["prompt"]["source:1"]["inputs"]["value"]
|
||||
imported = data["prompt"][link[0]]
|
||||
assert imported["class_type"] == "MediaImportV1"
|
||||
metadata = json.loads(imported["inputs"]["artifact_json"])
|
||||
assert "transport" not in metadata
|
||||
assert (self.root / imported["inputs"]["filename"]).read_bytes() == b"input bytes"
|
||||
if self.mode == "ambiguous":
|
||||
return web.Response(status=500)
|
||||
if self.mode == "partial":
|
||||
return web.json_response({"prompt_id": "prompt:1", "node_errors": {"export:1": {"errors": ["bad"]}}})
|
||||
return web.json_response({"prompt_id": "prompt:1", "node_errors": {}})
|
||||
|
||||
async def history(self, req):
|
||||
self.polls += 1
|
||||
if self.history_failures:
|
||||
self.history_failures -= 1
|
||||
return web.Response(status=503)
|
||||
if self.mode in {"queued", "running", "unrelated", "ambiguous"}:
|
||||
return web.json_response({})
|
||||
if self.data["media"]["artifacts"]:
|
||||
assert self.staged(), "staging must survive until terminal confirmation"
|
||||
self.active = False
|
||||
self.terminal_seen = True
|
||||
if self.mode in {"error", "interrupted"}:
|
||||
status = {"completed": False, "status_str": "error", "messages": [["execution_interrupted", {}]]}
|
||||
else:
|
||||
status = {"completed": True, "status_str": "success", "messages": []}
|
||||
return web.json_response({"prompt:1": {"status": status, "outputs": {} if self.mode == "missing" else {
|
||||
"export:1": {"media_exports": [self.record()]}}}})
|
||||
|
||||
async def view(self, req):
|
||||
self.views.append(dict(req.query))
|
||||
assert self.terminal_seen
|
||||
payload = self.payload if self.mode != "corrupt" else b"X" * len(self.payload)
|
||||
if self.mode == "oversized":
|
||||
payload += b"X"
|
||||
return web.Response(body=payload)
|
||||
|
||||
def entry(self, pid="prompt:1"):
|
||||
extra = dict(self.submitted[0].get("extra_data", {}), client_id=self.submitted[0]["client_id"])
|
||||
return [1, pid, self.submitted[0]["prompt"], extra, ["export:1"]]
|
||||
|
||||
async def queue(self, req):
|
||||
if req.method == "POST":
|
||||
body = await req.json()
|
||||
assert body == {"delete": ["prompt:1"]}
|
||||
self.deletes.append(body)
|
||||
assert self.staged()
|
||||
self.mode = "deleted"
|
||||
self.active = False
|
||||
return web.Response(status=200)
|
||||
running = [self.entry()] if self.mode == "running" else []
|
||||
if self.mode == "unrelated":
|
||||
running = [self.entry("not-our-prompt")]
|
||||
pending = [self.entry()] if self.mode in {"queued", "ambiguous"} else []
|
||||
return web.json_response({"queue_running": running, "queue_pending": pending})
|
||||
|
||||
async def interrupt(self, req):
|
||||
assert self.mode == "running"
|
||||
assert self.staged()
|
||||
self.interrupts.append(await req.json())
|
||||
self.mode = "interrupted"
|
||||
return web.Response(status=200)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def serve(comfy):
|
||||
app = web.Application()
|
||||
app.add_routes([web.post("/prompt", comfy.submit), web.get("/history/{prompt_id}", comfy.history),
|
||||
web.get("/view", comfy.view), web.get("/queue", comfy.queue),
|
||||
web.post("/queue", comfy.queue), web.post("/interrupt", comfy.interrupt)])
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
try:
|
||||
yield f"http://127.0.0.1:{port}"
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
def run(coro):
|
||||
return asyncio.run(coro)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind,index", [("IMAGE", 0), ("AUDIO", 1), ("MASK", 2), ("STRING", 3), ("VIDEO", 4)])
|
||||
def test_loopback_submit_poll_view_cleanup_and_caller_ownership(tmp_path, kind, index):
|
||||
async def check():
|
||||
data = request(kind)
|
||||
original = copy.deepcopy(data)
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url, aiohttp.ClientSession() as session:
|
||||
result = await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, session=session)
|
||||
assert not session.closed
|
||||
contract.validate_response(result["media"], request_id="request:1", bindings=data["media"]["outputs"])
|
||||
assert await contract.resolve_artifact(result["media"]["artifacts"][0]) == comfy.payload
|
||||
assert len(comfy.submitted) == 1
|
||||
assert comfy.submitted[0]["partial_execution_targets"] == ["export:1"]
|
||||
assert comfy.submitted[0]["prompt"]["source:1"]["inputs"]["value"][1] == index
|
||||
assert comfy.views == [{"filename": "a" * 32 + ".bin", "subfolder": "distributed-media-v1/" + hashlib.sha256(b"request:1").hexdigest(), "type": "temp"}]
|
||||
assert data == original
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode,match", [("corrupt", "checksum"), ("missing", "Missing"), ("error", "execution"), ("oversized", "length")])
|
||||
def test_terminal_failures_clean_only_after_terminal(tmp_path, mode, match):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = mode
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises((ValueError, RuntimeError), match=match):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert len(comfy.submitted) == 1
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", [
|
||||
lambda d: d["media"].update(version=True),
|
||||
lambda d: d["media"].update(profile="legacy"),
|
||||
lambda d: d["media"].update(request_id="../escape"),
|
||||
lambda d: d["media"]["outputs"][0].update(node_id="missing"),
|
||||
lambda d: d["workflow"]["export:1"].update(class_type="SaveImage"),
|
||||
lambda d: d["workflow"]["export:1"]["inputs"].update(request_id="other"),
|
||||
lambda d: d["workflow"]["export:1"]["inputs"].update(binding_json="{}"),
|
||||
lambda d: d["workflow"]["export:1"]["inputs"].update(images=["source:1", 0]),
|
||||
lambda d: d["workflow"]["export:1"]["inputs"].update(texts="literal"),
|
||||
lambda d: d["workflow"].update(unrelated={"class_type": "SaveImage", "inputs": {}}),
|
||||
lambda d: d["workflow"].update(unsafe={"class_type": "MediaImportV1", "inputs": {}}),
|
||||
lambda d: d["media"]["inputs"].append(copy.deepcopy(d["media"]["inputs"][0])),
|
||||
lambda d: d["media"]["inputs"][0].update(node_id="missing"),
|
||||
lambda d: d["media"]["inputs"][0].update(input_name="missing"),
|
||||
lambda d: d["media"].update(inputs=[]),
|
||||
lambda d: d["media"]["artifacts"][0].update(sha256="0" * 64),
|
||||
])
|
||||
def test_rejection_before_submit(tmp_path, mutation):
|
||||
async def check():
|
||||
data = request()
|
||||
mutation(data)
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.submitted
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", [
|
||||
lambda r: r.update(request_id="other"),
|
||||
lambda r: r["binding"].update(count=2),
|
||||
lambda r: r["binding"].update(collector_id="other"),
|
||||
lambda r: r["items"][0].update(item_index=1),
|
||||
lambda r: r["artifacts"][0]["transport"].update(filename="../escape"),
|
||||
lambda r: r["artifacts"][0]["transport"].update(subfolder="distributed-media-v1/other"),
|
||||
lambda r: r["artifacts"][0]["transport"].update(folder_type="input"),
|
||||
lambda r: r["artifacts"][0].update(byte_length=contract.MAX_ARTIFACT_BYTES + 1),
|
||||
])
|
||||
def test_bad_result_manifest_rejected_before_any_file_read(tmp_path, mutation):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.record_mutation = mutation
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.views
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_idempotent_history_retry_without_resubmission(tmp_path):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.history_failures = 2
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert comfy.polls == 3
|
||||
assert len(comfy.submitted) == 1
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["queued"])
|
||||
def test_deadline_cancels_only_own_prompt_before_cleanup(tmp_path, monkeypatch, mode):
|
||||
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = mode
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(TimeoutError):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, timeout_seconds=0.04)
|
||||
assert len(comfy.submitted) == 1
|
||||
assert bool(comfy.deletes) == (mode == "queued")
|
||||
assert bool(comfy.interrupts) == (mode == "running")
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_unconfirmed_terminal_retains_files_and_never_interrupts_unrelated(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
|
||||
monkeypatch.setattr(handler, "CANCEL_GRACE_SECONDS", 0.06)
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = "unrelated"
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(RuntimeError, match="retained"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, timeout_seconds=0.04)
|
||||
assert comfy.staged()
|
||||
assert not comfy.interrupts
|
||||
assert not comfy.deletes
|
||||
assert len(comfy.submitted) == 1
|
||||
run(check())
|
||||
|
||||
|
||||
def test_ambiguous_post_is_not_retried_and_own_queue_entry_is_reconciled(tmp_path):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = "ambiguous"
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(RuntimeError):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert len(comfy.submitted) == 1
|
||||
assert len(comfy.deletes) == 1
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_task_cancellation_fails_closed_and_preserves_caller_session(tmp_path):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = "running"
|
||||
async with serve(comfy) as url, aiohttp.ClientSession() as session:
|
||||
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, session=session))
|
||||
await comfy.prompt_seen.wait()
|
||||
await asyncio.sleep(0.01)
|
||||
task.cancel()
|
||||
with pytest.raises(handler.MediaExecutionUncertain, match="prompt-scoped interruption unavailable"):
|
||||
await task
|
||||
assert not session.closed
|
||||
assert not comfy.interrupts
|
||||
assert comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_legacy_images_are_owned_and_rewritten_without_overwriting_user_files(tmp_path):
|
||||
async def check():
|
||||
data = request(incoming=False)
|
||||
data["workflow"]["source:1"] = {"class_type": "LoadImage", "inputs": {"image": "same.png"}}
|
||||
data["images"] = [{"name": "same.png", "image": base64.b64encode(b"new image").decode()}]
|
||||
(tmp_path / "same.png").write_bytes(b"user image")
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
|
||||
filename = comfy.submitted[0]["prompt"]["source:1"]["inputs"]["image"]
|
||||
assert filename.startswith("distributed-media-v1/")
|
||||
assert filename.endswith(".bin")
|
||||
assert (tmp_path / "same.png").read_bytes() == b"user image"
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_url_input_uses_bounded_injected_downloader(tmp_path):
|
||||
async def check():
|
||||
data = request()
|
||||
data["media"]["artifacts"][0]["transport"] = {"type": "url", "url": "https://example.org/a", "expires_at": "2099-01-01T00:00:00Z"}
|
||||
calls = []
|
||||
async def download(url, *, max_bytes):
|
||||
calls.append((url, max_bytes))
|
||||
return b"input bytes"
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=download)
|
||||
assert calls == [("https://example.org/a", len(b"input bytes"))]
|
||||
run(check())
|
||||
|
||||
|
||||
def test_large_result_requires_publisher_and_validates_returned_transport(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path, payload=b"Z" * 100_000)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="publisher"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.views
|
||||
assert not comfy.staged()
|
||||
published = []
|
||||
async def publish(metadata, payload):
|
||||
assert "transport" not in metadata
|
||||
assert payload == comfy.payload
|
||||
published.append(metadata)
|
||||
return {"type": "url", "url": "https://example.org/result", "expires_at": "2099-01-01T00:00:00Z"}
|
||||
async with serve(comfy) as url:
|
||||
result = await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, publish=publish)
|
||||
assert len(published) == 1
|
||||
assert result["media"]["artifacts"][0]["transport"]["type"] == "url"
|
||||
assert len(json.dumps(result).encode()) <= 100_000
|
||||
async def bad_publish(metadata, payload):
|
||||
return {"type": "url", "url": "http://example.org/result", "expires_at": "2099-01-01T00:00:00Z"}
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, publish=bad_publish)
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", ["../photo.png", "folder/photo.png", "folder\\photo.png", ".", "..", "/photo.png"])
|
||||
def test_legacy_basename_rejected_without_submit(tmp_path, name):
|
||||
async def check():
|
||||
data = request(incoming=False)
|
||||
data["images"] = [{"name": name, "image": "YQ=="}]
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.submitted
|
||||
run(check())
|
||||
|
||||
|
||||
def test_staging_rejects_symlink_without_writing_or_submit(tmp_path):
|
||||
async def check():
|
||||
outside = tmp_path / "unrelated"
|
||||
outside.mkdir()
|
||||
(tmp_path / "distributed-media-v1").symlink_to(outside, target_is_directory=True)
|
||||
comfy = Comfy(tmp_path)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="symlink"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.submitted
|
||||
assert not list(outside.iterdir())
|
||||
run(check())
|
||||
|
||||
|
||||
def test_cleanup_preserves_other_runs_in_same_request_namespace(tmp_path):
|
||||
async def check():
|
||||
directory = tmp_path / "distributed-media-v1" / hashlib.sha256(b"request:1").hexdigest()
|
||||
directory.mkdir(parents=True)
|
||||
sibling = directory / ("b" * 32 + ".bin")
|
||||
sibling.write_bytes(b"other run")
|
||||
comfy = Comfy(tmp_path)
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert sibling.read_bytes() == b"other run"
|
||||
assert comfy.staged() == [sibling]
|
||||
run(check())
|
||||
|
||||
|
||||
def test_input_total_limit_checked_before_downloader_or_submit(tmp_path):
|
||||
async def check():
|
||||
data = request()
|
||||
data["media"]["artifacts"] = []
|
||||
data["media"]["inputs"] = []
|
||||
for index in range(3):
|
||||
artifact = contract.artifact_from_payload(f"input:{index}", "STRING", {}, b"",
|
||||
transport={"type": "url", "url": "https://example.org/input", "expires_at": "2099-01-01T00:00:00Z"})
|
||||
artifact["byte_length"] = contract.MAX_ARTIFACT_BYTES
|
||||
data["media"]["artifacts"].append(artifact)
|
||||
name = f"value:{index}"
|
||||
data["workflow"]["source:1"]["inputs"][name] = "literal"
|
||||
data["media"]["inputs"].append({"node_id": "source:1", "input_name": name, "artifact_id": artifact["id"]})
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async def forbidden_download(*args, **kwargs):
|
||||
pytest.fail("oversized request reached downloader")
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="total byte"):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=forbidden_download)
|
||||
assert not comfy.submitted
|
||||
run(check())
|
||||
|
||||
|
||||
def test_output_total_limit_checked_before_fetch(tmp_path):
|
||||
def excessive(record):
|
||||
first = record["artifacts"][0]
|
||||
record["artifacts"] = []
|
||||
record["items"] = []
|
||||
record["binding"]["count"] = 3
|
||||
for index, char in enumerate("abc"):
|
||||
artifact = copy.deepcopy(first)
|
||||
artifact["id"] = char * 32
|
||||
artifact["byte_length"] = contract.MAX_ARTIFACT_BYTES
|
||||
artifact["transport"]["filename"] = artifact["id"] + ".bin"
|
||||
record["artifacts"].append(artifact)
|
||||
record["items"].append({"binding_id": "binding:1", "item_index": index, "artifact_id": artifact["id"]})
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.record_mutation = excessive
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="total byte"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.views
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_partial_target_acceptance_is_not_reported_as_success(tmp_path):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = "partial"
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(RuntimeError, match="incomplete"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert len(comfy.submitted) == 1
|
||||
assert not comfy.views
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
# Native routes are compiled from the pinned source, without importing ComfyUI
|
||||
# (which would initialize the GPU). All dependencies and HTTP are CPU-only.
|
||||
@pytest.mark.parametrize("current,requested,race,expected", [
|
||||
("own", "own", False, ["own"]),
|
||||
("other", "own", False, []),
|
||||
("own", "own", True, ["successor"]),
|
||||
])
|
||||
def test_pinned_native_interrupt_route_has_snapshot_race(current, requested, race, expected):
|
||||
import ast
|
||||
import logging
|
||||
source = Path("/home/robert/ComfyUI_dev/server.py")
|
||||
if not source.is_file():
|
||||
pytest.skip("Pinned ComfyUI source unavailable")
|
||||
tree = ast.parse(source.read_text())
|
||||
routes = [node for node in ast.walk(tree) if isinstance(node, ast.AsyncFunctionDef)
|
||||
and any("/interrupt" in ast.unparse(d) for d in node.decorator_list)]
|
||||
assert len(routes) == 1
|
||||
route = routes[0]
|
||||
route.decorator_list = []
|
||||
interrupts = []
|
||||
class Queue:
|
||||
running = current
|
||||
def get_current_queue(self):
|
||||
snapshot = [[1, self.running, {}, {}, []]]
|
||||
if race:
|
||||
self.running = "successor"
|
||||
return snapshot, []
|
||||
queue = Queue()
|
||||
scope = {"json": json, "logging": logging, "web": web,
|
||||
"self": types.SimpleNamespace(prompt_queue=queue),
|
||||
"nodes": types.SimpleNamespace(interrupt_processing=lambda: interrupts.append(queue.running))}
|
||||
exec(compile(ast.Module(body=[route], type_ignores=[]), str(source), "exec"), scope)
|
||||
async def check():
|
||||
async def body():
|
||||
return {"prompt_id": requested}
|
||||
response = await scope[route.name](types.SimpleNamespace(json=body))
|
||||
assert response.status == 200
|
||||
run(check())
|
||||
assert interrupts == expected
|
||||
|
||||
|
||||
def exporter_file(root, *, name="a" * 32 + ".bin", payload=b"result bytes"):
|
||||
directory = root / "distributed-media-v1" / hashlib.sha256(b"request:1").hexdigest()
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
path = directory / name
|
||||
path.write_bytes(payload)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["success", "corrupt", "oversized", "error", "interrupted", "partial"])
|
||||
def test_owned_export_cleanup_after_terminal_success_or_failure(tmp_path, mode):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
sibling = exporter_file(temp, name="b" * 32 + ".bin", payload=b"unrelated")
|
||||
ordinary = exporter_file(temp, name="keep.txt", payload=b"ordinary")
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
comfy.mode = mode
|
||||
view = comfy.view
|
||||
async def reading(req):
|
||||
assert owned.exists(), "cleanup raced artifact collection"
|
||||
return await view(req)
|
||||
comfy.view = reading
|
||||
async with serve(comfy) as url:
|
||||
if mode == "success":
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
else:
|
||||
with pytest.raises((ValueError, RuntimeError)):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert comfy.terminal_seen
|
||||
assert not owned.exists()
|
||||
assert sibling.read_bytes() == b"unrelated"
|
||||
assert ordinary.read_bytes() == b"ordinary"
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_omitted_temp_root_leaves_exporter_lifecycle_ownership(tmp_path):
|
||||
async def check():
|
||||
owned = exporter_file(tmp_path / "temp")
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root)
|
||||
assert owned.read_bytes() == comfy.payload
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["running", "unrelated"])
|
||||
def test_uncertain_exports_and_inputs_retained_without_interrupt(tmp_path, monkeypatch, mode):
|
||||
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
|
||||
monkeypatch.setattr(handler, "CANCEL_GRACE_SECONDS", 0.04)
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
comfy.mode = mode
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(handler.MediaExecutionUncertain):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root,
|
||||
temp_dir=temp, timeout_seconds=0.03)
|
||||
assert owned.exists()
|
||||
assert comfy.staged()
|
||||
assert not comfy.interrupts
|
||||
assert not comfy.deletes
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("component", ["root", "namespace", "request", "file"])
|
||||
def test_export_cleanup_never_traverses_symlinks(tmp_path, component):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
selected = {"root": temp, "namespace": owned.parent.parent,
|
||||
"request": owned.parent, "file": owned}[component]
|
||||
outside = tmp_path / "outside"
|
||||
selected.rename(outside)
|
||||
selected.symlink_to(outside, target_is_directory=component != "file")
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert selected.is_symlink()
|
||||
assert owned.read_bytes() == comfy.payload
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", [
|
||||
lambda r: r.update(request_id="other"),
|
||||
lambda r: r["binding"].update(collector_id="other"),
|
||||
lambda r: r["artifacts"][0]["transport"].update(filename="../escape"),
|
||||
lambda r: r["artifacts"][0]["transport"].update(subfolder="distributed-media-v1/other"),
|
||||
])
|
||||
def test_unidentified_exports_are_not_deleted(tmp_path, mutation):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
comfy.record_mutation = mutation
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert owned.exists()
|
||||
assert not comfy.views
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["required", "publish"])
|
||||
def test_publication_failures_cleanup_identified_exports(tmp_path, monkeypatch, failure):
|
||||
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
|
||||
owned = exporter_file(temp, payload=comfy.payload)
|
||||
async def publish(metadata, payload):
|
||||
assert owned.exists()
|
||||
raise ValueError("publisher failed")
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="publisher"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root,
|
||||
temp_dir=temp, publish=publish if failure == "publish" else None)
|
||||
assert not owned.exists()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_repeated_cancel_drains_collection_before_cleanup(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
|
||||
owned = exporter_file(temp, payload=comfy.payload)
|
||||
started, draining, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||||
async def publish(metadata, payload):
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
finally:
|
||||
draining.set()
|
||||
await release.wait()
|
||||
assert owned.exists(), "cleanup while publisher still reading"
|
||||
async with serve(comfy) as url:
|
||||
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url,
|
||||
input_dir=comfy.root, temp_dir=temp, publish=publish))
|
||||
await asyncio.wait_for(started.wait(), 1)
|
||||
task.cancel()
|
||||
await asyncio.wait_for(draining.wait(), 1)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert owned.exists()
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert not owned.exists()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cancel_caller", [False, True])
|
||||
def test_stubborn_publisher_teardown_is_bounded_and_retains_files(tmp_path, monkeypatch, cancel_caller):
|
||||
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
|
||||
monkeypatch.setattr(handler, "DRAIN_GRACE_SECONDS", 0.03)
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
|
||||
owned = exporter_file(temp, payload=comfy.payload)
|
||||
started, draining, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||||
async def publish(metadata, payload):
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
finally:
|
||||
draining.set()
|
||||
await release.wait()
|
||||
assert owned.exists()
|
||||
async with serve(comfy) as url:
|
||||
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url,
|
||||
input_dir=comfy.root, temp_dir=temp, publish=publish,
|
||||
timeout_seconds=5 if cancel_caller else 0.15))
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), 1)
|
||||
if cancel_caller:
|
||||
task.cancel()
|
||||
await asyncio.wait_for(draining.wait(), 1)
|
||||
if cancel_caller:
|
||||
task.cancel() # must not reset the grace or recancel the reader
|
||||
done, _ = await asyncio.wait({task}, timeout=0.5)
|
||||
assert task in done, "callback suppressed cancellation indefinitely"
|
||||
with pytest.raises(handler.MediaExecutionUncertain, match="teardown"):
|
||||
task.result()
|
||||
assert comfy.terminal_seen and owned.exists() and comfy.staged()
|
||||
assert handler._UNCERTAIN_WORK
|
||||
pending = list(handler._UNCERTAIN_WORK)
|
||||
release.set()
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
await asyncio.sleep(0)
|
||||
assert not handler._UNCERTAIN_WORK
|
||||
assert owned.exists() and comfy.staged(), "late callback must not retire files"
|
||||
finally:
|
||||
release.set()
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
run(check())
|
||||
|
||||
|
||||
def test_export_replaced_during_collection_is_preserved(tmp_path):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
view = comfy.view
|
||||
async def replace(req):
|
||||
replacement = owned.with_suffix(".replacement")
|
||||
replacement.write_bytes(b"unrelated replacement")
|
||||
replacement.replace(owned)
|
||||
return await view(req)
|
||||
comfy.view = replace
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert owned.read_bytes() == b"unrelated replacement"
|
||||
run(check())
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", ["success", "error"])
|
||||
def test_partial_history_cleans_only_valid_identified_exporters(tmp_path, status):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
unidentified = exporter_file(temp, name="c" * 32 + ".bin")
|
||||
data = request()
|
||||
second = dict(data["media"]["outputs"][0], binding_id="binding:2", node_id="export:2", collector_id="collector:2")
|
||||
data["media"]["outputs"].append(second)
|
||||
data["workflow"]["export:2"] = {"class_type": "MediaExportV1", "inputs": {
|
||||
"request_id": "request:1", "binding_json": json.dumps(second), "texts": ["source:1", 0]}}
|
||||
comfy = Comfy(tmp_path / "input", data)
|
||||
comfy.mode = status
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises((ValueError, RuntimeError)):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert not owned.exists()
|
||||
assert unidentified.exists()
|
||||
assert not comfy.views
|
||||
run(check())
|
||||
|
||||
|
||||
def test_missing_history_exports_remain_lifecycle_owned(tmp_path):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
comfy.mode = "missing"
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="Missing"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert owned.exists()
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("replacement", ["symlink", "directory", "hardlink"])
|
||||
def test_nonregular_or_linked_export_files_are_preserved(tmp_path, replacement):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
outside = tmp_path / "outside"
|
||||
outside.write_bytes(b"unrelated")
|
||||
if replacement == "hardlink":
|
||||
alias = tmp_path / "alias"
|
||||
alias.hardlink_to(owned)
|
||||
else:
|
||||
owned.unlink()
|
||||
if replacement == "symlink":
|
||||
owned.symlink_to(outside)
|
||||
else:
|
||||
owned.mkdir()
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert owned.exists()
|
||||
assert outside.read_bytes() == b"unrelated"
|
||||
run(check())
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shared", [False, True])
|
||||
def test_direct_exporter_input_prunes_only_obsolete_loader_ancestors(tmp_path, shared):
|
||||
async def check():
|
||||
data = request("VIDEO")
|
||||
data["workflow"]["source:1"] = {"class_type": "LoadVideo", "inputs": {"file": "unavailable.mp4"}}
|
||||
data["media"]["inputs"] = [{"node_id": "export:1", "input_name": "video", "artifact_id": "input:1"}]
|
||||
if shared:
|
||||
second = dict(data["media"]["outputs"][0], binding_id="binding:2", node_id="export:2", collector_id="collector:2")
|
||||
data["media"]["outputs"].append(second)
|
||||
data["workflow"]["export:2"] = {"class_type": "MediaExportV1", "inputs": {
|
||||
"request_id": "request:1", "binding_json": json.dumps(second), "video": ["source:1", 0]}}
|
||||
fake_data = copy.deepcopy(data)
|
||||
fake_data["media"].update(artifacts=[], inputs=[])
|
||||
comfy = Comfy(tmp_path, fake_data)
|
||||
history = comfy.history
|
||||
async def full_history(req):
|
||||
response = await history(req)
|
||||
body = json.loads(response.text)
|
||||
if shared:
|
||||
record = comfy.record()
|
||||
record["binding"] = dict(second, count=1)
|
||||
record["artifacts"][0]["id"] = "b" * 32
|
||||
record["artifacts"][0]["transport"]["filename"] = "b" * 32 + ".bin"
|
||||
record["items"] = [{"binding_id": second["binding_id"], "item_index": 0, "artifact_id": "b" * 32}]
|
||||
body["prompt:1"]["outputs"]["export:2"] = {"media_exports": [record]}
|
||||
return web.json_response(body)
|
||||
comfy.history = full_history
|
||||
async with serve(comfy) as url:
|
||||
result = await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
|
||||
contract.validate_response(result["media"], request_id="request:1", bindings=data["media"]["outputs"])
|
||||
submitted = comfy.submitted[0]["prompt"]
|
||||
imported, slot = submitted["export:1"]["inputs"]["video"]
|
||||
assert slot == 4
|
||||
assert submitted[imported]["class_type"] == "MediaImportV1"
|
||||
assert json.loads(submitted[imported]["inputs"]["artifact_json"])["kind"] == "VIDEO"
|
||||
assert ("source:1" in submitted) is shared
|
||||
assert set(submitted) == ({"export:1", imported, "source:1", "export:2"} if shared else {"export:1", imported})
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target,kind", [("request_id", "VIDEO"), ("binding_json", "VIDEO"), ("video", "STRING")])
|
||||
def test_direct_exporter_input_rejects_controls_and_wrong_kind(tmp_path, target, kind):
|
||||
async def check():
|
||||
data = request("VIDEO")
|
||||
data["media"]["inputs"] = [{"node_id": "export:1", "input_name": target, "artifact_id": "input:1"}]
|
||||
data["media"]["artifacts"][0]["kind"] = kind
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.submitted
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_slow_download_times_out_before_any_prompt(tmp_path):
|
||||
async def check():
|
||||
data = request()
|
||||
data["media"]["artifacts"][0]["transport"] = {"type": "url", "url": "https://example.org/input", "expires_at": "2099-01-01T00:00:00Z"}
|
||||
async def download(*args, **kwargs):
|
||||
await asyncio.sleep(60)
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(TimeoutError):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=download, timeout_seconds=0.02)
|
||||
assert not comfy.submitted
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
@@ -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,83 @@
|
||||
"""Participant identity, typed compilation and legacy Collector pruning."""
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE = "media_transform_testpkg"
|
||||
for suffix in ("", ".api", ".api.orchestration", ".utils", ".serverless_bridge"):
|
||||
module = types.ModuleType(PACKAGE + suffix)
|
||||
module.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
|
||||
sys.modules[module.__name__] = module
|
||||
pt = importlib.import_module(PACKAGE + ".api.orchestration.prompt_transform")
|
||||
compiler = importlib.import_module(PACKAGE + ".api.orchestration.serverless_media_prompt")
|
||||
SOCKETS = ("images", "audio")
|
||||
|
||||
|
||||
def prompt(sockets=SOCKETS):
|
||||
nodes = {
|
||||
"source": {"class_type": "MediaSource", "inputs": {}},
|
||||
"collector": {"class_type": "DistributedCollector", "inputs": {
|
||||
socket: ["source", SOCKETS.index(socket)] for socket in sockets}},
|
||||
}
|
||||
for socket in sockets:
|
||||
nodes["sink:" + socket] = {"class_type": "Capture", "inputs": {"value": ["collector", SOCKETS.index(socket)]}}
|
||||
return nodes
|
||||
|
||||
|
||||
def prepare(graph, participant, delegate):
|
||||
index = pt.PromptIndex(graph)
|
||||
jobs = pt.generate_job_id_map(index, "run")
|
||||
assert jobs == {"collector": "run_collector"}
|
||||
return pt.apply_participant_overrides(index.copy_prompt(), participant,
|
||||
["serverless:pool:1", "serverless:pool:2"], jobs, "http://master", delegate, index)
|
||||
|
||||
|
||||
def test_typed_worker_preparation_compiles_both_bindings_without_source_mutation():
|
||||
graph = prompt()
|
||||
before = copy.deepcopy(graph)
|
||||
worker = prepare(graph, "serverless:pool:2", True)
|
||||
compiled, bindings = compiler.compile_media_prompt(worker, "request:2")
|
||||
assert len(bindings) == 2
|
||||
assert {row["participant_id"] for row in bindings} == {"serverless:pool:2"}
|
||||
assert {row["job_id"] for row in bindings} == {"run_collector"}
|
||||
assert {row["socket"] for row in bindings} == set(SOCKETS)
|
||||
assert set(compiled) == {"source"} | {row["node_id"] for row in bindings}
|
||||
assert graph == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sockets", [("images",), ("audio",), SOCKETS])
|
||||
def test_delegate_master_keeps_collector_socket_contract_and_legacy_placeholder(sockets):
|
||||
graph = prompt(sockets)
|
||||
prepared = prepare(graph, "master", True)
|
||||
before = copy.deepcopy(prepared)
|
||||
delegated = pt.prepare_delegate_master_prompt(prepared, ["collector"])
|
||||
assert "source" not in delegated
|
||||
assert {"collector"} | {"sink:" + s for s in sockets} <= set(delegated)
|
||||
inputs = delegated["collector"]["inputs"]
|
||||
assert "audio" not in inputs and "expected_media" not in inputs
|
||||
placeholders = [key for key, node in delegated.items() if node['class_type'] == 'DistributedEmptyImage']
|
||||
if "images" in sockets:
|
||||
assert len(placeholders) == 1
|
||||
assert inputs['images'] == [placeholders[0], 0]
|
||||
else:
|
||||
assert not placeholders and 'images' not in inputs
|
||||
assert inputs["delegate_only"] is True
|
||||
assert json.loads(inputs["enabled_worker_ids"]) == ["serverless:pool:1", "serverless:pool:2"]
|
||||
assert prepared == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("socket", SOCKETS)
|
||||
def test_participating_master_retains_exact_media_links(socket):
|
||||
graph = prompt()
|
||||
original = copy.deepcopy(graph)
|
||||
master = prepare(graph, "master", False)
|
||||
assert master["collector"]["inputs"][socket] == graph["collector"]["inputs"][socket]
|
||||
assert master["collector"]["inputs"]["delegate_only"] is False
|
||||
assert "worker_id" not in master["collector"]["inputs"]
|
||||
assert graph == original
|
||||
@@ -0,0 +1,552 @@
|
||||
"""Offline queue integration: real compilers, packers and wire validators.
|
||||
|
||||
Only host dispatch/probing/queueing and the manager lifecycle are mocked. No
|
||||
provider requests, services or GPU work are permitted by these tests.
|
||||
"""
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
import uuid
|
||||
import wave
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import aiohttp
|
||||
import pytest
|
||||
# Preload before the shared integration loader snapshots sys.modules: extension
|
||||
# modules cannot safely be unloaded and imported again between fixture contexts.
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from test_serverless_integration import load_modules, module, PKG
|
||||
|
||||
|
||||
PROFILE = "comfy-media-v1"
|
||||
|
||||
|
||||
def node(class_type, **inputs):
|
||||
return {"class_type": class_type, "inputs": inputs}
|
||||
|
||||
|
||||
def collector_prompt():
|
||||
return {
|
||||
"source": node("EmptyImage", width=32, height=32, batch_size=1, color=0),
|
||||
"collector": node("DistributedCollector", images=["source", 0], load_balance=False),
|
||||
"save": node("SaveImage", images=["collector", 0], filename_prefix="local"),
|
||||
}
|
||||
|
||||
|
||||
def legacy_prompt():
|
||||
return {
|
||||
"source": node("EmptyImage", width=32, height=32, batch_size=1, color=0),
|
||||
"collector": node("DistributedCollector", images=["source", 0]),
|
||||
"save": node("SaveImage", images=["collector", 0], filename_prefix="local"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def harness(monkeypatch):
|
||||
q, modules, ps = load_modules()
|
||||
with patch.dict(sys.modules, modules):
|
||||
pool = {"id": "pool", "name": "Media", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "offline-endpoint", "enabled": True, "jobs_per_run": 2,
|
||||
"timeout_seconds": 30, "media_profile": PROFILE}
|
||||
config = {"workers": [pool], "settings": {"master_delegate_only": True}}
|
||||
q.load_config = lambda: config
|
||||
q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate))
|
||||
q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 7, "node_errors": {}})
|
||||
manager = types.SimpleNamespace(start=AsyncMock(), cancel=AsyncMock(), ensure_capacity=Mock())
|
||||
sys.modules[PKG + ".api.orchestration.serverless_runtime"] = module("runtime", get_manager=lambda: manager)
|
||||
# These are the actual modules, not permissive serializer/packing stubs.
|
||||
client = importlib.import_module(PKG + ".api.orchestration.runpod_client")
|
||||
compiler = importlib.import_module(PKG + ".api.orchestration.serverless_media_prompt")
|
||||
packer = importlib.import_module(PKG + ".api.orchestration.serverless_typed_inputs")
|
||||
validator = importlib.import_module(PKG + ".api.orchestration.serverless_typed_results")
|
||||
handler = importlib.import_module(PKG + ".serverless_bridge.handler")
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError("Offline queue tests must not access a network or GPU")
|
||||
monkeypatch.setattr(aiohttp, "ClientSession", forbidden)
|
||||
monkeypatch.setattr(socket, "create_connection", forbidden)
|
||||
monkeypatch.setattr(socket.socket, "connect", forbidden)
|
||||
monkeypatch.setattr(socket.socket, "connect_ex", forbidden)
|
||||
monkeypatch.setattr(torch.cuda, "init", forbidden)
|
||||
monkeypatch.setattr(torch.cuda, "current_device", forbidden)
|
||||
h = types.SimpleNamespace(q=q, ps=ps, pool=pool, config=config, manager=manager,
|
||||
client=client, compiler=compiler, packer=packer,
|
||||
validator=validator, handler=handler)
|
||||
yield h
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def media_paths(tmp_path, monkeypatch):
|
||||
paths = types.SimpleNamespace(get_annotated_filepath=lambda name: str(tmp_path / name))
|
||||
for kind in ("input", "output", "temp"):
|
||||
setattr(paths, "get_" + kind + "_directory", lambda: str(tmp_path))
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", paths)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def run(h, prompt=None, **kwargs):
|
||||
return asyncio.run(h.q.orchestrate_distributed_execution(
|
||||
prompt if prompt is not None else collector_prompt(), None, "client", **kwargs))
|
||||
|
||||
|
||||
def assert_no_dispatch(h, *, no_probe=False):
|
||||
h.q.queue_prompt_payload.assert_not_awaited()
|
||||
h.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
h.manager.start.assert_not_awaited()
|
||||
h.manager.cancel.assert_not_awaited()
|
||||
assert h.ps.distributed_pending_jobs == {}
|
||||
if no_probe:
|
||||
h.q.select_active_workers.assert_not_awaited()
|
||||
h.q.select_least_busy_worker.assert_not_awaited()
|
||||
h.q.fetch_worker_has_node.assert_not_awaited()
|
||||
h.q.sync_worker_media.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("delegate", [True, False])
|
||||
def test_real_compile_pack_serialize_and_canonical_bindings(harness, media_paths, delegate):
|
||||
h = harness
|
||||
image_path = media_paths / "picture.png"
|
||||
Image.new("RGBA", (2, 3), color=(17, 29, 41, 53)).save(image_path)
|
||||
image_bytes = image_path.read_bytes()
|
||||
prompt = collector_prompt()
|
||||
prompt["image"] = node("LoadImage", image="picture.png")
|
||||
prompt["audio"] = node("GeneratedAudio")
|
||||
prompt["collector"]["inputs"].update(images=["image", 0], audio=["audio", 0])
|
||||
prompt["collector2"] = node("DistributedCollector", audio=["audio", 0])
|
||||
prompt["unused"] = node("PreviewImage", images=["image", 0])
|
||||
source = copy.deepcopy(prompt)
|
||||
with patch.object(h.compiler, "compile_media_prompt", wraps=h.compiler.compile_media_prompt) as compile_call, \
|
||||
patch.object(h.packer, "prepare_media_inputs", wraps=h.packer.prepare_media_inputs) as pack_call, \
|
||||
patch.object(h.client, "serialize_request", wraps=h.client.serialize_request) as serialize_call, \
|
||||
patch.object(h.q, "prune_prompt_for_worker", side_effect=AssertionError("typed graph must bypass legacy pruning")):
|
||||
result = run(h, prompt, delegate_master=delegate)
|
||||
assert result == ("master-prompt", 7, 2, {})
|
||||
assert compile_call.call_count == pack_call.await_count == 2
|
||||
assert serialize_call.call_count >= 2
|
||||
assert prompt == source
|
||||
h.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
jobs = h.manager.start.call_args.args[2]
|
||||
assert [j["worker"]["id"] for j in jobs] == ["serverless:pool:1", "serverless:pool:2"]
|
||||
assert len({j["media"]["request_id"] for j in jobs}) == 2
|
||||
master = h.q.queue_prompt_payload.call_args.args[0]
|
||||
for job in jobs:
|
||||
media = job["media"]
|
||||
assert str(uuid.UUID(media["request_id"])) == media["request_id"]
|
||||
assert uuid.UUID(media["request_id"]).version == 4
|
||||
assert media["profile"] == PROFILE and media["version"] == 1
|
||||
assert media["outputs"] == job["manifest"]
|
||||
assert media["outputs"] is not job["manifest"]
|
||||
assert media["artifacts"] == media["inputs"] == []
|
||||
assert {b["socket"] for b in job["manifest"]} == {"images", "audio"}
|
||||
assert len(job["manifest"]) == 3
|
||||
assert base64.b64decode(job["images"][0]["image"]) == image_bytes
|
||||
assert job["workflow"]["image"]["inputs"]["image"] == job["images"][0]["name"]
|
||||
assert not ({"save", "unused", "collector", "collector2"} & set(job["workflow"]))
|
||||
body = h.client.serialize_request(job["workflow"], job["images"], 30, media=media)
|
||||
wire = json.loads(body)
|
||||
assert wire["input"] == {"workflow": job["workflow"], "images": job["images"], "media": media}
|
||||
assert wire["policy"] == {"executionTimeout": 30000, "ttl": 30000}
|
||||
_, _, _, targets = h.handler._validate_request(wire["input"])
|
||||
assert targets == [b["node_id"] for b in job["manifest"]]
|
||||
assert h.validator.validate_prepared_job(job) is True
|
||||
for binding in job["manifest"]:
|
||||
assert binding["participant_id"] == job["worker"]["id"]
|
||||
assert binding["job_id"] == master[binding["collector_id"]]["inputs"]["multi_job_id"]
|
||||
controls = job["workflow"][binding["node_id"]]["inputs"]
|
||||
assert json.loads(controls["binding_json"]) == binding
|
||||
assert controls["request_id"] == media["request_id"]
|
||||
assert master["collector"]["inputs"]["delegate_only"] is delegate
|
||||
assert "expected_media" not in master["collector"]["inputs"]
|
||||
if delegate:
|
||||
assert not {"source", "audio"} & set(master)
|
||||
# LoadImage is an existing explicitly retained lightweight source; the
|
||||
# Collector is nevertheless detached from it and uses the placeholder.
|
||||
assert master["image"] == prompt["image"]
|
||||
assert "audio" not in master["collector"]["inputs"]
|
||||
placeholder = master["collector"]["inputs"]["images"][0]
|
||||
assert master[placeholder]["class_type"] == "DistributedEmptyImage"
|
||||
else:
|
||||
assert master["collector"]["inputs"]["audio"] == ["audio", 0]
|
||||
assert "audio" in master and "image" in master
|
||||
|
||||
|
||||
@pytest.mark.parametrize("audio_only", [False, True])
|
||||
def test_delegate_prunes_image_audio_sources_and_audio_only_needs_no_placeholder(harness, audio_only):
|
||||
prompt = collector_prompt()
|
||||
prompt["audio"] = node("GeneratedAudio")
|
||||
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
|
||||
if audio_only:
|
||||
del prompt["collector"]["inputs"]["images"]
|
||||
prompt["save"] = node("AudioSink", audio=["collector", 1])
|
||||
original = copy.deepcopy(prompt)
|
||||
run(harness, prompt)
|
||||
master = harness.q.queue_prompt_payload.call_args.args[0]
|
||||
inputs = master["collector"]["inputs"]
|
||||
assert "audio" not in inputs and "expected_media" not in inputs
|
||||
assert "source" not in master and "audio" not in master
|
||||
if audio_only:
|
||||
assert "images" not in inputs and set(master) == {"collector", "save"}
|
||||
else:
|
||||
assert master[inputs["images"][0]]["class_type"] == "DistributedEmptyImage"
|
||||
for job in harness.manager.start.call_args.args[2]:
|
||||
assert {b["socket"] for b in job["manifest"]} == ({"audio"} if audio_only else {"images", "audio"})
|
||||
assert "audio" in job["workflow"]
|
||||
assert ("source" in job["workflow"]) is not audio_only
|
||||
assert harness.validator.validate_prepared_job(job) is True
|
||||
assert prompt == original
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selected", ["master", "worker"])
|
||||
def test_typed_load_balance_selects_master_or_one_serverless_slot(harness, selected):
|
||||
h = harness
|
||||
prompt = collector_prompt()
|
||||
prompt["collector"]["inputs"]["load_balance"] = "true"
|
||||
h.q.select_least_busy_worker.side_effect = lambda candidates, **kw: next(
|
||||
w for w in candidates if (w["id"] == "master") == (selected == "master"))
|
||||
result = run(h, prompt, delegate_master=False)
|
||||
h.q.select_least_busy_worker.assert_awaited_once()
|
||||
master = h.q.queue_prompt_payload.call_args.args[0]
|
||||
if selected == "master":
|
||||
assert result[2] == 0
|
||||
h.manager.start.assert_not_awaited()
|
||||
assert master["collector"]["inputs"]["delegate_only"] is False
|
||||
assert "source" in master
|
||||
else:
|
||||
assert result[2] == 1
|
||||
assert len(h.manager.start.call_args.args[2]) == 1
|
||||
assert master["collector"]["inputs"]["delegate_only"] is True
|
||||
assert "source" not in master
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", [None, "legacy-image-v1"])
|
||||
def test_exact_rejects_legacy_pool_before_probe(harness, profile):
|
||||
h = harness
|
||||
prompt = collector_prompt()
|
||||
prompt["collector"]["inputs"]["transfer_format"] = "exact"
|
||||
if profile is None:
|
||||
h.pool.pop("media_profile")
|
||||
else:
|
||||
h.pool["media_profile"] = profile
|
||||
with pytest.raises(ValueError, match="exact requires comfy-media-v1"):
|
||||
run(h, prompt)
|
||||
assert_no_dispatch(h, no_probe=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", [None, "auto", "comfy-media-v2", False])
|
||||
def test_invalid_explicit_profiles_fail_closed(harness, profile):
|
||||
harness.pool["media_profile"] = profile
|
||||
with pytest.raises(ValueError, match="media_profile"):
|
||||
run(harness)
|
||||
assert_no_dispatch(harness, no_probe=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("direct_type", ["local", "remote", "cloud"])
|
||||
@pytest.mark.parametrize("explicit", [False, True])
|
||||
def test_direct_workers_participate_alongside_typed_pools(harness, direct_type, explicit):
|
||||
h = harness
|
||||
h.config["workers"].append({"id": "direct", "type": direct_type, "host": "offline.invalid",
|
||||
"enabled": not explicit, "port": 8188})
|
||||
kwargs = {"enabled_worker_ids": ["pool", "direct"]} if explicit else {}
|
||||
prompt = collector_prompt()
|
||||
original = copy.deepcopy(prompt)
|
||||
assert run(h, prompt, delegate_master=False, **kwargs)[2] == 3
|
||||
assert prompt == original
|
||||
h.q.dispatch_worker_prompt.assert_awaited_once()
|
||||
master = h.q.queue_prompt_payload.call_args.args[0]
|
||||
assert json.loads(master["collector"]["inputs"]["enabled_worker_ids"]) == [
|
||||
"serverless:pool:1", "serverless:pool:2", "direct"]
|
||||
jobs = h.manager.start.call_args.args[2]
|
||||
assert len(jobs) == 2 and all(j["media"]["profile"] == PROFILE for j in jobs)
|
||||
if direct_type != "local":
|
||||
h.q.sync_worker_media.assert_awaited_once()
|
||||
|
||||
|
||||
def test_disabled_or_unselected_workers_do_not_participate(harness):
|
||||
h = harness
|
||||
h.config["workers"] += [
|
||||
{"id": "direct", "type": "local", "enabled": False},
|
||||
dict(h.pool, id="legacy", media_profile="legacy-image-v1", enabled=False),
|
||||
]
|
||||
assert run(h)[2] == 2
|
||||
h.pool["enabled"] = False
|
||||
assert run(h, enabled_worker_ids=["pool"])[2] == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("workers", [True, False])
|
||||
def test_image_list_cannot_mix_with_collector_even_master_only(harness, workers):
|
||||
if not workers:
|
||||
harness.config["workers"] = []
|
||||
prompt = collector_prompt()
|
||||
prompt["incompatible"] = node("DistributedImageListStart")
|
||||
with pytest.raises(ValueError, match="cannot be mixed"):
|
||||
run(harness, prompt, delegate_master=False)
|
||||
assert_no_dispatch(harness, no_probe=True)
|
||||
|
||||
|
||||
def test_mixed_direct_typed_and_legacy_profiles_keep_each_wire_contract(harness):
|
||||
h = harness
|
||||
h.config["workers"] += [dict(h.pool, id="legacy", media_profile="legacy-image-v1"),
|
||||
{"id": "direct", "type": "local", "enabled": True}]
|
||||
assert run(h, collector_prompt())[2] == 5
|
||||
h.q.dispatch_worker_prompt.assert_awaited_once()
|
||||
jobs = h.manager.start.call_args.args[2]
|
||||
assert len(jobs) == 4
|
||||
for job in jobs:
|
||||
if job["worker"]["media_profile"] == PROFILE:
|
||||
assert job["media"]["outputs"] == job["manifest"]
|
||||
assert h.validator.validate_prepared_job(job) is True
|
||||
else:
|
||||
assert "media" not in job
|
||||
assert any(n["class_type"] == "SaveImage" for n in job["workflow"].values())
|
||||
|
||||
|
||||
def test_typed_master_only_is_supported_without_serverless_work(harness):
|
||||
h = harness
|
||||
h.config["workers"] = []
|
||||
assert run(h, delegate_master=False)[2] == 0
|
||||
master = h.q.queue_prompt_payload.call_args.args[0]
|
||||
assert master["collector"]["inputs"]["enabled_worker_ids"] == "[]"
|
||||
assert master["collector"]["inputs"]["images"] == ["source", 0]
|
||||
h.manager.start.assert_not_awaited()
|
||||
h.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("offline", [False, True])
|
||||
def test_collector_delegate_without_workers_preserves_master_fallback(harness, offline):
|
||||
h = harness
|
||||
if offline:
|
||||
h.q.select_active_workers.side_effect = lambda *a, **kw: ([], False)
|
||||
else:
|
||||
h.config["workers"] = []
|
||||
assert run(h)[2] == 0
|
||||
master = h.q.queue_prompt_payload.call_args.args[0]
|
||||
assert master["collector"]["inputs"]["delegate_only"] is False
|
||||
assert master["collector"]["inputs"]["images"] == ["source", 0]
|
||||
h.manager.start.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [" pool ", "pool white", "pool\n"])
|
||||
def test_pool_media_identity_is_not_repaired(harness, bad):
|
||||
harness.pool["id"] = bad
|
||||
with pytest.raises(ValueError, match="identifier"):
|
||||
run(harness)
|
||||
assert_no_dispatch(harness, no_probe=True)
|
||||
|
||||
|
||||
def test_integer_source_node_identity_is_not_normalized(harness):
|
||||
prompt = collector_prompt()
|
||||
prompt[1] = prompt.pop("source")
|
||||
prompt["collector"]["inputs"]["images"] = [1, 0]
|
||||
with pytest.raises(ValueError, match="identifier"):
|
||||
run(harness, prompt)
|
||||
assert_no_dispatch(harness, no_probe=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("broken", ["dangling-unused", "literal", "missing-file", "bad-loader-controls", "oversize"])
|
||||
def test_compilation_packing_and_size_errors_precede_all_dispatch(harness, media_paths, broken):
|
||||
h = harness
|
||||
prompt = collector_prompt()
|
||||
if broken == "dangling-unused":
|
||||
prompt["unused"] = node("PreviewImage", images=["missing", 0])
|
||||
elif broken == "literal":
|
||||
prompt["collector"]["inputs"]["images"] = "must use a linked image node"
|
||||
elif broken in ("missing-file", "bad-loader-controls"):
|
||||
prompt["audio"] = node("LoadAudio", audio="missing.wav")
|
||||
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
|
||||
if broken == "bad-loader-controls":
|
||||
prompt["audio"]["inputs"]["duration"] = 5
|
||||
else:
|
||||
prompt["source"]["inputs"]["label"] = "雪" * 1000
|
||||
original = copy.deepcopy(prompt)
|
||||
with patch.object(h.client, "MAX_REQUEST_BYTES", 100 if broken == "oversize" else h.client.MAX_REQUEST_BYTES):
|
||||
with pytest.raises((ValueError, OSError)):
|
||||
run(h, prompt)
|
||||
assert prompt == original
|
||||
assert_no_dispatch(h)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("corruption", ["manifest", "participant", "request-control", "binding-control", "profile", "input-ref"])
|
||||
def test_whole_batch_identity_preflight_before_master_and_manager(harness, corruption):
|
||||
h = harness
|
||||
real_prepare = h.q._prepare_worker_payload
|
||||
async def corrupt_later(*args, **kwargs):
|
||||
worker, prepared = await real_prepare(*args, **kwargs)
|
||||
if worker["id"].endswith(":2"):
|
||||
binding = prepared["manifest"][0]
|
||||
if corruption == "manifest":
|
||||
binding["job_id"] = "wrong-job"
|
||||
elif corruption == "participant":
|
||||
prepared["worker"] = dict(worker, id="different-participant")
|
||||
elif corruption == "request-control":
|
||||
prepared["workflow"][binding["node_id"]]["inputs"]["request_id"] = "different-request"
|
||||
elif corruption == "binding-control":
|
||||
prepared["workflow"][binding["node_id"]]["inputs"]["binding_json"] = "{}"
|
||||
elif corruption == "profile":
|
||||
prepared["worker"] = dict(worker, media_profile="legacy-image-v1")
|
||||
else:
|
||||
prepared["media"]["inputs"] = [{"node_id": "missing", "input_name": "texts", "artifact_id": "unknown"}]
|
||||
return worker, prepared
|
||||
with patch.object(h.q, "_prepare_worker_payload", side_effect=corrupt_later):
|
||||
with pytest.raises(ValueError):
|
||||
run(h)
|
||||
h.manager.ensure_capacity.assert_not_called()
|
||||
assert_no_dispatch(h)
|
||||
|
||||
|
||||
def test_batch_preflight_also_precedes_regular_direct_dispatch(harness):
|
||||
h = harness
|
||||
h.pool.pop("media_profile")
|
||||
h.config["workers"].append({"id": "direct", "type": "local", "enabled": True})
|
||||
real_prepare = h.q._prepare_worker_payload
|
||||
async def corrupt_later(*args, **kwargs):
|
||||
worker, prepared = await real_prepare(*args, **kwargs)
|
||||
if worker["id"] == "serverless:pool:2":
|
||||
prepared["worker"] = dict(worker, media_profile=PROFILE)
|
||||
return worker, prepared
|
||||
with patch.object(h.q, "_prepare_worker_payload", side_effect=corrupt_later):
|
||||
with pytest.raises(ValueError, match="mismatch"):
|
||||
run(h, legacy_prompt())
|
||||
assert_no_dispatch(h)
|
||||
|
||||
|
||||
def test_capacity_preflight_precedes_collector_state_and_master(harness):
|
||||
h = harness
|
||||
h.manager.ensure_capacity.side_effect = ValueError("capacity exhausted")
|
||||
with pytest.raises(ValueError, match="capacity"):
|
||||
run(h)
|
||||
assert_no_dispatch(h)
|
||||
|
||||
|
||||
def test_master_queue_failure_cleans_queues_and_never_starts_manager(harness):
|
||||
h = harness
|
||||
h.q.queue_prompt_payload.side_effect = ValueError("invalid master graph")
|
||||
with pytest.raises(ValueError, match="invalid master"):
|
||||
run(h)
|
||||
assert h.ps.distributed_pending_jobs == {}
|
||||
h.manager.start.assert_not_awaited()
|
||||
h.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
|
||||
|
||||
def test_manager_setup_failure_emits_one_terminal_per_collector_participant(harness):
|
||||
h = harness
|
||||
prompt = collector_prompt()
|
||||
prompt["audio"] = node("GeneratedAudio")
|
||||
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
|
||||
h.manager.start.side_effect = ValueError("journal failure")
|
||||
with pytest.raises(ValueError, match="journal failure"):
|
||||
run(h, prompt)
|
||||
h.manager.cancel.assert_awaited_once()
|
||||
queued = next(iter(h.ps.distributed_pending_jobs.values()))
|
||||
assert queued.qsize() == 2 # two participants, not four per-socket completions
|
||||
messages = [queued.get_nowait(), queued.get_nowait()]
|
||||
assert {m["worker_id"] for m in messages} == {"serverless:pool:1", "serverless:pool:2"}
|
||||
assert all(m["is_last"] is True and m["serverless_error"] for m in messages)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", [None, "legacy-image-v1"])
|
||||
def test_legacy_envelope_and_compiler_are_unchanged(harness, profile):
|
||||
h = harness
|
||||
if profile is None:
|
||||
h.pool.pop("media_profile")
|
||||
else:
|
||||
h.pool["media_profile"] = profile
|
||||
prompt = legacy_prompt()
|
||||
original = copy.deepcopy(prompt)
|
||||
with patch.object(h.compiler, "compile_media_prompt", side_effect=AssertionError("legacy must not use typed compiler")), \
|
||||
patch.object(h.packer, "prepare_media_inputs", side_effect=AssertionError("legacy must not use typed packer")):
|
||||
assert run(h, prompt)[2] == 2
|
||||
for job in h.manager.start.call_args.args[2]:
|
||||
assert "media" not in job
|
||||
assert any(n["class_type"] == "SaveImage" for n in job["workflow"].values())
|
||||
expected = {"input": {"workflow": job["workflow"], "images": job["images"]},
|
||||
"policy": {"executionTimeout": 30000, "ttl": 30000}}
|
||||
assert h.client.serialize_request(job["workflow"], job["images"], 30) == json.dumps(
|
||||
expected, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode()
|
||||
assert prompt == original
|
||||
|
||||
|
||||
def test_legacy_no_worker_delegate_fallback_is_preserved(harness):
|
||||
h = harness
|
||||
h.config["workers"] = []
|
||||
assert run(h, legacy_prompt())[2] == 0
|
||||
assert h.q.queue_prompt_payload.call_args.args[0]["collector"]["inputs"]["delegate_only"] is False
|
||||
|
||||
|
||||
def test_native_audio_compile_pack_serialized_input_cpu(harness, media_paths):
|
||||
"""Exercise real PyAV + core PCM conversion, not a mocked packing result."""
|
||||
if os.environ.get("_MEDIA_QUEUE_CORE_TEST") != "1":
|
||||
core = os.environ.get("COMFYUI_SOURCE_ROOT")
|
||||
if not core:
|
||||
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU native audio queue integration")
|
||||
env = dict(os.environ)
|
||||
env.pop("PYTHONPATH", None)
|
||||
env.pop("COMFYUI_SOURCE_ROOT", None)
|
||||
scratch = Path.home() / ".hermes/cache/scratch"
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
env.update(_MEDIA_QUEUE_CORE_TEST="1", CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="",
|
||||
PYTHONDONTWRITEBYTECODE="1", TMPDIR=str(scratch))
|
||||
bootstrap = """
|
||||
import sys
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
import torch
|
||||
|
||||
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('CPU audio integration may not initialize a GPU')
|
||||
torch.cuda.is_available = lambda: False
|
||||
torch.cuda.device_count = lambda: 0
|
||||
torch.cuda.init = forbidden
|
||||
torch.cuda.current_device = forbidden
|
||||
torch.cuda._lazy_init = forbidden
|
||||
import pytest
|
||||
raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider', '-k',
|
||||
'test_native_audio_compile_pack_serialized_input_cpu']))
|
||||
"""
|
||||
result = subprocess.run([sys.executable, "-c", bootstrap, core, str(Path(__file__).resolve())],
|
||||
env=env, capture_output=True, text=True, timeout=120)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert "1 passed" in result.stdout
|
||||
return
|
||||
h = harness
|
||||
h.pool["jobs_per_run"] = 1
|
||||
path = media_paths / "tone.wav"
|
||||
pcm = torch.tensor([-100, 0, 100, 32767, -32768, 7], dtype=torch.int16)
|
||||
with wave.open(str(path), "wb") as stream:
|
||||
stream.setnchannels(1)
|
||||
stream.setsampwidth(2)
|
||||
stream.setframerate(8000)
|
||||
stream.writeframes(pcm.numpy().tobytes())
|
||||
original_bytes = path.read_bytes()
|
||||
prompt = collector_prompt()
|
||||
prompt["audio"] = node("LoadAudio", audio="tone.wav")
|
||||
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
|
||||
original = copy.deepcopy(prompt)
|
||||
run(h, prompt)
|
||||
job = h.manager.start.call_args.args[2][0]
|
||||
wire = json.loads(h.client.serialize_request(job["workflow"], job["images"], 30, media=job["media"]))
|
||||
artifacts = wire["input"]["media"]["artifacts"]
|
||||
refs = wire["input"]["media"]["inputs"]
|
||||
assert len(artifacts) == len(refs) == 1
|
||||
binding = next(b for b in job["manifest"] if b["socket"] == "audio")
|
||||
assert refs == [{"node_id": binding["node_id"], "input_name": "audio", "artifact_id": artifacts[0]["id"]}]
|
||||
contract = importlib.import_module(PKG + ".serverless_bridge.contract")
|
||||
codec = importlib.import_module(PKG + ".serverless_bridge.codec")
|
||||
payload = asyncio.run(contract.resolve_artifact(artifacts[0]))
|
||||
value = codec.decode_payload("AUDIO", artifacts[0]["descriptor"], payload)
|
||||
assert value["sample_rate"] == 8000
|
||||
assert value["waveform"].shape == (1, 1, 6)
|
||||
assert torch.equal(value["waveform"], (pcm.to(torch.float32) / 32768).reshape(1, 1, -1))
|
||||
assert path.read_bytes() == original_bytes
|
||||
assert prompt == original
|
||||
assert h.validator.validate_prepared_job(job) is True
|
||||
@@ -0,0 +1,96 @@
|
||||
"""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_collector_typed_media import collector_module, native_media, assert_exact
|
||||
|
||||
FORMATS = ('legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1', 'exact')
|
||||
|
||||
|
||||
def compile_request(fmt='legacy_png', quality=90):
|
||||
prompt = _prompt()
|
||||
prompt['source:1'] = {'class_type': 'Source', 'inputs': {}}
|
||||
prompt['collector:1']['inputs'].update(transfer_format=fmt, transfer_quality=quality,
|
||||
audio=['source:1', 1])
|
||||
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)
|
||||
assert {b['socket'] for b in data['media']['outputs']} == {'images', 'audio'}
|
||||
for binding in data['media']['outputs']:
|
||||
inputs = data['workflow'][binding['node_id']]['inputs']
|
||||
if binding['kind'] == 'IMAGE' and fmt != 'exact':
|
||||
assert inputs['transfer_format'] == ('png' if fmt == 'legacy_png' else fmt)
|
||||
assert inputs['transfer_quality'] == 73
|
||||
else:
|
||||
assert 'transfer_format' not in inputs and 'transfer_quality' not in inputs
|
||||
handler._validate_request(data)
|
||||
|
||||
|
||||
def test_saved_collector_defaults_compile_to_png_not_exact():
|
||||
workflow, bindings = compiler.compile_media_prompt(_prompt(), 'request:1')
|
||||
inputs = workflow[bindings[0]['node_id']]['inputs']
|
||||
assert (inputs['transfer_format'], inputs['transfer_quality']) == ('png', 90)
|
||||
|
||||
|
||||
@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)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('socket,kind', [('audio', 'AUDIO'), ('masks', 'MASK'), ('texts', 'STRING'), ('video', 'VIDEO')])
|
||||
def test_handler_rejects_image_controls_on_nonimage_bridge_exports(socket, kind):
|
||||
# The bridge still supports all five kinds; these are direct bridge exports,
|
||||
# not nonexistent sockets on DistributedCollector.
|
||||
data = compile_request('exact')
|
||||
binding = next(b for b in data['media']['outputs'] if b['kind'] == 'AUDIO')
|
||||
binding.update(socket=socket, kind=kind)
|
||||
inputs = data['workflow'][binding['node_id']]['inputs']
|
||||
inputs.pop('audio')
|
||||
inputs.update({socket: ['source:1', 1], 'binding_json': json.dumps(binding)})
|
||||
handler._validate_request(data)
|
||||
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(collector_module, fmt):
|
||||
media = native_media()
|
||||
actual = collector_module.DistributedCollectorNode().run(
|
||||
images=[media['images']], audio=[media['audio']], pass_through=[True],
|
||||
transfer_format=[fmt], transfer_quality=[73])
|
||||
assert actual[0] is media['images']
|
||||
assert actual[1] is media['audio']
|
||||
assert_exact(actual[0], media['images'])
|
||||
|
||||
|
||||
def test_collector_widget_contract_keeps_saved_defaults_and_appends_exact(collector_module):
|
||||
optional = collector_module.DistributedCollectorNode.INPUT_TYPES()['optional']
|
||||
assert list(optional) == ['images', 'audio', 'transfer_format', 'transfer_quality']
|
||||
assert tuple(optional['transfer_format'][0]) == FORMATS
|
||||
assert optional['transfer_format'][1]['default'] == 'legacy_png'
|
||||
assert optional['transfer_quality'][1]['default'] == 90
|
||||
@@ -0,0 +1,562 @@
|
||||
"""Native ComfyUI VIDEO tests, isolated from the custom-node import graph.
|
||||
|
||||
The outer pytest process starts one bounded CPU-only subprocess. Inside it the
|
||||
real pinned core is imported *after* cli_args.args.cpu=True; no fake VideoInput
|
||||
classes, media transcodes, providers, server, or GPU queries are needed.
|
||||
"""
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
if os.environ.get("_DISTRIBUTED_VIDEO_CODEC_CORE_TEST") != "1":
|
||||
def test_native_video_codec_in_cpu_subprocess():
|
||||
configured = os.environ.get("COMFYUI_SOURCE_ROOT", os.environ.get("COMFYUI_VIDEO_TEST_ROOT"))
|
||||
if configured is None:
|
||||
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU native VIDEO acceptance")
|
||||
core = Path(configured).expanduser().resolve()
|
||||
assert configured and (core / "comfy_api/latest/_input_impl/video_types.py").is_file(), (
|
||||
"Configured COMFYUI_SOURCE_ROOT (or COMFYUI_VIDEO_TEST_ROOT) must provide native video classes"
|
||||
)
|
||||
env = dict(os.environ)
|
||||
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
env["TMPDIR"] = str(scratch)
|
||||
for name in ("PYTHONPATH", "COMFYUI_SOURCE_ROOT"):
|
||||
env.pop(name, None)
|
||||
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="",
|
||||
_DISTRIBUTED_VIDEO_CODEC_CORE_TEST="1", PYTHONDONTWRITEBYTECODE="1")
|
||||
bootstrap = """
|
||||
import sys
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
import torch
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('VIDEO tests must not query or initialize a GPU')
|
||||
# Third-party core imports probe availability even in CPU mode: answer without
|
||||
# executing a GPU query, then forbid even these probes during codec execution.
|
||||
torch.cuda.is_available = lambda: False
|
||||
torch.cuda.device_count = lambda: 0
|
||||
for name in ('init', '_lazy_init', 'current_device'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
import comfy_api.latest._input_impl.video_types
|
||||
for name in ('is_available', 'device_count'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
import pytest
|
||||
raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider']))
|
||||
"""
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", bootstrap, str(core), str(Path(__file__).resolve())],
|
||||
env=env, text=True, capture_output=True, timeout=150,
|
||||
cwd=Path(__file__).resolve().parents[1],
|
||||
)
|
||||
print(result.stdout)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
else:
|
||||
import copy
|
||||
from fractions import Fraction
|
||||
import importlib
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
import types
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
from comfy_api.latest._input import VideoInput
|
||||
from comfy_api.latest._input_impl.video_types import (
|
||||
VideoFromComponents, VideoFromFile, VideoFromList,
|
||||
)
|
||||
from comfy_api.latest._util import VideoComponents
|
||||
|
||||
# Only this sibling package is synthetic; all native media types are real.
|
||||
package_name = "_distributed_video_codec_test"
|
||||
package = types.ModuleType(package_name)
|
||||
package.__path__ = [str(Path(__file__).resolve().parents[1] / "serverless_bridge")]
|
||||
sys.modules[package_name] = package
|
||||
codec = importlib.import_module(package_name + ".video_codec")
|
||||
|
||||
def audio(dtype=torch.float64):
|
||||
return {"waveform": torch.tensor([[[0.123456789, -0.0, -1.5, 2.0]]], dtype=dtype, device="cpu"),
|
||||
"sample_rate": 48000}
|
||||
|
||||
def components(dtype=torch.float64, *, metadata=None, alpha=True, sound=True):
|
||||
image = torch.arange(36, device="cpu", dtype=torch.float64).reshape(2, 2, 3, 3) / 997
|
||||
image[0, 0, 0, 0] = -0.0
|
||||
return VideoFromComponents(VideoComponents(
|
||||
images=image.to(dtype), frame_rate=Fraction(30000, 1001),
|
||||
audio=audio(dtype) if sound else None,
|
||||
alpha=torch.full((2, 2, 3), 0.123456789, dtype=dtype, device="cpu") if alpha else None,
|
||||
metadata=metadata,
|
||||
), bit_depth=10, color_space="HDR")
|
||||
|
||||
def exact_tensor(actual, expected):
|
||||
assert actual.device.type == "cpu"
|
||||
assert actual.dtype == expected.dtype
|
||||
assert actual.shape == expected.shape
|
||||
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
|
||||
|
||||
def file_value(data=b"opaque original media bytes", **kwargs):
|
||||
return VideoFromFile(io.BytesIO(data), **kwargs)
|
||||
|
||||
def forbid(*args, **kwargs):
|
||||
raise AssertionError("Unexpected materialization/transcoding/allocation")
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
||||
@pytest.mark.parametrize("alpha,sound", [(False, False), (True, False), (False, True), (True, True)])
|
||||
def test_components_exact_roundtrip(dtype, alpha, sound, monkeypatch):
|
||||
metadata = {"title": "雪 🎞", "flags": [True, None, -4, 0.125], "nested": {"empty": []}}
|
||||
source = components(dtype, metadata=metadata, alpha=alpha, sound=sound)
|
||||
expected = source.get_components()
|
||||
monkeypatch.setattr(VideoFromComponents, "get_components", forbid)
|
||||
monkeypatch.setattr(VideoFromComponents, "save_to", forbid)
|
||||
descriptor, payload = codec.encode_video(source)
|
||||
assert descriptor["encoding"] == "native-video-v1"
|
||||
assert descriptor["representation"] == "components"
|
||||
assert descriptor["frame_rate"] == {"numerator": 30000, "denominator": 1001}
|
||||
# This is the actual JSON wire descriptor, not a Python-only shortcut.
|
||||
descriptor = json.loads(json.dumps(descriptor, allow_nan=False))
|
||||
result = codec.decode_video(descriptor, payload)
|
||||
assert type(result) is VideoFromComponents
|
||||
assert isinstance(result, VideoInput)
|
||||
actual = vars(result)["_VideoFromComponents__components"]
|
||||
exact_tensor(actual.images, expected.images)
|
||||
assert actual.frame_rate == expected.frame_rate
|
||||
assert type(actual.frame_rate) is Fraction
|
||||
assert actual.metadata == metadata
|
||||
assert actual.metadata is not metadata
|
||||
assert vars(result)["_VideoFromComponents__bit_depth"] == 10
|
||||
assert vars(result)["_VideoFromComponents__color_space"] == "HDR"
|
||||
if alpha:
|
||||
exact_tensor(actual.alpha, expected.alpha)
|
||||
else:
|
||||
assert actual.alpha is None
|
||||
if sound:
|
||||
exact_tensor(actual.audio["waveform"], expected.audio["waveform"])
|
||||
assert actual.audio["sample_rate"] == expected.audio["sample_rate"]
|
||||
else:
|
||||
assert actual.audio is None
|
||||
actual.images.fill_(99)
|
||||
assert not torch.equal(actual.images, expected.images)
|
||||
actual.metadata["nested"]["empty"].append("new")
|
||||
assert descriptor["metadata"]["nested"]["empty"] == []
|
||||
assert metadata["nested"]["empty"] == []
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
||||
def test_native_bhw1_alpha_shape_is_preserved(dtype):
|
||||
source = components(dtype)
|
||||
original = source.get_components()
|
||||
vars(source)["_VideoFromComponents__components"].alpha = original.alpha.unsqueeze(-1)
|
||||
descriptor, payload = codec.encode_video(source)
|
||||
result = codec.decode_video(descriptor, payload).get_components()
|
||||
exact_tensor(result.alpha, original.alpha.unsqueeze(-1))
|
||||
assert descriptor["alpha"]["descriptor"]["shape"] == [2, 2, 3, 1]
|
||||
|
||||
@pytest.mark.parametrize("start,duration,crop", [
|
||||
(-1.25, -2.5, (-3, 5, 7, -9)), (0, 0, None), (-0.0, 0.0, (0, 0, 2, 2)),
|
||||
])
|
||||
def test_file_preserves_bytes_cursor_signed_trim_crop_and_ownership(start, duration, crop, monkeypatch):
|
||||
source = file_value(start_time=start, duration=duration, crop=crop)
|
||||
stream = vars(source)["_VideoFromFile__file"]
|
||||
stream.seek(7)
|
||||
monkeypatch.setattr(VideoFromFile, "get_components", forbid)
|
||||
monkeypatch.setattr(VideoFromFile, "save_to", forbid)
|
||||
descriptor, payload = codec.encode_video(source)
|
||||
assert stream.tell() == 7
|
||||
assert payload == stream.getvalue()
|
||||
assert descriptor["start_time"] == start
|
||||
assert descriptor["duration"] == duration
|
||||
result = codec.decode_video(json.loads(json.dumps(descriptor)), payload)
|
||||
second = codec.decode_video(descriptor, payload)
|
||||
attrs = vars(result)
|
||||
assert type(result) is VideoFromFile
|
||||
assert attrs["_VideoFromFile__start_time"] == start
|
||||
assert type(attrs["_VideoFromFile__start_time"]) is type(start)
|
||||
assert attrs["_VideoFromFile__duration"] == duration
|
||||
assert attrs["_VideoFromFile__crop"] == crop
|
||||
if type(start) is float:
|
||||
assert math.copysign(1, attrs["_VideoFromFile__start_time"]) == math.copysign(1, start)
|
||||
decoded = attrs["_VideoFromFile__file"]
|
||||
assert type(decoded) is io.BytesIO
|
||||
assert decoded is not stream
|
||||
stream.close()
|
||||
assert decoded.getvalue() == payload
|
||||
decoded.write(b"changed")
|
||||
assert vars(second)["_VideoFromFile__file"].getvalue() == payload
|
||||
|
||||
def test_real_tiny_pyav_video_is_still_readable_after_source_closed():
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, "w", format="matroska") as container:
|
||||
stream = container.add_stream("ffv1", rate=Fraction(24, 1))
|
||||
# Width 32 avoids a native fillborders filter limitation on tiny
|
||||
# unaligned YUV frames; this is still a tiny CPU-only fixture.
|
||||
stream.width = stream.height = 32
|
||||
stream.pix_fmt = "yuv420p"
|
||||
for value in (16, 64, 128):
|
||||
frame = av.VideoFrame.from_ndarray(np.full((32, 32, 3), value, dtype=np.uint8), format="rgb24")
|
||||
for packet in stream.encode(frame):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode():
|
||||
container.mux(packet)
|
||||
data = buffer.getvalue()
|
||||
source = VideoFromFile(buffer, crop=(0, 0, 2, 2))
|
||||
descriptor, payload = codec.encode_video(source)
|
||||
expected = source.get_components()
|
||||
result = codec.decode_video(descriptor, payload)
|
||||
buffer.close()
|
||||
assert payload == data
|
||||
actual = result.get_components()
|
||||
exact_tensor(actual.images, expected.images)
|
||||
assert actual.frame_rate == expected.frame_rate
|
||||
assert tuple(actual.images.shape) == (3, 2, 2, 3)
|
||||
|
||||
def test_segments_order_complete_audio_no_transcode_and_independent_buffers(monkeypatch):
|
||||
values = [file_value(b"first", start_time=-2), file_value(b"second", duration=-4, crop=(1, 2, 3, 4))]
|
||||
source = VideoFromList([VideoFromList(values[:1]), values[1]], complete_audio=audio())
|
||||
monkeypatch.setattr(VideoFromFile, "get_components", forbid)
|
||||
monkeypatch.setattr(VideoFromFile, "save_to", forbid)
|
||||
monkeypatch.setattr(VideoFromList, "get_components", forbid)
|
||||
monkeypatch.setattr(VideoFromList, "save_to", forbid)
|
||||
descriptor, payload = codec.encode_video(source)
|
||||
assert descriptor["representation"] == "segments"
|
||||
assert payload.startswith(b"firstsecond")
|
||||
assert len(descriptor["segments"]) == 2
|
||||
decoded = codec.decode_video(json.loads(json.dumps(descriptor)), payload)
|
||||
assert type(decoded) is VideoFromList
|
||||
assert [vars(v)["_VideoFromFile__file"].getvalue() for v in decoded.videos] == [b"first", b"second"]
|
||||
exact_tensor(decoded.complete_audio["waveform"], source.complete_audio["waveform"])
|
||||
for old, new in zip(source.videos, decoded.videos):
|
||||
assert old is not new
|
||||
vars(old)["_VideoFromFile__file"].close()
|
||||
assert not vars(new)["_VideoFromFile__file"].closed
|
||||
assert vars(decoded.videos[0])["_VideoFromFile__start_time"] == -2
|
||||
assert vars(decoded.videos[1])["_VideoFromFile__crop"] == (1, 2, 3, 4)
|
||||
|
||||
def test_path_roots_symlinks_and_no_path_in_wire(tmp_path, monkeypatch):
|
||||
allowed = tmp_path / "allowed"
|
||||
allowed.mkdir()
|
||||
media = allowed / "private-name.mkv"
|
||||
media.write_bytes(b"media bytes")
|
||||
sibling = tmp_path / "allowed-sibling"
|
||||
sibling.mkdir()
|
||||
secret = sibling / "secret.mkv"
|
||||
secret.write_bytes(b"private")
|
||||
(allowed / "escape.mkv").symlink_to(secret)
|
||||
(allowed / "inside.mkv").symlink_to(media)
|
||||
descriptor, payload = codec.encode_video(VideoFromFile(str(media)), allowed_roots=[allowed])
|
||||
assert payload == b"media bytes"
|
||||
assert "private-name" not in json.dumps(descriptor)
|
||||
assert codec.encode_video(VideoFromFile(str(allowed / "inside.mkv")), allowed_roots=[allowed])[1] == payload
|
||||
for path in (secret, allowed / "escape.mkv", allowed, allowed / "missing"):
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(VideoFromFile(str(path)), allowed_roots=[allowed])
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(VideoFromFile(str(media)), allowed_roots=[])
|
||||
import folder_paths
|
||||
for name in ("get_input_directory", "get_output_directory", "get_temp_directory"):
|
||||
monkeypatch.setattr(folder_paths, name, lambda: str(allowed))
|
||||
assert codec.encode_video(VideoFromFile(str(media)))[1] == payload
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(VideoFromFile(str(secret)))
|
||||
result = codec.decode_video(descriptor, payload)
|
||||
media.unlink()
|
||||
assert vars(result)["_VideoFromFile__file"].getvalue() == b"media bytes"
|
||||
|
||||
def test_path_special_files_and_symlink_swap_are_rejected(tmp_path, monkeypatch):
|
||||
fifo = tmp_path / "fifo"
|
||||
os.mkfifo(fifo)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(VideoFromFile(str(fifo)), allowed_roots=[tmp_path])
|
||||
allowed = tmp_path / "allowed"
|
||||
allowed.mkdir()
|
||||
actual = allowed / "subdir"
|
||||
actual.mkdir()
|
||||
media = actual / "clip.mkv"
|
||||
media.write_bytes(b"allowed")
|
||||
outside = tmp_path / "outside"
|
||||
outside.mkdir()
|
||||
(outside / "clip.mkv").write_bytes(b"forbidden")
|
||||
real_resolve = codec.Path.resolve
|
||||
swapped = False
|
||||
|
||||
def swap_after_resolution(path, *args, **kwargs):
|
||||
nonlocal swapped
|
||||
resolved = real_resolve(path, *args, **kwargs)
|
||||
if path == media and not swapped:
|
||||
swapped = True
|
||||
actual.rename(allowed / "old")
|
||||
actual.symlink_to(outside, target_is_directory=True)
|
||||
return resolved
|
||||
|
||||
monkeypatch.setattr(codec.Path, "resolve", swap_after_resolution)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(VideoFromFile(str(media)), allowed_roots=[allowed])
|
||||
assert swapped
|
||||
|
||||
def test_file_size_budget_checked_before_read(tmp_path, monkeypatch):
|
||||
path = tmp_path / "large.mkv"
|
||||
path.write_bytes(b"12345")
|
||||
monkeypatch.setattr(codec.os, "read", forbid)
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.encode_video(VideoFromFile(str(path)), max_bytes=4, allowed_roots=[tmp_path])
|
||||
|
||||
@pytest.mark.parametrize("limit", [True, False, -1, 1.5, None])
|
||||
def test_invalid_limits(limit):
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(file_value(), max_bytes=limit)
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video({}, b"", max_bytes=limit)
|
||||
|
||||
def test_aggregate_budget_and_exact_boundary():
|
||||
value = VideoFromList([file_value(b"12"), file_value(b"345")])
|
||||
descriptor, payload = codec.encode_video(value, max_bytes=5)
|
||||
assert payload == b"12345"
|
||||
codec.decode_video(descriptor, payload, max_bytes=5)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value, max_bytes=4)
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload, max_bytes=4)
|
||||
stream = vars(value.videos[0])["_VideoFromFile__file"]
|
||||
stream.seek(1)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value.videos[0], max_bytes=1)
|
||||
assert stream.tell() == 1
|
||||
|
||||
@pytest.mark.parametrize("kind", ["components", "file", "segments"])
|
||||
def test_unknown_subclasses_never_materialized(kind):
|
||||
bases = {"components": VideoFromComponents, "file": VideoFromFile, "segments": VideoFromList}
|
||||
cls = type("UnknownVideo", (bases[kind],), {"get_components": forbid, "save_to": forbid})
|
||||
value = object.__new__(cls)
|
||||
with pytest.raises(ValueError, match="class|type|unsupported|Unsupported"):
|
||||
codec.encode_video(value)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(object())
|
||||
|
||||
@pytest.mark.parametrize("kind", ["components", "file", "segments"])
|
||||
def test_bridge_rejects_subclasses_without_materializing(kind, monkeypatch):
|
||||
nodes = importlib.import_module(package_name + ".nodes")
|
||||
bases = {"components": VideoFromComponents, "file": VideoFromFile, "segments": VideoFromList}
|
||||
cls = type("UnknownVideo", (bases[kind],), {"get_components": forbid, "save_to": forbid})
|
||||
monkeypatch.setattr(torch.Tensor, "cpu", forbid)
|
||||
with pytest.raises(ValueError, match="Unsupported native VIDEO class"):
|
||||
nodes.encode_value("VIDEO", object.__new__(cls))
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
||||
@pytest.mark.parametrize("kind", ["components", "segments", "file"])
|
||||
def test_bridge_stages_cpu_tensors_without_materialization(dtype, kind, monkeypatch):
|
||||
nodes = importlib.import_module(package_name + ".nodes")
|
||||
if kind == "components":
|
||||
source = components(dtype, metadata={"nested": ["雪", None, True]})
|
||||
state = vars(source)["_VideoFromComponents__components"]
|
||||
state.alpha = state.alpha.unsqueeze(-1)
|
||||
tensors = [state.images, state.alpha, state.audio["waveform"]]
|
||||
elif kind == "segments":
|
||||
source = VideoFromList([file_value(start_time=-1.5, duration=-2, crop=(-1, 2, 3, 4))],
|
||||
complete_audio=audio(dtype))
|
||||
tensors = [source.complete_audio["waveform"]]
|
||||
else:
|
||||
source = file_value(start_time=-0.0, duration=-2, crop=(-1, 2, 3, 4))
|
||||
tensors = []
|
||||
for tensor in tensors:
|
||||
tensor.requires_grad_(True)
|
||||
expected = codec.encode_video(source)
|
||||
detached, staged = [], []
|
||||
real_detach, real_cpu = torch.Tensor.detach, torch.Tensor.cpu
|
||||
|
||||
def detach(tensor):
|
||||
detached.append(id(tensor))
|
||||
return real_detach(tensor)
|
||||
|
||||
def cpu(tensor, *args, **kwargs):
|
||||
assert not tensor.requires_grad
|
||||
staged.append(tensor)
|
||||
return real_cpu(tensor, *args, **kwargs)
|
||||
|
||||
for cls in (VideoFromComponents, VideoFromFile, VideoFromList):
|
||||
monkeypatch.setattr(cls, "get_components", forbid)
|
||||
monkeypatch.setattr(cls, "save_to", forbid)
|
||||
monkeypatch.setattr(torch.Tensor, "detach", detach)
|
||||
monkeypatch.setattr(torch.Tensor, "cpu", cpu)
|
||||
actual = nodes.encode_value("VIDEO", source)
|
||||
assert len(staged) == len(tensors)
|
||||
assert all(id(tensor) in detached for tensor in tensors)
|
||||
assert all(tensor.requires_grad for tensor in tensors)
|
||||
assert actual == expected
|
||||
decoded = nodes.decode_value("VIDEO", *actual)
|
||||
assert type(decoded) is type(source)
|
||||
assert codec.encode_video(decoded) == expected
|
||||
assert len(staged) == len(tensors) # Pure codec never stages implicitly.
|
||||
|
||||
@pytest.mark.parametrize("field", ["images", "alpha", "audio", "complete_audio"])
|
||||
def test_pure_video_rejects_non_cpu_storage_without_transfer(field, monkeypatch):
|
||||
if field == "complete_audio":
|
||||
source = VideoFromList([file_value()], complete_audio=audio())
|
||||
target = source.complete_audio["waveform"]
|
||||
else:
|
||||
source = components()
|
||||
target = getattr(vars(source)["_VideoFromComponents__components"], field)
|
||||
if field == "audio":
|
||||
target = target["waveform"]
|
||||
# A property spy on real CPU tensors exercises the guard without a GPU
|
||||
# allocation, subclass, device discovery or CUDA initialization.
|
||||
real_device = torch.Tensor.device
|
||||
monkeypatch.setattr(torch.Tensor, "device", property(
|
||||
lambda tensor: types.SimpleNamespace(type="cuda") if tensor is target
|
||||
else real_device.__get__(tensor)))
|
||||
monkeypatch.setattr(torch.Tensor, "cpu", forbid)
|
||||
with pytest.raises(ValueError, match="must be on CPU"):
|
||||
codec.encode_video(source)
|
||||
|
||||
@pytest.mark.parametrize("field,bad", [
|
||||
("_VideoFromFile__start_time", True), ("_VideoFromFile__start_time", float("nan")),
|
||||
("_VideoFromFile__duration", float("inf")), ("_VideoFromFile__crop", (1, 2, True, 4)),
|
||||
("_VideoFromFile__crop", [1, 2, 3, 4]), ("_VideoFromFile__crop", (1, 2, 3)),
|
||||
("_VideoFromFile__file", b"not a BytesIO"),
|
||||
])
|
||||
def test_invalid_native_file_attributes(field, bad):
|
||||
value = file_value()
|
||||
setattr(value, field, bad)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
|
||||
def test_missing_attributes_and_invalid_component_state():
|
||||
value = file_value()
|
||||
del vars(value)["_VideoFromFile__duration"]
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
for field, bad in [("frame_rate", 24.0), ("frame_rate", Fraction(0)),
|
||||
("alpha", torch.ones(1, 1, 1, device="cpu"))]:
|
||||
value = components()
|
||||
setattr(vars(value)["_VideoFromComponents__components"], field, bad)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
|
||||
@pytest.mark.parametrize("metadata", [
|
||||
{"value": float("nan")}, {"value": float("inf")}, {1: "non-string key"},
|
||||
{"value": (1, 2)}, {"value": object()}, {"large": "x" * 65536},
|
||||
])
|
||||
def test_invalid_metadata_rejected_both_directions(metadata):
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(components(metadata=metadata))
|
||||
descriptor, payload = codec.encode_video(components())
|
||||
descriptor["metadata"] = metadata
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
def test_metadata_depth_cycles_and_exact_size():
|
||||
metadata = {}
|
||||
current = metadata
|
||||
for _ in range(40):
|
||||
current["child"] = {}
|
||||
current = current["child"]
|
||||
for value in (metadata,):
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(components(metadata=value))
|
||||
cyclic = {}
|
||||
cyclic["self"] = cyclic
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(components(metadata=cyclic))
|
||||
exact = {"x": "a" * (65536 - 8)}
|
||||
assert len(json.dumps(exact, separators=(",", ":")).encode()) == 65536
|
||||
descriptor, payload = codec.encode_video(components(metadata=exact))
|
||||
assert codec.decode_video(descriptor, payload).get_components().metadata == exact
|
||||
|
||||
@pytest.mark.parametrize("field,bad", [
|
||||
("encoding", "other"), ("representation", "other"), ("byte_length", True),
|
||||
("byte_length", -1), ("byte_length", 12345), ("start_time", False),
|
||||
("duration", float("inf")), ("crop", [0, 0, True, 4]), ("crop", [1, 2, 3]),
|
||||
])
|
||||
def test_malformed_file_descriptors(field, bad):
|
||||
descriptor, payload = codec.encode_video(file_value())
|
||||
descriptor[field] = bad
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
@pytest.mark.parametrize("field,bad", [("offset", True), ("offset", -1), ("offset", 1),
|
||||
("length", True), ("length", -1), ("length", 2**80)])
|
||||
def test_invalid_offsets_lengths(field, bad):
|
||||
descriptor, payload = codec.encode_video(file_value())
|
||||
descriptor["data"][field] = bad
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
def test_truncation_trailing_gaps_overlap_and_unknown_fields():
|
||||
value = VideoFromList([file_value(b"123"), file_value(b"456")])
|
||||
original, payload = codec.encode_video(value)
|
||||
for wire in (payload[:-1], payload + b"x"):
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(original, wire)
|
||||
for offset in (2, 4):
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["segments"][1]["offset"] = offset
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["surprise"] = "bad"
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["segments"][0]["descriptor"]["byte_length"] += 1
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["segments"][0]["descriptor"] = descriptor
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
def test_all_descriptors_validated_before_tensor_allocation(monkeypatch):
|
||||
original, payload = codec.encode_video(components())
|
||||
monkeypatch.setattr(codec, "decode_payload", forbid)
|
||||
cases = []
|
||||
for key, value in [("shape", [2**80, 2, 3]), ("shape", [True, 2, 3]),
|
||||
("dtype", "pickle"), ("layout", "BHWC")]:
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["alpha"]["descriptor"][key] = value
|
||||
cases.append(descriptor)
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["audio"]["descriptor"]["sample_rate"] = True
|
||||
cases.append(descriptor)
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["alpha"]["offset"] += 1
|
||||
cases.append(descriptor)
|
||||
for descriptor in cases:
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
@pytest.mark.parametrize("field,bad", [("numerator", True), ("denominator", True),
|
||||
("denominator", 0), ("numerator", -3),
|
||||
("denominator", -1001)])
|
||||
def test_invalid_fraction_fields(field, bad):
|
||||
descriptor, payload = codec.encode_video(components())
|
||||
descriptor["frame_rate"][field] = bad
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
def test_segment_limit_and_exact_maximum(monkeypatch):
|
||||
value = VideoFromList([file_value(b"x") for _ in range(64)])
|
||||
descriptor, payload = codec.encode_video(value)
|
||||
result = codec.decode_video(descriptor, payload)
|
||||
assert len(result.videos) == 64
|
||||
value.videos.append(file_value(b"x"))
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
descriptor["segments"].append(copy.deepcopy(descriptor["segments"][-1]))
|
||||
monkeypatch.setattr(codec, "decode_payload", forbid)
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
value.videos = []
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
value.videos = [components()]
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
|
||||
@pytest.mark.parametrize("descriptor,payload", [(None, b""), ([], b""), ({}, bytearray()), ({}, "")])
|
||||
def test_invalid_top_level_types(descriptor, payload):
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
@@ -60,6 +60,29 @@ class ValidateServerlessWorkerTests(unittest.TestCase):
|
||||
self.assertIsNot(result, worker)
|
||||
self.assertEqual(worker, {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"})
|
||||
|
||||
def test_omitted_media_profile_remains_absent_and_means_legacy(self):
|
||||
result = serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": "ep-123"}
|
||||
)
|
||||
self.assertNotIn("media_profile", result)
|
||||
self.assertEqual(result.get("media_profile", "legacy-image-v1"), "legacy-image-v1")
|
||||
|
||||
def test_explicit_media_profiles_are_retained(self):
|
||||
for profile in ("legacy-image-v1", "comfy-media-v1"):
|
||||
with self.subTest(profile=profile):
|
||||
result = serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": "ep-123", "media_profile": profile}
|
||||
)
|
||||
self.assertEqual(result["media_profile"], profile)
|
||||
|
||||
def test_media_profile_rejects_unknown_nonstring_and_null_values(self):
|
||||
for profile in ("", "auto", "comfy-media-v2", " comfy-media-v1", "COMFY-MEDIA-V1", None, True, 1, [], {}):
|
||||
with self.subTest(profile=profile):
|
||||
with self.assertRaisesRegex(ValueError, "media_profile"):
|
||||
serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": "ep-123", "media_profile": profile}
|
||||
)
|
||||
|
||||
def test_accepts_inclusive_numeric_bounds(self):
|
||||
for jobs, timeout in ((1, 30), (16, 86400)):
|
||||
with self.subTest(jobs=jobs, timeout=timeout):
|
||||
@@ -307,6 +330,7 @@ class PublicConfigTests(unittest.TestCase):
|
||||
"type": "serverless",
|
||||
"provider": "runpod",
|
||||
"endpoint_id": "ep-123",
|
||||
"media_profile": "comfy-media-v1",
|
||||
"api_key": "inline-secret",
|
||||
"clear_api_key": "another-secret",
|
||||
"credential": "credential-secret",
|
||||
@@ -334,6 +358,8 @@ class PublicConfigTests(unittest.TestCase):
|
||||
self.assertNotIn("api_key", local)
|
||||
self.assertNotIn("credentials_note", result["settings"])
|
||||
self.assertEqual(first["name"], "remote GPU")
|
||||
self.assertEqual(first["media_profile"], "comfy-media-v1")
|
||||
self.assertNotIn("media_profile", second)
|
||||
self.assertEqual(config, original)
|
||||
self.assertIsNot(result, config)
|
||||
self.assertIsNot(result["workers"], config["workers"])
|
||||
|
||||
@@ -0,0 +1,426 @@
|
||||
"""CPU-only graph/compiler tests; no ComfyUI package bootstrap is needed."""
|
||||
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE = "serverless_media_prompt_testpkg"
|
||||
# Real paths let sibling imports use production code, without running the
|
||||
# extension's ComfyUI-dependent __init__.py or API registration side effects.
|
||||
for suffix in ("", ".api", ".api.orchestration", ".serverless_bridge"):
|
||||
package = types.ModuleType(PACKAGE + suffix)
|
||||
package.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
|
||||
sys.modules[package.__name__] = package
|
||||
|
||||
compiler = importlib.import_module(f"{PACKAGE}.api.orchestration.serverless_media_prompt")
|
||||
legacy = importlib.import_module(f"{PACKAGE}.api.orchestration.serverless_prompt")
|
||||
contract = importlib.import_module(f"{PACKAGE}.serverless_bridge.contract")
|
||||
compile_media_prompt = compiler.compile_media_prompt
|
||||
SOCKETS = {"images": "IMAGE", "audio": "AUDIO"}
|
||||
|
||||
|
||||
def _node(class_type, **inputs):
|
||||
return {"class_type": class_type, "inputs": inputs}
|
||||
|
||||
|
||||
def _collector(**sockets):
|
||||
return _node(
|
||||
"DistributedCollector", multi_job_id="job/42", worker_id="worker_0", **sockets
|
||||
)
|
||||
|
||||
|
||||
def _prompt():
|
||||
return {
|
||||
"source:1": _node("LoadImage", image="example.png"),
|
||||
"collector:1": _collector(images=["source:1", 0]),
|
||||
}
|
||||
|
||||
|
||||
def _run(prompt):
|
||||
return compile_media_prompt(prompt, "request:1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("socket,kind", SOCKETS.items())
|
||||
def test_one_export_for_each_typed_socket_with_exact_canonical_binding(socket, kind):
|
||||
prompt = {"source": _node("Source"), "collector": _collector(**{socket: ["source", 0]})}
|
||||
graph, bindings = _run(prompt)
|
||||
assert len(bindings) == 1
|
||||
binding = bindings[0]
|
||||
assert set(binding) == contract.BINDING_FIELDS
|
||||
assert binding == {
|
||||
"binding_id": binding["binding_id"],
|
||||
"collector_id": "collector",
|
||||
"job_id": "job/42",
|
||||
"participant_id": "worker_0",
|
||||
"node_id": binding["node_id"],
|
||||
"socket": socket,
|
||||
"kind": kind,
|
||||
}
|
||||
assert all(isinstance(value, str) for value in binding.values())
|
||||
assert contract.validate_bindings(bindings) == {binding["binding_id"]: binding}
|
||||
assert graph[binding["node_id"]] == {
|
||||
"class_type": "MediaExportV1",
|
||||
"inputs": {
|
||||
"request_id": "request:1",
|
||||
"binding_json": json.dumps(binding),
|
||||
socket: ["source", 0],
|
||||
**({"transfer_format": "png", "transfer_quality": 90} if kind == "IMAGE" else {}),
|
||||
},
|
||||
}
|
||||
assert set(graph) == {"source", binding["node_id"]}
|
||||
|
||||
|
||||
def test_multiple_collectors_have_independent_image_audio_targets():
|
||||
prompt = {
|
||||
"shared": _node("MultiMediaSource"),
|
||||
"collector:z": _collector(audio=["shared", 1]),
|
||||
"collector:a": _collector(images=["shared", 0], audio=["shared", 1]),
|
||||
}
|
||||
prompt["collector:z"]["inputs"].update(multi_job_id="job:z", worker_id="participant:z")
|
||||
graph, bindings = _run(prompt)
|
||||
assert [(b["collector_id"], b["socket"]) for b in bindings] == [
|
||||
("collector:a", "images"), ("collector:a", "audio"), ("collector:z", "audio"),
|
||||
]
|
||||
assert len({b["binding_id"] for b in bindings}) == 3
|
||||
targets = {b["node_id"] for b in bindings}
|
||||
assert len(targets) == 3
|
||||
assert set(graph) == {"shared"} | targets
|
||||
for binding in bindings:
|
||||
export = graph[binding["node_id"]]
|
||||
inputs = prompt[binding["collector_id"]]["inputs"]
|
||||
assert export["inputs"][binding["socket"]] == inputs[binding["socket"]]
|
||||
assert binding["job_id"] == inputs["multi_job_id"]
|
||||
assert binding["participant_id"] == inputs["worker_id"]
|
||||
|
||||
|
||||
def test_prunes_descendants_and_unrelated_outputs_preserving_required_output_utility():
|
||||
prompt = _prompt()
|
||||
prompt.update({
|
||||
"required-output": _node("RequiredOutputUtility", images=["source:1", 0]),
|
||||
"consumer": _node("ImageTransform", image=["required-output", 0]),
|
||||
"descendant": _node("ImageTransform", image=["collector:1", 0]),
|
||||
"after": _node("SaveImage", images=["descendant", 0]),
|
||||
"unused-output": _node("SaveImage", images=["source:1", 0]),
|
||||
"unrelated": _node("LoadAudio"),
|
||||
})
|
||||
prompt["required-output"]["_meta"] = {"title": "Required OUTPUT_NODE utility"}
|
||||
prompt["collector:1"]["inputs"].update(images=["consumer", 0], load_balance=True, pass_through=False)
|
||||
graph, bindings = _run(prompt)
|
||||
assert set(graph) == {"source:1", "required-output", "consumer", bindings[0]["node_id"]}
|
||||
assert graph["required-output"] == prompt["required-output"]
|
||||
assert not {"load_balance", "pass_through"} & graph[bindings[0]["node_id"]]["inputs"].keys()
|
||||
|
||||
|
||||
def test_upstream_primitive_string_preserves_empty_and_link_looking_text():
|
||||
for text in ("", '["source:1", 0]', "hello\nworld"):
|
||||
prompt = {"text": _node("PrimitiveString", value=text),
|
||||
"source": _node("TextToImage", text=["text", 0]),
|
||||
"collector": _collector(images=["source", 0])}
|
||||
graph, bindings = _run(prompt)
|
||||
assert graph["text"] == prompt["text"]
|
||||
assert graph["source"]["inputs"]["text"] == ["text", 0]
|
||||
assert graph[bindings[0]["node_id"]]["inputs"]["images"] == ["source", 0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type,inputs", [
|
||||
("DistributedSeed", {}),
|
||||
("DistributedSeed", {"seed": 10, "is_worker": True, "worker_id": "worker_1"}),
|
||||
("DistributedSeed", {"seed": 10, "is_worker": False, "worker_id": "worker_1"}),
|
||||
("DistributedSeed", {"seed": 10, "is_worker": True, "worker_id": "bad"}),
|
||||
("DistributedValue", {"default_value": "3.9", "worker_values": '{"_type":"INT","1":"8.7"}', "is_worker": True, "worker_id": "worker_0"}),
|
||||
("DistributedValue", {"default_value": "3.9", "worker_values": '{"_type":"INT","2":""}', "is_worker": True, "worker_id": "worker_1"}),
|
||||
("DistributedValue", {"default_value": "fallback", "worker_values": "malformed"}),
|
||||
("DistributedValue", {"default_value": 1.25, "worker_values": '{"_type":"FLOAT","1":"bad"}', "is_worker": True, "worker_id": "0"}),
|
||||
("DistributedModelName", {"text": "model.safetensors"}),
|
||||
("DistributedModelName", {"text": {"name": "checkpoint"}}),
|
||||
("DistributedModelName", {"text": ["a", "b", "c"]}),
|
||||
])
|
||||
def test_upstream_scalar_values_match_legacy_compiler(class_type, inputs):
|
||||
prompt = _prompt()
|
||||
prompt["scalar"] = _node(class_type, **inputs)
|
||||
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 0]
|
||||
old_prompt = copy.deepcopy(prompt)
|
||||
old_prompt["collector:1"]["class_type"] = "DistributedCollector"
|
||||
expected, _ = legacy.compile_serverless_prompt(old_prompt)
|
||||
graph, _ = _run(prompt)
|
||||
assert graph["source:1"] == expected["source:1"]
|
||||
assert "scalar" not in graph
|
||||
|
||||
|
||||
def test_imports_same_scalar_resolver_and_resolves_primitive_chains():
|
||||
assert compiler._make_scalar_resolver is legacy._make_scalar_resolver
|
||||
prompt = _prompt()
|
||||
prompt.update({
|
||||
"int": _node("PrimitiveInt", value=25),
|
||||
"int-alias": _node("PrimitiveInt", value=["int", 0]),
|
||||
"float": _node("PrimitiveFloat", value=1.25),
|
||||
"bool": _node("PrimitiveBoolean", value=True),
|
||||
"string": _node("PrimitiveString", value="model.safetensors"),
|
||||
"seed": _node("DistributedSeed", seed=["int-alias", 0], is_worker=True, worker_id="worker_0"),
|
||||
"value": _node("DistributedValue", default_value=["float", 0], worker_values='{"_type":"FLOAT","1":"2.75"}', is_worker=True, worker_id="worker_0"),
|
||||
"enabled": _node("DistributedValue", default_value=["bool", 0]),
|
||||
"model": _node("DistributedModelName", text=["string", 0]),
|
||||
})
|
||||
prompt["source:1"]["inputs"].update(seed=["seed", 0], cfg=["value", 0], enabled=["enabled", 0], model=["model", 0])
|
||||
graph, bindings = _run(prompt)
|
||||
assert graph["source:1"]["inputs"] == {"image": "example.png", "seed": 26, "cfg": 2.75, "enabled": True, "model": "model.safetensors"}
|
||||
assert set(graph) == {"source:1", bindings[0]["node_id"]}
|
||||
|
||||
|
||||
def test_distributed_text_value_can_be_wrapped_by_primitive_string():
|
||||
prompt = {
|
||||
"value": _node("DistributedModelName", text="model.safetensors"),
|
||||
"text": _node("PrimitiveString", value=["value", 0]),
|
||||
"source": _node("TextToImage", text=["text", 0]),
|
||||
"collector": _collector(images=["source", 0]),
|
||||
}
|
||||
graph, bindings = _run(prompt)
|
||||
assert graph["text"]["inputs"]["value"] == "model.safetensors"
|
||||
assert graph["source"]["inputs"]["text"] == ["text", 0]
|
||||
assert graph[bindings[0]["node_id"]]["inputs"]["images"] == ["source", 0]
|
||||
assert "value" not in graph
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type,input_name", [
|
||||
("DistributedSeed", "seed"), ("DistributedValue", "default_value"), ("DistributedModelName", "text"),
|
||||
])
|
||||
def test_dynamic_scalar_resolution_stays_fail_closed(class_type, input_name):
|
||||
prompt = _prompt()
|
||||
prompt["dynamic"] = _node("ArbitraryValue")
|
||||
prompt["scalar"] = _node(class_type, **{input_name: ["dynamic", 0]})
|
||||
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 0]
|
||||
with pytest.raises(ValueError, match="literal or supported Primitive"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_namespaced_export_ids_avoid_existing_ids_even_on_pruned_nodes():
|
||||
prompt = _prompt()
|
||||
_, first = _run(prompt)
|
||||
original_id = first[0]["node_id"]
|
||||
assert ":" in original_id and not original_id.isdigit()
|
||||
prompt[original_id] = _node("ExistingUtility")
|
||||
_, second = _run(prompt)
|
||||
assert second[0]["node_id"] not in prompt
|
||||
assert second[0]["binding_id"] == first[0]["binding_id"]
|
||||
prompt[second[0]["node_id"]] = _node("AnotherUtility")
|
||||
graph, third = _run(prompt)
|
||||
assert third[0]["node_id"] not in prompt
|
||||
assert third == _run(dict(reversed(list(prompt.items()))))[1]
|
||||
assert set(graph) == {"source:1", third[0]["node_id"]}
|
||||
|
||||
|
||||
def test_long_valid_identifiers_remain_exact_and_generated_ids_fit_contract():
|
||||
collector_id = "namespace/" + "c" * 190
|
||||
prompt = {"source:1": _node("Source"), collector_id: _collector(images=["source:1", 0])}
|
||||
prompt[collector_id]["inputs"].update(multi_job_id="j" * 200, worker_id="p" * 200)
|
||||
_, bindings = _run(prompt)
|
||||
assert bindings[0]["collector_id"] == collector_id
|
||||
assert bindings[0]["job_id"] == "j" * 200
|
||||
assert bindings[0]["participant_id"] == "p" * 200
|
||||
contract.validate_bindings(bindings)
|
||||
|
||||
|
||||
def test_deterministic_independent_copies_and_request_identity():
|
||||
prompt = _prompt()
|
||||
prompt["source:1"]["_meta"] = {"nested": {"values": [1, 2, 3]}}
|
||||
original = copy.deepcopy(prompt)
|
||||
graph, bindings = _run(prompt)
|
||||
assert prompt == original
|
||||
assert (graph, bindings) == _run(prompt)
|
||||
assert (graph, bindings) == _run(dict(reversed(list(prompt.items()))))
|
||||
other_graph, other_bindings = compile_media_prompt(prompt, "request:2")
|
||||
assert other_bindings == bindings
|
||||
target = bindings[0]["node_id"]
|
||||
assert other_graph[target]["inputs"]["request_id"] == "request:2"
|
||||
graph["source:1"]["_meta"]["nested"]["values"].append(4)
|
||||
graph[target]["inputs"]["images"][0] = "changed"
|
||||
bindings[0]["collector_id"] = "changed"
|
||||
assert prompt == original
|
||||
assert json.loads(graph[target]["inputs"]["binding_json"])["collector_id"] == "collector:1"
|
||||
assert other_graph["source:1"] == original["source:1"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [None, [], {}, {"x": None}, {"x": []}, {"x": {}},
|
||||
{"x": {"class_type": None}}, {"x": {"class_type": ""}},
|
||||
{"x": {"class_type": "Source", "inputs": []}},
|
||||
{"x": {"class_type": "Source", "inputs": None}},
|
||||
{"x": {"class_type": "Source", "inputs": {1: "literal"}}},
|
||||
])
|
||||
def test_rejects_malformed_prompt_or_missing_collectors(bad):
|
||||
with pytest.raises(ValueError):
|
||||
_run(bad)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [None, 1, True, "", " bad", "bad ", "bad\\id", "bad\nid", "x" * 201])
|
||||
def test_rejects_malformed_ids_without_normalizing(bad):
|
||||
prompt = _prompt()
|
||||
prompt[bad] = _node("Source")
|
||||
with pytest.raises(ValueError):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [None, 1, True, "", " bad", "bad ", "bad\\id", "x" * 201])
|
||||
@pytest.mark.parametrize("field", ["multi_job_id", "worker_id", "request_id"])
|
||||
def test_rejects_malformed_binding_and_request_identity(field, bad):
|
||||
prompt = _prompt()
|
||||
request_id = "request:1"
|
||||
if field == "request_id":
|
||||
request_id = bad
|
||||
else:
|
||||
prompt["collector:1"]["inputs"][field] = bad
|
||||
with pytest.raises(ValueError):
|
||||
compile_media_prompt(prompt, request_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["multi_job_id", "worker_id"])
|
||||
def test_requires_prepared_literal_hidden_identity(field):
|
||||
prompt = _prompt()
|
||||
del prompt["collector:1"]["inputs"][field]
|
||||
with pytest.raises(ValueError):
|
||||
_run(prompt)
|
||||
prompt["collector:1"]["inputs"][field] = ["source:1", 0]
|
||||
with pytest.raises(ValueError):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("socket", SOCKETS)
|
||||
@pytest.mark.parametrize("value", [None, "literal", 5, True, [], ["source:1"], ["source:1", 0, 1], ("source:1", 0)])
|
||||
def test_every_present_media_socket_requires_an_api_link(socket, value):
|
||||
prompt = _prompt()
|
||||
prompt["collector:1"]["inputs"][socket] = value
|
||||
with pytest.raises(ValueError, match="link"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_each_collector_requires_at_least_one_media_socket():
|
||||
prompt = _prompt()
|
||||
prompt["empty"] = _collector(load_balance=True, pass_through=False)
|
||||
with pytest.raises(ValueError, match="connected|linked"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("link", [["missing", 0], ["source:1", -1], ["source:1", 0.0],
|
||||
["source:1", "0"], ["source:1", True], ["source:1", False], [1, 0],
|
||||
[True, 0], [None, 0], [[], 0], [{}, 0], [" source:1", 0],
|
||||
])
|
||||
@pytest.mark.parametrize("location", ["collector", "ordinary", "unrelated"])
|
||||
def test_rejects_invalid_links_everywhere_before_pruning(location, link):
|
||||
prompt = _prompt()
|
||||
prompt["1"] = _node("Source") # An integer/bool source must not be stringified to this ID.
|
||||
if location == "collector":
|
||||
prompt["collector:1"]["inputs"]["images"] = link
|
||||
elif location == "ordinary":
|
||||
prompt["source:1"]["inputs"]["parameter"] = link
|
||||
else:
|
||||
prompt["unrelated"] = _node("Utility", value=link)
|
||||
with pytest.raises(ValueError):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type", [
|
||||
"DistributedMediaCollector", "DistributedEmptyImage", "UltimateSDUpscaleDistributed",
|
||||
"DistributedImageListStart", "DistributedImageListEnd", "DistributedWorkerSink",
|
||||
"ImageBatchDivider", "AudioBatchDivider", "CustomDistributedNode", "custombatchdivider",
|
||||
])
|
||||
def test_rejects_unsupported_distributed_families_even_if_unused(class_type):
|
||||
prompt = _prompt()
|
||||
prompt["unsupported"] = _node(class_type)
|
||||
with pytest.raises(ValueError, match="unsupported"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type", ["DistributedSeed", "DistributedValue", "DistributedModelName", "PrimitiveInt", "PrimitiveString"])
|
||||
def test_rejects_invalid_known_scalar_output_slots(class_type):
|
||||
prompt = _prompt()
|
||||
prompt["scalar"] = _node(class_type)
|
||||
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 1]
|
||||
with pytest.raises(ValueError, match="slot"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("socket", SOCKETS)
|
||||
def test_rejects_direct_distributed_scalar_media_links(socket):
|
||||
prompt = {"scalar": _node("DistributedValue", default_value="text"), "collector": _collector(**{socket: ["scalar", 0]})}
|
||||
with pytest.raises(ValueError, match="scalar"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("indirect", [False, True])
|
||||
def test_rejects_collector_chains_across_different_media(indirect):
|
||||
prompt = _prompt()
|
||||
source = "collector:1"
|
||||
if indirect:
|
||||
prompt["processing"] = _node("ImageToAudio", images=[source, 0])
|
||||
source = "processing"
|
||||
prompt["collector:2"] = _collector(audio=[source, 0])
|
||||
with pytest.raises(ValueError, match="collector-to-collector"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["ordinary", "self", "unrelated", "primitive", "collector"])
|
||||
def test_rejects_cycles_in_entire_prompt(mode):
|
||||
prompt = _prompt()
|
||||
if mode == "ordinary":
|
||||
prompt["source:1"]["inputs"]["other"] = ["other", 0]
|
||||
prompt["other"] = _node("Utility", value=["source:1", 0])
|
||||
elif mode == "self":
|
||||
prompt["source:1"]["inputs"]["other"] = ["source:1", 0]
|
||||
elif mode == "unrelated":
|
||||
prompt["unused"] = _node("Utility", value=["unused", 0])
|
||||
elif mode == "primitive":
|
||||
prompt["p1"] = _node("PrimitiveInt", value=["p2", 0])
|
||||
prompt["p2"] = _node("PrimitiveInt", value=["p1", 0])
|
||||
else:
|
||||
prompt["source:1"]["inputs"]["other"] = ["collector:1", 0]
|
||||
with pytest.raises(ValueError, match="cycle"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_deep_acyclic_graph_does_not_depend_on_python_recursion_limit():
|
||||
prompt = {"n:0": _node("Source")}
|
||||
for index in range(1, 1500):
|
||||
prompt[f"n:{index}"] = _node("Utility", value=[f"n:{index - 1}", 0])
|
||||
prompt["collector"] = _collector(images=["n:1499", 0])
|
||||
graph, bindings = _run(prompt)
|
||||
assert len(graph) == 1501
|
||||
assert graph[bindings[0]["node_id"]]["inputs"]["images"] == ["n:1499", 0]
|
||||
|
||||
|
||||
def test_binding_count_respects_bridge_contract():
|
||||
prompt = {"source": _node("Source")}
|
||||
for index in range(contract.MAX_ARTIFACTS + 1):
|
||||
prompt[f"collector:{index}"] = _collector(images=["source", 0])
|
||||
with pytest.raises(ValueError, match="count"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_import_and_compile_do_not_import_comfy_or_tensor_libraries():
|
||||
script = '''
|
||||
import builtins, importlib, pathlib, sys, types
|
||||
root = pathlib.Path(sys.argv[1])
|
||||
name = "isolated_media_compile"
|
||||
for suffix in ("", ".api", ".api.orchestration", ".serverless_bridge"):
|
||||
package = types.ModuleType(name + suffix)
|
||||
package.__path__ = [str(root.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(root)]
|
||||
sys.modules[package.__name__] = package
|
||||
original_import = builtins.__import__
|
||||
def guarded_import(name, *args, **kwargs):
|
||||
if name.split(".")[0] in {"comfy", "comfy_api", "nodes", "server", "torch", "numpy"}:
|
||||
raise AssertionError("unexpected runtime import: " + name)
|
||||
return original_import(name, *args, **kwargs)
|
||||
builtins.__import__ = guarded_import
|
||||
module = importlib.import_module(name + ".api.orchestration.serverless_media_prompt")
|
||||
prompt = {"s": {"class_type": "Source", "inputs": {}}, "c": {"class_type": "DistributedCollector", "inputs": {"images": ["s", 0], "multi_job_id": "job", "worker_id": "worker"}}}
|
||||
graph, bindings = module.compile_media_prompt(prompt, "request")
|
||||
assert graph[bindings[0]["node_id"]]["class_type"] == "MediaExportV1"
|
||||
'''
|
||||
result = subprocess.run([sys.executable, "-c", script, str(ROOT)], text=True, capture_output=True)
|
||||
assert result.returncode == 0, result.stderr
|
||||
@@ -217,6 +217,78 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
|
||||
self.assertNotIn(obsolete, saved)
|
||||
self.assertEqual(self.secrets.get("pool-1"), secret)
|
||||
|
||||
async def test_media_profile_persists_and_is_public_without_changing_credentials(self):
|
||||
for profile in ("legacy-image-v1", "comfy-media-v1"):
|
||||
with self.subTest(profile=profile):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
"media_profile": profile, "api_key": "saved-key", "enabled": False,
|
||||
}))
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["media_profile"], profile)
|
||||
config_module.invalidate_config_cache()
|
||||
public = await config_routes.get_config_endpoint(_FakeRequest())
|
||||
self.assertEqual(public.status, 200)
|
||||
self.assertEqual(public.payload["workers"][0]["media_profile"], profile)
|
||||
self.assertTrue(public.payload["workers"][0]["api_key_configured"])
|
||||
self.assertNotIn("saved-key", json.dumps(public.payload))
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
|
||||
|
||||
async def test_old_clients_omit_media_profile_without_downgrading_saved_profile(self):
|
||||
for stored in ({}, {"media_profile": "comfy-media-v1"}):
|
||||
with self.subTest(stored=stored):
|
||||
worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", **stored}
|
||||
self.config["workers"] = [worker]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "name": "Old client save", "jobs_per_run": 2,
|
||||
}))
|
||||
self.assertEqual(response.status, 200)
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
self.assertEqual({k: v for k, v in saved.items() if k == "media_profile"}, stored)
|
||||
# The bulk route is also used by older clients with no profile field.
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "jobs_per_run": 3,
|
||||
}]}))
|
||||
self.assertEqual(response.status, 200)
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
self.assertEqual({k: v for k, v in saved.items() if k == "media_profile"}, stored)
|
||||
|
||||
async def test_invalid_media_profile_is_rejected_by_worker_and_bulk_routes(self):
|
||||
for profile in ("auto", "comfy-media-v2", None, False, 1, [], {}):
|
||||
with self.subTest(profile=profile):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
"media_profile": profile, "api_key": "never-save",
|
||||
}))
|
||||
self.assertEqual(response.status, 400)
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "media_profile": profile,
|
||||
}]}))
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
|
||||
async def test_active_pool_blocks_media_profile_changes_in_both_directions(self):
|
||||
for old_profile, new_profile in ((None, "comfy-media-v1"), ("comfy-media-v1", "legacy-image-v1")):
|
||||
with self.subTest(old_profile=old_profile):
|
||||
worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123"}
|
||||
if old_profile is not None:
|
||||
worker["media_profile"] = old_profile
|
||||
self.config["workers"] = [worker]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
original = self.read_config_file()
|
||||
with patch.object(config_routes, "_serverless_pool_is_active", return_value=True):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "media_profile": new_profile,
|
||||
}))
|
||||
self.assertEqual(response.status, 409)
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [dict(worker, media_profile=new_profile)],
|
||||
}))
|
||||
self.assertEqual(response.status, 409)
|
||||
self.assertEqual(self.read_config_file(), original)
|
||||
|
||||
async def test_enable_only_update_preserves_serverless_settings_and_saved_key(self):
|
||||
self.secrets.set("pool-1", "saved-secret")
|
||||
self.config["workers"] = [{
|
||||
@@ -302,13 +374,14 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
|
||||
"workers": [{
|
||||
"id": "pool-1", "type": "local", "port": 8188,
|
||||
"provider": "runpod", "endpoint_id": "legacy-ep", "api_key_configured": True,
|
||||
"media_profile": "comfy-media-v1",
|
||||
}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
for field in ("provider", "endpoint_id", "api_key_configured"):
|
||||
for field in ("provider", "endpoint_id", "api_key_configured", "media_profile"):
|
||||
self.assertNotIn(field, saved)
|
||||
|
||||
async def test_bulk_invalid_serverless_record_does_not_replace_configuration(self):
|
||||
@@ -341,6 +414,7 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "name": "Pool", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2, "timeout_seconds": 600,
|
||||
"media_profile": "comfy-media-v1",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
@@ -353,7 +427,7 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
worker = self.read_config_file()["workers"][0]
|
||||
self.assertEqual(worker["type"], "local")
|
||||
for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds"):
|
||||
for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "media_profile"):
|
||||
self.assertNotIn(serverless_field, worker)
|
||||
|
||||
async def test_enabled_toggle_allowed_during_active_run_but_endpoint_change_is_rejected(self):
|
||||
@@ -575,19 +649,41 @@ class PersistedPoolGuardTests(_ConfigRouteFixture):
|
||||
runtime_module.manager = self.manager
|
||||
else:
|
||||
self.manager.tasks["retained"] = types.SimpleNamespace(done=lambda: True)
|
||||
for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"type": "local", "port": 8188}):
|
||||
for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"type": "local", "port": 8188}, {"media_profile": "comfy-media-v1"}):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest(dict(fields, worker_id="pool-1")))
|
||||
self.assertEqual(response.status, 409, fields)
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
self.assertEqual(response.status, 409)
|
||||
worker = original["workers"][0]
|
||||
for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")]):
|
||||
for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")], [dict(worker, media_profile="comfy-media-v1")]):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": workers}))
|
||||
self.assertEqual(response.status, 409, workers)
|
||||
self.assertEqual(self.read_config_file(), original)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
|
||||
self.session_factory.assert_not_awaited()
|
||||
|
||||
async def test_retained_jobs_allow_same_or_omitted_media_profile(self):
|
||||
self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status="UNKNOWN")
|
||||
for profile in ("legacy-image-v1", "comfy-media-v1"):
|
||||
with self.subTest(profile=profile):
|
||||
worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123"}
|
||||
# An absent legacy field and an explicit legacy selection are equivalent.
|
||||
if profile != "legacy-image-v1":
|
||||
worker["media_profile"] = profile
|
||||
self.config["workers"] = [worker]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
for fields in ({"media_profile": profile}, {}):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "name": "Same profile", **fields,
|
||||
}))
|
||||
self.assertEqual(response.status, 200)
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", **fields,
|
||||
}]}))
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["media_profile"], profile)
|
||||
self.session_factory.assert_not_awaited()
|
||||
|
||||
async def test_retained_jobs_allow_images_per_run_and_legacy_concurrency_saves(self):
|
||||
# Upgrading from a config with concurrency=1 must not trip the busy guard,
|
||||
# and the image count stays editable while old jobs are unresolved.
|
||||
|
||||
@@ -0,0 +1,330 @@
|
||||
"""Bounded input packing; native acceptance is opt-in and CPU-only."""
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import importlib
|
||||
import io
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
import wave
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE = "_typed_inputs_tests"
|
||||
for suffix in ("", ".api", ".api.orchestration", ".serverless_bridge"):
|
||||
package = types.ModuleType(PACKAGE + suffix)
|
||||
package.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
|
||||
sys.modules[package.__name__] = package
|
||||
packer = importlib.import_module(PACKAGE + ".api.orchestration.serverless_typed_inputs")
|
||||
codec = importlib.import_module(PACKAGE + ".serverless_bridge.codec")
|
||||
|
||||
|
||||
def node(kind, **inputs):
|
||||
return {"class_type": kind, "inputs": inputs}
|
||||
|
||||
|
||||
def graph(kind="LoadAudio", filename="tone.wav"):
|
||||
key = "audio" if kind == "LoadAudio" else "file"
|
||||
socket = "audio" if kind == "LoadAudio" else "video"
|
||||
return {"source": node(kind, **{key: filename}),
|
||||
"consumer": node("Consumer", **{socket: ["source", 0]})}
|
||||
|
||||
|
||||
def run(workflow):
|
||||
return asyncio.run(packer.prepare_media_inputs(workflow))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def paths(tmp_path, monkeypatch):
|
||||
roots = {name: tmp_path / name for name in ("input", "output", "temp")}
|
||||
for path in roots.values():
|
||||
path.mkdir()
|
||||
def resolve(value):
|
||||
for name in roots:
|
||||
marker = " [" + name + "]"
|
||||
if value.endswith(marker):
|
||||
return str(roots[name] / value[:-len(marker)])
|
||||
return str(roots["input"] / value)
|
||||
folders = types.SimpleNamespace(get_annotated_filepath=resolve)
|
||||
for name, path in roots.items():
|
||||
setattr(folders, "get_" + name + "_directory", lambda path=path: str(path))
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", folders)
|
||||
return roots
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind,key", [("LoadAudio", "audio"), ("LoadVideo", "file")])
|
||||
@pytest.mark.parametrize("bad", [{}, None, [], {"filename": "a.wav"},
|
||||
{"audio": ["filename", 0]}, {"file": ["filename", 0]}])
|
||||
def test_rejects_non_native_or_linked_inputs_before_read(kind, key, bad, monkeypatch):
|
||||
workflow = graph(kind)
|
||||
workflow["source"]["inputs"] = bad
|
||||
read = Mock(side_effect=AssertionError("must not read"))
|
||||
monkeypatch.setattr(packer.video_codec, "_read_path", read)
|
||||
with pytest.raises(ValueError):
|
||||
run(workflow)
|
||||
read.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind,key", [("LoadAudio", "audio"), ("LoadVideo", "file")])
|
||||
@pytest.mark.parametrize("extra", ["start_time", "duration", "sample_rate", "crop", "upload"])
|
||||
def test_rejects_extra_loader_controls(kind, key, extra):
|
||||
workflow = graph(kind)
|
||||
workflow["source"]["inputs"][extra] = 0
|
||||
with pytest.raises(ValueError, match="controls|inputs"):
|
||||
run(workflow)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("slot", [1, -1, True, 0.0, "0", None])
|
||||
def test_rejects_nonzero_or_noninteger_output_slot(slot):
|
||||
workflow = graph()
|
||||
workflow["consumer"]["inputs"]["audio"] = ["source", slot]
|
||||
with pytest.raises(ValueError, match="output|slot"):
|
||||
run(workflow)
|
||||
|
||||
|
||||
def test_rejects_unused_loader_and_malformed_link():
|
||||
workflow = graph()
|
||||
workflow.pop("consumer")
|
||||
with pytest.raises(ValueError, match="consumer"):
|
||||
run(workflow)
|
||||
workflow = graph()
|
||||
workflow["consumer"]["inputs"]["audio"].append(0)
|
||||
with pytest.raises(ValueError):
|
||||
run(workflow)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key,value", [("audio", "private.wav"), ("file", "private.mp4"), ("video", "private.mov")])
|
||||
def test_unknown_media_loaders_still_use_strict_scan(key, value):
|
||||
with pytest.raises(ValueError):
|
||||
run({"source": node("CustomLoader", **{key: value})})
|
||||
|
||||
|
||||
def test_images_still_rewrite_only_the_private_graph(paths):
|
||||
payload = b"legacy staging preserves the original file bytes"
|
||||
(paths["input"] / "picture.png").write_bytes(payload)
|
||||
source = {"image": node("LoadImage", image="picture.png"),
|
||||
"export": node("MediaExportV1", images=["image", 0])}
|
||||
workflow = copy.deepcopy(source)
|
||||
images, artifacts, refs = run(workflow)
|
||||
assert artifacts == refs == []
|
||||
assert len(images) == 1
|
||||
assert base64.b64decode(images[0]["image"]) == payload
|
||||
assert workflow["image"]["inputs"]["image"] == images[0]["name"]
|
||||
assert workflow["export"] == source["export"]
|
||||
assert source["image"]["inputs"]["image"] == "picture.png"
|
||||
|
||||
|
||||
if os.environ.get("_DISTRIBUTED_TYPED_INPUT_CORE_TEST") != "1":
|
||||
def test_native_input_acceptance_cpu_subprocess():
|
||||
configured = os.environ.get("COMFYUI_SOURCE_ROOT")
|
||||
if configured is None:
|
||||
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU native input acceptance")
|
||||
core = Path(configured).expanduser().resolve()
|
||||
assert configured and (core / "comfy_extras/nodes_audio.py").is_file()
|
||||
env = dict(os.environ)
|
||||
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
env.pop("PYTHONPATH", None)
|
||||
env.pop("COMFYUI_SOURCE_ROOT", None)
|
||||
env.update(TMPDIR=str(scratch), CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="",
|
||||
PYTHONDONTWRITEBYTECODE="1", _DISTRIBUTED_TYPED_INPUT_CORE_TEST="1")
|
||||
bootstrap = """
|
||||
import sys
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
import torch
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('Input packing must not query or initialize a GPU')
|
||||
torch.cuda.is_available = lambda: False
|
||||
torch.cuda.device_count = lambda: 0
|
||||
for name in ('init', '_lazy_init', 'current_device'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
import comfy_extras.nodes_audio
|
||||
import comfy_api.latest._input_impl.video_types
|
||||
for name in ('is_available', 'device_count'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
import pytest
|
||||
raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider']))
|
||||
"""
|
||||
result = subprocess.run([sys.executable, "-c", bootstrap, str(core), str(Path(__file__).resolve())],
|
||||
cwd=ROOT, env=env, text=True, capture_output=True, timeout=150)
|
||||
print(result.stdout)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
else:
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
from comfy_extras import nodes_audio
|
||||
from comfy_api.latest._input_impl.video_types import VideoFromFile
|
||||
|
||||
def wav_bytes(channels=2, sample_width=2, rate=22050, variant=0):
|
||||
values = np.array([-32768, -12001 + variant, -1, 0, 1, 8193, 32767] * channels,
|
||||
dtype="<i2" if sample_width == 2 else "<i4")
|
||||
buffer = io.BytesIO()
|
||||
with wave.open(buffer, "wb") as out:
|
||||
out.setnchannels(channels)
|
||||
out.setsampwidth(sample_width)
|
||||
out.setframerate(rate)
|
||||
out.writeframes(values.tobytes())
|
||||
return buffer.getvalue()
|
||||
|
||||
def unpack(artifact):
|
||||
payload = base64.b64decode(artifact["transport"]["data_base64"])
|
||||
assert artifact["byte_length"] == len(payload)
|
||||
return payload
|
||||
|
||||
@pytest.mark.parametrize("root", ["input", "output", "temp"])
|
||||
@pytest.mark.parametrize("channels,width", [(1, 2), (2, 2), (2, 4)])
|
||||
def test_exact_native_audio_and_every_consumer_ref(paths, root, channels, width):
|
||||
source = paths[root] / "tone.wav"
|
||||
source.write_bytes(wav_bytes(channels, width))
|
||||
workflow = graph(filename="tone.wav [" + root + "]")
|
||||
workflow["export"] = node("MediaExportV1", audio=["source", 0])
|
||||
workflow["consumer"]["inputs"]["second_audio"] = ["source", 0]
|
||||
before = copy.deepcopy(workflow)
|
||||
images, artifacts, refs = run(workflow)
|
||||
assert images == [] and len(artifacts) == 1
|
||||
artifact = artifacts[0]
|
||||
expected, rate = nodes_audio.load(str(source))
|
||||
actual = codec.decode_payload("AUDIO", artifact["descriptor"], unpack(artifact))
|
||||
assert actual["sample_rate"] == rate == 22050
|
||||
assert actual["waveform"].dtype == expected.dtype
|
||||
assert torch.equal(actual["waveform"].view(torch.uint8), expected.unsqueeze(0).contiguous().view(torch.uint8))
|
||||
assert refs == [{"node_id": "consumer", "input_name": "audio", "artifact_id": artifact["id"]},
|
||||
{"node_id": "consumer", "input_name": "second_audio", "artifact_id": artifact["id"]},
|
||||
{"node_id": "export", "input_name": "audio", "artifact_id": artifact["id"]}]
|
||||
assert workflow == before
|
||||
|
||||
@pytest.mark.parametrize("root", ["input", "output", "temp"])
|
||||
def test_exact_video_bytes_without_container_decode(paths, root, monkeypatch):
|
||||
payload = b"opaque native video bytes\x00\xff"
|
||||
(paths[root] / "clip.mp4").write_bytes(payload)
|
||||
workflow = graph("LoadVideo", "clip.mp4 [" + root + "]")
|
||||
workflow["export"] = node("MediaExportV1", video=["source", 0])
|
||||
before = copy.deepcopy(workflow)
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError("must not open, decode or transcode video")
|
||||
monkeypatch.setattr(av, "open", forbidden)
|
||||
monkeypatch.setattr(VideoFromFile, "get_components", forbidden)
|
||||
images, artifacts, refs = run(workflow)
|
||||
assert images == [] and len(artifacts) == 1
|
||||
artifact = artifacts[0]
|
||||
assert unpack(artifact) == payload
|
||||
assert artifact["descriptor"] == {"encoding": "native-video-v1", "representation": "file",
|
||||
"byte_length": len(payload), "start_time": 0, "duration": 0, "crop": None,
|
||||
"data": {"offset": 0, "length": len(payload)}}
|
||||
assert refs == [{"node_id": name, "input_name": "video", "artifact_id": artifact["id"]}
|
||||
for name in ("consumer", "export")]
|
||||
native = packer.video_codec.decode_video(artifact["descriptor"], unpack(artifact))
|
||||
assert vars(native)["_VideoFromFile__file"].getvalue() == payload
|
||||
assert workflow == before
|
||||
|
||||
def test_content_dedupe_kind_separation_and_mixed_images(paths):
|
||||
data = wav_bytes()
|
||||
(paths["input"] / "a.wav").write_bytes(data)
|
||||
(paths["output"] / "b.wav").write_bytes(data)
|
||||
(paths["input"] / "image.png").write_bytes(b"image bytes")
|
||||
workflow = graph(filename="a.wav")
|
||||
workflow["same"] = node("LoadAudio", audio="b.wav [output]")
|
||||
workflow["video"] = node("LoadVideo", file="a.wav")
|
||||
workflow["image"] = node("LoadImage", image="image.png")
|
||||
workflow["export"] = node("MediaExportV1", audio=["same", 0], video=["video", 0], images=["image", 0])
|
||||
before = copy.deepcopy(workflow)
|
||||
images, artifacts, refs = run(workflow)
|
||||
assert len(images) == 1 and len(artifacts) == 2 and len(refs) == 3
|
||||
assert refs[0]["artifact_id"] == refs[1]["artifact_id"] != refs[2]["artifact_id"]
|
||||
assert {a["kind"] for a in artifacts} == {"AUDIO", "VIDEO"}
|
||||
before["image"]["inputs"]["image"] = images[0]["name"]
|
||||
assert workflow == before
|
||||
second = graph(filename="b.wav [output]")
|
||||
assert run(second)[1][0]["id"] == artifacts[0]["id"]
|
||||
|
||||
@pytest.mark.parametrize("kind", ["LoadAudio", "LoadVideo"])
|
||||
def test_path_rejections_and_per_file_bounds(paths, tmp_path, monkeypatch, kind):
|
||||
outside = tmp_path / "outside.wav"
|
||||
outside.write_bytes(wav_bytes())
|
||||
(paths["input"] / "escape.wav").symlink_to(outside)
|
||||
(paths["input"] / "dirlink").symlink_to(tmp_path, target_is_directory=True)
|
||||
os.mkfifo(paths["input"] / "fifo.wav")
|
||||
values = ["../outside.wav", str(outside), "escape.wav", "dirlink/outside.wav", "fifo.wav",
|
||||
"absent.wav", "tone.wav [unknown]", "https://example.com/a.wav", "input [temp]", "C:\\a.wav"]
|
||||
for value in values:
|
||||
with pytest.raises(ValueError):
|
||||
run(graph(kind, value))
|
||||
(paths["input"] / "large.wav").write_bytes(wav_bytes())
|
||||
monkeypatch.setattr(packer.contract, "MAX_ARTIFACT_BYTES", 8)
|
||||
with pytest.raises(ValueError, match="limit|size"):
|
||||
run(graph(kind, "large.wav"))
|
||||
|
||||
def test_aggregate_decoded_budget_and_duplicate_not_double_charged(paths, monkeypatch):
|
||||
payload = wav_bytes()
|
||||
(paths["input"] / "a.wav").write_bytes(payload)
|
||||
(paths["input"] / "b.wav").write_bytes(wav_bytes(variant=1))
|
||||
workflow = graph(filename="a.wav")
|
||||
workflow["other"] = node("LoadAudio", audio="a.wav")
|
||||
workflow["export"] = node("MediaExportV1", audio=["other", 0])
|
||||
# File source bytes have a separate aggregate bound; use a longer WAV
|
||||
# whose float output dominates its tiny header.
|
||||
with wave.open(io.BytesIO(payload), "rb") as wav:
|
||||
frames = wav.readframes(wav.getnframes())
|
||||
buffer = io.BytesIO()
|
||||
with wave.open(buffer, "wb") as wav:
|
||||
wav.setparams((2, 2, 22050, 0, "NONE", "not compressed"))
|
||||
wav.writeframes(frames * 100)
|
||||
(paths["input"] / "a.wav").write_bytes(buffer.getvalue())
|
||||
artifact = run(graph(filename="a.wav"))[1][0]
|
||||
monkeypatch.setattr(packer.contract, "MAX_TOTAL_BYTES", artifact["byte_length"])
|
||||
assert len(run(workflow)[1]) == 1
|
||||
workflow["other"]["inputs"]["audio"] = "b.wav"
|
||||
with pytest.raises(ValueError, match="limit|size|budget"):
|
||||
run(workflow)
|
||||
|
||||
def test_audio_charged_before_ndarray_cat_or_conversion(monkeypatch):
|
||||
frame = types.SimpleNamespace(samples=100, sample_rate=22050,
|
||||
layout=types.SimpleNamespace(channels=[1, 2]),
|
||||
format=types.SimpleNamespace(name="s16", bytes=2),
|
||||
planes=[types.SimpleNamespace(buffer_size=400)], to_ndarray=Mock())
|
||||
stream = types.SimpleNamespace(index=0, channels=2,
|
||||
codec_context=types.SimpleNamespace(sample_rate=22050))
|
||||
container = Mock()
|
||||
container.__enter__ = Mock(return_value=container)
|
||||
container.__exit__ = Mock(return_value=False)
|
||||
container.streams = types.SimpleNamespace(audio=[stream])
|
||||
container.decode = Mock(return_value=iter([frame]))
|
||||
monkeypatch.setattr(av, "open", Mock(return_value=container))
|
||||
cat = Mock(side_effect=AssertionError("no accumulation allowed"))
|
||||
monkeypatch.setattr(torch, "cat", cat)
|
||||
with pytest.raises(ValueError, match="limit|size|budget"):
|
||||
packer._decode_audio(b"input", 799)
|
||||
frame.to_ndarray.assert_not_called()
|
||||
cat.assert_not_called()
|
||||
frame.samples = 1
|
||||
frame.planes[0].buffer_size = 800
|
||||
container.decode.return_value = iter([frame])
|
||||
with pytest.raises(ValueError, match="limit|size|budget"):
|
||||
packer._decode_audio(b"input", 799)
|
||||
frame.to_ndarray.assert_not_called()
|
||||
|
||||
def test_audio_rejects_empty_invalid_and_channel_or_rate_drift(paths, monkeypatch):
|
||||
(paths["input"] / "bad.wav").write_bytes(b"not audio")
|
||||
with pytest.raises(ValueError):
|
||||
run(graph(filename="bad.wav"))
|
||||
stream = types.SimpleNamespace(index=0, channels=2, codec_context=types.SimpleNamespace(sample_rate=22050))
|
||||
for frames in ([], [types.SimpleNamespace(samples=1, sample_rate=44100,
|
||||
layout=types.SimpleNamespace(channels=[1, 2]))],
|
||||
[types.SimpleNamespace(samples=1, sample_rate=22050,
|
||||
layout=types.SimpleNamespace(channels=[1]))]):
|
||||
container = Mock()
|
||||
container.__enter__ = Mock(return_value=container)
|
||||
container.__exit__ = Mock(return_value=False)
|
||||
container.streams = types.SimpleNamespace(audio=[stream])
|
||||
container.decode = Mock(return_value=iter(frames))
|
||||
monkeypatch.setattr(av, "open", Mock(return_value=container))
|
||||
with pytest.raises(ValueError):
|
||||
packer._decode_audio(b"input", 1000)
|
||||
@@ -0,0 +1,319 @@
|
||||
"""Offline typed submission/delivery regressions; no provider or GPU access."""
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
import socket
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
# Load once before integration fixtures snapshot/restore sys.modules. PyTorch
|
||||
# cannot be re-imported after those fixtures unload its extension-backed modules.
|
||||
import torch
|
||||
|
||||
from test_runpod_client import Response, Session
|
||||
from test_serverless_runtime import runtime, Secrets
|
||||
|
||||
client = importlib.import_module(runtime.__package__ + ".runpod_client")
|
||||
bridge = runtime.__package__.split(".")[0] + ".serverless_bridge"
|
||||
contract = importlib.import_module(bridge + ".contract")
|
||||
|
||||
|
||||
def prepared(sockets=("texts",)):
|
||||
worker = {"id": "pool:1", "pool_id": "pool", "endpoint_id": "ep1",
|
||||
"timeout_seconds": 30, "type": "serverless", "media_profile": contract.PROFILE}
|
||||
bindings = [dict(binding_id=f"binding-{i}", collector_id="collector", job_id="collector-job",
|
||||
participant_id=worker["id"], node_id=f"export-{i}", socket=name,
|
||||
kind=contract.SOCKET_KINDS[name]) for i, name in enumerate(sockets)]
|
||||
media = dict(version=1, profile=contract.PROFILE, request_id="request-1",
|
||||
artifacts=[], inputs=[], outputs=copy.deepcopy(bindings))
|
||||
graph = {"source": {"class_type": "Source", "inputs": {"value": "café雪"}}}
|
||||
for binding in bindings:
|
||||
graph[binding["node_id"]] = {"class_type": "MediaExportV1", "inputs": {
|
||||
"request_id": media["request_id"], "binding_json": json.dumps(binding),
|
||||
binding["socket"]: ["source", 0]}}
|
||||
return dict(worker=worker, workflow=graph, images=[], manifest=bindings, media=media)
|
||||
|
||||
|
||||
def response(job):
|
||||
envelope = dict(version=1, profile=contract.PROFILE, request_id=job["media"]["request_id"],
|
||||
bindings=[], artifacts=[], items=[])
|
||||
for index, binding in enumerate(job["manifest"]):
|
||||
envelope["bindings"].append(dict(binding, count=2))
|
||||
for item_index, text in enumerate(("", " \n雪")):
|
||||
aid = f"artifact-{index}-{item_index}"
|
||||
descriptor, payload = {"encoding": "utf8-v1"}, text.encode()
|
||||
if binding["kind"] != "STRING":
|
||||
codec = importlib.import_module(bridge + ".codec")
|
||||
values = {"IMAGE": torch.arange(8, dtype=torch.float64).reshape(1, 1, 2, 4),
|
||||
"MASK": torch.arange(2, dtype=torch.float16).reshape(1, 1, 2),
|
||||
"AUDIO": {"waveform": torch.arange(4, dtype=torch.float32).reshape(1, 2, 2), "sample_rate": 48000}}
|
||||
descriptor, payload = codec.encode_payload(binding["kind"], values[binding["kind"]])
|
||||
artifact = contract.artifact_from_payload(aid, binding["kind"], descriptor, payload)
|
||||
envelope["artifacts"].append(artifact)
|
||||
envelope["items"].insert(0, dict(binding_id=binding["binding_id"], item_index=item_index, artifact_id=aid))
|
||||
return {"media": envelope}
|
||||
|
||||
|
||||
class TypedClientTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_exact_typed_bytes_and_legacy_unchanged(self):
|
||||
job = prepared()
|
||||
session = Session([Response(payload={"id": "job-1"})])
|
||||
api = client.RunpodClient(session, "ep1", "test-only")
|
||||
await api.submit(job["workflow"], [], 30, media=job["media"])
|
||||
expected = {"input": {"workflow": job["workflow"], "images": [], "media": job["media"]},
|
||||
"policy": {"executionTimeout": 30000, "ttl": 30000}}
|
||||
self.assertEqual(session.calls[0][2]["data"], json.dumps(expected, allow_nan=False,
|
||||
ensure_ascii=False, separators=(",", ":")).encode())
|
||||
self.assertEqual(len(session.calls), 1)
|
||||
expected["input"].pop("media")
|
||||
legacy = json.dumps(expected, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode()
|
||||
self.assertEqual(client.serialize_request(job["workflow"], [], 30), legacy)
|
||||
self.assertEqual(client.serialize_request(job["workflow"], [], 30, media=None), legacy)
|
||||
|
||||
async def test_invalid_envelopes_never_post(self):
|
||||
for field, value in [("version", True), ("profile", "legacy-image-v1"),
|
||||
("request_id", "../bad"), ("outputs", []), ("inputs", [{}])]:
|
||||
with self.subTest(field=field):
|
||||
job = prepared()
|
||||
job["media"][field] = value
|
||||
session = Session([])
|
||||
with self.assertRaises(ValueError):
|
||||
await client.RunpodClient(session, "ep1", "test-only").submit(
|
||||
job["workflow"], [], 30, media=job["media"])
|
||||
self.assertFalse(session.calls)
|
||||
session = Session([])
|
||||
with self.assertRaises(ValueError):
|
||||
await client.RunpodClient(session, "ep1", "test-only").submit({}, [], 30, media={})
|
||||
self.assertFalse(session.calls)
|
||||
|
||||
async def test_typed_size_boundary_uses_exact_body_before_post(self):
|
||||
job = prepared()
|
||||
size = len(client.serialize_request(job["workflow"], [], 30, media=job["media"]))
|
||||
session = Session([Response(payload={"id": "boundary"})])
|
||||
api = client.RunpodClient(session, "ep1", "test-only")
|
||||
with patch.object(client, "MAX_REQUEST_BYTES", size):
|
||||
await api.submit(job["workflow"], [], 30, media=job["media"])
|
||||
job["workflow"]["source"]["inputs"]["value"] += "x"
|
||||
with self.assertRaisesRegex(ValueError, "Serialized serverless request body"):
|
||||
await api.submit(job["workflow"], [], 30, media=job["media"])
|
||||
self.assertEqual(len(session.calls), 1)
|
||||
|
||||
|
||||
class TypedRuntimeTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector-job": asyncio.Queue()},
|
||||
distributed_jobs_lock=asyncio.Lock())
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), Path(self.temp.name) / "jobs.json",
|
||||
AsyncMock(return_value=object()))
|
||||
self.manager.poll_seconds = 0
|
||||
self.job = prepared()
|
||||
self.api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"),
|
||||
status=AsyncMock(return_value={"status": "COMPLETED", "output": response(self.job)}),
|
||||
cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
self.client_patch = patch.object(runtime, "RunpodClient", return_value=self.api)
|
||||
self.client_patch.start()
|
||||
|
||||
async def asyncTearDown(self):
|
||||
await self.manager.cancel()
|
||||
self.client_patch.stop()
|
||||
self.temp.cleanup()
|
||||
|
||||
async def run_job(self):
|
||||
await self.manager.start("execution", "prompt", [self.job])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
|
||||
async def test_preflight_all_jobs_before_any_journal_or_task(self):
|
||||
for mutation in ("profile", "missing_media", "manifest", "participant", "exporter", "empty_media"):
|
||||
with self.subTest(mutation=mutation):
|
||||
bad = prepared()
|
||||
if mutation == "profile":
|
||||
bad["worker"]["media_profile"] = "legacy-image-v1"
|
||||
elif mutation == "missing_media":
|
||||
bad.pop("media")
|
||||
elif mutation == "manifest":
|
||||
bad["manifest"][0]["job_id"] = "other"
|
||||
elif mutation == "participant":
|
||||
bad["worker"]["id"] = "other-participant"
|
||||
elif mutation == "exporter":
|
||||
bad["workflow"]["export-0"]["inputs"]["request_id"] = "other"
|
||||
else:
|
||||
bad["media"] = {}
|
||||
with self.assertRaises(ValueError):
|
||||
await self.manager.start("execution", "prompt", [self.job, bad])
|
||||
self.assertFalse(self.manager.ledger.all())
|
||||
self.assertFalse(self.manager.tasks)
|
||||
self.assertFalse(self.manager.runs)
|
||||
self.api.submit.assert_not_awaited()
|
||||
|
||||
async def test_text_roundtrip_id_journal_and_payload_cleanup(self):
|
||||
async def status(_job):
|
||||
self.assertEqual(self.manager.ledger.all()[0]["job_id"], "provider-job")
|
||||
return {"status": "COMPLETED", "output": response(self.job)}
|
||||
self.api.status.side_effect = status
|
||||
await self.run_job()
|
||||
self.api.submit.assert_awaited_once_with(self.job["workflow"], [], 30, media=self.job["media"])
|
||||
queue = self.server.distributed_pending_jobs["collector-job"]
|
||||
self.assertEqual(queue.get_nowait(), {"worker_id": "pool:1", "media": {"texts": ["", " \n雪"]}, "is_last": True})
|
||||
self.assertTrue(queue.empty())
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
|
||||
run = next(iter(self.manager.runs.values()))
|
||||
for field in ("workflow", "images", "media"):
|
||||
self.assertNotIn(field, run)
|
||||
self.assertNotIn("data_base64", self.manager.ledger.path.read_text())
|
||||
self.api.cancel.assert_not_awaited()
|
||||
|
||||
async def test_all_sockets_grouped_without_native_value_conversion(self):
|
||||
self.job = prepared(tuple(contract.SOCKET_KINDS))
|
||||
native = {binding["binding_id"]: [object(), {"native": object()}] for binding in self.job["manifest"]}
|
||||
with patch.object(contract, "decode_response", AsyncMock(return_value=native)) as decode:
|
||||
await self.run_job()
|
||||
decode.assert_awaited_once()
|
||||
self.assertEqual(decode.call_args.kwargs["request_id"], "request-1")
|
||||
self.assertEqual(decode.call_args.kwargs["bindings"], self.job["manifest"])
|
||||
queue = self.server.distributed_pending_jobs["collector-job"]
|
||||
item = queue.get_nowait()
|
||||
self.assertEqual(set(item), {"worker_id", "media", "is_last"})
|
||||
self.assertEqual(set(item["media"]), set(contract.SOCKET_KINDS))
|
||||
for binding in self.job["manifest"]:
|
||||
for actual, expected in zip(item["media"][binding["socket"]], native[binding["binding_id"]]):
|
||||
self.assertIs(actual, expected)
|
||||
self.assertTrue(queue.empty())
|
||||
|
||||
async def test_identity_and_late_decode_failure_deliver_only_one_error(self):
|
||||
self.job = prepared(("texts", "images"))
|
||||
for mutation in ("request", "binding", "checksum", "descriptor"):
|
||||
with self.subTest(mutation=mutation):
|
||||
output = response(self.job)
|
||||
if mutation == "request":
|
||||
output["media"]["request_id"] = "other"
|
||||
elif mutation == "binding":
|
||||
output["media"]["bindings"][-1]["participant_id"] = "other"
|
||||
elif mutation == "checksum":
|
||||
output["media"]["artifacts"][-1]["sha256"] = "0" * 64
|
||||
else:
|
||||
output["media"]["artifacts"][-1]["descriptor"] = {}
|
||||
self.api.status.return_value = {"status": "COMPLETED", "output": output}
|
||||
await self.run_job()
|
||||
queue = self.server.distributed_pending_jobs["collector-job"]
|
||||
item = queue.get_nowait()
|
||||
self.assertIn("serverless_error", item)
|
||||
self.assertNotIn("media", item)
|
||||
self.assertTrue(queue.empty())
|
||||
self.api.cancel.assert_not_awaited()
|
||||
|
||||
async def test_closed_queue_is_not_resurrected(self):
|
||||
self.server.distributed_pending_jobs.clear()
|
||||
await self.run_job()
|
||||
self.assertFalse(self.server.distributed_pending_jobs)
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
|
||||
|
||||
async def test_real_tensor_audio_and_mask_codecs_preserve_dtype_and_shape(self):
|
||||
self.job = prepared(("texts", "images", "audio", "masks"))
|
||||
self.api.status.return_value = {"status": "COMPLETED", "output": response(self.job)}
|
||||
await self.run_job()
|
||||
result = self.server.distributed_pending_jobs["collector-job"].get_nowait()["media"]
|
||||
self.assertEqual(result["texts"], ["", " \n雪"])
|
||||
self.assertEqual(result["images"][0].dtype, torch.float64)
|
||||
self.assertEqual(result["images"][0].shape, (1, 1, 2, 4))
|
||||
self.assertEqual(result["masks"][0].dtype, torch.float16)
|
||||
self.assertEqual(result["audio"][0]["sample_rate"], 48000)
|
||||
self.assertTrue(torch.equal(result["audio"][0]["waveform"], torch.arange(4, dtype=torch.float32).reshape(1, 2, 2)))
|
||||
self.assertTrue(all(len(values) == 2 for values in result.values()))
|
||||
|
||||
async def test_run_rechecks_profile_before_submission(self):
|
||||
await self.manager.start("execution", "prompt", [self.job])
|
||||
self.job["worker"]["media_profile"] = "legacy-image-v1"
|
||||
await self.manager.wait_for_execution("execution")
|
||||
self.api.submit.assert_not_awaited()
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "FAILED")
|
||||
|
||||
async def test_multiple_collectors_receive_one_message_each(self):
|
||||
self.job = prepared(("texts", "images"))
|
||||
self.job["manifest"][-1]["job_id"] = "second-job"
|
||||
self.job["media"]["outputs"] = copy.deepcopy(self.job["manifest"])
|
||||
self.job["workflow"]["export-1"]["inputs"]["binding_json"] = json.dumps(self.job["manifest"][-1])
|
||||
self.server.distributed_pending_jobs["second-job"] = asyncio.Queue()
|
||||
self.api.status.return_value = {"status": "COMPLETED", "output": response(self.job)}
|
||||
await self.run_job()
|
||||
for job_id, name in (("collector-job", "texts"), ("second-job", "images")):
|
||||
queue = self.server.distributed_pending_jobs[job_id]
|
||||
self.assertEqual(queue.qsize(), 1)
|
||||
item = queue.get_nowait()
|
||||
self.assertEqual(set(item["media"]), {name})
|
||||
self.assertTrue(item["is_last"])
|
||||
|
||||
async def test_prestart_cancel_releases_media_and_deduplicates_errors(self):
|
||||
self.job = prepared(("texts", "images"))
|
||||
await self.manager.start("execution", "prompt", [self.job])
|
||||
await self.manager.cancel(execution_id="execution")
|
||||
self.api.submit.assert_not_awaited()
|
||||
run = next(iter(self.manager.runs.values()))
|
||||
self.assertNotIn("media", run)
|
||||
self.assertEqual(self.server.distributed_pending_jobs["collector-job"].qsize(), 1)
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
|
||||
|
||||
async def test_ambiguous_submit_keeps_reservation_without_media(self):
|
||||
self.api.submit.side_effect = runtime.AmbiguousSubmission("Unknown submission")
|
||||
await self.run_job()
|
||||
self.api.submit.assert_awaited_once()
|
||||
self.api.cancel.assert_not_awaited()
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
self.assertNotIn("media", next(iter(self.manager.runs.values())))
|
||||
|
||||
async def test_legacy_still_uses_tensor_messages_and_no_media_kwarg(self):
|
||||
self.job.pop("media")
|
||||
self.job["worker"]["media_profile"] = "legacy-image-v1"
|
||||
self.job["manifest"] = [{"job_id": "collector-job", "prefix": "dist_x", "kind": "image"}]
|
||||
with patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["a", "b"]})):
|
||||
await self.run_job()
|
||||
self.api.submit.assert_awaited_once_with(self.job["workflow"], [], 30)
|
||||
queue = self.server.distributed_pending_jobs["collector-job"]
|
||||
self.assertEqual([queue.get_nowait(), queue.get_nowait()], [
|
||||
{"worker_id": "pool:1", "tensor": "a", "image_index": 0, "is_last": False},
|
||||
{"worker_id": "pool:1", "tensor": "b", "image_index": 1, "is_last": True}])
|
||||
|
||||
|
||||
class TypedDownloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_stream_bound_cap_and_redirect_policy_without_credentials(self):
|
||||
typed = importlib.import_module(runtime.__package__ + ".serverless_typed_results")
|
||||
class DownloadResponse(Response):
|
||||
headers = {}
|
||||
async def chunks(self, _size):
|
||||
yield b"ab"
|
||||
yield b"cd"
|
||||
@property
|
||||
def content(self):
|
||||
return types.SimpleNamespace(iter_chunked=self.chunks)
|
||||
class DownloadSession(Response):
|
||||
def get(self, url, **kwargs):
|
||||
self.calls.append((url, kwargs))
|
||||
return self.response
|
||||
session = DownloadSession()
|
||||
session.calls = []
|
||||
session.response = DownloadResponse()
|
||||
with patch.object(typed.aiohttp, "TCPConnector", return_value=object()), \
|
||||
patch.object(typed.aiohttp, "ClientSession", return_value=session) as factory:
|
||||
self.assertEqual(await typed.download_public_artifact("https://example.com/a", max_bytes=contract.MAX_ARTIFACT_BYTES), b"abcd")
|
||||
self.assertFalse(factory.call_args.kwargs["trust_env"])
|
||||
self.assertNotIn("headers", factory.call_args.kwargs)
|
||||
self.assertFalse(session.calls[0][1]["allow_redirects"])
|
||||
with self.assertRaises(ValueError):
|
||||
await typed.download_public_artifact("https://example.com/a", max_bytes=3)
|
||||
session.response.status = 302
|
||||
session.response.headers = {"Location": "https://127.0.0.1/private"}
|
||||
with self.assertRaises(ValueError):
|
||||
await typed.download_public_artifact("https://example.com/a", max_bytes=4)
|
||||
for url in ("http://example.com/a", "https://127.0.0.1/a", "https://user:pass@example.com/a"):
|
||||
with self.assertRaises(ValueError):
|
||||
await typed.download_public_artifact(url, max_bytes=4)
|
||||
resolver = typed._PublicResolver()
|
||||
with patch.object(resolver, "_getaddrinfo", AsyncMock(return_value=[
|
||||
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 443))])):
|
||||
with self.assertRaises(OSError):
|
||||
await resolver.resolve("example.com", 443)
|
||||
@@ -133,6 +133,7 @@ def test_quality_to_crf_mapping():
|
||||
|
||||
def test_normalize_format_rejects_unknown_values():
|
||||
assert tc.normalize_format(" H264 ") == "h264"
|
||||
assert tc.normalize_format(" EXACT ") == "exact"
|
||||
assert tc.normalize_format(None) == tc.DEFAULT_TRANSFER_FORMAT
|
||||
with pytest.raises(tc.TransferCodecError):
|
||||
tc.normalize_format("gif")
|
||||
@@ -158,6 +159,22 @@ def test_decode_chunk_validates_envelope_against_payload():
|
||||
tc.decode_chunk({**video_env, "count": 5}, video["parts"])
|
||||
|
||||
|
||||
def test_exact_is_appended_without_changing_saved_format_defaults():
|
||||
assert tc.TRANSFER_FORMATS == (
|
||||
"legacy_png", "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1", "exact",
|
||||
)
|
||||
assert tc.DEFAULT_TRANSFER_FORMAT == "legacy_png"
|
||||
assert tc.EXACT_FORMAT not in tc.IMAGE_FORMATS + tc.VIDEO_FORMATS
|
||||
|
||||
|
||||
def test_exact_does_not_enter_uint8_chunk_codec():
|
||||
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=1))
|
||||
with pytest.raises(tc.TransferCodecError, match="exact"):
|
||||
list(tc.encode_chunks(frames, "exact", 90, 10**9))
|
||||
with pytest.raises(tc.TransferCodecError, match="exact"):
|
||||
tc.decode_chunk({"format": "exact", "count": 1, "width": 91, "height": 67}, [b"raw tensor"])
|
||||
|
||||
|
||||
def test_legacy_format_is_not_chunk_encoded():
|
||||
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=1))
|
||||
with pytest.raises(tc.TransferCodecError):
|
||||
|
||||
@@ -98,6 +98,7 @@ def test_registered_collector_schema_exposes_optional_transfer_inputs(v3):
|
||||
fmt = inputs["transfer_format"]
|
||||
assert fmt.kind == "Combo"
|
||||
assert fmt.kwargs["options"] == list(codec.TRANSFER_FORMATS)
|
||||
assert fmt.kwargs["options"][-1] == "exact"
|
||||
assert fmt.kwargs["default"] == "legacy_png"
|
||||
assert fmt.kwargs["optional"] is True
|
||||
quality = inputs["transfer_quality"]
|
||||
@@ -107,13 +108,14 @@ def test_registered_collector_schema_exposes_optional_transfer_inputs(v3):
|
||||
assert schema.is_input_list is True
|
||||
|
||||
|
||||
def test_registered_collector_forwards_transfer_options(v3):
|
||||
@pytest.mark.parametrize("fmt", ["av1", "exact"])
|
||||
def test_registered_collector_forwards_transfer_options(v3, fmt):
|
||||
mod, _, calls = v3
|
||||
|
||||
mod.DistributedCollector.execute(images=["img"], audio=[None], transfer_format=["av1"], transfer_quality=[70],
|
||||
mod.DistributedCollector.execute(images=["img"], audio=[None], transfer_format=[fmt], transfer_quality=[70],
|
||||
multi_job_id=["job"], is_worker=[True])
|
||||
|
||||
assert calls[-1]["transfer_format"] == ["av1"]
|
||||
assert calls[-1]["transfer_format"] == [fmt]
|
||||
assert calls[-1]["transfer_quality"] == [70]
|
||||
assert calls[-1]["multi_job_id"] == ["job"]
|
||||
|
||||
|
||||
@@ -84,6 +84,14 @@ def validate_serverless_worker(worker: dict) -> dict:
|
||||
if "enabled" in result and type(result["enabled"]) is not bool:
|
||||
raise ServerlessConfigError("enabled must be a boolean")
|
||||
|
||||
# Omission preserves old saved records and means legacy-image-v1. Selecting
|
||||
# comfy-media-v1 declares an installed endpoint bridge, not health detection;
|
||||
# its schema/representations/limits are fixed by serverless_bridge/contract.py.
|
||||
if "media_profile" in result:
|
||||
profile = result["media_profile"]
|
||||
if not isinstance(profile, str) or profile not in ("legacy-image-v1", "comfy-media-v1"):
|
||||
raise ServerlessConfigError("media_profile must be legacy-image-v1 or comfy-media-v1")
|
||||
|
||||
for field, (minimum, maximum, default) in _NUMERIC_FIELDS.items():
|
||||
value = result.get(field, default)
|
||||
if type(value) is not int or not minimum <= value <= maximum:
|
||||
|
||||
+18
-5
@@ -17,9 +17,13 @@ jpeg, webp Lossy per-frame stills.
|
||||
h264, h265, av1
|
||||
Lossy video for the whole batch. Smallest by far for video
|
||||
frame batches, where consecutive frames are similar.
|
||||
exact Original floating-point tensors (raw little-endian bytes, dtype,
|
||||
channels and values preserved). Largest; opt-in. Encoded by the
|
||||
Collector/job_complete route with serverless_bridge.codec, never by
|
||||
encode_chunks/decode_chunk below.
|
||||
|
||||
All formats quantize to 8-bit exactly like the legacy path did, so the
|
||||
lossless formats reproduce the legacy master tensors bit-for-bit.
|
||||
All formats except exact quantize to 8-bit exactly like the legacy path did,
|
||||
so the lossless formats reproduce the legacy master tensors bit-for-bit.
|
||||
"""
|
||||
import io
|
||||
import math
|
||||
@@ -29,20 +33,25 @@ import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
LEGACY_FORMAT = "legacy_png"
|
||||
EXACT_FORMAT = "exact"
|
||||
IMAGE_FORMATS = ("png", "jpeg", "webp")
|
||||
VIDEO_FORMATS = ("ffv1", "h264", "h265", "av1")
|
||||
LOSSLESS_FORMATS = (LEGACY_FORMAT, "png", "ffv1")
|
||||
TRANSFER_FORMATS = (LEGACY_FORMAT, "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1")
|
||||
# Appended, so the pre-existing option order (and saved widget values) is unchanged.
|
||||
TRANSFER_FORMATS = (LEGACY_FORMAT, "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1", EXACT_FORMAT)
|
||||
|
||||
DEFAULT_TRANSFER_FORMAT = LEGACY_FORMAT
|
||||
DEFAULT_TRANSFER_QUALITY = 90
|
||||
|
||||
TRANSFER_FORMAT_TOOLTIP = (
|
||||
"How workers send images back to the master. "
|
||||
"How remote workers send images back to the master; the master's own images are never re-encoded. "
|
||||
"legacy_png: original uncompressed PNG (largest). "
|
||||
"png / ffv1: lossless, identical result. "
|
||||
"jpeg / webp: lossy stills. "
|
||||
"h264 / h265 / av1: lossy video of the whole batch, smallest for frame sequences."
|
||||
"h264 / h265 / av1: lossy video of the whole batch, smallest for frame sequences "
|
||||
"(h265 is experimental). "
|
||||
"exact: original floating-point values, no 8-bit rounding; much larger, opt-in. "
|
||||
"Serverless media pools send legacy_png as png (same pixels) and never substitute another format."
|
||||
)
|
||||
TRANSFER_QUALITY_TOOLTIP = (
|
||||
"Quality for lossy transfer formats (ignored by lossless ones). "
|
||||
@@ -298,6 +307,8 @@ def encode_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes=None):
|
||||
fmt = normalize_format(fmt)
|
||||
if fmt == LEGACY_FORMAT:
|
||||
raise TransferCodecError("legacy_png is sent through the JSON path, not encode_chunks")
|
||||
if fmt == EXACT_FORMAT:
|
||||
raise TransferCodecError("exact tensors are not encoded by encode_chunks")
|
||||
total = int(frames.shape[0])
|
||||
if total == 0:
|
||||
return
|
||||
@@ -378,6 +389,8 @@ def decode_chunk(envelope, parts):
|
||||
fmt = normalize_format(envelope.get("format"))
|
||||
if fmt == LEGACY_FORMAT:
|
||||
raise TransferCodecError("legacy_png chunks are not sent as multipart")
|
||||
if fmt == EXACT_FORMAT:
|
||||
raise TransferCodecError("exact chunks are decoded as tensors, not frames")
|
||||
count = envelope.get("count")
|
||||
width = envelope.get("width")
|
||||
height = envelope.get("height")
|
||||
|
||||
@@ -7,7 +7,7 @@ import { createWorkerSettingsForm } from "../ui/settingsForm.js";
|
||||
import { DistributedUI } from "../ui.js";
|
||||
import { checkWorkerStatus, formatServerlessStatus, getWorkerUrl, launchWorker, stopWorker, viewWorkerLog } from "../workerLifecycle.js";
|
||||
import { handleClearMemory, handleInterruptWorkers } from "../workerUtils.js";
|
||||
import { isRemoteWorker, normalizeRunpodEndpointId, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js";
|
||||
import { cancelWorkerSettings, isRemoteWorker, normalizeRunpodEndpointId, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js";
|
||||
import { performPreflightCheck } from "../executionUtils.js";
|
||||
|
||||
function createDocumentMock() {
|
||||
@@ -44,6 +44,7 @@ function createDocumentMock() {
|
||||
document: {
|
||||
createElement: makeElement,
|
||||
getElementById: (id) => elements.get(id) || null,
|
||||
querySelector: () => null,
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -125,19 +126,94 @@ describe("Runpod serverless worker panel", () => {
|
||||
expect(elements.get(`${field}-group-saved`).style.display).toBe("none");
|
||||
}
|
||||
expect(elements.get("runpod-text-saved").style.display).toBe("none");
|
||||
for (const field of ["endpoint-id", "api-key", "clear-api-key", "jobs-per-run", "timeout-seconds", "serverless-test"]) {
|
||||
for (const field of ["endpoint-id", "media-profile", "api-key", "clear-api-key", "jobs-per-run", "timeout-seconds", "serverless-test"]) {
|
||||
expect(elements.get(`${field}-group-saved`).style.display).toBe("flex");
|
||||
}
|
||||
// The separate concurrency setting was removed; a legacy value is ignored.
|
||||
expect(elements.has("concurrency-saved")).toBe(false);
|
||||
expect(elements.has("concurrency-group-saved")).toBe(false);
|
||||
expect(elements.get("endpoint-id-saved").value).toBe("ep_123");
|
||||
expect(elements.get("media-profile-saved").value).toBe("legacy-image-v1");
|
||||
expect(elements.get("jobs-per-run-saved").value).toBe("2");
|
||||
expect(elements.get("timeout-seconds-saved").value).toBe("900");
|
||||
expect(elements.get("api-key-saved").value).toBe("");
|
||||
expect(elements.get("api-key-saved").placeholder).toMatch(/leave blank/i);
|
||||
});
|
||||
|
||||
it("renders explicit endpoint profiles with installation and no-auto-detection guidance", () => {
|
||||
createWorkerSettingsForm(createFormUi(), {}, {
|
||||
id: "profile", name: "Runpod", type: "serverless", media_profile: "comfy-media-v1",
|
||||
});
|
||||
const select = elements.get("media-profile-profile");
|
||||
const group = elements.get("media-profile-group-profile");
|
||||
expect(group.children[0].textContent).toBe("Endpoint media profile");
|
||||
expect(group.children[0].htmlFor).toBe(select.id);
|
||||
expect(select.children.map((option) => [option.value, option.textContent])).toEqual([
|
||||
["legacy-image-v1", "Legacy images (existing handler)"],
|
||||
["comfy-media-v1", "Typed media v1 (requires endpoint bridge)"],
|
||||
]);
|
||||
expect(select.value).toBe("comfy-media-v1");
|
||||
expect(select.title).toMatch(/after installing.*endpoint bridge/i);
|
||||
expect(select.title).toMatch(/health checks do not auto-detect/i);
|
||||
expect(select.title).toMatch(/schema.*formats and limits are fixed/i);
|
||||
const type = elements.get("worker-type-profile");
|
||||
type.value = "remote";
|
||||
type.dispatchEvent();
|
||||
expect(group.style.display).toBe("none");
|
||||
type.value = "serverless";
|
||||
type.dispatchEvent();
|
||||
expect(group.style.display).toBe("flex");
|
||||
expect(select.value).toBe("comfy-media-v1");
|
||||
});
|
||||
|
||||
it.each(["legacy-image-v1", "comfy-media-v1"])("round-trips an explicit %s selection through the save payload and reload", async (profile) => {
|
||||
const worker = { id: "roundtrip", name: "Runpod", type: "serverless", endpoint_id: "ep_123", enabled: false };
|
||||
let persisted;
|
||||
const extension = makeSaveExtension(worker, vi.fn(async (id, payload) => {
|
||||
persisted = JSON.parse(JSON.stringify({ id, ...payload }));
|
||||
return {};
|
||||
}));
|
||||
const ui = createFormUi();
|
||||
const form = createWorkerSettingsForm(ui, extension, worker);
|
||||
elements.get("media-profile-roundtrip").value = profile;
|
||||
const saveButton = form.children.at(-1).children.find((button) => button.textContent === "Save");
|
||||
await saveButton.onclick();
|
||||
expect(persisted.media_profile).toBe(profile);
|
||||
expect(worker.media_profile).toBe(profile);
|
||||
expect(persisted).not.toHaveProperty("api_key");
|
||||
expect(persisted).not.toHaveProperty("concurrency");
|
||||
createWorkerSettingsForm(ui, extension, persisted);
|
||||
expect(elements.get("media-profile-roundtrip").value).toBe(profile);
|
||||
});
|
||||
|
||||
it.each([undefined, "comfy-media-v1"])("restores the saved profile on cancel (%s)", (profile) => {
|
||||
const worker = { id: "cancel", name: "Runpod", type: "serverless", media_profile: profile };
|
||||
const extension = makeSaveExtension(worker);
|
||||
createWorkerSettingsForm(createFormUi(), extension, worker);
|
||||
elements.get("media-profile-cancel").value = profile ? "legacy-image-v1" : "comfy-media-v1";
|
||||
cancelWorkerSettings(extension, worker.id);
|
||||
expect(elements.get("media-profile-cancel").value).toBe(profile ?? "legacy-image-v1");
|
||||
expect(extension.api.updateWorker).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("keeps the saved profile after a busy-guard error and rejects invalid selections", async () => {
|
||||
const worker = { id: "busy", name: "Runpod", type: "serverless", endpoint_id: "ep_123", media_profile: "comfy-media-v1" };
|
||||
const extension = makeSaveExtension(worker, vi.fn().mockRejectedValue(new Error("Finish or reconcile unresolved serverless jobs")));
|
||||
createWorkerSettingsForm(createFormUi(), extension, worker);
|
||||
elements.get("media-profile-busy").value = "legacy-image-v1";
|
||||
await saveWorkerSettings(extension, worker.id);
|
||||
expect(worker.media_profile).toBe("comfy-media-v1");
|
||||
expect(extension.app.extensionManager.toast.add).toHaveBeenLastCalledWith(
|
||||
expect.objectContaining({ summary: "Save Failed", detail: expect.stringMatching(/unresolved/) }),
|
||||
);
|
||||
elements.get("media-profile-busy").value = "auto";
|
||||
await saveWorkerSettings(extension, worker.id);
|
||||
expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
|
||||
expect(extension.app.extensionManager.toast.add).toHaveBeenLastCalledWith(
|
||||
expect.objectContaining({ summary: "Validation Error", detail: expect.stringMatching(/media profile/) }),
|
||||
);
|
||||
});
|
||||
|
||||
it("preserves a saved cloud port on first render and applies defaults only on type changes", () => {
|
||||
vi.spyOn(document, "getElementById").mockReturnValue(null);
|
||||
createWorkerSettingsForm(createFormUi(), {}, {
|
||||
@@ -218,6 +294,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
elements.set("name-w2", { value: "Runpod worker" });
|
||||
elements.set("worker-type-w2", { value: "serverless" });
|
||||
elements.set("endpoint-id-w2", { value: "ep_123" });
|
||||
elements.set("media-profile-w2", { value: "legacy-image-v1" });
|
||||
elements.set("api-key-w2", { value: "secret-token" });
|
||||
elements.set("clear-api-key-w2", { checked: false });
|
||||
elements.set("jobs-per-run-w2", { value: "4" });
|
||||
@@ -227,6 +304,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
|
||||
expect(extension.api.updateWorker).toHaveBeenCalledWith("w2", {
|
||||
name: "Runpod worker", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
|
||||
media_profile: "legacy-image-v1",
|
||||
api_key: "secret-token", jobs_per_run: 4, timeout_seconds: 2400, enabled: false,
|
||||
});
|
||||
expect(worker).toMatchObject({ type: "serverless", provider: "runpod", enabled: false, api_key_configured: true });
|
||||
@@ -244,6 +322,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
elements.set("name-legacy", { value: "Runpod" });
|
||||
elements.set("worker-type-legacy", { value: "serverless" });
|
||||
elements.set("endpoint-id-legacy", { value: "ep_123" });
|
||||
elements.set("media-profile-legacy", { value: "legacy-image-v1" });
|
||||
elements.set("api-key-legacy", { value: "" });
|
||||
elements.set("clear-api-key-legacy", { checked: false });
|
||||
elements.set("jobs-per-run-legacy", { value: "2" });
|
||||
@@ -278,6 +357,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
elements.set("name-url", { value: "Runpod" });
|
||||
elements.set("worker-type-url", { value: "serverless" });
|
||||
elements.set("endpoint-id-url", { value: "https://api.runpod.ai/v2/abc123xyz/runsync" });
|
||||
elements.set("media-profile-url", { value: "legacy-image-v1" });
|
||||
elements.set("api-key-url", { value: "" });
|
||||
elements.set("clear-api-key-url", { checked: false });
|
||||
elements.set("jobs-per-run-url", { value: "1" });
|
||||
@@ -294,7 +374,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
});
|
||||
|
||||
it("asks before a type change deletes the saved API key, and cancels cleanly", async () => {
|
||||
const worker = { id: "conv", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: true };
|
||||
const worker = { id: "conv", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: true, media_profile: "comfy-media-v1" };
|
||||
const extension = makeSaveExtension(worker);
|
||||
elements.set("name-conv", { value: "Runpod" });
|
||||
elements.set("worker-type-conv", { value: "remote" });
|
||||
@@ -307,11 +387,13 @@ describe("Runpod serverless worker panel", () => {
|
||||
expect(confirmSpy).toHaveBeenCalledWith(expect.stringMatching(/saved Runpod API key will be deleted/));
|
||||
expect(extension.api.updateWorker).not.toHaveBeenCalled();
|
||||
expect(worker.type).toBe("serverless");
|
||||
expect(worker.media_profile).toBe("comfy-media-v1");
|
||||
|
||||
confirmSpy.mockReturnValue(true);
|
||||
await saveWorkerSettings(extension, "conv");
|
||||
expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
|
||||
expect(extension.api.updateWorker.mock.calls[0][1].type).toBe("remote");
|
||||
expect(worker).not.toHaveProperty("media_profile");
|
||||
} finally {
|
||||
delete globalThis.confirm;
|
||||
}
|
||||
@@ -341,6 +423,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
elements.set("name-w-blank", { value: "Runpod" });
|
||||
elements.set("worker-type-w-blank", { value: "serverless" });
|
||||
elements.set("endpoint-id-w-blank", { value: "ep_123" });
|
||||
elements.set("media-profile-w-blank", { value: "legacy-image-v1" });
|
||||
elements.set("api-key-w-blank", { value: "" });
|
||||
elements.set("clear-api-key-w-blank", { checked: false });
|
||||
elements.set("jobs-per-run-w-blank", { value: "1" });
|
||||
@@ -361,6 +444,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
elements.set("name-w3", { value: "Runpod" });
|
||||
elements.set("worker-type-w3", { value: "serverless" });
|
||||
elements.set("endpoint-id-w3", { value: "ep_123" });
|
||||
elements.set("media-profile-w3", { value: "legacy-image-v1" });
|
||||
elements.set("api-key-w3", { value: "" });
|
||||
elements.set("clear-api-key-w3", { checked: true });
|
||||
elements.set("jobs-per-run-w3", { value: "1" });
|
||||
@@ -370,6 +454,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
|
||||
expect(extension.api.updateWorker).toHaveBeenCalledWith("w3", {
|
||||
name: "Runpod", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
|
||||
media_profile: "legacy-image-v1",
|
||||
clear_api_key: true, jobs_per_run: 1, timeout_seconds: 1800, enabled: true,
|
||||
});
|
||||
expect(worker.api_key_configured).toBe(false);
|
||||
@@ -527,9 +612,6 @@ describe("Runpod serverless worker panel", () => {
|
||||
});
|
||||
|
||||
it("tests only the saved Runpod endpoint and reports the backend message", async () => {
|
||||
const button = document.createElement("button");
|
||||
button.id = "test-serverless-w6";
|
||||
button.textContent = "Test Saved Endpoint";
|
||||
const extension = {
|
||||
api: {
|
||||
testServerlessWorker: vi.fn().mockResolvedValue({ ok: true, message: "Endpoint reachable" }),
|
||||
@@ -537,9 +619,15 @@ describe("Runpod serverless worker panel", () => {
|
||||
},
|
||||
app: { extensionManager: { toast: { add: vi.fn() } } },
|
||||
};
|
||||
const worker = { id: "w6", name: "Runpod", type: "serverless" };
|
||||
extension.config = { workers: [worker] };
|
||||
createWorkerSettingsForm(createFormUi(), extension, worker);
|
||||
const button = elements.get("test-serverless-w6");
|
||||
|
||||
await expect(testServerlessWorker(extension, "w6")).resolves.toEqual({ ok: true, message: "Endpoint reachable" });
|
||||
|
||||
expect(worker).not.toHaveProperty("media_profile");
|
||||
expect(elements.get("media-profile-w6").value).toBe("legacy-image-v1");
|
||||
expect(extension.api.testServerlessWorker).toHaveBeenCalledWith("w6");
|
||||
expect(extension.api.queueDistributed).not.toHaveBeenCalled();
|
||||
expect(extension.app.extensionManager.toast.add).toHaveBeenCalledWith(expect.objectContaining({
|
||||
|
||||
@@ -4,6 +4,9 @@ import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings,
|
||||
export const IMAGES_PER_RUN_TOOLTIP =
|
||||
"Separate Runpod jobs submitted each run, all at once. Runpod's max workers decides how many GPUs actually run them.";
|
||||
|
||||
const MEDIA_PROFILE_TOOLTIP =
|
||||
"Select Typed media v1 only after installing the ComfyUI-Distributed endpoint bridge. This declares bridge support; endpoint health checks do not auto-detect it. The v1 schema, media formats and limits are fixed.";
|
||||
|
||||
export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
const form = document.createElement("div");
|
||||
form.style.cssText = "display: flex; flex-direction: column; gap: 8px;";
|
||||
@@ -59,6 +62,7 @@ export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
const portInput = portGroup.input;
|
||||
const serverlessGroups = [
|
||||
endpointGroup.group,
|
||||
mediaProfileGroup,
|
||||
apiKeyGroup.group,
|
||||
clearApiKeyGroup,
|
||||
jobsPerRunGroup.group,
|
||||
@@ -137,6 +141,32 @@ export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
endpointGroup.group.style.display = "none";
|
||||
form.appendChild(endpointGroup.group);
|
||||
|
||||
const mediaProfileGroup = document.createElement("div");
|
||||
mediaProfileGroup.id = `media-profile-group-${worker.id}`;
|
||||
mediaProfileGroup.style.cssText = "display: none; flex-direction: column; gap: 4px;";
|
||||
mediaProfileGroup.title = MEDIA_PROFILE_TOOLTIP;
|
||||
const mediaProfileLabel = document.createElement("label");
|
||||
mediaProfileLabel.htmlFor = `media-profile-${worker.id}`;
|
||||
mediaProfileLabel.textContent = "Endpoint media profile";
|
||||
mediaProfileLabel.style.cssText = typeLabel.style.cssText;
|
||||
const mediaProfileSelect = document.createElement("select");
|
||||
mediaProfileSelect.id = mediaProfileLabel.htmlFor;
|
||||
mediaProfileSelect.style.cssText = typeSelect.style.cssText;
|
||||
mediaProfileSelect.title = MEDIA_PROFILE_TOOLTIP;
|
||||
for (const [value, label] of [
|
||||
["legacy-image-v1", "Legacy images (existing handler)"],
|
||||
["comfy-media-v1", "Typed media v1 (requires endpoint bridge)"],
|
||||
]) {
|
||||
const option = document.createElement("option");
|
||||
option.value = value;
|
||||
option.textContent = label;
|
||||
mediaProfileSelect.appendChild(option);
|
||||
}
|
||||
mediaProfileSelect.value = worker.media_profile ?? "legacy-image-v1";
|
||||
mediaProfileGroup.appendChild(mediaProfileLabel);
|
||||
mediaProfileGroup.appendChild(mediaProfileSelect);
|
||||
form.appendChild(mediaProfileGroup);
|
||||
|
||||
const apiKeyGroup = ui.createFormGroup("Runpod API Key (write-only):", "", `api-key-${worker.id}`, "password", worker.api_key_configured ? "Leave blank to keep saved API key" : "Runpod API key");
|
||||
apiKeyGroup.group.id = `api-key-group-${worker.id}`;
|
||||
apiKeyGroup.input.autocomplete = "new-password";
|
||||
|
||||
@@ -61,6 +61,7 @@ function showServerlessValidationError(extension, detail) {
|
||||
|
||||
async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
|
||||
const endpointId = normalizeRunpodEndpointId(document.getElementById(`endpoint-id-${workerId}`).value);
|
||||
const mediaProfile = document.getElementById(`media-profile-${workerId}`).value;
|
||||
const apiKey = document.getElementById(`api-key-${workerId}`).value || "";
|
||||
const clearApiKey = Boolean(document.getElementById(`clear-api-key-${workerId}`).checked);
|
||||
const jobsPerRun = Number(document.getElementById(`jobs-per-run-${workerId}`).value);
|
||||
@@ -81,6 +82,10 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
|
||||
);
|
||||
return;
|
||||
}
|
||||
if (!["legacy-image-v1", "comfy-media-v1"].includes(mediaProfile)) {
|
||||
showServerlessValidationError(extension, "Select a supported endpoint media profile");
|
||||
return;
|
||||
}
|
||||
if (!Number.isInteger(jobsPerRun) || jobsPerRun < 1 || jobsPerRun > 16) {
|
||||
showServerlessValidationError(extension, "Images per run must be a whole number from 1 to 16");
|
||||
return;
|
||||
@@ -95,6 +100,7 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
|
||||
type: "serverless",
|
||||
provider: "runpod",
|
||||
endpoint_id: endpointId,
|
||||
media_profile: mediaProfile,
|
||||
...(apiKey.length > 0 && !clearApiKey ? { api_key: apiKey } : {}),
|
||||
...(clearApiKey ? { clear_api_key: true } : {}),
|
||||
jobs_per_run: jobsPerRun,
|
||||
@@ -111,6 +117,7 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
|
||||
worker.type = "serverless";
|
||||
worker.provider = "runpod";
|
||||
worker.endpoint_id = endpointId;
|
||||
worker.media_profile = mediaProfile;
|
||||
worker.jobs_per_run = jobsPerRun;
|
||||
// Replaced by a fixed backend limit; drop any value loaded from an older config.
|
||||
delete worker.concurrency;
|
||||
@@ -310,6 +317,7 @@ export async function saveWorkerSettings(extension, workerId) {
|
||||
// Update local config
|
||||
worker.name = name.trim();
|
||||
worker.type = workerType;
|
||||
delete worker.media_profile;
|
||||
if (isRemote) {
|
||||
worker.host = hostTrimmed;
|
||||
delete worker.cuda_device;
|
||||
@@ -366,6 +374,8 @@ export function cancelWorkerSettings(extension, workerId) {
|
||||
document.getElementById(`args-${workerId}`).value = worker.extra_args || "";
|
||||
const endpointInput = document.getElementById(`endpoint-id-${workerId}`);
|
||||
if (endpointInput) endpointInput.value = worker.endpoint_id || "";
|
||||
const mediaProfileInput = document.getElementById(`media-profile-${workerId}`);
|
||||
if (mediaProfileInput) mediaProfileInput.value = worker.media_profile ?? "legacy-image-v1";
|
||||
const apiKeyInput = document.getElementById(`api-key-${workerId}`);
|
||||
if (apiKeyInput) {
|
||||
apiKeyInput.value = "";
|
||||
|
||||
Reference in New Issue
Block a user