Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b636fd5202 | ||
|
|
f0927f1166 |
+19
-5
@@ -97,10 +97,11 @@ _SECRET_FIELDS = {"api_key", "clear_api_key"}
|
||||
_SERVERLESS_LOCAL_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
|
||||
_SERVERLESS_ONLY_FIELDS = {
|
||||
"provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "api_key_configured",
|
||||
"media_profile",
|
||||
}
|
||||
_SERVERLESS_CONFIG_FIELDS = {
|
||||
"worker_id", "type", "provider", "endpoint_id", "jobs_per_run", "concurrency",
|
||||
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key",
|
||||
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key", "media_profile",
|
||||
}
|
||||
|
||||
|
||||
@@ -199,7 +200,13 @@ def _clean_bulk_type_conversions(old_workers, new_workers):
|
||||
if not isinstance(worker, dict) or worker.get("id") is None:
|
||||
continue
|
||||
old = old_by_id.get(str(worker["id"]))
|
||||
if old is None or old.get("type") != "serverless" or worker.get("type") == "serverless":
|
||||
if old is None or old.get("type") != "serverless":
|
||||
continue
|
||||
if worker.get("type") == "serverless":
|
||||
# Older clients do not know this field. Only an explicit selection
|
||||
# may change the endpoint contract, including on bulk saves.
|
||||
if "media_profile" not in worker and "media_profile" in old:
|
||||
worker["media_profile"] = old["media_profile"]
|
||||
continue
|
||||
converted = dict(worker)
|
||||
for field in _SERVERLESS_ONLY_FIELDS:
|
||||
@@ -218,7 +225,10 @@ def _bulk_serverless_changes(old_workers, new_workers):
|
||||
if new is None or new.get("type") != "serverless":
|
||||
_ensure_pool_idle(worker_id)
|
||||
secret_ids.add(worker_id)
|
||||
elif old.get("endpoint_id") != new.get("endpoint_id"):
|
||||
elif (
|
||||
old.get("endpoint_id") != new.get("endpoint_id")
|
||||
or old.get("media_profile", "legacy-image-v1") != new.get("media_profile", "legacy-image-v1")
|
||||
):
|
||||
_ensure_pool_idle(worker_id)
|
||||
elif new is not None and new.get("type") == "serverless":
|
||||
secret_ids.add(worker_id)
|
||||
@@ -401,7 +411,7 @@ async def _update_worker_endpoint(request, secret_changes):
|
||||
candidate.setdefault("provider", "runpod")
|
||||
candidate.setdefault("enabled", False)
|
||||
# Legacy "concurrency" is accepted (see _SERVERLESS_CONFIG_FIELDS) but ignored.
|
||||
for field in ("provider", "endpoint_id", "jobs_per_run", "timeout_seconds", "enabled", "name"):
|
||||
for field in ("provider", "endpoint_id", "jobs_per_run", "timeout_seconds", "enabled", "name", "media_profile"):
|
||||
if field in data:
|
||||
candidate[field] = data[field]
|
||||
candidate = _clean_serverless_worker(candidate)
|
||||
@@ -424,8 +434,12 @@ async def _update_worker_endpoint(request, secret_changes):
|
||||
current_type == "serverless"
|
||||
and existing.get("endpoint_id") != validated.get("endpoint_id")
|
||||
)
|
||||
media_profile_changed = (
|
||||
current_type == "serverless"
|
||||
and existing.get("media_profile", "legacy-image-v1") != validated.get("media_profile", "legacy-image-v1")
|
||||
)
|
||||
if current_type == "serverless" and (
|
||||
endpoint_changed or next_secret != prior_secret
|
||||
endpoint_changed or media_profile_changed or next_secret != prior_secret
|
||||
):
|
||||
_ensure_pool_idle(worker_id)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -18,7 +18,7 @@ class AmbiguousSubmission(RunpodError):
|
||||
"""Submission may have been accepted: never retry it automatically."""
|
||||
|
||||
|
||||
def serialize_request(workflow, images, timeout_seconds):
|
||||
def serialize_request(workflow, images, timeout_seconds, *, media=None):
|
||||
"""Serialize and size-check the exact JSON body used by Runpod ``/run``."""
|
||||
payload = {
|
||||
"input": {"workflow": workflow, "images": images},
|
||||
@@ -27,6 +27,10 @@ def serialize_request(workflow, images, timeout_seconds):
|
||||
"ttl": timeout_seconds * 1000,
|
||||
},
|
||||
}
|
||||
if media is not None:
|
||||
from ...serverless_bridge.handler import _validate_request
|
||||
payload["input"]["media"] = media
|
||||
_validate_request(payload["input"])
|
||||
body = json.dumps(
|
||||
payload,
|
||||
allow_nan=False,
|
||||
@@ -108,8 +112,8 @@ class RunpodClient:
|
||||
async def health(self):
|
||||
return await self._request("GET", "health")
|
||||
|
||||
async def submit(self, workflow, images, timeout_seconds):
|
||||
body = serialize_request(workflow, images, timeout_seconds)
|
||||
async def submit(self, workflow, images, timeout_seconds, *, media=None):
|
||||
body = serialize_request(workflow, images, timeout_seconds, media=media)
|
||||
result = await self._request("POST", "run", request_body=body)
|
||||
try:
|
||||
return checked_id(result.get("id"))
|
||||
|
||||
@@ -0,0 +1,199 @@
|
||||
"""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_bridge.image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
|
||||
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", {})
|
||||
validate_options(inputs.get("transfer_format", DEFAULT_FORMAT),
|
||||
inputs.get("transfer_quality", DEFAULT_QUALITY))
|
||||
for field in ("multi_job_id", "worker_id"):
|
||||
_validate_identifier(inputs.get(field), f"collector {collector_id!r} {field}")
|
||||
connected = [socket for socket in SOCKET_KINDS if socket in inputs]
|
||||
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]),
|
||||
},
|
||||
}
|
||||
if kind == "IMAGE":
|
||||
fmt, quality = validate_options(inputs.get("transfer_format", DEFAULT_FORMAT),
|
||||
inputs.get("transfer_quality", DEFAULT_QUALITY))
|
||||
if fmt != DEFAULT_FORMAT:
|
||||
exports[node_id]["inputs"].update(transfer_format=fmt, transfer_quality=quality)
|
||||
validate_bindings(bindings)
|
||||
return bindings, exports
|
||||
|
||||
|
||||
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
|
||||
@@ -146,6 +146,10 @@ class ServerlessManager:
|
||||
|
||||
async def start(self, execution_id, prompt_id, prepared_jobs):
|
||||
self.ensure_capacity(len(prepared_jobs))
|
||||
from .serverless_typed_results import validate_prepared_job
|
||||
# Validate the entire batch before creating even the first local slot.
|
||||
for prepared in prepared_jobs:
|
||||
validate_prepared_job(prepared)
|
||||
new_runs = []
|
||||
# Journal every local slot before starting any billed submission.
|
||||
for prepared in prepared_jobs:
|
||||
@@ -177,14 +181,30 @@ class ServerlessManager:
|
||||
return any(run["worker"]["id"] == participant_id and
|
||||
any(m["job_id"] == collector_job_id for m in run["manifest"]) for run in self.runs.values())
|
||||
|
||||
async def _deliver(self, run, images=None, error=None):
|
||||
async with self.server.distributed_jobs_lock:
|
||||
async def _deliver(self, run, images=None, error=None, *, media=None):
|
||||
if media is not None and not error:
|
||||
# Validate all destinations before exposing any partial success.
|
||||
for binding in run["manifest"]:
|
||||
if binding["participant_id"] != run["worker"]["id"]:
|
||||
raise ValueError("Serverless media participant identity mismatch")
|
||||
values = media[binding["job_id"]][binding["socket"]]
|
||||
if not isinstance(values, list) or not values:
|
||||
raise ValueError("Missing serverless media values")
|
||||
async with self.server.distributed_jobs_lock:
|
||||
delivered = set()
|
||||
for binding in run["manifest"]:
|
||||
job_id = binding["job_id"]
|
||||
if error or media is not None:
|
||||
if job_id in delivered:
|
||||
continue
|
||||
delivered.add(job_id)
|
||||
queue = self.server.distributed_pending_jobs.get(binding["job_id"])
|
||||
if queue is None:
|
||||
continue # A cancelled/finished collector must never be resurrected.
|
||||
if error:
|
||||
queue.put_nowait({"worker_id": run["worker"]["id"], "serverless_error": error, "is_last": True})
|
||||
elif media is not None:
|
||||
queue.put_nowait({"worker_id": run["worker"]["id"], "media": media[job_id], "is_last": True})
|
||||
else:
|
||||
batch = images[binding["job_id"]]
|
||||
for index, tensor in enumerate(batch):
|
||||
@@ -218,13 +238,18 @@ class ServerlessManager:
|
||||
submitted = False
|
||||
provider_terminal = False
|
||||
try:
|
||||
from .serverless_typed_results import validate_prepared_job, collect_typed_results
|
||||
typed = validate_prepared_job(run)
|
||||
async with asyncio.timeout(max(0, run["deadline"] - time.monotonic())):
|
||||
async with semaphore:
|
||||
await self._reserve_provider_capacity(run)
|
||||
client = await self._client(run["worker"])
|
||||
run["client"] = client
|
||||
submitted = True
|
||||
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
|
||||
if typed:
|
||||
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"], media=run["media"])
|
||||
else:
|
||||
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
|
||||
run["provider_job_id"] = job_id
|
||||
# Persist immediately: all subsequent actions use this exact ID.
|
||||
self.ledger.update(rid, job_id=job_id, status="IN_QUEUE", updated_at=time.time())
|
||||
@@ -240,8 +265,12 @@ class ServerlessManager:
|
||||
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
|
||||
raise RunpodError("Serverless workflow reported an output error; inspect the provider job")
|
||||
self.ledger.update(rid, status="DOWNLOADING", updated_at=time.time())
|
||||
images = await decode_results(output, run["manifest"])
|
||||
await self._deliver(run, images=images)
|
||||
if typed:
|
||||
media = await collect_typed_results(output, run["media"], run["manifest"])
|
||||
await self._deliver(run, media=media)
|
||||
else:
|
||||
images = await decode_results(output, run["manifest"])
|
||||
await self._deliver(run, images=images)
|
||||
self.ledger.update(rid, status="COMPLETED", updated_at=time.time())
|
||||
return
|
||||
self.ledger.update(rid, status=status, updated_at=time.time())
|
||||
@@ -276,6 +305,7 @@ class ServerlessManager:
|
||||
# base64 inputs, workflow metadata, or a client holding credentials.
|
||||
run.pop("workflow", None)
|
||||
run.pop("images", None)
|
||||
run.pop("media", None)
|
||||
run["client"] = None
|
||||
# A released local semaphore alone cannot release a durable reservation.
|
||||
await self._notify_capacity()
|
||||
@@ -359,6 +389,7 @@ class ServerlessManager:
|
||||
self.ledger.update(rid, status="CANCELLED", updated_at=time.time())
|
||||
run.pop("workflow", None)
|
||||
run.pop("images", None)
|
||||
run.pop("media", None)
|
||||
await self._deliver(run, error="Serverless job cancelled before submission")
|
||||
continue
|
||||
# Do not send two cancellation requests for the same live job in
|
||||
|
||||
@@ -0,0 +1,225 @@
|
||||
"""Bounded local native inputs for comfy-media-v1 requests.
|
||||
|
||||
Only stock LoadAudio/LoadVideo schemas are adapted. Original nodes and links
|
||||
remain in the private request graph for endpoint validation; only the existing
|
||||
image preparer rewrites filenames. No upload, transcode, resample or provider
|
||||
operation belongs here. The caller still enforces the complete request-body cap.
|
||||
"""
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from ...serverless_bridge import codec, contract, video_codec
|
||||
from . import serverless_media
|
||||
|
||||
|
||||
_LOADERS = {"LoadAudio": ("audio", "AUDIO"), "LoadVideo": ("file", "VIDEO")}
|
||||
Error = serverless_media.ServerlessMediaError
|
||||
|
||||
|
||||
def _loader_references(workflow):
|
||||
if not isinstance(workflow, dict):
|
||||
raise Error("Serverless workflow must be an API graph")
|
||||
loaders = {}
|
||||
for node_id, node in workflow.items():
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
class_type = node.get("class_type")
|
||||
if not isinstance(class_type, str) or class_type not in _LOADERS:
|
||||
continue
|
||||
contract.identifier(node_id)
|
||||
key, kind = _LOADERS[class_type]
|
||||
inputs = node.get("inputs")
|
||||
if not isinstance(inputs, dict) or set(inputs) != {key}:
|
||||
raise Error(f"{class_type} inputs must contain only its native filename; extra controls are unsupported")
|
||||
if not isinstance(inputs[key], str):
|
||||
raise Error(f"{class_type} requires a literal filename; linked inputs are unsupported")
|
||||
normalized, _suffix = serverless_media._normalize_media_reference(inputs[key])
|
||||
loaders[node_id] = (kind, normalized)
|
||||
|
||||
references = []
|
||||
used = set()
|
||||
for node_id, node in workflow.items():
|
||||
if not isinstance(node, dict) or not isinstance(node.get("inputs"), dict):
|
||||
continue
|
||||
for key, value in node["inputs"].items():
|
||||
if not isinstance(value, (list, tuple)) or not value or not isinstance(value[0], str):
|
||||
continue
|
||||
if value[0] not in loaders:
|
||||
continue
|
||||
if type(value) is not list or len(value) != 2 or type(value[1]) is not int or value[1] != 0:
|
||||
raise Error("Native media loader links require output slot 0")
|
||||
contract.identifier(node_id)
|
||||
contract.identifier(key)
|
||||
references.append((node_id, key, value[0]))
|
||||
used.add(value[0])
|
||||
if used != set(loaders):
|
||||
raise Error("Native media loader has no consumer")
|
||||
return loaders, references
|
||||
|
||||
|
||||
def _source_path(value, folder_paths):
|
||||
"""Resolve native annotations; authorization/opening is owned by _read_path."""
|
||||
annotation = serverless_media._MEDIA_ANNOTATION.search(value)
|
||||
unannotated = value[:annotation.start()] if annotation else value
|
||||
if Path(unannotated).is_absolute():
|
||||
return unannotated
|
||||
try:
|
||||
return folder_paths.get_annotated_filepath(value)
|
||||
except Exception:
|
||||
raise Error("Native media input path is missing or invalid") from None
|
||||
|
||||
|
||||
def _decode_audio(data, limit):
|
||||
"""Match core nodes_audio.load, bounding each frame *before* accumulation.
|
||||
|
||||
Float32 output bytes are charged even for narrow integer PCM. Native float64
|
||||
remains float64, as in f32_pcm (there is no silent dtype conversion). Raw
|
||||
retained arrays and each live AV plane allocation have a separate <= limit
|
||||
guard, so wider/padded raw frames cannot bypass the output budget. PyAV owns
|
||||
the one frame being decoded; it is checked before to_ndarray or torch.cat.
|
||||
Concatenation/conversion use bounded temporary copies, not an RSS guarantee.
|
||||
"""
|
||||
if limit <= 0:
|
||||
raise Error("AUDIO decoded input exceeds byte limit")
|
||||
try:
|
||||
import av
|
||||
import torch
|
||||
from comfy_extras.nodes_audio import f32_pcm
|
||||
except ImportError:
|
||||
raise Error("Native AUDIO decoding is unavailable") from None
|
||||
|
||||
try:
|
||||
with av.open(io.BytesIO(data)) as container:
|
||||
if not container.streams.audio:
|
||||
raise Error("Native AUDIO input has no audio stream")
|
||||
stream = container.streams.audio[0]
|
||||
rate = stream.codec_context.sample_rate
|
||||
channels = stream.channels
|
||||
if type(rate) is not int or rate <= 0 or type(channels) is not int or channels <= 0:
|
||||
raise Error("Native AUDIO stream shape or sample rate is invalid")
|
||||
frames = []
|
||||
output_bytes = raw_bytes = 0
|
||||
dtype = None
|
||||
for frame in container.decode(streams=stream.index):
|
||||
samples = frame.samples
|
||||
if (type(samples) is not int or samples <= 0 or frame.sample_rate != rate
|
||||
or len(frame.layout.channels) != channels):
|
||||
raise Error("Native AUDIO frame shape or sample rate changed")
|
||||
if frame.format.name not in {"s16", "s16p", "s32", "s32p", "flt", "fltp", "dbl", "dblp"}:
|
||||
raise Error("Unsupported native AUDIO PCM format")
|
||||
itemsize = frame.format.bytes
|
||||
count = samples * channels
|
||||
# Native f32_pcm promotes integer PCM to float32, not float64.
|
||||
decoded_bytes = count * max(4, itemsize)
|
||||
frame_bytes = count * itemsize
|
||||
plane_bytes = sum(plane.buffer_size for plane in frame.planes)
|
||||
if (decoded_bytes > limit - output_bytes or frame_bytes > limit - raw_bytes
|
||||
or plane_bytes > limit - raw_bytes):
|
||||
raise Error("AUDIO decoded input or raw frame exceeds byte limit")
|
||||
output_bytes += decoded_bytes
|
||||
raw_bytes += frame_bytes
|
||||
array = frame.to_ndarray()
|
||||
if array.ndim != 2 or array.size != count or array.nbytes != frame_bytes:
|
||||
raise Error("Native AUDIO frame shape is invalid")
|
||||
buf = torch.from_numpy(array)
|
||||
if buf.dtype not in (torch.int16, torch.int32, torch.float32, torch.float64):
|
||||
raise Error("Unsupported native AUDIO PCM dtype")
|
||||
if dtype is not None and buf.dtype != dtype:
|
||||
raise Error("Native AUDIO frame dtype changed")
|
||||
dtype = buf.dtype
|
||||
if buf.shape[0] != channels:
|
||||
buf = buf.view(-1, channels).t()
|
||||
if tuple(buf.shape) != (channels, samples):
|
||||
raise Error("Native AUDIO frame shape is invalid")
|
||||
frames.append(buf)
|
||||
if not frames:
|
||||
raise Error("Native AUDIO input has no decoded frames")
|
||||
waveform = f32_pcm(torch.cat(frames, dim=1)).unsqueeze(0)
|
||||
return {"waveform": waveform, "sample_rate": rate}
|
||||
except Error:
|
||||
raise
|
||||
except Exception:
|
||||
raise Error("Native AUDIO input could not be decoded safely") from None
|
||||
|
||||
|
||||
def _artifact_id(kind, descriptor, payload):
|
||||
digest = hashlib.sha256()
|
||||
digest.update(kind.encode("ascii") + b"\0")
|
||||
# Include interpretation (not just waveform bytes), notably sample_rate.
|
||||
digest.update(json.dumps(descriptor, sort_keys=True, separators=(",", ":"), allow_nan=False).encode())
|
||||
digest.update(b"\0")
|
||||
digest.update(payload)
|
||||
return "input-" + kind.lower() + "-" + digest.hexdigest()
|
||||
|
||||
|
||||
async def prepare_media_inputs(workflow: dict) -> tuple[list, list, list]:
|
||||
"""Return legacy images, inline typed artifacts, and every consumer reference.
|
||||
|
||||
Paths are read once into bounded bytes through the video adapter's no-follow
|
||||
directory-fd reader. Decoders receive BytesIO, never a second path open.
|
||||
Source files and unique decoded artifacts have independent aggregate caps.
|
||||
Repeated content of the same kind is decoded once and charged once.
|
||||
"""
|
||||
loaders, references = _loader_references(workflow)
|
||||
image_view = {node_id: node for node_id, node in workflow.items() if node_id not in loaders}
|
||||
# Retain the strict scan for every other loader before expensive decoding.
|
||||
serverless_media._iter_image_references(image_view)
|
||||
if not loaders:
|
||||
return await serverless_media.prepare_images(image_view), [], []
|
||||
try:
|
||||
import folder_paths
|
||||
except ImportError:
|
||||
raise Error("ComfyUI media directories are unavailable") from None
|
||||
roots = list(serverless_media._folder_roots(folder_paths).values())
|
||||
artifacts = {}
|
||||
source_artifacts = {}
|
||||
source_digests = set()
|
||||
filenames = {}
|
||||
loader_artifacts = {}
|
||||
file_total = decoded_total = 0
|
||||
for node_id, (kind, filename) in loaders.items():
|
||||
# Avoid retaining file buffers merely to deduplicate repeat references.
|
||||
if (kind, filename) in filenames:
|
||||
loader_artifacts[node_id] = filenames[kind, filename]
|
||||
continue
|
||||
path = _source_path(filename, folder_paths)
|
||||
data = video_codec._read_path(path, contract.MAX_ARTIFACT_BYTES, roots)
|
||||
source_digest = hashlib.sha256(data).hexdigest()
|
||||
if source_digest not in source_digests:
|
||||
if len(data) > contract.MAX_TOTAL_BYTES - file_total:
|
||||
raise Error("Native media source files exceed total byte limit")
|
||||
file_total += len(data)
|
||||
source_digests.add(source_digest)
|
||||
source_key = (kind, source_digest)
|
||||
if source_key not in source_artifacts:
|
||||
if len(artifacts) >= contract.MAX_ARTIFACTS:
|
||||
raise Error("Native media input artifact count exceeds limit")
|
||||
limit = min(contract.MAX_ARTIFACT_BYTES, contract.MAX_TOTAL_BYTES - decoded_total)
|
||||
if kind == "AUDIO":
|
||||
value = _decode_audio(data, limit)
|
||||
descriptor, payload = codec.encode_payload(kind, value, max_bytes=limit)
|
||||
del value
|
||||
else:
|
||||
from comfy_api.latest import InputImpl
|
||||
value = InputImpl.VideoFromFile(io.BytesIO(data))
|
||||
descriptor, payload = video_codec.encode_video(value, max_bytes=limit, allowed_roots=roots)
|
||||
del value
|
||||
aid = _artifact_id(kind, descriptor, payload)
|
||||
if aid not in artifacts:
|
||||
if len(payload) > limit:
|
||||
raise Error("Native media decoded inputs exceed total byte limit")
|
||||
artifacts[aid] = contract.artifact_from_payload(
|
||||
aid, kind, descriptor, payload, max_bytes=contract.MAX_ARTIFACT_BYTES)
|
||||
decoded_total += len(payload)
|
||||
source_artifacts[source_key] = aid
|
||||
del payload
|
||||
aid = source_artifacts[source_key]
|
||||
loader_artifacts[node_id] = filenames[kind, filename] = aid
|
||||
del data
|
||||
refs = [{"node_id": node_id, "input_name": key, "artifact_id": loader_artifacts[source]}
|
||||
for node_id, key, source in references]
|
||||
# Shared node dictionaries deliberately propagate only legacy image rewrites.
|
||||
images = await serverless_media.prepare_images(image_view)
|
||||
return images, list(artifacts.values()), refs
|
||||
@@ -0,0 +1,79 @@
|
||||
"""Explicit media-v1 preflight and atomic, lossless result materialization."""
|
||||
from urllib.parse import urljoin
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ...serverless_bridge import contract
|
||||
from ...serverless_bridge.contract import MAX_ARTIFACT_BYTES, PROFILE, validate_bindings
|
||||
from .serverless_media import (
|
||||
MAX_REDIRECTS, _CHUNK_SIZE, _DOWNLOAD_TIMEOUT_SECONDS,
|
||||
_PublicResolver, _validate_artifact_url,
|
||||
)
|
||||
|
||||
|
||||
def validate_prepared_job(prepared):
|
||||
"""Fail before journalling/submission; legacy jobs keep their original path."""
|
||||
profile = prepared["worker"].get("media_profile", "legacy-image-v1")
|
||||
media = prepared.get("media")
|
||||
if profile == "legacy-image-v1" and media is None:
|
||||
return False
|
||||
if profile != PROFILE or media is None:
|
||||
raise ValueError("Serverless media profile/request mismatch")
|
||||
from .runpod_client import serialize_request
|
||||
# The shared endpoint validator also checks exporter controls and input refs.
|
||||
serialize_request(prepared["workflow"], prepared["images"],
|
||||
prepared["worker"]["timeout_seconds"], media=media)
|
||||
expected = validate_bindings(prepared["manifest"])
|
||||
if expected != validate_bindings(media["outputs"]):
|
||||
raise ValueError("Serverless media manifest identity mismatch")
|
||||
if any(binding["participant_id"] != prepared["worker"]["id"] for binding in expected.values()):
|
||||
raise ValueError("Serverless media participant identity mismatch")
|
||||
return True
|
||||
|
||||
|
||||
async def download_public_artifact(url, *, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
"""Contract-sized public HTTPS download, never using a provider session/key."""
|
||||
if type(max_bytes) is not int or not 0 <= max_bytes <= MAX_ARTIFACT_BYTES:
|
||||
raise ValueError("Media artifact exceeds the permitted size")
|
||||
current_url = _validate_artifact_url(url)
|
||||
connector = aiohttp.TCPConnector(resolver=_PublicResolver(), use_dns_cache=False, limit=4)
|
||||
timeout = aiohttp.ClientTimeout(total=_DOWNLOAD_TIMEOUT_SECONDS, connect=10, sock_read=15)
|
||||
try:
|
||||
async with aiohttp.ClientSession(connector=connector, timeout=timeout, trust_env=False,
|
||||
cookie_jar=aiohttp.DummyCookieJar()) as session:
|
||||
for redirect_count in range(MAX_REDIRECTS + 1):
|
||||
async with session.get(current_url, allow_redirects=False) as response:
|
||||
if response.status in {301, 302, 303, 307, 308}:
|
||||
location = response.headers.get("Location")
|
||||
if not location or redirect_count >= MAX_REDIRECTS:
|
||||
raise ValueError("Media artifact redirect is invalid")
|
||||
current_url = _validate_artifact_url(urljoin(current_url, location))
|
||||
continue
|
||||
if response.status != 200:
|
||||
raise ValueError("Media artifact download failed")
|
||||
payload = bytearray()
|
||||
async for chunk in response.content.iter_chunked(_CHUNK_SIZE):
|
||||
if not isinstance(chunk, bytes) or len(payload) + len(chunk) > max_bytes:
|
||||
raise ValueError("Media artifact exceeds the permitted size")
|
||||
payload.extend(chunk)
|
||||
return bytes(payload)
|
||||
except Exception:
|
||||
# Signed URLs and transport exception details must not reach the ledger.
|
||||
raise ValueError("Media artifact download failed or exceeded its byte limit") from None
|
||||
raise ValueError("Media artifact download failed")
|
||||
|
||||
|
||||
async def collect_typed_results(output, media, manifest, *, download=None):
|
||||
"""Decode every binding before returning any collector's native value lists."""
|
||||
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
|
||||
raise ValueError("Invalid serverless media output")
|
||||
decoded = await contract.decode_response(output.get("media"), request_id=media["request_id"],
|
||||
bindings=manifest,
|
||||
download=download if download is not None else download_public_artifact)
|
||||
grouped = {}
|
||||
for binding in manifest:
|
||||
values = decoded[binding["binding_id"]]
|
||||
if not isinstance(values, list) or not values:
|
||||
raise ValueError("Missing serverless media values")
|
||||
grouped.setdefault(binding["job_id"], {}).setdefault(binding["socket"], []).extend(values)
|
||||
return grouped
|
||||
+100
-10
@@ -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"})
|
||||
|
||||
@@ -0,0 +1,158 @@
|
||||
# Typed serverless media (PR2)
|
||||
|
||||
`comfy-media-v1` adds typed list results, exact by default, for `IMAGE`, `AUDIO`, `MASK`,
|
||||
`STRING`, and native ComfyUI `VIDEO`. It is an **explicit endpoint profile**, not
|
||||
an automatically detected capability. Omitting `media_profile` still selects
|
||||
`legacy-image-v1`; existing image collectors and handlers keep their old contract.
|
||||
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.
|
||||
|
||||
## Image return format and quality
|
||||
|
||||
`transfer_format` defaults to **`exact`** and `transfer_quality` to **90**.
|
||||
These controls affect **remote IMAGE returns only**, not master-local images,
|
||||
uploaded inputs, MASK, AUDIO, STRING or native VIDEO values.
|
||||
|
||||
- `exact`: preserve floating-point dtype and values, with no image quantization.
|
||||
- `legacy_png`, `png`, `ffv1`: lossless **8-bit image pixels**, not exact float
|
||||
tensors. Image values are clipped to [0,1] and quantized as in the existing
|
||||
Collector. `legacy_png` uses PNG without compression; `png` compresses it.
|
||||
- `jpeg`, `webp`: lossy still-image compression.
|
||||
- `h264`, `h265`, `av1`: lossy video encoding of each IMAGE tensor's frame batch.
|
||||
This does not transcode the native `video` socket or combine separate list items.
|
||||
- Quality 1–100 affects lossy formats only, using the existing Collector's
|
||||
quality-to-CRF mapping for video. Dimensions, frame order and channel count
|
||||
are preserved; alpha is carried separately where needed, not silently dropped.
|
||||
|
||||
Both master and endpoint must have this bridge revision for compressed image
|
||||
returns. Pillow handles still formats and CPU/software PyAV encoders handle video;
|
||||
unavailable codecs fail clearly rather than silently changing the selection.
|
||||
Compression can reduce result egress and inline response size, but does not relax
|
||||
encoded/decoded byte budgets or the input request limit. Compressed IMAGE batches
|
||||
are capped at 240 frames and 8192 pixels per dimension, with a conservative
|
||||
256 MiB aggregate decode budget (including codec scratch and alpha). The original Collector's
|
||||
settings, default and fallback behavior are unchanged.
|
||||
|
||||
## Inputs and fidelity
|
||||
|
||||
Automatic local-file transfer supports stock `LoadImage`, plus explicit stock
|
||||
`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.
|
||||
@@ -0,0 +1,229 @@
|
||||
"""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 ..serverless_bridge.image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY,
|
||||
FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options)
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
|
||||
POLL_SECONDS = 0.1
|
||||
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)],
|
||||
io.Combo.Input("transfer_format", options=list(FORMATS), default=DEFAULT_FORMAT,
|
||||
optional=True, tooltip=FORMAT_TOOLTIP),
|
||||
io.Int.Input("transfer_quality", default=DEFAULT_QUALITY, min=1, max=100,
|
||||
optional=True, tooltip=QUALITY_TOOLTIP)],
|
||||
outputs=[kind.Output(display_name=socket, is_output_list=True)
|
||||
for socket, kind in zip(SOCKET_KINDS, types)],
|
||||
is_input_list=True, accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@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='[]', transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
|
||||
validate_options(_control(transfer_format, 'transfer_format', str),
|
||||
_control(transfer_quality, 'transfer_quality', int))
|
||||
media = dict(images=images, audio=audio, masks=masks, texts=texts, video=video)
|
||||
controls = dict(load_balance=load_balance, is_worker=is_worker,
|
||||
enabled_worker_ids=enabled_worker_ids, delegate_only=delegate_only,
|
||||
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)
|
||||
+2
-1
@@ -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]
|
||||
|
||||
@@ -0,0 +1,251 @@
|
||||
# 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 `<ComfyUI>/custom_nodes/serverless_bridge`.
|
||||
Use a compatible ComfyUI with `comfy_api.v0_0_2` and the native VIDEO classes
|
||||
used by [video_codec.py](video_codec.py). These exact-class/private-state
|
||||
adapters are tied to the supported core implementation, not arbitrary VIDEO
|
||||
plugins or future core layouts.
|
||||
2. Restart ComfyUI so `comfy_entrypoint()` registers the two nodes. Keep the
|
||||
workflow's own models and custom nodes installed separately.
|
||||
3. Make `<ComfyUI>/custom_nodes` importable by the handler process (for example,
|
||||
via its `PYTHONPATH`). The handler uses `aiohttp`; the nodes/codecs use
|
||||
ComfyUI's PyTorch/native media dependencies. The handler module itself does
|
||||
not import ComfyUI, PyTorch, Runpod, or Distributed bootstrap code.
|
||||
4. Give the handler filesystem access to **the same input directory** used by
|
||||
that ComfyUI instance, plus its trusted HTTP URL. These are operator settings,
|
||||
never request-controlled paths/URLs. Staging/cleanup requires POSIX directory
|
||||
descriptor and no-follow file operations.
|
||||
|
||||
Installing the nodes alone is insufficient: the endpoint must route media-v1
|
||||
inputs through this executor and implement the uncertainty policy below. A
|
||||
legacy image handler cannot interpret the typed envelope.
|
||||
|
||||
## Callable integration
|
||||
|
||||
```python
|
||||
from serverless_bridge.handler import execute_media_request
|
||||
|
||||
|
||||
async def execute_endpoint_input(
|
||||
input_data, *, comfy_url, input_dir, temp_dir=None,
|
||||
download=None, publish=None,
|
||||
):
|
||||
# Invoke only under a supervisor implementing the uncertainty policy below.
|
||||
# Exceptions intentionally propagate; this is NOT a provider worker loop.
|
||||
return await execute_media_request(
|
||||
input_data,
|
||||
comfy_url=comfy_url,
|
||||
input_dir=input_dir,
|
||||
timeout_seconds=1800,
|
||||
session=None,
|
||||
download=download,
|
||||
publish=publish,
|
||||
temp_dir=temp_dir,
|
||||
)
|
||||
```
|
||||
|
||||
Pass the provider job's **`input` object**, not the whole job. Return the resulting
|
||||
`{"media": ...}` object as provider output; Runpod exposes it as `output.media`.
|
||||
Do not add another `output` wrapper inside it. `timeout_seconds` must be finite
|
||||
and positive. If supplied, `session` is a caller-owned `aiohttp.ClientSession`
|
||||
for local Comfy HTTP and remains open; otherwise the executor creates/closes its
|
||||
own session. Input dictionaries are copied before graph rewriting.
|
||||
|
||||
**This callable is not a turnkey Runpod handler.** An outer SDK that catches every
|
||||
exception and immediately reuses the worker is unsafe. The supervisor must
|
||||
recognize `serverless_bridge.handler.MediaExecutionUncertain`, stop intake, and
|
||||
quarantine/reconcile or terminate the owning environment before reuse. Merely
|
||||
returning `{"error": ...}` or marking a provider job `FAILED` is not that policy.
|
||||
|
||||
## Request and response
|
||||
|
||||
A minimal 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": "<base64>" }`, must match stock `LoadImage` references, and are rewritten
|
||||
to owned staging paths. They are not an arbitrary filename-upload interface.
|
||||
|
||||
The returned `media` object has exactly `version`, `profile`, `request_id`,
|
||||
`bindings`, `artifacts`, and `items`:
|
||||
|
||||
- `bindings`: every declared binding, unchanged, plus `count` (1–64).
|
||||
- `artifacts`: the same artifact schema above with inline or URL transport.
|
||||
- `items`: `{ "binding_id": "...", "item_index": 0, "artifact_id": "..." }`
|
||||
records covering every index exactly once, ordered by `item_index` when decoded.
|
||||
|
||||
All identity, completeness, type, size, and checksum checks must pass; do not
|
||||
consume a partial result. The internal history transport `comfy-file` is local
|
||||
only and must not be returned to the master. The executor reads validated
|
||||
`media_exports` history records and bounded `/view` payloads before constructing
|
||||
the external envelope. [contract.py](contract.py) is the schema authority.
|
||||
|
||||
## Optional image return encoding
|
||||
|
||||
IMAGE `MediaExportV1` nodes may also contain literal `transfer_format` and
|
||||
`transfer_quality` controls. Defaults are `exact` and `90`; old exact requests
|
||||
omit the controls and retain their original representation. Other media exports
|
||||
must not carry image controls. The master compiler supplies them from the typed
|
||||
Collector's widgets; callers should not alter binding identity metadata.
|
||||
|
||||
Formats: `exact`, `legacy_png`, `png`, `ffv1`, `jpeg`, `webp`, `h264`, `h265`, `av1`.
|
||||
`exact` preserves float tensor bits. The image/video codecs quantize RGB/grayscale
|
||||
to 8-bit; PNG/FFV1 are lossless at that precision, while the other codecs are
|
||||
lossy. Quality is an integer 1–100 and affects only lossy formats. IMAGE batch
|
||||
order, shape and channel count are retained, and alpha is not discarded.
|
||||
Native VIDEO/MASK/AUDIO/STRING still use their original exact codecs.
|
||||
|
||||
Copy/update the complete bridge directory on both sides, including
|
||||
`image_transfer.py` and `image_codec.py`. Pillow and the selected CPU/software
|
||||
PyAV codec must be available; missing encoders fail without automatic fallback.
|
||||
Compressed IMAGE artifacts use `image-transfer-v1`; decode through the shared
|
||||
`codec.decode_payload`, which enforces payload and decoded-shape limits.
|
||||
See [image transfer choices](../docs/serverless-media.md#image-return-format-and-quality).
|
||||
|
||||
## Transport budgets and ownership
|
||||
|
||||
- Contract maxima: 256 MiB per artifact, 512 MiB aggregate artifact bytes,
|
||||
128 artifacts, 64 items per binding. Request input totals include decoded
|
||||
legacy-image bytes. Decoders validate tensor/VIDEO sizes and layouts before
|
||||
allocating native storage; these are payload limits, **not peak-RSS guarantees**.
|
||||
- The built-in master currently emits inline inputs only. Its complete Runpod
|
||||
request JSON is capped at **10,000,000 bytes**, including base64, graph and
|
||||
policy. The larger contract limits do not bypass this cap; there is no
|
||||
built-in large-input/object-storage uploader.
|
||||
- The executor's result JSON is capped at **10,000,000 bytes**. Inline results
|
||||
are the default. When needed, `publish` is called for selected larger
|
||||
artifacts to fit that budget; supplying it does not force all results to URLs.
|
||||
Without a publisher, an oversized response fails instead of being truncated.
|
||||
- Optional **async** `download(url, *, max_bytes)` returns `bytes`. The executor
|
||||
supplies the declared artifact length and checks exact length and SHA-256.
|
||||
The callback must enforce public HTTPS, DNS/IP restrictions, redirect checks,
|
||||
and bounded streaming **before allocation**; syntactic contract validation
|
||||
does not provide SSRF protection. No downloader is installed by default.
|
||||
- Optional **async** `publish(metadata, payload)` uploads the exact `bytes` and
|
||||
returns exactly `{"type": "url", "url": "https://...", "expires_at": "..."}`.
|
||||
`metadata` is the artifact record without `transport`. Expiry must be a future,
|
||||
timezone-aware ISO timestamp; URLs use HTTPS port 443, no userinfo/fragment,
|
||||
and at most 8192 characters. Keep objects available long enough for collection.
|
||||
The callback/operator owns upload credentials, object expiry and cleanup,
|
||||
including orphan uploads after failure. There is no bridge storage backend.
|
||||
- The master uses a separate bounded public-HTTPS downloader without provider
|
||||
credentials. Never send provider credentials to artifact hosts or log signed
|
||||
URLs. Endpoint callbacks must enforce equivalent protections themselves.
|
||||
|
||||
Callbacks must cooperate with async cancellation, finish/abort their own I/O,
|
||||
and not detach readers/uploads that outlive the callback. Synchronous code that
|
||||
blocks the event loop cannot be hard-timed-out by this executor.
|
||||
|
||||
## Cancellation, quarantine, and file retention
|
||||
|
||||
`timeout_seconds` covers staging, submission, polling, downloads and publication.
|
||||
On failure/cancellation, there is a separate **10-second callback-drain grace**
|
||||
and, when needed, a **10-second prompt-reconciliation grace**. These are recovery
|
||||
allowances, not evidence that execution stopped. Ambiguous `/prompt` submission
|
||||
is never retried. An identified pending prompt may be deleted from `/queue`,
|
||||
with exact read-back; running prompts are never sent `/interrupt`. The pinned
|
||||
core route has a snapshot/global-interrupt race even when given a prompt ID.
|
||||
|
||||
`MediaExecutionUncertain` means prompt acceptance/termination **or callback
|
||||
teardown** is unconfirmed. Even successful terminal Comfy history is insufficient
|
||||
if an async downloader/publisher ignores cancellation. The bridge retains input
|
||||
and exporter files in that case; late callback exit does not automatically
|
||||
remove retained files. Before accepting another job on that environment:
|
||||
|
||||
1. Quarantine it and stop new intake; do not automatically resubmit the request.
|
||||
2. Reconcile the local prompt and all readers, or terminate the owning execution
|
||||
environment when their termination cannot be established.
|
||||
3. Only then clean retained files and allow reuse/replacement. A provider terminal
|
||||
job status alone cannot certify a surviving local Comfy process or callback.
|
||||
|
||||
Staged inputs use random exclusive files under
|
||||
`input_dir/distributed-media-v1/<sha256(request_id)>/`. Cleanup removes only files
|
||||
created by that invocation, after confirmed terminal state (or before submission)
|
||||
and reader teardown; it never recursively deletes a shared namespace.
|
||||
|
||||
`temp_dir=None` leaves exporter retention to the Comfy temp-directory/environment
|
||||
lifecycle. Optionally pass the **trusted local Comfy temp root** to enable
|
||||
best-effort cleanup of validated, invocation-owned exporter UUID `.bin` files
|
||||
under the corresponding hashed namespace, only after terminal history and
|
||||
collection/publisher teardown. Symlinked, hard-linked, replaced, missing, or
|
||||
unidentifiable files are not claimed; directories and all unidentified leftovers
|
||||
remain lifecycle-owned. This is not a temp-directory sweep. Independently decoded
|
||||
VIDEO objects do not depend on either staging or exporter files.
|
||||
|
||||
Verification scope: source and offline tests. No live provider/GPU validation is
|
||||
claimed, and this directory is not a tested turnkey provider deployment.
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Standalone typed-media bridge; copy this directory into endpoint custom_nodes.
|
||||
|
||||
Importing codecs does not load ComfyUI, providers or Distributed bootstrap.
|
||||
"""
|
||||
|
||||
async def comfy_entrypoint():
|
||||
from comfy_api.v0_0_2 import ComfyExtension
|
||||
from .nodes import MediaImportV1, MediaExportV1
|
||||
|
||||
class MediaBridgeExtension(ComfyExtension):
|
||||
async def get_node_list(self):
|
||||
return [MediaImportV1, MediaExportV1]
|
||||
|
||||
return MediaBridgeExtension()
|
||||
@@ -0,0 +1,189 @@
|
||||
"""Lossless, bounded CPU payloads shared with the standalone endpoint bridge.
|
||||
|
||||
Descriptors contain only kind-specific metadata. Tensor bytes are contiguous,
|
||||
little-endian floating-point values, never a pickle or an image/audio container.
|
||||
The optional byte budget can lower, but never raise, MAX_PAYLOAD_BYTES.
|
||||
"""
|
||||
|
||||
import ctypes
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
MAX_PAYLOAD_BYTES = 256 * 1024 * 1024
|
||||
|
||||
_DTYPES = {
|
||||
"float16": (torch.float16, 2),
|
||||
"bfloat16": (torch.bfloat16, 2),
|
||||
"float32": (torch.float32, 4),
|
||||
"float64": (torch.float64, 8),
|
||||
}
|
||||
_DTYPE_NAMES = {dtype: name for name, (dtype, _) in _DTYPES.items()}
|
||||
_LAYOUTS = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}
|
||||
|
||||
|
||||
def _validate_kind(kind):
|
||||
if type(kind) is not str or kind not in ("IMAGE", "MASK", "STRING", "AUDIO"):
|
||||
raise ValueError("Unsupported payload kind")
|
||||
|
||||
|
||||
def _byte_limit(max_bytes):
|
||||
if type(max_bytes) is not int or max_bytes < 0:
|
||||
raise ValueError("max_bytes must be a nonnegative integer")
|
||||
return min(max_bytes, MAX_PAYLOAD_BYTES)
|
||||
|
||||
|
||||
def _sample_rate(value):
|
||||
if type(value) is not int or value <= 0:
|
||||
raise ValueError("AUDIO sample_rate must be a positive integer")
|
||||
return value
|
||||
|
||||
|
||||
def _tensor_size(kind, shape, itemsize, limit):
|
||||
if type(shape) is not list or len(shape) != len(_LAYOUTS[kind]):
|
||||
raise ValueError("Tensor shape has an invalid rank")
|
||||
if any(type(dimension) is not int or dimension <= 0 for dimension in shape):
|
||||
raise ValueError("Tensor shape dimensions must be positive integers")
|
||||
if kind == "IMAGE" and shape[-1] not in (1, 3, 4):
|
||||
raise ValueError("IMAGE shape must have 1, 3 or 4 channels")
|
||||
size = itemsize
|
||||
for dimension in shape:
|
||||
# Divide before multiplying, including for arbitrary-size Python ints.
|
||||
if dimension > limit // size:
|
||||
raise ValueError("Tensor payload exceeds byte limit")
|
||||
size *= dimension
|
||||
return size
|
||||
|
||||
|
||||
def _native_little_endian(payload, itemsize):
|
||||
"""Convert native to wire byte order, or vice versa (an involution)."""
|
||||
if sys.byteorder == "little":
|
||||
return payload
|
||||
swapped = bytearray(len(payload))
|
||||
for offset in range(itemsize):
|
||||
swapped[offset::itemsize] = payload[itemsize - 1 - offset::itemsize]
|
||||
return bytes(swapped)
|
||||
|
||||
|
||||
def _encode_string(value, limit):
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("STRING value must be text")
|
||||
# Check the exact UTF-8 size before allocating the encoded payload.
|
||||
if len(value) > limit:
|
||||
raise ValueError("STRING payload exceeds byte limit")
|
||||
size = 0
|
||||
for character in value:
|
||||
codepoint = ord(character)
|
||||
if 0xD800 <= codepoint <= 0xDFFF:
|
||||
raise ValueError("STRING value must be valid UTF-8 text")
|
||||
size += 1 + (codepoint > 0x7F) + (codepoint > 0x7FF) + (codepoint > 0xFFFF)
|
||||
if size > limit:
|
||||
raise ValueError("STRING payload exceeds byte limit")
|
||||
return {"encoding": "utf8-v1"}, value.encode("utf-8")
|
||||
|
||||
|
||||
def encode_payload(kind: str, value, *, max_bytes: int = MAX_PAYLOAD_BYTES) -> tuple[dict, bytes]:
|
||||
"""Encode a typed value without dtype conversion, normalization or GPU use.
|
||||
|
||||
IMAGE/MASK accept dense CPU tensors; AUDIO accepts a waveform CPU tensor
|
||||
plus a positive integer sample_rate. STRING accepts exact Unicode text.
|
||||
Invalid values and exceeded budgets raise nonsensitive ValueError messages.
|
||||
"""
|
||||
_validate_kind(kind)
|
||||
limit = _byte_limit(max_bytes)
|
||||
if kind == "STRING":
|
||||
return _encode_string(value, limit)
|
||||
|
||||
sample_rate = None
|
||||
if kind == "AUDIO":
|
||||
if not isinstance(value, dict) or set(value) != {"waveform", "sample_rate"}:
|
||||
raise ValueError("AUDIO value must contain waveform and sample_rate")
|
||||
sample_rate = _sample_rate(value["sample_rate"])
|
||||
value = value["waveform"]
|
||||
if not isinstance(value, torch.Tensor):
|
||||
raise ValueError("Tensor value must be a torch.Tensor")
|
||||
if value.device.type != "cpu":
|
||||
raise ValueError("Tensor value must be on CPU")
|
||||
if value.layout != torch.strided or value.is_nested:
|
||||
raise ValueError("Tensor value must have a dense strided layout")
|
||||
if value.dtype not in _DTYPE_NAMES:
|
||||
raise ValueError("Unsupported tensor dtype")
|
||||
|
||||
dtype_name = _DTYPE_NAMES[value.dtype]
|
||||
_, itemsize = _DTYPES[dtype_name]
|
||||
shape = list(value.shape)
|
||||
size = _tensor_size(kind, shape, itemsize, limit)
|
||||
tensor = value.detach().resolve_neg().contiguous()
|
||||
if not torch.isfinite(tensor).all().item():
|
||||
raise ValueError("Tensor values must be finite")
|
||||
# A uint16 view preserves BF16 bits without conversion through float32.
|
||||
if tensor.dtype == torch.bfloat16:
|
||||
tensor = tensor.view(torch.uint16)
|
||||
# The live, contiguous CPU tensor owns this pointer for the entire copy.
|
||||
payload = _native_little_endian(ctypes.string_at(tensor.data_ptr(), size), itemsize)
|
||||
descriptor = {
|
||||
"encoding": "tensor-le-v1", "dtype": dtype_name,
|
||||
"shape": shape, "layout": _LAYOUTS[kind],
|
||||
}
|
||||
if kind == "AUDIO":
|
||||
descriptor["sample_rate"] = sample_rate
|
||||
return descriptor, payload
|
||||
|
||||
|
||||
def decode_payload(kind: str, descriptor: dict, payload: bytes, *, max_bytes: int = MAX_PAYLOAD_BYTES):
|
||||
"""Decode validated wire bytes into independent contiguous CPU tensors/text.
|
||||
|
||||
All metadata, shape products and byte lengths are checked before tensor
|
||||
allocation. No GPU device discovery, transfers or object deserialization
|
||||
occur. AUDIO returns {'waveform': tensor, 'sample_rate': int}.
|
||||
"""
|
||||
_validate_kind(kind)
|
||||
limit = _byte_limit(max_bytes)
|
||||
if not isinstance(descriptor, dict):
|
||||
raise ValueError("Payload descriptor must be a dictionary")
|
||||
if not isinstance(payload, bytes):
|
||||
raise ValueError("Payload must be bytes")
|
||||
if len(payload) > limit:
|
||||
raise ValueError("Payload exceeds byte limit")
|
||||
if kind == "IMAGE" and descriptor.get("encoding") == "image-transfer-v1":
|
||||
# Lazy and IMAGE-only: exact tensors and other kinds stay unchanged.
|
||||
from .image_codec import decode_image
|
||||
return decode_image(descriptor, payload, max_bytes=limit)
|
||||
if kind == "STRING":
|
||||
if set(descriptor) != {"encoding"} or descriptor["encoding"] != "utf8-v1":
|
||||
raise ValueError("Invalid STRING descriptor encoding or fields")
|
||||
try:
|
||||
return payload.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
raise ValueError("STRING payload must be valid UTF-8 bytes") from None
|
||||
|
||||
fields = {"encoding", "dtype", "shape", "layout"}
|
||||
if kind == "AUDIO":
|
||||
fields.add("sample_rate")
|
||||
if set(descriptor) != fields:
|
||||
raise ValueError("Invalid tensor descriptor fields")
|
||||
if descriptor["encoding"] != "tensor-le-v1":
|
||||
raise ValueError("Unsupported tensor descriptor encoding")
|
||||
if descriptor["layout"] != _LAYOUTS[kind]:
|
||||
raise ValueError("Invalid tensor descriptor layout")
|
||||
dtype_name = descriptor["dtype"]
|
||||
if type(dtype_name) is not str or dtype_name not in _DTYPES:
|
||||
raise ValueError("Unsupported tensor descriptor dtype")
|
||||
dtype, itemsize = _DTYPES[dtype_name]
|
||||
size = _tensor_size(kind, descriptor["shape"], itemsize, limit)
|
||||
if len(payload) != size:
|
||||
raise ValueError("Tensor payload length does not match shape and dtype")
|
||||
sample_rate = _sample_rate(descriptor["sample_rate"]) if kind == "AUDIO" else None
|
||||
|
||||
# frombuffer is always CPU. Copy immutable wire bytes into a writable private
|
||||
# buffer, then clone so the returned tensor owns its storage independently.
|
||||
buffer = bytearray(_native_little_endian(payload, itemsize))
|
||||
storage_dtype = torch.uint16 if dtype == torch.bfloat16 else dtype
|
||||
tensor = torch.frombuffer(buffer, dtype=storage_dtype).view(dtype)
|
||||
if not torch.isfinite(tensor).all().item():
|
||||
raise ValueError("Tensor payload values must be finite")
|
||||
tensor = tensor.clone().reshape(descriptor["shape"])
|
||||
if kind == "AUDIO":
|
||||
return {"waveform": tensor, "sample_rate": sample_rate}
|
||||
return tensor
|
||||
@@ -0,0 +1,249 @@
|
||||
"""Versioned, bounded media envelopes shared by master and endpoint bridge.
|
||||
|
||||
No provider or ComfyUI bootstrap imports. Network I/O is supplied by the caller,
|
||||
which must enforce public-HTTPS resolution, redirect and streaming byte limits.
|
||||
"""
|
||||
import base64
|
||||
import binascii
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from datetime import datetime, timezone
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
PROFILE = "comfy-media-v1"
|
||||
VERSION = 1
|
||||
MAX_ARTIFACT_BYTES = 256 * 1024 * 1024
|
||||
MAX_TOTAL_BYTES = 512 * 1024 * 1024
|
||||
MAX_ARTIFACTS = 128
|
||||
MAX_ITEMS = 64
|
||||
MAX_DESCRIPTOR_BYTES = 128 * 1024
|
||||
SOCKET_KINDS = {"images": "IMAGE", "audio": "AUDIO", "masks": "MASK", "texts": "STRING", "video": "VIDEO"}
|
||||
BINDING_FIELDS = {"binding_id", "collector_id", "job_id", "participant_id", "node_id", "socket", "kind"}
|
||||
_ID = re.compile(r"[A-Za-z0-9_.:/-]{1,200}\Z")
|
||||
_HASH = re.compile(r"[a-f0-9]{64}\Z")
|
||||
|
||||
|
||||
def identifier(value):
|
||||
if not isinstance(value, str) or not _ID.fullmatch(value):
|
||||
raise ValueError("Invalid media identifier")
|
||||
return value
|
||||
|
||||
|
||||
def bounded_int(value, maximum, *, minimum=0):
|
||||
if type(value) is not int or not minimum <= value <= maximum:
|
||||
raise ValueError("Invalid media count or size")
|
||||
return value
|
||||
|
||||
|
||||
def validate_json(value, *, max_bytes=MAX_DESCRIPTOR_BYTES):
|
||||
"""Bound metadata depth and size without accepting non-finite JSON numbers."""
|
||||
def walk(item, depth=0):
|
||||
if depth > 12:
|
||||
raise ValueError("Media metadata nesting exceeds limit")
|
||||
if isinstance(item, dict):
|
||||
if len(item) > 4096 or any(not isinstance(k, str) for k in item):
|
||||
raise ValueError("Invalid media metadata object")
|
||||
for child in item.values():
|
||||
walk(child, depth + 1)
|
||||
elif isinstance(item, list):
|
||||
if len(item) > 4096:
|
||||
raise ValueError("Media metadata list exceeds limit")
|
||||
for child in item:
|
||||
walk(child, depth + 1)
|
||||
elif item is not None and type(item) not in (str, int, float, bool):
|
||||
raise ValueError("Invalid media metadata value")
|
||||
walk(value)
|
||||
try:
|
||||
raw = json.dumps(value, allow_nan=False, separators=(",", ":")).encode("utf-8")
|
||||
except (ValueError, TypeError, UnicodeError, RecursionError):
|
||||
raise ValueError("Invalid media JSON") from None
|
||||
if len(raw) > max_bytes:
|
||||
raise ValueError("Media JSON exceeds byte limit")
|
||||
return raw
|
||||
|
||||
|
||||
def validate_binding(binding):
|
||||
if not isinstance(binding, dict) or set(binding) != BINDING_FIELDS:
|
||||
raise ValueError("Invalid media binding")
|
||||
for field in BINDING_FIELDS - {"socket", "kind"}:
|
||||
identifier(binding[field])
|
||||
if SOCKET_KINDS.get(binding["socket"]) != binding["kind"]:
|
||||
raise ValueError("Media socket/type mismatch")
|
||||
return binding
|
||||
|
||||
|
||||
def validate_bindings(bindings):
|
||||
if not isinstance(bindings, list) or not 1 <= len(bindings) <= MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid media bindings count")
|
||||
result = {}
|
||||
destinations = set()
|
||||
for binding in bindings:
|
||||
validate_binding(binding)
|
||||
bid = binding["binding_id"]
|
||||
destination = (binding["job_id"], binding["participant_id"], binding["collector_id"], binding["socket"])
|
||||
if bid in result or destination in destinations:
|
||||
raise ValueError("Duplicate media binding")
|
||||
result[bid] = binding
|
||||
destinations.add(destination)
|
||||
return result
|
||||
|
||||
|
||||
def _validate_url_transport(transport):
|
||||
if set(transport) != {"type", "url", "expires_at"}:
|
||||
raise ValueError("Invalid media URL transport")
|
||||
url = transport["url"]
|
||||
if not isinstance(url, str) or len(url) > 8192:
|
||||
raise ValueError("Invalid media URL")
|
||||
try:
|
||||
parsed = urlsplit(url)
|
||||
if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.fragment:
|
||||
raise ValueError()
|
||||
if parsed.port not in (None, 443):
|
||||
raise ValueError()
|
||||
expiry = datetime.fromisoformat(transport["expires_at"].replace("Z", "+00:00"))
|
||||
if expiry.tzinfo is None or expiry <= datetime.now(timezone.utc):
|
||||
raise ValueError()
|
||||
except (ValueError, TypeError, AttributeError):
|
||||
raise ValueError("Invalid or expired media URL") from None
|
||||
|
||||
|
||||
def validate_artifact(artifact, *, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
fields = {"id", "kind", "descriptor", "byte_length", "sha256", "transport"}
|
||||
if not isinstance(artifact, dict) or set(artifact) != fields:
|
||||
raise ValueError("Invalid media artifact")
|
||||
identifier(artifact["id"])
|
||||
if artifact["kind"] not in SOCKET_KINDS.values():
|
||||
raise ValueError("Unsupported media kind")
|
||||
if not isinstance(artifact["descriptor"], dict):
|
||||
raise ValueError("Invalid media descriptor")
|
||||
validate_json(artifact["descriptor"])
|
||||
size = bounded_int(artifact["byte_length"], max_bytes)
|
||||
digest = artifact["sha256"]
|
||||
if not isinstance(digest, str) or not _HASH.fullmatch(digest):
|
||||
raise ValueError("Invalid media checksum")
|
||||
transport = artifact["transport"]
|
||||
if not isinstance(transport, dict):
|
||||
raise ValueError("Invalid media transport")
|
||||
if transport.get("type") == "inline":
|
||||
if set(transport) != {"type", "data_base64"} or not isinstance(transport["data_base64"], str):
|
||||
raise ValueError("Invalid inline media transport")
|
||||
if len(transport["data_base64"]) != 4 * ((size + 2) // 3):
|
||||
raise ValueError("Inline media length mismatch")
|
||||
elif transport.get("type") == "url":
|
||||
_validate_url_transport(transport)
|
||||
else:
|
||||
raise ValueError("Unsupported media transport")
|
||||
return artifact
|
||||
|
||||
|
||||
def artifact_from_payload(artifact_id, kind, descriptor, payload, *, transport=None, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
if not isinstance(payload, bytes):
|
||||
raise ValueError("Media payload must be bytes")
|
||||
bounded_int(len(payload), max_bytes)
|
||||
result = {"id": artifact_id, "kind": kind, "descriptor": descriptor,
|
||||
"byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(),
|
||||
"transport": transport if transport is not None else {
|
||||
"type": "inline", "data_base64": base64.b64encode(payload).decode("ascii")}}
|
||||
return validate_artifact(result, max_bytes=max_bytes)
|
||||
|
||||
|
||||
async def resolve_artifact(artifact, *, download=None, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
"""download(url, *, max_bytes) must return bounded bytes; never log its URL."""
|
||||
validate_artifact(artifact, max_bytes=max_bytes)
|
||||
transport = artifact["transport"]
|
||||
if transport["type"] == "inline":
|
||||
try:
|
||||
payload = base64.b64decode(transport["data_base64"], validate=True)
|
||||
except (ValueError, binascii.Error):
|
||||
raise ValueError("Invalid media base64") from None
|
||||
else:
|
||||
if download is None:
|
||||
raise ValueError("Media URL download is not configured")
|
||||
payload = await download(transport["url"], max_bytes=artifact["byte_length"])
|
||||
if not isinstance(payload, bytes) or len(payload) != artifact["byte_length"]:
|
||||
raise ValueError("Media byte length mismatch")
|
||||
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
|
||||
raise ValueError("Media checksum mismatch")
|
||||
return payload
|
||||
|
||||
|
||||
def validate_response(response, *, request_id, bindings, max_total_bytes=MAX_TOTAL_BYTES):
|
||||
"""Validate ALL identity/completeness before downloading or delivering anything."""
|
||||
expected = validate_bindings(bindings)
|
||||
identifier(request_id)
|
||||
if not isinstance(response, dict) or set(response) != {"version", "profile", "request_id", "bindings", "artifacts", "items"}:
|
||||
raise ValueError("Invalid media response envelope")
|
||||
if type(response["version"]) is not int or response["version"] != VERSION or response["profile"] != PROFILE:
|
||||
raise ValueError("Unsupported media response version/profile")
|
||||
if response["request_id"] != request_id:
|
||||
raise ValueError("Media request identity mismatch")
|
||||
returned = response["bindings"]
|
||||
if not isinstance(returned, list) or len(returned) != len(expected):
|
||||
raise ValueError("Missing or extra media bindings")
|
||||
counts = {}
|
||||
for row in returned:
|
||||
if not isinstance(row, dict) or set(row) != BINDING_FIELDS | {"count"}:
|
||||
raise ValueError("Invalid result binding")
|
||||
binding = {key: row[key] for key in BINDING_FIELDS}
|
||||
validate_binding(binding)
|
||||
bid = binding["binding_id"]
|
||||
if bid in counts or expected.get(bid) != binding:
|
||||
raise ValueError("Media binding identity mismatch")
|
||||
# Connected outputs must produce at least one value. Empty STRING is a value.
|
||||
counts[bid] = bounded_int(row["count"], MAX_ITEMS, minimum=1)
|
||||
artifacts = response["artifacts"]
|
||||
if not isinstance(artifacts, list) or not 1 <= len(artifacts) <= MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid media artifacts count")
|
||||
indexed = {}
|
||||
total = 0
|
||||
for artifact in artifacts:
|
||||
validate_artifact(artifact)
|
||||
aid = artifact["id"]
|
||||
if aid in indexed:
|
||||
raise ValueError("Duplicate media artifact")
|
||||
total += artifact["byte_length"]
|
||||
if total > max_total_bytes:
|
||||
raise ValueError("Media response exceeds total byte limit")
|
||||
indexed[aid] = artifact
|
||||
items = response["items"]
|
||||
if not isinstance(items, list) or len(items) != sum(counts.values()):
|
||||
raise ValueError("Missing or extra media items")
|
||||
ordered = {bid: [None] * count for bid, count in counts.items()}
|
||||
used = set()
|
||||
for item in items:
|
||||
if not isinstance(item, dict) or set(item) != {"binding_id", "item_index", "artifact_id"}:
|
||||
raise ValueError("Invalid media item")
|
||||
bid, aid = identifier(item["binding_id"]), identifier(item["artifact_id"])
|
||||
if bid not in counts or aid not in indexed:
|
||||
raise ValueError("Unknown media item identity")
|
||||
index = bounded_int(item["item_index"], counts[bid] - 1)
|
||||
if ordered[bid][index] is not None or indexed[aid]["kind"] != expected[bid]["kind"]:
|
||||
raise ValueError("Duplicate or mistyped media item")
|
||||
ordered[bid][index] = indexed[aid]
|
||||
used.add(aid)
|
||||
if used != set(indexed) or any(value is None for values in ordered.values() for value in values):
|
||||
raise ValueError("Incomplete media result manifest")
|
||||
return ordered
|
||||
|
||||
|
||||
async def decode_response(response, *, request_id, bindings, download=None, max_total_bytes=MAX_TOTAL_BYTES):
|
||||
ordered = validate_response(response, request_id=request_id, bindings=bindings, max_total_bytes=max_total_bytes)
|
||||
# Resolve and decode before exposing a single result to the collector queue.
|
||||
from .codec import decode_payload
|
||||
cache = {}
|
||||
decoded = {}
|
||||
for bid, artifacts in ordered.items():
|
||||
values = []
|
||||
for artifact in artifacts:
|
||||
aid = artifact["id"]
|
||||
if aid not in cache:
|
||||
payload = await resolve_artifact(artifact, download=download)
|
||||
if artifact["kind"] == "VIDEO":
|
||||
from .video_codec import decode_video
|
||||
cache[aid] = decode_video(artifact["descriptor"], payload)
|
||||
else:
|
||||
cache[aid] = decode_payload(artifact["kind"], artifact["descriptor"], payload)
|
||||
values.append(cache[aid])
|
||||
decoded[bid] = values
|
||||
return decoded
|
||||
@@ -0,0 +1,794 @@
|
||||
"""Pure HTTP executor for media-v1 requests against a local ComfyUI instance.
|
||||
|
||||
No ComfyUI, tensor, or provider imports. ``download(url, *, max_bytes)`` must
|
||||
provide public-HTTPS/redirect/DNS protection and bounded streaming. Optional
|
||||
``publish(metadata, payload)`` returns a contract-valid, expiring HTTPS transport.
|
||||
|
||||
Only identified invocation-owned files are deleted. Optional trusted local
|
||||
``temp_dir`` enables exporter cleanup after terminal history and collection.
|
||||
Without it, exporter retention belongs to the ComfyUI temp-directory lifecycle.
|
||||
The pinned native /interrupt route is not atomic: running-prompt cancellation
|
||||
fails closed and retains files instead of risking an unrelated successor prompt.
|
||||
"""
|
||||
import asyncio
|
||||
import base64
|
||||
import binascii
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import stat
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from urllib.parse import quote, urlsplit
|
||||
|
||||
import aiohttp
|
||||
|
||||
from .contract import (
|
||||
BINDING_FIELDS, MAX_ARTIFACT_BYTES, MAX_ARTIFACTS, MAX_TOTAL_BYTES,
|
||||
PROFILE, VERSION, artifact_from_payload, identifier, resolve_artifact,
|
||||
validate_artifact, validate_bindings, validate_json, validate_response,
|
||||
)
|
||||
|
||||
MAX_RESPONSE_JSON_BYTES = 10_000_000
|
||||
MAX_HTTP_JSON_BYTES = 32 * 1024 * 1024
|
||||
POLL_INTERVAL_SECONDS = 0.25
|
||||
CANCEL_GRACE_SECONDS = 10.0
|
||||
DRAIN_GRACE_SECONDS = 10.0
|
||||
HTTP_TIMEOUT_SECONDS = 15.0
|
||||
_UNCERTAIN_WORK = set()
|
||||
GET_ATTEMPTS = 3
|
||||
_IMPORT_SLOTS = {"IMAGE": 0, "AUDIO": 1, "MASK": 2, "STRING": 3, "VIDEO": 4}
|
||||
# Used ONLY to run the shared manifest validator before /view. Never fetched.
|
||||
_VALIDATION_TRANSPORT = {"type": "url", "url": "https://media-validation.invalid/artifact",
|
||||
"expires_at": "9999-01-01T00:00:00Z"}
|
||||
|
||||
|
||||
class MediaExecutionUncertain(RuntimeError):
|
||||
"""The local prompt may still execute; staged files were deliberately retained."""
|
||||
|
||||
|
||||
class _HTTPError(RuntimeError):
|
||||
def __init__(self, status):
|
||||
super().__init__(f"Local Comfy HTTP request failed (status {status})")
|
||||
self.status = status
|
||||
|
||||
|
||||
def _scope_id(value):
|
||||
identifier(value)
|
||||
if any(part in ("", ".", "..") for part in value.split("/")):
|
||||
raise ValueError("Unsafe media scope identifier")
|
||||
return value
|
||||
|
||||
|
||||
def _json_bytes(value):
|
||||
return json.dumps(value, allow_nan=False).encode("utf-8")
|
||||
|
||||
|
||||
def _link(value, workflow):
|
||||
if (not isinstance(value, list) or len(value) != 2 or
|
||||
not isinstance(value[0], str) or value[0] not in workflow or
|
||||
type(value[1]) is not int or value[1] < 0):
|
||||
raise ValueError("Invalid workflow socket link")
|
||||
return value[0]
|
||||
|
||||
|
||||
def _validate_request(input_data):
|
||||
if (not isinstance(input_data, dict) or not {"workflow", "media"} <= set(input_data)
|
||||
or set(input_data) - {"workflow", "media", "images"}):
|
||||
raise ValueError("Invalid media request envelope")
|
||||
media = input_data["media"]
|
||||
if (not isinstance(media, dict) or set(media) !=
|
||||
{"version", "profile", "request_id", "artifacts", "inputs", "outputs"}):
|
||||
raise ValueError("Invalid media request controls")
|
||||
if type(media["version"]) is not int or media["version"] != VERSION or media["profile"] != PROFILE:
|
||||
raise ValueError("Unsupported media request version/profile")
|
||||
_scope_id(media["request_id"])
|
||||
bindings = validate_bindings(media["outputs"])
|
||||
targets = []
|
||||
for binding in bindings.values():
|
||||
for field in BINDING_FIELDS - {"socket", "kind"}:
|
||||
_scope_id(binding[field])
|
||||
targets.append(binding["node_id"])
|
||||
if len(set(targets)) != len(targets):
|
||||
raise ValueError("Duplicate media exporter node")
|
||||
|
||||
workflow = input_data["workflow"]
|
||||
if not isinstance(workflow, dict) or not workflow:
|
||||
raise ValueError("Invalid API workflow")
|
||||
validate_json(workflow, max_bytes=16 * 1024 * 1024)
|
||||
dependencies = {}
|
||||
for node_id, node in workflow.items():
|
||||
_scope_id(node_id)
|
||||
if (not isinstance(node, dict) or not isinstance(node.get("class_type"), str)
|
||||
or not node["class_type"] or not isinstance(node.get("inputs"), dict)):
|
||||
raise ValueError("Invalid API workflow node")
|
||||
if node["class_type"] == "MediaImportV1":
|
||||
raise ValueError("Request may not supply its own media import filenames")
|
||||
if node["class_type"] == "MediaExportV1" and node_id not in targets:
|
||||
raise ValueError("Undeclared media exporter")
|
||||
dependencies[node_id] = set()
|
||||
for name, value in node["inputs"].items():
|
||||
_scope_id(name)
|
||||
if isinstance(value, list):
|
||||
dependencies[node_id].add(_link(value, workflow))
|
||||
for binding in bindings.values():
|
||||
node = workflow.get(binding["node_id"])
|
||||
if node is None or node["class_type"] != "MediaExportV1":
|
||||
raise ValueError("Missing declared MediaExportV1 node")
|
||||
controls = node["inputs"]
|
||||
from .image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
|
||||
required = {"request_id", "binding_json", binding["socket"]}
|
||||
optional = {"transfer_format", "transfer_quality"} if binding["kind"] == "IMAGE" else set()
|
||||
if not required <= set(controls) or set(controls) - required - optional:
|
||||
raise ValueError("Media exporter must have exactly its declared socket and controls")
|
||||
validate_options(controls.get("transfer_format", DEFAULT_FORMAT),
|
||||
controls.get("transfer_quality", DEFAULT_QUALITY))
|
||||
if controls["request_id"] != media["request_id"]:
|
||||
raise ValueError("Media exporter request identity mismatch")
|
||||
text = controls["binding_json"]
|
||||
if not isinstance(text, str) or len(text) > 128 * 1024:
|
||||
raise ValueError("Invalid media exporter binding JSON")
|
||||
try:
|
||||
declared = json.loads(text)
|
||||
except (ValueError, RecursionError):
|
||||
raise ValueError("Invalid media exporter binding JSON") from None
|
||||
if declared != binding:
|
||||
raise ValueError("Media exporter binding identity mismatch")
|
||||
_link(controls[binding["socket"]], workflow)
|
||||
# The execution roots are exactly the declared exporters, not every class
|
||||
# with OUTPUT_NODE=True. Required upstream output utilities may remain.
|
||||
visited = set()
|
||||
pending = list(targets)
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id not in visited:
|
||||
visited.add(node_id)
|
||||
pending.extend(dependencies[node_id])
|
||||
if visited != set(workflow):
|
||||
raise ValueError("Workflow contains nodes outside declared output roots")
|
||||
if any(set(targets) & sources for sources in dependencies.values()):
|
||||
raise ValueError("A media exporter cannot be an upstream dependency")
|
||||
# Kahn traversal bounds cycle checking without relying on recursion depth.
|
||||
counts = {node: len(sources) for node, sources in dependencies.items()}
|
||||
consumers = {node: [] for node in workflow}
|
||||
for node, sources in dependencies.items():
|
||||
for source in sources:
|
||||
consumers[source].append(node)
|
||||
pending = [node for node, count in counts.items() if not count]
|
||||
processed = 0
|
||||
while pending:
|
||||
processed += 1
|
||||
for consumer in consumers[pending.pop()]:
|
||||
counts[consumer] -= 1
|
||||
if not counts[consumer]:
|
||||
pending.append(consumer)
|
||||
if processed != len(workflow):
|
||||
raise ValueError("Workflow contains a dependency cycle")
|
||||
|
||||
artifacts = media["artifacts"]
|
||||
if not isinstance(artifacts, list) or len(artifacts) > MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid input artifacts count")
|
||||
indexed, total = {}, 0
|
||||
for artifact in artifacts:
|
||||
validate_artifact(artifact)
|
||||
aid = _scope_id(artifact["id"])
|
||||
if aid in indexed:
|
||||
raise ValueError("Duplicate input artifact")
|
||||
indexed[aid] = artifact
|
||||
total += artifact["byte_length"]
|
||||
inputs = media["inputs"]
|
||||
if not isinstance(inputs, list) or len(inputs) > 4096:
|
||||
raise ValueError("Invalid input references")
|
||||
destinations, used = set(), set()
|
||||
exporters = {binding["node_id"]: binding for binding in bindings.values()}
|
||||
for ref in inputs:
|
||||
if not isinstance(ref, dict) or set(ref) != {"node_id", "input_name", "artifact_id"}:
|
||||
raise ValueError("Invalid media input reference")
|
||||
for value in ref.values():
|
||||
_scope_id(value)
|
||||
node_id, name, aid = ref["node_id"], ref["input_name"], ref["artifact_id"]
|
||||
if (node_id not in workflow or name not in workflow[node_id]["inputs"]
|
||||
or aid not in indexed or (node_id, name) in destinations):
|
||||
raise ValueError("Unknown or duplicate media input reference")
|
||||
binding = exporters.get(node_id)
|
||||
if binding is not None and (name != binding["socket"] or indexed[aid]["kind"] != binding["kind"]):
|
||||
raise ValueError("Exporter input must match its declared media socket and kind")
|
||||
destinations.add((node_id, name))
|
||||
used.add(aid)
|
||||
if used != set(indexed):
|
||||
raise ValueError("Every input artifact must be used")
|
||||
|
||||
images = input_data.get("images", [])
|
||||
if not isinstance(images, list) or len(images) > MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid legacy images count")
|
||||
names = set()
|
||||
legacy = []
|
||||
for image in images:
|
||||
if not isinstance(image, dict) or set(image) != {"name", "image"}:
|
||||
raise ValueError("Invalid legacy image")
|
||||
name, encoded = image["name"], image["image"]
|
||||
if (not isinstance(name, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_. -]{0,254}", name)
|
||||
or name in names or not isinstance(encoded, str)):
|
||||
raise ValueError("Invalid or duplicate legacy image basename")
|
||||
names.add(name)
|
||||
if len(encoded) > 4 * ((MAX_ARTIFACT_BYTES + 2) // 3):
|
||||
raise ValueError("Legacy image exceeds byte limit")
|
||||
# Exact decoded length, including padding, before allocating bytes.
|
||||
size = (len(encoded) // 4) * 3 - (len(encoded) - len(encoded.rstrip("=")))
|
||||
if len(encoded) % 4 or size < 0 or size > MAX_ARTIFACT_BYTES:
|
||||
raise ValueError("Invalid legacy image base64 length")
|
||||
total += size
|
||||
matches = [node_id for node_id, node in workflow.items()
|
||||
if node["class_type"] == "LoadImage" and node["inputs"].get("image") == name]
|
||||
if not matches or any((node_id, "image") in destinations for node_id in matches):
|
||||
raise ValueError("Unused or conflicting legacy image binding")
|
||||
legacy.append((encoded, size, matches))
|
||||
if total > MAX_TOTAL_BYTES:
|
||||
raise ValueError("Media inputs exceed total byte limit")
|
||||
return copy.deepcopy(workflow), copy.deepcopy(media), legacy, targets
|
||||
|
||||
|
||||
class _Staging:
|
||||
"""Exclusive random files under a hashed namespace; never delete siblings."""
|
||||
def __init__(self, input_dir, request_id):
|
||||
self.root = Path(input_dir).resolve()
|
||||
self.key = hashlib.sha256(request_id.encode()).hexdigest()
|
||||
self.relative = f"distributed-media-v1/{self.key}"
|
||||
self.names = []
|
||||
self.fd = None
|
||||
|
||||
def _open(self):
|
||||
self.root.mkdir(parents=True, exist_ok=True)
|
||||
directory = self.root
|
||||
for component in ("distributed-media-v1", self.key):
|
||||
directory = directory / component
|
||||
if directory.is_symlink():
|
||||
raise ValueError("Unsafe media staging directory symlink")
|
||||
directory.mkdir(exist_ok=True, mode=0o700)
|
||||
if not directory.resolve().is_relative_to(self.root):
|
||||
raise ValueError("Unsafe media staging directory")
|
||||
self.fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
|
||||
|
||||
def write(self, payload):
|
||||
if self.fd is None:
|
||||
self._open()
|
||||
name = uuid.uuid4().hex + ".bin"
|
||||
fd = os.open(name, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600, dir_fd=self.fd)
|
||||
self.names.append(name)
|
||||
with os.fdopen(fd, "wb") as stream:
|
||||
stream.write(payload)
|
||||
return f"{self.relative}/{name}"
|
||||
|
||||
def close(self, *, cleanup):
|
||||
if self.fd is not None:
|
||||
try:
|
||||
if cleanup:
|
||||
for name in self.names:
|
||||
try:
|
||||
os.unlink(name, dir_fd=self.fd)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
finally:
|
||||
os.close(self.fd)
|
||||
self.fd = None
|
||||
if cleanup:
|
||||
# Never recursively remove a shared namespace or another run.
|
||||
for directory in (self.root / self.relative, self.root / "distributed-media-v1"):
|
||||
try:
|
||||
directory.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
class _OwnedExports:
|
||||
"""Best-effort retirement of validated terminal-history files, never a sweep.
|
||||
|
||||
Walk every directory component with openat/O_NOFOLLOW, then pin the directory
|
||||
fd and each regular file's identity. Missing, linked, replaced, or unidentifiable
|
||||
files stay with the endpoint lifecycle. Never use history-supplied path strings.
|
||||
"""
|
||||
def __init__(self, temp_dir, request_id):
|
||||
self.root = Path(temp_dir).absolute() if temp_dir is not None else None
|
||||
self.key = hashlib.sha256(request_id.encode()).hexdigest()
|
||||
self.fd = None
|
||||
self.files = {}
|
||||
|
||||
def _open(self):
|
||||
if self.fd is not None:
|
||||
return True
|
||||
if self.root is None or ".." in self.root.parts:
|
||||
return False
|
||||
flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
|
||||
fd = None
|
||||
try:
|
||||
fd = os.open(self.root.anchor, flags)
|
||||
for component in (*self.root.parts[1:], "distributed-media-v1", self.key):
|
||||
child = os.open(component, flags, dir_fd=fd)
|
||||
os.close(fd)
|
||||
fd = child
|
||||
self.fd, fd = fd, None
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
finally:
|
||||
if fd is not None:
|
||||
os.close(fd)
|
||||
|
||||
def _identity(self, name):
|
||||
info = os.stat(name, dir_fd=self.fd, follow_symlinks=False)
|
||||
if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1:
|
||||
return None
|
||||
return (info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns)
|
||||
|
||||
def identify(self, entry, media):
|
||||
if self.root is None or not isinstance(entry.get("outputs"), dict):
|
||||
return
|
||||
# A failed prompt may contain valid completed exporters alongside missing
|
||||
# or malformed ones. Claim only independently validated declared records.
|
||||
for binding in media["outputs"]:
|
||||
node_id = binding["node_id"]
|
||||
try:
|
||||
_, files = _manifest({"outputs": {node_id: entry["outputs"].get(node_id)}},
|
||||
dict(media, outputs=[binding]))
|
||||
except ValueError:
|
||||
continue
|
||||
if not files or not self._open():
|
||||
continue
|
||||
for artifact in files:
|
||||
name = artifact["id"] + ".bin" # _manifest enforces UUID hex.
|
||||
try:
|
||||
identity = self._identity(name)
|
||||
except OSError:
|
||||
continue
|
||||
if identity is not None:
|
||||
self.files.setdefault(name, identity)
|
||||
|
||||
def close(self, *, cleanup):
|
||||
if self.fd is None:
|
||||
return
|
||||
try:
|
||||
if cleanup:
|
||||
for name, identity in self.files.items():
|
||||
try:
|
||||
if self._identity(name) == identity:
|
||||
os.unlink(name, dir_fd=self.fd)
|
||||
except OSError:
|
||||
pass # Endpoint lifecycle owns inaccessible leftovers.
|
||||
finally:
|
||||
os.close(self.fd)
|
||||
self.fd = None
|
||||
# Deliberately retain directories: they may belong to concurrent runs.
|
||||
|
||||
|
||||
async def _body(response, maximum):
|
||||
length = response.content_length
|
||||
if length is not None and length > maximum:
|
||||
raise ValueError("Local Comfy response exceeds declared byte length")
|
||||
body = bytearray()
|
||||
async for chunk in response.content.iter_chunked(64 * 1024):
|
||||
if len(body) + len(chunk) > maximum:
|
||||
raise ValueError("Local Comfy response exceeds declared byte length")
|
||||
body.extend(chunk)
|
||||
return bytes(body)
|
||||
|
||||
|
||||
class _Execution:
|
||||
def __init__(self, session, url, stage, media, targets, deadline, exports):
|
||||
self.session, self.url, self.stage = session, url, stage
|
||||
self.media, self.targets = media, targets
|
||||
self.deadline = deadline
|
||||
self.run_id = uuid.uuid4().hex
|
||||
self.prompt_id = None
|
||||
self.attempted = False
|
||||
self.terminal = False
|
||||
self.exports = exports
|
||||
self.interruption_blocked = False
|
||||
|
||||
def remaining(self):
|
||||
remaining = self.deadline - asyncio.get_running_loop().time()
|
||||
if remaining <= 0:
|
||||
raise asyncio.TimeoutError("Media execution deadline expired")
|
||||
return remaining
|
||||
|
||||
async def json_request(self, method, path, *, body=None, permit_rejection=False):
|
||||
async with self.session.request(method, self.url + path, json=body,
|
||||
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining())),
|
||||
allow_redirects=False) as response:
|
||||
raw = await _body(response, MAX_HTTP_JSON_BYTES)
|
||||
# Native queue-delete and interrupt routes return an empty 200.
|
||||
if response.status == 200 and not raw and method == "POST" and path in {"/queue", "/interrupt"}:
|
||||
return {}
|
||||
try:
|
||||
value = json.loads(raw)
|
||||
except (ValueError, UnicodeError, RecursionError):
|
||||
if response.status != 200:
|
||||
raise _HTTPError(response.status) from None
|
||||
raise ValueError("Invalid local Comfy JSON response") from None
|
||||
if response.status != 200:
|
||||
# Native /prompt rejects the entire graph with a 400 error.
|
||||
# Other failures may hide acceptance and are never retried.
|
||||
if (permit_rejection and response.status == 400 and isinstance(value, dict)
|
||||
and value.get("error") and "prompt_id" not in value):
|
||||
self.terminal = True
|
||||
raise _HTTPError(response.status)
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError("Invalid local Comfy JSON object")
|
||||
return value
|
||||
|
||||
async def get_json(self, path):
|
||||
for attempt in range(GET_ATTEMPTS):
|
||||
try:
|
||||
return await self.json_request("GET", path)
|
||||
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
||||
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
|
||||
raise
|
||||
if attempt == GET_ATTEMPTS - 1:
|
||||
raise
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
|
||||
|
||||
async def history(self):
|
||||
response = await self.get_json("/history/" + quote(self.prompt_id, safe=""))
|
||||
entry = response.get(self.prompt_id)
|
||||
if entry is None:
|
||||
return None
|
||||
if not isinstance(entry, dict) or not isinstance(entry.get("status"), dict):
|
||||
raise ValueError("Malformed Comfy prompt history")
|
||||
status = entry["status"]
|
||||
messages = status.get("messages", [])
|
||||
if not isinstance(messages, list):
|
||||
raise ValueError("Malformed Comfy history messages")
|
||||
failed = status.get("status_str") in {"error", "interrupted"} or any(
|
||||
isinstance(message, list) and message and message[0] in {"execution_error", "execution_interrupted"}
|
||||
for message in messages)
|
||||
if failed or status.get("completed") is True:
|
||||
self.terminal = True
|
||||
self.exports.identify(entry, self.media)
|
||||
if failed or status.get("status_str") != "success":
|
||||
raise RuntimeError("Local Comfy prompt execution failed")
|
||||
return entry
|
||||
return None
|
||||
|
||||
async def submit_and_wait(self, workflow):
|
||||
body = {"prompt": workflow, "client_id": self.run_id,
|
||||
"extra_data": {"distributed_media_run_id": self.run_id},
|
||||
"partial_execution_targets": list(self.targets)}
|
||||
self.remaining() # CPU staging must not spend past the absolute deadline.
|
||||
self.attempted = True
|
||||
response = await self.json_request("POST", "/prompt", body=body, permit_rejection=True)
|
||||
self.prompt_id = _scope_id(response.get("prompt_id"))
|
||||
if response.get("node_errors"):
|
||||
raise RuntimeError("Local Comfy accepted an incomplete set of execution targets")
|
||||
while True:
|
||||
entry = await self.history()
|
||||
if entry is not None:
|
||||
return entry
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
||||
|
||||
async def settle(self):
|
||||
"""Bounded recovery; caller shields this from task cancellation."""
|
||||
deleted = False
|
||||
while True:
|
||||
if self.prompt_id:
|
||||
try:
|
||||
await self.history()
|
||||
except RuntimeError:
|
||||
if not self.terminal:
|
||||
raise
|
||||
if self.terminal:
|
||||
return
|
||||
queue = await self.get_json("/queue")
|
||||
groups = []
|
||||
for key in ("queue_running", "queue_pending"):
|
||||
rows = queue.get(key)
|
||||
if not isinstance(rows, list):
|
||||
raise ValueError("Malformed Comfy queue")
|
||||
group = {}
|
||||
for row in rows:
|
||||
if not isinstance(row, list) or len(row) < 4 or not isinstance(row[3], dict):
|
||||
raise ValueError("Malformed Comfy queue entry")
|
||||
pid = _scope_id(row[1])
|
||||
if pid in group:
|
||||
raise ValueError("Duplicate Comfy queue identity")
|
||||
group[pid] = row[3]
|
||||
groups.append(group)
|
||||
running, pending = groups
|
||||
if self.prompt_id is None:
|
||||
own = [pid for group in groups for pid, extra in group.items()
|
||||
if extra.get("distributed_media_run_id") == self.run_id
|
||||
and extra.get("client_id") == self.run_id]
|
||||
if len(own) > 1:
|
||||
raise ValueError("Ambiguous Comfy queue ownership")
|
||||
if own:
|
||||
self.prompt_id = own[0]
|
||||
if self.prompt_id in pending and not deleted:
|
||||
await self.json_request("POST", "/queue", body={"delete": [self.prompt_id]})
|
||||
deleted = True
|
||||
elif self.prompt_id in running:
|
||||
# Pinned ComfyUI server.py checks a snapshot then globally
|
||||
# interrupts, outside the queue mutex. Supplying prompt_id does
|
||||
# NOT close that race. No safe capability negotiation exists on
|
||||
# this route; do not invoke it, even after our own queue read.
|
||||
self.interruption_blocked = True
|
||||
raise RuntimeError("Native prompt-scoped interruption unavailable")
|
||||
elif deleted and self.prompt_id not in running and self.prompt_id not in pending:
|
||||
self.terminal = True # queue deletion confirmed by exact read-back
|
||||
return
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
||||
|
||||
async def fetch(self, artifact):
|
||||
transport = artifact["transport"]
|
||||
params = {"filename": transport["filename"], "subfolder": transport["subfolder"], "type": "temp"}
|
||||
for attempt in range(GET_ATTEMPTS):
|
||||
try:
|
||||
async with self.session.get(self.url + "/view", params=params, allow_redirects=False,
|
||||
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining()))) as response:
|
||||
if response.status != 200:
|
||||
raise _HTTPError(response.status)
|
||||
payload = await _body(response, artifact["byte_length"])
|
||||
break
|
||||
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
||||
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
|
||||
raise
|
||||
if attempt == GET_ATTEMPTS - 1:
|
||||
raise
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
|
||||
if len(payload) != artifact["byte_length"]:
|
||||
raise ValueError("Media output byte length mismatch")
|
||||
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
|
||||
raise ValueError("Media output checksum mismatch")
|
||||
return payload
|
||||
|
||||
|
||||
def _manifest(entry, media):
|
||||
outputs = entry.get("outputs")
|
||||
if not isinstance(outputs, dict):
|
||||
raise ValueError("Missing media history outputs")
|
||||
expected_ids = {binding["node_id"] for binding in media["outputs"]}
|
||||
if any(node_id not in expected_ids and isinstance(value, dict) and "media_exports" in value
|
||||
for node_id, value in outputs.items()):
|
||||
raise ValueError("Unexpected media export history")
|
||||
response = {"version": VERSION, "profile": PROFILE, "request_id": media["request_id"],
|
||||
"bindings": [], "artifacts": [], "items": []}
|
||||
files = []
|
||||
subfolder = "distributed-media-v1/" + hashlib.sha256(media["request_id"].encode()).hexdigest()
|
||||
for binding in media["outputs"]:
|
||||
node = outputs.get(binding["node_id"], {})
|
||||
records = node.get("media_exports") if isinstance(node, dict) else None
|
||||
if not isinstance(records, list) or len(records) != 1:
|
||||
raise ValueError("Missing or duplicate media export result")
|
||||
record = records[0]
|
||||
if (not isinstance(record, dict) or set(record) != {"request_id", "binding", "artifacts", "items"}
|
||||
or record["request_id"] != media["request_id"]):
|
||||
raise ValueError("Media export request identity mismatch")
|
||||
if not isinstance(record["artifacts"], list) or len(record["artifacts"]) > MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid exported artifacts count")
|
||||
converted = []
|
||||
for artifact in record["artifacts"]:
|
||||
if not isinstance(artifact, dict) or not isinstance(artifact.get("transport"), dict):
|
||||
raise ValueError("Invalid exported artifact")
|
||||
transport = artifact["transport"]
|
||||
aid = artifact.get("id")
|
||||
if (not isinstance(aid, str) or not re.fullmatch(r"[a-f0-9]{32}", aid)
|
||||
or transport != {"type": "comfy-file", "filename": aid + ".bin",
|
||||
"subfolder": subfolder, "folder_type": "temp"}):
|
||||
raise ValueError("Unsafe or unowned exported media file")
|
||||
converted.append(dict(artifact, transport=dict(_VALIDATION_TRANSPORT)))
|
||||
partial = dict(response, bindings=[record["binding"]], artifacts=converted, items=record["items"])
|
||||
validate_response(partial, request_id=media["request_id"], bindings=[binding])
|
||||
response["bindings"].append(copy.deepcopy(record["binding"]))
|
||||
response["artifacts"].extend(converted)
|
||||
response["items"].extend(copy.deepcopy(record["items"]))
|
||||
files.extend(copy.deepcopy(record["artifacts"]))
|
||||
# Validate the whole manifest BEFORE the first /view or publisher callback.
|
||||
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
|
||||
return response, files
|
||||
|
||||
|
||||
def _publication_plan(response, files, publish):
|
||||
skeleton = dict(response, artifacts=[dict(a, transport={"type": "inline", "data_base64": ""}) for a in files])
|
||||
base_size = len(_json_bytes({"media": skeleton}))
|
||||
sizes = {a["id"]: 4 * ((a["byte_length"] + 2) // 3) for a in files}
|
||||
predicted = base_size + sum(sizes.values())
|
||||
published = set()
|
||||
if predicted > MAX_RESPONSE_JSON_BYTES and publish is None:
|
||||
raise ValueError("Media response exceeds inline JSON limit; a publisher is required")
|
||||
# Reserve enough space for a maximum-size escaped URL, including metadata.
|
||||
# Small remaining inline values preserve their exact bytes and ordering.
|
||||
for aid in sorted(sizes, key=sizes.get, reverse=True):
|
||||
if predicted <= MAX_RESPONSE_JSON_BYTES:
|
||||
break
|
||||
if sizes[aid] > 65_536:
|
||||
published.add(aid)
|
||||
predicted -= sizes[aid] - 65_536
|
||||
if predicted > MAX_RESPONSE_JSON_BYTES:
|
||||
raise ValueError("Media response metadata exceeds JSON limit")
|
||||
return published
|
||||
|
||||
|
||||
async def _run(execution, workflow, legacy, download, publish):
|
||||
media = execution.media
|
||||
imports = {}
|
||||
for artifact in media["artifacts"]:
|
||||
execution.remaining()
|
||||
payload = await resolve_artifact(artifact, download=download)
|
||||
execution.remaining()
|
||||
filename = execution.stage.write(payload)
|
||||
node_id = "distributed-media-import:" + uuid.uuid4().hex
|
||||
while node_id in workflow:
|
||||
node_id = "distributed-media-import:" + uuid.uuid4().hex
|
||||
metadata = {key: value for key, value in artifact.items() if key != "transport"}
|
||||
workflow[node_id] = {"class_type": "MediaImportV1", "inputs": {
|
||||
"artifact_json": json.dumps(metadata, allow_nan=False), "filename": filename}}
|
||||
imports[artifact["id"]] = [node_id, _IMPORT_SLOTS[artifact["kind"]]]
|
||||
for ref in media["inputs"]:
|
||||
workflow[ref["node_id"]]["inputs"][ref["input_name"]] = list(imports[ref["artifact_id"]])
|
||||
for encoded, size, matches in legacy:
|
||||
execution.remaining()
|
||||
try:
|
||||
payload = base64.b64decode(encoded, validate=True)
|
||||
except (ValueError, binascii.Error):
|
||||
raise ValueError("Invalid legacy image base64") from None
|
||||
if len(payload) != size:
|
||||
raise ValueError("Legacy image byte length mismatch")
|
||||
execution.remaining()
|
||||
filename = execution.stage.write(payload)
|
||||
for node_id in matches:
|
||||
workflow[node_id]["inputs"]["image"] = filename
|
||||
# Replacements can disconnect file loaders whose original endpoint-local
|
||||
# paths are unavailable. Keep exactly the post-rewrite exporter closure;
|
||||
# shared ancestors still needed by another branch must survive.
|
||||
reachable, pending = set(), list(execution.targets)
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id in reachable:
|
||||
continue
|
||||
reachable.add(node_id)
|
||||
pending.extend(_link(value, workflow) for value in workflow[node_id]["inputs"].values()
|
||||
if isinstance(value, list))
|
||||
workflow = {node_id: node for node_id, node in workflow.items() if node_id in reachable}
|
||||
entry = await execution.submit_and_wait(workflow)
|
||||
response, files = _manifest(entry, media)
|
||||
published = _publication_plan(response, files, publish)
|
||||
artifacts = []
|
||||
for artifact in files:
|
||||
payload = await execution.fetch(artifact)
|
||||
transport = None
|
||||
if artifact["id"] in published:
|
||||
metadata = {key: copy.deepcopy(value) for key, value in artifact.items() if key != "transport"}
|
||||
transport = await publish(metadata, payload)
|
||||
if not isinstance(transport, dict) or transport.get("type") != "url":
|
||||
raise ValueError("Publisher must return an expiring HTTPS URL transport")
|
||||
artifacts.append(artifact_from_payload(artifact["id"], artifact["kind"], artifact["descriptor"],
|
||||
payload, transport=transport))
|
||||
response["artifacts"] = artifacts
|
||||
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
|
||||
result = {"media": response}
|
||||
if len(_json_bytes(result)) > MAX_RESPONSE_JSON_BYTES:
|
||||
raise ValueError("Media response exceeds JSON byte limit")
|
||||
execution.remaining()
|
||||
return result
|
||||
|
||||
|
||||
def _release_uncertain_work(task):
|
||||
_UNCERTAIN_WORK.discard(task)
|
||||
try:
|
||||
task.result()
|
||||
except BaseException:
|
||||
pass
|
||||
|
||||
|
||||
async def _drain_work(work):
|
||||
"""Cancel once; a misbehaving callback cannot make teardown unbounded."""
|
||||
if not work.done():
|
||||
work.cancel()
|
||||
deadline = asyncio.get_running_loop().time() + DRAIN_GRACE_SECONDS
|
||||
while not work.done():
|
||||
remaining = deadline - asyncio.get_running_loop().time()
|
||||
if remaining <= 0:
|
||||
# Keep a strong reference until it exits, but never delete its files
|
||||
# asynchronously: uncertain ownership belongs to endpoint recovery.
|
||||
_UNCERTAIN_WORK.add(work)
|
||||
work.add_done_callback(_release_uncertain_work)
|
||||
return False
|
||||
try:
|
||||
# Unlike wait_for, wait's timeout does not wait for cancellation.
|
||||
await asyncio.wait({work}, timeout=remaining)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
_release_uncertain_work(work)
|
||||
return True
|
||||
|
||||
|
||||
async def execute_media_request(input_data, *, comfy_url, input_dir, temp_dir=None,
|
||||
timeout_seconds=1800, session=None, download=None, publish=None) -> dict:
|
||||
"""Validate, stage, submit once, collect exact media, and retire owned inputs.
|
||||
|
||||
``timeout_seconds`` covers staging, HTTP, polling, and publishing. A separate
|
||||
bounded cancellation grace is allowed ONLY to establish the local prompt is
|
||||
terminal. On uncertain acceptance/termination, raise MediaExecutionUncertain
|
||||
and retain input files rather than racing an executing prompt. Never retry an
|
||||
ambiguous submission, even when the exception is a timeout.
|
||||
Async callback teardown has its own bounded grace; non-cooperative callbacks
|
||||
also raise MediaExecutionUncertain and retain files, even after terminal history.
|
||||
|
||||
``temp_dir`` is a trusted endpoint setting for this ComfyUI's local temp root,
|
||||
never a request-supplied path. Only validated exporter UUID.bin files beneath
|
||||
distributed-media-v1/sha256(request_id) are retired, after all collection and
|
||||
publisher work has unwound. Unidentified files remain lifecycle-owned. With
|
||||
``temp_dir=None`` all exporter files remain lifecycle-owned as before.
|
||||
|
||||
The pinned native /interrupt route has a check/global-interrupt race. Running
|
||||
prompts therefore fail closed with MediaExecutionUncertain, without posting
|
||||
/interrupt. A prompt-scoped request body alone is not a safety guarantee.
|
||||
"""
|
||||
if (type(timeout_seconds) not in (int, float) or not math.isfinite(timeout_seconds)
|
||||
or timeout_seconds <= 0):
|
||||
raise ValueError("Invalid media execution timeout")
|
||||
if not isinstance(comfy_url, str):
|
||||
raise ValueError("Invalid local Comfy URL")
|
||||
try:
|
||||
parsed = urlsplit(comfy_url)
|
||||
if (parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username
|
||||
or parsed.password or parsed.query or parsed.fragment):
|
||||
raise ValueError()
|
||||
parsed.port
|
||||
except ValueError:
|
||||
raise ValueError("Invalid local Comfy URL") from None
|
||||
deadline = asyncio.get_running_loop().time() + timeout_seconds
|
||||
workflow, media, legacy, targets = _validate_request(input_data)
|
||||
stage = _Staging(input_dir, media["request_id"])
|
||||
exports = _OwnedExports(temp_dir, media["request_id"])
|
||||
owned_session = session is None
|
||||
if owned_session:
|
||||
session = aiohttp.ClientSession()
|
||||
execution = _Execution(session, comfy_url.rstrip("/"), stage, media, targets, deadline, exports)
|
||||
work = None
|
||||
try:
|
||||
remaining = execution.remaining()
|
||||
work = asyncio.create_task(_run(execution, workflow, legacy, download, publish))
|
||||
return await asyncio.wait_for(asyncio.shield(work), remaining)
|
||||
except BaseException as exc:
|
||||
# Retire files only after readers unwind. Cooperative async callbacks
|
||||
# normally drain immediately; an uncooperative one requires quarantine.
|
||||
if work is not None and not await _drain_work(work):
|
||||
raise MediaExecutionUncertain(
|
||||
"Media callback teardown unconfirmed; owned files retained; "
|
||||
"quarantine this endpoint before reuse; do not resubmit automatically"
|
||||
) from exc
|
||||
if execution.attempted and not execution.terminal:
|
||||
execution.deadline = asyncio.get_running_loop().time() + CANCEL_GRACE_SECONDS
|
||||
recovery = asyncio.create_task(asyncio.wait_for(execution.settle(), CANCEL_GRACE_SECONDS))
|
||||
# Shield the bounded recovery even from repeated task.cancel() calls.
|
||||
while not recovery.done():
|
||||
try:
|
||||
await asyncio.shield(recovery)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
except Exception:
|
||||
break
|
||||
try:
|
||||
recovery.result()
|
||||
except BaseException:
|
||||
pass
|
||||
if not execution.terminal:
|
||||
raise MediaExecutionUncertain(
|
||||
"Local Comfy prompt terminal state unconfirmed; owned staged files retained "
|
||||
f"under {stage.relative}; prompt_id={execution.prompt_id or 'unknown'}; "
|
||||
"do not resubmit this request automatically"
|
||||
+ ("; native prompt-scoped interruption unavailable (non-atomic /interrupt route)"
|
||||
if execution.interruption_blocked else "")
|
||||
) from exc
|
||||
raise
|
||||
finally:
|
||||
try:
|
||||
try:
|
||||
exports.close(cleanup=(work is None or work.done()) and execution.terminal)
|
||||
finally:
|
||||
stage.close(cleanup=(work is None or work.done()) and
|
||||
(not execution.attempted or execution.terminal))
|
||||
finally:
|
||||
if owned_session:
|
||||
await session.close()
|
||||
@@ -0,0 +1,634 @@
|
||||
"""Bounded optional IMAGE returns, independent of ComfyUI and Distributed.
|
||||
|
||||
Exact mode is the existing tensor wire format, unchanged. All other modes clip
|
||||
and truncate RGB/luma to uint8 (as the Collector does), normalize in float32,
|
||||
then cast to the source dtype. RGBA always carries an exact tensor alpha
|
||||
sidecar, including PNG/WebP. ``legacy_png`` means uncompressed binary PNG here,
|
||||
not the old Collector's JSON/base64 transport.
|
||||
|
||||
The compressed wire format concatenates length-delimited stills or one Matroska
|
||||
video, followed by optional alpha bytes. No archive, path, URL or pickle is
|
||||
accepted. Limits cover wire bytes, descriptor structure, output tensor bytes,
|
||||
and aggregate padded pixel buffers, not process RSS/native codec overhead.
|
||||
Video preflight reads EBML track dimensions and block counts before creating a
|
||||
decoder. We extract packets ourselves, avoiding native container stream probing
|
||||
(which can decode before caller checks). FFmpeg max_pixels also bounds
|
||||
bitstream-driven resolution changes, including native buffer alignment.
|
||||
Software codecs only, one thread, no fallback to another encoder/format.
|
||||
"""
|
||||
|
||||
import io
|
||||
import struct
|
||||
import zlib
|
||||
from fractions import Fraction
|
||||
|
||||
import torch
|
||||
|
||||
from . import codec
|
||||
from .codec import MAX_PAYLOAD_BYTES
|
||||
from .image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
|
||||
|
||||
MAX_FRAMES = 240
|
||||
MAX_DIMENSION = 8192
|
||||
MAX_DESCRIPTOR_BYTES = 64 * 1024
|
||||
MAX_CONTAINER_ELEMENTS = 16384
|
||||
_VIDEO_CODECS = {
|
||||
"ffv1": ("ffv1", None, "bgr0"),
|
||||
"h264": ("libx264", 51, "yuv420p"),
|
||||
"h265": ("libx265", 51, "yuv420p"),
|
||||
"av1": ("libsvtav1", 63, "yuv420p"),
|
||||
}
|
||||
_CODEC_IDS = {"ffv1": b"V_FFV1", "h264": b"V_MPEG4/ISO/AVC",
|
||||
"h265": b"V_MPEGH/ISO/HEVC", "av1": b"V_AV1"}
|
||||
|
||||
|
||||
class _BoundedIO(io.BytesIO):
|
||||
"""Enforce the final file size on every write, including muxer flush/seek."""
|
||||
def __init__(self, limit):
|
||||
super().__init__()
|
||||
self.limit = limit
|
||||
|
||||
def write(self, data):
|
||||
if self.tell() + len(data) > self.limit:
|
||||
raise ValueError("Encoded IMAGE exceeds byte limit")
|
||||
return super().write(data)
|
||||
|
||||
def seek(self, offset, whence=0):
|
||||
if whence == 0:
|
||||
position = offset
|
||||
elif whence == 1:
|
||||
position = self.tell() + offset
|
||||
elif whence == 2:
|
||||
position = self.getbuffer().nbytes + offset
|
||||
else:
|
||||
raise ValueError("Invalid image buffer seek")
|
||||
if not 0 <= position <= self.limit:
|
||||
raise ValueError("Encoded IMAGE exceeds byte limit")
|
||||
return super().seek(offset, whence)
|
||||
|
||||
|
||||
def _quality_to_crf(quality, maximum):
|
||||
return int(round(maximum * (1.0 - quality / 100.0) ** 0.6))
|
||||
|
||||
|
||||
def _padded_size(fmt, height, width):
|
||||
if fmt in ("h264", "h265", "av1"):
|
||||
height, width = height + height % 2, width + width % 2
|
||||
# libx265 / SVT-AV1 cannot safely encode arbitrarily small pictures.
|
||||
if fmt in ("h265", "av1"):
|
||||
height, width = max(64, height), max(64, width)
|
||||
return height, width
|
||||
|
||||
|
||||
def _shape_budget(shape, fmt, limit, itemsize=4, alpha_size=0, coded_frames=None):
|
||||
size = codec._tensor_size("IMAGE", shape, itemsize, limit)
|
||||
count, height, width, _ = shape
|
||||
if count > MAX_FRAMES or max(height, width) > MAX_DIMENSION:
|
||||
raise ValueError("IMAGE frame count or dimensions exceed limit")
|
||||
ph, pw = _padded_size(fmt, height, width)
|
||||
# Reserve output and padded RGBA-sized decoded buffers BEFORE allocation.
|
||||
bh, bw = ((ph + 63) // 64 * 64, (pw + 63) // 64 * 64) if fmt in _VIDEO_CODECS else (ph, pw)
|
||||
# Slice/OBU counts conservatively bound native frames, including AV1 hidden
|
||||
# frames and multiple outputs from one packet. Reserve 16 bytes per pixel
|
||||
# for hostile high-bit-depth streams, not only our usual 8-bit RGB frames.
|
||||
native_count = max(count, coded_frames or count)
|
||||
if native_count > MAX_FRAMES * 4:
|
||||
raise ValueError("IMAGE coded frame count exceeds limit")
|
||||
pixels = codec._tensor_size("IMAGE", [native_count, bh, bw, 4], 4, limit)
|
||||
if size + pixels + alpha_size > limit:
|
||||
raise ValueError("IMAGE decoded aggregate exceeds byte limit")
|
||||
return ph, pw
|
||||
|
||||
|
||||
def _descriptor_budget(value):
|
||||
"""Bound the JSON-shaped tree without serializing attacker-sized values."""
|
||||
remaining = MAX_DESCRIPTOR_BYTES
|
||||
def walk(obj, depth):
|
||||
nonlocal remaining
|
||||
if depth > 5:
|
||||
raise ValueError("IMAGE descriptor nesting exceeds limit")
|
||||
remaining -= 8
|
||||
if remaining < 0:
|
||||
raise ValueError("IMAGE descriptor exceeds byte limit")
|
||||
if obj is None:
|
||||
return
|
||||
if type(obj) is str:
|
||||
remaining -= len(obj) * 6
|
||||
elif type(obj) is int:
|
||||
if obj.bit_length() > 63:
|
||||
raise ValueError("IMAGE descriptor integer exceeds limit")
|
||||
remaining -= 20
|
||||
elif type(obj) is list:
|
||||
if len(obj) > MAX_FRAMES:
|
||||
raise ValueError("IMAGE descriptor list exceeds limit")
|
||||
for item in obj:
|
||||
walk(item, depth + 1)
|
||||
elif type(obj) is dict:
|
||||
if len(obj) > 10:
|
||||
raise ValueError("IMAGE descriptor fields exceed limit")
|
||||
for key, item in obj.items():
|
||||
if type(key) is not str:
|
||||
raise ValueError("Invalid IMAGE descriptor key")
|
||||
walk(key, depth + 1)
|
||||
walk(item, depth + 1)
|
||||
else:
|
||||
raise ValueError("Invalid IMAGE descriptor value")
|
||||
if remaining < 0:
|
||||
raise ValueError("IMAGE descriptor exceeds byte limit")
|
||||
walk(value, 0)
|
||||
|
||||
|
||||
def _alpha_metadata(alpha, shape, limit):
|
||||
if shape[-1] != 4:
|
||||
if alpha is not None:
|
||||
raise ValueError("Unexpected IMAGE alpha sidecar")
|
||||
return 0, torch.float32
|
||||
if type(alpha) is not dict or set(alpha) != {"descriptor", "size"}:
|
||||
raise ValueError("Missing or invalid IMAGE alpha sidecar")
|
||||
d = alpha["descriptor"]
|
||||
if type(d) is not dict or set(d) != {"encoding", "dtype", "shape", "layout"}:
|
||||
raise ValueError("Invalid IMAGE alpha descriptor")
|
||||
if d["encoding"] != "tensor-le-v1" or d["layout"] != "BHWC" or d["shape"] != shape[:3] + [1]:
|
||||
raise ValueError("Invalid IMAGE alpha shape or encoding")
|
||||
if type(d["dtype"]) is not str or d["dtype"] not in codec._DTYPES:
|
||||
raise ValueError("Invalid IMAGE alpha dtype")
|
||||
dtype, itemsize = codec._DTYPES[d["dtype"]]
|
||||
size = codec._tensor_size("IMAGE", d["shape"], itemsize, limit)
|
||||
if type(alpha["size"]) is not int or alpha["size"] != size:
|
||||
raise ValueError("Invalid IMAGE alpha size")
|
||||
return size, torch.float64 if dtype == torch.float64 else torch.float32
|
||||
|
||||
|
||||
def _import_av():
|
||||
try:
|
||||
import av
|
||||
except ImportError:
|
||||
raise ValueError("PyAV is required for IMAGE video transfer") from None
|
||||
return av
|
||||
|
||||
|
||||
def _encode_still(frame, fmt, quality, limit):
|
||||
from PIL import Image
|
||||
sink = _BoundedIO(limit)
|
||||
if frame.shape[-1] == 1:
|
||||
frame = frame[..., 0]
|
||||
image = Image.fromarray(frame)
|
||||
if fmt in ("legacy_png", "png"):
|
||||
image.save(sink, format="PNG", compress_level=0 if fmt == "legacy_png" else 4)
|
||||
elif fmt == "jpeg":
|
||||
image.save(sink, format="JPEG", quality=quality, subsampling=0 if quality >= 90 else 2)
|
||||
else:
|
||||
image.save(sink, format="WEBP", quality=quality, method=4)
|
||||
return sink.getvalue()
|
||||
|
||||
|
||||
def _encode_video(frames, fmt, quality, height, width, limit):
|
||||
import numpy as np
|
||||
av = _import_av()
|
||||
encoder, maximum, pix_fmt = _VIDEO_CODECS[fmt]
|
||||
try:
|
||||
av.codec.Codec(encoder, "w")
|
||||
except Exception:
|
||||
raise ValueError("IMAGE software encoder unavailable: " + encoder) from None
|
||||
sink = _BoundedIO(limit)
|
||||
options = {}
|
||||
if maximum is not None:
|
||||
options["crf"] = str(_quality_to_crf(quality, maximum))
|
||||
if fmt == "h264":
|
||||
options["preset"] = "medium"
|
||||
elif fmt == "h265":
|
||||
options.update(preset="fast", **{"x265-params": "log-level=error:pools=none:frame-threads=1"})
|
||||
elif fmt == "av1":
|
||||
options.update(preset="8", **{"svtav1-params": "lp=1"})
|
||||
else:
|
||||
options.update(level="3", slicecrc="1")
|
||||
try:
|
||||
with av.open(sink, mode="w", format="matroska") as container:
|
||||
stream = container.add_stream(encoder, rate=24)
|
||||
stream.width, stream.height = width, height
|
||||
stream.pix_fmt = pix_fmt
|
||||
stream.codec_context.thread_count = 1
|
||||
stream.options = options
|
||||
for index, frame in enumerate(frames):
|
||||
if frame.shape[-1] == 1:
|
||||
frame = np.repeat(frame, 3, axis=-1)
|
||||
h, w = frame.shape[:2]
|
||||
frame = np.pad(frame, ((0, height - h), (0, width - w), (0, 0)), mode="edge")
|
||||
vf = av.VideoFrame.from_ndarray(frame, format="rgb24")
|
||||
vf.pts = index
|
||||
vf.time_base = Fraction(1, 24)
|
||||
for packet in stream.encode(vf):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode():
|
||||
container.mux(packet)
|
||||
data = sink.getvalue()
|
||||
if len(data) > limit:
|
||||
raise ValueError("Encoded IMAGE exceeds byte limit")
|
||||
return data
|
||||
except av.error.FFmpegError:
|
||||
raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None
|
||||
except ValueError:
|
||||
raise
|
||||
except Exception:
|
||||
raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None
|
||||
|
||||
|
||||
def encode_image(value, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY,
|
||||
max_bytes=MAX_PAYLOAD_BYTES):
|
||||
"""Return (descriptor, bytes); reject unsupported inputs, never fall back."""
|
||||
fmt, quality = validate_options(transfer_format, transfer_quality)
|
||||
if fmt == "exact":
|
||||
return codec.encode_payload("IMAGE", value, max_bytes=max_bytes)
|
||||
limit = codec._byte_limit(max_bytes)
|
||||
if not isinstance(value, torch.Tensor) or value.device.type != "cpu":
|
||||
raise ValueError("IMAGE must be a CPU tensor")
|
||||
if value.layout != torch.strided or value.is_nested:
|
||||
raise ValueError("IMAGE must have dense strided layout")
|
||||
if value.dtype not in codec._DTYPE_NAMES:
|
||||
raise ValueError("Unsupported IMAGE dtype")
|
||||
shape = list(value.shape)
|
||||
codec._tensor_size("IMAGE", shape, value.element_size(), limit)
|
||||
alpha_size = shape[0] * shape[1] * shape[2] * value.element_size() if shape[-1] == 4 else 0
|
||||
itemsize = value.element_size()
|
||||
ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size)
|
||||
if not torch.isfinite(value).all().item():
|
||||
raise ValueError("IMAGE values must be finite")
|
||||
value = value.detach().resolve_neg()
|
||||
alpha, alpha_payload = None, b""
|
||||
if shape[-1] == 4:
|
||||
ad, alpha_payload = codec.encode_payload("IMAGE", value[..., 3:4], max_bytes=limit)
|
||||
alpha = {"descriptor": ad, "size": len(alpha_payload)}
|
||||
import numpy as np
|
||||
def frames():
|
||||
for frame in value:
|
||||
arr = frame[..., :3].float().numpy()
|
||||
with np.errstate(over="ignore", invalid="ignore"):
|
||||
yield np.ascontiguousarray(np.clip(arr * 255.0, 0, 255).astype(np.uint8))
|
||||
budget = limit - len(alpha_payload)
|
||||
if fmt in _VIDEO_CODECS:
|
||||
parts = [_encode_video(frames(), fmt, quality, ph, pw, budget)]
|
||||
header = _video_header(parts[0], fmt, ph, pw, shape[0])
|
||||
coded_frames = _coded_frame_bound(parts[0], fmt, header)
|
||||
_shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames)
|
||||
else:
|
||||
parts = []
|
||||
for frame in frames():
|
||||
part = _encode_still(frame, fmt, quality, budget)
|
||||
budget -= len(part)
|
||||
parts.append(part)
|
||||
descriptor = {"encoding": "image-transfer-v1", "format": fmt, "quality": quality,
|
||||
"shape": shape, "layout": "BHWC", "dtype": codec._DTYPE_NAMES[value.dtype],
|
||||
"parts": [len(p) for p in parts], "alpha": alpha}
|
||||
_descriptor_budget(descriptor)
|
||||
payload = b"".join(parts) + alpha_payload
|
||||
if len(payload) > limit:
|
||||
raise ValueError("Encoded IMAGE exceeds byte limit")
|
||||
return descriptor, payload
|
||||
|
||||
|
||||
def _still_header(data, fmt):
|
||||
"""Read dimensions before Pillow (WebP may allocate pixels during open)."""
|
||||
if fmt in ("png", "legacy_png"):
|
||||
if data[:8] != b"\x89PNG\r\n\x1a\n":
|
||||
raise ValueError("Invalid IMAGE PNG header")
|
||||
offset, dimensions, saw_data, chunks = 8, None, False, 0
|
||||
while offset < len(data):
|
||||
chunks += 1
|
||||
if chunks > MAX_CONTAINER_ELEMENTS:
|
||||
raise ValueError("IMAGE PNG chunk count exceeds limit")
|
||||
if offset + 12 > len(data):
|
||||
raise ValueError("Truncated IMAGE PNG chunk")
|
||||
length = int.from_bytes(data[offset:offset + 4], "big")
|
||||
end = offset + 12 + length
|
||||
if end > len(data):
|
||||
raise ValueError("Truncated IMAGE PNG chunk")
|
||||
tag = data[offset + 4:offset + 8]
|
||||
body = data[offset + 8:end - 4]
|
||||
if zlib.crc32(data[offset + 4:end - 4]) != int.from_bytes(data[end - 4:end], "big"):
|
||||
raise ValueError("Invalid IMAGE PNG checksum")
|
||||
if tag == b"IHDR" and offset == 8 and length == 13:
|
||||
width, height, depth, color, compression, filtering, interlace = struct.unpack(">IIBBBBB", body)
|
||||
if depth != 8 or color not in (0, 2) or (compression, filtering, interlace) != (0, 0, 0):
|
||||
raise ValueError("Unsupported IMAGE PNG layout")
|
||||
dimensions = (height, width)
|
||||
elif tag == b"IDAT" and dimensions:
|
||||
saw_data = True
|
||||
elif tag == b"IEND" and not length and end == len(data) and saw_data:
|
||||
return dimensions
|
||||
else:
|
||||
raise ValueError("Unsupported IMAGE PNG chunk")
|
||||
offset = end
|
||||
raise ValueError("Truncated IMAGE PNG")
|
||||
if fmt == "webp":
|
||||
# Our encoder writes a single lossy VP8 still. Reject animation/alpha/
|
||||
# extended WebP before WebPAnimDecoder can allocate its canvas.
|
||||
if len(data) < 30 or data[:4] != b"RIFF" or data[8:16] != b"WEBPVP8 ":
|
||||
raise ValueError("Unsupported IMAGE WebP header")
|
||||
size = int.from_bytes(data[16:20], "little")
|
||||
if int.from_bytes(data[4:8], "little") + 8 != len(data) or 20 + size + size % 2 != len(data):
|
||||
raise ValueError("Invalid IMAGE WebP length")
|
||||
if data[23:26] != b"\x9d\x01\x2a":
|
||||
raise ValueError("Invalid IMAGE WebP frame")
|
||||
width, height = struct.unpack("<HH", data[26:30])
|
||||
return height & 0x3fff, width & 0x3fff
|
||||
if data[:2] != b"\xff\xd8" or data[-2:] != b"\xff\xd9":
|
||||
raise ValueError("Invalid IMAGE JPEG header or trailer")
|
||||
offset, dimensions = 2, None
|
||||
while offset + 4 <= len(data):
|
||||
if data[offset] != 0xff:
|
||||
raise ValueError("Invalid IMAGE JPEG marker")
|
||||
marker = data[offset + 1]
|
||||
size = int.from_bytes(data[offset + 2:offset + 4], "big")
|
||||
end = offset + 2 + size
|
||||
if size < 2 or end > len(data):
|
||||
raise ValueError("Truncated IMAGE JPEG marker")
|
||||
if marker in (0xc0, 0xc2):
|
||||
if dimensions or size < 8 or data[offset + 4] != 8 or data[offset + 9] not in (1, 3):
|
||||
raise ValueError("Unsupported IMAGE JPEG layout")
|
||||
dimensions = struct.unpack(">HH", data[offset + 5:offset + 9])
|
||||
elif marker == 0xda:
|
||||
if dimensions is None:
|
||||
raise ValueError("Missing IMAGE JPEG dimensions")
|
||||
return dimensions
|
||||
offset = end
|
||||
raise ValueError("Missing IMAGE JPEG pixels")
|
||||
|
||||
|
||||
def _video_header(data, fmt, height, width, count):
|
||||
"""Strict bounded EBML preflight: one track, no lacing/attachments/encodings.
|
||||
|
||||
Native stream probing is not the first dimension check: it may instantiate a
|
||||
decoder. Validate the Matroska header and count blocks without decompressing.
|
||||
"""
|
||||
nodes = 0
|
||||
def vint(pos, end, identifier=False):
|
||||
if pos >= end or data[pos] == 0:
|
||||
raise ValueError("Invalid IMAGE EBML integer")
|
||||
length = 9 - data[pos].bit_length()
|
||||
if length > (4 if identifier else 8) or pos + length > end:
|
||||
raise ValueError("Truncated IMAGE EBML integer")
|
||||
value = int.from_bytes(data[pos:pos + length], "big")
|
||||
if not identifier:
|
||||
value &= (1 << (7 * length)) - 1
|
||||
if value == (1 << (7 * length)) - 1:
|
||||
raise ValueError("Unknown-sized IMAGE EBML elements are forbidden")
|
||||
return value, pos + length
|
||||
def elements(start, end):
|
||||
nonlocal nodes
|
||||
while start < end:
|
||||
nodes += 1
|
||||
if nodes > MAX_CONTAINER_ELEMENTS:
|
||||
raise ValueError("IMAGE container element count exceeds limit")
|
||||
tag, pos = vint(start, end, True)
|
||||
length, body = vint(pos, end)
|
||||
finish = body + length
|
||||
if finish > end:
|
||||
raise ValueError("Truncated IMAGE EBML element")
|
||||
yield tag, body, finish
|
||||
start = finish
|
||||
def integer(start, end):
|
||||
if not 1 <= end - start <= 8:
|
||||
raise ValueError("Invalid IMAGE EBML numeric field")
|
||||
return int.from_bytes(data[start:end], "big")
|
||||
roots = list(elements(0, len(data)))
|
||||
if len(roots) != 2 or [r[0] for r in roots] != [0x1a45dfa3, 0x18538067]:
|
||||
raise ValueError("Invalid IMAGE Matroska structure")
|
||||
tracks, blocks, packets = [], 0, []
|
||||
for tag, start, end in elements(roots[1][1], roots[1][2]):
|
||||
if tag == 0x1654ae6b:
|
||||
for entry, es, ee in elements(start, end):
|
||||
if entry in (0xbf, 0xec):
|
||||
continue
|
||||
if entry != 0xae:
|
||||
raise ValueError("Unexpected IMAGE track element")
|
||||
fields, dims = {}, {}
|
||||
for field, fs, fe in elements(es, ee):
|
||||
if field in fields:
|
||||
raise ValueError("Duplicate IMAGE track field")
|
||||
fields[field] = (fs, fe)
|
||||
if field == 0x6d80:
|
||||
raise ValueError("IMAGE track encodings are forbidden")
|
||||
if field == 0xe0:
|
||||
for dim, ds, de in elements(fs, fe):
|
||||
if dim in dims:
|
||||
raise ValueError("Duplicate IMAGE video field")
|
||||
dims[dim] = (ds, de)
|
||||
if any(key not in fields for key in (0xd7, 0x83, 0x86)) or any(key not in dims for key in (0xb0, 0xba)):
|
||||
raise ValueError("Incomplete IMAGE video track")
|
||||
if integer(*fields[0xd7]) != 1 or integer(*fields[0x83]) != 1:
|
||||
raise ValueError("Unsupported IMAGE video track")
|
||||
if fields[0x86][1] - fields[0x86][0] > 32:
|
||||
raise ValueError("IMAGE codec identifier exceeds limit")
|
||||
codec_id = bytes(data[slice(*fields[0x86])])
|
||||
if codec_id != _CODEC_IDS[fmt] and not (fmt == "ffv1" and codec_id == b"V_MS/VFW/FOURCC"):
|
||||
raise ValueError("IMAGE video codec mismatch")
|
||||
if integer(*dims[0xb0]) != width or integer(*dims[0xba]) != height:
|
||||
raise ValueError("IMAGE video header dimensions mismatch")
|
||||
if codec_id == b"V_MS/VFW/FOURCC":
|
||||
cs, ce = fields.get(0x63a2, (0, 0))
|
||||
if ce - cs < 40 or data[cs + 16:cs + 20] != b"FFV1":
|
||||
raise ValueError("Invalid IMAGE FFV1 codec header")
|
||||
if struct.unpack("<ii", data[cs + 4:cs + 12]) != (width, height):
|
||||
raise ValueError("IMAGE FFV1 header dimensions mismatch")
|
||||
tracks.append(fields)
|
||||
elif tag == 0x1f43b675:
|
||||
def block(bs, be):
|
||||
nonlocal blocks
|
||||
track, pos = vint(bs, be)
|
||||
if track != 1 or pos + 3 >= be or data[pos + 2] & 0x06:
|
||||
raise ValueError("Invalid or laced IMAGE video block")
|
||||
blocks += 1
|
||||
if blocks > count:
|
||||
raise ValueError("IMAGE video frame count exceeds descriptor")
|
||||
packets.append((pos + 3, be))
|
||||
for field, fs, fe in elements(start, end):
|
||||
if field == 0xa3:
|
||||
block(fs, fe)
|
||||
elif field == 0xa0:
|
||||
group = [e for e in elements(fs, fe) if e[0] == 0xa1]
|
||||
if len(group) != 1:
|
||||
raise ValueError("Invalid IMAGE video block group")
|
||||
block(group[0][1], group[0][2])
|
||||
elif field not in (0xe7, 0xec, 0xbf):
|
||||
raise ValueError("Unsupported IMAGE cluster field")
|
||||
elif tag not in (0x114d9b74, 0x1549a966, 0x1c53bb6b, 0x1254c367, 0xec, 0xbf):
|
||||
raise ValueError("Unsupported IMAGE Matroska element")
|
||||
if len(tracks) != 1 or blocks != count:
|
||||
raise ValueError("IMAGE video track or frame count mismatch")
|
||||
cs, ce = tracks[0].get(0x63a2, (0, 0))
|
||||
if not 0 < ce - cs <= MAX_DESCRIPTOR_BYTES:
|
||||
raise ValueError("IMAGE codec configuration exceeds limit or is missing")
|
||||
if bytes(data[slice(*tracks[0][0x86])]) == b"V_MS/VFW/FOURCC":
|
||||
cs += 40 # Skip BITMAPINFOHEADER; FFV1 receives only its codec config.
|
||||
return bytes(data[cs:ce]), packets
|
||||
|
||||
|
||||
def _coded_frame_bound(data, fmt, header):
|
||||
"""Conservative per-packet allocation bound before native decode.
|
||||
|
||||
FFV1 has one picture per packet. Every AVC/HEVC picture needs at least one
|
||||
VCL NAL; every AV1 output (including show-existing) needs a frame/header OBU.
|
||||
Counting slices overestimates pictures, which is safe. Do not trust a
|
||||
Matroska block count alone: a malicious packet may contain several pictures.
|
||||
"""
|
||||
config, packets = header
|
||||
if fmt == "ffv1":
|
||||
return len(packets)
|
||||
total, units = 0, 0
|
||||
if fmt in ("h264", "h265"):
|
||||
index = 4 if fmt == "h264" else 21
|
||||
if len(config) <= index or config[0] != 1:
|
||||
raise ValueError("Invalid IMAGE video configuration")
|
||||
length_size = (config[index] & 3) + 1
|
||||
for start, end in packets:
|
||||
while start < end:
|
||||
units += 1
|
||||
if units > MAX_CONTAINER_ELEMENTS:
|
||||
raise ValueError("IMAGE coded unit count exceeds limit")
|
||||
if fmt in ("h264", "h265"):
|
||||
if start + length_size > end:
|
||||
raise ValueError("Truncated IMAGE NAL length")
|
||||
length = int.from_bytes(data[start:start + length_size], "big")
|
||||
start += length_size
|
||||
if length < (1 if fmt == "h264" else 2) or length > end - start:
|
||||
raise ValueError("Truncated IMAGE NAL unit")
|
||||
kind = (data[start] & 31) if fmt == "h264" else (data[start] >> 1) & 63
|
||||
total += (1 <= kind <= 5) if fmt == "h264" else (kind <= 31)
|
||||
else:
|
||||
tag = data[start]
|
||||
start += 1
|
||||
if tag & 0x81 or not tag & 2:
|
||||
raise ValueError("Unsupported IMAGE AV1 OBU header")
|
||||
kind = (tag >> 3) & 15
|
||||
if tag & 4:
|
||||
start += 1
|
||||
length = 0
|
||||
for shift in range(0, 56, 7):
|
||||
if start >= end:
|
||||
raise ValueError("Truncated IMAGE AV1 OBU length")
|
||||
value = data[start]
|
||||
start += 1
|
||||
length |= (value & 127) << shift
|
||||
if not value & 128:
|
||||
break
|
||||
else:
|
||||
raise ValueError("IMAGE AV1 OBU length exceeds limit")
|
||||
if length > end - start:
|
||||
raise ValueError("Truncated IMAGE AV1 OBU")
|
||||
total += kind in (3, 6)
|
||||
start += length
|
||||
if total > MAX_FRAMES * 4:
|
||||
raise ValueError("IMAGE coded frame count exceeds limit")
|
||||
if total < len(packets):
|
||||
raise ValueError("IMAGE video has missing coded frames")
|
||||
return total
|
||||
|
||||
|
||||
def _decode_video(data, fmt, height, width, count, header):
|
||||
av = _import_av()
|
||||
maximum = ((height + 63) // 64 * 64) * ((width + 63) // 64 * 64)
|
||||
ctx = av.CodecContext.create("libdav1d" if fmt == "av1" else "hevc" if fmt == "h265" else fmt, "r")
|
||||
ctx.width, ctx.height = width, height
|
||||
ctx.extradata, packets = header
|
||||
ctx.thread_count = 1
|
||||
ctx.options = {"max_pixels": str(maximum), "threads": "1", "err_detect": "explode"}
|
||||
ctx.open()
|
||||
if "max_pixels" in ctx.options:
|
||||
raise ValueError("IMAGE decoder does not support max_pixels safety limit")
|
||||
if (ctx.width, ctx.height) != (width, height):
|
||||
raise ValueError("IMAGE bitstream dimensions mismatch")
|
||||
decoded = 0
|
||||
for bounds in packets + [None]:
|
||||
packet = av.Packet(data[slice(*bounds)]) if bounds is not None else None
|
||||
for frame in ctx.decode(packet):
|
||||
if decoded >= count or (frame.width, frame.height) != (width, height):
|
||||
raise ValueError("IMAGE decoded dimensions or frame count mismatch")
|
||||
decoded += 1
|
||||
yield frame.to_ndarray(format="rgb24")
|
||||
if decoded != count:
|
||||
raise ValueError("IMAGE decoded frame count mismatch")
|
||||
|
||||
|
||||
def decode_image(descriptor, payload, *, max_bytes=MAX_PAYLOAD_BYTES):
|
||||
"""Decode to an independent contiguous CPU tensor after bounded preflight."""
|
||||
limit = codec._byte_limit(max_bytes)
|
||||
if not isinstance(descriptor, dict):
|
||||
raise ValueError("IMAGE descriptor must be a dictionary")
|
||||
if descriptor.get("encoding") == "tensor-le-v1":
|
||||
return codec.decode_payload("IMAGE", descriptor, payload, max_bytes=limit)
|
||||
_descriptor_budget(descriptor)
|
||||
if set(descriptor) != {"encoding", "format", "quality", "shape", "layout", "dtype", "parts", "alpha"}:
|
||||
raise ValueError("Invalid IMAGE descriptor fields")
|
||||
if descriptor["encoding"] != "image-transfer-v1" or descriptor["layout"] != "BHWC":
|
||||
raise ValueError("Invalid IMAGE descriptor encoding or layout")
|
||||
fmt, _ = validate_options(descriptor["format"], descriptor["quality"])
|
||||
if fmt == "exact":
|
||||
raise ValueError("Exact IMAGE requires the tensor encoding")
|
||||
if not isinstance(payload, bytes):
|
||||
raise ValueError("IMAGE payload must be bytes")
|
||||
if len(payload) > limit:
|
||||
raise ValueError("IMAGE payload exceeds byte limit")
|
||||
shape = descriptor["shape"]
|
||||
if type(descriptor["dtype"]) is not str or descriptor["dtype"] not in codec._DTYPES:
|
||||
raise ValueError("Unsupported IMAGE descriptor dtype")
|
||||
dtype, itemsize = codec._DTYPES[descriptor["dtype"]]
|
||||
codec._tensor_size("IMAGE", shape, itemsize, limit)
|
||||
alpha_size, _ = _alpha_metadata(descriptor["alpha"], shape, limit)
|
||||
if alpha_size and descriptor["alpha"]["descriptor"]["dtype"] != descriptor["dtype"]:
|
||||
raise ValueError("IMAGE alpha dtype mismatch")
|
||||
ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size)
|
||||
count, height, width, channels = shape
|
||||
lengths = descriptor["parts"]
|
||||
if type(lengths) is not list or len(lengths) != (1 if fmt in _VIDEO_CODECS else count):
|
||||
raise ValueError("IMAGE part count mismatch")
|
||||
total = alpha_size
|
||||
for length in lengths:
|
||||
if type(length) is not int or not 0 < length <= limit - total:
|
||||
raise ValueError("IMAGE part length exceeds byte limit")
|
||||
total += length
|
||||
if total != len(payload):
|
||||
raise ValueError("IMAGE payload length does not match parts")
|
||||
parts, offset = [], 0
|
||||
for length in lengths:
|
||||
parts.append(memoryview(payload)[offset:offset + length])
|
||||
offset += length
|
||||
try:
|
||||
if fmt in _VIDEO_CODECS:
|
||||
video_header = _video_header(parts[0], fmt, ph, pw, count)
|
||||
coded_frames = _coded_frame_bound(parts[0], fmt, video_header)
|
||||
_shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames)
|
||||
else:
|
||||
for part in parts:
|
||||
if _still_header(part, fmt) != (height, width):
|
||||
raise ValueError("IMAGE header dimensions mismatch")
|
||||
# All headers and aggregate budgets are checked before output allocation.
|
||||
alpha = None
|
||||
if alpha_size:
|
||||
alpha = codec.decode_payload("IMAGE", descriptor["alpha"]["descriptor"], payload[offset:], max_bytes=limit)
|
||||
result = torch.empty(shape, dtype=dtype, device="cpu")
|
||||
if fmt in _VIDEO_CODECS:
|
||||
frames = _decode_video(parts[0], fmt, ph, pw, count, video_header)
|
||||
else:
|
||||
def stills():
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
for part in parts:
|
||||
with Image.open(io.BytesIO(part)) as image:
|
||||
if image.size != (width, height) or getattr(image, "n_frames", 1) != 1:
|
||||
raise ValueError("IMAGE still dimensions or frame count mismatch")
|
||||
if image.mode not in ("L", "RGB"):
|
||||
raise ValueError("Unsupported IMAGE still pixel mode")
|
||||
yield np.array(image.convert("RGB"), dtype=np.uint8)
|
||||
frames = stills()
|
||||
color = 3 if channels == 4 else channels
|
||||
for index, frame in enumerate(frames):
|
||||
# Float32 division followed by the source dtype cast reproduces
|
||||
# Collector normalized pixels without narrowing exact alpha.
|
||||
result[index, ..., :color] = torch.from_numpy(frame[:height, :width, :color].copy()).float().div_(255)
|
||||
if alpha is not None:
|
||||
result[..., 3:4] = alpha
|
||||
return result
|
||||
except ValueError:
|
||||
raise
|
||||
except Exception:
|
||||
raise ValueError("Malformed or unsupported IMAGE compressed payload") from None
|
||||
@@ -0,0 +1,30 @@
|
||||
"""Image-return policy shared by the master and standalone endpoint bridge.
|
||||
|
||||
This module is deliberately dependency-free so graph compilation does not import
|
||||
Torch, Pillow, PyAV or ComfyUI. Existing Collector settings stay unchanged.
|
||||
"""
|
||||
|
||||
FORMATS = ("exact", "legacy_png", "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1")
|
||||
DEFAULT_FORMAT = "exact"
|
||||
DEFAULT_QUALITY = 90
|
||||
FORMAT_TOOLTIP = (
|
||||
"Image returns only; master-local images and other media stay unchanged. "
|
||||
"exact: original floating-point tensor, no quantization. "
|
||||
"legacy_png / png / ffv1: lossless 8-bit pixels, not exact float tensors. "
|
||||
"jpeg / webp: lossy stills; h264 / h265 / av1: lossy frame-batch video. "
|
||||
"Alpha is preserved separately when the codec cannot carry it. "
|
||||
"Requires matching transfer support on the worker/endpoint."
|
||||
)
|
||||
QUALITY_TOOLTIP = (
|
||||
"Quality 1–100 for lossy image returns; ignored by exact and lossless modes. "
|
||||
"JPEG/WebP use it directly; video formats use the existing Collector CRF mapping."
|
||||
)
|
||||
|
||||
|
||||
def validate_options(transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
|
||||
"""Reject malformed controls rather than repairing or silently substituting."""
|
||||
if type(transfer_format) is not str or transfer_format not in FORMATS:
|
||||
raise ValueError("Unsupported image transfer_format")
|
||||
if type(transfer_quality) is not int or not 1 <= transfer_quality <= 100:
|
||||
raise ValueError("Image transfer_quality must be an integer from 1 to 100")
|
||||
return transfer_format, transfer_quality
|
||||
@@ -0,0 +1,175 @@
|
||||
"""Internal ComfyUI boundary nodes for the standalone media-v1 handler."""
|
||||
import hashlib
|
||||
import json
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import folder_paths
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
from .codec import decode_payload, encode_payload
|
||||
from .image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY,
|
||||
FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options)
|
||||
from .contract import (
|
||||
MAX_ARTIFACT_BYTES, MAX_ITEMS, MAX_TOTAL_BYTES, SOCKET_KINDS,
|
||||
bounded_int, identifier, validate_binding, validate_json,
|
||||
)
|
||||
|
||||
|
||||
def _scalar(value):
|
||||
if isinstance(value, list):
|
||||
if len(value) != 1:
|
||||
raise ValueError("Media bridge control must have exactly one value")
|
||||
return value[0]
|
||||
return value
|
||||
|
||||
|
||||
def _json(text):
|
||||
if not isinstance(text, str) or len(text) > 256 * 1024:
|
||||
raise ValueError("Invalid media bridge JSON")
|
||||
try:
|
||||
value = json.loads(text)
|
||||
except (ValueError, TypeError, RecursionError):
|
||||
raise ValueError("Invalid media bridge JSON") from None
|
||||
validate_json(value)
|
||||
return value
|
||||
|
||||
|
||||
def _cpu(kind, value):
|
||||
# Runtime transfer only, never sampling or device discovery. Pure codecs
|
||||
# deliberately reject non-CPU storage rather than making hidden transfers.
|
||||
if kind in ("IMAGE", "MASK"):
|
||||
return value.detach().cpu()
|
||||
if kind == "AUDIO":
|
||||
return {"waveform": value["waveform"].detach().cpu(), "sample_rate": value["sample_rate"]}
|
||||
return value
|
||||
|
||||
|
||||
def encode_value(kind, value, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
|
||||
fmt, quality = validate_options(transfer_format, transfer_quality)
|
||||
if kind == "IMAGE" and fmt != DEFAULT_FORMAT:
|
||||
from .image_codec import encode_image
|
||||
return encode_image(_cpu(kind, value), transfer_format=fmt, transfer_quality=quality)
|
||||
if kind == "VIDEO":
|
||||
from .video_codec import encode_video_cpu
|
||||
return encode_video_cpu(value)
|
||||
return encode_payload(kind, _cpu(kind, value))
|
||||
|
||||
|
||||
def decode_value(kind, descriptor, payload):
|
||||
if kind == "VIDEO":
|
||||
from .video_codec import decode_video
|
||||
return decode_video(descriptor, payload)
|
||||
return decode_payload(kind, descriptor, payload)
|
||||
|
||||
|
||||
def export_values(request_id, binding, values, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
|
||||
"""Write bounded owned payloads; history contains metadata, never base64."""
|
||||
identifier(request_id)
|
||||
validate_binding(binding)
|
||||
fmt, quality = validate_options(transfer_format, transfer_quality)
|
||||
if binding["kind"] != "IMAGE" and fmt != DEFAULT_FORMAT:
|
||||
raise ValueError("Image transfer controls apply only to IMAGE exports")
|
||||
if not isinstance(values, list) or not 1 <= len(values) <= MAX_ITEMS:
|
||||
raise ValueError("Connected media output has no values or exceeds item limit")
|
||||
request_key = hashlib.sha256(request_id.encode()).hexdigest()
|
||||
subfolder = f"distributed-media-v1/{request_key}"
|
||||
root = Path(folder_paths.get_temp_directory()).resolve()
|
||||
directory = (root / subfolder).resolve()
|
||||
if not directory.is_relative_to(root):
|
||||
raise ValueError("Unsafe media output directory")
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
written, artifacts, items = [], [], []
|
||||
total = 0
|
||||
try:
|
||||
for index, value in enumerate(values):
|
||||
descriptor, payload = encode_value(binding["kind"], value,
|
||||
transfer_format=fmt, transfer_quality=quality)
|
||||
validate_json(descriptor)
|
||||
total += len(payload)
|
||||
if total > MAX_TOTAL_BYTES:
|
||||
raise ValueError("Media output exceeds total byte limit")
|
||||
aid = uuid.uuid4().hex
|
||||
filename = aid + ".bin"
|
||||
path = directory / filename
|
||||
with path.open("xb") as stream:
|
||||
written.append(path)
|
||||
stream.write(payload)
|
||||
artifacts.append({"id": aid, "kind": binding["kind"], "descriptor": descriptor,
|
||||
"byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(),
|
||||
"transport": {"type": "comfy-file", "filename": filename,
|
||||
"subfolder": subfolder, "folder_type": "temp"}})
|
||||
items.append({"binding_id": binding["binding_id"], "item_index": index, "artifact_id": aid})
|
||||
except BaseException:
|
||||
for path in written:
|
||||
path.unlink(missing_ok=True)
|
||||
raise
|
||||
return {"request_id": request_id, "binding": dict(binding, count=len(values)),
|
||||
"artifacts": artifacts, "items": items}
|
||||
|
||||
|
||||
def import_value(artifact, filename):
|
||||
"""Read a handler-staged, checksummed file within native media roots only."""
|
||||
fields = {"id", "kind", "descriptor", "byte_length", "sha256"}
|
||||
if not isinstance(artifact, dict) or set(artifact) != fields:
|
||||
raise ValueError("Invalid staged media metadata")
|
||||
identifier(artifact["id"])
|
||||
if artifact["kind"] not in SOCKET_KINDS.values():
|
||||
raise ValueError("Invalid staged media kind")
|
||||
size = bounded_int(artifact["byte_length"], MAX_ARTIFACT_BYTES)
|
||||
if not isinstance(filename, str):
|
||||
raise ValueError("Invalid staged media filename")
|
||||
roots = [Path(getter()).resolve() for getter in (
|
||||
folder_paths.get_input_directory, folder_paths.get_output_directory, folder_paths.get_temp_directory)]
|
||||
path = Path(folder_paths.get_annotated_filepath(filename)).resolve()
|
||||
if not any(path.is_relative_to(root) for root in roots) or not path.is_file():
|
||||
raise ValueError("Staged media file is outside allowed roots or unavailable")
|
||||
if path.stat().st_size != size:
|
||||
raise ValueError("Staged media file length mismatch")
|
||||
with path.open("rb") as stream:
|
||||
payload = stream.read(size + 1)
|
||||
if len(payload) != size or hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
|
||||
raise ValueError("Staged media file failed integrity validation")
|
||||
return decode_value(artifact["kind"], artifact["descriptor"], payload)
|
||||
|
||||
|
||||
class MediaImportV1(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(node_id="MediaImportV1", category="distributed/internal", is_deprecated=True,
|
||||
inputs=[io.String.Input("artifact_json"), io.String.Input("filename")],
|
||||
outputs=[io.Image.Output(), io.Audio.Output(), io.Mask.Output(), io.String.Output(), io.Video.Output()])
|
||||
|
||||
@classmethod
|
||||
def execute(cls, artifact_json, filename):
|
||||
artifact = _json(artifact_json)
|
||||
value = import_value(artifact, filename)
|
||||
return io.NodeOutput(*(value if kind == artifact["kind"] else None for kind in SOCKET_KINDS.values()))
|
||||
|
||||
|
||||
class MediaExportV1(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(node_id="MediaExportV1", category="distributed/internal", is_deprecated=True,
|
||||
inputs=[io.String.Input("request_id"), io.String.Input("binding_json"),
|
||||
io.Image.Input("images", optional=True), io.Audio.Input("audio", optional=True),
|
||||
io.Mask.Input("masks", optional=True), io.String.Input("texts", optional=True, force_input=True),
|
||||
io.Video.Input("video", optional=True),
|
||||
io.Combo.Input("transfer_format", options=list(FORMATS), default=DEFAULT_FORMAT,
|
||||
optional=True, tooltip=FORMAT_TOOLTIP),
|
||||
io.Int.Input("transfer_quality", default=DEFAULT_QUALITY, min=1, max=100,
|
||||
optional=True, tooltip=QUALITY_TOOLTIP)],
|
||||
outputs=[], is_output_node=True, is_input_list=True)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, request_id, binding_json, images=None, audio=None, masks=None, texts=None, video=None,
|
||||
transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
|
||||
binding = _json(_scalar(binding_json))
|
||||
validate_binding(binding)
|
||||
provided = dict(images=images, audio=audio, masks=masks, texts=texts, video=video)
|
||||
if {name for name, value in provided.items() if value is not None} != {binding["socket"]}:
|
||||
raise ValueError("Media export must receive exactly its declared typed socket")
|
||||
exported = export_values(_scalar(request_id), binding, provided[binding["socket"]],
|
||||
transfer_format=_scalar(transfer_format),
|
||||
transfer_quality=_scalar(transfer_quality))
|
||||
return io.NodeOutput(ui={"media_exports": [exported]})
|
||||
@@ -0,0 +1,482 @@
|
||||
"""Lossless transport of pinned native ComfyUI VIDEO representations.
|
||||
|
||||
No save_to/get_components fallback is allowed: an unknown class must not silently
|
||||
transcode. Tensor components use the shared CPU codec. File bytes are opaque and
|
||||
retain signed trim/crop parameters; decoded files own independent BytesIOs.
|
||||
VideoFromList's native constructor flattens its inputs, so wire segments are a
|
||||
bounded, flat list of exact VideoFromFile objects, never recursive playlists.
|
||||
|
||||
Offsets are relative to the containing VIDEO payload. Canonical part order is
|
||||
images/alpha/audio, or segments/complete_audio. Every byte is covered exactly
|
||||
once. All descriptor structure and tensor sizes are preflighted before decoding
|
||||
any media storage. File containers are not parsed until a native consumer reads
|
||||
them; this adapter does not attempt to validate or repair container codecs.
|
||||
"""
|
||||
from fractions import Fraction
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
import stat
|
||||
|
||||
from .codec import MAX_PAYLOAD_BYTES, decode_payload, encode_payload
|
||||
|
||||
|
||||
MAX_SEGMENTS = 64
|
||||
MAX_METADATA_BYTES = 64 * 1024
|
||||
MAX_METADATA_DEPTH = 16
|
||||
_ENCODING = "native-video-v1"
|
||||
_BASE_FIELDS = {"encoding", "representation", "byte_length"}
|
||||
_COMPONENT_FIELDS = {"images", "alpha", "audio", "frame_rate", "metadata", "bit_depth", "color_space"}
|
||||
_FILE_FIELDS = {"data", "start_time", "duration", "crop"}
|
||||
_SEGMENT_FIELDS = {"segments", "complete_audio"}
|
||||
_TENSOR_BYTES = {"float16": 2, "bfloat16": 2, "float32": 4, "float64": 8}
|
||||
|
||||
|
||||
def _limit(max_bytes):
|
||||
if type(max_bytes) is not int or max_bytes < 0:
|
||||
raise ValueError("VIDEO max_bytes must be a nonnegative integer")
|
||||
return min(max_bytes, MAX_PAYLOAD_BYTES)
|
||||
|
||||
|
||||
def _native_types():
|
||||
# Lazy imports keep non-VIDEO users independent of ComfyUI and PyAV.
|
||||
from comfy_api.latest._input_impl.video_types import (
|
||||
VideoFromComponents, VideoFromFile, VideoFromList,
|
||||
)
|
||||
from comfy_api.latest._util import VideoComponents
|
||||
return VideoFromComponents, VideoFromFile, VideoFromList, VideoComponents
|
||||
|
||||
|
||||
def _fields(value, fields, label):
|
||||
if type(value) is not dict or set(value) != fields:
|
||||
raise ValueError(f"Invalid VIDEO {label} fields")
|
||||
|
||||
|
||||
def _integer(value, *, positive=False):
|
||||
if type(value) is not int or value < (1 if positive else 0):
|
||||
raise ValueError("VIDEO length, offset and dimensions must be integers in range")
|
||||
return value
|
||||
|
||||
|
||||
def _number(value):
|
||||
if type(value) not in (int, float) or (type(value) is float and not math.isfinite(value)):
|
||||
raise ValueError("VIDEO trim values must be finite numbers")
|
||||
return value
|
||||
|
||||
|
||||
def _crop(value, *, wire=False):
|
||||
if value is None:
|
||||
return None
|
||||
if type(value) is not (list if wire else tuple) or len(value) != 4:
|
||||
raise ValueError("VIDEO crop must contain four integers")
|
||||
if any(type(part) is not int for part in value):
|
||||
raise ValueError("VIDEO crop must contain four integers")
|
||||
# Native crop normalization is a consumer concern, not a transport operation.
|
||||
return list(value) if not wire else tuple(value)
|
||||
|
||||
|
||||
def _frame_rate(value):
|
||||
_fields(value, {"numerator", "denominator"}, "frame_rate")
|
||||
_integer(value["numerator"], positive=True)
|
||||
_integer(value["denominator"], positive=True)
|
||||
|
||||
|
||||
def _encoding_options(bit_depth, color_space):
|
||||
_integer(bit_depth, positive=True)
|
||||
if type(color_space) is not str or color_space not in ("sRGB", "HDR", "HDR PQ"):
|
||||
raise ValueError("Unsupported native VIDEO color_space")
|
||||
|
||||
|
||||
def _metadata(value):
|
||||
"""Validate bounded JSON before allocating its serialized form; return a copy."""
|
||||
if value is None:
|
||||
return None
|
||||
if type(value) is not dict:
|
||||
raise ValueError("VIDEO metadata must be a JSON object or null")
|
||||
remaining = MAX_METADATA_BYTES
|
||||
|
||||
def charge(size):
|
||||
nonlocal remaining
|
||||
remaining -= size
|
||||
if remaining < 0:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
|
||||
def visit(item, depth):
|
||||
if depth > MAX_METADATA_DEPTH:
|
||||
raise ValueError("VIDEO metadata exceeds depth limit")
|
||||
kind = type(item)
|
||||
if kind in (dict, list):
|
||||
# Each member costs at least one JSON byte. Also bounds iteration
|
||||
# before walking potentially huge containers or cyclic structures.
|
||||
if len(item) > remaining:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
charge(2 + max(0, len(item) - 1))
|
||||
if kind is dict:
|
||||
for key, child in item.items():
|
||||
if type(key) is not str:
|
||||
raise ValueError("VIDEO metadata keys must be strings")
|
||||
visit(key, depth + 1)
|
||||
charge(1)
|
||||
visit(child, depth + 1)
|
||||
else:
|
||||
for child in item:
|
||||
visit(child, depth + 1)
|
||||
elif kind in (str, int, float, bool) or item is None:
|
||||
if kind is str and len(item) > remaining:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
# Avoid converting an unbounded Python integer to decimal first.
|
||||
if kind is int and item.bit_length() > MAX_METADATA_BYTES * 4:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
if kind is float and not math.isfinite(item):
|
||||
raise ValueError("VIDEO metadata numbers must be finite")
|
||||
try:
|
||||
charge(len(json.dumps(item, ensure_ascii=True, allow_nan=False)))
|
||||
except (ValueError, OverflowError):
|
||||
raise ValueError("Invalid or oversized VIDEO metadata scalar") from None
|
||||
else:
|
||||
raise ValueError("VIDEO metadata must contain only JSON values")
|
||||
|
||||
visit(value, 0)
|
||||
return json.loads(json.dumps(value, ensure_ascii=True, allow_nan=False, separators=(",", ":")))
|
||||
|
||||
|
||||
def _attrs(value, names):
|
||||
attributes = vars(value)
|
||||
if not all(name in attributes for name in names):
|
||||
raise ValueError("Native VIDEO attributes do not match the pinned adapter")
|
||||
return [attributes[name] for name in names]
|
||||
|
||||
|
||||
def _roots(allowed_roots):
|
||||
if allowed_roots is None:
|
||||
import folder_paths
|
||||
allowed_roots = [folder_paths.get_input_directory(), folder_paths.get_output_directory(),
|
||||
folder_paths.get_temp_directory()]
|
||||
if type(allowed_roots) not in (list, tuple):
|
||||
raise ValueError("VIDEO allowed_roots must be a list or tuple")
|
||||
try:
|
||||
roots = [Path(root).resolve() for root in allowed_roots]
|
||||
except (OSError, TypeError, ValueError, RuntimeError):
|
||||
raise ValueError("Invalid VIDEO allowed roots") from None
|
||||
return roots
|
||||
|
||||
|
||||
def _read_path(value, limit, allowed_roots):
|
||||
"""Resolve authorization, reject special files, and bound all reads.
|
||||
|
||||
Open each resolved path component with NOFOLLOW using directory fds, so an
|
||||
intermediate symlink swapped in after resolve() cannot bypass authorization.
|
||||
NONBLOCK prevents an attacker swapping in a FIFO from hanging the worker.
|
||||
"""
|
||||
roots = _roots(allowed_roots)
|
||||
directory_fd = file_fd = None
|
||||
try:
|
||||
path = Path(value).resolve(strict=True)
|
||||
if not any(path.is_relative_to(root) for root in roots):
|
||||
raise ValueError("VIDEO file is outside allowed roots")
|
||||
directory_flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
|
||||
directory_fd = os.open(path.anchor, directory_flags)
|
||||
for part in path.parts[1:-1]:
|
||||
next_fd = os.open(part, directory_flags, dir_fd=directory_fd)
|
||||
os.close(directory_fd)
|
||||
directory_fd = next_fd
|
||||
file_fd = os.open(path.name, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=directory_fd)
|
||||
info = os.fstat(file_fd)
|
||||
if not stat.S_ISREG(info.st_mode):
|
||||
raise ValueError("VIDEO source must be a regular file")
|
||||
if info.st_size > limit:
|
||||
raise ValueError("VIDEO file exceeds byte limit")
|
||||
# Read at most the stat size plus a sentinel; reject growth or truncation.
|
||||
chunks = []
|
||||
count = 0
|
||||
while count <= info.st_size:
|
||||
chunk = os.read(file_fd, min(1024 * 1024, info.st_size + 1 - count))
|
||||
if not chunk:
|
||||
break
|
||||
chunks.append(chunk)
|
||||
count += len(chunk)
|
||||
if count != info.st_size:
|
||||
raise ValueError("VIDEO source changed size during bounded read")
|
||||
return b"".join(chunks)
|
||||
except (OSError, RuntimeError):
|
||||
raise ValueError("VIDEO file cannot be read safely") from None
|
||||
finally:
|
||||
if file_fd is not None:
|
||||
os.close(file_fd)
|
||||
if directory_fd is not None:
|
||||
os.close(directory_fd)
|
||||
|
||||
|
||||
def _file_bytes(value, limit, allowed_roots):
|
||||
if type(value) is io.BytesIO:
|
||||
try:
|
||||
with value.getbuffer() as view:
|
||||
if view.nbytes > limit:
|
||||
raise ValueError("VIDEO file exceeds byte limit")
|
||||
return bytes(view)
|
||||
except ValueError:
|
||||
raise ValueError("VIDEO BytesIO is closed or exceeds byte limit") from None
|
||||
if isinstance(value, (str, os.PathLike)):
|
||||
return _read_path(value, limit, allowed_roots)
|
||||
raise ValueError("VIDEO file must be a local path or exact BytesIO")
|
||||
|
||||
|
||||
def encode_video(value, *, max_bytes=MAX_PAYLOAD_BYTES, allowed_roots=None):
|
||||
"""Encode an exact native VIDEO without opening/transcoding its container.
|
||||
|
||||
Paths are authorized only under allowed_roots (when supplied), otherwise
|
||||
under ComfyUI input/output/temp roots. Budgets may lower, never raise, the
|
||||
shared maximum. Unsupported native state raises ValueError, not a fallback.
|
||||
Non-CPU tensors are rejected; runtime staging belongs to the bridge boundary.
|
||||
"""
|
||||
return _encode_video(value, max_bytes=max_bytes, allowed_roots=allowed_roots, stage_cpu=False)
|
||||
|
||||
|
||||
def encode_video_cpu(value, *, max_bytes=MAX_PAYLOAD_BYTES, allowed_roots=None):
|
||||
"""Bridge-only boundary: detach/stage native tensors, then losslessly encode.
|
||||
|
||||
Share the exact-type adapter with encode_video rather than reconstructing
|
||||
native videos, calling get_components, or duplicating private attributes.
|
||||
Images, alpha and both audio representations move without dtype conversion;
|
||||
metadata/options and opaque file bytes/trim/crop take the unchanged path.
|
||||
"""
|
||||
return _encode_video(value, max_bytes=max_bytes, allowed_roots=allowed_roots, stage_cpu=True)
|
||||
|
||||
|
||||
def _stage_tensor_cpu(value):
|
||||
import torch
|
||||
if not isinstance(value, torch.Tensor):
|
||||
raise ValueError("Tensor value must be a torch.Tensor")
|
||||
return value.detach().cpu()
|
||||
|
||||
|
||||
def _encode_video(value, *, max_bytes, allowed_roots, stage_cpu):
|
||||
limit = _limit(max_bytes)
|
||||
components_class, file_class, list_class, components_type = _native_types()
|
||||
descriptor = {"encoding": _ENCODING}
|
||||
chunks = []
|
||||
size = 0
|
||||
|
||||
def add(payload, nested=None):
|
||||
nonlocal size
|
||||
if len(payload) > limit - size:
|
||||
raise ValueError("VIDEO payload exceeds byte limit")
|
||||
slot = {"offset": size, "length": len(payload)}
|
||||
if nested is not None:
|
||||
slot["descriptor"] = nested
|
||||
chunks.append(payload)
|
||||
size += len(payload)
|
||||
return slot
|
||||
|
||||
def tensor(kind, tensor_value):
|
||||
if stage_cpu:
|
||||
if kind == "AUDIO":
|
||||
if not isinstance(tensor_value, dict) or set(tensor_value) != {"waveform", "sample_rate"}:
|
||||
raise ValueError("AUDIO value must contain waveform and sample_rate")
|
||||
tensor_value = {**tensor_value, "waveform": _stage_tensor_cpu(tensor_value["waveform"])}
|
||||
else:
|
||||
tensor_value = _stage_tensor_cpu(tensor_value)
|
||||
nested, payload = encode_payload(kind, tensor_value, max_bytes=limit - size)
|
||||
return add(payload, nested)
|
||||
|
||||
if type(value) is components_class:
|
||||
component, bit_depth, color_space = _attrs(value, [
|
||||
"_VideoFromComponents__components", "_VideoFromComponents__bit_depth",
|
||||
"_VideoFromComponents__color_space",
|
||||
])
|
||||
if type(component) is not components_type:
|
||||
raise ValueError("Unsupported native VIDEO components class")
|
||||
images, fps, audio, metadata, alpha = _attrs(component, [
|
||||
"images", "frame_rate", "audio", "metadata", "alpha",
|
||||
])
|
||||
if type(fps) is not Fraction or fps <= 0:
|
||||
raise ValueError("VIDEO frame_rate must be a positive Fraction")
|
||||
_encoding_options(bit_depth, color_space)
|
||||
descriptor.update(representation="components", bit_depth=bit_depth, color_space=color_space,
|
||||
frame_rate={"numerator": fps.numerator, "denominator": fps.denominator},
|
||||
metadata=_metadata(metadata))
|
||||
descriptor["images"] = tensor("IMAGE", images)
|
||||
# Native get_components() returns BHW1 alpha; MASK nodes use BHW.
|
||||
# Select an existing lossless tensor layout without squeezing/converting.
|
||||
alpha_kind = "IMAGE" if getattr(alpha, "ndim", None) == 4 else "MASK"
|
||||
descriptor["alpha"] = tensor(alpha_kind, alpha) if alpha is not None else None
|
||||
if alpha is not None:
|
||||
_alpha_header(descriptor["alpha"], descriptor["images"])
|
||||
descriptor["audio"] = tensor("AUDIO", audio) if audio is not None else None
|
||||
elif type(value) is file_class:
|
||||
source, start, duration, crop = _attrs(value, [
|
||||
"_VideoFromFile__file", "_VideoFromFile__start_time", "_VideoFromFile__duration",
|
||||
"_VideoFromFile__crop",
|
||||
])
|
||||
descriptor.update(representation="file", start_time=_number(start), duration=_number(duration),
|
||||
crop=_crop(crop))
|
||||
descriptor["data"] = add(_file_bytes(source, limit, allowed_roots))
|
||||
elif type(value) is list_class:
|
||||
videos, complete_audio = _attrs(value, ["videos", "complete_audio"])
|
||||
if type(videos) is not list or not 1 <= len(videos) <= MAX_SEGMENTS:
|
||||
raise ValueError("VIDEO must have between 1 and 64 segments")
|
||||
if any(type(segment) is not file_class for segment in videos):
|
||||
raise ValueError("VIDEO segments must be exact native VideoFromFile objects")
|
||||
descriptor.update(representation="segments", segments=[])
|
||||
for segment in videos:
|
||||
nested, payload = encode_video(segment, max_bytes=limit - size, allowed_roots=allowed_roots)
|
||||
descriptor["segments"].append(add(payload, nested))
|
||||
descriptor["complete_audio"] = tensor("AUDIO", complete_audio) if complete_audio is not None else None
|
||||
else:
|
||||
raise ValueError("Unsupported native VIDEO class")
|
||||
descriptor["byte_length"] = size
|
||||
return descriptor, b"".join(chunks)
|
||||
|
||||
|
||||
def _tensor_header(kind, descriptor, length):
|
||||
"""Allocation-free preflight of the shared tensor-le-v1 descriptor.
|
||||
|
||||
The shared decoder still validates the bytes/finite values. Doing its small
|
||||
structural check here first prevents allocation for an early part when a
|
||||
later part's descriptor is already known to be invalid.
|
||||
"""
|
||||
fields = {"encoding", "dtype", "shape", "layout"}
|
||||
if kind == "AUDIO":
|
||||
fields.add("sample_rate")
|
||||
_fields(descriptor, fields, "tensor")
|
||||
layout = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}[kind]
|
||||
if type(descriptor["encoding"]) is not str or descriptor["encoding"] != "tensor-le-v1":
|
||||
raise ValueError("Invalid VIDEO tensor encoding")
|
||||
if type(descriptor["layout"]) is not str or descriptor["layout"] != layout:
|
||||
raise ValueError("Invalid VIDEO tensor layout")
|
||||
dtype = descriptor["dtype"]
|
||||
if type(dtype) is not str or dtype not in _TENSOR_BYTES:
|
||||
raise ValueError("Invalid VIDEO tensor dtype")
|
||||
shape = descriptor["shape"]
|
||||
if type(shape) is not list or len(shape) != len(layout):
|
||||
raise ValueError("Invalid VIDEO tensor rank")
|
||||
size = _TENSOR_BYTES[dtype]
|
||||
for dimension in shape:
|
||||
_integer(dimension, positive=True)
|
||||
if dimension > length // size:
|
||||
raise ValueError("VIDEO tensor shape exceeds byte limit")
|
||||
size *= dimension
|
||||
if kind == "IMAGE" and shape[-1] not in (1, 3, 4):
|
||||
raise ValueError("Invalid VIDEO image channels")
|
||||
if size != length:
|
||||
raise ValueError("VIDEO tensor length does not match shape")
|
||||
if kind == "AUDIO":
|
||||
_integer(descriptor["sample_rate"], positive=True)
|
||||
|
||||
|
||||
def _alpha_header(alpha, images):
|
||||
_fields(alpha, {"offset", "length", "descriptor"}, "alpha part")
|
||||
descriptor = alpha["descriptor"]
|
||||
if type(descriptor) is not dict:
|
||||
raise ValueError("Invalid VIDEO alpha descriptor")
|
||||
layout = descriptor.get("layout")
|
||||
if type(layout) is not str or layout not in ("BHW", "BHWC"):
|
||||
raise ValueError("Invalid VIDEO alpha layout")
|
||||
kind = "MASK" if layout == "BHW" else "IMAGE"
|
||||
_tensor_header(kind, descriptor, _integer(alpha["length"]))
|
||||
expected = images["descriptor"]["shape"][:3] + ([1] if kind == "IMAGE" else [])
|
||||
if descriptor["shape"] != expected:
|
||||
raise ValueError("VIDEO alpha shape must match images")
|
||||
return kind
|
||||
|
||||
|
||||
def _validate_video(descriptor, size, *, file_only=False):
|
||||
if type(descriptor) is not dict:
|
||||
raise ValueError("VIDEO descriptor must be a dictionary")
|
||||
representation = descriptor.get("representation")
|
||||
if type(representation) is not str or representation not in ("components", "file", "segments"):
|
||||
raise ValueError("Unsupported VIDEO representation")
|
||||
if file_only and representation != "file":
|
||||
# Depth is at most one: native lists already flatten their segments.
|
||||
raise ValueError("VIDEO segments cannot contain nested lists or components")
|
||||
fields = {"components": _COMPONENT_FIELDS, "file": _FILE_FIELDS, "segments": _SEGMENT_FIELDS}[representation]
|
||||
_fields(descriptor, _BASE_FIELDS | fields, "descriptor")
|
||||
if type(descriptor["encoding"]) is not str or descriptor["encoding"] != _ENCODING:
|
||||
raise ValueError("Unsupported VIDEO encoding")
|
||||
if _integer(descriptor["byte_length"]) != size:
|
||||
raise ValueError("VIDEO payload length does not match descriptor")
|
||||
result = dict(descriptor)
|
||||
cursor = 0
|
||||
|
||||
def part(slot, kind=None, video=False):
|
||||
nonlocal cursor
|
||||
fields = {"offset", "length", "descriptor"} if kind or video else {"offset", "length"}
|
||||
_fields(slot, fields, "part")
|
||||
offset = _integer(slot["offset"])
|
||||
length = _integer(slot["length"])
|
||||
if offset != cursor or length > size - cursor:
|
||||
raise ValueError("VIDEO parts must not overlap, leave gaps or exceed bounds")
|
||||
cursor += length
|
||||
if kind:
|
||||
_tensor_header(kind, slot["descriptor"], length)
|
||||
if video:
|
||||
return {**slot, "descriptor": _validate_video(slot["descriptor"], length, file_only=True)}
|
||||
return slot
|
||||
|
||||
if representation == "components":
|
||||
_frame_rate(descriptor["frame_rate"])
|
||||
_encoding_options(descriptor["bit_depth"], descriptor["color_space"])
|
||||
result["metadata"] = _metadata(descriptor["metadata"])
|
||||
part(descriptor["images"], "IMAGE")
|
||||
if descriptor["alpha"] is not None:
|
||||
part(descriptor["alpha"], _alpha_header(descriptor["alpha"], descriptor["images"]))
|
||||
if descriptor["audio"] is not None:
|
||||
part(descriptor["audio"], "AUDIO")
|
||||
elif representation == "file":
|
||||
_number(descriptor["start_time"])
|
||||
_number(descriptor["duration"])
|
||||
_crop(descriptor["crop"], wire=True)
|
||||
part(descriptor["data"])
|
||||
else:
|
||||
segments = descriptor["segments"]
|
||||
if type(segments) is not list or not 1 <= len(segments) <= MAX_SEGMENTS:
|
||||
raise ValueError("VIDEO must have between 1 and 64 segments")
|
||||
result["segments"] = [part(segment, video=True) for segment in segments]
|
||||
if descriptor["complete_audio"] is not None:
|
||||
part(descriptor["complete_audio"], "AUDIO")
|
||||
if cursor != size:
|
||||
raise ValueError("VIDEO payload contains trailing bytes")
|
||||
return result
|
||||
|
||||
|
||||
def _decode_validated(descriptor, payload, types):
|
||||
components_class, file_class, list_class, components_type = types
|
||||
|
||||
def data(slot):
|
||||
return payload[slot["offset"]:slot["offset"] + slot["length"]]
|
||||
|
||||
def tensor(kind, slot):
|
||||
if slot is None:
|
||||
return None
|
||||
return decode_payload(kind, slot["descriptor"], data(slot), max_bytes=slot["length"])
|
||||
|
||||
if descriptor["representation"] == "file":
|
||||
# VideoFromFile retains the BytesIO, unlike a temporary filesystem path.
|
||||
return file_class(io.BytesIO(data(descriptor["data"])), start_time=descriptor["start_time"],
|
||||
duration=descriptor["duration"], crop=_crop(descriptor["crop"], wire=True))
|
||||
if descriptor["representation"] == "components":
|
||||
fps = descriptor["frame_rate"]
|
||||
return components_class(components_type(
|
||||
images=tensor("IMAGE", descriptor["images"]),
|
||||
frame_rate=Fraction(fps["numerator"], fps["denominator"]),
|
||||
audio=tensor("AUDIO", descriptor["audio"]),
|
||||
metadata=descriptor["metadata"],
|
||||
alpha=tensor("IMAGE" if descriptor["alpha"] is not None and
|
||||
descriptor["alpha"]["descriptor"]["layout"] == "BHWC" else "MASK",
|
||||
descriptor["alpha"]),
|
||||
), bit_depth=descriptor["bit_depth"], color_space=descriptor["color_space"])
|
||||
videos = [_decode_validated(slot["descriptor"], data(slot), types) for slot in descriptor["segments"]]
|
||||
return list_class(videos, complete_audio=tensor("AUDIO", descriptor["complete_audio"]))
|
||||
|
||||
|
||||
def decode_video(descriptor, payload, *, max_bytes=MAX_PAYLOAD_BYTES):
|
||||
"""Decode a bounded VIDEO into independently owned native CPU media objects."""
|
||||
limit = _limit(max_bytes)
|
||||
if type(payload) is not bytes:
|
||||
raise ValueError("VIDEO payload must be bytes")
|
||||
if len(payload) > limit:
|
||||
raise ValueError("VIDEO payload exceeds byte limit")
|
||||
validated = _validate_video(descriptor, len(payload))
|
||||
return _decode_validated(validated, payload, _native_types())
|
||||
@@ -0,0 +1,469 @@
|
||||
"""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="<i2")
|
||||
with wave.open(str(wav), "wb") as out:
|
||||
out.setnchannels(2)
|
||||
out.setsampwidth(2)
|
||||
out.setframerate(22050)
|
||||
out.writeframes(pcm.tobytes())
|
||||
waveform, sample_rate = nodes_audio.load(str(wav))
|
||||
clip = native_root / "native.mkv"
|
||||
clip.write_bytes(tiny_video(3))
|
||||
native_expected = copy.copy(fixtures["worker-a"])
|
||||
native_expected["audio"] = [{"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}]
|
||||
native_expected["video"] = [VideoFromFile(str(clip))]
|
||||
expectations = {}
|
||||
for key, values in (("worker-a", native_expected), ("worker-b", fixtures["worker-b"])):
|
||||
expectations[key] = {socket: [video_codec.encode_video(v, allowed_roots=[native_root]) if socket == "video" else v
|
||||
for v in values[socket]] for socket in sockets}
|
||||
|
||||
# Native routes only, bound to a new ephemeral loopback listener. No fake
|
||||
# history, permissive validator, bridge/codec mock, or existing service.
|
||||
allowed_paths = {"/prompt", "/history/{prompt_id}", "/view", "/queue"}
|
||||
ps.app.add_routes([route for route in ps.routes if route.path in allowed_paths])
|
||||
completed = []
|
||||
stop = asyncio.Event()
|
||||
|
||||
async def queue_worker():
|
||||
while not stop.is_set():
|
||||
item = await asyncio.to_thread(ps.prompt_queue.get, 0.1)
|
||||
if item is None:
|
||||
continue
|
||||
queued, item_id = item
|
||||
assert len(queued) in (5, 6), "Unsupported core PromptQueue tuple; adapt acceptance worker explicitly"
|
||||
executor = execution.PromptExecutor(ps, cache_args={"ram": 0, "ram_inactive": 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, queued[2], queued[1], queued[3], queued[4])
|
||||
status = execution.PromptQueue.ExecutionStatus(
|
||||
status_str="success" if executor.success else "error", completed=executor.success,
|
||||
messages=executor.status_messages)
|
||||
ps.prompt_queue.task_done(item_id, executor.history_result, status=status)
|
||||
completed.append((queued, executor.success, executor.status_messages))
|
||||
|
||||
@asynccontextmanager
|
||||
async def serve():
|
||||
runner = web.AppRunner(ps.app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
ALLOWED_CONNECTIONS.add(("127.0.0.1", port))
|
||||
worker = asyncio.create_task(queue_worker())
|
||||
try:
|
||||
yield f"http://127.0.0.1:{port}"
|
||||
finally:
|
||||
stop.set()
|
||||
await asyncio.wait_for(worker, 10)
|
||||
await runner.cleanup()
|
||||
ALLOWED_CONNECTIONS.clear()
|
||||
|
||||
grouped_results = {}
|
||||
async with serve() as url:
|
||||
for key in ("worker-a", "worker-b"):
|
||||
participant = "serverless:" + key
|
||||
request_id = "acceptance-" + key
|
||||
graph = {"source": node("AcceptanceMediaSource", key=key),
|
||||
"collector": node("DistributedMediaCollector", load_balance=False,
|
||||
multi_job_id="job-main", worker_id=participant,
|
||||
**{socket: ["source", index] for index, socket in enumerate(sockets)}),
|
||||
"text-collector": node("DistributedMediaCollector", load_balance=False,
|
||||
multi_job_id="job-text", worker_id=participant, texts=["source", 3])}
|
||||
if key == "worker-a":
|
||||
graph.update({"wav": node("LoadAudio", audio=wav.name), "clip": node("LoadVideo", file=clip.name)})
|
||||
graph["collector"]["inputs"].update(audio=["wav", 0], video=["clip", 0])
|
||||
before = copy.deepcopy(graph)
|
||||
workflow, manifest = compiler.compile_media_prompt(graph, request_id)
|
||||
assert graph == before
|
||||
assert len(manifest) == 6 and {b["kind"] for b in manifest} == set(kinds)
|
||||
folder_paths.set_input_directory(str(native_root))
|
||||
workflow_before = copy.deepcopy(workflow)
|
||||
images, artifacts, refs = await packer.prepare_media_inputs(workflow)
|
||||
assert images == [] and workflow == workflow_before
|
||||
if key == "worker-a":
|
||||
assert {a["kind"] for a in artifacts} == {"AUDIO", "VIDEO"} and len(artifacts) == 2
|
||||
assert len(refs) == 2
|
||||
assert {ref["input_name"] for ref in refs} == {"audio", "video"}
|
||||
assert all(workflow[ref["node_id"]]["class_type"] == "MediaExportV1" for ref in refs)
|
||||
# Endpoint cannot fall back to the master's local files/loaders.
|
||||
wav.unlink()
|
||||
clip.unlink()
|
||||
print("NATIVE_INPUTS_AND_DIRECT_EXPORT_REFS_OK WAV + native file VIDEO")
|
||||
else:
|
||||
assert not artifacts and not refs
|
||||
folder_paths.set_input_directory(str(base / "input"))
|
||||
media = {"version": 1, "profile": contract.PROFILE, "request_id": request_id,
|
||||
"artifacts": artifacts, "inputs": refs, "outputs": manifest}
|
||||
body = json.loads(serializer(workflow, images, 30, media=media))
|
||||
assert decoder.validate_prepared_job({"worker": {"id": participant, "media_profile": contract.PROFILE,
|
||||
"timeout_seconds": 30}, "workflow": workflow, "images": images,
|
||||
"media": media, "manifest": manifest}) is True
|
||||
before = copy.deepcopy(body["input"])
|
||||
result = await handler.execute_media_request(body["input"], comfy_url=url,
|
||||
input_dir=base / "input", temp_dir=base / "temp", timeout_seconds=30)
|
||||
assert body["input"] == before
|
||||
assert completed and completed[-1][1], completed
|
||||
queued = completed[-1][0]
|
||||
assert set(queued[4]) == {b["node_id"] for b in manifest}
|
||||
submitted = queued[2]
|
||||
if key == "worker-a":
|
||||
assert "wav" not in submitted and "clip" not in submitted
|
||||
imports = [n for n in submitted.values() if n["class_type"] == "MediaImportV1"]
|
||||
assert len(imports) == 2
|
||||
assert {json.loads(n["inputs"]["artifact_json"])["kind"] for n in imports} == {"AUDIO", "VIDEO"}
|
||||
assert not list((base / "input").rglob("*.bin"))
|
||||
assert not list((base / "temp").rglob("*.bin"))
|
||||
if key == "worker-b":
|
||||
# Retire the originals as well as handler-owned staging/export
|
||||
# files before constructing the result's lazy native videos.
|
||||
for original in fixtures[key]["video"][1].videos:
|
||||
vars(original)["_VideoFromFile__file"].close()
|
||||
# Exercise the wire format, not an in-memory handler-only object.
|
||||
result = json.loads(json.dumps(result, allow_nan=False))
|
||||
grouped = await decoder.collect_typed_results(result, media, manifest)
|
||||
assert set(grouped) == {"job-main", "job-text"}
|
||||
assert set(grouped["job-main"]) == set(sockets)
|
||||
assert grouped["job-text"]["texts"] == expectations[key]["texts"]
|
||||
for socket, kind in zip(sockets, kinds):
|
||||
values, expected = grouped["job-main"][socket], expectations[key][socket]
|
||||
assert len(values) == len(expected), (key, socket)
|
||||
for actual, wanted in zip(values, expected):
|
||||
exact(kind, actual, wanted)
|
||||
if key == "worker-b":
|
||||
original = fixtures[key]["video"][0].get_components()
|
||||
decoded = grouped["job-main"]["video"][0].get_components()
|
||||
for field in ("images", "alpha"):
|
||||
assert getattr(decoded, field).data_ptr() != getattr(original, field).data_ptr()
|
||||
assert decoded.audio["waveform"].data_ptr() != original.audio["waveform"].data_ptr()
|
||||
assert decoded.metadata == original.metadata and decoded.metadata is not original.metadata
|
||||
grouped_results[participant] = grouped
|
||||
# Late invalid artifact must reject the entire grouped response.
|
||||
corrupt = copy.deepcopy(result)
|
||||
corrupt["media"]["artifacts"][-1]["sha256"] = "0" * 64
|
||||
try:
|
||||
await decoder.collect_typed_results(corrupt, media, manifest)
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError("Result decoder accepted a corrupt final artifact")
|
||||
assert ps.distributed_pending_jobs == {}
|
||||
# Exercise user-selected image return encoding through the actual
|
||||
# compiler, HTTP handler, V3 exporter and atomic response decoder.
|
||||
for fmt in ("png", "jpeg", "h264"):
|
||||
key = "encoded-" + fmt
|
||||
fixtures[key] = native_fixtures(3)
|
||||
# Keep full alpha precision and use nontrivial image dimensions.
|
||||
torch.manual_seed(41)
|
||||
rgba = torch.rand((2, 64, 64, 4), dtype=torch.float64)
|
||||
fixtures[key]["images"] = [rgba]
|
||||
request_id = "request-" + key
|
||||
graph = {"source": node("AcceptanceMediaSource", key=key),
|
||||
"collector": node("DistributedMediaCollector", transfer_format=fmt,
|
||||
transfer_quality=73, multi_job_id="job-" + key,
|
||||
worker_id="serverless:" + key,
|
||||
**{socket: ["source", i] for i, socket in enumerate(sockets)})}
|
||||
original = copy.deepcopy(graph)
|
||||
workflow, manifest = compiler.compile_media_prompt(graph, request_id)
|
||||
media = {"version": 1, "profile": contract.PROFILE, "request_id": request_id,
|
||||
"artifacts": [], "inputs": [], "outputs": manifest}
|
||||
body = json.loads(serializer(workflow, [], 30, media=media))
|
||||
result = await handler.execute_media_request(body["input"], comfy_url=url,
|
||||
input_dir=base / "input", temp_dir=base / "temp", timeout_seconds=30)
|
||||
assert graph == original
|
||||
grouped = await decoder.collect_typed_results(result, media, manifest)
|
||||
actual = grouped["job-" + key]
|
||||
image = actual["images"][0]
|
||||
assert image.shape == rgba.shape and image.dtype == rgba.dtype
|
||||
assert torch.equal(image[..., 3], rgba[..., 3]), "alpha must not be discarded/narrowed"
|
||||
assert not torch.equal(image[..., :3], rgba[..., :3]), "selection must actually encode"
|
||||
if fmt == "png":
|
||||
pixels = (rgba[..., :3].float().clamp(0, 1) * 255).to(torch.uint8)
|
||||
assert torch.equal(image[..., :3], (pixels.float() / 255).to(rgba.dtype))
|
||||
for socket, kind in zip(sockets[1:], kinds[1:]):
|
||||
for got, wanted in zip(actual[socket], fixtures[key][socket]):
|
||||
exact(kind, got, video_codec.encode_video(wanted) if kind == "VIDEO" else wanted)
|
||||
assert not list((base / "temp").rglob("*.bin"))
|
||||
print("IMAGE_TRANSFER_HTTP_OK png jpeg h264; alpha and other media exact")
|
||||
print("REAL_HTTP_BRIDGE_RESULTS_OK all five kinds; 6 bindings/request; native VIDEO remains readable after cleanup")
|
||||
print("ATOMIC_DECODE_REJECTION_OK corrupt final artifact rejected before delivery")
|
||||
|
||||
participants = list(grouped_results)
|
||||
for delegate in (False, True):
|
||||
tag = "delegate" if delegate else "master"
|
||||
job = "collect-" + tag
|
||||
queue = asyncio.Queue()
|
||||
for participant in reversed(participants):
|
||||
await queue.put({"worker_id": participant, "media": grouped_results[participant]["job-main"], "is_last": True})
|
||||
ps.distributed_pending_jobs[job] = queue
|
||||
ordered = ([] if delegate else [fixtures["local"]]) + [grouped_results[p]["job-main"] for p in participants]
|
||||
expected_sinks[tag] = {socket: [item for values in ordered for item in values[socket]] for socket in sockets}
|
||||
graph = {"collector": node("DistributedMediaCollector", load_balance=False, multi_job_id=job,
|
||||
enabled_worker_ids=json.dumps(participants), expected_media=json.dumps(sockets), delegate_only=delegate),
|
||||
"assert": node("AcceptanceMediaAssert", tag=tag,
|
||||
**{socket: ["collector", index] for index, socket in enumerate(sockets)})}
|
||||
if not delegate:
|
||||
graph["source"] = node("AcceptanceMediaSource", key="local")
|
||||
graph["collector"]["inputs"].update({socket: ["source", index] for index, socket in enumerate(sockets)})
|
||||
# The second group from the very same complete response must reach its
|
||||
# own Collector, rather than being folded into the first collector job.
|
||||
text_job, text_tag = job + "-text", tag + "-text"
|
||||
text_queue = asyncio.Queue()
|
||||
for participant in reversed(participants):
|
||||
await text_queue.put({"worker_id": participant,
|
||||
"media": grouped_results[participant]["job-text"], "is_last": True})
|
||||
ps.distributed_pending_jobs[text_job] = text_queue
|
||||
text_ordered = ([] if delegate else [fixtures["local"]]) + [grouped_results[p]["job-text"] for p in participants]
|
||||
expected_sinks[text_tag] = {"texts": [item for values in text_ordered for item in values["texts"]]}
|
||||
graph["text-collector"] = node("DistributedMediaCollector", load_balance=False, multi_job_id=text_job,
|
||||
enabled_worker_ids=json.dumps(participants), expected_media='["texts"]', delegate_only=delegate)
|
||||
graph["text-assert"] = node("AcceptanceMediaAssert", tag=text_tag, texts=["text-collector", 3])
|
||||
if not delegate:
|
||||
graph["text-collector"]["inputs"]["texts"] = ["source", 3]
|
||||
valid = await execution.validate_prompt(job, graph, None)
|
||||
assert valid[0] and not valid[3], valid
|
||||
executor = execution.PromptExecutor(ps, cache_args={"ram": 0, "ram_inactive": 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, graph, job, {}, valid[2])
|
||||
assert executor.success, executor.status_messages
|
||||
assert executor.history_result["outputs"]["assert"]["accepted"] == [seen_sinks[tag]]
|
||||
assert executor.history_result["outputs"]["text-assert"]["accepted"] == [seen_sinks[text_tag]]
|
||||
assert job not in ps.distributed_pending_jobs and text_job not in ps.distributed_pending_jobs
|
||||
assert source_calls == ["worker-a", "worker-b", "encoded-png", "encoded-jpeg", "encoded-h264", "local"], source_calls
|
||||
assert set(seen_sinks) == {"master", "delegate", "master-text", "delegate-text"}
|
||||
assert not list((base / "output").rglob("*"))
|
||||
assert not torch.cuda.is_initialized()
|
||||
print("MASTER_AND_DELEGATE_ORDER_OK", json.dumps(seen_sinks, sort_keys=True))
|
||||
print("MEDIA_END_TO_END_CPU_OK real PromptExecutor/HTTP; no queue-orchestration or provider acceptance claimed")
|
||||
|
||||
|
||||
ALLOWED_CONNECTIONS = set()
|
||||
|
||||
|
||||
def network_guard(event, args):
|
||||
"""Fail closed on provider/normal-instance traffic, including during imports."""
|
||||
if event == "socket.connect":
|
||||
sock, address = args
|
||||
if sock.family in (socket.AF_INET, socket.AF_INET6):
|
||||
assert address in ALLOWED_CONNECTIONS, "Acceptance forbids non-owned network connections"
|
||||
elif event == "socket.bind":
|
||||
sock, address = args
|
||||
if sock.family in (socket.AF_INET, socket.AF_INET6):
|
||||
assert address == ("127.0.0.1", 0), "Acceptance binds only ephemeral loopback"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.addaudithook(network_guard)
|
||||
asyncio.run(asyncio.wait_for(check(bootstrap()), timeout=120))
|
||||
@@ -45,7 +45,7 @@ ADDED_OPTIONAL_INPUTS = {
|
||||
|
||||
# Nodes added after the V1 baseline; not part of saved-workflow contracts.
|
||||
# DistributedWorkerSink is internal (dev-only) and only inserted by the master.
|
||||
ADDED_NODES = {'DistributedWorkerSink', 'DistributedImageListStart', 'DistributedImageListEnd'}
|
||||
ADDED_NODES = {'DistributedWorkerSink', 'DistributedImageListStart', 'DistributedImageListEnd', 'DistributedMediaCollector'}
|
||||
|
||||
|
||||
def normalized_input(value):
|
||||
@@ -366,10 +366,14 @@ async def main():
|
||||
extension = await module.comfy_entrypoint()
|
||||
classes = await extension.get_node_list()
|
||||
mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
|
||||
assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES) == 11
|
||||
assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES)
|
||||
assert set(mapping) == set(BASELINE) | ADDED_NODES
|
||||
for node_id in ADDED_NODES:
|
||||
assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id]
|
||||
media_info = mapping['DistributedMediaCollector'].GET_NODE_INFO_V1()
|
||||
assert media_info['output'] == ['IMAGE', 'AUDIO', 'MASK', 'STRING', 'VIDEO'], media_info
|
||||
assert all(media_info['output_is_list']), media_info
|
||||
assert media_info['input']['optional']['texts'][1]['forceInput'] is True
|
||||
sink_info = mapping['DistributedWorkerSink'].GET_NODE_INFO_V1()
|
||||
assert sink_info['output_node'] and sink_info['dev_only'], sink_info
|
||||
for node_id, cls in mapping.items():
|
||||
|
||||
@@ -4,6 +4,7 @@ import types
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
@@ -19,6 +20,13 @@ def _register_real_transfer_codec(package_name, utils_pkg):
|
||||
|
||||
|
||||
def _load_collector_module():
|
||||
# Restore shared dependencies even when a different test imports this loader.
|
||||
# Keep real modules imported during loading (notably extension modules) cached.
|
||||
with pytest.MonkeyPatch.context() as monkeypatch:
|
||||
return _load_collector_module_with_stubs(monkeypatch)
|
||||
|
||||
|
||||
def _load_collector_module_with_stubs(monkeypatch):
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
|
||||
package_name = "dist_collector_list_testpkg"
|
||||
|
||||
@@ -53,7 +61,7 @@ def _load_collector_module():
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
||||
sys.modules["server"] = server_module
|
||||
monkeypatch.setitem(sys.modules, "server", server_module)
|
||||
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
model_management = types.ModuleType("comfy.model_management")
|
||||
@@ -77,13 +85,13 @@ def _load_collector_module():
|
||||
|
||||
comfy_utils.ProgressBar = ProgressBar
|
||||
comfy_module.utils = comfy_utils
|
||||
sys.modules["comfy"] = comfy_module
|
||||
sys.modules["comfy.model_management"] = model_management
|
||||
sys.modules["comfy.utils"] = comfy_utils
|
||||
monkeypatch.setitem(sys.modules, "comfy", comfy_module)
|
||||
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
|
||||
monkeypatch.setitem(sys.modules, "comfy.utils", comfy_utils)
|
||||
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
monkeypatch.setitem(sys.modules, "aiohttp", aiohttp_module)
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
"""Real ComfyUI V3 bridge mapping, file ownership and exact CPU round trips."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def test_real_media_bridge_nodes():
|
||||
core = os.environ.get("COMFYUI_SOURCE_ROOT")
|
||||
if core is None:
|
||||
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU media bridge acceptance")
|
||||
assert core and (Path(core) / "comfy_api/latest/_input_impl/video_types.py").is_file(), (
|
||||
"Configured COMFYUI_SOURCE_ROOT must provide native video classes"
|
||||
)
|
||||
probe = r'''
|
||||
import asyncio, importlib.util, io, json, os, sys, tempfile
|
||||
from fractions import Fraction
|
||||
import av
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
import torch
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('Bridge tests must not query or initialize a GPU')
|
||||
torch.cuda.is_available = lambda: False
|
||||
torch.cuda.device_count = lambda: 0
|
||||
for name in ('init', '_lazy_init', 'current_device'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
import execution
|
||||
import folder_paths
|
||||
from comfy_api.latest._input_impl.video_types import VideoFromComponents, VideoFromFile, VideoFromList
|
||||
from comfy_api.latest._util import VideoComponents
|
||||
for name in ('is_available', 'device_count'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
spec = importlib.util.spec_from_file_location('typed_bridge_probe', Path(sys.argv[2])/'serverless_bridge/__init__.py', submodule_search_locations=[str(Path(sys.argv[2])/'serverless_bridge')])
|
||||
bridge = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = bridge
|
||||
spec.loader.exec_module(bridge)
|
||||
from typed_bridge_probe.nodes import MediaImportV1, MediaExportV1, export_values, import_value
|
||||
from typed_bridge_probe.contract import SOCKET_KINDS
|
||||
from typed_bridge_probe.video_codec import encode_video
|
||||
|
||||
def tiny_video():
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, 'w', format='matroska') as container:
|
||||
stream = container.add_stream('ffv1', rate=Fraction(24, 1))
|
||||
stream.width = stream.height = 32
|
||||
stream.pix_fmt = 'yuv420p'
|
||||
for level in (16, 64, 128):
|
||||
frame = av.VideoFrame.from_ndarray(np.full((32, 32, 3), level, dtype=np.uint8), format='rgb24')
|
||||
for packet in stream.encode(frame): container.mux(packet)
|
||||
for packet in stream.encode(): container.mux(packet)
|
||||
return buffer.getvalue()
|
||||
|
||||
def exact_tensor(actual, expected):
|
||||
assert actual.device.type == 'cpu' and actual.dtype == expected.dtype
|
||||
assert actual.shape == expected.shape
|
||||
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
|
||||
|
||||
def exact_video(actual, expected, wire):
|
||||
assert type(actual) is type(expected)
|
||||
assert encode_video(actual) == wire
|
||||
if type(actual) is VideoFromComponents:
|
||||
a = vars(actual)['_VideoFromComponents__components']
|
||||
e = vars(expected)['_VideoFromComponents__components']
|
||||
assert type(a) is VideoComponents and type(a.frame_rate) is Fraction
|
||||
assert a.frame_rate == e.frame_rate and a.metadata == e.metadata
|
||||
assert a.metadata is not e.metadata
|
||||
for name in ('images', 'alpha'): exact_tensor(getattr(a, name), getattr(e, name))
|
||||
exact_tensor(a.audio['waveform'], e.audio['waveform'])
|
||||
assert a.audio['sample_rate'] == e.audio['sample_rate']
|
||||
assert actual.get_bit_depth() == 10 and actual.get_color_space() == 'HDR PQ'
|
||||
else:
|
||||
values = actual.videos if type(actual) is VideoFromList else [actual]
|
||||
for value in values:
|
||||
buffer = vars(value)['_VideoFromFile__file']
|
||||
assert type(buffer) is io.BytesIO and not buffer.closed
|
||||
buffer.seek(0)
|
||||
with av.open(buffer) as container:
|
||||
assert len(list(container.decode(video=0))) == 3
|
||||
if type(actual) is VideoFromList:
|
||||
exact_tensor(actual.complete_audio['waveform'], expected.complete_audio['waveform'])
|
||||
|
||||
async def check():
|
||||
extension = await bridge.comfy_entrypoint()
|
||||
assert await extension.get_node_list() == [MediaImportV1, MediaExportV1]
|
||||
imp = MediaImportV1.GET_NODE_INFO_V1()
|
||||
exp = MediaExportV1.GET_NODE_INFO_V1()
|
||||
assert imp['output'] == ['IMAGE','AUDIO','MASK','STRING','VIDEO'], imp
|
||||
assert exp['output_node'] and MediaExportV1.define_schema().is_input_list, exp
|
||||
with tempfile.TemporaryDirectory(prefix='typed-media-', dir=os.environ.get('TMPDIR')) as scratch:
|
||||
base = Path(scratch)
|
||||
roots = {key: base/key for key in ('input','output','temp')}
|
||||
for path in roots.values(): path.mkdir()
|
||||
def annotated(name):
|
||||
for key, root in roots.items():
|
||||
suffix = f' [{key}]'
|
||||
if name.endswith(suffix): return str(root/name[:-len(suffix)])
|
||||
return str(roots['input']/name)
|
||||
with patch.object(folder_paths,'get_input_directory',lambda: str(roots['input'])), patch.object(folder_paths,'get_output_directory',lambda: str(roots['output'])), patch.object(folder_paths,'get_temp_directory',lambda: str(roots['temp'])), patch.object(folder_paths,'get_annotated_filepath',annotated):
|
||||
examples = {'images': torch.rand(2,2,3,4,dtype=torch.float64),
|
||||
'masks': torch.tensor([[[0.123456789,0.987654321]]],dtype=torch.float64),
|
||||
'audio': {'waveform':torch.randn(1,6,31,dtype=torch.float64),'sample_rate':44100},
|
||||
'texts': ' café\n🎥\r\n'}
|
||||
media = tiny_video()
|
||||
sources = [roots['input']/f'source-{n}.mkv' for n in range(3)]
|
||||
for path in sources: path.write_bytes(media)
|
||||
images = torch.arange(36, dtype=torch.float64, device='cpu').reshape(2,2,3,3) / 997
|
||||
images[0,0,0,0] = -0.0
|
||||
component_video = VideoFromComponents(VideoComponents(
|
||||
images=images, frame_rate=Fraction(30000,1001),
|
||||
alpha=torch.full((2,2,3,1),0.123456789,dtype=torch.float64,device='cpu'),
|
||||
audio=examples['audio'], metadata={'title':'雪', 'nested':[True,None,-4,0.125]}),
|
||||
bit_depth=10,color_space='HDR PQ')
|
||||
videos = [component_video,
|
||||
VideoFromFile(str(sources[0]),start_time=-0.0,duration=0,crop=(0,0,2,2)),
|
||||
VideoFromList([VideoFromFile(str(sources[1]),start_time=-1.25,duration=-2.5,crop=(-3,5,7,-9)),
|
||||
VideoFromFile(str(sources[2]),start_time=0,duration=0)],
|
||||
complete_audio=examples['audio'])]
|
||||
for socket,value in [*examples.items(), *[('video',v) for v in videos]]:
|
||||
kind = SOCKET_KINDS[socket]
|
||||
expected_wire = encode_video(value) if kind == 'VIDEO' else None
|
||||
binding = dict(binding_id='b-'+socket,collector_id='collector',job_id='job',participant_id='worker',node_id='export-'+socket,socket=socket,kind=kind)
|
||||
controls = {'request_id':'request','binding_json':json.dumps(binding)}
|
||||
data, missing, v3 = execution.get_input_data(controls,MediaExportV1,'export-'+socket)
|
||||
assert not missing, missing
|
||||
data[socket] = [value]
|
||||
results = await execution._async_map_node_over_list('probe','export-'+socket,MediaExportV1,data,MediaExportV1.FUNCTION,v3_data=v3)
|
||||
results = await execution.resolve_map_node_over_list_results(results)
|
||||
raw = results[0].ui
|
||||
exported = raw['media_exports'][0]
|
||||
assert exported['binding']['count'] == 1
|
||||
artifact = exported['artifacts'][0]
|
||||
transport = artifact.pop('transport')
|
||||
filename = transport['subfolder']+'/'+transport['filename']+' [temp]'
|
||||
decoded = import_value(artifact,filename)
|
||||
if kind in ('IMAGE','MASK'):
|
||||
assert decoded.dtype == value.dtype and torch.equal(decoded,value)
|
||||
elif kind == 'AUDIO':
|
||||
assert decoded['sample_rate'] == value['sample_rate']
|
||||
assert torch.equal(decoded['waveform'],value['waveform'])
|
||||
elif kind == 'VIDEO': exact_video(decoded,value,expected_wire)
|
||||
else: assert decoded == value
|
||||
mapped, missing, hidden = execution.get_input_data({'artifact_json':json.dumps(artifact),'filename':filename},MediaImportV1,'import')
|
||||
got = await execution._async_map_node_over_list('probe','import',MediaImportV1,mapped,MediaImportV1.FUNCTION,v3_data=hidden)
|
||||
got = await execution.resolve_map_node_over_list_results(got)
|
||||
assert len(got[0].result) == 5
|
||||
assert all(x is None for n,x in enumerate(got[0].result) if n != list(SOCKET_KINDS).index(socket))
|
||||
bad = dict(artifact,sha256='0'*64)
|
||||
try: import_value(bad,filename)
|
||||
except ValueError: pass
|
||||
else: raise AssertionError('checksum accepted')
|
||||
try: import_value(artifact,str(base/'outside.bin'))
|
||||
except ValueError: pass
|
||||
else: raise AssertionError('outside path accepted')
|
||||
if kind == 'VIDEO':
|
||||
mapped_video = got[0].result[list(SOCKET_KINDS).index(socket)]
|
||||
# Remove both the exported artifact and the original input
|
||||
# containers before checking independently owned BytesIOs.
|
||||
Path(annotated(filename)).unlink()
|
||||
originals = value.videos if type(value) is VideoFromList else [value]
|
||||
for original in originals:
|
||||
if type(original) is VideoFromFile:
|
||||
Path(vars(original)['_VideoFromFile__file']).unlink()
|
||||
exact_video(decoded,value,expected_wire)
|
||||
exact_video(mapped_video,value,expected_wire)
|
||||
if type(value) is VideoFromFile:
|
||||
assert tuple(decoded.get_components().images.shape) == (3,2,2,3)
|
||||
text_binding = dict(binding_id='empty',collector_id='collector',job_id='job',participant_id='worker',node_id='export-empty',socket='texts',kind='STRING')
|
||||
out = export_values('request',text_binding,[''])
|
||||
assert out['artifacts'][0]['byte_length'] == 0
|
||||
before = set(roots['temp'].rglob('*.bin'))
|
||||
try: export_values('request',text_binding,['first',None])
|
||||
except ValueError: pass
|
||||
else: raise AssertionError('invalid text accepted')
|
||||
assert set(roots['temp'].rglob('*.bin')) == before
|
||||
print('MEDIA_BRIDGE_CPU_OK')
|
||||
asyncio.run(check())
|
||||
'''
|
||||
env = dict(os.environ)
|
||||
env.pop("PYTHONPATH", None)
|
||||
env.pop("COMFYUI_SOURCE_ROOT", None)
|
||||
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="", PYTHONDONTWRITEBYTECODE="1",
|
||||
TMPDIR=str(scratch))
|
||||
result = subprocess.run([sys.executable, "-c", probe, str(Path(core).resolve()), str(ROOT)],
|
||||
capture_output=True, text=True, timeout=90, env=env)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert "MEDIA_BRIDGE_CPU_OK" in result.stdout
|
||||
@@ -0,0 +1,328 @@
|
||||
import importlib.util
|
||||
import math
|
||||
import struct
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
MODULE_PATH = Path(__file__).resolve().parents[1] / "serverless_bridge" / "codec.py"
|
||||
spec = importlib.util.spec_from_file_location("dist_media_codec_test", MODULE_PATH)
|
||||
codec = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(codec)
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
|
||||
SHAPES = [
|
||||
("IMAGE", (2, 3, 5, 1), "BHWC"),
|
||||
("IMAGE", (2, 3, 5, 3), "BHWC"),
|
||||
("IMAGE", (2, 3, 5, 4), "BHWC"),
|
||||
("MASK", (2, 3, 5), "BHW"),
|
||||
("AUDIO", (2, 1, 17), "BCT"),
|
||||
("AUDIO", (2, 6, 17), "BCT"),
|
||||
]
|
||||
|
||||
|
||||
def audio(tensor, sample_rate=48000):
|
||||
return {"waveform": tensor, "sample_rate": sample_rate}
|
||||
|
||||
|
||||
def descriptor(_kind="MASK", **changes):
|
||||
result = {
|
||||
"encoding": "tensor-le-v1", "dtype": "float32",
|
||||
"shape": [1, 1, 1], "layout": "BHW",
|
||||
}
|
||||
if _kind == "IMAGE":
|
||||
result.update(shape=[1, 1, 1, 3], layout="BHWC")
|
||||
elif _kind == "AUDIO":
|
||||
result.update(layout="BCT", sample_rate=48000)
|
||||
result.update(changes)
|
||||
return result
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def forbid_gpu_queries(monkeypatch):
|
||||
def forbidden(*args, **kwargs):
|
||||
pytest.fail("codec must not query or initialize a GPU")
|
||||
for name in ("is_available", "init", "_lazy_init", "device_count"):
|
||||
monkeypatch.setattr(torch.cuda, name, forbidden)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("kind,shape,layout", SHAPES)
|
||||
def test_exact_roundtrip_preserves_dtype_shape_channels_and_bits(dtype, kind, shape, layout):
|
||||
count = math.prod(shape)
|
||||
tensor = (torch.arange(count, device="cpu", dtype=torch.float64) / 997 - 0.125).to(dtype).reshape(shape)
|
||||
tensor.flatten()[0] = -0.0
|
||||
tensor.requires_grad_(True)
|
||||
value = audio(tensor) if kind == "AUDIO" else tensor
|
||||
|
||||
metadata, payload = codec.encode_payload(kind, value)
|
||||
expected = {
|
||||
"encoding": "tensor-le-v1", "dtype": str(dtype).removeprefix("torch."),
|
||||
"shape": list(shape), "layout": layout,
|
||||
}
|
||||
if kind == "AUDIO":
|
||||
expected["sample_rate"] = 48000
|
||||
assert metadata == expected
|
||||
assert type(payload) is bytes
|
||||
assert len(payload) == count * tensor.element_size()
|
||||
result = codec.decode_payload(kind, metadata, payload)
|
||||
if kind == "AUDIO":
|
||||
assert set(result) == {"waveform", "sample_rate"}
|
||||
assert result["sample_rate"] == 48000
|
||||
result = result["waveform"]
|
||||
assert result.device.type == "cpu"
|
||||
assert result.is_contiguous()
|
||||
assert not result.requires_grad
|
||||
assert result.dtype == dtype
|
||||
assert result.shape == shape
|
||||
assert torch.equal(result.view(torch.uint8), tensor.detach().view(torch.uint8))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype,fmt,expected_bits", [
|
||||
(torch.float16, "e", None), (torch.bfloat16, "H", [0x3F80, 0x8000, 0xC020]),
|
||||
(torch.float32, "f", None), (torch.float64, "d", None),
|
||||
])
|
||||
def test_wire_bytes_are_little_endian(dtype, fmt, expected_bits):
|
||||
values = [1.0, -0.0, -2.5]
|
||||
tensor = torch.tensor(values, dtype=dtype, device="cpu").reshape(1, 1, 3)
|
||||
metadata, payload = codec.encode_payload("MASK", tensor)
|
||||
expected = struct.pack("<3" + fmt, *(expected_bits if expected_bits else values))
|
||||
assert payload == expected
|
||||
assert torch.equal(codec.decode_payload("MASK", metadata, expected).view(torch.uint8), tensor.view(torch.uint8))
|
||||
|
||||
|
||||
def test_bfloat16_extremes_and_subnormal_bits_are_lossless():
|
||||
bits = torch.tensor([0, 0x8000, 1, 0x8001, 0x3F81, 0x7F7F, 0xFF7F], dtype=torch.uint16, device="cpu")
|
||||
value = bits.view(torch.bfloat16).reshape(1, 1, -1)
|
||||
metadata, payload = codec.encode_payload("MASK", value)
|
||||
assert payload == struct.pack("<7H", *bits.tolist())
|
||||
assert torch.equal(codec.decode_payload("MASK", metadata, payload).view(torch.uint16).flatten(), bits)
|
||||
|
||||
|
||||
def test_fractional_mask_is_not_quantized_clipped_or_normalized():
|
||||
value = torch.tensor([0.1234567890123, 1.0000000000001, -2.5], dtype=torch.float64, device="cpu").reshape(1, 1, 3)
|
||||
metadata, payload = codec.encode_payload("MASK", value)
|
||||
assert torch.equal(codec.decode_payload("MASK", metadata, payload), value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text", ["", " ", "\t\n ", "café 雪 🎵", "e\u0301", "\x00\ufeff\r\n"])
|
||||
def test_strings_preserve_exact_utf8(text):
|
||||
metadata, payload = codec.encode_payload("STRING", text)
|
||||
assert metadata == {"encoding": "utf8-v1"}
|
||||
assert payload == text.encode("utf-8")
|
||||
assert codec.decode_payload("STRING", metadata, payload) == text
|
||||
|
||||
|
||||
def test_noncontiguous_and_lazy_negative_tensors_roundtrip():
|
||||
value = torch.arange(24, dtype=torch.float32, device="cpu").reshape(2, 3, 4).transpose(1, 2)
|
||||
for tensor in (value, torch._neg_view(value)):
|
||||
metadata, payload = codec.encode_payload("MASK", tensor)
|
||||
result = codec.decode_payload("MASK", metadata, payload)
|
||||
assert result.is_contiguous()
|
||||
assert torch.equal(result, tensor)
|
||||
|
||||
|
||||
def test_decode_is_independent_and_ignores_default_device(monkeypatch):
|
||||
source = torch.tensor([0.25], device="cpu").reshape(1, 1, 1)
|
||||
metadata, payload = codec.encode_payload("MASK", source)
|
||||
buffers = []
|
||||
original = torch.frombuffer
|
||||
def capture(buffer, *args, **kwargs):
|
||||
buffers.append(buffer)
|
||||
return original(buffer, *args, **kwargs)
|
||||
monkeypatch.setattr(torch, "frombuffer", capture)
|
||||
with torch.device("meta"):
|
||||
first = codec.decode_payload("MASK", metadata, payload)
|
||||
second = codec.decode_payload("MASK", metadata, payload)
|
||||
assert first.device.type == second.device.type == "cpu"
|
||||
first.fill_(9)
|
||||
source.fill_(7)
|
||||
for buffer in buffers:
|
||||
if isinstance(buffer, bytearray):
|
||||
buffer[:] = b"\x00" * len(buffer)
|
||||
assert second.item() == 0.25
|
||||
assert codec.decode_payload("MASK", metadata, payload).item() == 0.25
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", ["image", " MASK", "LATENT", "", None, 1, [], {}])
|
||||
def test_unknown_kinds_raise_value_error(kind):
|
||||
with pytest.raises(ValueError, match="kind"):
|
||||
codec.encode_payload(kind, "")
|
||||
with pytest.raises(ValueError, match="kind"):
|
||||
codec.decode_payload(kind, {}, b"")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind,value", [
|
||||
("STRING", None), ("STRING", b"text"), ("STRING", 123),
|
||||
("IMAGE", [[[[0.5]]]]), ("MASK", None),
|
||||
("AUDIO", torch.ones(1, 1, 1, device="cpu")), ("AUDIO", {}),
|
||||
("AUDIO", {"waveform": "invalid", "sample_rate": 48000}),
|
||||
("AUDIO", {"waveform": torch.ones(1, 1, 1, device="cpu")}),
|
||||
])
|
||||
def test_unsupported_values_raise_value_error(kind, value):
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_payload(kind, value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.uint8, torch.int64, torch.bool, torch.complex64])
|
||||
def test_unsupported_tensor_dtype(dtype):
|
||||
with pytest.raises(ValueError, match="dtype"):
|
||||
codec.encode_payload("MASK", torch.ones(1, 1, 1, dtype=dtype, device="cpu"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind,shape", [
|
||||
("IMAGE", (1, 2, 3)), ("IMAGE", (1, 1, 1, 2)), ("IMAGE", (1, 1, 1, 5)),
|
||||
("MASK", (1, 2)), ("MASK", (1, 2, 3, 1)), ("AUDIO", (2, 4)),
|
||||
("IMAGE", (1, 0, 2, 3)), ("MASK", (0, 1, 1)), ("AUDIO", (1, 0, 4)),
|
||||
("AUDIO", (1, 2, 0)),
|
||||
])
|
||||
def test_invalid_tensor_shape(kind, shape):
|
||||
tensor = torch.empty(shape, dtype=torch.float32, device="cpu")
|
||||
with pytest.raises(ValueError, match="shape|rank|channel|dimension"):
|
||||
codec.encode_payload(kind, audio(tensor) if kind == "AUDIO" else tensor)
|
||||
|
||||
|
||||
def test_rejects_non_cpu_and_sparse_without_transferring(monkeypatch):
|
||||
meta = torch.empty((1, 1, 1), device="meta")
|
||||
sparse = torch.ones((1, 1, 1), device="cpu").to_sparse()
|
||||
def forbidden(*args, **kwargs):
|
||||
pytest.fail("codec must not move tensors between devices")
|
||||
monkeypatch.setattr(torch.Tensor, "cpu", forbidden)
|
||||
monkeypatch.setattr(torch.Tensor, "to", forbidden)
|
||||
for tensor in (meta, sparse):
|
||||
with pytest.raises(ValueError, match="CPU|layout|strided"):
|
||||
codec.encode_payload("MASK", tensor)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sample_rate", [True, False, 0, -1, 44100.0, "44100", None])
|
||||
def test_invalid_sample_rate(sample_rate):
|
||||
with pytest.raises(ValueError, match="sample_rate"):
|
||||
codec.encode_payload("AUDIO", audio(torch.ones(1, 1, 1, device="cpu"), sample_rate))
|
||||
with pytest.raises(ValueError, match="sample_rate"):
|
||||
codec.decode_payload("AUDIO", descriptor("AUDIO", sample_rate=sample_rate), b"\x00" * 4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [None, [], "descriptor", {}, {"encoding": "utf8-v1"}])
|
||||
def test_malformed_tensor_descriptor(bad):
|
||||
with pytest.raises(ValueError, match="descriptor"):
|
||||
codec.decode_payload("MASK", bad, b"\x00" * 4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("changes", [
|
||||
{"encoding": "pickle"}, {"encoding": None}, {"encoding": []},
|
||||
{"dtype": "uint8"}, {"dtype": torch.float32}, {"dtype": []},
|
||||
{"layout": "BHWC"}, {"layout": None}, {"layout": []},
|
||||
{"shape": (1, 1, 1)}, {"shape": None}, {"shape": "1,1,1"},
|
||||
{"shape": [1, 1]}, {"shape": [1, 1, 1, 1]}, {"shape": [1, 0, 1]},
|
||||
{"shape": [1, -1, 1]}, {"shape": [1, True, 1]}, {"shape": [1, 1.0, 1]},
|
||||
{"shape": [1, "1", 1]}, {"shape": [1, [], 1]}, {"kind": "MASK"},
|
||||
])
|
||||
def test_invalid_tensor_descriptor_fields(changes):
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_payload("MASK", descriptor(**changes), b"\x00" * 4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metadata", [{}, {"encoding": "tensor-le-v1"}, {"encoding": "utf8-v1", "kind": "STRING"}, {"encoding": []}])
|
||||
def test_invalid_string_descriptor(metadata):
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_payload("STRING", metadata, b"")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [b"", b"\x00" * 3, b"\x00" * 5, b"\x00" * 8])
|
||||
def test_shape_length_mismatch(payload):
|
||||
with pytest.raises(ValueError, match="length"):
|
||||
codec.decode_payload("MASK", descriptor(), payload)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [None, "text", [0, 0, 0, 0], 4])
|
||||
def test_nonbyte_payload_rejected(payload):
|
||||
with pytest.raises(ValueError, match="bytes"):
|
||||
codec.decode_payload("MASK", descriptor(), payload)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")])
|
||||
def test_nonfinite_values_fail_on_encode_and_decode(dtype, value):
|
||||
tensor = torch.tensor([value], dtype=dtype, device="cpu").reshape(1, 1, 1)
|
||||
with pytest.raises(ValueError, match="finite"):
|
||||
codec.encode_payload("MASK", tensor)
|
||||
fmt = {torch.float16: "e", torch.float32: "f", torch.float64: "d"}
|
||||
if dtype == torch.bfloat16:
|
||||
payload = struct.pack("<H", 0x7FC0 if math.isnan(value) else (0x7F80 if value > 0 else 0xFF80))
|
||||
else:
|
||||
payload = struct.pack("<" + fmt[dtype], value)
|
||||
metadata = descriptor(dtype=str(dtype).removeprefix("torch."))
|
||||
with pytest.raises(ValueError, match="finite"):
|
||||
codec.decode_payload("MASK", metadata, payload)
|
||||
|
||||
|
||||
def test_invalid_utf8_errors_do_not_echo_content():
|
||||
with pytest.raises(ValueError) as error:
|
||||
codec.encode_payload("STRING", "private-value\ud800")
|
||||
assert "private-value" not in str(error.value)
|
||||
with pytest.raises(ValueError) as error:
|
||||
codec.decode_payload("STRING", {"encoding": "utf8-v1"}, b"private-value\xff")
|
||||
assert "private-value" not in str(error.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("max_bytes", [True, False, -1, 4.0, "4", None])
|
||||
def test_invalid_byte_budget(max_bytes):
|
||||
with pytest.raises(ValueError, match="max_bytes"):
|
||||
codec.encode_payload("STRING", "", max_bytes=max_bytes)
|
||||
with pytest.raises(ValueError, match="max_bytes"):
|
||||
codec.decode_payload("STRING", {"encoding": "utf8-v1"}, b"", max_bytes=max_bytes)
|
||||
|
||||
|
||||
def test_exact_byte_budget_and_zero_budget_for_empty_string():
|
||||
tensor = torch.ones(1, 1, 1, device="cpu")
|
||||
metadata, payload = codec.encode_payload("MASK", tensor, max_bytes=4)
|
||||
assert codec.decode_payload("MASK", metadata, payload, max_bytes=4).item() == 1
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.encode_payload("MASK", tensor, max_bytes=3)
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.decode_payload("MASK", metadata, payload, max_bytes=3)
|
||||
metadata, payload = codec.encode_payload("STRING", "", max_bytes=0)
|
||||
assert codec.decode_payload("STRING", metadata, payload, max_bytes=0) == ""
|
||||
for text in ("a", "é", "🎵"):
|
||||
byte_count = len(text.encode("utf-8"))
|
||||
metadata, payload = codec.encode_payload("STRING", text, max_bytes=byte_count)
|
||||
assert codec.decode_payload("STRING", metadata, payload, max_bytes=byte_count) == text
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.encode_payload("STRING", text, max_bytes=byte_count - 1)
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.decode_payload("STRING", metadata, payload, max_bytes=byte_count - 1)
|
||||
|
||||
|
||||
def test_oversized_shape_rejected_before_tensor_allocation(monkeypatch):
|
||||
assert codec.MAX_PAYLOAD_BYTES == 256 * 1024 * 1024
|
||||
def forbidden(*args, **kwargs):
|
||||
pytest.fail("oversized shape must fail before tensor allocation")
|
||||
for name in ("frombuffer", "empty", "tensor", "zeros", "isfinite"):
|
||||
monkeypatch.setattr(torch, name, forbidden)
|
||||
for shape in ([1, 1, codec.MAX_PAYLOAD_BYTES], [2**100, 2**100, 1]):
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.decode_payload("MASK", descriptor(shape=shape), b"\x00" * 4)
|
||||
|
||||
|
||||
def test_oversized_encode_rejected_before_materialization(monkeypatch):
|
||||
huge_view = torch.ones(1, device="cpu").expand(1, 1, codec.MAX_PAYLOAD_BYTES)
|
||||
def forbidden(*args, **kwargs):
|
||||
pytest.fail("oversized tensor must fail before materialization or scanning")
|
||||
monkeypatch.setattr(torch.Tensor, "contiguous", forbidden)
|
||||
monkeypatch.setattr(torch, "isfinite", forbidden)
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.encode_payload("MASK", huge_view)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("itemsize", [2, 4, 8])
|
||||
def test_native_little_endian_conversion_handles_big_endian(monkeypatch, itemsize):
|
||||
raw = bytes(range(itemsize * 3))
|
||||
monkeypatch.setattr(codec, "sys", types.SimpleNamespace(byteorder="little"))
|
||||
assert codec._native_little_endian(raw, itemsize) == raw
|
||||
monkeypatch.setattr(codec, "sys", types.SimpleNamespace(byteorder="big"))
|
||||
expected = b"".join(raw[i:i + itemsize][::-1] for i in range(0, len(raw), itemsize))
|
||||
assert codec._native_little_endian(raw, itemsize) == expected
|
||||
assert codec._native_little_endian(expected, itemsize) == raw
|
||||
@@ -0,0 +1,314 @@
|
||||
"""CPU-only typed collection and optional native V3 mapping acceptance."""
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
SOCKETS = ('images', 'audio', 'masks', 'texts', 'video')
|
||||
WORKERS = ['serverless:pool:1', 'serverless:pool:2']
|
||||
|
||||
|
||||
class Manager:
|
||||
def __init__(self, *, pending=False, owned=True):
|
||||
self.pending, self.owned = pending, owned
|
||||
self.cancelled = []
|
||||
|
||||
def owns(self, job, worker):
|
||||
return self.owned
|
||||
|
||||
def is_pending(self, job, worker):
|
||||
return self.pending
|
||||
|
||||
async def cancel(self, **kwargs):
|
||||
self.cancelled.append(kwargs)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def module(monkeypatch):
|
||||
package = '_media_collector_test'
|
||||
for name, path in ((package, ROOT), (package + '.nodes', ROOT / 'nodes'),
|
||||
(package + '.utils', ROOT / 'utils'),
|
||||
(package + '.serverless_bridge', ROOT / 'serverless_bridge')):
|
||||
stub = types.ModuleType(name)
|
||||
stub.__path__ = [str(path)]
|
||||
monkeypatch.setitem(sys.modules, name, stub)
|
||||
io = types.SimpleNamespace(ComfyNode=object, NodeOutput=lambda *items: types.SimpleNamespace(result=items))
|
||||
api = types.ModuleType('comfy_api.v0_0_2')
|
||||
api.io = io
|
||||
monkeypatch.setitem(sys.modules, 'comfy_api', types.ModuleType('comfy_api'))
|
||||
monkeypatch.setitem(sys.modules, 'comfy_api.v0_0_2', api)
|
||||
management = types.ModuleType('comfy.model_management')
|
||||
management.throw_exception_if_processing_interrupted = lambda: None
|
||||
comfy = types.ModuleType('comfy')
|
||||
comfy.model_management = management
|
||||
monkeypatch.setitem(sys.modules, 'comfy', comfy)
|
||||
monkeypatch.setitem(sys.modules, 'comfy.model_management', management)
|
||||
ps = types.SimpleNamespace(distributed_pending_jobs={}, distributed_jobs_lock=asyncio.Lock())
|
||||
server = types.ModuleType('server')
|
||||
server.PromptServer = types.SimpleNamespace(instance=ps)
|
||||
monkeypatch.setitem(sys.modules, 'server', server)
|
||||
helper = types.ModuleType(package + '.utils.async_helpers')
|
||||
helper.run_async_in_server_loop = asyncio.run
|
||||
monkeypatch.setitem(sys.modules, helper.__name__, helper)
|
||||
spec = importlib.util.spec_from_file_location(package + '.nodes.media_collector', ROOT / 'nodes/media_collector.py')
|
||||
loaded = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, loaded)
|
||||
spec.loader.exec_module(loaded)
|
||||
monkeypatch.setattr(loaded, 'POLL_SECONDS', 0.005)
|
||||
monkeypatch.setattr(loaded, 'START_GRACE_SECONDS', 0.025)
|
||||
loaded.test_server = ps
|
||||
loaded.test_management = management
|
||||
return loaded
|
||||
|
||||
|
||||
def native_media(offset=0):
|
||||
return {
|
||||
'images': [torch.full((1, 2 + offset, 3, 4), offset, dtype=torch.float64),
|
||||
torch.zeros((2, 1, 1, 3), dtype=torch.float16)],
|
||||
'audio': [{'waveform': torch.zeros((1, 1 + offset, 8), dtype=torch.float64),
|
||||
'sample_rate': 44100 + offset}],
|
||||
'masks': [torch.zeros((1, 2 + offset, 3), dtype=torch.float64)],
|
||||
'texts': ['', ' café\n🎥\r\n'],
|
||||
'video': [object()],
|
||||
}
|
||||
|
||||
|
||||
def setup(module, messages=(), *, manager=None):
|
||||
queue = asyncio.Queue()
|
||||
for message in messages:
|
||||
queue.put_nowait(message)
|
||||
ps = module.test_server
|
||||
ps.distributed_pending_jobs = {'job': queue, 'unrelated': asyncio.Queue()}
|
||||
ps.distributed_serverless = manager if manager is not None else Manager()
|
||||
return ps, queue
|
||||
|
||||
|
||||
def message(worker, media):
|
||||
return {'worker_id': worker, 'media': media, 'is_last': True}
|
||||
|
||||
|
||||
def collect(module, **kwargs):
|
||||
return module.DistributedMediaCollector.execute(**kwargs).result
|
||||
|
||||
|
||||
def remote(module, **kwargs):
|
||||
controls = dict(multi_job_id=['job'], enabled_worker_ids=[json.dumps(WORKERS)],
|
||||
expected_media=[json.dumps(['texts'])], delegate_only=[True])
|
||||
controls.update(kwargs)
|
||||
return collect(module, **controls)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('controls', [{}, {'pass_through': [True]}, {'enabled_worker_ids': ['[]']}])
|
||||
def test_all_five_passthrough_preserves_lists_and_native_values(module, controls):
|
||||
values = native_media()
|
||||
result = collect(module, **values, **controls)
|
||||
assert len(result) == 5
|
||||
for index, socket in enumerate(SOCKETS):
|
||||
assert result[index] is values[socket]
|
||||
|
||||
|
||||
def test_reversed_arrivals_master_then_workers_then_items(module):
|
||||
master, first, second = native_media(), native_media(1), native_media(2)
|
||||
ps, _ = setup(module, [message(WORKERS[1], second), message(WORKERS[0], first)])
|
||||
result = remote(module, **master, delegate_only=[False], expected_media=[json.dumps(list(SOCKETS))])
|
||||
for index, socket in enumerate(SOCKETS):
|
||||
expected = master[socket] + first[socket] + second[socket]
|
||||
assert len(result[index]) == len(expected)
|
||||
assert all(actual is original for actual, original in zip(result[index], expected))
|
||||
assert list(ps.distributed_pending_jobs) == ['unrelated']
|
||||
assert ps.distributed_serverless.cancelled == []
|
||||
|
||||
|
||||
def test_delegate_only_and_empty_string(module):
|
||||
ps, _ = setup(module, [message(WORKERS[1], {'texts': ['b']}), message(WORKERS[0], {'texts': ['']})])
|
||||
result = remote(module)
|
||||
assert result == ([], [], [], ['', 'b'], [])
|
||||
assert 'job' not in ps.distributed_pending_jobs
|
||||
|
||||
|
||||
@pytest.mark.parametrize('kwargs', [{}, {'texts': []}, {'texts': ''}, {'texts': ['x'], 'is_worker': [True]},
|
||||
{'texts': ['x'], 'enabled_worker_ids': ['["direct-worker"]']},
|
||||
{'texts': ['x'], 'expected_media': ['["images"]']},
|
||||
{'texts': ['x'], 'pass_through': [True, False]}, {'delegate_only': True},
|
||||
{'texts': ['x'], 'enabled_worker_ids': ['["serverless:pool:1","serverless:pool:1"]']},
|
||||
{'texts': ['x'], 'load_balance': ['false']}])
|
||||
def test_rejects_missing_media_or_invalid_controls(module, kwargs):
|
||||
with pytest.raises((ValueError, RuntimeError)):
|
||||
collect(module, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('bad', [
|
||||
{'worker_id': 'serverless:other:1', 'media': {'texts': ['x']}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': ['x']}, 'is_last': False},
|
||||
{'worker_id': WORKERS[0], 'media': {'images': [object()]}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': []}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': 'x'}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': [None]}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': ['x'] * 65}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': ['x'], 'unknown': ['x']}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'serverless_error': 'sanitized failure', 'is_last': True},
|
||||
])
|
||||
def test_rejects_bad_delivery_without_partial_output_and_cleans(module, bad):
|
||||
ps, _ = setup(module, [bad])
|
||||
with pytest.raises((ValueError, RuntimeError)):
|
||||
remote(module)
|
||||
assert list(ps.distributed_pending_jobs) == ['unrelated']
|
||||
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
|
||||
|
||||
|
||||
@pytest.mark.parametrize('after_complete', [False, True])
|
||||
def test_duplicate_delivery_rejected_including_final_queue_drain(module, after_complete):
|
||||
first = message(WORKERS[0], {'texts': ['a']})
|
||||
second = message(WORKERS[1], {'texts': ['b']})
|
||||
messages = [first, second, first] if after_complete else [first, first, second]
|
||||
ps, _ = setup(module, messages)
|
||||
with pytest.raises(RuntimeError, match='[Dd]uplicate'):
|
||||
remote(module)
|
||||
assert 'job' not in ps.distributed_pending_jobs
|
||||
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
|
||||
|
||||
|
||||
@pytest.mark.parametrize('manager', [None, Manager(owned=False), Manager(pending=False)])
|
||||
def test_missing_manager_or_finished_participant_fails_bounded(module, manager):
|
||||
ps, _ = setup(module)
|
||||
ps.distributed_serverless = manager
|
||||
with pytest.raises(RuntimeError, match='[Ss]erverless'):
|
||||
remote(module)
|
||||
assert 'job' not in ps.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_control_error_also_cleans_exact_job(module):
|
||||
ps, _ = setup(module)
|
||||
with pytest.raises(ValueError):
|
||||
remote(module, expected_media='["unknown"]')
|
||||
assert list(ps.distributed_pending_jobs) == ['unrelated']
|
||||
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
|
||||
|
||||
|
||||
def test_launch_race_and_pending_deadline_owned_by_manager(module, monkeypatch):
|
||||
ps, queue = setup(module)
|
||||
ps.distributed_serverless = None
|
||||
async def runner(coro):
|
||||
async def produce():
|
||||
await asyncio.sleep(0.01)
|
||||
ps.distributed_serverless = Manager(pending=True)
|
||||
await asyncio.sleep(0.035)
|
||||
queue.put_nowait(message(WORKERS[1], {'texts': ['b']}))
|
||||
queue.put_nowait(message(WORKERS[0], {'texts': ['a']}))
|
||||
task = asyncio.create_task(produce())
|
||||
result = await coro
|
||||
await task
|
||||
return result
|
||||
monkeypatch.setattr(module, 'run_async_in_server_loop', lambda coro: asyncio.run(runner(coro)))
|
||||
assert remote(module)[3] == ['a', 'b']
|
||||
assert ps.distributed_serverless.cancelled == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize('cancel_task', [False, True])
|
||||
def test_cancellation_while_blocked_cleans_and_propagates(module, monkeypatch, cancel_task):
|
||||
ps, _ = setup(module, manager=Manager(pending=True))
|
||||
class Interrupted(Exception):
|
||||
pass
|
||||
async def runner(coro):
|
||||
task = asyncio.create_task(coro)
|
||||
await asyncio.sleep(0.012)
|
||||
if cancel_task:
|
||||
task.cancel()
|
||||
else:
|
||||
def interrupt():
|
||||
raise Interrupted('stop')
|
||||
module.test_management.throw_exception_if_processing_interrupted = interrupt
|
||||
return await task
|
||||
monkeypatch.setattr(module, 'run_async_in_server_loop', lambda coro: asyncio.run(runner(coro)))
|
||||
with pytest.raises(asyncio.CancelledError if cancel_task else Interrupted):
|
||||
remote(module)
|
||||
assert list(ps.distributed_pending_jobs) == ['unrelated']
|
||||
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
|
||||
|
||||
|
||||
def test_replaced_queue_not_removed(module, monkeypatch):
|
||||
ps, _ = setup(module, manager=Manager(pending=True))
|
||||
replacement = asyncio.Queue()
|
||||
async def runner(coro):
|
||||
task = asyncio.create_task(coro)
|
||||
await asyncio.sleep(0.012)
|
||||
ps.distributed_pending_jobs['job'] = replacement
|
||||
return await task
|
||||
monkeypatch.setattr(module, 'run_async_in_server_loop', lambda coro: asyncio.run(runner(coro)))
|
||||
with pytest.raises(RuntimeError, match='queue'):
|
||||
remote(module)
|
||||
assert ps.distributed_pending_jobs['job'] is replacement
|
||||
|
||||
|
||||
def test_native_v3_schema_and_list_mapping_cpu():
|
||||
core = os.environ.get('COMFYUI_SOURCE_ROOT')
|
||||
if not core:
|
||||
pytest.skip('Set COMFYUI_SOURCE_ROOT for CPU-only native V3 mapping')
|
||||
probe = r'''
|
||||
import asyncio, importlib.util, sys, types
|
||||
from pathlib import Path
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
import execution
|
||||
import torch
|
||||
import server
|
||||
# No HTTP listener or PromptServer bootstrap is needed for schema/list mapping.
|
||||
server.PromptServer.instance = types.SimpleNamespace()
|
||||
root = Path(sys.argv[2])
|
||||
for name, path in [('typed_collector', root), ('typed_collector.nodes', root/'nodes'),
|
||||
('typed_collector.utils', root/'utils'),
|
||||
('typed_collector.serverless_bridge', root/'serverless_bridge')]:
|
||||
package = types.ModuleType(name)
|
||||
package.__path__ = [str(path)]
|
||||
sys.modules[name] = package
|
||||
spec = importlib.util.spec_from_file_location('typed_collector.nodes.media_collector', root/'nodes/media_collector.py')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
cls = module.DistributedMediaCollector
|
||||
info = cls.GET_NODE_INFO_V1()
|
||||
assert info['output'] == ['IMAGE','AUDIO','MASK','STRING','VIDEO'], info
|
||||
assert info['output_is_list'] == [True]*5, info
|
||||
assert cls.define_schema().is_input_list and cls.define_schema().accept_all_inputs
|
||||
for name, kind in zip(['images','audio','masks','texts','video'], info['output']):
|
||||
field = info['input']['optional'][name]
|
||||
assert field[0] == kind and field[1]['forceInput'] is True, field
|
||||
assert info['input']['required']['load_balance'][1]['default'] is False
|
||||
async def check():
|
||||
values = {'load_balance':False, 'pass_through':True}
|
||||
prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe')
|
||||
assert not missing, missing
|
||||
image = torch.zeros((1,2,3,4), dtype=torch.float64)
|
||||
audio = {'waveform':torch.zeros((1,2,7)), 'sample_rate':48000}
|
||||
mask = torch.ones((1,2,3), dtype=torch.float64)
|
||||
from comfy_api.input_impl import VideoFromComponents
|
||||
from comfy_api.util import VideoComponents
|
||||
from fractions import Fraction
|
||||
video = VideoFromComponents(VideoComponents(images=image[:,:,:,:3].float(), frame_rate=Fraction(24,1)))
|
||||
native = {'images':[image], 'audio':[audio], 'masks':[mask], 'texts':[''], 'video':[video]}
|
||||
prepared.update(native)
|
||||
results = await execution._async_map_node_over_list('cpu','probe',cls,prepared,cls.FUNCTION,v3_data=hidden)
|
||||
results = await execution.resolve_map_node_over_list_results(results)
|
||||
assert len(results) == 1
|
||||
for index, name in enumerate(native):
|
||||
assert results[0].result[index] is native[name]
|
||||
assert results[0].result[4][0] is video
|
||||
print('MEDIA_COLLECTOR_CPU_OK')
|
||||
asyncio.run(check())
|
||||
'''
|
||||
env = {**os.environ, 'CUDA_VISIBLE_DEVICES': '', 'HIP_VISIBLE_DEVICES': '', 'PYTHONDONTWRITEBYTECODE': '1'}
|
||||
env.pop('PYTHONPATH', None)
|
||||
result = subprocess.run([sys.executable, '-c', probe, core, str(ROOT)], env=env,
|
||||
capture_output=True, text=True, timeout=90)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert 'MEDIA_COLLECTOR_CPU_OK' in result.stdout
|
||||
@@ -0,0 +1,152 @@
|
||||
"""Offline contract tests: malformed results never reach transport/delivery."""
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
spec = importlib.util.spec_from_file_location("media_contract_test", ROOT / "serverless_bridge" / "contract.py")
|
||||
media = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(media)
|
||||
|
||||
|
||||
def binding(**changes):
|
||||
return dict(binding_id="binding-1", collector_id="100", job_id="job-1",
|
||||
participant_id="serverless:worker:0", node_id="export-1", socket="texts", kind="STRING", **changes)
|
||||
|
||||
|
||||
def response():
|
||||
artifact = media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b" \n")
|
||||
return {"version": 1, "profile": media.PROFILE, "request_id": "request-1",
|
||||
"bindings": [dict(binding(), count=1)], "artifacts": [artifact],
|
||||
"items": [{"binding_id": "binding-1", "item_index": 0, "artifact_id": "a"}]}
|
||||
|
||||
|
||||
def validate(value):
|
||||
return media.validate_response(value, request_id="request-1", bindings=[binding()])
|
||||
|
||||
|
||||
def test_valid_manifest_and_exact_inline_payload():
|
||||
value = response()
|
||||
ordered = validate(value)
|
||||
assert asyncio.run(media.resolve_artifact(ordered["binding-1"][0])) == b" \n"
|
||||
|
||||
|
||||
def test_order_uses_indices_not_arrival():
|
||||
value = response()
|
||||
value["bindings"][0]["count"] = 2
|
||||
value["artifacts"].append(media.artifact_from_payload("b", "STRING", {"encoding": "utf8-v1"}, b"second"))
|
||||
value["items"] = [{"binding_id": "binding-1", "item_index": 1, "artifact_id": "b"}, value["items"][0]]
|
||||
assert [item["id"] for item in validate(value)["binding-1"]] == ["a", "b"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field,value", [("version", True), ("version", 2), ("profile", "legacy"), ("request_id", "another")])
|
||||
def test_envelope_identity(field, value):
|
||||
obj = response()
|
||||
obj[field] = value
|
||||
with pytest.raises(ValueError):
|
||||
validate(obj)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["collector_id", "participant_id", "job_id", "node_id", "binding_id"])
|
||||
def test_binding_identity(field):
|
||||
value = response()
|
||||
value["bindings"][0][field] = "another"
|
||||
with pytest.raises(ValueError):
|
||||
validate(value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("count", [0, -1, True, 1.0, "1", 65])
|
||||
def test_invalid_count(count):
|
||||
value = response()
|
||||
value["bindings"][0]["count"] = count
|
||||
with pytest.raises(ValueError):
|
||||
validate(value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", ["missing_binding", "missing_item", "extra_item", "duplicate_binding", "duplicate_artifact", "unused_artifact", "wrong_type", "wrong_index"])
|
||||
def test_incomplete_or_ambiguous_manifest(mutation):
|
||||
value = response()
|
||||
if mutation == "missing_binding":
|
||||
value["bindings"] = []
|
||||
elif mutation == "missing_item":
|
||||
value["items"] = []
|
||||
elif mutation == "extra_item":
|
||||
value["items"] *= 2
|
||||
elif mutation == "duplicate_binding":
|
||||
value["bindings"] *= 2
|
||||
elif mutation == "duplicate_artifact":
|
||||
value["artifacts"] *= 2
|
||||
elif mutation == "unused_artifact":
|
||||
artifact = copy.deepcopy(value["artifacts"][0])
|
||||
artifact["id"] = "unused"
|
||||
value["artifacts"].append(artifact)
|
||||
elif mutation == "wrong_type":
|
||||
value["artifacts"][0]["kind"] = "MASK"
|
||||
elif mutation == "wrong_index":
|
||||
value["items"][0]["item_index"] = True
|
||||
with pytest.raises(ValueError):
|
||||
validate(value)
|
||||
|
||||
|
||||
def test_duplicate_item_indices():
|
||||
value = response()
|
||||
value["bindings"][0]["count"] = 2
|
||||
value["items"] *= 2
|
||||
with pytest.raises(ValueError, match="Duplicate"):
|
||||
validate(value)
|
||||
|
||||
|
||||
def test_total_budget():
|
||||
with pytest.raises(ValueError, match="total byte"):
|
||||
media.validate_response(response(), request_id="request-1", bindings=[binding()], max_total_bytes=1)
|
||||
|
||||
|
||||
def test_checksum_and_base64_validation():
|
||||
artifact = response()["artifacts"][0]
|
||||
artifact["sha256"] = "0" * 64
|
||||
with pytest.raises(ValueError, match="checksum"):
|
||||
asyncio.run(media.resolve_artifact(artifact))
|
||||
artifact["transport"]["data_base64"] = "!!!!"
|
||||
with pytest.raises(ValueError, match="base64"):
|
||||
asyncio.run(media.resolve_artifact(artifact))
|
||||
|
||||
|
||||
def test_empty_text_is_a_value():
|
||||
artifact = media.artifact_from_payload("empty", "STRING", {"encoding": "utf8-v1"}, b"")
|
||||
assert asyncio.run(media.resolve_artifact(artifact)) == b""
|
||||
|
||||
|
||||
@pytest.mark.parametrize("url,expires", [("http://example.com/a", "2099-01-01T00:00:00Z"),
|
||||
("https://user:secret@example.com/a", "2099-01-01T00:00:00Z"),
|
||||
("https://example.com/a", "2000-01-01T00:00:00Z"),
|
||||
("https://example.com/a", "2099-01-01"),
|
||||
("https://example.com:1234/a", "2099-01-01T00:00:00Z")])
|
||||
def test_invalid_url_envelopes(url, expires):
|
||||
with pytest.raises(ValueError):
|
||||
media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b"a",
|
||||
transport={"type": "url", "url": url, "expires_at": expires})
|
||||
|
||||
|
||||
def test_url_resolution_is_bounded_and_hash_verified():
|
||||
artifact = media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b"hello",
|
||||
transport={"type": "url", "url": "https://example.com/a", "expires_at": "2099-01-01T00:00:00Z"})
|
||||
calls = []
|
||||
async def download(url, *, max_bytes):
|
||||
calls.append((url, max_bytes))
|
||||
return b"hello"
|
||||
assert asyncio.run(media.resolve_artifact(artifact, download=download)) == b"hello"
|
||||
assert calls == [("https://example.com/a", 5)]
|
||||
|
||||
|
||||
def test_descriptor_depth_and_nonfinite_rejected():
|
||||
value = {}
|
||||
cursor = value
|
||||
for _ in range(15):
|
||||
cursor["next"] = {}
|
||||
cursor = cursor["next"]
|
||||
for descriptor in [value, {"x": float("nan")}, {"x": object()}]:
|
||||
with pytest.raises(ValueError):
|
||||
media.artifact_from_payload("a", "STRING", descriptor, b"")
|
||||
@@ -0,0 +1,47 @@
|
||||
"""Opt-in, isolated CPU acceptance of the complete typed-media data path.
|
||||
|
||||
COMFYUI_SOURCE_ROOT must name an existing compatible checkout; run with that
|
||||
checkout's Python environment. No installation, normal instance, or provider is
|
||||
used. The child owns its loopback server and every writable Comfy directory.
|
||||
"""
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def test_typed_media_end_to_end_cpu():
|
||||
configured = os.environ.get("COMFYUI_SOURCE_ROOT")
|
||||
if configured is None:
|
||||
pytest.skip("Set COMFYUI_SOURCE_ROOT for real CPU typed-media end-to-end acceptance")
|
||||
core = Path(configured).expanduser().resolve()
|
||||
assert configured and all((core / name).is_file() for name in (
|
||||
"execution.py", "server.py", "comfy/cli_args.py",
|
||||
"comfy_api/latest/_input_impl/video_types.py",
|
||||
)), "Configured COMFYUI_SOURCE_ROOT must provide PromptExecutor and native V3 VIDEO"
|
||||
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
env = dict(os.environ)
|
||||
for name in ("PYTHONPATH", "COMFYUI_SOURCE_ROOT"):
|
||||
env.pop(name, None)
|
||||
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="", ROCR_VISIBLE_DEVICES="",
|
||||
PYTHONDONTWRITEBYTECODE="1", PYTHONHASHSEED="0", COMFYUI_IS_WORKER="1",
|
||||
OMP_NUM_THREADS="1", MKL_NUM_THREADS="1")
|
||||
with tempfile.TemporaryDirectory(prefix="media-e2e-", dir=scratch) as run_root:
|
||||
env["TMPDIR"] = run_root
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(ROOT / "tests/helpers/media_end_to_end_check.py"),
|
||||
str(core), run_root], cwd=run_root, env=env,
|
||||
capture_output=True, text=True, timeout=150,
|
||||
)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
for marker in ("NATIVE_INPUTS_AND_DIRECT_EXPORT_REFS_OK", "REAL_HTTP_BRIDGE_RESULTS_OK",
|
||||
"MASTER_AND_DELEGATE_ORDER_OK", "ATOMIC_DECODE_REJECTION_OK",
|
||||
"MEDIA_END_TO_END_CPU_OK"):
|
||||
assert marker in result.stdout, result.stdout + result.stderr
|
||||
print(result.stdout)
|
||||
@@ -0,0 +1,877 @@
|
||||
"""CPU-only endpoint adapter tests against an actual loopback Comfy HTTP fake."""
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import hashlib
|
||||
import importlib
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE = "media_handler_testpkg"
|
||||
package = types.ModuleType(PACKAGE)
|
||||
package.__path__ = [str(ROOT / "serverless_bridge")]
|
||||
sys.modules[PACKAGE] = package
|
||||
handler = importlib.import_module(PACKAGE + ".handler")
|
||||
contract = importlib.import_module(PACKAGE + ".contract")
|
||||
|
||||
|
||||
def request(kind="STRING", *, incoming=True):
|
||||
socket = {value: key for key, value in contract.SOCKET_KINDS.items()}[kind]
|
||||
binding = dict(binding_id="binding:1", collector_id="collector:1", job_id="job/42",
|
||||
participant_id="worker:0", node_id="export:1", socket=socket, kind=kind)
|
||||
artifact = contract.artifact_from_payload("input:1", kind, {"encoding": "test"}, b"input bytes")
|
||||
return {"workflow": {
|
||||
"source:1": {"class_type": "Source", "inputs": {"value": "literal"}},
|
||||
"export:1": {"class_type": "MediaExportV1", "inputs": {
|
||||
"request_id": "request:1", "binding_json": json.dumps(binding), socket: ["source:1", 0]}},
|
||||
}, "media": {"version": 1, "profile": contract.PROFILE, "request_id": "request:1",
|
||||
"artifacts": [artifact] if incoming else [],
|
||||
"inputs": [{"node_id": "source:1", "input_name": "value", "artifact_id": "input:1"}] if incoming else [],
|
||||
"outputs": [binding]}}
|
||||
|
||||
|
||||
class Comfy:
|
||||
def __init__(self, root, data=None, *, payload=b"result bytes"):
|
||||
self.root = root
|
||||
self.data = data or request()
|
||||
self.payload = payload
|
||||
self.submitted = []
|
||||
self.views = []
|
||||
self.deletes = []
|
||||
self.interrupts = []
|
||||
self.polls = 0
|
||||
self.active = False
|
||||
self.mode = "success"
|
||||
self.history_failures = 0
|
||||
self.record_mutation = None
|
||||
self.prompt_seen = asyncio.Event()
|
||||
self.terminal_seen = False
|
||||
|
||||
def staged(self):
|
||||
return list(self.root.glob("distributed-media-v1/*/*.bin"))
|
||||
|
||||
def record(self):
|
||||
binding = self.data["media"]["outputs"][0]
|
||||
aid = "a" * 32
|
||||
artifact = contract.artifact_from_payload(aid, binding["kind"], {"encoding": "test"}, self.payload)
|
||||
artifact["transport"] = {"type": "comfy-file", "filename": aid + ".bin",
|
||||
"subfolder": "distributed-media-v1/" + hashlib.sha256(b"request:1").hexdigest(),
|
||||
"folder_type": "temp"}
|
||||
value = {"request_id": "request:1", "binding": dict(binding, count=1), "artifacts": [artifact],
|
||||
"items": [{"binding_id": binding["binding_id"], "item_index": 0, "artifact_id": aid}]}
|
||||
if self.record_mutation:
|
||||
self.record_mutation(value)
|
||||
return value
|
||||
|
||||
async def submit(self, req):
|
||||
data = await req.json()
|
||||
self.submitted.append(data)
|
||||
self.active = True
|
||||
self.prompt_seen.set()
|
||||
if self.data["media"]["artifacts"]:
|
||||
assert self.staged()
|
||||
link = data["prompt"]["source:1"]["inputs"]["value"]
|
||||
imported = data["prompt"][link[0]]
|
||||
assert imported["class_type"] == "MediaImportV1"
|
||||
metadata = json.loads(imported["inputs"]["artifact_json"])
|
||||
assert "transport" not in metadata
|
||||
assert (self.root / imported["inputs"]["filename"]).read_bytes() == b"input bytes"
|
||||
if self.mode == "ambiguous":
|
||||
return web.Response(status=500)
|
||||
if self.mode == "partial":
|
||||
return web.json_response({"prompt_id": "prompt:1", "node_errors": {"export:1": {"errors": ["bad"]}}})
|
||||
return web.json_response({"prompt_id": "prompt:1", "node_errors": {}})
|
||||
|
||||
async def history(self, req):
|
||||
self.polls += 1
|
||||
if self.history_failures:
|
||||
self.history_failures -= 1
|
||||
return web.Response(status=503)
|
||||
if self.mode in {"queued", "running", "unrelated", "ambiguous"}:
|
||||
return web.json_response({})
|
||||
if self.data["media"]["artifacts"]:
|
||||
assert self.staged(), "staging must survive until terminal confirmation"
|
||||
self.active = False
|
||||
self.terminal_seen = True
|
||||
if self.mode in {"error", "interrupted"}:
|
||||
status = {"completed": False, "status_str": "error", "messages": [["execution_interrupted", {}]]}
|
||||
else:
|
||||
status = {"completed": True, "status_str": "success", "messages": []}
|
||||
return web.json_response({"prompt:1": {"status": status, "outputs": {} if self.mode == "missing" else {
|
||||
"export:1": {"media_exports": [self.record()]}}}})
|
||||
|
||||
async def view(self, req):
|
||||
self.views.append(dict(req.query))
|
||||
assert self.terminal_seen
|
||||
payload = self.payload if self.mode != "corrupt" else b"X" * len(self.payload)
|
||||
if self.mode == "oversized":
|
||||
payload += b"X"
|
||||
return web.Response(body=payload)
|
||||
|
||||
def entry(self, pid="prompt:1"):
|
||||
extra = dict(self.submitted[0].get("extra_data", {}), client_id=self.submitted[0]["client_id"])
|
||||
return [1, pid, self.submitted[0]["prompt"], extra, ["export:1"]]
|
||||
|
||||
async def queue(self, req):
|
||||
if req.method == "POST":
|
||||
body = await req.json()
|
||||
assert body == {"delete": ["prompt:1"]}
|
||||
self.deletes.append(body)
|
||||
assert self.staged()
|
||||
self.mode = "deleted"
|
||||
self.active = False
|
||||
return web.Response(status=200)
|
||||
running = [self.entry()] if self.mode == "running" else []
|
||||
if self.mode == "unrelated":
|
||||
running = [self.entry("not-our-prompt")]
|
||||
pending = [self.entry()] if self.mode in {"queued", "ambiguous"} else []
|
||||
return web.json_response({"queue_running": running, "queue_pending": pending})
|
||||
|
||||
async def interrupt(self, req):
|
||||
assert self.mode == "running"
|
||||
assert self.staged()
|
||||
self.interrupts.append(await req.json())
|
||||
self.mode = "interrupted"
|
||||
return web.Response(status=200)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def serve(comfy):
|
||||
app = web.Application()
|
||||
app.add_routes([web.post("/prompt", comfy.submit), web.get("/history/{prompt_id}", comfy.history),
|
||||
web.get("/view", comfy.view), web.get("/queue", comfy.queue),
|
||||
web.post("/queue", comfy.queue), web.post("/interrupt", comfy.interrupt)])
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
try:
|
||||
yield f"http://127.0.0.1:{port}"
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
def run(coro):
|
||||
return asyncio.run(coro)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind,index", [("IMAGE", 0), ("AUDIO", 1), ("MASK", 2), ("STRING", 3), ("VIDEO", 4)])
|
||||
def test_loopback_submit_poll_view_cleanup_and_caller_ownership(tmp_path, kind, index):
|
||||
async def check():
|
||||
data = request(kind)
|
||||
original = copy.deepcopy(data)
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url, aiohttp.ClientSession() as session:
|
||||
result = await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, session=session)
|
||||
assert not session.closed
|
||||
contract.validate_response(result["media"], request_id="request:1", bindings=data["media"]["outputs"])
|
||||
assert await contract.resolve_artifact(result["media"]["artifacts"][0]) == comfy.payload
|
||||
assert len(comfy.submitted) == 1
|
||||
assert comfy.submitted[0]["partial_execution_targets"] == ["export:1"]
|
||||
assert comfy.submitted[0]["prompt"]["source:1"]["inputs"]["value"][1] == index
|
||||
assert comfy.views == [{"filename": "a" * 32 + ".bin", "subfolder": "distributed-media-v1/" + hashlib.sha256(b"request:1").hexdigest(), "type": "temp"}]
|
||||
assert data == original
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode,match", [("corrupt", "checksum"), ("missing", "Missing"), ("error", "execution"), ("oversized", "length")])
|
||||
def test_terminal_failures_clean_only_after_terminal(tmp_path, mode, match):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = mode
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises((ValueError, RuntimeError), match=match):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert len(comfy.submitted) == 1
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", [
|
||||
lambda d: d["media"].update(version=True),
|
||||
lambda d: d["media"].update(profile="legacy"),
|
||||
lambda d: d["media"].update(request_id="../escape"),
|
||||
lambda d: d["media"]["outputs"][0].update(node_id="missing"),
|
||||
lambda d: d["workflow"]["export:1"].update(class_type="SaveImage"),
|
||||
lambda d: d["workflow"]["export:1"]["inputs"].update(request_id="other"),
|
||||
lambda d: d["workflow"]["export:1"]["inputs"].update(binding_json="{}"),
|
||||
lambda d: d["workflow"]["export:1"]["inputs"].update(images=["source:1", 0]),
|
||||
lambda d: d["workflow"]["export:1"]["inputs"].update(texts="literal"),
|
||||
lambda d: d["workflow"].update(unrelated={"class_type": "SaveImage", "inputs": {}}),
|
||||
lambda d: d["workflow"].update(unsafe={"class_type": "MediaImportV1", "inputs": {}}),
|
||||
lambda d: d["media"]["inputs"].append(copy.deepcopy(d["media"]["inputs"][0])),
|
||||
lambda d: d["media"]["inputs"][0].update(node_id="missing"),
|
||||
lambda d: d["media"]["inputs"][0].update(input_name="missing"),
|
||||
lambda d: d["media"].update(inputs=[]),
|
||||
lambda d: d["media"]["artifacts"][0].update(sha256="0" * 64),
|
||||
])
|
||||
def test_rejection_before_submit(tmp_path, mutation):
|
||||
async def check():
|
||||
data = request()
|
||||
mutation(data)
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.submitted
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", [
|
||||
lambda r: r.update(request_id="other"),
|
||||
lambda r: r["binding"].update(count=2),
|
||||
lambda r: r["binding"].update(collector_id="other"),
|
||||
lambda r: r["items"][0].update(item_index=1),
|
||||
lambda r: r["artifacts"][0]["transport"].update(filename="../escape"),
|
||||
lambda r: r["artifacts"][0]["transport"].update(subfolder="distributed-media-v1/other"),
|
||||
lambda r: r["artifacts"][0]["transport"].update(folder_type="input"),
|
||||
lambda r: r["artifacts"][0].update(byte_length=contract.MAX_ARTIFACT_BYTES + 1),
|
||||
])
|
||||
def test_bad_result_manifest_rejected_before_any_file_read(tmp_path, mutation):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.record_mutation = mutation
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.views
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_idempotent_history_retry_without_resubmission(tmp_path):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.history_failures = 2
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert comfy.polls == 3
|
||||
assert len(comfy.submitted) == 1
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["queued"])
|
||||
def test_deadline_cancels_only_own_prompt_before_cleanup(tmp_path, monkeypatch, mode):
|
||||
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = mode
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(TimeoutError):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, timeout_seconds=0.04)
|
||||
assert len(comfy.submitted) == 1
|
||||
assert bool(comfy.deletes) == (mode == "queued")
|
||||
assert bool(comfy.interrupts) == (mode == "running")
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_unconfirmed_terminal_retains_files_and_never_interrupts_unrelated(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
|
||||
monkeypatch.setattr(handler, "CANCEL_GRACE_SECONDS", 0.06)
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = "unrelated"
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(RuntimeError, match="retained"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, timeout_seconds=0.04)
|
||||
assert comfy.staged()
|
||||
assert not comfy.interrupts
|
||||
assert not comfy.deletes
|
||||
assert len(comfy.submitted) == 1
|
||||
run(check())
|
||||
|
||||
|
||||
def test_ambiguous_post_is_not_retried_and_own_queue_entry_is_reconciled(tmp_path):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = "ambiguous"
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(RuntimeError):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert len(comfy.submitted) == 1
|
||||
assert len(comfy.deletes) == 1
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_task_cancellation_fails_closed_and_preserves_caller_session(tmp_path):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = "running"
|
||||
async with serve(comfy) as url, aiohttp.ClientSession() as session:
|
||||
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, session=session))
|
||||
await comfy.prompt_seen.wait()
|
||||
await asyncio.sleep(0.01)
|
||||
task.cancel()
|
||||
with pytest.raises(handler.MediaExecutionUncertain, match="prompt-scoped interruption unavailable"):
|
||||
await task
|
||||
assert not session.closed
|
||||
assert not comfy.interrupts
|
||||
assert comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_legacy_images_are_owned_and_rewritten_without_overwriting_user_files(tmp_path):
|
||||
async def check():
|
||||
data = request(incoming=False)
|
||||
data["workflow"]["source:1"] = {"class_type": "LoadImage", "inputs": {"image": "same.png"}}
|
||||
data["images"] = [{"name": "same.png", "image": base64.b64encode(b"new image").decode()}]
|
||||
(tmp_path / "same.png").write_bytes(b"user image")
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
|
||||
filename = comfy.submitted[0]["prompt"]["source:1"]["inputs"]["image"]
|
||||
assert filename.startswith("distributed-media-v1/")
|
||||
assert filename.endswith(".bin")
|
||||
assert (tmp_path / "same.png").read_bytes() == b"user image"
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_url_input_uses_bounded_injected_downloader(tmp_path):
|
||||
async def check():
|
||||
data = request()
|
||||
data["media"]["artifacts"][0]["transport"] = {"type": "url", "url": "https://example.org/a", "expires_at": "2099-01-01T00:00:00Z"}
|
||||
calls = []
|
||||
async def download(url, *, max_bytes):
|
||||
calls.append((url, max_bytes))
|
||||
return b"input bytes"
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=download)
|
||||
assert calls == [("https://example.org/a", len(b"input bytes"))]
|
||||
run(check())
|
||||
|
||||
|
||||
def test_large_result_requires_publisher_and_validates_returned_transport(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path, payload=b"Z" * 100_000)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="publisher"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.views
|
||||
assert not comfy.staged()
|
||||
published = []
|
||||
async def publish(metadata, payload):
|
||||
assert "transport" not in metadata
|
||||
assert payload == comfy.payload
|
||||
published.append(metadata)
|
||||
return {"type": "url", "url": "https://example.org/result", "expires_at": "2099-01-01T00:00:00Z"}
|
||||
async with serve(comfy) as url:
|
||||
result = await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, publish=publish)
|
||||
assert len(published) == 1
|
||||
assert result["media"]["artifacts"][0]["transport"]["type"] == "url"
|
||||
assert len(json.dumps(result).encode()) <= 100_000
|
||||
async def bad_publish(metadata, payload):
|
||||
return {"type": "url", "url": "http://example.org/result", "expires_at": "2099-01-01T00:00:00Z"}
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, publish=bad_publish)
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", ["../photo.png", "folder/photo.png", "folder\\photo.png", ".", "..", "/photo.png"])
|
||||
def test_legacy_basename_rejected_without_submit(tmp_path, name):
|
||||
async def check():
|
||||
data = request(incoming=False)
|
||||
data["images"] = [{"name": name, "image": "YQ=="}]
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.submitted
|
||||
run(check())
|
||||
|
||||
|
||||
def test_staging_rejects_symlink_without_writing_or_submit(tmp_path):
|
||||
async def check():
|
||||
outside = tmp_path / "unrelated"
|
||||
outside.mkdir()
|
||||
(tmp_path / "distributed-media-v1").symlink_to(outside, target_is_directory=True)
|
||||
comfy = Comfy(tmp_path)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="symlink"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.submitted
|
||||
assert not list(outside.iterdir())
|
||||
run(check())
|
||||
|
||||
|
||||
def test_cleanup_preserves_other_runs_in_same_request_namespace(tmp_path):
|
||||
async def check():
|
||||
directory = tmp_path / "distributed-media-v1" / hashlib.sha256(b"request:1").hexdigest()
|
||||
directory.mkdir(parents=True)
|
||||
sibling = directory / ("b" * 32 + ".bin")
|
||||
sibling.write_bytes(b"other run")
|
||||
comfy = Comfy(tmp_path)
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert sibling.read_bytes() == b"other run"
|
||||
assert comfy.staged() == [sibling]
|
||||
run(check())
|
||||
|
||||
|
||||
def test_input_total_limit_checked_before_downloader_or_submit(tmp_path):
|
||||
async def check():
|
||||
data = request()
|
||||
data["media"]["artifacts"] = []
|
||||
data["media"]["inputs"] = []
|
||||
for index in range(3):
|
||||
artifact = contract.artifact_from_payload(f"input:{index}", "STRING", {}, b"",
|
||||
transport={"type": "url", "url": "https://example.org/input", "expires_at": "2099-01-01T00:00:00Z"})
|
||||
artifact["byte_length"] = contract.MAX_ARTIFACT_BYTES
|
||||
data["media"]["artifacts"].append(artifact)
|
||||
name = f"value:{index}"
|
||||
data["workflow"]["source:1"]["inputs"][name] = "literal"
|
||||
data["media"]["inputs"].append({"node_id": "source:1", "input_name": name, "artifact_id": artifact["id"]})
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async def forbidden_download(*args, **kwargs):
|
||||
pytest.fail("oversized request reached downloader")
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="total byte"):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=forbidden_download)
|
||||
assert not comfy.submitted
|
||||
run(check())
|
||||
|
||||
|
||||
def test_output_total_limit_checked_before_fetch(tmp_path):
|
||||
def excessive(record):
|
||||
first = record["artifacts"][0]
|
||||
record["artifacts"] = []
|
||||
record["items"] = []
|
||||
record["binding"]["count"] = 3
|
||||
for index, char in enumerate("abc"):
|
||||
artifact = copy.deepcopy(first)
|
||||
artifact["id"] = char * 32
|
||||
artifact["byte_length"] = contract.MAX_ARTIFACT_BYTES
|
||||
artifact["transport"]["filename"] = artifact["id"] + ".bin"
|
||||
record["artifacts"].append(artifact)
|
||||
record["items"].append({"binding_id": "binding:1", "item_index": index, "artifact_id": artifact["id"]})
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.record_mutation = excessive
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="total byte"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.views
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_partial_target_acceptance_is_not_reported_as_success(tmp_path):
|
||||
async def check():
|
||||
comfy = Comfy(tmp_path)
|
||||
comfy.mode = "partial"
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(RuntimeError, match="incomplete"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
|
||||
assert len(comfy.submitted) == 1
|
||||
assert not comfy.views
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
# Native routes are compiled from the pinned source, without importing ComfyUI
|
||||
# (which would initialize the GPU). All dependencies and HTTP are CPU-only.
|
||||
@pytest.mark.parametrize("current,requested,race,expected", [
|
||||
("own", "own", False, ["own"]),
|
||||
("other", "own", False, []),
|
||||
("own", "own", True, ["successor"]),
|
||||
])
|
||||
def test_pinned_native_interrupt_route_has_snapshot_race(current, requested, race, expected):
|
||||
import ast
|
||||
import logging
|
||||
source = Path("/home/robert/ComfyUI_dev/server.py")
|
||||
if not source.is_file():
|
||||
pytest.skip("Pinned ComfyUI source unavailable")
|
||||
tree = ast.parse(source.read_text())
|
||||
routes = [node for node in ast.walk(tree) if isinstance(node, ast.AsyncFunctionDef)
|
||||
and any("/interrupt" in ast.unparse(d) for d in node.decorator_list)]
|
||||
assert len(routes) == 1
|
||||
route = routes[0]
|
||||
route.decorator_list = []
|
||||
interrupts = []
|
||||
class Queue:
|
||||
running = current
|
||||
def get_current_queue(self):
|
||||
snapshot = [[1, self.running, {}, {}, []]]
|
||||
if race:
|
||||
self.running = "successor"
|
||||
return snapshot, []
|
||||
queue = Queue()
|
||||
scope = {"json": json, "logging": logging, "web": web,
|
||||
"self": types.SimpleNamespace(prompt_queue=queue),
|
||||
"nodes": types.SimpleNamespace(interrupt_processing=lambda: interrupts.append(queue.running))}
|
||||
exec(compile(ast.Module(body=[route], type_ignores=[]), str(source), "exec"), scope)
|
||||
async def check():
|
||||
async def body():
|
||||
return {"prompt_id": requested}
|
||||
response = await scope[route.name](types.SimpleNamespace(json=body))
|
||||
assert response.status == 200
|
||||
run(check())
|
||||
assert interrupts == expected
|
||||
|
||||
|
||||
def exporter_file(root, *, name="a" * 32 + ".bin", payload=b"result bytes"):
|
||||
directory = root / "distributed-media-v1" / hashlib.sha256(b"request:1").hexdigest()
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
path = directory / name
|
||||
path.write_bytes(payload)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["success", "corrupt", "oversized", "error", "interrupted", "partial"])
|
||||
def test_owned_export_cleanup_after_terminal_success_or_failure(tmp_path, mode):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
sibling = exporter_file(temp, name="b" * 32 + ".bin", payload=b"unrelated")
|
||||
ordinary = exporter_file(temp, name="keep.txt", payload=b"ordinary")
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
comfy.mode = mode
|
||||
view = comfy.view
|
||||
async def reading(req):
|
||||
assert owned.exists(), "cleanup raced artifact collection"
|
||||
return await view(req)
|
||||
comfy.view = reading
|
||||
async with serve(comfy) as url:
|
||||
if mode == "success":
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
else:
|
||||
with pytest.raises((ValueError, RuntimeError)):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert comfy.terminal_seen
|
||||
assert not owned.exists()
|
||||
assert sibling.read_bytes() == b"unrelated"
|
||||
assert ordinary.read_bytes() == b"ordinary"
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_omitted_temp_root_leaves_exporter_lifecycle_ownership(tmp_path):
|
||||
async def check():
|
||||
owned = exporter_file(tmp_path / "temp")
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root)
|
||||
assert owned.read_bytes() == comfy.payload
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["running", "unrelated"])
|
||||
def test_uncertain_exports_and_inputs_retained_without_interrupt(tmp_path, monkeypatch, mode):
|
||||
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
|
||||
monkeypatch.setattr(handler, "CANCEL_GRACE_SECONDS", 0.04)
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
comfy.mode = mode
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(handler.MediaExecutionUncertain):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root,
|
||||
temp_dir=temp, timeout_seconds=0.03)
|
||||
assert owned.exists()
|
||||
assert comfy.staged()
|
||||
assert not comfy.interrupts
|
||||
assert not comfy.deletes
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("component", ["root", "namespace", "request", "file"])
|
||||
def test_export_cleanup_never_traverses_symlinks(tmp_path, component):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
selected = {"root": temp, "namespace": owned.parent.parent,
|
||||
"request": owned.parent, "file": owned}[component]
|
||||
outside = tmp_path / "outside"
|
||||
selected.rename(outside)
|
||||
selected.symlink_to(outside, target_is_directory=component != "file")
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert selected.is_symlink()
|
||||
assert owned.read_bytes() == comfy.payload
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", [
|
||||
lambda r: r.update(request_id="other"),
|
||||
lambda r: r["binding"].update(collector_id="other"),
|
||||
lambda r: r["artifacts"][0]["transport"].update(filename="../escape"),
|
||||
lambda r: r["artifacts"][0]["transport"].update(subfolder="distributed-media-v1/other"),
|
||||
])
|
||||
def test_unidentified_exports_are_not_deleted(tmp_path, mutation):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
comfy.record_mutation = mutation
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert owned.exists()
|
||||
assert not comfy.views
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["required", "publish"])
|
||||
def test_publication_failures_cleanup_identified_exports(tmp_path, monkeypatch, failure):
|
||||
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
|
||||
owned = exporter_file(temp, payload=comfy.payload)
|
||||
async def publish(metadata, payload):
|
||||
assert owned.exists()
|
||||
raise ValueError("publisher failed")
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="publisher"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root,
|
||||
temp_dir=temp, publish=publish if failure == "publish" else None)
|
||||
assert not owned.exists()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_repeated_cancel_drains_collection_before_cleanup(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
|
||||
owned = exporter_file(temp, payload=comfy.payload)
|
||||
started, draining, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||||
async def publish(metadata, payload):
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
finally:
|
||||
draining.set()
|
||||
await release.wait()
|
||||
assert owned.exists(), "cleanup while publisher still reading"
|
||||
async with serve(comfy) as url:
|
||||
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url,
|
||||
input_dir=comfy.root, temp_dir=temp, publish=publish))
|
||||
await asyncio.wait_for(started.wait(), 1)
|
||||
task.cancel()
|
||||
await asyncio.wait_for(draining.wait(), 1)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert owned.exists()
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert not owned.exists()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cancel_caller", [False, True])
|
||||
def test_stubborn_publisher_teardown_is_bounded_and_retains_files(tmp_path, monkeypatch, cancel_caller):
|
||||
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
|
||||
monkeypatch.setattr(handler, "DRAIN_GRACE_SECONDS", 0.03)
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
|
||||
owned = exporter_file(temp, payload=comfy.payload)
|
||||
started, draining, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||||
async def publish(metadata, payload):
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
finally:
|
||||
draining.set()
|
||||
await release.wait()
|
||||
assert owned.exists()
|
||||
async with serve(comfy) as url:
|
||||
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url,
|
||||
input_dir=comfy.root, temp_dir=temp, publish=publish,
|
||||
timeout_seconds=5 if cancel_caller else 0.15))
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), 1)
|
||||
if cancel_caller:
|
||||
task.cancel()
|
||||
await asyncio.wait_for(draining.wait(), 1)
|
||||
if cancel_caller:
|
||||
task.cancel() # must not reset the grace or recancel the reader
|
||||
done, _ = await asyncio.wait({task}, timeout=0.5)
|
||||
assert task in done, "callback suppressed cancellation indefinitely"
|
||||
with pytest.raises(handler.MediaExecutionUncertain, match="teardown"):
|
||||
task.result()
|
||||
assert comfy.terminal_seen and owned.exists() and comfy.staged()
|
||||
assert handler._UNCERTAIN_WORK
|
||||
pending = list(handler._UNCERTAIN_WORK)
|
||||
release.set()
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
await asyncio.sleep(0)
|
||||
assert not handler._UNCERTAIN_WORK
|
||||
assert owned.exists() and comfy.staged(), "late callback must not retire files"
|
||||
finally:
|
||||
release.set()
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
run(check())
|
||||
|
||||
|
||||
def test_export_replaced_during_collection_is_preserved(tmp_path):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
view = comfy.view
|
||||
async def replace(req):
|
||||
replacement = owned.with_suffix(".replacement")
|
||||
replacement.write_bytes(b"unrelated replacement")
|
||||
replacement.replace(owned)
|
||||
return await view(req)
|
||||
comfy.view = replace
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert owned.read_bytes() == b"unrelated replacement"
|
||||
run(check())
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", ["success", "error"])
|
||||
def test_partial_history_cleans_only_valid_identified_exporters(tmp_path, status):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
unidentified = exporter_file(temp, name="c" * 32 + ".bin")
|
||||
data = request()
|
||||
second = dict(data["media"]["outputs"][0], binding_id="binding:2", node_id="export:2", collector_id="collector:2")
|
||||
data["media"]["outputs"].append(second)
|
||||
data["workflow"]["export:2"] = {"class_type": "MediaExportV1", "inputs": {
|
||||
"request_id": "request:1", "binding_json": json.dumps(second), "texts": ["source:1", 0]}}
|
||||
comfy = Comfy(tmp_path / "input", data)
|
||||
comfy.mode = status
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises((ValueError, RuntimeError)):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert not owned.exists()
|
||||
assert unidentified.exists()
|
||||
assert not comfy.views
|
||||
run(check())
|
||||
|
||||
|
||||
def test_missing_history_exports_remain_lifecycle_owned(tmp_path):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
comfy.mode = "missing"
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError, match="Missing"):
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert owned.exists()
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("replacement", ["symlink", "directory", "hardlink"])
|
||||
def test_nonregular_or_linked_export_files_are_preserved(tmp_path, replacement):
|
||||
async def check():
|
||||
temp = tmp_path / "temp"
|
||||
owned = exporter_file(temp)
|
||||
outside = tmp_path / "outside"
|
||||
outside.write_bytes(b"unrelated")
|
||||
if replacement == "hardlink":
|
||||
alias = tmp_path / "alias"
|
||||
alias.hardlink_to(owned)
|
||||
else:
|
||||
owned.unlink()
|
||||
if replacement == "symlink":
|
||||
owned.symlink_to(outside)
|
||||
else:
|
||||
owned.mkdir()
|
||||
comfy = Comfy(tmp_path / "input")
|
||||
async with serve(comfy) as url:
|
||||
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
|
||||
assert owned.exists()
|
||||
assert outside.read_bytes() == b"unrelated"
|
||||
run(check())
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shared", [False, True])
|
||||
def test_direct_exporter_input_prunes_only_obsolete_loader_ancestors(tmp_path, shared):
|
||||
async def check():
|
||||
data = request("VIDEO")
|
||||
data["workflow"]["source:1"] = {"class_type": "LoadVideo", "inputs": {"file": "unavailable.mp4"}}
|
||||
data["media"]["inputs"] = [{"node_id": "export:1", "input_name": "video", "artifact_id": "input:1"}]
|
||||
if shared:
|
||||
second = dict(data["media"]["outputs"][0], binding_id="binding:2", node_id="export:2", collector_id="collector:2")
|
||||
data["media"]["outputs"].append(second)
|
||||
data["workflow"]["export:2"] = {"class_type": "MediaExportV1", "inputs": {
|
||||
"request_id": "request:1", "binding_json": json.dumps(second), "video": ["source:1", 0]}}
|
||||
fake_data = copy.deepcopy(data)
|
||||
fake_data["media"].update(artifacts=[], inputs=[])
|
||||
comfy = Comfy(tmp_path, fake_data)
|
||||
history = comfy.history
|
||||
async def full_history(req):
|
||||
response = await history(req)
|
||||
body = json.loads(response.text)
|
||||
if shared:
|
||||
record = comfy.record()
|
||||
record["binding"] = dict(second, count=1)
|
||||
record["artifacts"][0]["id"] = "b" * 32
|
||||
record["artifacts"][0]["transport"]["filename"] = "b" * 32 + ".bin"
|
||||
record["items"] = [{"binding_id": second["binding_id"], "item_index": 0, "artifact_id": "b" * 32}]
|
||||
body["prompt:1"]["outputs"]["export:2"] = {"media_exports": [record]}
|
||||
return web.json_response(body)
|
||||
comfy.history = full_history
|
||||
async with serve(comfy) as url:
|
||||
result = await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
|
||||
contract.validate_response(result["media"], request_id="request:1", bindings=data["media"]["outputs"])
|
||||
submitted = comfy.submitted[0]["prompt"]
|
||||
imported, slot = submitted["export:1"]["inputs"]["video"]
|
||||
assert slot == 4
|
||||
assert submitted[imported]["class_type"] == "MediaImportV1"
|
||||
assert json.loads(submitted[imported]["inputs"]["artifact_json"])["kind"] == "VIDEO"
|
||||
assert ("source:1" in submitted) is shared
|
||||
assert set(submitted) == ({"export:1", imported, "source:1", "export:2"} if shared else {"export:1", imported})
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target,kind", [("request_id", "VIDEO"), ("binding_json", "VIDEO"), ("video", "STRING")])
|
||||
def test_direct_exporter_input_rejects_controls_and_wrong_kind(tmp_path, target, kind):
|
||||
async def check():
|
||||
data = request("VIDEO")
|
||||
data["media"]["inputs"] = [{"node_id": "export:1", "input_name": target, "artifact_id": "input:1"}]
|
||||
data["media"]["artifacts"][0]["kind"] = kind
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(ValueError):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
|
||||
assert not comfy.submitted
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
|
||||
|
||||
def test_slow_download_times_out_before_any_prompt(tmp_path):
|
||||
async def check():
|
||||
data = request()
|
||||
data["media"]["artifacts"][0]["transport"] = {"type": "url", "url": "https://example.org/input", "expires_at": "2099-01-01T00:00:00Z"}
|
||||
async def download(*args, **kwargs):
|
||||
await asyncio.sleep(60)
|
||||
comfy = Comfy(tmp_path, data)
|
||||
async with serve(comfy) as url:
|
||||
with pytest.raises(TimeoutError):
|
||||
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=download, timeout_seconds=0.02)
|
||||
assert not comfy.submitted
|
||||
assert not comfy.staged()
|
||||
run(check())
|
||||
@@ -0,0 +1,330 @@
|
||||
"""Real CPU-only IMAGE transfers; no ComfyUI/bootstrap/provider imports."""
|
||||
import importlib
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
# Load a private package, not Distributed's root __init__.
|
||||
ROOT = Path(__file__).resolve().parents[1] / "serverless_bridge"
|
||||
PACKAGE = "image_codec_tests_bridge"
|
||||
pkg = types.ModuleType(PACKAGE)
|
||||
pkg.__path__ = [str(ROOT)]
|
||||
sys.modules[PACKAGE] = pkg
|
||||
codec = importlib.import_module(PACKAGE + ".codec")
|
||||
images = importlib.import_module(PACKAGE + ".image_codec")
|
||||
FORMATS = ("legacy_png", "png", "jpeg", "webp", "ffv1", "h264", "h265", "av1")
|
||||
LOSSLESS = ("legacy_png", "png", "ffv1")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_gpu(monkeypatch):
|
||||
def forbidden(*a, **kw):
|
||||
pytest.fail("GPU query or initialization is forbidden")
|
||||
for name in ("is_available", "init", "_lazy_init", "device_count"):
|
||||
monkeypatch.setattr(torch.cuda, name, forbidden)
|
||||
|
||||
|
||||
def available(fmt):
|
||||
if fmt in images._VIDEO_CODECS:
|
||||
av = pytest.importorskip("av")
|
||||
name = images._VIDEO_CODECS[fmt][0]
|
||||
if name not in av.codecs_available:
|
||||
pytest.skip("Software encoder unavailable: " + name)
|
||||
|
||||
|
||||
def sample(channels=3, dtype=torch.float32, shape=(3, 65, 67)):
|
||||
g = torch.Generator(device="cpu").manual_seed(42)
|
||||
return torch.rand((*shape, channels), generator=g, device="cpu", dtype=dtype)
|
||||
|
||||
|
||||
def quantized(value):
|
||||
return torch.from_numpy(np.clip(value.float().numpy() * 255, 0, 255).astype(np.uint8)).float() / 255
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
||||
def test_exact_unchanged(dtype):
|
||||
value = sample(4, dtype, (2, 3, 5))
|
||||
value.flatten()[0] = -0.0
|
||||
expected = codec.encode_payload("IMAGE", value)
|
||||
assert images.encode_image(value) == expected
|
||||
result = images.decode_image(*expected)
|
||||
assert torch.equal(result.view(torch.uint8), value.view(torch.uint8))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", FORMATS)
|
||||
@pytest.mark.parametrize("channels", [1, 3, 4])
|
||||
def test_real_roundtrips_shape_order_quantization_alpha(fmt, channels):
|
||||
available(fmt)
|
||||
value = sample(channels)
|
||||
# Large inter-frame difference also verifies frame ordering.
|
||||
value[0] *= 0.2
|
||||
value[2] = value[2] * 0.2 + 0.8
|
||||
if channels == 4:
|
||||
value[..., 3] = torch.linspace(-0.1, 1.1, value[..., 3].numel()).reshape(value[..., 3].shape)
|
||||
descriptor, payload = images.encode_image(value, transfer_format=fmt)
|
||||
assert descriptor["encoding"] == "image-transfer-v1"
|
||||
result = codec.decode_payload("IMAGE", descriptor, payload)
|
||||
assert result.shape == value.shape
|
||||
assert result.device.type == "cpu" and result.is_contiguous()
|
||||
color = 3 if channels == 4 else channels
|
||||
expected = quantized(value[..., :color])
|
||||
if fmt in LOSSLESS:
|
||||
assert torch.equal(result[..., :color], expected)
|
||||
else:
|
||||
# Random RGB noise loses chroma detail in 4:2:0 even at high quality.
|
||||
assert (result[..., :color] - expected).abs().mean() < 0.12
|
||||
assert result[0, ..., :color].mean() < result[2, ..., :color].mean()
|
||||
if channels == 4:
|
||||
assert torch.equal(result[..., 3], value[..., 3])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["png", "ffv1", "jpeg", "webp", "h264", "h265", "av1"])
|
||||
def test_float64_alpha_is_exact(fmt):
|
||||
available(fmt)
|
||||
value = sample(4, torch.float64, (1, 65, 67))
|
||||
descriptor, payload = images.encode_image(value, transfer_format=fmt)
|
||||
out = images.decode_image(descriptor, payload)
|
||||
assert out.dtype == torch.float64
|
||||
assert torch.equal(out[..., 3].contiguous().view(torch.uint8), value[..., 3].contiguous().view(torch.uint8))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["jpeg", "webp", "h264", "h265", "av1"])
|
||||
def test_lossy_quality_changes_final_payload_size(fmt):
|
||||
available(fmt)
|
||||
value = sample(shape=(2, 96, 96))
|
||||
low = images.encode_image(value, transfer_format=fmt, transfer_quality=15)[1]
|
||||
high = images.encode_image(value, transfer_format=fmt, transfer_quality=95)[1]
|
||||
assert len(high) > len(low)
|
||||
|
||||
|
||||
def test_truncation_clipping_grayscale_noncontiguous():
|
||||
value = torch.tensor([-0.1, 0.5, 1.1, 0.12345], device="cpu").reshape(1, 2, 2, 1).transpose(1, 2)
|
||||
out = images.decode_image(*images.encode_image(value, transfer_format="png"))
|
||||
assert torch.equal(out, quantized(value))
|
||||
assert out[0, 0, 1, 0] == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [None, [], torch.ones(2, 3, 3), torch.ones(1, 2, 3, 2), torch.ones(1, 2, 3, 0), torch.ones(1, 2, 3, 3, dtype=torch.uint8), torch.ones(1, 2, 3, 3, device="meta"), torch.ones(1, 2, 3, 3).to_sparse()])
|
||||
def test_invalid_values(value):
|
||||
with pytest.raises(ValueError):
|
||||
images.encode_image(value, transfer_format="png")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [float("nan"), float("inf"), -float("inf")])
|
||||
def test_nonfinite_rejected(bad):
|
||||
with pytest.raises(ValueError, match="finite"):
|
||||
images.encode_image(torch.full((1, 1, 1, 3), bad), transfer_format="png")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("options", [{"transfer_format": "PNG"}, {"transfer_format": []}, {"transfer_quality": True}, {"transfer_quality": 0}, {"transfer_quality": 101}, {"transfer_quality": 1.5}, {"max_bytes": True}, {"max_bytes": -1}])
|
||||
def test_strict_options(options):
|
||||
with pytest.raises(ValueError):
|
||||
images.encode_image(sample(shape=(1, 1, 1)), **options)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", [{"extra": 1}, {"format": "exact"}, {"quality": True}, {"layout": "BCHW"}, {"shape": [1, 2**60, 1, 3]}, {"shape": [True, 1, 1, 3]}, {"parts": [0]}, {"parts": [True]}, {"parts": [2**60]}, {"alpha": {"extra": 0}}])
|
||||
def test_bad_descriptors(change):
|
||||
descriptor, payload = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
|
||||
descriptor.update(change)
|
||||
with pytest.raises(ValueError):
|
||||
images.decode_image(descriptor, payload)
|
||||
|
||||
|
||||
def test_truncated_appended_and_mismatched_parts():
|
||||
d, p = images.encode_image(sample(shape=(2, 3, 5)), transfer_format="png")
|
||||
for payload in (p[:-1], p + b"x"):
|
||||
with pytest.raises(ValueError):
|
||||
images.decode_image(d, payload)
|
||||
for change in ({"shape": [1, 3, 5, 3]}, {"parts": [len(p)]}, {"shape": [2, 2, 5, 3]}):
|
||||
bad = dict(d, **change)
|
||||
with pytest.raises(ValueError):
|
||||
images.decode_image(bad, p)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["png", "jpeg", "webp", "ffv1", "h264", "h265", "av1"])
|
||||
def test_payload_and_decoded_budget(fmt):
|
||||
available(fmt)
|
||||
value = sample(shape=(1, 65, 67))
|
||||
d, p = images.encode_image(value, transfer_format=fmt)
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
images.decode_image(d, p, max_bytes=100)
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
images.encode_image(value, transfer_format=fmt, max_bytes=100)
|
||||
|
||||
|
||||
def test_oversized_shape_before_materialization(monkeypatch):
|
||||
value = torch.ones(1).expand(1, 8192, 8192, 4)
|
||||
def forbidden(*a, **kw):
|
||||
pytest.fail("Materialized an oversized tensor")
|
||||
monkeypatch.setattr(torch.Tensor, "contiguous", forbidden)
|
||||
monkeypatch.setattr(torch, "isfinite", forbidden)
|
||||
with pytest.raises(ValueError):
|
||||
images.encode_image(value, transfer_format="png")
|
||||
|
||||
|
||||
def test_header_dimensions_checked_before_pixels(monkeypatch):
|
||||
d, p = images.encode_image(sample(shape=(1, 9, 11)), transfer_format="png")
|
||||
d["shape"] = [1, 1, 1, 3]
|
||||
def forbidden(*a, **kw):
|
||||
pytest.fail("Pixels decoded before header dimensions were checked")
|
||||
monkeypatch.setattr(Image.Image, "load", forbidden)
|
||||
with pytest.raises(ValueError):
|
||||
images.decode_image(d, p)
|
||||
|
||||
|
||||
def test_other_kinds_do_not_dispatch():
|
||||
d, p = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
|
||||
for kind in ("MASK", "AUDIO", "STRING"):
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_payload(kind, d, p)
|
||||
|
||||
|
||||
def test_missing_video_encoder_is_explicit(monkeypatch):
|
||||
monkeypatch.setitem(images._VIDEO_CODECS, "h264", ("missing_software_encoder", 51, "yuv420p"))
|
||||
with pytest.raises(ValueError, match="encoder|unavailable"):
|
||||
images.encode_image(sample(shape=(1, 64, 64)), transfer_format="h264")
|
||||
|
||||
|
||||
def test_bounded_sink_catches_final_flush_bytes():
|
||||
sink = images._BoundedIO(4)
|
||||
sink.write(b"1234")
|
||||
with pytest.raises(ValueError, match="limit"):
|
||||
sink.write(b"5")
|
||||
assert sink.getvalue() == b"1234"
|
||||
|
||||
|
||||
def test_crf_mapping():
|
||||
assert images._quality_to_crf(90, 51) == 13
|
||||
assert images._quality_to_crf(90, 63) == 16
|
||||
assert images._quality_to_crf(100, 63) == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", FORMATS)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
||||
def test_compressed_preserves_source_dtype_and_exact_alpha(fmt, dtype):
|
||||
available(fmt)
|
||||
value = sample(4, dtype, (1, 17, 19))
|
||||
value[..., 3].flatten()[0] = -0.0
|
||||
d, p = images.encode_image(value, transfer_format=fmt)
|
||||
with torch.device("meta"):
|
||||
result = images.decode_image(d, p)
|
||||
assert result.device.type == "cpu" and result.dtype == dtype
|
||||
assert d["dtype"] == str(dtype).removeprefix("torch.")
|
||||
assert torch.equal(result[..., 3].contiguous().view(torch.uint8), value[..., 3].contiguous().view(torch.uint8))
|
||||
if fmt in LOSSLESS:
|
||||
assert torch.equal(result[..., :3], quantized(value[..., :3]).to(dtype))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
|
||||
def test_video_count_and_header_fail_before_native_decoder(fmt, monkeypatch):
|
||||
available(fmt)
|
||||
d, p = images.encode_image(sample(shape=(2, 65, 67)), transfer_format=fmt)
|
||||
def forbidden(*a, **kw):
|
||||
pytest.fail("Native decoder reached before header validation")
|
||||
monkeypatch.setattr(images, "_decode_video", forbidden)
|
||||
for shape in ([1, 65, 67, 3], [3, 65, 67, 3], [2, 32, 32, 3]):
|
||||
bad = dict(d, shape=shape)
|
||||
with pytest.raises(ValueError):
|
||||
images.decode_image(bad, p)
|
||||
for payload in (p[:-1], p[:40], p + b"extra"):
|
||||
bad = dict(d, parts=[len(payload)])
|
||||
with pytest.raises(ValueError):
|
||||
images.decode_image(bad, payload)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
|
||||
def test_video_does_not_use_native_container_probe(fmt, monkeypatch):
|
||||
available(fmt)
|
||||
d, p = images.encode_image(sample(shape=(2, 65, 67)), transfer_format=fmt)
|
||||
av = images._import_av()
|
||||
def forbidden(*a, **kw):
|
||||
pytest.fail("Unbounded native stream probing is forbidden")
|
||||
monkeypatch.setattr(av, "open", forbidden)
|
||||
assert images.decode_image(d, p).shape == (2, 65, 67, 3)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
|
||||
def test_real_video_flush_bound(fmt):
|
||||
available(fmt)
|
||||
frames = [np.zeros((64, 64, 3), dtype=np.uint8)]
|
||||
good = images._encode_video(iter(frames), fmt, 90, 64, 64, 100000)
|
||||
with pytest.raises(ValueError, match="limit"):
|
||||
images._encode_video(iter(frames), fmt, 90, 64, 64, len(good) - 1)
|
||||
|
||||
|
||||
def test_alpha_invalid_dtype_nonfinite_and_extra_fields():
|
||||
import copy
|
||||
d, p = images.encode_image(sample(4, shape=(1, 3, 5)), transfer_format="png")
|
||||
for field, value in (("dtype", "float64"), ("shape", [1, 3, 6, 1]), ("extra", 1)):
|
||||
bad = copy.deepcopy(d)
|
||||
bad["alpha"]["descriptor"][field] = value
|
||||
with pytest.raises(ValueError):
|
||||
images.decode_image(bad, p)
|
||||
import struct
|
||||
bad = p[:-4] + struct.pack("<f", float("nan"))
|
||||
with pytest.raises(ValueError, match="finite"):
|
||||
images.decode_image(d, bad)
|
||||
|
||||
|
||||
def test_descriptor_structure_bound_and_frame_limit():
|
||||
d, p = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
|
||||
for bad in (dict(d, quality="x" * 70000), dict(d, parts=[1] * 241), dict(d, shape=[241, 1, 1, 3]), dict(d, dtype="x")):
|
||||
with pytest.raises(ValueError):
|
||||
images.decode_image(bad, p)
|
||||
with pytest.raises(ValueError):
|
||||
images.encode_image(torch.zeros(241, 1, 1, 3), transfer_format="png")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["png", "jpeg", "webp"])
|
||||
def test_all_still_headers_checked_before_pillow_open(fmt, monkeypatch):
|
||||
d, p = images.encode_image(sample(shape=(1, 31, 33)), transfer_format=fmt)
|
||||
d["shape"] = [1, 2, 2, 3]
|
||||
def forbidden(*a, **kw):
|
||||
pytest.fail("Pillow opened a mismatched header")
|
||||
monkeypatch.setattr(Image, "open", forbidden)
|
||||
with pytest.raises(ValueError, match="dimensions"):
|
||||
images.decode_image(d, p)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
|
||||
def test_video_longer_batch_order_and_coded_allocation_bound(fmt):
|
||||
available(fmt)
|
||||
value = torch.linspace(0, 1, 16).reshape(16, 1, 1, 1).expand(16, 64, 64, 3)
|
||||
d, p = images.encode_image(value, transfer_format=fmt)
|
||||
header = images._video_header(p, fmt, 64, 64, 16)
|
||||
coded = images._coded_frame_bound(p, fmt, header)
|
||||
assert 16 <= coded <= images.MAX_FRAMES * 4
|
||||
result = images.decode_image(d, p)
|
||||
assert torch.all((result.mean(dim=(1, 2, 3)) - value.mean(dim=(1, 2, 3))).abs() < 0.03)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["h264", "h265", "av1"])
|
||||
def test_packet_unit_bounds_before_native_decode(fmt):
|
||||
available(fmt)
|
||||
d, p = images.encode_image(sample(shape=(1, 64, 64)), transfer_format=fmt)
|
||||
config, bounds = images._video_header(p, fmt, 64, 64, 1)
|
||||
start, end = bounds[0]
|
||||
packet = p[start:end]
|
||||
with pytest.raises(ValueError):
|
||||
images._coded_frame_bound(packet[:-1], fmt, (config, [(0, len(packet) - 1)]))
|
||||
# An adversarial single block can contain many coded frames: bound it too.
|
||||
enormous = packet * (images.MAX_FRAMES * 4 + 1)
|
||||
with pytest.raises(ValueError, match="limit"):
|
||||
images._coded_frame_bound(enormous, fmt, (config, [(0, len(enormous))]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["h264", "h265", "av1"])
|
||||
def test_bitstream_resolution_cannot_bypass_native_pixel_limit(fmt, monkeypatch):
|
||||
available(fmt)
|
||||
d, p = images.encode_image(sample(shape=(1, 256, 256)), transfer_format=fmt)
|
||||
header = images._video_header(p, fmt, 256, 256, 1)
|
||||
av = images._import_av()
|
||||
# A lying container is not enough to disable max_pixels in the decoder.
|
||||
# The bitstream still contains 256x256 pixels but the budget permits 64x64.
|
||||
with pytest.raises((ValueError, av.error.FFmpegError)):
|
||||
list(images._decode_video(p, fmt, 64, 64, 1, header))
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Pure integration of participant identity, typed compilation and delegate pruning."""
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE = "media_transform_testpkg"
|
||||
for suffix in ("", ".api", ".api.orchestration", ".utils", ".serverless_bridge"):
|
||||
module = types.ModuleType(PACKAGE + suffix)
|
||||
module.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
|
||||
sys.modules[module.__name__] = module
|
||||
pt = importlib.import_module(PACKAGE + ".api.orchestration.prompt_transform")
|
||||
compiler = importlib.import_module(PACKAGE + ".api.orchestration.serverless_media_prompt")
|
||||
contract = importlib.import_module(PACKAGE + ".serverless_bridge.contract")
|
||||
|
||||
|
||||
def prompt():
|
||||
nodes = {
|
||||
"source": {"class_type": "MediaSource", "inputs": {}},
|
||||
"collector": {"class_type": "DistributedMediaCollector", "inputs": {
|
||||
socket: ["source", index] for index, socket in enumerate(contract.SOCKET_KINDS)}},
|
||||
}
|
||||
for index, socket in enumerate(contract.SOCKET_KINDS):
|
||||
nodes["sink:" + socket] = {"class_type": "Capture", "inputs": {"value": ["collector", index]}}
|
||||
return nodes
|
||||
|
||||
|
||||
def prepare(graph, participant, delegate):
|
||||
index = pt.PromptIndex(graph)
|
||||
jobs = pt.generate_job_id_map(index, "run")
|
||||
assert jobs == {"collector": "run_collector"}
|
||||
return pt.apply_participant_overrides(index.copy_prompt(), participant,
|
||||
["serverless:pool:1", "serverless:pool:2"], jobs, "http://master", delegate, index)
|
||||
|
||||
|
||||
def test_typed_worker_preparation_compiles_all_bindings_without_source_mutation():
|
||||
graph = prompt()
|
||||
before = copy.deepcopy(graph)
|
||||
worker = prepare(graph, "serverless:pool:2", True)
|
||||
compiled, bindings = compiler.compile_media_prompt(worker, "request:2")
|
||||
assert len(bindings) == 5
|
||||
assert {row["participant_id"] for row in bindings} == {"serverless:pool:2"}
|
||||
assert {row["job_id"] for row in bindings} == {"run_collector"}
|
||||
assert {row["socket"] for row in bindings} == set(contract.SOCKET_KINDS)
|
||||
assert set(compiled) == {"source"} | {row["node_id"] for row in bindings}
|
||||
assert graph == before
|
||||
|
||||
|
||||
def test_typed_delegate_master_keeps_outputs_without_sources_or_placeholders():
|
||||
graph = prompt()
|
||||
prepared = prepare(graph, "master", True)
|
||||
before = copy.deepcopy(prepared)
|
||||
delegated = pt.prepare_delegate_master_prompt(prepared, ["collector"])
|
||||
assert set(delegated) == {"collector"} | {"sink:" + s for s in contract.SOCKET_KINDS}
|
||||
inputs = delegated["collector"]["inputs"]
|
||||
assert not set(contract.SOCKET_KINDS) & set(inputs)
|
||||
assert json.loads(inputs["expected_media"]) == list(contract.SOCKET_KINDS)
|
||||
assert inputs["delegate_only"] is True
|
||||
assert prepared == before
|
||||
|
||||
|
||||
def test_text_source_is_not_retained_as_lightweight_delegate_scalar(monkeypatch):
|
||||
graph = {
|
||||
"text": {"class_type": "PrimitiveString", "inputs": {"value": "local text"}},
|
||||
"collector": {"class_type": "DistributedMediaCollector", "inputs": {"texts": ["text", 0]}},
|
||||
"sink": {"class_type": "Capture", "inputs": {"value": ["collector", 3]}},
|
||||
}
|
||||
monkeypatch.setattr(pt, "_delegate_master_input_is_safe_scalar", lambda *args: True)
|
||||
monkeypatch.setattr(pt, "_is_delegate_master_retained_upstream_node", lambda *args: True)
|
||||
delegated = pt.prepare_delegate_master_prompt(prepare(graph, "master", True), ["collector"])
|
||||
assert "text" not in delegated
|
||||
assert "texts" not in delegated["collector"]["inputs"]
|
||||
assert json.loads(delegated["collector"]["inputs"]["expected_media"]) == ["texts"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("socket", contract.SOCKET_KINDS)
|
||||
def test_participating_master_retains_exact_media_links(socket):
|
||||
graph = prompt()
|
||||
original = copy.deepcopy(graph)
|
||||
master = prepare(graph, "master", False)
|
||||
assert master["collector"]["inputs"][socket] == graph["collector"]["inputs"][socket]
|
||||
assert master["collector"]["inputs"]["delegate_only"] is False
|
||||
assert "worker_id" not in master["collector"]["inputs"]
|
||||
assert graph == original
|
||||
@@ -0,0 +1,531 @@
|
||||
"""Offline queue integration: real compilers, packers and wire validators.
|
||||
|
||||
Only host dispatch/probing/queueing and the manager lifecycle are mocked. No
|
||||
provider requests, services or GPU work are permitted by these tests.
|
||||
"""
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
import uuid
|
||||
import wave
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import aiohttp
|
||||
import pytest
|
||||
# Preload before the shared integration loader snapshots sys.modules: extension
|
||||
# modules cannot safely be unloaded and imported again between fixture contexts.
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from test_serverless_integration import load_modules, module, PKG
|
||||
|
||||
|
||||
PROFILE = "comfy-media-v1"
|
||||
|
||||
|
||||
def node(class_type, **inputs):
|
||||
return {"class_type": class_type, "inputs": inputs}
|
||||
|
||||
|
||||
def text_prompt():
|
||||
return {
|
||||
"text": node("PrimitiveString", value=" café雪\n"),
|
||||
"collector": node("DistributedMediaCollector", texts=["text", 0], load_balance=False),
|
||||
"save": node("TextSink", texts=["collector", 3]),
|
||||
}
|
||||
|
||||
|
||||
def legacy_prompt():
|
||||
return {
|
||||
"source": node("EmptyImage", width=32, height=32, batch_size=1, color=0),
|
||||
"collector": node("DistributedCollector", images=["source", 0]),
|
||||
"save": node("SaveImage", images=["collector", 0], filename_prefix="local"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def harness(monkeypatch):
|
||||
q, modules, ps = load_modules()
|
||||
with patch.dict(sys.modules, modules):
|
||||
pool = {"id": "pool", "name": "Media", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "offline-endpoint", "enabled": True, "jobs_per_run": 2,
|
||||
"timeout_seconds": 30, "media_profile": PROFILE}
|
||||
config = {"workers": [pool], "settings": {"master_delegate_only": True}}
|
||||
q.load_config = lambda: config
|
||||
q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate))
|
||||
q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 7, "node_errors": {}})
|
||||
manager = types.SimpleNamespace(start=AsyncMock(), cancel=AsyncMock(), ensure_capacity=Mock())
|
||||
sys.modules[PKG + ".api.orchestration.serverless_runtime"] = module("runtime", get_manager=lambda: manager)
|
||||
# These are the actual modules, not permissive serializer/packing stubs.
|
||||
client = importlib.import_module(PKG + ".api.orchestration.runpod_client")
|
||||
compiler = importlib.import_module(PKG + ".api.orchestration.serverless_media_prompt")
|
||||
packer = importlib.import_module(PKG + ".api.orchestration.serverless_typed_inputs")
|
||||
validator = importlib.import_module(PKG + ".api.orchestration.serverless_typed_results")
|
||||
handler = importlib.import_module(PKG + ".serverless_bridge.handler")
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError("Offline queue tests must not access a network or GPU")
|
||||
monkeypatch.setattr(aiohttp, "ClientSession", forbidden)
|
||||
monkeypatch.setattr(socket, "create_connection", forbidden)
|
||||
monkeypatch.setattr(socket.socket, "connect", forbidden)
|
||||
monkeypatch.setattr(socket.socket, "connect_ex", forbidden)
|
||||
monkeypatch.setattr(torch.cuda, "init", forbidden)
|
||||
monkeypatch.setattr(torch.cuda, "current_device", forbidden)
|
||||
h = types.SimpleNamespace(q=q, ps=ps, pool=pool, config=config, manager=manager,
|
||||
client=client, compiler=compiler, packer=packer,
|
||||
validator=validator, handler=handler)
|
||||
yield h
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def media_paths(tmp_path, monkeypatch):
|
||||
paths = types.SimpleNamespace(get_annotated_filepath=lambda name: str(tmp_path / name))
|
||||
for kind in ("input", "output", "temp"):
|
||||
setattr(paths, "get_" + kind + "_directory", lambda: str(tmp_path))
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", paths)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def run(h, prompt=None, **kwargs):
|
||||
return asyncio.run(h.q.orchestrate_distributed_execution(
|
||||
prompt if prompt is not None else text_prompt(), None, "client", **kwargs))
|
||||
|
||||
|
||||
def assert_no_dispatch(h, *, no_probe=False):
|
||||
h.q.queue_prompt_payload.assert_not_awaited()
|
||||
h.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
h.manager.start.assert_not_awaited()
|
||||
h.manager.cancel.assert_not_awaited()
|
||||
assert h.ps.distributed_pending_jobs == {}
|
||||
if no_probe:
|
||||
h.q.select_active_workers.assert_not_awaited()
|
||||
h.q.select_least_busy_worker.assert_not_awaited()
|
||||
h.q.fetch_worker_has_node.assert_not_awaited()
|
||||
h.q.sync_worker_media.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("delegate", [True, False])
|
||||
def test_real_compile_pack_serialize_and_canonical_bindings(harness, media_paths, delegate):
|
||||
h = harness
|
||||
image_path = media_paths / "picture.png"
|
||||
Image.new("RGBA", (2, 3), color=(17, 29, 41, 53)).save(image_path)
|
||||
image_bytes = image_path.read_bytes()
|
||||
prompt = text_prompt()
|
||||
prompt["image"] = node("LoadImage", image="picture.png")
|
||||
prompt["collector"]["inputs"].update(images=["image", 0], masks=["image", 1])
|
||||
prompt["collector2"] = node("DistributedMediaCollector", texts=["text", 0])
|
||||
prompt["unused"] = node("PreviewImage", images=["image", 0])
|
||||
source = copy.deepcopy(prompt)
|
||||
with patch.object(h.compiler, "compile_media_prompt", wraps=h.compiler.compile_media_prompt) as compile_call, \
|
||||
patch.object(h.packer, "prepare_media_inputs", wraps=h.packer.prepare_media_inputs) as pack_call, \
|
||||
patch.object(h.client, "serialize_request", wraps=h.client.serialize_request) as serialize_call, \
|
||||
patch.object(h.q, "prune_prompt_for_worker", side_effect=AssertionError("typed graph must bypass legacy pruning")):
|
||||
result = run(h, prompt, delegate_master=delegate)
|
||||
assert result == ("master-prompt", 7, 2, {})
|
||||
assert compile_call.call_count == pack_call.await_count == 2
|
||||
assert serialize_call.call_count >= 2
|
||||
assert prompt == source
|
||||
h.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
jobs = h.manager.start.call_args.args[2]
|
||||
assert [j["worker"]["id"] for j in jobs] == ["serverless:pool:1", "serverless:pool:2"]
|
||||
assert len({j["media"]["request_id"] for j in jobs}) == 2
|
||||
master = h.q.queue_prompt_payload.call_args.args[0]
|
||||
for job in jobs:
|
||||
media = job["media"]
|
||||
assert str(uuid.UUID(media["request_id"])) == media["request_id"]
|
||||
assert uuid.UUID(media["request_id"]).version == 4
|
||||
assert media["profile"] == PROFILE and media["version"] == 1
|
||||
assert media["outputs"] == job["manifest"]
|
||||
assert media["outputs"] is not job["manifest"]
|
||||
assert media["artifacts"] == media["inputs"] == []
|
||||
assert {b["socket"] for b in job["manifest"]} == {"texts", "images", "masks"}
|
||||
assert len(job["manifest"]) == 4
|
||||
assert base64.b64decode(job["images"][0]["image"]) == image_bytes
|
||||
assert job["workflow"]["image"]["inputs"]["image"] == job["images"][0]["name"]
|
||||
assert not ({"save", "unused", "collector", "collector2"} & set(job["workflow"]))
|
||||
body = h.client.serialize_request(job["workflow"], job["images"], 30, media=media)
|
||||
wire = json.loads(body)
|
||||
assert wire["input"] == {"workflow": job["workflow"], "images": job["images"], "media": media}
|
||||
assert wire["policy"] == {"executionTimeout": 30000, "ttl": 30000}
|
||||
_, _, _, targets = h.handler._validate_request(wire["input"])
|
||||
assert targets == [b["node_id"] for b in job["manifest"]]
|
||||
assert h.validator.validate_prepared_job(job) is True
|
||||
for binding in job["manifest"]:
|
||||
assert binding["participant_id"] == job["worker"]["id"]
|
||||
assert binding["job_id"] == master[binding["collector_id"]]["inputs"]["multi_job_id"]
|
||||
controls = job["workflow"][binding["node_id"]]["inputs"]
|
||||
assert json.loads(controls["binding_json"]) == binding
|
||||
assert controls["request_id"] == media["request_id"]
|
||||
assert master["collector"]["inputs"]["delegate_only"] is delegate
|
||||
assert json.loads(master["collector"]["inputs"]["expected_media"]) == ["images", "masks", "texts"]
|
||||
if delegate:
|
||||
assert "text" not in master and "image" not in master
|
||||
assert not any(name in master["collector"]["inputs"] for name in ("images", "audio", "masks", "texts", "video"))
|
||||
assert not any(n["class_type"] == "DistributedEmptyImage" for n in master.values())
|
||||
else:
|
||||
assert master["collector"]["inputs"]["texts"] == ["text", 0]
|
||||
assert "text" in master and "image" in master
|
||||
|
||||
|
||||
def test_delegate_disconnects_every_media_socket(harness):
|
||||
prompt = text_prompt()
|
||||
for index, name in enumerate(("images", "audio", "masks", "video")):
|
||||
source_id = "source-" + name
|
||||
prompt[source_id] = node("GeneratedMedia", seed=index)
|
||||
prompt["collector"]["inputs"][name] = [source_id, 0]
|
||||
original = copy.deepcopy(prompt)
|
||||
run(harness, prompt)
|
||||
master = harness.q.queue_prompt_payload.call_args.args[0]
|
||||
inputs = master["collector"]["inputs"]
|
||||
sockets = {"images", "audio", "masks", "texts", "video"}
|
||||
assert set(json.loads(inputs["expected_media"])) == sockets
|
||||
assert not sockets.intersection(inputs)
|
||||
assert set(master) == {"collector", "save"}
|
||||
for job in harness.manager.start.call_args.args[2]:
|
||||
assert {b["socket"] for b in job["manifest"]} == sockets
|
||||
assert {"text", "source-images", "source-audio", "source-masks", "source-video"} <= set(job["workflow"])
|
||||
assert prompt == original
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selected", ["master", "worker"])
|
||||
def test_typed_load_balance_selects_master_or_one_serverless_slot(harness, selected):
|
||||
h = harness
|
||||
prompt = text_prompt()
|
||||
prompt["collector"]["inputs"]["load_balance"] = "true"
|
||||
h.q.select_least_busy_worker.side_effect = lambda candidates, **kw: next(
|
||||
w for w in candidates if (w["id"] == "master") == (selected == "master"))
|
||||
result = run(h, prompt, delegate_master=False)
|
||||
h.q.select_least_busy_worker.assert_awaited_once()
|
||||
master = h.q.queue_prompt_payload.call_args.args[0]
|
||||
if selected == "master":
|
||||
assert result[2] == 0
|
||||
h.manager.start.assert_not_awaited()
|
||||
assert master["collector"]["inputs"]["delegate_only"] is False
|
||||
assert "text" in master
|
||||
else:
|
||||
assert result[2] == 1
|
||||
assert len(h.manager.start.call_args.args[2]) == 1
|
||||
assert master["collector"]["inputs"]["delegate_only"] is True
|
||||
assert "text" not in master
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", [None, "legacy-image-v1", PROFILE])
|
||||
def test_explicit_profile_compatibility_rejected_before_probe(harness, profile):
|
||||
h = harness
|
||||
prompt = legacy_prompt() if profile == PROFILE else text_prompt()
|
||||
if profile is None:
|
||||
h.pool.pop("media_profile")
|
||||
else:
|
||||
h.pool["media_profile"] = profile
|
||||
with pytest.raises(ValueError, match="requires|require"):
|
||||
run(h, prompt)
|
||||
assert_no_dispatch(h, no_probe=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", [None, "auto", "comfy-media-v2", False])
|
||||
def test_invalid_explicit_profiles_fail_closed(harness, profile):
|
||||
harness.pool["media_profile"] = profile
|
||||
with pytest.raises(ValueError, match="media_profile"):
|
||||
run(harness)
|
||||
assert_no_dispatch(harness, no_probe=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("direct_type", ["local", "remote", "cloud"])
|
||||
@pytest.mark.parametrize("explicit", [False, True])
|
||||
def test_direct_workers_never_silently_skipped_or_fallback(harness, direct_type, explicit):
|
||||
h = harness
|
||||
h.config["workers"].append({"id": "direct", "type": direct_type, "host": "offline.invalid",
|
||||
"enabled": not explicit, "port": 8188})
|
||||
prompt = text_prompt()
|
||||
prompt["collector"]["inputs"]["load_balance"] = True
|
||||
kwargs = {"enabled_worker_ids": ["pool", "direct"]} if explicit else {}
|
||||
with pytest.raises(ValueError, match="direct workers"):
|
||||
run(h, prompt, delegate_master=False, **kwargs)
|
||||
assert_no_dispatch(h, no_probe=True)
|
||||
|
||||
|
||||
def test_disabled_or_unselected_incompatible_workers_do_not_block(harness):
|
||||
h = harness
|
||||
h.config["workers"] += [
|
||||
{"id": "direct", "type": "local", "enabled": False},
|
||||
dict(h.pool, id="legacy", media_profile="legacy-image-v1", enabled=False),
|
||||
]
|
||||
assert run(h)[2] == 2
|
||||
h.pool["enabled"] = False
|
||||
assert run(h, enabled_worker_ids=["pool"])[2] == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("other", ["DistributedCollector", "UltimateSDUpscaleDistributed", "DistributedImageListStart"])
|
||||
@pytest.mark.parametrize("workers", [True, False])
|
||||
def test_mixed_collectors_or_image_list_rejected_even_master_only(harness, other, workers):
|
||||
h = harness
|
||||
if not workers:
|
||||
h.config["workers"] = []
|
||||
prompt = text_prompt()
|
||||
prompt["incompatible"] = node(other)
|
||||
with pytest.raises(ValueError, match="cannot be mixed"):
|
||||
run(h, prompt, delegate_master=False)
|
||||
assert_no_dispatch(h, no_probe=True)
|
||||
|
||||
|
||||
def test_mixed_pool_profiles_rejected_before_load_balancing(harness):
|
||||
h = harness
|
||||
h.config["workers"].append(dict(h.pool, id="legacy", media_profile="legacy-image-v1"))
|
||||
prompt = text_prompt()
|
||||
prompt["collector"]["inputs"]["load_balance"] = True
|
||||
with pytest.raises(ValueError, match="requires"):
|
||||
run(h, prompt, delegate_master=False)
|
||||
assert_no_dispatch(h, no_probe=True)
|
||||
|
||||
|
||||
def test_typed_master_only_is_supported_without_serverless_work(harness):
|
||||
h = harness
|
||||
h.config["workers"] = []
|
||||
assert run(h, delegate_master=False)[2] == 0
|
||||
master = h.q.queue_prompt_payload.call_args.args[0]
|
||||
assert master["collector"]["inputs"]["enabled_worker_ids"] == "[]"
|
||||
assert master["collector"]["inputs"]["texts"] == ["text", 0]
|
||||
h.manager.start.assert_not_awaited()
|
||||
h.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("offline", [False, True])
|
||||
def test_typed_delegate_requires_available_worker_without_master_fallback(harness, offline):
|
||||
h = harness
|
||||
if offline:
|
||||
h.q.select_active_workers.side_effect = lambda *a, **kw: ([], False)
|
||||
else:
|
||||
h.config["workers"] = []
|
||||
with pytest.raises(RuntimeError, match="comfy-media-v1"):
|
||||
run(h)
|
||||
assert_no_dispatch(h)
|
||||
|
||||
|
||||
def test_unknown_explicit_typed_worker_is_rejected(harness):
|
||||
with pytest.raises(ValueError, match="Unknown explicitly selected"):
|
||||
run(harness, enabled_worker_ids=["unknown"], delegate_master=False)
|
||||
assert_no_dispatch(harness, no_probe=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [" pool ", "pool white", "pool\n"])
|
||||
def test_pool_media_identity_is_not_repaired(harness, bad):
|
||||
harness.pool["id"] = bad
|
||||
with pytest.raises(ValueError, match="identifier"):
|
||||
run(harness)
|
||||
assert_no_dispatch(harness, no_probe=True)
|
||||
|
||||
|
||||
def test_integer_source_node_identity_is_not_normalized(harness):
|
||||
prompt = text_prompt()
|
||||
prompt[1] = prompt.pop("text")
|
||||
prompt["collector"]["inputs"]["texts"] = [1, 0]
|
||||
with pytest.raises(ValueError, match="identifier"):
|
||||
run(harness, prompt)
|
||||
assert_no_dispatch(harness, no_probe=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("broken", ["dangling-unused", "literal", "missing-file", "bad-loader-controls", "oversize"])
|
||||
def test_compilation_packing_and_size_errors_precede_all_dispatch(harness, media_paths, broken):
|
||||
h = harness
|
||||
prompt = text_prompt()
|
||||
if broken == "dangling-unused":
|
||||
prompt["unused"] = node("PreviewImage", images=["missing", 0])
|
||||
elif broken == "literal":
|
||||
prompt["collector"]["inputs"]["texts"] = "must use a linked string node"
|
||||
elif broken in ("missing-file", "bad-loader-controls"):
|
||||
prompt["audio"] = node("LoadAudio", audio="missing.wav")
|
||||
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
|
||||
if broken == "bad-loader-controls":
|
||||
prompt["audio"]["inputs"]["duration"] = 5
|
||||
else:
|
||||
prompt["text"]["inputs"]["value"] = "雪" * 1000
|
||||
original = copy.deepcopy(prompt)
|
||||
with patch.object(h.client, "MAX_REQUEST_BYTES", 100 if broken == "oversize" else h.client.MAX_REQUEST_BYTES):
|
||||
with pytest.raises((ValueError, OSError)):
|
||||
run(h, prompt)
|
||||
assert prompt == original
|
||||
assert_no_dispatch(h)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("corruption", ["manifest", "participant", "request-control", "binding-control", "profile", "input-ref"])
|
||||
def test_whole_batch_identity_preflight_before_master_and_manager(harness, corruption):
|
||||
h = harness
|
||||
real_prepare = h.q._prepare_worker_payload
|
||||
async def corrupt_later(*args, **kwargs):
|
||||
worker, prepared = await real_prepare(*args, **kwargs)
|
||||
if worker["id"].endswith(":2"):
|
||||
binding = prepared["manifest"][0]
|
||||
if corruption == "manifest":
|
||||
binding["job_id"] = "wrong-job"
|
||||
elif corruption == "participant":
|
||||
prepared["worker"] = dict(worker, id="different-participant")
|
||||
elif corruption == "request-control":
|
||||
prepared["workflow"][binding["node_id"]]["inputs"]["request_id"] = "different-request"
|
||||
elif corruption == "binding-control":
|
||||
prepared["workflow"][binding["node_id"]]["inputs"]["binding_json"] = "{}"
|
||||
elif corruption == "profile":
|
||||
prepared["worker"] = dict(worker, media_profile="legacy-image-v1")
|
||||
else:
|
||||
prepared["media"]["inputs"] = [{"node_id": "missing", "input_name": "texts", "artifact_id": "unknown"}]
|
||||
return worker, prepared
|
||||
with patch.object(h.q, "_prepare_worker_payload", side_effect=corrupt_later):
|
||||
with pytest.raises(ValueError):
|
||||
run(h)
|
||||
h.manager.ensure_capacity.assert_not_called()
|
||||
assert_no_dispatch(h)
|
||||
|
||||
|
||||
def test_batch_preflight_also_precedes_regular_direct_dispatch(harness):
|
||||
h = harness
|
||||
h.pool.pop("media_profile")
|
||||
h.config["workers"].append({"id": "direct", "type": "local", "enabled": True})
|
||||
real_prepare = h.q._prepare_worker_payload
|
||||
async def corrupt_later(*args, **kwargs):
|
||||
worker, prepared = await real_prepare(*args, **kwargs)
|
||||
if worker["id"] == "serverless:pool:2":
|
||||
prepared["worker"] = dict(worker, media_profile=PROFILE)
|
||||
return worker, prepared
|
||||
with patch.object(h.q, "_prepare_worker_payload", side_effect=corrupt_later):
|
||||
with pytest.raises(ValueError, match="mismatch"):
|
||||
run(h, legacy_prompt())
|
||||
assert_no_dispatch(h)
|
||||
|
||||
|
||||
def test_capacity_preflight_precedes_collector_state_and_master(harness):
|
||||
h = harness
|
||||
h.manager.ensure_capacity.side_effect = ValueError("capacity exhausted")
|
||||
with pytest.raises(ValueError, match="capacity"):
|
||||
run(h)
|
||||
assert_no_dispatch(h)
|
||||
|
||||
|
||||
def test_master_queue_failure_cleans_queues_and_never_starts_manager(harness):
|
||||
h = harness
|
||||
h.q.queue_prompt_payload.side_effect = ValueError("invalid master graph")
|
||||
with pytest.raises(ValueError, match="invalid master"):
|
||||
run(h)
|
||||
assert h.ps.distributed_pending_jobs == {}
|
||||
h.manager.start.assert_not_awaited()
|
||||
h.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
|
||||
|
||||
def test_manager_setup_failure_emits_one_terminal_per_collector_participant(harness):
|
||||
h = harness
|
||||
prompt = text_prompt()
|
||||
prompt["image"] = node("EmptyImage", width=2, height=2, batch_size=1, color=0)
|
||||
prompt["collector"]["inputs"]["images"] = ["image", 0]
|
||||
h.manager.start.side_effect = ValueError("journal failure")
|
||||
with pytest.raises(ValueError, match="journal failure"):
|
||||
run(h, prompt)
|
||||
h.manager.cancel.assert_awaited_once()
|
||||
queued = next(iter(h.ps.distributed_pending_jobs.values()))
|
||||
assert queued.qsize() == 2 # two participants, not four per-socket completions
|
||||
messages = [queued.get_nowait(), queued.get_nowait()]
|
||||
assert {m["worker_id"] for m in messages} == {"serverless:pool:1", "serverless:pool:2"}
|
||||
assert all(m["is_last"] is True and m["serverless_error"] for m in messages)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", [None, "legacy-image-v1"])
|
||||
def test_legacy_envelope_and_compiler_are_unchanged(harness, profile):
|
||||
h = harness
|
||||
if profile is None:
|
||||
h.pool.pop("media_profile")
|
||||
else:
|
||||
h.pool["media_profile"] = profile
|
||||
prompt = legacy_prompt()
|
||||
original = copy.deepcopy(prompt)
|
||||
with patch.object(h.compiler, "compile_media_prompt", side_effect=AssertionError("legacy must not use typed compiler")), \
|
||||
patch.object(h.packer, "prepare_media_inputs", side_effect=AssertionError("legacy must not use typed packer")):
|
||||
assert run(h, prompt)[2] == 2
|
||||
for job in h.manager.start.call_args.args[2]:
|
||||
assert "media" not in job
|
||||
assert any(n["class_type"] == "SaveImage" for n in job["workflow"].values())
|
||||
expected = {"input": {"workflow": job["workflow"], "images": job["images"]},
|
||||
"policy": {"executionTimeout": 30000, "ttl": 30000}}
|
||||
assert h.client.serialize_request(job["workflow"], job["images"], 30) == json.dumps(
|
||||
expected, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode()
|
||||
assert prompt == original
|
||||
|
||||
|
||||
def test_legacy_no_worker_delegate_fallback_is_preserved(harness):
|
||||
h = harness
|
||||
h.config["workers"] = []
|
||||
assert run(h, legacy_prompt())[2] == 0
|
||||
assert h.q.queue_prompt_payload.call_args.args[0]["collector"]["inputs"]["delegate_only"] is False
|
||||
|
||||
|
||||
def test_native_audio_compile_pack_serialized_input_cpu(harness, media_paths):
|
||||
"""Exercise real PyAV + core PCM conversion, not a mocked packing result."""
|
||||
if os.environ.get("_MEDIA_QUEUE_CORE_TEST") != "1":
|
||||
core = os.environ.get("COMFYUI_SOURCE_ROOT")
|
||||
if not core:
|
||||
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU native audio queue integration")
|
||||
env = dict(os.environ)
|
||||
env.pop("PYTHONPATH", None)
|
||||
env.pop("COMFYUI_SOURCE_ROOT", None)
|
||||
scratch = Path.home() / ".hermes/cache/scratch"
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
env.update(_MEDIA_QUEUE_CORE_TEST="1", CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="",
|
||||
PYTHONDONTWRITEBYTECODE="1", TMPDIR=str(scratch))
|
||||
bootstrap = """
|
||||
import sys
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
import torch
|
||||
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('CPU audio integration may not initialize a GPU')
|
||||
torch.cuda.is_available = lambda: False
|
||||
torch.cuda.device_count = lambda: 0
|
||||
torch.cuda.init = forbidden
|
||||
torch.cuda.current_device = forbidden
|
||||
torch.cuda._lazy_init = forbidden
|
||||
import pytest
|
||||
raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider', '-k',
|
||||
'test_native_audio_compile_pack_serialized_input_cpu']))
|
||||
"""
|
||||
result = subprocess.run([sys.executable, "-c", bootstrap, core, str(Path(__file__).resolve())],
|
||||
env=env, capture_output=True, text=True, timeout=120)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert "1 passed" in result.stdout
|
||||
return
|
||||
h = harness
|
||||
h.pool["jobs_per_run"] = 1
|
||||
path = media_paths / "tone.wav"
|
||||
pcm = torch.tensor([-100, 0, 100, 32767, -32768, 7], dtype=torch.int16)
|
||||
with wave.open(str(path), "wb") as stream:
|
||||
stream.setnchannels(1)
|
||||
stream.setsampwidth(2)
|
||||
stream.setframerate(8000)
|
||||
stream.writeframes(pcm.numpy().tobytes())
|
||||
original_bytes = path.read_bytes()
|
||||
prompt = text_prompt()
|
||||
prompt["audio"] = node("LoadAudio", audio="tone.wav")
|
||||
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
|
||||
original = copy.deepcopy(prompt)
|
||||
run(h, prompt)
|
||||
job = h.manager.start.call_args.args[2][0]
|
||||
wire = json.loads(h.client.serialize_request(job["workflow"], job["images"], 30, media=job["media"]))
|
||||
artifacts = wire["input"]["media"]["artifacts"]
|
||||
refs = wire["input"]["media"]["inputs"]
|
||||
assert len(artifacts) == len(refs) == 1
|
||||
binding = next(b for b in job["manifest"] if b["socket"] == "audio")
|
||||
assert refs == [{"node_id": binding["node_id"], "input_name": "audio", "artifact_id": artifacts[0]["id"]}]
|
||||
contract = importlib.import_module(PKG + ".serverless_bridge.contract")
|
||||
codec = importlib.import_module(PKG + ".serverless_bridge.codec")
|
||||
payload = asyncio.run(contract.resolve_artifact(artifacts[0]))
|
||||
value = codec.decode_payload("AUDIO", artifacts[0]["descriptor"], payload)
|
||||
assert value["sample_rate"] == 8000
|
||||
assert value["waveform"].shape == (1, 1, 6)
|
||||
assert torch.equal(value["waveform"], (pcm.to(torch.float32) / 32768).reshape(1, 1, -1))
|
||||
assert path.read_bytes() == original_bytes
|
||||
assert prompt == original
|
||||
assert h.validator.validate_prepared_job(job) is True
|
||||
@@ -0,0 +1,76 @@
|
||||
"""Collector controls -> compiler -> endpoint validation, no network or GPU."""
|
||||
import copy
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from test_serverless_media_prompt import compiler, _prompt
|
||||
from test_media_handler import handler
|
||||
from test_media_collector import module, collect, native_media
|
||||
|
||||
FORMATS = ('exact', 'legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1')
|
||||
|
||||
|
||||
def compile_request(fmt='exact', quality=90):
|
||||
prompt = _prompt()
|
||||
prompt['source:1'] = {'class_type': 'Source', 'inputs': {}}
|
||||
prompt['collector:1']['inputs'].update(transfer_format=fmt, transfer_quality=quality,
|
||||
masks=['source:1', 1], texts=['source:1', 2])
|
||||
original = copy.deepcopy(prompt)
|
||||
workflow, outputs = compiler.compile_media_prompt(prompt, 'request:1')
|
||||
assert prompt == original
|
||||
return {'workflow': workflow, 'images': [], 'media': {
|
||||
'version': 1, 'profile': 'comfy-media-v1', 'request_id': 'request:1',
|
||||
'artifacts': [], 'inputs': [], 'outputs': outputs}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize('fmt', FORMATS)
|
||||
def test_compiler_and_handler_preserve_image_policy_only(fmt):
|
||||
data = compile_request(fmt, 73)
|
||||
for binding in data['media']['outputs']:
|
||||
inputs = data['workflow'][binding['node_id']]['inputs']
|
||||
if binding['kind'] == 'IMAGE' and fmt != 'exact':
|
||||
assert inputs['transfer_format'] == fmt
|
||||
assert inputs['transfer_quality'] == 73
|
||||
else:
|
||||
assert 'transfer_format' not in inputs and 'transfer_quality' not in inputs
|
||||
handler._validate_request(data)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('fmt,quality', [('JPEG',90),('',90),('jpeg',True),('jpeg',0),
|
||||
('jpeg',101),('jpeg','90'),(['source:1',0],90)])
|
||||
def test_compiler_rejects_invalid_or_linked_controls(fmt, quality):
|
||||
with pytest.raises(ValueError):
|
||||
compile_request(fmt, quality)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('field,value', [('transfer_format','unknown'),('transfer_quality',0),
|
||||
('transfer_quality',True),('transfer_quality','90')])
|
||||
def test_handler_rejects_tampered_policy(field, value):
|
||||
data = compile_request('jpeg')
|
||||
binding = next(b for b in data['media']['outputs'] if b['kind'] == 'IMAGE')
|
||||
data['workflow'][binding['node_id']]['inputs'][field] = value
|
||||
with pytest.raises(ValueError):
|
||||
handler._validate_request(data)
|
||||
|
||||
|
||||
def test_handler_rejects_image_controls_on_mask_export():
|
||||
data = compile_request()
|
||||
binding = next(b for b in data['media']['outputs'] if b['kind'] == 'MASK')
|
||||
data['workflow'][binding['node_id']]['inputs']['transfer_format'] = 'jpeg'
|
||||
with pytest.raises(ValueError):
|
||||
handler._validate_request(data)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('fmt', FORMATS)
|
||||
def test_master_passthrough_stays_native_for_every_format(module, fmt):
|
||||
media = native_media()
|
||||
actual = collect(module, **media, transfer_format=[fmt], transfer_quality=[73])
|
||||
for values, original in zip(actual, media.values()):
|
||||
assert values is original
|
||||
|
||||
|
||||
@pytest.mark.parametrize('quality', [True, 0, 101, '90', [90,91]])
|
||||
def test_node_rejects_bad_quality_even_without_remote_participants(module, quality):
|
||||
with pytest.raises(ValueError):
|
||||
collect(module, texts=['hello'], transfer_quality=quality)
|
||||
@@ -0,0 +1,562 @@
|
||||
"""Native ComfyUI VIDEO tests, isolated from the custom-node import graph.
|
||||
|
||||
The outer pytest process starts one bounded CPU-only subprocess. Inside it the
|
||||
real pinned core is imported *after* cli_args.args.cpu=True; no fake VideoInput
|
||||
classes, media transcodes, providers, server, or GPU queries are needed.
|
||||
"""
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
if os.environ.get("_DISTRIBUTED_VIDEO_CODEC_CORE_TEST") != "1":
|
||||
def test_native_video_codec_in_cpu_subprocess():
|
||||
configured = os.environ.get("COMFYUI_SOURCE_ROOT", os.environ.get("COMFYUI_VIDEO_TEST_ROOT"))
|
||||
if configured is None:
|
||||
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU native VIDEO acceptance")
|
||||
core = Path(configured).expanduser().resolve()
|
||||
assert configured and (core / "comfy_api/latest/_input_impl/video_types.py").is_file(), (
|
||||
"Configured COMFYUI_SOURCE_ROOT (or COMFYUI_VIDEO_TEST_ROOT) must provide native video classes"
|
||||
)
|
||||
env = dict(os.environ)
|
||||
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
env["TMPDIR"] = str(scratch)
|
||||
for name in ("PYTHONPATH", "COMFYUI_SOURCE_ROOT"):
|
||||
env.pop(name, None)
|
||||
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="",
|
||||
_DISTRIBUTED_VIDEO_CODEC_CORE_TEST="1", PYTHONDONTWRITEBYTECODE="1")
|
||||
bootstrap = """
|
||||
import sys
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
import torch
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('VIDEO tests must not query or initialize a GPU')
|
||||
# Third-party core imports probe availability even in CPU mode: answer without
|
||||
# executing a GPU query, then forbid even these probes during codec execution.
|
||||
torch.cuda.is_available = lambda: False
|
||||
torch.cuda.device_count = lambda: 0
|
||||
for name in ('init', '_lazy_init', 'current_device'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
import comfy_api.latest._input_impl.video_types
|
||||
for name in ('is_available', 'device_count'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
import pytest
|
||||
raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider']))
|
||||
"""
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", bootstrap, str(core), str(Path(__file__).resolve())],
|
||||
env=env, text=True, capture_output=True, timeout=150,
|
||||
cwd=Path(__file__).resolve().parents[1],
|
||||
)
|
||||
print(result.stdout)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
else:
|
||||
import copy
|
||||
from fractions import Fraction
|
||||
import importlib
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
import types
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
from comfy_api.latest._input import VideoInput
|
||||
from comfy_api.latest._input_impl.video_types import (
|
||||
VideoFromComponents, VideoFromFile, VideoFromList,
|
||||
)
|
||||
from comfy_api.latest._util import VideoComponents
|
||||
|
||||
# Only this sibling package is synthetic; all native media types are real.
|
||||
package_name = "_distributed_video_codec_test"
|
||||
package = types.ModuleType(package_name)
|
||||
package.__path__ = [str(Path(__file__).resolve().parents[1] / "serverless_bridge")]
|
||||
sys.modules[package_name] = package
|
||||
codec = importlib.import_module(package_name + ".video_codec")
|
||||
|
||||
def audio(dtype=torch.float64):
|
||||
return {"waveform": torch.tensor([[[0.123456789, -0.0, -1.5, 2.0]]], dtype=dtype, device="cpu"),
|
||||
"sample_rate": 48000}
|
||||
|
||||
def components(dtype=torch.float64, *, metadata=None, alpha=True, sound=True):
|
||||
image = torch.arange(36, device="cpu", dtype=torch.float64).reshape(2, 2, 3, 3) / 997
|
||||
image[0, 0, 0, 0] = -0.0
|
||||
return VideoFromComponents(VideoComponents(
|
||||
images=image.to(dtype), frame_rate=Fraction(30000, 1001),
|
||||
audio=audio(dtype) if sound else None,
|
||||
alpha=torch.full((2, 2, 3), 0.123456789, dtype=dtype, device="cpu") if alpha else None,
|
||||
metadata=metadata,
|
||||
), bit_depth=10, color_space="HDR")
|
||||
|
||||
def exact_tensor(actual, expected):
|
||||
assert actual.device.type == "cpu"
|
||||
assert actual.dtype == expected.dtype
|
||||
assert actual.shape == expected.shape
|
||||
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
|
||||
|
||||
def file_value(data=b"opaque original media bytes", **kwargs):
|
||||
return VideoFromFile(io.BytesIO(data), **kwargs)
|
||||
|
||||
def forbid(*args, **kwargs):
|
||||
raise AssertionError("Unexpected materialization/transcoding/allocation")
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
||||
@pytest.mark.parametrize("alpha,sound", [(False, False), (True, False), (False, True), (True, True)])
|
||||
def test_components_exact_roundtrip(dtype, alpha, sound, monkeypatch):
|
||||
metadata = {"title": "雪 🎞", "flags": [True, None, -4, 0.125], "nested": {"empty": []}}
|
||||
source = components(dtype, metadata=metadata, alpha=alpha, sound=sound)
|
||||
expected = source.get_components()
|
||||
monkeypatch.setattr(VideoFromComponents, "get_components", forbid)
|
||||
monkeypatch.setattr(VideoFromComponents, "save_to", forbid)
|
||||
descriptor, payload = codec.encode_video(source)
|
||||
assert descriptor["encoding"] == "native-video-v1"
|
||||
assert descriptor["representation"] == "components"
|
||||
assert descriptor["frame_rate"] == {"numerator": 30000, "denominator": 1001}
|
||||
# This is the actual JSON wire descriptor, not a Python-only shortcut.
|
||||
descriptor = json.loads(json.dumps(descriptor, allow_nan=False))
|
||||
result = codec.decode_video(descriptor, payload)
|
||||
assert type(result) is VideoFromComponents
|
||||
assert isinstance(result, VideoInput)
|
||||
actual = vars(result)["_VideoFromComponents__components"]
|
||||
exact_tensor(actual.images, expected.images)
|
||||
assert actual.frame_rate == expected.frame_rate
|
||||
assert type(actual.frame_rate) is Fraction
|
||||
assert actual.metadata == metadata
|
||||
assert actual.metadata is not metadata
|
||||
assert vars(result)["_VideoFromComponents__bit_depth"] == 10
|
||||
assert vars(result)["_VideoFromComponents__color_space"] == "HDR"
|
||||
if alpha:
|
||||
exact_tensor(actual.alpha, expected.alpha)
|
||||
else:
|
||||
assert actual.alpha is None
|
||||
if sound:
|
||||
exact_tensor(actual.audio["waveform"], expected.audio["waveform"])
|
||||
assert actual.audio["sample_rate"] == expected.audio["sample_rate"]
|
||||
else:
|
||||
assert actual.audio is None
|
||||
actual.images.fill_(99)
|
||||
assert not torch.equal(actual.images, expected.images)
|
||||
actual.metadata["nested"]["empty"].append("new")
|
||||
assert descriptor["metadata"]["nested"]["empty"] == []
|
||||
assert metadata["nested"]["empty"] == []
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
||||
def test_native_bhw1_alpha_shape_is_preserved(dtype):
|
||||
source = components(dtype)
|
||||
original = source.get_components()
|
||||
vars(source)["_VideoFromComponents__components"].alpha = original.alpha.unsqueeze(-1)
|
||||
descriptor, payload = codec.encode_video(source)
|
||||
result = codec.decode_video(descriptor, payload).get_components()
|
||||
exact_tensor(result.alpha, original.alpha.unsqueeze(-1))
|
||||
assert descriptor["alpha"]["descriptor"]["shape"] == [2, 2, 3, 1]
|
||||
|
||||
@pytest.mark.parametrize("start,duration,crop", [
|
||||
(-1.25, -2.5, (-3, 5, 7, -9)), (0, 0, None), (-0.0, 0.0, (0, 0, 2, 2)),
|
||||
])
|
||||
def test_file_preserves_bytes_cursor_signed_trim_crop_and_ownership(start, duration, crop, monkeypatch):
|
||||
source = file_value(start_time=start, duration=duration, crop=crop)
|
||||
stream = vars(source)["_VideoFromFile__file"]
|
||||
stream.seek(7)
|
||||
monkeypatch.setattr(VideoFromFile, "get_components", forbid)
|
||||
monkeypatch.setattr(VideoFromFile, "save_to", forbid)
|
||||
descriptor, payload = codec.encode_video(source)
|
||||
assert stream.tell() == 7
|
||||
assert payload == stream.getvalue()
|
||||
assert descriptor["start_time"] == start
|
||||
assert descriptor["duration"] == duration
|
||||
result = codec.decode_video(json.loads(json.dumps(descriptor)), payload)
|
||||
second = codec.decode_video(descriptor, payload)
|
||||
attrs = vars(result)
|
||||
assert type(result) is VideoFromFile
|
||||
assert attrs["_VideoFromFile__start_time"] == start
|
||||
assert type(attrs["_VideoFromFile__start_time"]) is type(start)
|
||||
assert attrs["_VideoFromFile__duration"] == duration
|
||||
assert attrs["_VideoFromFile__crop"] == crop
|
||||
if type(start) is float:
|
||||
assert math.copysign(1, attrs["_VideoFromFile__start_time"]) == math.copysign(1, start)
|
||||
decoded = attrs["_VideoFromFile__file"]
|
||||
assert type(decoded) is io.BytesIO
|
||||
assert decoded is not stream
|
||||
stream.close()
|
||||
assert decoded.getvalue() == payload
|
||||
decoded.write(b"changed")
|
||||
assert vars(second)["_VideoFromFile__file"].getvalue() == payload
|
||||
|
||||
def test_real_tiny_pyav_video_is_still_readable_after_source_closed():
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, "w", format="matroska") as container:
|
||||
stream = container.add_stream("ffv1", rate=Fraction(24, 1))
|
||||
# Width 32 avoids a native fillborders filter limitation on tiny
|
||||
# unaligned YUV frames; this is still a tiny CPU-only fixture.
|
||||
stream.width = stream.height = 32
|
||||
stream.pix_fmt = "yuv420p"
|
||||
for value in (16, 64, 128):
|
||||
frame = av.VideoFrame.from_ndarray(np.full((32, 32, 3), value, dtype=np.uint8), format="rgb24")
|
||||
for packet in stream.encode(frame):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode():
|
||||
container.mux(packet)
|
||||
data = buffer.getvalue()
|
||||
source = VideoFromFile(buffer, crop=(0, 0, 2, 2))
|
||||
descriptor, payload = codec.encode_video(source)
|
||||
expected = source.get_components()
|
||||
result = codec.decode_video(descriptor, payload)
|
||||
buffer.close()
|
||||
assert payload == data
|
||||
actual = result.get_components()
|
||||
exact_tensor(actual.images, expected.images)
|
||||
assert actual.frame_rate == expected.frame_rate
|
||||
assert tuple(actual.images.shape) == (3, 2, 2, 3)
|
||||
|
||||
def test_segments_order_complete_audio_no_transcode_and_independent_buffers(monkeypatch):
|
||||
values = [file_value(b"first", start_time=-2), file_value(b"second", duration=-4, crop=(1, 2, 3, 4))]
|
||||
source = VideoFromList([VideoFromList(values[:1]), values[1]], complete_audio=audio())
|
||||
monkeypatch.setattr(VideoFromFile, "get_components", forbid)
|
||||
monkeypatch.setattr(VideoFromFile, "save_to", forbid)
|
||||
monkeypatch.setattr(VideoFromList, "get_components", forbid)
|
||||
monkeypatch.setattr(VideoFromList, "save_to", forbid)
|
||||
descriptor, payload = codec.encode_video(source)
|
||||
assert descriptor["representation"] == "segments"
|
||||
assert payload.startswith(b"firstsecond")
|
||||
assert len(descriptor["segments"]) == 2
|
||||
decoded = codec.decode_video(json.loads(json.dumps(descriptor)), payload)
|
||||
assert type(decoded) is VideoFromList
|
||||
assert [vars(v)["_VideoFromFile__file"].getvalue() for v in decoded.videos] == [b"first", b"second"]
|
||||
exact_tensor(decoded.complete_audio["waveform"], source.complete_audio["waveform"])
|
||||
for old, new in zip(source.videos, decoded.videos):
|
||||
assert old is not new
|
||||
vars(old)["_VideoFromFile__file"].close()
|
||||
assert not vars(new)["_VideoFromFile__file"].closed
|
||||
assert vars(decoded.videos[0])["_VideoFromFile__start_time"] == -2
|
||||
assert vars(decoded.videos[1])["_VideoFromFile__crop"] == (1, 2, 3, 4)
|
||||
|
||||
def test_path_roots_symlinks_and_no_path_in_wire(tmp_path, monkeypatch):
|
||||
allowed = tmp_path / "allowed"
|
||||
allowed.mkdir()
|
||||
media = allowed / "private-name.mkv"
|
||||
media.write_bytes(b"media bytes")
|
||||
sibling = tmp_path / "allowed-sibling"
|
||||
sibling.mkdir()
|
||||
secret = sibling / "secret.mkv"
|
||||
secret.write_bytes(b"private")
|
||||
(allowed / "escape.mkv").symlink_to(secret)
|
||||
(allowed / "inside.mkv").symlink_to(media)
|
||||
descriptor, payload = codec.encode_video(VideoFromFile(str(media)), allowed_roots=[allowed])
|
||||
assert payload == b"media bytes"
|
||||
assert "private-name" not in json.dumps(descriptor)
|
||||
assert codec.encode_video(VideoFromFile(str(allowed / "inside.mkv")), allowed_roots=[allowed])[1] == payload
|
||||
for path in (secret, allowed / "escape.mkv", allowed, allowed / "missing"):
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(VideoFromFile(str(path)), allowed_roots=[allowed])
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(VideoFromFile(str(media)), allowed_roots=[])
|
||||
import folder_paths
|
||||
for name in ("get_input_directory", "get_output_directory", "get_temp_directory"):
|
||||
monkeypatch.setattr(folder_paths, name, lambda: str(allowed))
|
||||
assert codec.encode_video(VideoFromFile(str(media)))[1] == payload
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(VideoFromFile(str(secret)))
|
||||
result = codec.decode_video(descriptor, payload)
|
||||
media.unlink()
|
||||
assert vars(result)["_VideoFromFile__file"].getvalue() == b"media bytes"
|
||||
|
||||
def test_path_special_files_and_symlink_swap_are_rejected(tmp_path, monkeypatch):
|
||||
fifo = tmp_path / "fifo"
|
||||
os.mkfifo(fifo)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(VideoFromFile(str(fifo)), allowed_roots=[tmp_path])
|
||||
allowed = tmp_path / "allowed"
|
||||
allowed.mkdir()
|
||||
actual = allowed / "subdir"
|
||||
actual.mkdir()
|
||||
media = actual / "clip.mkv"
|
||||
media.write_bytes(b"allowed")
|
||||
outside = tmp_path / "outside"
|
||||
outside.mkdir()
|
||||
(outside / "clip.mkv").write_bytes(b"forbidden")
|
||||
real_resolve = codec.Path.resolve
|
||||
swapped = False
|
||||
|
||||
def swap_after_resolution(path, *args, **kwargs):
|
||||
nonlocal swapped
|
||||
resolved = real_resolve(path, *args, **kwargs)
|
||||
if path == media and not swapped:
|
||||
swapped = True
|
||||
actual.rename(allowed / "old")
|
||||
actual.symlink_to(outside, target_is_directory=True)
|
||||
return resolved
|
||||
|
||||
monkeypatch.setattr(codec.Path, "resolve", swap_after_resolution)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(VideoFromFile(str(media)), allowed_roots=[allowed])
|
||||
assert swapped
|
||||
|
||||
def test_file_size_budget_checked_before_read(tmp_path, monkeypatch):
|
||||
path = tmp_path / "large.mkv"
|
||||
path.write_bytes(b"12345")
|
||||
monkeypatch.setattr(codec.os, "read", forbid)
|
||||
with pytest.raises(ValueError, match="limit|budget"):
|
||||
codec.encode_video(VideoFromFile(str(path)), max_bytes=4, allowed_roots=[tmp_path])
|
||||
|
||||
@pytest.mark.parametrize("limit", [True, False, -1, 1.5, None])
|
||||
def test_invalid_limits(limit):
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(file_value(), max_bytes=limit)
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video({}, b"", max_bytes=limit)
|
||||
|
||||
def test_aggregate_budget_and_exact_boundary():
|
||||
value = VideoFromList([file_value(b"12"), file_value(b"345")])
|
||||
descriptor, payload = codec.encode_video(value, max_bytes=5)
|
||||
assert payload == b"12345"
|
||||
codec.decode_video(descriptor, payload, max_bytes=5)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value, max_bytes=4)
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload, max_bytes=4)
|
||||
stream = vars(value.videos[0])["_VideoFromFile__file"]
|
||||
stream.seek(1)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value.videos[0], max_bytes=1)
|
||||
assert stream.tell() == 1
|
||||
|
||||
@pytest.mark.parametrize("kind", ["components", "file", "segments"])
|
||||
def test_unknown_subclasses_never_materialized(kind):
|
||||
bases = {"components": VideoFromComponents, "file": VideoFromFile, "segments": VideoFromList}
|
||||
cls = type("UnknownVideo", (bases[kind],), {"get_components": forbid, "save_to": forbid})
|
||||
value = object.__new__(cls)
|
||||
with pytest.raises(ValueError, match="class|type|unsupported|Unsupported"):
|
||||
codec.encode_video(value)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(object())
|
||||
|
||||
@pytest.mark.parametrize("kind", ["components", "file", "segments"])
|
||||
def test_bridge_rejects_subclasses_without_materializing(kind, monkeypatch):
|
||||
nodes = importlib.import_module(package_name + ".nodes")
|
||||
bases = {"components": VideoFromComponents, "file": VideoFromFile, "segments": VideoFromList}
|
||||
cls = type("UnknownVideo", (bases[kind],), {"get_components": forbid, "save_to": forbid})
|
||||
monkeypatch.setattr(torch.Tensor, "cpu", forbid)
|
||||
with pytest.raises(ValueError, match="Unsupported native VIDEO class"):
|
||||
nodes.encode_value("VIDEO", object.__new__(cls))
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
||||
@pytest.mark.parametrize("kind", ["components", "segments", "file"])
|
||||
def test_bridge_stages_cpu_tensors_without_materialization(dtype, kind, monkeypatch):
|
||||
nodes = importlib.import_module(package_name + ".nodes")
|
||||
if kind == "components":
|
||||
source = components(dtype, metadata={"nested": ["雪", None, True]})
|
||||
state = vars(source)["_VideoFromComponents__components"]
|
||||
state.alpha = state.alpha.unsqueeze(-1)
|
||||
tensors = [state.images, state.alpha, state.audio["waveform"]]
|
||||
elif kind == "segments":
|
||||
source = VideoFromList([file_value(start_time=-1.5, duration=-2, crop=(-1, 2, 3, 4))],
|
||||
complete_audio=audio(dtype))
|
||||
tensors = [source.complete_audio["waveform"]]
|
||||
else:
|
||||
source = file_value(start_time=-0.0, duration=-2, crop=(-1, 2, 3, 4))
|
||||
tensors = []
|
||||
for tensor in tensors:
|
||||
tensor.requires_grad_(True)
|
||||
expected = codec.encode_video(source)
|
||||
detached, staged = [], []
|
||||
real_detach, real_cpu = torch.Tensor.detach, torch.Tensor.cpu
|
||||
|
||||
def detach(tensor):
|
||||
detached.append(id(tensor))
|
||||
return real_detach(tensor)
|
||||
|
||||
def cpu(tensor, *args, **kwargs):
|
||||
assert not tensor.requires_grad
|
||||
staged.append(tensor)
|
||||
return real_cpu(tensor, *args, **kwargs)
|
||||
|
||||
for cls in (VideoFromComponents, VideoFromFile, VideoFromList):
|
||||
monkeypatch.setattr(cls, "get_components", forbid)
|
||||
monkeypatch.setattr(cls, "save_to", forbid)
|
||||
monkeypatch.setattr(torch.Tensor, "detach", detach)
|
||||
monkeypatch.setattr(torch.Tensor, "cpu", cpu)
|
||||
actual = nodes.encode_value("VIDEO", source)
|
||||
assert len(staged) == len(tensors)
|
||||
assert all(id(tensor) in detached for tensor in tensors)
|
||||
assert all(tensor.requires_grad for tensor in tensors)
|
||||
assert actual == expected
|
||||
decoded = nodes.decode_value("VIDEO", *actual)
|
||||
assert type(decoded) is type(source)
|
||||
assert codec.encode_video(decoded) == expected
|
||||
assert len(staged) == len(tensors) # Pure codec never stages implicitly.
|
||||
|
||||
@pytest.mark.parametrize("field", ["images", "alpha", "audio", "complete_audio"])
|
||||
def test_pure_video_rejects_non_cpu_storage_without_transfer(field, monkeypatch):
|
||||
if field == "complete_audio":
|
||||
source = VideoFromList([file_value()], complete_audio=audio())
|
||||
target = source.complete_audio["waveform"]
|
||||
else:
|
||||
source = components()
|
||||
target = getattr(vars(source)["_VideoFromComponents__components"], field)
|
||||
if field == "audio":
|
||||
target = target["waveform"]
|
||||
# A property spy on real CPU tensors exercises the guard without a GPU
|
||||
# allocation, subclass, device discovery or CUDA initialization.
|
||||
real_device = torch.Tensor.device
|
||||
monkeypatch.setattr(torch.Tensor, "device", property(
|
||||
lambda tensor: types.SimpleNamespace(type="cuda") if tensor is target
|
||||
else real_device.__get__(tensor)))
|
||||
monkeypatch.setattr(torch.Tensor, "cpu", forbid)
|
||||
with pytest.raises(ValueError, match="must be on CPU"):
|
||||
codec.encode_video(source)
|
||||
|
||||
@pytest.mark.parametrize("field,bad", [
|
||||
("_VideoFromFile__start_time", True), ("_VideoFromFile__start_time", float("nan")),
|
||||
("_VideoFromFile__duration", float("inf")), ("_VideoFromFile__crop", (1, 2, True, 4)),
|
||||
("_VideoFromFile__crop", [1, 2, 3, 4]), ("_VideoFromFile__crop", (1, 2, 3)),
|
||||
("_VideoFromFile__file", b"not a BytesIO"),
|
||||
])
|
||||
def test_invalid_native_file_attributes(field, bad):
|
||||
value = file_value()
|
||||
setattr(value, field, bad)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
|
||||
def test_missing_attributes_and_invalid_component_state():
|
||||
value = file_value()
|
||||
del vars(value)["_VideoFromFile__duration"]
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
for field, bad in [("frame_rate", 24.0), ("frame_rate", Fraction(0)),
|
||||
("alpha", torch.ones(1, 1, 1, device="cpu"))]:
|
||||
value = components()
|
||||
setattr(vars(value)["_VideoFromComponents__components"], field, bad)
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
|
||||
@pytest.mark.parametrize("metadata", [
|
||||
{"value": float("nan")}, {"value": float("inf")}, {1: "non-string key"},
|
||||
{"value": (1, 2)}, {"value": object()}, {"large": "x" * 65536},
|
||||
])
|
||||
def test_invalid_metadata_rejected_both_directions(metadata):
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(components(metadata=metadata))
|
||||
descriptor, payload = codec.encode_video(components())
|
||||
descriptor["metadata"] = metadata
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
def test_metadata_depth_cycles_and_exact_size():
|
||||
metadata = {}
|
||||
current = metadata
|
||||
for _ in range(40):
|
||||
current["child"] = {}
|
||||
current = current["child"]
|
||||
for value in (metadata,):
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(components(metadata=value))
|
||||
cyclic = {}
|
||||
cyclic["self"] = cyclic
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(components(metadata=cyclic))
|
||||
exact = {"x": "a" * (65536 - 8)}
|
||||
assert len(json.dumps(exact, separators=(",", ":")).encode()) == 65536
|
||||
descriptor, payload = codec.encode_video(components(metadata=exact))
|
||||
assert codec.decode_video(descriptor, payload).get_components().metadata == exact
|
||||
|
||||
@pytest.mark.parametrize("field,bad", [
|
||||
("encoding", "other"), ("representation", "other"), ("byte_length", True),
|
||||
("byte_length", -1), ("byte_length", 12345), ("start_time", False),
|
||||
("duration", float("inf")), ("crop", [0, 0, True, 4]), ("crop", [1, 2, 3]),
|
||||
])
|
||||
def test_malformed_file_descriptors(field, bad):
|
||||
descriptor, payload = codec.encode_video(file_value())
|
||||
descriptor[field] = bad
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
@pytest.mark.parametrize("field,bad", [("offset", True), ("offset", -1), ("offset", 1),
|
||||
("length", True), ("length", -1), ("length", 2**80)])
|
||||
def test_invalid_offsets_lengths(field, bad):
|
||||
descriptor, payload = codec.encode_video(file_value())
|
||||
descriptor["data"][field] = bad
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
def test_truncation_trailing_gaps_overlap_and_unknown_fields():
|
||||
value = VideoFromList([file_value(b"123"), file_value(b"456")])
|
||||
original, payload = codec.encode_video(value)
|
||||
for wire in (payload[:-1], payload + b"x"):
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(original, wire)
|
||||
for offset in (2, 4):
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["segments"][1]["offset"] = offset
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["surprise"] = "bad"
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["segments"][0]["descriptor"]["byte_length"] += 1
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["segments"][0]["descriptor"] = descriptor
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
def test_all_descriptors_validated_before_tensor_allocation(monkeypatch):
|
||||
original, payload = codec.encode_video(components())
|
||||
monkeypatch.setattr(codec, "decode_payload", forbid)
|
||||
cases = []
|
||||
for key, value in [("shape", [2**80, 2, 3]), ("shape", [True, 2, 3]),
|
||||
("dtype", "pickle"), ("layout", "BHWC")]:
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["alpha"]["descriptor"][key] = value
|
||||
cases.append(descriptor)
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["audio"]["descriptor"]["sample_rate"] = True
|
||||
cases.append(descriptor)
|
||||
descriptor = copy.deepcopy(original)
|
||||
descriptor["alpha"]["offset"] += 1
|
||||
cases.append(descriptor)
|
||||
for descriptor in cases:
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
@pytest.mark.parametrize("field,bad", [("numerator", True), ("denominator", True),
|
||||
("denominator", 0), ("numerator", -3),
|
||||
("denominator", -1001)])
|
||||
def test_invalid_fraction_fields(field, bad):
|
||||
descriptor, payload = codec.encode_video(components())
|
||||
descriptor["frame_rate"][field] = bad
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
|
||||
def test_segment_limit_and_exact_maximum(monkeypatch):
|
||||
value = VideoFromList([file_value(b"x") for _ in range(64)])
|
||||
descriptor, payload = codec.encode_video(value)
|
||||
result = codec.decode_video(descriptor, payload)
|
||||
assert len(result.videos) == 64
|
||||
value.videos.append(file_value(b"x"))
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
descriptor["segments"].append(copy.deepcopy(descriptor["segments"][-1]))
|
||||
monkeypatch.setattr(codec, "decode_payload", forbid)
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
value.videos = []
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
value.videos = [components()]
|
||||
with pytest.raises(ValueError):
|
||||
codec.encode_video(value)
|
||||
|
||||
@pytest.mark.parametrize("descriptor,payload", [(None, b""), ([], b""), ({}, bytearray()), ({}, "")])
|
||||
def test_invalid_top_level_types(descriptor, payload):
|
||||
with pytest.raises(ValueError):
|
||||
codec.decode_video(descriptor, payload)
|
||||
@@ -60,6 +60,29 @@ class ValidateServerlessWorkerTests(unittest.TestCase):
|
||||
self.assertIsNot(result, worker)
|
||||
self.assertEqual(worker, {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"})
|
||||
|
||||
def test_omitted_media_profile_remains_absent_and_means_legacy(self):
|
||||
result = serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": "ep-123"}
|
||||
)
|
||||
self.assertNotIn("media_profile", result)
|
||||
self.assertEqual(result.get("media_profile", "legacy-image-v1"), "legacy-image-v1")
|
||||
|
||||
def test_explicit_media_profiles_are_retained(self):
|
||||
for profile in ("legacy-image-v1", "comfy-media-v1"):
|
||||
with self.subTest(profile=profile):
|
||||
result = serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": "ep-123", "media_profile": profile}
|
||||
)
|
||||
self.assertEqual(result["media_profile"], profile)
|
||||
|
||||
def test_media_profile_rejects_unknown_nonstring_and_null_values(self):
|
||||
for profile in ("", "auto", "comfy-media-v2", " comfy-media-v1", "COMFY-MEDIA-V1", None, True, 1, [], {}):
|
||||
with self.subTest(profile=profile):
|
||||
with self.assertRaisesRegex(ValueError, "media_profile"):
|
||||
serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": "ep-123", "media_profile": profile}
|
||||
)
|
||||
|
||||
def test_accepts_inclusive_numeric_bounds(self):
|
||||
for jobs, timeout in ((1, 30), (16, 86400)):
|
||||
with self.subTest(jobs=jobs, timeout=timeout):
|
||||
@@ -307,6 +330,7 @@ class PublicConfigTests(unittest.TestCase):
|
||||
"type": "serverless",
|
||||
"provider": "runpod",
|
||||
"endpoint_id": "ep-123",
|
||||
"media_profile": "comfy-media-v1",
|
||||
"api_key": "inline-secret",
|
||||
"clear_api_key": "another-secret",
|
||||
"credential": "credential-secret",
|
||||
@@ -334,6 +358,8 @@ class PublicConfigTests(unittest.TestCase):
|
||||
self.assertNotIn("api_key", local)
|
||||
self.assertNotIn("credentials_note", result["settings"])
|
||||
self.assertEqual(first["name"], "remote GPU")
|
||||
self.assertEqual(first["media_profile"], "comfy-media-v1")
|
||||
self.assertNotIn("media_profile", second)
|
||||
self.assertEqual(config, original)
|
||||
self.assertIsNot(result, config)
|
||||
self.assertIsNot(result["workers"], config["workers"])
|
||||
|
||||
@@ -0,0 +1,421 @@
|
||||
"""CPU-only graph/compiler tests; no ComfyUI package bootstrap is needed."""
|
||||
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE = "serverless_media_prompt_testpkg"
|
||||
# Real paths let sibling imports use production code, without running the
|
||||
# extension's ComfyUI-dependent __init__.py or API registration side effects.
|
||||
for suffix in ("", ".api", ".api.orchestration", ".serverless_bridge"):
|
||||
package = types.ModuleType(PACKAGE + suffix)
|
||||
package.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
|
||||
sys.modules[package.__name__] = package
|
||||
|
||||
compiler = importlib.import_module(f"{PACKAGE}.api.orchestration.serverless_media_prompt")
|
||||
legacy = importlib.import_module(f"{PACKAGE}.api.orchestration.serverless_prompt")
|
||||
contract = importlib.import_module(f"{PACKAGE}.serverless_bridge.contract")
|
||||
compile_media_prompt = compiler.compile_media_prompt
|
||||
SOCKETS = {"images": "IMAGE", "audio": "AUDIO", "masks": "MASK", "texts": "STRING", "video": "VIDEO"}
|
||||
|
||||
|
||||
def _node(class_type, **inputs):
|
||||
return {"class_type": class_type, "inputs": inputs}
|
||||
|
||||
|
||||
def _collector(**sockets):
|
||||
return _node(
|
||||
"DistributedMediaCollector", multi_job_id="job/42", worker_id="worker_0", **sockets
|
||||
)
|
||||
|
||||
|
||||
def _prompt():
|
||||
return {
|
||||
"source:1": _node("LoadImage", image="example.png"),
|
||||
"collector:1": _collector(images=["source:1", 0]),
|
||||
}
|
||||
|
||||
|
||||
def _run(prompt):
|
||||
return compile_media_prompt(prompt, "request:1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("socket,kind", SOCKETS.items())
|
||||
def test_one_export_for_each_typed_socket_with_exact_canonical_binding(socket, kind):
|
||||
prompt = {"source": _node("Source"), "collector": _collector(**{socket: ["source", 0]})}
|
||||
graph, bindings = _run(prompt)
|
||||
assert len(bindings) == 1
|
||||
binding = bindings[0]
|
||||
assert set(binding) == contract.BINDING_FIELDS
|
||||
assert binding == {
|
||||
"binding_id": binding["binding_id"],
|
||||
"collector_id": "collector",
|
||||
"job_id": "job/42",
|
||||
"participant_id": "worker_0",
|
||||
"node_id": binding["node_id"],
|
||||
"socket": socket,
|
||||
"kind": kind,
|
||||
}
|
||||
assert all(isinstance(value, str) for value in binding.values())
|
||||
assert contract.validate_bindings(bindings) == {binding["binding_id"]: binding}
|
||||
assert graph[binding["node_id"]] == {
|
||||
"class_type": "MediaExportV1",
|
||||
"inputs": {
|
||||
"request_id": "request:1",
|
||||
"binding_json": json.dumps(binding),
|
||||
socket: ["source", 0],
|
||||
},
|
||||
}
|
||||
assert set(graph) == {"source", binding["node_id"]}
|
||||
|
||||
|
||||
def test_mixed_collectors_all_five_kinds_have_independent_dispatch_targets():
|
||||
prompt = {
|
||||
"shared": _node("MultiMediaSource"),
|
||||
"collector:z": _collector(audio=["shared", 1], video=["shared", 4]),
|
||||
"collector:a": _collector(images=["shared", 0], masks=["shared", 2], texts=["shared", 3]),
|
||||
}
|
||||
prompt["collector:z"]["inputs"].update(multi_job_id="job:z", worker_id="participant:z")
|
||||
graph, bindings = _run(prompt)
|
||||
assert [(b["collector_id"], b["socket"]) for b in bindings] == [
|
||||
("collector:a", "images"), ("collector:a", "masks"), ("collector:a", "texts"),
|
||||
("collector:z", "audio"), ("collector:z", "video"),
|
||||
]
|
||||
assert len({b["binding_id"] for b in bindings}) == 5
|
||||
targets = {b["node_id"] for b in bindings}
|
||||
assert len(targets) == 5
|
||||
assert set(graph) == {"shared"} | targets
|
||||
for binding in bindings:
|
||||
export = graph[binding["node_id"]]
|
||||
inputs = prompt[binding["collector_id"]]["inputs"]
|
||||
assert export["inputs"][binding["socket"]] == inputs[binding["socket"]]
|
||||
assert binding["job_id"] == inputs["multi_job_id"]
|
||||
assert binding["participant_id"] == inputs["worker_id"]
|
||||
|
||||
|
||||
def test_prunes_descendants_and_unrelated_outputs_preserving_required_output_utility():
|
||||
prompt = _prompt()
|
||||
prompt.update({
|
||||
"required-output": _node("RequiredOutputUtility", images=["source:1", 0]),
|
||||
"consumer": _node("ImageTransform", image=["required-output", 0]),
|
||||
"descendant": _node("ImageTransform", image=["collector:1", 0]),
|
||||
"after": _node("SaveImage", images=["descendant", 0]),
|
||||
"unused-output": _node("SaveImage", images=["source:1", 0]),
|
||||
"unrelated": _node("LoadAudio"),
|
||||
})
|
||||
prompt["required-output"]["_meta"] = {"title": "Required OUTPUT_NODE utility"}
|
||||
prompt["collector:1"]["inputs"].update(images=["consumer", 0], load_balance=True, pass_through=False)
|
||||
graph, bindings = _run(prompt)
|
||||
assert set(graph) == {"source:1", "required-output", "consumer", bindings[0]["node_id"]}
|
||||
assert graph["required-output"] == prompt["required-output"]
|
||||
assert not {"load_balance", "pass_through"} & graph[bindings[0]["node_id"]]["inputs"].keys()
|
||||
|
||||
|
||||
def test_stock_primitive_string_preserves_empty_and_link_looking_text():
|
||||
for text in ("", '["source:1", 0]', "hello\nworld"):
|
||||
prompt = {"text": _node("PrimitiveString", value=text), "collector": _collector(texts=["text", 0])}
|
||||
graph, bindings = _run(prompt)
|
||||
assert graph["text"] == prompt["text"]
|
||||
assert graph[bindings[0]["node_id"]]["inputs"]["texts"] == ["text", 0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type,inputs", [
|
||||
("DistributedSeed", {}),
|
||||
("DistributedSeed", {"seed": 10, "is_worker": True, "worker_id": "worker_1"}),
|
||||
("DistributedSeed", {"seed": 10, "is_worker": False, "worker_id": "worker_1"}),
|
||||
("DistributedSeed", {"seed": 10, "is_worker": True, "worker_id": "bad"}),
|
||||
("DistributedValue", {"default_value": "3.9", "worker_values": '{"_type":"INT","1":"8.7"}', "is_worker": True, "worker_id": "worker_0"}),
|
||||
("DistributedValue", {"default_value": "3.9", "worker_values": '{"_type":"INT","2":""}', "is_worker": True, "worker_id": "worker_1"}),
|
||||
("DistributedValue", {"default_value": "fallback", "worker_values": "malformed"}),
|
||||
("DistributedValue", {"default_value": 1.25, "worker_values": '{"_type":"FLOAT","1":"bad"}', "is_worker": True, "worker_id": "0"}),
|
||||
("DistributedModelName", {"text": "model.safetensors"}),
|
||||
("DistributedModelName", {"text": {"name": "checkpoint"}}),
|
||||
("DistributedModelName", {"text": ["a", "b", "c"]}),
|
||||
])
|
||||
def test_upstream_scalar_values_match_legacy_compiler(class_type, inputs):
|
||||
prompt = _prompt()
|
||||
prompt["scalar"] = _node(class_type, **inputs)
|
||||
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 0]
|
||||
old_prompt = copy.deepcopy(prompt)
|
||||
old_prompt["collector:1"]["class_type"] = "DistributedCollector"
|
||||
expected, _ = legacy.compile_serverless_prompt(old_prompt)
|
||||
graph, _ = _run(prompt)
|
||||
assert graph["source:1"] == expected["source:1"]
|
||||
assert "scalar" not in graph
|
||||
|
||||
|
||||
def test_imports_same_scalar_resolver_and_resolves_primitive_chains():
|
||||
assert compiler._make_scalar_resolver is legacy._make_scalar_resolver
|
||||
prompt = _prompt()
|
||||
prompt.update({
|
||||
"int": _node("PrimitiveInt", value=25),
|
||||
"int-alias": _node("PrimitiveInt", value=["int", 0]),
|
||||
"float": _node("PrimitiveFloat", value=1.25),
|
||||
"bool": _node("PrimitiveBoolean", value=True),
|
||||
"string": _node("PrimitiveString", value="model.safetensors"),
|
||||
"seed": _node("DistributedSeed", seed=["int-alias", 0], is_worker=True, worker_id="worker_0"),
|
||||
"value": _node("DistributedValue", default_value=["float", 0], worker_values='{"_type":"FLOAT","1":"2.75"}', is_worker=True, worker_id="worker_0"),
|
||||
"enabled": _node("DistributedValue", default_value=["bool", 0]),
|
||||
"model": _node("DistributedModelName", text=["string", 0]),
|
||||
})
|
||||
prompt["source:1"]["inputs"].update(seed=["seed", 0], cfg=["value", 0], enabled=["enabled", 0], model=["model", 0])
|
||||
graph, bindings = _run(prompt)
|
||||
assert graph["source:1"]["inputs"] == {"image": "example.png", "seed": 26, "cfg": 2.75, "enabled": True, "model": "model.safetensors"}
|
||||
assert set(graph) == {"source:1", bindings[0]["node_id"]}
|
||||
|
||||
|
||||
def test_distributed_text_value_can_be_wrapped_by_primitive_string():
|
||||
prompt = {
|
||||
"value": _node("DistributedModelName", text="model.safetensors"),
|
||||
"text": _node("PrimitiveString", value=["value", 0]),
|
||||
"collector": _collector(texts=["text", 0]),
|
||||
}
|
||||
graph, bindings = _run(prompt)
|
||||
assert graph["text"]["inputs"]["value"] == "model.safetensors"
|
||||
assert graph[bindings[0]["node_id"]]["inputs"]["texts"] == ["text", 0]
|
||||
assert "value" not in graph
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type,input_name", [
|
||||
("DistributedSeed", "seed"), ("DistributedValue", "default_value"), ("DistributedModelName", "text"),
|
||||
])
|
||||
def test_dynamic_scalar_resolution_stays_fail_closed(class_type, input_name):
|
||||
prompt = _prompt()
|
||||
prompt["dynamic"] = _node("ArbitraryValue")
|
||||
prompt["scalar"] = _node(class_type, **{input_name: ["dynamic", 0]})
|
||||
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 0]
|
||||
with pytest.raises(ValueError, match="literal or supported Primitive"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_namespaced_export_ids_avoid_existing_ids_even_on_pruned_nodes():
|
||||
prompt = _prompt()
|
||||
_, first = _run(prompt)
|
||||
original_id = first[0]["node_id"]
|
||||
assert ":" in original_id and not original_id.isdigit()
|
||||
prompt[original_id] = _node("ExistingUtility")
|
||||
_, second = _run(prompt)
|
||||
assert second[0]["node_id"] not in prompt
|
||||
assert second[0]["binding_id"] == first[0]["binding_id"]
|
||||
prompt[second[0]["node_id"]] = _node("AnotherUtility")
|
||||
graph, third = _run(prompt)
|
||||
assert third[0]["node_id"] not in prompt
|
||||
assert third == _run(dict(reversed(list(prompt.items()))))[1]
|
||||
assert set(graph) == {"source:1", third[0]["node_id"]}
|
||||
|
||||
|
||||
def test_long_valid_identifiers_remain_exact_and_generated_ids_fit_contract():
|
||||
collector_id = "namespace/" + "c" * 190
|
||||
prompt = {"source:1": _node("Source"), collector_id: _collector(images=["source:1", 0])}
|
||||
prompt[collector_id]["inputs"].update(multi_job_id="j" * 200, worker_id="p" * 200)
|
||||
_, bindings = _run(prompt)
|
||||
assert bindings[0]["collector_id"] == collector_id
|
||||
assert bindings[0]["job_id"] == "j" * 200
|
||||
assert bindings[0]["participant_id"] == "p" * 200
|
||||
contract.validate_bindings(bindings)
|
||||
|
||||
|
||||
def test_deterministic_independent_copies_and_request_identity():
|
||||
prompt = _prompt()
|
||||
prompt["source:1"]["_meta"] = {"nested": {"values": [1, 2, 3]}}
|
||||
original = copy.deepcopy(prompt)
|
||||
graph, bindings = _run(prompt)
|
||||
assert prompt == original
|
||||
assert (graph, bindings) == _run(prompt)
|
||||
assert (graph, bindings) == _run(dict(reversed(list(prompt.items()))))
|
||||
other_graph, other_bindings = compile_media_prompt(prompt, "request:2")
|
||||
assert other_bindings == bindings
|
||||
target = bindings[0]["node_id"]
|
||||
assert other_graph[target]["inputs"]["request_id"] == "request:2"
|
||||
graph["source:1"]["_meta"]["nested"]["values"].append(4)
|
||||
graph[target]["inputs"]["images"][0] = "changed"
|
||||
bindings[0]["collector_id"] = "changed"
|
||||
assert prompt == original
|
||||
assert json.loads(graph[target]["inputs"]["binding_json"])["collector_id"] == "collector:1"
|
||||
assert other_graph["source:1"] == original["source:1"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [None, [], {}, {"x": None}, {"x": []}, {"x": {}},
|
||||
{"x": {"class_type": None}}, {"x": {"class_type": ""}},
|
||||
{"x": {"class_type": "Source", "inputs": []}},
|
||||
{"x": {"class_type": "Source", "inputs": None}},
|
||||
{"x": {"class_type": "Source", "inputs": {1: "literal"}}},
|
||||
])
|
||||
def test_rejects_malformed_prompt_or_missing_collectors(bad):
|
||||
with pytest.raises(ValueError):
|
||||
_run(bad)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [None, 1, True, "", " bad", "bad ", "bad\\id", "bad\nid", "x" * 201])
|
||||
def test_rejects_malformed_ids_without_normalizing(bad):
|
||||
prompt = _prompt()
|
||||
prompt[bad] = _node("Source")
|
||||
with pytest.raises(ValueError):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [None, 1, True, "", " bad", "bad ", "bad\\id", "x" * 201])
|
||||
@pytest.mark.parametrize("field", ["multi_job_id", "worker_id", "request_id"])
|
||||
def test_rejects_malformed_binding_and_request_identity(field, bad):
|
||||
prompt = _prompt()
|
||||
request_id = "request:1"
|
||||
if field == "request_id":
|
||||
request_id = bad
|
||||
else:
|
||||
prompt["collector:1"]["inputs"][field] = bad
|
||||
with pytest.raises(ValueError):
|
||||
compile_media_prompt(prompt, request_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["multi_job_id", "worker_id"])
|
||||
def test_requires_prepared_literal_hidden_identity(field):
|
||||
prompt = _prompt()
|
||||
del prompt["collector:1"]["inputs"][field]
|
||||
with pytest.raises(ValueError):
|
||||
_run(prompt)
|
||||
prompt["collector:1"]["inputs"][field] = ["source:1", 0]
|
||||
with pytest.raises(ValueError):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("socket", SOCKETS)
|
||||
@pytest.mark.parametrize("value", [None, "literal", 5, True, [], ["source:1"], ["source:1", 0, 1], ("source:1", 0)])
|
||||
def test_every_present_media_socket_requires_an_api_link(socket, value):
|
||||
prompt = _prompt()
|
||||
prompt["collector:1"]["inputs"][socket] = value
|
||||
with pytest.raises(ValueError, match="link"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_each_collector_requires_at_least_one_media_socket():
|
||||
prompt = _prompt()
|
||||
prompt["empty"] = _collector(load_balance=True, pass_through=False)
|
||||
with pytest.raises(ValueError, match="connected|linked"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("link", [["missing", 0], ["source:1", -1], ["source:1", 0.0],
|
||||
["source:1", "0"], ["source:1", True], ["source:1", False], [1, 0],
|
||||
[True, 0], [None, 0], [[], 0], [{}, 0], [" source:1", 0],
|
||||
])
|
||||
@pytest.mark.parametrize("location", ["collector", "ordinary", "unrelated"])
|
||||
def test_rejects_invalid_links_everywhere_before_pruning(location, link):
|
||||
prompt = _prompt()
|
||||
prompt["1"] = _node("Source") # An integer/bool source must not be stringified to this ID.
|
||||
if location == "collector":
|
||||
prompt["collector:1"]["inputs"]["images"] = link
|
||||
elif location == "ordinary":
|
||||
prompt["source:1"]["inputs"]["parameter"] = link
|
||||
else:
|
||||
prompt["unrelated"] = _node("Utility", value=link)
|
||||
with pytest.raises(ValueError):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type", [
|
||||
"DistributedCollector", "DistributedEmptyImage", "UltimateSDUpscaleDistributed",
|
||||
"DistributedImageListStart", "DistributedImageListEnd", "DistributedWorkerSink",
|
||||
"ImageBatchDivider", "AudioBatchDivider", "CustomDistributedNode", "custombatchdivider",
|
||||
])
|
||||
def test_rejects_unsupported_distributed_families_even_if_unused(class_type):
|
||||
prompt = _prompt()
|
||||
prompt["unsupported"] = _node(class_type)
|
||||
with pytest.raises(ValueError, match="unsupported"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type", ["DistributedSeed", "DistributedValue", "DistributedModelName", "PrimitiveInt", "PrimitiveString"])
|
||||
def test_rejects_invalid_known_scalar_output_slots(class_type):
|
||||
prompt = _prompt()
|
||||
prompt["scalar"] = _node(class_type)
|
||||
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 1]
|
||||
with pytest.raises(ValueError, match="slot"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("socket", SOCKETS)
|
||||
def test_rejects_direct_distributed_scalar_media_links(socket):
|
||||
prompt = {"scalar": _node("DistributedValue", default_value="text"), "collector": _collector(**{socket: ["scalar", 0]})}
|
||||
with pytest.raises(ValueError, match="scalar"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("indirect", [False, True])
|
||||
def test_rejects_collector_chains_across_different_media(indirect):
|
||||
prompt = _prompt()
|
||||
source = "collector:1"
|
||||
if indirect:
|
||||
prompt["processing"] = _node("ImageToVideo", images=[source, 0])
|
||||
source = "processing"
|
||||
prompt["collector:2"] = _collector(video=[source, 0])
|
||||
with pytest.raises(ValueError, match="collector-to-collector"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["ordinary", "self", "unrelated", "primitive", "collector"])
|
||||
def test_rejects_cycles_in_entire_prompt(mode):
|
||||
prompt = _prompt()
|
||||
if mode == "ordinary":
|
||||
prompt["source:1"]["inputs"]["other"] = ["other", 0]
|
||||
prompt["other"] = _node("Utility", value=["source:1", 0])
|
||||
elif mode == "self":
|
||||
prompt["source:1"]["inputs"]["other"] = ["source:1", 0]
|
||||
elif mode == "unrelated":
|
||||
prompt["unused"] = _node("Utility", value=["unused", 0])
|
||||
elif mode == "primitive":
|
||||
prompt["p1"] = _node("PrimitiveInt", value=["p2", 0])
|
||||
prompt["p2"] = _node("PrimitiveInt", value=["p1", 0])
|
||||
else:
|
||||
prompt["source:1"]["inputs"]["other"] = ["collector:1", 0]
|
||||
with pytest.raises(ValueError, match="cycle"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_deep_acyclic_graph_does_not_depend_on_python_recursion_limit():
|
||||
prompt = {"n:0": _node("Source")}
|
||||
for index in range(1, 1500):
|
||||
prompt[f"n:{index}"] = _node("Utility", value=[f"n:{index - 1}", 0])
|
||||
prompt["collector"] = _collector(images=["n:1499", 0])
|
||||
graph, bindings = _run(prompt)
|
||||
assert len(graph) == 1501
|
||||
assert graph[bindings[0]["node_id"]]["inputs"]["images"] == ["n:1499", 0]
|
||||
|
||||
|
||||
def test_binding_count_respects_bridge_contract():
|
||||
prompt = {"source": _node("Source")}
|
||||
for index in range(contract.MAX_ARTIFACTS + 1):
|
||||
prompt[f"collector:{index}"] = _collector(images=["source", 0])
|
||||
with pytest.raises(ValueError, match="count"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_import_and_compile_do_not_import_comfy_or_tensor_libraries():
|
||||
script = '''
|
||||
import builtins, importlib, pathlib, sys, types
|
||||
root = pathlib.Path(sys.argv[1])
|
||||
name = "isolated_media_compile"
|
||||
for suffix in ("", ".api", ".api.orchestration", ".serverless_bridge"):
|
||||
package = types.ModuleType(name + suffix)
|
||||
package.__path__ = [str(root.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(root)]
|
||||
sys.modules[package.__name__] = package
|
||||
original_import = builtins.__import__
|
||||
def guarded_import(name, *args, **kwargs):
|
||||
if name.split(".")[0] in {"comfy", "comfy_api", "nodes", "server", "torch", "numpy"}:
|
||||
raise AssertionError("unexpected runtime import: " + name)
|
||||
return original_import(name, *args, **kwargs)
|
||||
builtins.__import__ = guarded_import
|
||||
module = importlib.import_module(name + ".api.orchestration.serverless_media_prompt")
|
||||
prompt = {"s": {"class_type": "Source", "inputs": {}}, "c": {"class_type": "DistributedMediaCollector", "inputs": {"texts": ["s", 0], "multi_job_id": "job", "worker_id": "worker"}}}
|
||||
graph, bindings = module.compile_media_prompt(prompt, "request")
|
||||
assert graph[bindings[0]["node_id"]]["class_type"] == "MediaExportV1"
|
||||
'''
|
||||
result = subprocess.run([sys.executable, "-c", script, str(ROOT)], text=True, capture_output=True)
|
||||
assert result.returncode == 0, result.stderr
|
||||
@@ -217,6 +217,78 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
|
||||
self.assertNotIn(obsolete, saved)
|
||||
self.assertEqual(self.secrets.get("pool-1"), secret)
|
||||
|
||||
async def test_media_profile_persists_and_is_public_without_changing_credentials(self):
|
||||
for profile in ("legacy-image-v1", "comfy-media-v1"):
|
||||
with self.subTest(profile=profile):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
"media_profile": profile, "api_key": "saved-key", "enabled": False,
|
||||
}))
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["media_profile"], profile)
|
||||
config_module.invalidate_config_cache()
|
||||
public = await config_routes.get_config_endpoint(_FakeRequest())
|
||||
self.assertEqual(public.status, 200)
|
||||
self.assertEqual(public.payload["workers"][0]["media_profile"], profile)
|
||||
self.assertTrue(public.payload["workers"][0]["api_key_configured"])
|
||||
self.assertNotIn("saved-key", json.dumps(public.payload))
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
|
||||
|
||||
async def test_old_clients_omit_media_profile_without_downgrading_saved_profile(self):
|
||||
for stored in ({}, {"media_profile": "comfy-media-v1"}):
|
||||
with self.subTest(stored=stored):
|
||||
worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", **stored}
|
||||
self.config["workers"] = [worker]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "name": "Old client save", "jobs_per_run": 2,
|
||||
}))
|
||||
self.assertEqual(response.status, 200)
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
self.assertEqual({k: v for k, v in saved.items() if k == "media_profile"}, stored)
|
||||
# The bulk route is also used by older clients with no profile field.
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "jobs_per_run": 3,
|
||||
}]}))
|
||||
self.assertEqual(response.status, 200)
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
self.assertEqual({k: v for k, v in saved.items() if k == "media_profile"}, stored)
|
||||
|
||||
async def test_invalid_media_profile_is_rejected_by_worker_and_bulk_routes(self):
|
||||
for profile in ("auto", "comfy-media-v2", None, False, 1, [], {}):
|
||||
with self.subTest(profile=profile):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
"media_profile": profile, "api_key": "never-save",
|
||||
}))
|
||||
self.assertEqual(response.status, 400)
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "media_profile": profile,
|
||||
}]}))
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
|
||||
async def test_active_pool_blocks_media_profile_changes_in_both_directions(self):
|
||||
for old_profile, new_profile in ((None, "comfy-media-v1"), ("comfy-media-v1", "legacy-image-v1")):
|
||||
with self.subTest(old_profile=old_profile):
|
||||
worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123"}
|
||||
if old_profile is not None:
|
||||
worker["media_profile"] = old_profile
|
||||
self.config["workers"] = [worker]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
original = self.read_config_file()
|
||||
with patch.object(config_routes, "_serverless_pool_is_active", return_value=True):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "media_profile": new_profile,
|
||||
}))
|
||||
self.assertEqual(response.status, 409)
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [dict(worker, media_profile=new_profile)],
|
||||
}))
|
||||
self.assertEqual(response.status, 409)
|
||||
self.assertEqual(self.read_config_file(), original)
|
||||
|
||||
async def test_enable_only_update_preserves_serverless_settings_and_saved_key(self):
|
||||
self.secrets.set("pool-1", "saved-secret")
|
||||
self.config["workers"] = [{
|
||||
@@ -302,13 +374,14 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
|
||||
"workers": [{
|
||||
"id": "pool-1", "type": "local", "port": 8188,
|
||||
"provider": "runpod", "endpoint_id": "legacy-ep", "api_key_configured": True,
|
||||
"media_profile": "comfy-media-v1",
|
||||
}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
for field in ("provider", "endpoint_id", "api_key_configured"):
|
||||
for field in ("provider", "endpoint_id", "api_key_configured", "media_profile"):
|
||||
self.assertNotIn(field, saved)
|
||||
|
||||
async def test_bulk_invalid_serverless_record_does_not_replace_configuration(self):
|
||||
@@ -341,6 +414,7 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "name": "Pool", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2, "timeout_seconds": 600,
|
||||
"media_profile": "comfy-media-v1",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
@@ -353,7 +427,7 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
worker = self.read_config_file()["workers"][0]
|
||||
self.assertEqual(worker["type"], "local")
|
||||
for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds"):
|
||||
for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "media_profile"):
|
||||
self.assertNotIn(serverless_field, worker)
|
||||
|
||||
async def test_enabled_toggle_allowed_during_active_run_but_endpoint_change_is_rejected(self):
|
||||
@@ -575,19 +649,41 @@ class PersistedPoolGuardTests(_ConfigRouteFixture):
|
||||
runtime_module.manager = self.manager
|
||||
else:
|
||||
self.manager.tasks["retained"] = types.SimpleNamespace(done=lambda: True)
|
||||
for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"type": "local", "port": 8188}):
|
||||
for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"type": "local", "port": 8188}, {"media_profile": "comfy-media-v1"}):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest(dict(fields, worker_id="pool-1")))
|
||||
self.assertEqual(response.status, 409, fields)
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
self.assertEqual(response.status, 409)
|
||||
worker = original["workers"][0]
|
||||
for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")]):
|
||||
for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")], [dict(worker, media_profile="comfy-media-v1")]):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": workers}))
|
||||
self.assertEqual(response.status, 409, workers)
|
||||
self.assertEqual(self.read_config_file(), original)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
|
||||
self.session_factory.assert_not_awaited()
|
||||
|
||||
async def test_retained_jobs_allow_same_or_omitted_media_profile(self):
|
||||
self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status="UNKNOWN")
|
||||
for profile in ("legacy-image-v1", "comfy-media-v1"):
|
||||
with self.subTest(profile=profile):
|
||||
worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123"}
|
||||
# An absent legacy field and an explicit legacy selection are equivalent.
|
||||
if profile != "legacy-image-v1":
|
||||
worker["media_profile"] = profile
|
||||
self.config["workers"] = [worker]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
for fields in ({"media_profile": profile}, {}):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "name": "Same profile", **fields,
|
||||
}))
|
||||
self.assertEqual(response.status, 200)
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", **fields,
|
||||
}]}))
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["media_profile"], profile)
|
||||
self.session_factory.assert_not_awaited()
|
||||
|
||||
async def test_retained_jobs_allow_images_per_run_and_legacy_concurrency_saves(self):
|
||||
# Upgrading from a config with concurrency=1 must not trip the busy guard,
|
||||
# and the image count stays editable while old jobs are unresolved.
|
||||
|
||||
@@ -0,0 +1,330 @@
|
||||
"""Bounded input packing; native acceptance is opt-in and CPU-only."""
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import importlib
|
||||
import io
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
import wave
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE = "_typed_inputs_tests"
|
||||
for suffix in ("", ".api", ".api.orchestration", ".serverless_bridge"):
|
||||
package = types.ModuleType(PACKAGE + suffix)
|
||||
package.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
|
||||
sys.modules[package.__name__] = package
|
||||
packer = importlib.import_module(PACKAGE + ".api.orchestration.serverless_typed_inputs")
|
||||
codec = importlib.import_module(PACKAGE + ".serverless_bridge.codec")
|
||||
|
||||
|
||||
def node(kind, **inputs):
|
||||
return {"class_type": kind, "inputs": inputs}
|
||||
|
||||
|
||||
def graph(kind="LoadAudio", filename="tone.wav"):
|
||||
key = "audio" if kind == "LoadAudio" else "file"
|
||||
socket = "audio" if kind == "LoadAudio" else "video"
|
||||
return {"source": node(kind, **{key: filename}),
|
||||
"consumer": node("Consumer", **{socket: ["source", 0]})}
|
||||
|
||||
|
||||
def run(workflow):
|
||||
return asyncio.run(packer.prepare_media_inputs(workflow))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def paths(tmp_path, monkeypatch):
|
||||
roots = {name: tmp_path / name for name in ("input", "output", "temp")}
|
||||
for path in roots.values():
|
||||
path.mkdir()
|
||||
def resolve(value):
|
||||
for name in roots:
|
||||
marker = " [" + name + "]"
|
||||
if value.endswith(marker):
|
||||
return str(roots[name] / value[:-len(marker)])
|
||||
return str(roots["input"] / value)
|
||||
folders = types.SimpleNamespace(get_annotated_filepath=resolve)
|
||||
for name, path in roots.items():
|
||||
setattr(folders, "get_" + name + "_directory", lambda path=path: str(path))
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", folders)
|
||||
return roots
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind,key", [("LoadAudio", "audio"), ("LoadVideo", "file")])
|
||||
@pytest.mark.parametrize("bad", [{}, None, [], {"filename": "a.wav"},
|
||||
{"audio": ["filename", 0]}, {"file": ["filename", 0]}])
|
||||
def test_rejects_non_native_or_linked_inputs_before_read(kind, key, bad, monkeypatch):
|
||||
workflow = graph(kind)
|
||||
workflow["source"]["inputs"] = bad
|
||||
read = Mock(side_effect=AssertionError("must not read"))
|
||||
monkeypatch.setattr(packer.video_codec, "_read_path", read)
|
||||
with pytest.raises(ValueError):
|
||||
run(workflow)
|
||||
read.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind,key", [("LoadAudio", "audio"), ("LoadVideo", "file")])
|
||||
@pytest.mark.parametrize("extra", ["start_time", "duration", "sample_rate", "crop", "upload"])
|
||||
def test_rejects_extra_loader_controls(kind, key, extra):
|
||||
workflow = graph(kind)
|
||||
workflow["source"]["inputs"][extra] = 0
|
||||
with pytest.raises(ValueError, match="controls|inputs"):
|
||||
run(workflow)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("slot", [1, -1, True, 0.0, "0", None])
|
||||
def test_rejects_nonzero_or_noninteger_output_slot(slot):
|
||||
workflow = graph()
|
||||
workflow["consumer"]["inputs"]["audio"] = ["source", slot]
|
||||
with pytest.raises(ValueError, match="output|slot"):
|
||||
run(workflow)
|
||||
|
||||
|
||||
def test_rejects_unused_loader_and_malformed_link():
|
||||
workflow = graph()
|
||||
workflow.pop("consumer")
|
||||
with pytest.raises(ValueError, match="consumer"):
|
||||
run(workflow)
|
||||
workflow = graph()
|
||||
workflow["consumer"]["inputs"]["audio"].append(0)
|
||||
with pytest.raises(ValueError):
|
||||
run(workflow)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key,value", [("audio", "private.wav"), ("file", "private.mp4"), ("video", "private.mov")])
|
||||
def test_unknown_media_loaders_still_use_strict_scan(key, value):
|
||||
with pytest.raises(ValueError):
|
||||
run({"source": node("CustomLoader", **{key: value})})
|
||||
|
||||
|
||||
def test_images_still_rewrite_only_the_private_graph(paths):
|
||||
payload = b"legacy staging preserves the original file bytes"
|
||||
(paths["input"] / "picture.png").write_bytes(payload)
|
||||
source = {"image": node("LoadImage", image="picture.png"),
|
||||
"export": node("MediaExportV1", images=["image", 0])}
|
||||
workflow = copy.deepcopy(source)
|
||||
images, artifacts, refs = run(workflow)
|
||||
assert artifacts == refs == []
|
||||
assert len(images) == 1
|
||||
assert base64.b64decode(images[0]["image"]) == payload
|
||||
assert workflow["image"]["inputs"]["image"] == images[0]["name"]
|
||||
assert workflow["export"] == source["export"]
|
||||
assert source["image"]["inputs"]["image"] == "picture.png"
|
||||
|
||||
|
||||
if os.environ.get("_DISTRIBUTED_TYPED_INPUT_CORE_TEST") != "1":
|
||||
def test_native_input_acceptance_cpu_subprocess():
|
||||
configured = os.environ.get("COMFYUI_SOURCE_ROOT")
|
||||
if configured is None:
|
||||
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU native input acceptance")
|
||||
core = Path(configured).expanduser().resolve()
|
||||
assert configured and (core / "comfy_extras/nodes_audio.py").is_file()
|
||||
env = dict(os.environ)
|
||||
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
env.pop("PYTHONPATH", None)
|
||||
env.pop("COMFYUI_SOURCE_ROOT", None)
|
||||
env.update(TMPDIR=str(scratch), CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="",
|
||||
PYTHONDONTWRITEBYTECODE="1", _DISTRIBUTED_TYPED_INPUT_CORE_TEST="1")
|
||||
bootstrap = """
|
||||
import sys
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
import torch
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('Input packing must not query or initialize a GPU')
|
||||
torch.cuda.is_available = lambda: False
|
||||
torch.cuda.device_count = lambda: 0
|
||||
for name in ('init', '_lazy_init', 'current_device'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
import comfy_extras.nodes_audio
|
||||
import comfy_api.latest._input_impl.video_types
|
||||
for name in ('is_available', 'device_count'):
|
||||
setattr(torch.cuda, name, forbidden)
|
||||
import pytest
|
||||
raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider']))
|
||||
"""
|
||||
result = subprocess.run([sys.executable, "-c", bootstrap, str(core), str(Path(__file__).resolve())],
|
||||
cwd=ROOT, env=env, text=True, capture_output=True, timeout=150)
|
||||
print(result.stdout)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
else:
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
from comfy_extras import nodes_audio
|
||||
from comfy_api.latest._input_impl.video_types import VideoFromFile
|
||||
|
||||
def wav_bytes(channels=2, sample_width=2, rate=22050, variant=0):
|
||||
values = np.array([-32768, -12001 + variant, -1, 0, 1, 8193, 32767] * channels,
|
||||
dtype="<i2" if sample_width == 2 else "<i4")
|
||||
buffer = io.BytesIO()
|
||||
with wave.open(buffer, "wb") as out:
|
||||
out.setnchannels(channels)
|
||||
out.setsampwidth(sample_width)
|
||||
out.setframerate(rate)
|
||||
out.writeframes(values.tobytes())
|
||||
return buffer.getvalue()
|
||||
|
||||
def unpack(artifact):
|
||||
payload = base64.b64decode(artifact["transport"]["data_base64"])
|
||||
assert artifact["byte_length"] == len(payload)
|
||||
return payload
|
||||
|
||||
@pytest.mark.parametrize("root", ["input", "output", "temp"])
|
||||
@pytest.mark.parametrize("channels,width", [(1, 2), (2, 2), (2, 4)])
|
||||
def test_exact_native_audio_and_every_consumer_ref(paths, root, channels, width):
|
||||
source = paths[root] / "tone.wav"
|
||||
source.write_bytes(wav_bytes(channels, width))
|
||||
workflow = graph(filename="tone.wav [" + root + "]")
|
||||
workflow["export"] = node("MediaExportV1", audio=["source", 0])
|
||||
workflow["consumer"]["inputs"]["second_audio"] = ["source", 0]
|
||||
before = copy.deepcopy(workflow)
|
||||
images, artifacts, refs = run(workflow)
|
||||
assert images == [] and len(artifacts) == 1
|
||||
artifact = artifacts[0]
|
||||
expected, rate = nodes_audio.load(str(source))
|
||||
actual = codec.decode_payload("AUDIO", artifact["descriptor"], unpack(artifact))
|
||||
assert actual["sample_rate"] == rate == 22050
|
||||
assert actual["waveform"].dtype == expected.dtype
|
||||
assert torch.equal(actual["waveform"].view(torch.uint8), expected.unsqueeze(0).contiguous().view(torch.uint8))
|
||||
assert refs == [{"node_id": "consumer", "input_name": "audio", "artifact_id": artifact["id"]},
|
||||
{"node_id": "consumer", "input_name": "second_audio", "artifact_id": artifact["id"]},
|
||||
{"node_id": "export", "input_name": "audio", "artifact_id": artifact["id"]}]
|
||||
assert workflow == before
|
||||
|
||||
@pytest.mark.parametrize("root", ["input", "output", "temp"])
|
||||
def test_exact_video_bytes_without_container_decode(paths, root, monkeypatch):
|
||||
payload = b"opaque native video bytes\x00\xff"
|
||||
(paths[root] / "clip.mp4").write_bytes(payload)
|
||||
workflow = graph("LoadVideo", "clip.mp4 [" + root + "]")
|
||||
workflow["export"] = node("MediaExportV1", video=["source", 0])
|
||||
before = copy.deepcopy(workflow)
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError("must not open, decode or transcode video")
|
||||
monkeypatch.setattr(av, "open", forbidden)
|
||||
monkeypatch.setattr(VideoFromFile, "get_components", forbidden)
|
||||
images, artifacts, refs = run(workflow)
|
||||
assert images == [] and len(artifacts) == 1
|
||||
artifact = artifacts[0]
|
||||
assert unpack(artifact) == payload
|
||||
assert artifact["descriptor"] == {"encoding": "native-video-v1", "representation": "file",
|
||||
"byte_length": len(payload), "start_time": 0, "duration": 0, "crop": None,
|
||||
"data": {"offset": 0, "length": len(payload)}}
|
||||
assert refs == [{"node_id": name, "input_name": "video", "artifact_id": artifact["id"]}
|
||||
for name in ("consumer", "export")]
|
||||
native = packer.video_codec.decode_video(artifact["descriptor"], unpack(artifact))
|
||||
assert vars(native)["_VideoFromFile__file"].getvalue() == payload
|
||||
assert workflow == before
|
||||
|
||||
def test_content_dedupe_kind_separation_and_mixed_images(paths):
|
||||
data = wav_bytes()
|
||||
(paths["input"] / "a.wav").write_bytes(data)
|
||||
(paths["output"] / "b.wav").write_bytes(data)
|
||||
(paths["input"] / "image.png").write_bytes(b"image bytes")
|
||||
workflow = graph(filename="a.wav")
|
||||
workflow["same"] = node("LoadAudio", audio="b.wav [output]")
|
||||
workflow["video"] = node("LoadVideo", file="a.wav")
|
||||
workflow["image"] = node("LoadImage", image="image.png")
|
||||
workflow["export"] = node("MediaExportV1", audio=["same", 0], video=["video", 0], images=["image", 0])
|
||||
before = copy.deepcopy(workflow)
|
||||
images, artifacts, refs = run(workflow)
|
||||
assert len(images) == 1 and len(artifacts) == 2 and len(refs) == 3
|
||||
assert refs[0]["artifact_id"] == refs[1]["artifact_id"] != refs[2]["artifact_id"]
|
||||
assert {a["kind"] for a in artifacts} == {"AUDIO", "VIDEO"}
|
||||
before["image"]["inputs"]["image"] = images[0]["name"]
|
||||
assert workflow == before
|
||||
second = graph(filename="b.wav [output]")
|
||||
assert run(second)[1][0]["id"] == artifacts[0]["id"]
|
||||
|
||||
@pytest.mark.parametrize("kind", ["LoadAudio", "LoadVideo"])
|
||||
def test_path_rejections_and_per_file_bounds(paths, tmp_path, monkeypatch, kind):
|
||||
outside = tmp_path / "outside.wav"
|
||||
outside.write_bytes(wav_bytes())
|
||||
(paths["input"] / "escape.wav").symlink_to(outside)
|
||||
(paths["input"] / "dirlink").symlink_to(tmp_path, target_is_directory=True)
|
||||
os.mkfifo(paths["input"] / "fifo.wav")
|
||||
values = ["../outside.wav", str(outside), "escape.wav", "dirlink/outside.wav", "fifo.wav",
|
||||
"absent.wav", "tone.wav [unknown]", "https://example.com/a.wav", "input [temp]", "C:\\a.wav"]
|
||||
for value in values:
|
||||
with pytest.raises(ValueError):
|
||||
run(graph(kind, value))
|
||||
(paths["input"] / "large.wav").write_bytes(wav_bytes())
|
||||
monkeypatch.setattr(packer.contract, "MAX_ARTIFACT_BYTES", 8)
|
||||
with pytest.raises(ValueError, match="limit|size"):
|
||||
run(graph(kind, "large.wav"))
|
||||
|
||||
def test_aggregate_decoded_budget_and_duplicate_not_double_charged(paths, monkeypatch):
|
||||
payload = wav_bytes()
|
||||
(paths["input"] / "a.wav").write_bytes(payload)
|
||||
(paths["input"] / "b.wav").write_bytes(wav_bytes(variant=1))
|
||||
workflow = graph(filename="a.wav")
|
||||
workflow["other"] = node("LoadAudio", audio="a.wav")
|
||||
workflow["export"] = node("MediaExportV1", audio=["other", 0])
|
||||
# File source bytes have a separate aggregate bound; use a longer WAV
|
||||
# whose float output dominates its tiny header.
|
||||
with wave.open(io.BytesIO(payload), "rb") as wav:
|
||||
frames = wav.readframes(wav.getnframes())
|
||||
buffer = io.BytesIO()
|
||||
with wave.open(buffer, "wb") as wav:
|
||||
wav.setparams((2, 2, 22050, 0, "NONE", "not compressed"))
|
||||
wav.writeframes(frames * 100)
|
||||
(paths["input"] / "a.wav").write_bytes(buffer.getvalue())
|
||||
artifact = run(graph(filename="a.wav"))[1][0]
|
||||
monkeypatch.setattr(packer.contract, "MAX_TOTAL_BYTES", artifact["byte_length"])
|
||||
assert len(run(workflow)[1]) == 1
|
||||
workflow["other"]["inputs"]["audio"] = "b.wav"
|
||||
with pytest.raises(ValueError, match="limit|size|budget"):
|
||||
run(workflow)
|
||||
|
||||
def test_audio_charged_before_ndarray_cat_or_conversion(monkeypatch):
|
||||
frame = types.SimpleNamespace(samples=100, sample_rate=22050,
|
||||
layout=types.SimpleNamespace(channels=[1, 2]),
|
||||
format=types.SimpleNamespace(name="s16", bytes=2),
|
||||
planes=[types.SimpleNamespace(buffer_size=400)], to_ndarray=Mock())
|
||||
stream = types.SimpleNamespace(index=0, channels=2,
|
||||
codec_context=types.SimpleNamespace(sample_rate=22050))
|
||||
container = Mock()
|
||||
container.__enter__ = Mock(return_value=container)
|
||||
container.__exit__ = Mock(return_value=False)
|
||||
container.streams = types.SimpleNamespace(audio=[stream])
|
||||
container.decode = Mock(return_value=iter([frame]))
|
||||
monkeypatch.setattr(av, "open", Mock(return_value=container))
|
||||
cat = Mock(side_effect=AssertionError("no accumulation allowed"))
|
||||
monkeypatch.setattr(torch, "cat", cat)
|
||||
with pytest.raises(ValueError, match="limit|size|budget"):
|
||||
packer._decode_audio(b"input", 799)
|
||||
frame.to_ndarray.assert_not_called()
|
||||
cat.assert_not_called()
|
||||
frame.samples = 1
|
||||
frame.planes[0].buffer_size = 800
|
||||
container.decode.return_value = iter([frame])
|
||||
with pytest.raises(ValueError, match="limit|size|budget"):
|
||||
packer._decode_audio(b"input", 799)
|
||||
frame.to_ndarray.assert_not_called()
|
||||
|
||||
def test_audio_rejects_empty_invalid_and_channel_or_rate_drift(paths, monkeypatch):
|
||||
(paths["input"] / "bad.wav").write_bytes(b"not audio")
|
||||
with pytest.raises(ValueError):
|
||||
run(graph(filename="bad.wav"))
|
||||
stream = types.SimpleNamespace(index=0, channels=2, codec_context=types.SimpleNamespace(sample_rate=22050))
|
||||
for frames in ([], [types.SimpleNamespace(samples=1, sample_rate=44100,
|
||||
layout=types.SimpleNamespace(channels=[1, 2]))],
|
||||
[types.SimpleNamespace(samples=1, sample_rate=22050,
|
||||
layout=types.SimpleNamespace(channels=[1]))]):
|
||||
container = Mock()
|
||||
container.__enter__ = Mock(return_value=container)
|
||||
container.__exit__ = Mock(return_value=False)
|
||||
container.streams = types.SimpleNamespace(audio=[stream])
|
||||
container.decode = Mock(return_value=iter(frames))
|
||||
monkeypatch.setattr(av, "open", Mock(return_value=container))
|
||||
with pytest.raises(ValueError):
|
||||
packer._decode_audio(b"input", 1000)
|
||||
@@ -0,0 +1,319 @@
|
||||
"""Offline typed submission/delivery regressions; no provider or GPU access."""
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
import socket
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
# Load once before integration fixtures snapshot/restore sys.modules. PyTorch
|
||||
# cannot be re-imported after those fixtures unload its extension-backed modules.
|
||||
import torch
|
||||
|
||||
from test_runpod_client import Response, Session
|
||||
from test_serverless_runtime import runtime, Secrets
|
||||
|
||||
client = importlib.import_module(runtime.__package__ + ".runpod_client")
|
||||
bridge = runtime.__package__.split(".")[0] + ".serverless_bridge"
|
||||
contract = importlib.import_module(bridge + ".contract")
|
||||
|
||||
|
||||
def prepared(sockets=("texts",)):
|
||||
worker = {"id": "pool:1", "pool_id": "pool", "endpoint_id": "ep1",
|
||||
"timeout_seconds": 30, "type": "serverless", "media_profile": contract.PROFILE}
|
||||
bindings = [dict(binding_id=f"binding-{i}", collector_id="collector", job_id="collector-job",
|
||||
participant_id=worker["id"], node_id=f"export-{i}", socket=name,
|
||||
kind=contract.SOCKET_KINDS[name]) for i, name in enumerate(sockets)]
|
||||
media = dict(version=1, profile=contract.PROFILE, request_id="request-1",
|
||||
artifacts=[], inputs=[], outputs=copy.deepcopy(bindings))
|
||||
graph = {"source": {"class_type": "Source", "inputs": {"value": "café雪"}}}
|
||||
for binding in bindings:
|
||||
graph[binding["node_id"]] = {"class_type": "MediaExportV1", "inputs": {
|
||||
"request_id": media["request_id"], "binding_json": json.dumps(binding),
|
||||
binding["socket"]: ["source", 0]}}
|
||||
return dict(worker=worker, workflow=graph, images=[], manifest=bindings, media=media)
|
||||
|
||||
|
||||
def response(job):
|
||||
envelope = dict(version=1, profile=contract.PROFILE, request_id=job["media"]["request_id"],
|
||||
bindings=[], artifacts=[], items=[])
|
||||
for index, binding in enumerate(job["manifest"]):
|
||||
envelope["bindings"].append(dict(binding, count=2))
|
||||
for item_index, text in enumerate(("", " \n雪")):
|
||||
aid = f"artifact-{index}-{item_index}"
|
||||
descriptor, payload = {"encoding": "utf8-v1"}, text.encode()
|
||||
if binding["kind"] != "STRING":
|
||||
codec = importlib.import_module(bridge + ".codec")
|
||||
values = {"IMAGE": torch.arange(8, dtype=torch.float64).reshape(1, 1, 2, 4),
|
||||
"MASK": torch.arange(2, dtype=torch.float16).reshape(1, 1, 2),
|
||||
"AUDIO": {"waveform": torch.arange(4, dtype=torch.float32).reshape(1, 2, 2), "sample_rate": 48000}}
|
||||
descriptor, payload = codec.encode_payload(binding["kind"], values[binding["kind"]])
|
||||
artifact = contract.artifact_from_payload(aid, binding["kind"], descriptor, payload)
|
||||
envelope["artifacts"].append(artifact)
|
||||
envelope["items"].insert(0, dict(binding_id=binding["binding_id"], item_index=item_index, artifact_id=aid))
|
||||
return {"media": envelope}
|
||||
|
||||
|
||||
class TypedClientTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_exact_typed_bytes_and_legacy_unchanged(self):
|
||||
job = prepared()
|
||||
session = Session([Response(payload={"id": "job-1"})])
|
||||
api = client.RunpodClient(session, "ep1", "test-only")
|
||||
await api.submit(job["workflow"], [], 30, media=job["media"])
|
||||
expected = {"input": {"workflow": job["workflow"], "images": [], "media": job["media"]},
|
||||
"policy": {"executionTimeout": 30000, "ttl": 30000}}
|
||||
self.assertEqual(session.calls[0][2]["data"], json.dumps(expected, allow_nan=False,
|
||||
ensure_ascii=False, separators=(",", ":")).encode())
|
||||
self.assertEqual(len(session.calls), 1)
|
||||
expected["input"].pop("media")
|
||||
legacy = json.dumps(expected, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode()
|
||||
self.assertEqual(client.serialize_request(job["workflow"], [], 30), legacy)
|
||||
self.assertEqual(client.serialize_request(job["workflow"], [], 30, media=None), legacy)
|
||||
|
||||
async def test_invalid_envelopes_never_post(self):
|
||||
for field, value in [("version", True), ("profile", "legacy-image-v1"),
|
||||
("request_id", "../bad"), ("outputs", []), ("inputs", [{}])]:
|
||||
with self.subTest(field=field):
|
||||
job = prepared()
|
||||
job["media"][field] = value
|
||||
session = Session([])
|
||||
with self.assertRaises(ValueError):
|
||||
await client.RunpodClient(session, "ep1", "test-only").submit(
|
||||
job["workflow"], [], 30, media=job["media"])
|
||||
self.assertFalse(session.calls)
|
||||
session = Session([])
|
||||
with self.assertRaises(ValueError):
|
||||
await client.RunpodClient(session, "ep1", "test-only").submit({}, [], 30, media={})
|
||||
self.assertFalse(session.calls)
|
||||
|
||||
async def test_typed_size_boundary_uses_exact_body_before_post(self):
|
||||
job = prepared()
|
||||
size = len(client.serialize_request(job["workflow"], [], 30, media=job["media"]))
|
||||
session = Session([Response(payload={"id": "boundary"})])
|
||||
api = client.RunpodClient(session, "ep1", "test-only")
|
||||
with patch.object(client, "MAX_REQUEST_BYTES", size):
|
||||
await api.submit(job["workflow"], [], 30, media=job["media"])
|
||||
job["workflow"]["source"]["inputs"]["value"] += "x"
|
||||
with self.assertRaisesRegex(ValueError, "Serialized serverless request body"):
|
||||
await api.submit(job["workflow"], [], 30, media=job["media"])
|
||||
self.assertEqual(len(session.calls), 1)
|
||||
|
||||
|
||||
class TypedRuntimeTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector-job": asyncio.Queue()},
|
||||
distributed_jobs_lock=asyncio.Lock())
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), Path(self.temp.name) / "jobs.json",
|
||||
AsyncMock(return_value=object()))
|
||||
self.manager.poll_seconds = 0
|
||||
self.job = prepared()
|
||||
self.api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"),
|
||||
status=AsyncMock(return_value={"status": "COMPLETED", "output": response(self.job)}),
|
||||
cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
self.client_patch = patch.object(runtime, "RunpodClient", return_value=self.api)
|
||||
self.client_patch.start()
|
||||
|
||||
async def asyncTearDown(self):
|
||||
await self.manager.cancel()
|
||||
self.client_patch.stop()
|
||||
self.temp.cleanup()
|
||||
|
||||
async def run_job(self):
|
||||
await self.manager.start("execution", "prompt", [self.job])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
|
||||
async def test_preflight_all_jobs_before_any_journal_or_task(self):
|
||||
for mutation in ("profile", "missing_media", "manifest", "participant", "exporter", "empty_media"):
|
||||
with self.subTest(mutation=mutation):
|
||||
bad = prepared()
|
||||
if mutation == "profile":
|
||||
bad["worker"]["media_profile"] = "legacy-image-v1"
|
||||
elif mutation == "missing_media":
|
||||
bad.pop("media")
|
||||
elif mutation == "manifest":
|
||||
bad["manifest"][0]["job_id"] = "other"
|
||||
elif mutation == "participant":
|
||||
bad["worker"]["id"] = "other-participant"
|
||||
elif mutation == "exporter":
|
||||
bad["workflow"]["export-0"]["inputs"]["request_id"] = "other"
|
||||
else:
|
||||
bad["media"] = {}
|
||||
with self.assertRaises(ValueError):
|
||||
await self.manager.start("execution", "prompt", [self.job, bad])
|
||||
self.assertFalse(self.manager.ledger.all())
|
||||
self.assertFalse(self.manager.tasks)
|
||||
self.assertFalse(self.manager.runs)
|
||||
self.api.submit.assert_not_awaited()
|
||||
|
||||
async def test_text_roundtrip_id_journal_and_payload_cleanup(self):
|
||||
async def status(_job):
|
||||
self.assertEqual(self.manager.ledger.all()[0]["job_id"], "provider-job")
|
||||
return {"status": "COMPLETED", "output": response(self.job)}
|
||||
self.api.status.side_effect = status
|
||||
await self.run_job()
|
||||
self.api.submit.assert_awaited_once_with(self.job["workflow"], [], 30, media=self.job["media"])
|
||||
queue = self.server.distributed_pending_jobs["collector-job"]
|
||||
self.assertEqual(queue.get_nowait(), {"worker_id": "pool:1", "media": {"texts": ["", " \n雪"]}, "is_last": True})
|
||||
self.assertTrue(queue.empty())
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
|
||||
run = next(iter(self.manager.runs.values()))
|
||||
for field in ("workflow", "images", "media"):
|
||||
self.assertNotIn(field, run)
|
||||
self.assertNotIn("data_base64", self.manager.ledger.path.read_text())
|
||||
self.api.cancel.assert_not_awaited()
|
||||
|
||||
async def test_all_sockets_grouped_without_native_value_conversion(self):
|
||||
self.job = prepared(tuple(contract.SOCKET_KINDS))
|
||||
native = {binding["binding_id"]: [object(), {"native": object()}] for binding in self.job["manifest"]}
|
||||
with patch.object(contract, "decode_response", AsyncMock(return_value=native)) as decode:
|
||||
await self.run_job()
|
||||
decode.assert_awaited_once()
|
||||
self.assertEqual(decode.call_args.kwargs["request_id"], "request-1")
|
||||
self.assertEqual(decode.call_args.kwargs["bindings"], self.job["manifest"])
|
||||
queue = self.server.distributed_pending_jobs["collector-job"]
|
||||
item = queue.get_nowait()
|
||||
self.assertEqual(set(item), {"worker_id", "media", "is_last"})
|
||||
self.assertEqual(set(item["media"]), set(contract.SOCKET_KINDS))
|
||||
for binding in self.job["manifest"]:
|
||||
for actual, expected in zip(item["media"][binding["socket"]], native[binding["binding_id"]]):
|
||||
self.assertIs(actual, expected)
|
||||
self.assertTrue(queue.empty())
|
||||
|
||||
async def test_identity_and_late_decode_failure_deliver_only_one_error(self):
|
||||
self.job = prepared(("texts", "images"))
|
||||
for mutation in ("request", "binding", "checksum", "descriptor"):
|
||||
with self.subTest(mutation=mutation):
|
||||
output = response(self.job)
|
||||
if mutation == "request":
|
||||
output["media"]["request_id"] = "other"
|
||||
elif mutation == "binding":
|
||||
output["media"]["bindings"][-1]["participant_id"] = "other"
|
||||
elif mutation == "checksum":
|
||||
output["media"]["artifacts"][-1]["sha256"] = "0" * 64
|
||||
else:
|
||||
output["media"]["artifacts"][-1]["descriptor"] = {}
|
||||
self.api.status.return_value = {"status": "COMPLETED", "output": output}
|
||||
await self.run_job()
|
||||
queue = self.server.distributed_pending_jobs["collector-job"]
|
||||
item = queue.get_nowait()
|
||||
self.assertIn("serverless_error", item)
|
||||
self.assertNotIn("media", item)
|
||||
self.assertTrue(queue.empty())
|
||||
self.api.cancel.assert_not_awaited()
|
||||
|
||||
async def test_closed_queue_is_not_resurrected(self):
|
||||
self.server.distributed_pending_jobs.clear()
|
||||
await self.run_job()
|
||||
self.assertFalse(self.server.distributed_pending_jobs)
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
|
||||
|
||||
async def test_real_tensor_audio_and_mask_codecs_preserve_dtype_and_shape(self):
|
||||
self.job = prepared(("texts", "images", "audio", "masks"))
|
||||
self.api.status.return_value = {"status": "COMPLETED", "output": response(self.job)}
|
||||
await self.run_job()
|
||||
result = self.server.distributed_pending_jobs["collector-job"].get_nowait()["media"]
|
||||
self.assertEqual(result["texts"], ["", " \n雪"])
|
||||
self.assertEqual(result["images"][0].dtype, torch.float64)
|
||||
self.assertEqual(result["images"][0].shape, (1, 1, 2, 4))
|
||||
self.assertEqual(result["masks"][0].dtype, torch.float16)
|
||||
self.assertEqual(result["audio"][0]["sample_rate"], 48000)
|
||||
self.assertTrue(torch.equal(result["audio"][0]["waveform"], torch.arange(4, dtype=torch.float32).reshape(1, 2, 2)))
|
||||
self.assertTrue(all(len(values) == 2 for values in result.values()))
|
||||
|
||||
async def test_run_rechecks_profile_before_submission(self):
|
||||
await self.manager.start("execution", "prompt", [self.job])
|
||||
self.job["worker"]["media_profile"] = "legacy-image-v1"
|
||||
await self.manager.wait_for_execution("execution")
|
||||
self.api.submit.assert_not_awaited()
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "FAILED")
|
||||
|
||||
async def test_multiple_collectors_receive_one_message_each(self):
|
||||
self.job = prepared(("texts", "images"))
|
||||
self.job["manifest"][-1]["job_id"] = "second-job"
|
||||
self.job["media"]["outputs"] = copy.deepcopy(self.job["manifest"])
|
||||
self.job["workflow"]["export-1"]["inputs"]["binding_json"] = json.dumps(self.job["manifest"][-1])
|
||||
self.server.distributed_pending_jobs["second-job"] = asyncio.Queue()
|
||||
self.api.status.return_value = {"status": "COMPLETED", "output": response(self.job)}
|
||||
await self.run_job()
|
||||
for job_id, name in (("collector-job", "texts"), ("second-job", "images")):
|
||||
queue = self.server.distributed_pending_jobs[job_id]
|
||||
self.assertEqual(queue.qsize(), 1)
|
||||
item = queue.get_nowait()
|
||||
self.assertEqual(set(item["media"]), {name})
|
||||
self.assertTrue(item["is_last"])
|
||||
|
||||
async def test_prestart_cancel_releases_media_and_deduplicates_errors(self):
|
||||
self.job = prepared(("texts", "images"))
|
||||
await self.manager.start("execution", "prompt", [self.job])
|
||||
await self.manager.cancel(execution_id="execution")
|
||||
self.api.submit.assert_not_awaited()
|
||||
run = next(iter(self.manager.runs.values()))
|
||||
self.assertNotIn("media", run)
|
||||
self.assertEqual(self.server.distributed_pending_jobs["collector-job"].qsize(), 1)
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
|
||||
|
||||
async def test_ambiguous_submit_keeps_reservation_without_media(self):
|
||||
self.api.submit.side_effect = runtime.AmbiguousSubmission("Unknown submission")
|
||||
await self.run_job()
|
||||
self.api.submit.assert_awaited_once()
|
||||
self.api.cancel.assert_not_awaited()
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
self.assertNotIn("media", next(iter(self.manager.runs.values())))
|
||||
|
||||
async def test_legacy_still_uses_tensor_messages_and_no_media_kwarg(self):
|
||||
self.job.pop("media")
|
||||
self.job["worker"]["media_profile"] = "legacy-image-v1"
|
||||
self.job["manifest"] = [{"job_id": "collector-job", "prefix": "dist_x", "kind": "image"}]
|
||||
with patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["a", "b"]})):
|
||||
await self.run_job()
|
||||
self.api.submit.assert_awaited_once_with(self.job["workflow"], [], 30)
|
||||
queue = self.server.distributed_pending_jobs["collector-job"]
|
||||
self.assertEqual([queue.get_nowait(), queue.get_nowait()], [
|
||||
{"worker_id": "pool:1", "tensor": "a", "image_index": 0, "is_last": False},
|
||||
{"worker_id": "pool:1", "tensor": "b", "image_index": 1, "is_last": True}])
|
||||
|
||||
|
||||
class TypedDownloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_stream_bound_cap_and_redirect_policy_without_credentials(self):
|
||||
typed = importlib.import_module(runtime.__package__ + ".serverless_typed_results")
|
||||
class DownloadResponse(Response):
|
||||
headers = {}
|
||||
async def chunks(self, _size):
|
||||
yield b"ab"
|
||||
yield b"cd"
|
||||
@property
|
||||
def content(self):
|
||||
return types.SimpleNamespace(iter_chunked=self.chunks)
|
||||
class DownloadSession(Response):
|
||||
def get(self, url, **kwargs):
|
||||
self.calls.append((url, kwargs))
|
||||
return self.response
|
||||
session = DownloadSession()
|
||||
session.calls = []
|
||||
session.response = DownloadResponse()
|
||||
with patch.object(typed.aiohttp, "TCPConnector", return_value=object()), \
|
||||
patch.object(typed.aiohttp, "ClientSession", return_value=session) as factory:
|
||||
self.assertEqual(await typed.download_public_artifact("https://example.com/a", max_bytes=contract.MAX_ARTIFACT_BYTES), b"abcd")
|
||||
self.assertFalse(factory.call_args.kwargs["trust_env"])
|
||||
self.assertNotIn("headers", factory.call_args.kwargs)
|
||||
self.assertFalse(session.calls[0][1]["allow_redirects"])
|
||||
with self.assertRaises(ValueError):
|
||||
await typed.download_public_artifact("https://example.com/a", max_bytes=3)
|
||||
session.response.status = 302
|
||||
session.response.headers = {"Location": "https://127.0.0.1/private"}
|
||||
with self.assertRaises(ValueError):
|
||||
await typed.download_public_artifact("https://example.com/a", max_bytes=4)
|
||||
for url in ("http://example.com/a", "https://127.0.0.1/a", "https://user:pass@example.com/a"):
|
||||
with self.assertRaises(ValueError):
|
||||
await typed.download_public_artifact(url, max_bytes=4)
|
||||
resolver = typed._PublicResolver()
|
||||
with patch.object(resolver, "_getaddrinfo", AsyncMock(return_value=[
|
||||
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 443))])):
|
||||
with self.assertRaises(OSError):
|
||||
await resolver.resolve("example.com", 443)
|
||||
@@ -70,6 +70,7 @@ def v3(monkeypatch):
|
||||
module(f"{PKG}.nodes.distributed_upscale", UltimateSDUpscaleDistributed=object)
|
||||
module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime)
|
||||
module(f"{PKG}.nodes.image_list", DistributedImageListStart=object, DistributedImageListEnd=object)
|
||||
module(f"{PKG}.nodes.media_collector", DistributedMediaCollector=object)
|
||||
module("comfy", __path__=[])
|
||||
module("comfy.samplers", KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["normal"]))
|
||||
module("comfy_api", __path__=[])
|
||||
|
||||
@@ -84,6 +84,14 @@ def validate_serverless_worker(worker: dict) -> dict:
|
||||
if "enabled" in result and type(result["enabled"]) is not bool:
|
||||
raise ServerlessConfigError("enabled must be a boolean")
|
||||
|
||||
# Omission preserves old saved records and means legacy-image-v1. Selecting
|
||||
# comfy-media-v1 declares an installed endpoint bridge, not health detection;
|
||||
# its schema/representations/limits are fixed by serverless_bridge/contract.py.
|
||||
if "media_profile" in result:
|
||||
profile = result["media_profile"]
|
||||
if not isinstance(profile, str) or profile not in ("legacy-image-v1", "comfy-media-v1"):
|
||||
raise ServerlessConfigError("media_profile must be legacy-image-v1 or comfy-media-v1")
|
||||
|
||||
for field, (minimum, maximum, default) in _NUMERIC_FIELDS.items():
|
||||
value = result.get(field, default)
|
||||
if type(value) is not int or not minimum <= value <= maximum:
|
||||
|
||||
@@ -7,7 +7,7 @@ import { createWorkerSettingsForm } from "../ui/settingsForm.js";
|
||||
import { DistributedUI } from "../ui.js";
|
||||
import { checkWorkerStatus, formatServerlessStatus, getWorkerUrl, launchWorker, stopWorker, viewWorkerLog } from "../workerLifecycle.js";
|
||||
import { handleClearMemory, handleInterruptWorkers } from "../workerUtils.js";
|
||||
import { isRemoteWorker, normalizeRunpodEndpointId, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js";
|
||||
import { cancelWorkerSettings, isRemoteWorker, normalizeRunpodEndpointId, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js";
|
||||
import { performPreflightCheck } from "../executionUtils.js";
|
||||
|
||||
function createDocumentMock() {
|
||||
@@ -44,6 +44,7 @@ function createDocumentMock() {
|
||||
document: {
|
||||
createElement: makeElement,
|
||||
getElementById: (id) => elements.get(id) || null,
|
||||
querySelector: () => null,
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -125,19 +126,94 @@ describe("Runpod serverless worker panel", () => {
|
||||
expect(elements.get(`${field}-group-saved`).style.display).toBe("none");
|
||||
}
|
||||
expect(elements.get("runpod-text-saved").style.display).toBe("none");
|
||||
for (const field of ["endpoint-id", "api-key", "clear-api-key", "jobs-per-run", "timeout-seconds", "serverless-test"]) {
|
||||
for (const field of ["endpoint-id", "media-profile", "api-key", "clear-api-key", "jobs-per-run", "timeout-seconds", "serverless-test"]) {
|
||||
expect(elements.get(`${field}-group-saved`).style.display).toBe("flex");
|
||||
}
|
||||
// The separate concurrency setting was removed; a legacy value is ignored.
|
||||
expect(elements.has("concurrency-saved")).toBe(false);
|
||||
expect(elements.has("concurrency-group-saved")).toBe(false);
|
||||
expect(elements.get("endpoint-id-saved").value).toBe("ep_123");
|
||||
expect(elements.get("media-profile-saved").value).toBe("legacy-image-v1");
|
||||
expect(elements.get("jobs-per-run-saved").value).toBe("2");
|
||||
expect(elements.get("timeout-seconds-saved").value).toBe("900");
|
||||
expect(elements.get("api-key-saved").value).toBe("");
|
||||
expect(elements.get("api-key-saved").placeholder).toMatch(/leave blank/i);
|
||||
});
|
||||
|
||||
it("renders explicit endpoint profiles with installation and no-auto-detection guidance", () => {
|
||||
createWorkerSettingsForm(createFormUi(), {}, {
|
||||
id: "profile", name: "Runpod", type: "serverless", media_profile: "comfy-media-v1",
|
||||
});
|
||||
const select = elements.get("media-profile-profile");
|
||||
const group = elements.get("media-profile-group-profile");
|
||||
expect(group.children[0].textContent).toBe("Endpoint media profile");
|
||||
expect(group.children[0].htmlFor).toBe(select.id);
|
||||
expect(select.children.map((option) => [option.value, option.textContent])).toEqual([
|
||||
["legacy-image-v1", "Legacy images (existing handler)"],
|
||||
["comfy-media-v1", "Typed media v1 (requires endpoint bridge)"],
|
||||
]);
|
||||
expect(select.value).toBe("comfy-media-v1");
|
||||
expect(select.title).toMatch(/after installing.*endpoint bridge/i);
|
||||
expect(select.title).toMatch(/health checks do not auto-detect/i);
|
||||
expect(select.title).toMatch(/schema.*formats and limits are fixed/i);
|
||||
const type = elements.get("worker-type-profile");
|
||||
type.value = "remote";
|
||||
type.dispatchEvent();
|
||||
expect(group.style.display).toBe("none");
|
||||
type.value = "serverless";
|
||||
type.dispatchEvent();
|
||||
expect(group.style.display).toBe("flex");
|
||||
expect(select.value).toBe("comfy-media-v1");
|
||||
});
|
||||
|
||||
it.each(["legacy-image-v1", "comfy-media-v1"])("round-trips an explicit %s selection through the save payload and reload", async (profile) => {
|
||||
const worker = { id: "roundtrip", name: "Runpod", type: "serverless", endpoint_id: "ep_123", enabled: false };
|
||||
let persisted;
|
||||
const extension = makeSaveExtension(worker, vi.fn(async (id, payload) => {
|
||||
persisted = JSON.parse(JSON.stringify({ id, ...payload }));
|
||||
return {};
|
||||
}));
|
||||
const ui = createFormUi();
|
||||
const form = createWorkerSettingsForm(ui, extension, worker);
|
||||
elements.get("media-profile-roundtrip").value = profile;
|
||||
const saveButton = form.children.at(-1).children.find((button) => button.textContent === "Save");
|
||||
await saveButton.onclick();
|
||||
expect(persisted.media_profile).toBe(profile);
|
||||
expect(worker.media_profile).toBe(profile);
|
||||
expect(persisted).not.toHaveProperty("api_key");
|
||||
expect(persisted).not.toHaveProperty("concurrency");
|
||||
createWorkerSettingsForm(ui, extension, persisted);
|
||||
expect(elements.get("media-profile-roundtrip").value).toBe(profile);
|
||||
});
|
||||
|
||||
it.each([undefined, "comfy-media-v1"])("restores the saved profile on cancel (%s)", (profile) => {
|
||||
const worker = { id: "cancel", name: "Runpod", type: "serverless", media_profile: profile };
|
||||
const extension = makeSaveExtension(worker);
|
||||
createWorkerSettingsForm(createFormUi(), extension, worker);
|
||||
elements.get("media-profile-cancel").value = profile ? "legacy-image-v1" : "comfy-media-v1";
|
||||
cancelWorkerSettings(extension, worker.id);
|
||||
expect(elements.get("media-profile-cancel").value).toBe(profile ?? "legacy-image-v1");
|
||||
expect(extension.api.updateWorker).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("keeps the saved profile after a busy-guard error and rejects invalid selections", async () => {
|
||||
const worker = { id: "busy", name: "Runpod", type: "serverless", endpoint_id: "ep_123", media_profile: "comfy-media-v1" };
|
||||
const extension = makeSaveExtension(worker, vi.fn().mockRejectedValue(new Error("Finish or reconcile unresolved serverless jobs")));
|
||||
createWorkerSettingsForm(createFormUi(), extension, worker);
|
||||
elements.get("media-profile-busy").value = "legacy-image-v1";
|
||||
await saveWorkerSettings(extension, worker.id);
|
||||
expect(worker.media_profile).toBe("comfy-media-v1");
|
||||
expect(extension.app.extensionManager.toast.add).toHaveBeenLastCalledWith(
|
||||
expect.objectContaining({ summary: "Save Failed", detail: expect.stringMatching(/unresolved/) }),
|
||||
);
|
||||
elements.get("media-profile-busy").value = "auto";
|
||||
await saveWorkerSettings(extension, worker.id);
|
||||
expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
|
||||
expect(extension.app.extensionManager.toast.add).toHaveBeenLastCalledWith(
|
||||
expect.objectContaining({ summary: "Validation Error", detail: expect.stringMatching(/media profile/) }),
|
||||
);
|
||||
});
|
||||
|
||||
it("preserves a saved cloud port on first render and applies defaults only on type changes", () => {
|
||||
vi.spyOn(document, "getElementById").mockReturnValue(null);
|
||||
createWorkerSettingsForm(createFormUi(), {}, {
|
||||
@@ -218,6 +294,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
elements.set("name-w2", { value: "Runpod worker" });
|
||||
elements.set("worker-type-w2", { value: "serverless" });
|
||||
elements.set("endpoint-id-w2", { value: "ep_123" });
|
||||
elements.set("media-profile-w2", { value: "legacy-image-v1" });
|
||||
elements.set("api-key-w2", { value: "secret-token" });
|
||||
elements.set("clear-api-key-w2", { checked: false });
|
||||
elements.set("jobs-per-run-w2", { value: "4" });
|
||||
@@ -227,6 +304,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
|
||||
expect(extension.api.updateWorker).toHaveBeenCalledWith("w2", {
|
||||
name: "Runpod worker", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
|
||||
media_profile: "legacy-image-v1",
|
||||
api_key: "secret-token", jobs_per_run: 4, timeout_seconds: 2400, enabled: false,
|
||||
});
|
||||
expect(worker).toMatchObject({ type: "serverless", provider: "runpod", enabled: false, api_key_configured: true });
|
||||
@@ -244,6 +322,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
elements.set("name-legacy", { value: "Runpod" });
|
||||
elements.set("worker-type-legacy", { value: "serverless" });
|
||||
elements.set("endpoint-id-legacy", { value: "ep_123" });
|
||||
elements.set("media-profile-legacy", { value: "legacy-image-v1" });
|
||||
elements.set("api-key-legacy", { value: "" });
|
||||
elements.set("clear-api-key-legacy", { checked: false });
|
||||
elements.set("jobs-per-run-legacy", { value: "2" });
|
||||
@@ -278,6 +357,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
elements.set("name-url", { value: "Runpod" });
|
||||
elements.set("worker-type-url", { value: "serverless" });
|
||||
elements.set("endpoint-id-url", { value: "https://api.runpod.ai/v2/abc123xyz/runsync" });
|
||||
elements.set("media-profile-url", { value: "legacy-image-v1" });
|
||||
elements.set("api-key-url", { value: "" });
|
||||
elements.set("clear-api-key-url", { checked: false });
|
||||
elements.set("jobs-per-run-url", { value: "1" });
|
||||
@@ -294,7 +374,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
});
|
||||
|
||||
it("asks before a type change deletes the saved API key, and cancels cleanly", async () => {
|
||||
const worker = { id: "conv", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: true };
|
||||
const worker = { id: "conv", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: true, media_profile: "comfy-media-v1" };
|
||||
const extension = makeSaveExtension(worker);
|
||||
elements.set("name-conv", { value: "Runpod" });
|
||||
elements.set("worker-type-conv", { value: "remote" });
|
||||
@@ -307,11 +387,13 @@ describe("Runpod serverless worker panel", () => {
|
||||
expect(confirmSpy).toHaveBeenCalledWith(expect.stringMatching(/saved Runpod API key will be deleted/));
|
||||
expect(extension.api.updateWorker).not.toHaveBeenCalled();
|
||||
expect(worker.type).toBe("serverless");
|
||||
expect(worker.media_profile).toBe("comfy-media-v1");
|
||||
|
||||
confirmSpy.mockReturnValue(true);
|
||||
await saveWorkerSettings(extension, "conv");
|
||||
expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
|
||||
expect(extension.api.updateWorker.mock.calls[0][1].type).toBe("remote");
|
||||
expect(worker).not.toHaveProperty("media_profile");
|
||||
} finally {
|
||||
delete globalThis.confirm;
|
||||
}
|
||||
@@ -341,6 +423,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
elements.set("name-w-blank", { value: "Runpod" });
|
||||
elements.set("worker-type-w-blank", { value: "serverless" });
|
||||
elements.set("endpoint-id-w-blank", { value: "ep_123" });
|
||||
elements.set("media-profile-w-blank", { value: "legacy-image-v1" });
|
||||
elements.set("api-key-w-blank", { value: "" });
|
||||
elements.set("clear-api-key-w-blank", { checked: false });
|
||||
elements.set("jobs-per-run-w-blank", { value: "1" });
|
||||
@@ -361,6 +444,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
elements.set("name-w3", { value: "Runpod" });
|
||||
elements.set("worker-type-w3", { value: "serverless" });
|
||||
elements.set("endpoint-id-w3", { value: "ep_123" });
|
||||
elements.set("media-profile-w3", { value: "legacy-image-v1" });
|
||||
elements.set("api-key-w3", { value: "" });
|
||||
elements.set("clear-api-key-w3", { checked: true });
|
||||
elements.set("jobs-per-run-w3", { value: "1" });
|
||||
@@ -370,6 +454,7 @@ describe("Runpod serverless worker panel", () => {
|
||||
|
||||
expect(extension.api.updateWorker).toHaveBeenCalledWith("w3", {
|
||||
name: "Runpod", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
|
||||
media_profile: "legacy-image-v1",
|
||||
clear_api_key: true, jobs_per_run: 1, timeout_seconds: 1800, enabled: true,
|
||||
});
|
||||
expect(worker.api_key_configured).toBe(false);
|
||||
@@ -527,9 +612,6 @@ describe("Runpod serverless worker panel", () => {
|
||||
});
|
||||
|
||||
it("tests only the saved Runpod endpoint and reports the backend message", async () => {
|
||||
const button = document.createElement("button");
|
||||
button.id = "test-serverless-w6";
|
||||
button.textContent = "Test Saved Endpoint";
|
||||
const extension = {
|
||||
api: {
|
||||
testServerlessWorker: vi.fn().mockResolvedValue({ ok: true, message: "Endpoint reachable" }),
|
||||
@@ -537,9 +619,15 @@ describe("Runpod serverless worker panel", () => {
|
||||
},
|
||||
app: { extensionManager: { toast: { add: vi.fn() } } },
|
||||
};
|
||||
const worker = { id: "w6", name: "Runpod", type: "serverless" };
|
||||
extension.config = { workers: [worker] };
|
||||
createWorkerSettingsForm(createFormUi(), extension, worker);
|
||||
const button = elements.get("test-serverless-w6");
|
||||
|
||||
await expect(testServerlessWorker(extension, "w6")).resolves.toEqual({ ok: true, message: "Endpoint reachable" });
|
||||
|
||||
expect(worker).not.toHaveProperty("media_profile");
|
||||
expect(elements.get("media-profile-w6").value).toBe("legacy-image-v1");
|
||||
expect(extension.api.testServerlessWorker).toHaveBeenCalledWith("w6");
|
||||
expect(extension.api.queueDistributed).not.toHaveBeenCalled();
|
||||
expect(extension.app.extensionManager.toast.add).toHaveBeenCalledWith(expect.objectContaining({
|
||||
|
||||
@@ -4,6 +4,9 @@ import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings,
|
||||
export const IMAGES_PER_RUN_TOOLTIP =
|
||||
"Separate Runpod jobs submitted each run, all at once. Runpod's max workers decides how many GPUs actually run them.";
|
||||
|
||||
const MEDIA_PROFILE_TOOLTIP =
|
||||
"Select Typed media v1 only after installing the ComfyUI-Distributed endpoint bridge. This declares bridge support; endpoint health checks do not auto-detect it. The v1 schema, media formats and limits are fixed.";
|
||||
|
||||
export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
const form = document.createElement("div");
|
||||
form.style.cssText = "display: flex; flex-direction: column; gap: 8px;";
|
||||
@@ -59,6 +62,7 @@ export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
const portInput = portGroup.input;
|
||||
const serverlessGroups = [
|
||||
endpointGroup.group,
|
||||
mediaProfileGroup,
|
||||
apiKeyGroup.group,
|
||||
clearApiKeyGroup,
|
||||
jobsPerRunGroup.group,
|
||||
@@ -137,6 +141,32 @@ export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
endpointGroup.group.style.display = "none";
|
||||
form.appendChild(endpointGroup.group);
|
||||
|
||||
const mediaProfileGroup = document.createElement("div");
|
||||
mediaProfileGroup.id = `media-profile-group-${worker.id}`;
|
||||
mediaProfileGroup.style.cssText = "display: none; flex-direction: column; gap: 4px;";
|
||||
mediaProfileGroup.title = MEDIA_PROFILE_TOOLTIP;
|
||||
const mediaProfileLabel = document.createElement("label");
|
||||
mediaProfileLabel.htmlFor = `media-profile-${worker.id}`;
|
||||
mediaProfileLabel.textContent = "Endpoint media profile";
|
||||
mediaProfileLabel.style.cssText = typeLabel.style.cssText;
|
||||
const mediaProfileSelect = document.createElement("select");
|
||||
mediaProfileSelect.id = mediaProfileLabel.htmlFor;
|
||||
mediaProfileSelect.style.cssText = typeSelect.style.cssText;
|
||||
mediaProfileSelect.title = MEDIA_PROFILE_TOOLTIP;
|
||||
for (const [value, label] of [
|
||||
["legacy-image-v1", "Legacy images (existing handler)"],
|
||||
["comfy-media-v1", "Typed media v1 (requires endpoint bridge)"],
|
||||
]) {
|
||||
const option = document.createElement("option");
|
||||
option.value = value;
|
||||
option.textContent = label;
|
||||
mediaProfileSelect.appendChild(option);
|
||||
}
|
||||
mediaProfileSelect.value = worker.media_profile ?? "legacy-image-v1";
|
||||
mediaProfileGroup.appendChild(mediaProfileLabel);
|
||||
mediaProfileGroup.appendChild(mediaProfileSelect);
|
||||
form.appendChild(mediaProfileGroup);
|
||||
|
||||
const apiKeyGroup = ui.createFormGroup("Runpod API Key (write-only):", "", `api-key-${worker.id}`, "password", worker.api_key_configured ? "Leave blank to keep saved API key" : "Runpod API key");
|
||||
apiKeyGroup.group.id = `api-key-group-${worker.id}`;
|
||||
apiKeyGroup.input.autocomplete = "new-password";
|
||||
|
||||
@@ -61,6 +61,7 @@ function showServerlessValidationError(extension, detail) {
|
||||
|
||||
async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
|
||||
const endpointId = normalizeRunpodEndpointId(document.getElementById(`endpoint-id-${workerId}`).value);
|
||||
const mediaProfile = document.getElementById(`media-profile-${workerId}`).value;
|
||||
const apiKey = document.getElementById(`api-key-${workerId}`).value || "";
|
||||
const clearApiKey = Boolean(document.getElementById(`clear-api-key-${workerId}`).checked);
|
||||
const jobsPerRun = Number(document.getElementById(`jobs-per-run-${workerId}`).value);
|
||||
@@ -81,6 +82,10 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
|
||||
);
|
||||
return;
|
||||
}
|
||||
if (!["legacy-image-v1", "comfy-media-v1"].includes(mediaProfile)) {
|
||||
showServerlessValidationError(extension, "Select a supported endpoint media profile");
|
||||
return;
|
||||
}
|
||||
if (!Number.isInteger(jobsPerRun) || jobsPerRun < 1 || jobsPerRun > 16) {
|
||||
showServerlessValidationError(extension, "Images per run must be a whole number from 1 to 16");
|
||||
return;
|
||||
@@ -95,6 +100,7 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
|
||||
type: "serverless",
|
||||
provider: "runpod",
|
||||
endpoint_id: endpointId,
|
||||
media_profile: mediaProfile,
|
||||
...(apiKey.length > 0 && !clearApiKey ? { api_key: apiKey } : {}),
|
||||
...(clearApiKey ? { clear_api_key: true } : {}),
|
||||
jobs_per_run: jobsPerRun,
|
||||
@@ -111,6 +117,7 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
|
||||
worker.type = "serverless";
|
||||
worker.provider = "runpod";
|
||||
worker.endpoint_id = endpointId;
|
||||
worker.media_profile = mediaProfile;
|
||||
worker.jobs_per_run = jobsPerRun;
|
||||
// Replaced by a fixed backend limit; drop any value loaded from an older config.
|
||||
delete worker.concurrency;
|
||||
@@ -310,6 +317,7 @@ export async function saveWorkerSettings(extension, workerId) {
|
||||
// Update local config
|
||||
worker.name = name.trim();
|
||||
worker.type = workerType;
|
||||
delete worker.media_profile;
|
||||
if (isRemote) {
|
||||
worker.host = hostTrimmed;
|
||||
delete worker.cuda_device;
|
||||
@@ -366,6 +374,8 @@ export function cancelWorkerSettings(extension, workerId) {
|
||||
document.getElementById(`args-${workerId}`).value = worker.extra_args || "";
|
||||
const endpointInput = document.getElementById(`endpoint-id-${workerId}`);
|
||||
if (endpointInput) endpointInput.value = worker.endpoint_id || "";
|
||||
const mediaProfileInput = document.getElementById(`media-profile-${workerId}`);
|
||||
if (mediaProfileInput) mediaProfileInput.value = worker.media_profile ?? "legacy-image-v1";
|
||||
const apiKeyInput = document.getElementById(`api-key-${workerId}`);
|
||||
if (apiKeyInput) {
|
||||
apiKeyInput.value = "";
|
||||
|
||||
Reference in New Issue
Block a user