From f0927f1166ebfda0b4c17199ec0cfc81d7342559 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 13 Oct 2026 23:53:58 +0000 Subject: [PATCH] feat: add versioned lossless serverless media transport --- api/config_routes.py | 24 +- api/orchestration/prompt_transform.py | 24 +- api/orchestration/runpod_client.py | 10 +- api/orchestration/serverless_media_prompt.py | 191 ++++ api/orchestration/serverless_runtime.py | 41 +- api/orchestration/serverless_typed_inputs.py | 225 +++++ api/orchestration/serverless_typed_results.py | 79 ++ api/queue_orchestration.py | 110 ++- docs/serverless-media.md | 132 +++ nodes/media_collector.py | 221 +++++ nodes/v3.py | 3 +- serverless_bridge/README.md | 229 +++++ serverless_bridge/__init__.py | 14 + serverless_bridge/codec.py | 185 ++++ serverless_bridge/contract.py | 249 +++++ serverless_bridge/handler.py | 789 ++++++++++++++++ serverless_bridge/nodes.py | 157 ++++ serverless_bridge/video_codec.py | 482 ++++++++++ tests/helpers/media_end_to_end_check.py | 432 +++++++++ tests/helpers/v3_runtime_check.py | 8 +- tests/test_collector_list_inputs.py | 18 +- tests/test_media_bridge_nodes.py | 198 ++++ tests/test_media_codec.py | 328 +++++++ tests/test_media_collector.py | 314 +++++++ tests/test_media_contract.py | 152 +++ tests/test_media_end_to_end.py | 47 + tests/test_media_handler.py | 877 ++++++++++++++++++ tests/test_media_prompt_transform.py | 89 ++ tests/test_media_queue_orchestration.py | 531 +++++++++++ tests/test_media_video_codec.py | 562 +++++++++++ tests/test_serverless_config.py | 26 + tests/test_serverless_media_prompt.py | 421 +++++++++ tests/test_serverless_routes.py | 104 ++- tests/test_serverless_typed_inputs.py | 330 +++++++ tests/test_serverless_typed_runtime.py | 319 +++++++ tests/test_v3_collector_transfer.py | 1 + utils/serverless_config.py | 8 + web/tests/serverlessWorkers.test.js | 100 +- web/ui/settingsForm.js | 30 + web/workerSettings.js | 10 + 40 files changed, 8027 insertions(+), 43 deletions(-) create mode 100644 api/orchestration/serverless_media_prompt.py create mode 100644 api/orchestration/serverless_typed_inputs.py create mode 100644 api/orchestration/serverless_typed_results.py create mode 100644 docs/serverless-media.md create mode 100644 nodes/media_collector.py create mode 100644 serverless_bridge/README.md create mode 100644 serverless_bridge/__init__.py create mode 100644 serverless_bridge/codec.py create mode 100644 serverless_bridge/contract.py create mode 100644 serverless_bridge/handler.py create mode 100644 serverless_bridge/nodes.py create mode 100644 serverless_bridge/video_codec.py create mode 100644 tests/helpers/media_end_to_end_check.py create mode 100644 tests/test_media_bridge_nodes.py create mode 100644 tests/test_media_codec.py create mode 100644 tests/test_media_collector.py create mode 100644 tests/test_media_contract.py create mode 100644 tests/test_media_end_to_end.py create mode 100644 tests/test_media_handler.py create mode 100644 tests/test_media_prompt_transform.py create mode 100644 tests/test_media_queue_orchestration.py create mode 100644 tests/test_media_video_codec.py create mode 100644 tests/test_serverless_media_prompt.py create mode 100644 tests/test_serverless_typed_inputs.py create mode 100644 tests/test_serverless_typed_runtime.py diff --git a/api/config_routes.py b/api/config_routes.py index 2ea58b7..29590f1 100644 --- a/api/config_routes.py +++ b/api/config_routes.py @@ -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) diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 318b22d..d288da3 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -401,6 +401,20 @@ def prune_prompt_for_worker(prompt_obj, use_worker_sink=False): def prepare_delegate_master_prompt(prompt_obj, collector_ids): """Prune master prompt so it only executes post-collector nodes in delegate mode.""" + # Typed collectors need no fabricated IMAGE placeholder. Sever all media + # links, including STRING scalar branches, before retained-config analysis. + # Their expected_media control was frozen by participant overrides. + if any((prompt_obj.get(cid) or {}).get("class_type") == "DistributedMediaCollector" + for cid in collector_ids): + from ...serverless_bridge.contract import SOCKET_KINDS + prompt_obj = json.loads(json.dumps(prompt_obj)) + for cid in collector_ids: + node = prompt_obj.get(cid) or {} + if node.get("class_type") == "DistributedMediaCollector": + inputs = node.setdefault("inputs", {}) + inputs.setdefault("expected_media", json.dumps([name for name in SOCKET_KINDS if name in inputs])) + for name in SOCKET_KINDS: + inputs.pop(name, None) downstream = _find_downstream_nodes(prompt_obj, collector_ids) nodes_to_keep = set(collector_ids) nodes_to_keep.update(downstream) @@ -505,6 +519,7 @@ def generate_job_id_map(prompt_index, prefix): job_map = {} distributed_nodes = ( prompt_index.nodes_for_class("DistributedCollector") + + prompt_index.nodes_for_class("DistributedMediaCollector") + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed") + prompt_index.nodes_for_class("DistributedImageListStart") ) @@ -537,8 +552,10 @@ def _override_collector_nodes( enabled_json, delegate_master, ): - """Configure DistributedCollector nodes for master or worker role.""" - for node_id in prompt_index.nodes_for_class("DistributedCollector"): + """Configure legacy and typed collectors without changing their media.""" + collector_ids = (prompt_index.nodes_for_class("DistributedCollector") + + prompt_index.nodes_for_class("DistributedMediaCollector")) + for node_id in collector_ids: node = prompt_copy.get(node_id) if not isinstance(node, dict): continue @@ -551,6 +568,9 @@ def _override_collector_nodes( inputs["multi_job_id"] = job_id_map.get(node_id, node_id) inputs["is_worker"] = not is_master inputs["enabled_worker_ids"] = enabled_json + if node.get("class_type") == "DistributedMediaCollector": + from ...serverless_bridge.contract import SOCKET_KINDS + inputs["expected_media"] = json.dumps([socket for socket in SOCKET_KINDS if socket in inputs]) if is_master: inputs["delegate_only"] = bool(delegate_master) inputs.pop("master_url", None) diff --git a/api/orchestration/runpod_client.py b/api/orchestration/runpod_client.py index dbf7c1a..ee7e44a 100644 --- a/api/orchestration/runpod_client.py +++ b/api/orchestration/runpod_client.py @@ -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")) diff --git a/api/orchestration/serverless_media_prompt.py b/api/orchestration/serverless_media_prompt.py new file mode 100644 index 0000000..17087e5 --- /dev/null +++ b/api/orchestration/serverless_media_prompt.py @@ -0,0 +1,191 @@ +"""Pure compilation of prepared typed-media collectors for serverless workers. + +The legacy image compiler remains separate. This compiler emits one MediaExportV1 +execution target per connected socket, 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_prompt import ( + _PRIMITIVE_NODES, + _REMOVED_SCALAR_NODES, + _is_link, + _is_unsupported_distributed, + _make_scalar_resolver, + _prune_graph, +) + + +_COLLECTOR = "DistributedMediaCollector" + + +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 class_type == "DistributedCollector" or _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 DistributedMediaCollector 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(SOCKET_KINDS): + 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", {}) + for field in ("multi_job_id", "worker_id"): + _validate_identifier(inputs.get(field), f"collector {collector_id!r} {field}") + connected = [socket for socket in SOCKET_KINDS if socket in inputs] + 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, kind in SOCKET_KINDS.items(): + 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]), + }, + } + 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 DistributedMediaCollector must have prepared literal ``multi_job_id`` + and ``worker_id`` identities and at least one linked media input. Present + media inputs (including texts) must be API links; literal text belongs in a + PrimitiveString node. Nonmedia collector 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 diff --git a/api/orchestration/serverless_runtime.py b/api/orchestration/serverless_runtime.py index a28c7b5..5379e63 100644 --- a/api/orchestration/serverless_runtime.py +++ b/api/orchestration/serverless_runtime.py @@ -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 diff --git a/api/orchestration/serverless_typed_inputs.py b/api/orchestration/serverless_typed_inputs.py new file mode 100644 index 0000000..191635a --- /dev/null +++ b/api/orchestration/serverless_typed_inputs.py @@ -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 diff --git a/api/orchestration/serverless_typed_results.py b/api/orchestration/serverless_typed_results.py new file mode 100644 index 0000000..cc4d718 --- /dev/null +++ b/api/orchestration/serverless_typed_results.py @@ -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 diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index d10872a..b6a411e 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -1,4 +1,5 @@ import asyncio +import copy import time import uuid @@ -85,6 +86,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}")) @@ -147,13 +152,47 @@ def _is_load_balance_enabled(value): def _prompt_requests_load_balance(prompt_index): - for node_id in prompt_index.nodes_for_class("DistributedCollector"): + collector_ids = (prompt_index.nodes_for_class("DistributedCollector") + + prompt_index.nodes_for_class("DistributedMediaCollector")) + for node_id in collector_ids: inputs = prompt_index.inputs_by_node.get(node_id, {}) if _is_load_balance_enabled(inputs.get("load_balance", False)): return True return False +def _validate_media_workers(prompt_index, workers): + """Reject incompatible selections before probing, pruning or dispatching. + + Typed collectors currently have only a master/serverless transport. Check + all selected workers, not just reachable or load-balanced participants, so + a bad selection cannot silently turn into a different execution mode. + """ + typed = bool(prompt_index.nodes_for_class("DistributedMediaCollector")) + if typed and (prompt_index.nodes_for_class("DistributedCollector") + or prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")): + raise ValueError( + "DistributedMediaCollector cannot be mixed with DistributedCollector or " + "UltimateSDUpscaleDistributed in one prompt" + ) + for worker in workers: + if worker.get("type") != "serverless": + if typed: + raise ValueError( + "DistributedMediaCollector does not support selected direct workers; " + "use master execution or comfy-media-v1 serverless pools" + ) + 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 typed != (profile == "comfy-media-v1"): + raise ValueError( + "DistributedMediaCollector requires comfy-media-v1 serverless pools; " + "legacy collectors require legacy-image-v1" + ) + + async def _prepare_worker_payload( worker, prompt_index, @@ -169,11 +208,34 @@ 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(prompt_index, [worker]) 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 legacy pruner only recognizes Collector/Upscale roots. + # The typed compiler validates the complete private graph and + # prunes it itself, retaining every connected typed 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, @@ -266,21 +328,28 @@ async def orchestrate_distributed_execution( ) = _resolve_orchestration_limits(config) requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None prompt_index = PromptIndex(prompt_obj) + typed_media_mode = bool(prompt_index.nodes_for_class("DistributedMediaCollector")) + if typed_media_mode: + from ..serverless_bridge.contract import identifier + # PromptIndex's JSON copy would otherwise turn integer keys into strings. + for node_id in prompt_obj: + identifier(node_id) image_list_start_ids = prompt_index.nodes_for_class("DistributedImageListStart") image_list_mode = bool(image_list_start_ids) 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) other_distributed = ( prompt_index.nodes_for_class("DistributedCollector") + + prompt_index.nodes_for_class("DistributedMediaCollector") + 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" + "Distributed Image List v1 cannot be mixed with Distributed Collector, " + "DistributedMediaCollector or Ultimate SD Upscale nodes in one prompt" ) + validate_image_list_regions(prompt_obj) configured_serverless = [ str(worker.get("id")) for worker in config.get("workers", []) if worker.get("enabled", False) @@ -300,7 +369,14 @@ async def orchestrate_distributed_execution( "serverless pools are unsupported" ) else: + if typed_media_mode and requested_ids is not None: + configured_ids = {worker.get("id") for worker in config.get("workers", [])} + if any(worker_id not in configured_ids for worker_id in requested_ids): + raise ValueError("Unknown explicitly selected media worker ID") workers = _resolve_enabled_workers(config, requested_ids) + _validate_media_workers(prompt_index, workers) + if typed_media_mode and delegate_master and not workers: + raise RuntimeError("Delegate-only DistributedMediaCollector requires a selected comfy-media-v1 pool") load_balance_requested = _prompt_requests_load_balance(prompt_index) trace_debug( execution_trace_id, @@ -316,8 +392,8 @@ async def orchestrate_distributed_execution( ), ) - # Keep the legacy Collector fallback, but never override image-list - # orchestrator-only mode merely because its direct workers are unavailable. + # Keep the legacy Collector fallback. Image-list and typed-media delegate + # modes must not silently enable master participation. if not image_list_mode and not workers and delegate_master: trace_debug( execution_trace_id, @@ -332,6 +408,8 @@ async def orchestrate_distributed_execution( trace_execution_id=execution_trace_id, probe_concurrency=worker_probe_concurrency, ) + if typed_media_mode and delegate_master and not active_workers: + raise RuntimeError("No selected comfy-media-v1 worker is reachable for delegate-only DistributedMediaCollector") if image_list_mode: from .orchestration.dispatch import preflight_image_list_workers if not active_workers and delegate_master: @@ -430,7 +508,8 @@ async def orchestrate_distributed_execution( ) if delegate_master and not image_list_mode: - collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector") + collector_ids = (find_nodes_by_class(master_prompt, "DistributedCollector") + + find_nodes_by_class(master_prompt, "DistributedMediaCollector")) upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed") if upscale_nodes: debug_log( @@ -477,6 +556,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)) @@ -522,8 +610,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"}) diff --git a/docs/serverless-media.md b/docs/serverless-media.md new file mode 100644 index 0000000..b55bcc4 --- /dev/null +++ b/docs/serverless-media.md @@ -0,0 +1,132 @@ +# Typed serverless media (PR2) + +`comfy-media-v1` adds typed, lossless list results for `IMAGE`, `AUDIO`, `MASK`, +`STRING`, and native ComfyUI `VIDEO`. It is an **explicit endpoint profile**, not +an automatically detected capability. Omitting `media_profile` still selects +`legacy-image-v1`; existing image collectors and handlers keep their old contract. +See [legacy serverless setup](serverless-workers.md) for credentials and recovery. + +## Use a typed collector + +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 Media Collector** (`DistributedMediaCollector`), connect at + least one typed input, and connect its outputs to the appropriate downstream + nodes. For literal text, connect a `PrimitiveString` node to `texts`. +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 has five optional inputs and five **list outputs**, in this order: + +| Socket / output index | Native type | Each list item | +| --- | --- | --- | +| `images` / 0 | `IMAGE` | One BHWC tensor, including its batch dimension | +| `audio` / 1 | `AUDIO` | `{ "waveform": BCT_tensor, "sample_rate": positive_int }` | +| `masks` / 2 | `MASK` | One BHW tensor | +| `texts` / 3 | `STRING` | Exact text, including whitespace or an empty string | +| `video` / 4 | `VIDEO` | One supported native video object | + +For each connected socket, order is **master first, then the prepared participant +order, then that participant's item order**—never completion order. Delegate-only +mode omits the master. An unconnected socket returns `[]`. A connected remote +socket must return 1–64 items; an empty string is an item, an empty list is not. +Lists are not concatenated into one tensor, resampled, resized, or aligned across +sockets. Different items may have different valid shapes or audio sample rates. +The entire remote response is validated and decoded before any of its results +are delivered; missing, duplicate, mismatched, or corrupt results fail collection. + +## 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`). + +### Native VIDEO is not 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. +- Remote typed transport is **serverless only**. Direct-worker typed transport + is not implemented. Do not mix a typed collector with legacy image collectors, + unsupported distributed/batch nodes, or collector-to-collector dependencies. + Downstream processing stays on the master. Image-list scheduling is separate + PR3 work, not part of this contract. + +## 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 and its offline tests; they do not claim +live Runpod/provider or GPU end-to-end validation. diff --git a/nodes/media_collector.py b/nodes/media_collector.py new file mode 100644 index 0000000..b994c9c --- /dev/null +++ b/nodes/media_collector.py @@ -0,0 +1,221 @@ +"""Typed, lossless list collection for serverless media participants. + +The master compiler supplies expected_media before pruning delegate inputs. +The provider manager owns remote deadlines and delivers one atomic native-media +message per participant. This node never serializes, concatenates or resamples. +""" +import asyncio +import json +import time + +import comfy.model_management +from comfy_api.v0_0_2 import io +import server + +from ..serverless_bridge.contract import MAX_ITEMS, SOCKET_KINDS, identifier +from ..utils.async_helpers import run_async_in_server_loop + +POLL_SECONDS = 0.1 +START_GRACE_SECONDS = 2.0 + + +def _control(value, name, kind): + # Only execution controls are scalars wrapped by ComfyUI's list mapper. + if isinstance(value, list): + if len(value) != 1: + raise ValueError(f'{name} must be a singleton control') + value = value[0] + if type(value) is not kind: + raise ValueError(f'Invalid {name} control') + return value + + +def _names(value, name): + try: + result = json.loads(_control(value, name, str)) + except (ValueError, TypeError): + raise ValueError(f'{name} must be a JSON string list') from None + if (not isinstance(result, list) or any(not isinstance(item, str) for item in result) + or len(result) != len(set(result))): + raise ValueError(f'{name} must contain unique strings') + return result + + +def _validate_media(media, expected): + if not isinstance(media, dict) or set(media) != set(expected): + raise ValueError('Media sockets do not match expected_media') + for socket, items in media.items(): + if (socket not in SOCKET_KINDS or not isinstance(items, list) + or not 1 <= len(items) <= MAX_ITEMS or any(item is None for item in items)): + raise ValueError('Connected media sockets require nonempty bounded item lists') + + +def _prepare(media, controls): + worker = _control(controls['is_worker'], 'is_worker', bool) + delegate = _control(controls['delegate_only'], 'delegate_only', bool) + passthrough = _control(controls['pass_through'], 'pass_through', bool) + _control(controls['load_balance'], 'load_balance', bool) # Routing is master-owned. + if worker: + raise ValueError('DistributedMediaCollector does not support direct-worker execution') + participants = _names(controls['enabled_worker_ids'], 'enabled_worker_ids') + for participant in participants: + identifier(participant) + if not participant.startswith('serverless:') or not participant[len('serverless:'):]: + raise ValueError('DistributedMediaCollector requires serverless participant IDs') + expected = _names(controls['expected_media'], 'expected_media') + if any(socket not in SOCKET_KINDS for socket in expected): + raise ValueError('Unknown expected_media socket') + local = {socket: items for socket, items in media.items() if items is not None} + _validate_media(local, local) + expected = expected or list(local) + if not expected: + raise ValueError('At least one connected media socket is required') + if not delegate or passthrough: + _validate_media(local, expected) + if delegate and not passthrough and not participants: + raise ValueError('Delegate-only collection requires serverless participants') + return local, participants, expected, delegate, passthrough + + +def _outputs(media): + return tuple(media.get(socket, []) for socket in SOCKET_KINDS) + + +def _consume(message, participants, expected, received): + if not isinstance(message, dict): + raise RuntimeError('Invalid serverless media delivery') + participant = message.get('worker_id') + if not isinstance(participant, str) or participant not in participants: + raise RuntimeError('Unexpected serverless participant') + if participant in received: + raise RuntimeError('Duplicate serverless participant delivery') + if message.get('is_last') is not True: + raise RuntimeError('Serverless media delivery must be atomic and final') + if 'serverless_error' in message: + if (set(message) != {'worker_id', 'serverless_error', 'is_last'} + or not isinstance(message['serverless_error'], str) or not message['serverless_error']): + raise RuntimeError('Invalid serverless error delivery') + # Do not echo provider-controlled text, credentials or signed URLs. + raise RuntimeError(f'Serverless participant {participant} failed to produce media') + if set(message) != {'worker_id', 'media', 'is_last'}: + raise RuntimeError('Invalid serverless media delivery fields') + _validate_media(message['media'], expected) + received[participant] = message['media'] + + +async def _collect(job, media, controls): + ps = server.PromptServer.instance + queue = None + try: + # Reuse the master's precreated queue, including early completions. + async with ps.distributed_jobs_lock: + queue = ps.distributed_pending_jobs.setdefault(job, asyncio.Queue()) + local, participants, expected, delegate, passthrough = _prepare(media, controls) + if passthrough or not participants: + return _outputs(local) + received = {} + started = time.monotonic() + while len(received) < len(participants): + comfy.model_management.throw_exception_if_processing_interrupted() + async with ps.distributed_jobs_lock: + if ps.distributed_pending_jobs.get(job) is not queue: + raise RuntimeError('Serverless collector queue was removed or replaced') + try: + message = await asyncio.wait_for(queue.get(), timeout=POLL_SECONDS) + except asyncio.TimeoutError: + comfy.model_management.throw_exception_if_processing_interrupted() + # wait_for cancellation may race an atomic manager delivery. + if not queue.empty(): + continue + manager = getattr(ps, 'distributed_serverless', None) + for participant in participants: + if participant in received: + continue + if manager is not None and manager.is_pending(job, participant): + continue # The manager enforces its remote absolute deadline. + if manager is not None and manager.owns(job, participant): + raise RuntimeError('Serverless participant stopped without delivering media') + # The master queues its prompt before manager.start journals slots. + if time.monotonic() - started >= START_GRACE_SECONDS: + raise RuntimeError('Serverless manager did not start the expected participant') + continue + _consume(message, participants, expected, received) + + comfy.model_management.throw_exception_if_processing_interrupted() + # Drain and close atomically with the manager's delivery lock. A queued + # duplicate after the final expected message must not yield success. + async with ps.distributed_jobs_lock: + if ps.distributed_pending_jobs.get(job) is not queue: + raise RuntimeError('Serverless collector queue was removed or replaced') + while not queue.empty(): + _consume(queue.get_nowait(), participants, expected, received) + del ps.distributed_pending_jobs[job] + ordered = [] if delegate else [local] + ordered.extend(received[participant] for participant in participants) + return tuple([item for source in ordered for item in source.get(socket, [])] + for socket in SOCKET_KINDS) + except BaseException: + # Remove only our queue before cancellation can trigger final deliveries. + async with ps.distributed_jobs_lock: + if queue is not None and ps.distributed_pending_jobs.get(job) is queue: + del ps.distributed_pending_jobs[job] + manager = getattr(ps, 'distributed_serverless', None) + if manager is not None: + try: + await manager.cancel(collector_job_id=job) + except BaseException: + pass # Preserve the original failure/interruption, not cleanup's. + raise + finally: + async with ps.distributed_jobs_lock: + if queue is not None and ps.distributed_pending_jobs.get(job) is queue: + del ps.distributed_pending_jobs[job] + + +async def _bridge_result(coro): + # The legacy sync bridge catches Exception, not asyncio.CancelledError. + # Carry BaseException as data so cancellation is re-raised on the caller. + try: + return await coro, None + except BaseException as error: + return None, error + + +class DistributedMediaCollector(io.ComfyNode): + @classmethod + def define_schema(cls): + types = (io.Image, io.Audio, io.Mask, io.String, io.Video) + return io.Schema( + node_id='DistributedMediaCollector', display_name='Distributed Media Collector', + category='utils', + inputs=[io.Boolean.Input('load_balance', default=False), + # Native non-widget inputs expose forceInput via extra_dict. + *[kind.Input(socket, optional=True, + **({'force_input': True} if socket == 'texts' + else {'extra_dict': {'forceInput': True}})) + for socket, kind in zip(SOCKET_KINDS, types)]], + outputs=[kind.Output(display_name=socket, is_output_list=True) + for socket, kind in zip(SOCKET_KINDS, types)], + is_input_list=True, accept_all_inputs=True, + ) + + @classmethod + def execute(cls, images=None, audio=None, masks=None, texts=None, video=None, + load_balance=False, multi_job_id='', is_worker=False, + enabled_worker_ids='[]', delegate_only=False, pass_through=False, + expected_media='[]'): + media = dict(images=images, audio=audio, masks=masks, texts=texts, video=video) + controls = dict(load_balance=load_balance, is_worker=is_worker, + enabled_worker_ids=enabled_worker_ids, delegate_only=delegate_only, + pass_through=pass_through, expected_media=expected_media) + job = _control(multi_job_id, 'multi_job_id', str) + if not job: + local, participants, _, _, passthrough = _prepare(media, controls) + if participants and not passthrough: + raise ValueError('Serverless collection requires multi_job_id') + return io.NodeOutput(*_outputs(local)) + identifier(job) + result, error = run_async_in_server_loop(_bridge_result(_collect(job, media, controls))) + if error is not None: + raise error + return io.NodeOutput(*result) diff --git a/nodes/v3.py b/nodes/v3.py index 9097da8..63ea542 100644 --- a/nodes/v3.py +++ b/nodes/v3.py @@ -11,6 +11,7 @@ from .collector import DistributedCollectorNode as _CollectorRuntime from ..utils import transfer_codec as _transfer_codec from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime from .image_list import DistributedImageListEnd, DistributedImageListStart +from .media_collector import DistributedMediaCollector class DistributedSeed(io.ComfyNode): @@ -212,4 +213,4 @@ class DistributedWorkerSink(io.ComfyNode): NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue, ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed, - DistributedWorkerSink, DistributedImageListStart, DistributedImageListEnd] + DistributedWorkerSink, DistributedImageListStart, DistributedImageListEnd, DistributedMediaCollector] diff --git a/serverless_bridge/README.md b/serverless_bridge/README.md new file mode 100644 index 0000000..c2bc06a --- /dev/null +++ b/serverless_bridge/README.md @@ -0,0 +1,229 @@ +# 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. + +For collector setup, list ordering, supported native VIDEO states, and PR2 +limitations, see [typed serverless media](../docs/serverless-media.md). + +## Install in the endpoint image + +1. Copy this complete directory to `/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 `/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 text request builder (no files or model required by this graph): + +```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": "" }`, 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. + +## 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//`. 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. + +Verification scope: source and offline tests. No live provider/GPU validation is +claimed, and this directory is not a tested turnkey provider deployment. diff --git a/serverless_bridge/__init__.py b/serverless_bridge/__init__.py new file mode 100644 index 0000000..fe3cbfe --- /dev/null +++ b/serverless_bridge/__init__.py @@ -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() diff --git a/serverless_bridge/codec.py b/serverless_bridge/codec.py new file mode 100644 index 0000000..cac8c13 --- /dev/null +++ b/serverless_bridge/codec.py @@ -0,0 +1,185 @@ +"""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 == "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 diff --git a/serverless_bridge/contract.py b/serverless_bridge/contract.py new file mode 100644 index 0000000..ef3082f --- /dev/null +++ b/serverless_bridge/contract.py @@ -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 diff --git a/serverless_bridge/handler.py b/serverless_bridge/handler.py new file mode 100644 index 0000000..3c28caf --- /dev/null +++ b/serverless_bridge/handler.py @@ -0,0 +1,789 @@ +"""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"] + if set(controls) != {"request_id", "binding_json", binding["socket"]}: + raise ValueError("Media exporter must have exactly its declared socket and controls") + 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() diff --git a/serverless_bridge/nodes.py b/serverless_bridge/nodes.py new file mode 100644 index 0000000..5c61328 --- /dev/null +++ b/serverless_bridge/nodes.py @@ -0,0 +1,157 @@ +"""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 .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): + 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): + """Write bounded owned payloads; history contains metadata, never base64.""" + identifier(request_id) + validate_binding(binding) + 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) + 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)], 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): + 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"]]) + return io.NodeOutput(ui={"media_exports": [exported]}) diff --git a/serverless_bridge/video_codec.py b/serverless_bridge/video_codec.py new file mode 100644 index 0000000..61ddcd8 --- /dev/null +++ b/serverless_bridge/video_codec.py @@ -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()) diff --git a/tests/helpers/media_end_to_end_check.py b/tests/helpers/media_end_to_end_check.py new file mode 100644 index 0000000..e70b160 --- /dev/null +++ b/tests/helpers/media_end_to_end_check.py @@ -0,0 +1,432 @@ +"""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.media_collector").DistributedMediaCollector + 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) + 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) + # Collector must preserve native objects and ordered item lists, + # not merge batches, resample, convert types, or encode again. + assert all(a is e for a, e in zip(media[socket], expected[socket])), (tag, socket) + 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 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 + variant} + 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="