790 lines
37 KiB
Python
790 lines
37 KiB
Python
"""Pure HTTP executor for media-v1 requests against a local ComfyUI instance.
|
|
|
|
No ComfyUI, tensor, or provider imports. ``download(url, *, max_bytes)`` must
|
|
provide public-HTTPS/redirect/DNS protection and bounded streaming. Optional
|
|
``publish(metadata, payload)`` returns a contract-valid, expiring HTTPS transport.
|
|
|
|
Only identified invocation-owned files are deleted. Optional trusted local
|
|
``temp_dir`` enables exporter cleanup after terminal history and collection.
|
|
Without it, exporter retention belongs to the ComfyUI temp-directory lifecycle.
|
|
The pinned native /interrupt route is not atomic: running-prompt cancellation
|
|
fails closed and retains files instead of risking an unrelated successor prompt.
|
|
"""
|
|
import asyncio
|
|
import base64
|
|
import binascii
|
|
import copy
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import re
|
|
import stat
|
|
import uuid
|
|
from pathlib import Path
|
|
from urllib.parse import quote, urlsplit
|
|
|
|
import aiohttp
|
|
|
|
from .contract import (
|
|
BINDING_FIELDS, MAX_ARTIFACT_BYTES, MAX_ARTIFACTS, MAX_TOTAL_BYTES,
|
|
PROFILE, VERSION, artifact_from_payload, identifier, resolve_artifact,
|
|
validate_artifact, validate_bindings, validate_json, validate_response,
|
|
)
|
|
|
|
MAX_RESPONSE_JSON_BYTES = 10_000_000
|
|
MAX_HTTP_JSON_BYTES = 32 * 1024 * 1024
|
|
POLL_INTERVAL_SECONDS = 0.25
|
|
CANCEL_GRACE_SECONDS = 10.0
|
|
DRAIN_GRACE_SECONDS = 10.0
|
|
HTTP_TIMEOUT_SECONDS = 15.0
|
|
_UNCERTAIN_WORK = set()
|
|
GET_ATTEMPTS = 3
|
|
_IMPORT_SLOTS = {"IMAGE": 0, "AUDIO": 1, "MASK": 2, "STRING": 3, "VIDEO": 4}
|
|
# Used ONLY to run the shared manifest validator before /view. Never fetched.
|
|
_VALIDATION_TRANSPORT = {"type": "url", "url": "https://media-validation.invalid/artifact",
|
|
"expires_at": "9999-01-01T00:00:00Z"}
|
|
|
|
|
|
class MediaExecutionUncertain(RuntimeError):
|
|
"""The local prompt may still execute; staged files were deliberately retained."""
|
|
|
|
|
|
class _HTTPError(RuntimeError):
|
|
def __init__(self, status):
|
|
super().__init__(f"Local Comfy HTTP request failed (status {status})")
|
|
self.status = status
|
|
|
|
|
|
def _scope_id(value):
|
|
identifier(value)
|
|
if any(part in ("", ".", "..") for part in value.split("/")):
|
|
raise ValueError("Unsafe media scope identifier")
|
|
return value
|
|
|
|
|
|
def _json_bytes(value):
|
|
return json.dumps(value, allow_nan=False).encode("utf-8")
|
|
|
|
|
|
def _link(value, workflow):
|
|
if (not isinstance(value, list) or len(value) != 2 or
|
|
not isinstance(value[0], str) or value[0] not in workflow or
|
|
type(value[1]) is not int or value[1] < 0):
|
|
raise ValueError("Invalid workflow socket link")
|
|
return value[0]
|
|
|
|
|
|
def _validate_request(input_data):
|
|
if (not isinstance(input_data, dict) or not {"workflow", "media"} <= set(input_data)
|
|
or set(input_data) - {"workflow", "media", "images"}):
|
|
raise ValueError("Invalid media request envelope")
|
|
media = input_data["media"]
|
|
if (not isinstance(media, dict) or set(media) !=
|
|
{"version", "profile", "request_id", "artifacts", "inputs", "outputs"}):
|
|
raise ValueError("Invalid media request controls")
|
|
if type(media["version"]) is not int or media["version"] != VERSION or media["profile"] != PROFILE:
|
|
raise ValueError("Unsupported media request version/profile")
|
|
_scope_id(media["request_id"])
|
|
bindings = validate_bindings(media["outputs"])
|
|
targets = []
|
|
for binding in bindings.values():
|
|
for field in BINDING_FIELDS - {"socket", "kind"}:
|
|
_scope_id(binding[field])
|
|
targets.append(binding["node_id"])
|
|
if len(set(targets)) != len(targets):
|
|
raise ValueError("Duplicate media exporter node")
|
|
|
|
workflow = input_data["workflow"]
|
|
if not isinstance(workflow, dict) or not workflow:
|
|
raise ValueError("Invalid API workflow")
|
|
validate_json(workflow, max_bytes=16 * 1024 * 1024)
|
|
dependencies = {}
|
|
for node_id, node in workflow.items():
|
|
_scope_id(node_id)
|
|
if (not isinstance(node, dict) or not isinstance(node.get("class_type"), str)
|
|
or not node["class_type"] or not isinstance(node.get("inputs"), dict)):
|
|
raise ValueError("Invalid API workflow node")
|
|
if node["class_type"] == "MediaImportV1":
|
|
raise ValueError("Request may not supply its own media import filenames")
|
|
if node["class_type"] == "MediaExportV1" and node_id not in targets:
|
|
raise ValueError("Undeclared media exporter")
|
|
dependencies[node_id] = set()
|
|
for name, value in node["inputs"].items():
|
|
_scope_id(name)
|
|
if isinstance(value, list):
|
|
dependencies[node_id].add(_link(value, workflow))
|
|
for binding in bindings.values():
|
|
node = workflow.get(binding["node_id"])
|
|
if node is None or node["class_type"] != "MediaExportV1":
|
|
raise ValueError("Missing declared MediaExportV1 node")
|
|
controls = node["inputs"]
|
|
if set(controls) != {"request_id", "binding_json", binding["socket"]}:
|
|
raise ValueError("Media exporter must have exactly its declared socket and controls")
|
|
if controls["request_id"] != media["request_id"]:
|
|
raise ValueError("Media exporter request identity mismatch")
|
|
text = controls["binding_json"]
|
|
if not isinstance(text, str) or len(text) > 128 * 1024:
|
|
raise ValueError("Invalid media exporter binding JSON")
|
|
try:
|
|
declared = json.loads(text)
|
|
except (ValueError, RecursionError):
|
|
raise ValueError("Invalid media exporter binding JSON") from None
|
|
if declared != binding:
|
|
raise ValueError("Media exporter binding identity mismatch")
|
|
_link(controls[binding["socket"]], workflow)
|
|
# The execution roots are exactly the declared exporters, not every class
|
|
# with OUTPUT_NODE=True. Required upstream output utilities may remain.
|
|
visited = set()
|
|
pending = list(targets)
|
|
while pending:
|
|
node_id = pending.pop()
|
|
if node_id not in visited:
|
|
visited.add(node_id)
|
|
pending.extend(dependencies[node_id])
|
|
if visited != set(workflow):
|
|
raise ValueError("Workflow contains nodes outside declared output roots")
|
|
if any(set(targets) & sources for sources in dependencies.values()):
|
|
raise ValueError("A media exporter cannot be an upstream dependency")
|
|
# Kahn traversal bounds cycle checking without relying on recursion depth.
|
|
counts = {node: len(sources) for node, sources in dependencies.items()}
|
|
consumers = {node: [] for node in workflow}
|
|
for node, sources in dependencies.items():
|
|
for source in sources:
|
|
consumers[source].append(node)
|
|
pending = [node for node, count in counts.items() if not count]
|
|
processed = 0
|
|
while pending:
|
|
processed += 1
|
|
for consumer in consumers[pending.pop()]:
|
|
counts[consumer] -= 1
|
|
if not counts[consumer]:
|
|
pending.append(consumer)
|
|
if processed != len(workflow):
|
|
raise ValueError("Workflow contains a dependency cycle")
|
|
|
|
artifacts = media["artifacts"]
|
|
if not isinstance(artifacts, list) or len(artifacts) > MAX_ARTIFACTS:
|
|
raise ValueError("Invalid input artifacts count")
|
|
indexed, total = {}, 0
|
|
for artifact in artifacts:
|
|
validate_artifact(artifact)
|
|
aid = _scope_id(artifact["id"])
|
|
if aid in indexed:
|
|
raise ValueError("Duplicate input artifact")
|
|
indexed[aid] = artifact
|
|
total += artifact["byte_length"]
|
|
inputs = media["inputs"]
|
|
if not isinstance(inputs, list) or len(inputs) > 4096:
|
|
raise ValueError("Invalid input references")
|
|
destinations, used = set(), set()
|
|
exporters = {binding["node_id"]: binding for binding in bindings.values()}
|
|
for ref in inputs:
|
|
if not isinstance(ref, dict) or set(ref) != {"node_id", "input_name", "artifact_id"}:
|
|
raise ValueError("Invalid media input reference")
|
|
for value in ref.values():
|
|
_scope_id(value)
|
|
node_id, name, aid = ref["node_id"], ref["input_name"], ref["artifact_id"]
|
|
if (node_id not in workflow or name not in workflow[node_id]["inputs"]
|
|
or aid not in indexed or (node_id, name) in destinations):
|
|
raise ValueError("Unknown or duplicate media input reference")
|
|
binding = exporters.get(node_id)
|
|
if binding is not None and (name != binding["socket"] or indexed[aid]["kind"] != binding["kind"]):
|
|
raise ValueError("Exporter input must match its declared media socket and kind")
|
|
destinations.add((node_id, name))
|
|
used.add(aid)
|
|
if used != set(indexed):
|
|
raise ValueError("Every input artifact must be used")
|
|
|
|
images = input_data.get("images", [])
|
|
if not isinstance(images, list) or len(images) > MAX_ARTIFACTS:
|
|
raise ValueError("Invalid legacy images count")
|
|
names = set()
|
|
legacy = []
|
|
for image in images:
|
|
if not isinstance(image, dict) or set(image) != {"name", "image"}:
|
|
raise ValueError("Invalid legacy image")
|
|
name, encoded = image["name"], image["image"]
|
|
if (not isinstance(name, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_. -]{0,254}", name)
|
|
or name in names or not isinstance(encoded, str)):
|
|
raise ValueError("Invalid or duplicate legacy image basename")
|
|
names.add(name)
|
|
if len(encoded) > 4 * ((MAX_ARTIFACT_BYTES + 2) // 3):
|
|
raise ValueError("Legacy image exceeds byte limit")
|
|
# Exact decoded length, including padding, before allocating bytes.
|
|
size = (len(encoded) // 4) * 3 - (len(encoded) - len(encoded.rstrip("=")))
|
|
if len(encoded) % 4 or size < 0 or size > MAX_ARTIFACT_BYTES:
|
|
raise ValueError("Invalid legacy image base64 length")
|
|
total += size
|
|
matches = [node_id for node_id, node in workflow.items()
|
|
if node["class_type"] == "LoadImage" and node["inputs"].get("image") == name]
|
|
if not matches or any((node_id, "image") in destinations for node_id in matches):
|
|
raise ValueError("Unused or conflicting legacy image binding")
|
|
legacy.append((encoded, size, matches))
|
|
if total > MAX_TOTAL_BYTES:
|
|
raise ValueError("Media inputs exceed total byte limit")
|
|
return copy.deepcopy(workflow), copy.deepcopy(media), legacy, targets
|
|
|
|
|
|
class _Staging:
|
|
"""Exclusive random files under a hashed namespace; never delete siblings."""
|
|
def __init__(self, input_dir, request_id):
|
|
self.root = Path(input_dir).resolve()
|
|
self.key = hashlib.sha256(request_id.encode()).hexdigest()
|
|
self.relative = f"distributed-media-v1/{self.key}"
|
|
self.names = []
|
|
self.fd = None
|
|
|
|
def _open(self):
|
|
self.root.mkdir(parents=True, exist_ok=True)
|
|
directory = self.root
|
|
for component in ("distributed-media-v1", self.key):
|
|
directory = directory / component
|
|
if directory.is_symlink():
|
|
raise ValueError("Unsafe media staging directory symlink")
|
|
directory.mkdir(exist_ok=True, mode=0o700)
|
|
if not directory.resolve().is_relative_to(self.root):
|
|
raise ValueError("Unsafe media staging directory")
|
|
self.fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
|
|
|
|
def write(self, payload):
|
|
if self.fd is None:
|
|
self._open()
|
|
name = uuid.uuid4().hex + ".bin"
|
|
fd = os.open(name, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600, dir_fd=self.fd)
|
|
self.names.append(name)
|
|
with os.fdopen(fd, "wb") as stream:
|
|
stream.write(payload)
|
|
return f"{self.relative}/{name}"
|
|
|
|
def close(self, *, cleanup):
|
|
if self.fd is not None:
|
|
try:
|
|
if cleanup:
|
|
for name in self.names:
|
|
try:
|
|
os.unlink(name, dir_fd=self.fd)
|
|
except FileNotFoundError:
|
|
pass
|
|
finally:
|
|
os.close(self.fd)
|
|
self.fd = None
|
|
if cleanup:
|
|
# Never recursively remove a shared namespace or another run.
|
|
for directory in (self.root / self.relative, self.root / "distributed-media-v1"):
|
|
try:
|
|
directory.rmdir()
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
class _OwnedExports:
|
|
"""Best-effort retirement of validated terminal-history files, never a sweep.
|
|
|
|
Walk every directory component with openat/O_NOFOLLOW, then pin the directory
|
|
fd and each regular file's identity. Missing, linked, replaced, or unidentifiable
|
|
files stay with the endpoint lifecycle. Never use history-supplied path strings.
|
|
"""
|
|
def __init__(self, temp_dir, request_id):
|
|
self.root = Path(temp_dir).absolute() if temp_dir is not None else None
|
|
self.key = hashlib.sha256(request_id.encode()).hexdigest()
|
|
self.fd = None
|
|
self.files = {}
|
|
|
|
def _open(self):
|
|
if self.fd is not None:
|
|
return True
|
|
if self.root is None or ".." in self.root.parts:
|
|
return False
|
|
flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
|
|
fd = None
|
|
try:
|
|
fd = os.open(self.root.anchor, flags)
|
|
for component in (*self.root.parts[1:], "distributed-media-v1", self.key):
|
|
child = os.open(component, flags, dir_fd=fd)
|
|
os.close(fd)
|
|
fd = child
|
|
self.fd, fd = fd, None
|
|
return True
|
|
except OSError:
|
|
return False
|
|
finally:
|
|
if fd is not None:
|
|
os.close(fd)
|
|
|
|
def _identity(self, name):
|
|
info = os.stat(name, dir_fd=self.fd, follow_symlinks=False)
|
|
if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1:
|
|
return None
|
|
return (info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns)
|
|
|
|
def identify(self, entry, media):
|
|
if self.root is None or not isinstance(entry.get("outputs"), dict):
|
|
return
|
|
# A failed prompt may contain valid completed exporters alongside missing
|
|
# or malformed ones. Claim only independently validated declared records.
|
|
for binding in media["outputs"]:
|
|
node_id = binding["node_id"]
|
|
try:
|
|
_, files = _manifest({"outputs": {node_id: entry["outputs"].get(node_id)}},
|
|
dict(media, outputs=[binding]))
|
|
except ValueError:
|
|
continue
|
|
if not files or not self._open():
|
|
continue
|
|
for artifact in files:
|
|
name = artifact["id"] + ".bin" # _manifest enforces UUID hex.
|
|
try:
|
|
identity = self._identity(name)
|
|
except OSError:
|
|
continue
|
|
if identity is not None:
|
|
self.files.setdefault(name, identity)
|
|
|
|
def close(self, *, cleanup):
|
|
if self.fd is None:
|
|
return
|
|
try:
|
|
if cleanup:
|
|
for name, identity in self.files.items():
|
|
try:
|
|
if self._identity(name) == identity:
|
|
os.unlink(name, dir_fd=self.fd)
|
|
except OSError:
|
|
pass # Endpoint lifecycle owns inaccessible leftovers.
|
|
finally:
|
|
os.close(self.fd)
|
|
self.fd = None
|
|
# Deliberately retain directories: they may belong to concurrent runs.
|
|
|
|
|
|
async def _body(response, maximum):
|
|
length = response.content_length
|
|
if length is not None and length > maximum:
|
|
raise ValueError("Local Comfy response exceeds declared byte length")
|
|
body = bytearray()
|
|
async for chunk in response.content.iter_chunked(64 * 1024):
|
|
if len(body) + len(chunk) > maximum:
|
|
raise ValueError("Local Comfy response exceeds declared byte length")
|
|
body.extend(chunk)
|
|
return bytes(body)
|
|
|
|
|
|
class _Execution:
|
|
def __init__(self, session, url, stage, media, targets, deadline, exports):
|
|
self.session, self.url, self.stage = session, url, stage
|
|
self.media, self.targets = media, targets
|
|
self.deadline = deadline
|
|
self.run_id = uuid.uuid4().hex
|
|
self.prompt_id = None
|
|
self.attempted = False
|
|
self.terminal = False
|
|
self.exports = exports
|
|
self.interruption_blocked = False
|
|
|
|
def remaining(self):
|
|
remaining = self.deadline - asyncio.get_running_loop().time()
|
|
if remaining <= 0:
|
|
raise asyncio.TimeoutError("Media execution deadline expired")
|
|
return remaining
|
|
|
|
async def json_request(self, method, path, *, body=None, permit_rejection=False):
|
|
async with self.session.request(method, self.url + path, json=body,
|
|
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining())),
|
|
allow_redirects=False) as response:
|
|
raw = await _body(response, MAX_HTTP_JSON_BYTES)
|
|
# Native queue-delete and interrupt routes return an empty 200.
|
|
if response.status == 200 and not raw and method == "POST" and path in {"/queue", "/interrupt"}:
|
|
return {}
|
|
try:
|
|
value = json.loads(raw)
|
|
except (ValueError, UnicodeError, RecursionError):
|
|
if response.status != 200:
|
|
raise _HTTPError(response.status) from None
|
|
raise ValueError("Invalid local Comfy JSON response") from None
|
|
if response.status != 200:
|
|
# Native /prompt rejects the entire graph with a 400 error.
|
|
# Other failures may hide acceptance and are never retried.
|
|
if (permit_rejection and response.status == 400 and isinstance(value, dict)
|
|
and value.get("error") and "prompt_id" not in value):
|
|
self.terminal = True
|
|
raise _HTTPError(response.status)
|
|
if not isinstance(value, dict):
|
|
raise ValueError("Invalid local Comfy JSON object")
|
|
return value
|
|
|
|
async def get_json(self, path):
|
|
for attempt in range(GET_ATTEMPTS):
|
|
try:
|
|
return await self.json_request("GET", path)
|
|
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
|
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
|
|
raise
|
|
if attempt == GET_ATTEMPTS - 1:
|
|
raise
|
|
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
|
|
|
|
async def history(self):
|
|
response = await self.get_json("/history/" + quote(self.prompt_id, safe=""))
|
|
entry = response.get(self.prompt_id)
|
|
if entry is None:
|
|
return None
|
|
if not isinstance(entry, dict) or not isinstance(entry.get("status"), dict):
|
|
raise ValueError("Malformed Comfy prompt history")
|
|
status = entry["status"]
|
|
messages = status.get("messages", [])
|
|
if not isinstance(messages, list):
|
|
raise ValueError("Malformed Comfy history messages")
|
|
failed = status.get("status_str") in {"error", "interrupted"} or any(
|
|
isinstance(message, list) and message and message[0] in {"execution_error", "execution_interrupted"}
|
|
for message in messages)
|
|
if failed or status.get("completed") is True:
|
|
self.terminal = True
|
|
self.exports.identify(entry, self.media)
|
|
if failed or status.get("status_str") != "success":
|
|
raise RuntimeError("Local Comfy prompt execution failed")
|
|
return entry
|
|
return None
|
|
|
|
async def submit_and_wait(self, workflow):
|
|
body = {"prompt": workflow, "client_id": self.run_id,
|
|
"extra_data": {"distributed_media_run_id": self.run_id},
|
|
"partial_execution_targets": list(self.targets)}
|
|
self.remaining() # CPU staging must not spend past the absolute deadline.
|
|
self.attempted = True
|
|
response = await self.json_request("POST", "/prompt", body=body, permit_rejection=True)
|
|
self.prompt_id = _scope_id(response.get("prompt_id"))
|
|
if response.get("node_errors"):
|
|
raise RuntimeError("Local Comfy accepted an incomplete set of execution targets")
|
|
while True:
|
|
entry = await self.history()
|
|
if entry is not None:
|
|
return entry
|
|
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
|
|
|
async def settle(self):
|
|
"""Bounded recovery; caller shields this from task cancellation."""
|
|
deleted = False
|
|
while True:
|
|
if self.prompt_id:
|
|
try:
|
|
await self.history()
|
|
except RuntimeError:
|
|
if not self.terminal:
|
|
raise
|
|
if self.terminal:
|
|
return
|
|
queue = await self.get_json("/queue")
|
|
groups = []
|
|
for key in ("queue_running", "queue_pending"):
|
|
rows = queue.get(key)
|
|
if not isinstance(rows, list):
|
|
raise ValueError("Malformed Comfy queue")
|
|
group = {}
|
|
for row in rows:
|
|
if not isinstance(row, list) or len(row) < 4 or not isinstance(row[3], dict):
|
|
raise ValueError("Malformed Comfy queue entry")
|
|
pid = _scope_id(row[1])
|
|
if pid in group:
|
|
raise ValueError("Duplicate Comfy queue identity")
|
|
group[pid] = row[3]
|
|
groups.append(group)
|
|
running, pending = groups
|
|
if self.prompt_id is None:
|
|
own = [pid for group in groups for pid, extra in group.items()
|
|
if extra.get("distributed_media_run_id") == self.run_id
|
|
and extra.get("client_id") == self.run_id]
|
|
if len(own) > 1:
|
|
raise ValueError("Ambiguous Comfy queue ownership")
|
|
if own:
|
|
self.prompt_id = own[0]
|
|
if self.prompt_id in pending and not deleted:
|
|
await self.json_request("POST", "/queue", body={"delete": [self.prompt_id]})
|
|
deleted = True
|
|
elif self.prompt_id in running:
|
|
# Pinned ComfyUI server.py checks a snapshot then globally
|
|
# interrupts, outside the queue mutex. Supplying prompt_id does
|
|
# NOT close that race. No safe capability negotiation exists on
|
|
# this route; do not invoke it, even after our own queue read.
|
|
self.interruption_blocked = True
|
|
raise RuntimeError("Native prompt-scoped interruption unavailable")
|
|
elif deleted and self.prompt_id not in running and self.prompt_id not in pending:
|
|
self.terminal = True # queue deletion confirmed by exact read-back
|
|
return
|
|
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
|
|
|
async def fetch(self, artifact):
|
|
transport = artifact["transport"]
|
|
params = {"filename": transport["filename"], "subfolder": transport["subfolder"], "type": "temp"}
|
|
for attempt in range(GET_ATTEMPTS):
|
|
try:
|
|
async with self.session.get(self.url + "/view", params=params, allow_redirects=False,
|
|
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining()))) as response:
|
|
if response.status != 200:
|
|
raise _HTTPError(response.status)
|
|
payload = await _body(response, artifact["byte_length"])
|
|
break
|
|
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
|
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
|
|
raise
|
|
if attempt == GET_ATTEMPTS - 1:
|
|
raise
|
|
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
|
|
if len(payload) != artifact["byte_length"]:
|
|
raise ValueError("Media output byte length mismatch")
|
|
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
|
|
raise ValueError("Media output checksum mismatch")
|
|
return payload
|
|
|
|
|
|
def _manifest(entry, media):
|
|
outputs = entry.get("outputs")
|
|
if not isinstance(outputs, dict):
|
|
raise ValueError("Missing media history outputs")
|
|
expected_ids = {binding["node_id"] for binding in media["outputs"]}
|
|
if any(node_id not in expected_ids and isinstance(value, dict) and "media_exports" in value
|
|
for node_id, value in outputs.items()):
|
|
raise ValueError("Unexpected media export history")
|
|
response = {"version": VERSION, "profile": PROFILE, "request_id": media["request_id"],
|
|
"bindings": [], "artifacts": [], "items": []}
|
|
files = []
|
|
subfolder = "distributed-media-v1/" + hashlib.sha256(media["request_id"].encode()).hexdigest()
|
|
for binding in media["outputs"]:
|
|
node = outputs.get(binding["node_id"], {})
|
|
records = node.get("media_exports") if isinstance(node, dict) else None
|
|
if not isinstance(records, list) or len(records) != 1:
|
|
raise ValueError("Missing or duplicate media export result")
|
|
record = records[0]
|
|
if (not isinstance(record, dict) or set(record) != {"request_id", "binding", "artifacts", "items"}
|
|
or record["request_id"] != media["request_id"]):
|
|
raise ValueError("Media export request identity mismatch")
|
|
if not isinstance(record["artifacts"], list) or len(record["artifacts"]) > MAX_ARTIFACTS:
|
|
raise ValueError("Invalid exported artifacts count")
|
|
converted = []
|
|
for artifact in record["artifacts"]:
|
|
if not isinstance(artifact, dict) or not isinstance(artifact.get("transport"), dict):
|
|
raise ValueError("Invalid exported artifact")
|
|
transport = artifact["transport"]
|
|
aid = artifact.get("id")
|
|
if (not isinstance(aid, str) or not re.fullmatch(r"[a-f0-9]{32}", aid)
|
|
or transport != {"type": "comfy-file", "filename": aid + ".bin",
|
|
"subfolder": subfolder, "folder_type": "temp"}):
|
|
raise ValueError("Unsafe or unowned exported media file")
|
|
converted.append(dict(artifact, transport=dict(_VALIDATION_TRANSPORT)))
|
|
partial = dict(response, bindings=[record["binding"]], artifacts=converted, items=record["items"])
|
|
validate_response(partial, request_id=media["request_id"], bindings=[binding])
|
|
response["bindings"].append(copy.deepcopy(record["binding"]))
|
|
response["artifacts"].extend(converted)
|
|
response["items"].extend(copy.deepcopy(record["items"]))
|
|
files.extend(copy.deepcopy(record["artifacts"]))
|
|
# Validate the whole manifest BEFORE the first /view or publisher callback.
|
|
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
|
|
return response, files
|
|
|
|
|
|
def _publication_plan(response, files, publish):
|
|
skeleton = dict(response, artifacts=[dict(a, transport={"type": "inline", "data_base64": ""}) for a in files])
|
|
base_size = len(_json_bytes({"media": skeleton}))
|
|
sizes = {a["id"]: 4 * ((a["byte_length"] + 2) // 3) for a in files}
|
|
predicted = base_size + sum(sizes.values())
|
|
published = set()
|
|
if predicted > MAX_RESPONSE_JSON_BYTES and publish is None:
|
|
raise ValueError("Media response exceeds inline JSON limit; a publisher is required")
|
|
# Reserve enough space for a maximum-size escaped URL, including metadata.
|
|
# Small remaining inline values preserve their exact bytes and ordering.
|
|
for aid in sorted(sizes, key=sizes.get, reverse=True):
|
|
if predicted <= MAX_RESPONSE_JSON_BYTES:
|
|
break
|
|
if sizes[aid] > 65_536:
|
|
published.add(aid)
|
|
predicted -= sizes[aid] - 65_536
|
|
if predicted > MAX_RESPONSE_JSON_BYTES:
|
|
raise ValueError("Media response metadata exceeds JSON limit")
|
|
return published
|
|
|
|
|
|
async def _run(execution, workflow, legacy, download, publish):
|
|
media = execution.media
|
|
imports = {}
|
|
for artifact in media["artifacts"]:
|
|
execution.remaining()
|
|
payload = await resolve_artifact(artifact, download=download)
|
|
execution.remaining()
|
|
filename = execution.stage.write(payload)
|
|
node_id = "distributed-media-import:" + uuid.uuid4().hex
|
|
while node_id in workflow:
|
|
node_id = "distributed-media-import:" + uuid.uuid4().hex
|
|
metadata = {key: value for key, value in artifact.items() if key != "transport"}
|
|
workflow[node_id] = {"class_type": "MediaImportV1", "inputs": {
|
|
"artifact_json": json.dumps(metadata, allow_nan=False), "filename": filename}}
|
|
imports[artifact["id"]] = [node_id, _IMPORT_SLOTS[artifact["kind"]]]
|
|
for ref in media["inputs"]:
|
|
workflow[ref["node_id"]]["inputs"][ref["input_name"]] = list(imports[ref["artifact_id"]])
|
|
for encoded, size, matches in legacy:
|
|
execution.remaining()
|
|
try:
|
|
payload = base64.b64decode(encoded, validate=True)
|
|
except (ValueError, binascii.Error):
|
|
raise ValueError("Invalid legacy image base64") from None
|
|
if len(payload) != size:
|
|
raise ValueError("Legacy image byte length mismatch")
|
|
execution.remaining()
|
|
filename = execution.stage.write(payload)
|
|
for node_id in matches:
|
|
workflow[node_id]["inputs"]["image"] = filename
|
|
# Replacements can disconnect file loaders whose original endpoint-local
|
|
# paths are unavailable. Keep exactly the post-rewrite exporter closure;
|
|
# shared ancestors still needed by another branch must survive.
|
|
reachable, pending = set(), list(execution.targets)
|
|
while pending:
|
|
node_id = pending.pop()
|
|
if node_id in reachable:
|
|
continue
|
|
reachable.add(node_id)
|
|
pending.extend(_link(value, workflow) for value in workflow[node_id]["inputs"].values()
|
|
if isinstance(value, list))
|
|
workflow = {node_id: node for node_id, node in workflow.items() if node_id in reachable}
|
|
entry = await execution.submit_and_wait(workflow)
|
|
response, files = _manifest(entry, media)
|
|
published = _publication_plan(response, files, publish)
|
|
artifacts = []
|
|
for artifact in files:
|
|
payload = await execution.fetch(artifact)
|
|
transport = None
|
|
if artifact["id"] in published:
|
|
metadata = {key: copy.deepcopy(value) for key, value in artifact.items() if key != "transport"}
|
|
transport = await publish(metadata, payload)
|
|
if not isinstance(transport, dict) or transport.get("type") != "url":
|
|
raise ValueError("Publisher must return an expiring HTTPS URL transport")
|
|
artifacts.append(artifact_from_payload(artifact["id"], artifact["kind"], artifact["descriptor"],
|
|
payload, transport=transport))
|
|
response["artifacts"] = artifacts
|
|
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
|
|
result = {"media": response}
|
|
if len(_json_bytes(result)) > MAX_RESPONSE_JSON_BYTES:
|
|
raise ValueError("Media response exceeds JSON byte limit")
|
|
execution.remaining()
|
|
return result
|
|
|
|
|
|
def _release_uncertain_work(task):
|
|
_UNCERTAIN_WORK.discard(task)
|
|
try:
|
|
task.result()
|
|
except BaseException:
|
|
pass
|
|
|
|
|
|
async def _drain_work(work):
|
|
"""Cancel once; a misbehaving callback cannot make teardown unbounded."""
|
|
if not work.done():
|
|
work.cancel()
|
|
deadline = asyncio.get_running_loop().time() + DRAIN_GRACE_SECONDS
|
|
while not work.done():
|
|
remaining = deadline - asyncio.get_running_loop().time()
|
|
if remaining <= 0:
|
|
# Keep a strong reference until it exits, but never delete its files
|
|
# asynchronously: uncertain ownership belongs to endpoint recovery.
|
|
_UNCERTAIN_WORK.add(work)
|
|
work.add_done_callback(_release_uncertain_work)
|
|
return False
|
|
try:
|
|
# Unlike wait_for, wait's timeout does not wait for cancellation.
|
|
await asyncio.wait({work}, timeout=remaining)
|
|
except asyncio.CancelledError:
|
|
continue
|
|
_release_uncertain_work(work)
|
|
return True
|
|
|
|
|
|
async def execute_media_request(input_data, *, comfy_url, input_dir, temp_dir=None,
|
|
timeout_seconds=1800, session=None, download=None, publish=None) -> dict:
|
|
"""Validate, stage, submit once, collect exact media, and retire owned inputs.
|
|
|
|
``timeout_seconds`` covers staging, HTTP, polling, and publishing. A separate
|
|
bounded cancellation grace is allowed ONLY to establish the local prompt is
|
|
terminal. On uncertain acceptance/termination, raise MediaExecutionUncertain
|
|
and retain input files rather than racing an executing prompt. Never retry an
|
|
ambiguous submission, even when the exception is a timeout.
|
|
Async callback teardown has its own bounded grace; non-cooperative callbacks
|
|
also raise MediaExecutionUncertain and retain files, even after terminal history.
|
|
|
|
``temp_dir`` is a trusted endpoint setting for this ComfyUI's local temp root,
|
|
never a request-supplied path. Only validated exporter UUID.bin files beneath
|
|
distributed-media-v1/sha256(request_id) are retired, after all collection and
|
|
publisher work has unwound. Unidentified files remain lifecycle-owned. With
|
|
``temp_dir=None`` all exporter files remain lifecycle-owned as before.
|
|
|
|
The pinned native /interrupt route has a check/global-interrupt race. Running
|
|
prompts therefore fail closed with MediaExecutionUncertain, without posting
|
|
/interrupt. A prompt-scoped request body alone is not a safety guarantee.
|
|
"""
|
|
if (type(timeout_seconds) not in (int, float) or not math.isfinite(timeout_seconds)
|
|
or timeout_seconds <= 0):
|
|
raise ValueError("Invalid media execution timeout")
|
|
if not isinstance(comfy_url, str):
|
|
raise ValueError("Invalid local Comfy URL")
|
|
try:
|
|
parsed = urlsplit(comfy_url)
|
|
if (parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username
|
|
or parsed.password or parsed.query or parsed.fragment):
|
|
raise ValueError()
|
|
parsed.port
|
|
except ValueError:
|
|
raise ValueError("Invalid local Comfy URL") from None
|
|
deadline = asyncio.get_running_loop().time() + timeout_seconds
|
|
workflow, media, legacy, targets = _validate_request(input_data)
|
|
stage = _Staging(input_dir, media["request_id"])
|
|
exports = _OwnedExports(temp_dir, media["request_id"])
|
|
owned_session = session is None
|
|
if owned_session:
|
|
session = aiohttp.ClientSession()
|
|
execution = _Execution(session, comfy_url.rstrip("/"), stage, media, targets, deadline, exports)
|
|
work = None
|
|
try:
|
|
remaining = execution.remaining()
|
|
work = asyncio.create_task(_run(execution, workflow, legacy, download, publish))
|
|
return await asyncio.wait_for(asyncio.shield(work), remaining)
|
|
except BaseException as exc:
|
|
# Retire files only after readers unwind. Cooperative async callbacks
|
|
# normally drain immediately; an uncooperative one requires quarantine.
|
|
if work is not None and not await _drain_work(work):
|
|
raise MediaExecutionUncertain(
|
|
"Media callback teardown unconfirmed; owned files retained; "
|
|
"quarantine this endpoint before reuse; do not resubmit automatically"
|
|
) from exc
|
|
if execution.attempted and not execution.terminal:
|
|
execution.deadline = asyncio.get_running_loop().time() + CANCEL_GRACE_SECONDS
|
|
recovery = asyncio.create_task(asyncio.wait_for(execution.settle(), CANCEL_GRACE_SECONDS))
|
|
# Shield the bounded recovery even from repeated task.cancel() calls.
|
|
while not recovery.done():
|
|
try:
|
|
await asyncio.shield(recovery)
|
|
except asyncio.CancelledError:
|
|
continue
|
|
except Exception:
|
|
break
|
|
try:
|
|
recovery.result()
|
|
except BaseException:
|
|
pass
|
|
if not execution.terminal:
|
|
raise MediaExecutionUncertain(
|
|
"Local Comfy prompt terminal state unconfirmed; owned staged files retained "
|
|
f"under {stage.relative}; prompt_id={execution.prompt_id or 'unknown'}; "
|
|
"do not resubmit this request automatically"
|
|
+ ("; native prompt-scoped interruption unavailable (non-atomic /interrupt route)"
|
|
if execution.interruption_blocked else "")
|
|
) from exc
|
|
raise
|
|
finally:
|
|
try:
|
|
try:
|
|
exports.close(cleanup=(work is None or work.done()) and execution.terminal)
|
|
finally:
|
|
stage.close(cleanup=(work is None or work.done()) and
|
|
(not execution.attempted or execution.terminal))
|
|
finally:
|
|
if owned_session:
|
|
await session.close()
|