From 487f5b9c8562ebb8a048935e49768ae5b396984d Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 28 Feb 2026 01:03:23 +0000 Subject: [PATCH] refactor: centralize audio payload codec and harden worker dispatch compatibility --- api/job_routes.py | 47 +------------ api/worker_routes.py | 11 ++- nodes/collector.py | 27 +------- tests/api/test_distributed_queue.py | 44 ++++++++++++ tests/api/test_worker_routes.py | 8 +++ utils/audio_payload.py | 103 ++++++++++++++++++++++++++++ 6 files changed, 169 insertions(+), 71 deletions(-) create mode 100644 utils/audio_payload.py diff --git a/api/job_routes.py b/api/job_routes.py index f4f9ef5..3843101 100644 --- a/api/job_routes.py +++ b/api/job_routes.py @@ -9,7 +9,6 @@ import time from aiohttp import web import server import torch -import numpy as np from PIL import Image from ..utils.logging import debug_log @@ -135,51 +134,9 @@ def _decode_canonical_png_tensor(image_payload): def _decode_audio_payload(audio_payload): """Decode canonical audio payload into an AUDIO dict.""" - if audio_payload is None: - return None - if not isinstance(audio_payload, dict): - raise ValueError("Field 'audio' must be an object when provided.") + from ..utils.audio_payload import decode_audio_payload - encoded = audio_payload.get("data") - shape = audio_payload.get("shape") - sample_rate = audio_payload.get("sample_rate", 44100) - - if not isinstance(encoded, str) or not encoded.strip(): - raise ValueError("Field 'audio.data' must be a non-empty base64 string.") - if not isinstance(shape, list) or len(shape) != 3: - raise ValueError("Field 'audio.shape' must be a 3-item list [batch, channels, samples].") - - try: - shape_tuple = tuple(int(dim) for dim in shape) - except (TypeError, ValueError) as exc: - raise ValueError("Field 'audio.shape' must contain integers.") from exc - if any(dim < 0 for dim in shape_tuple): - raise ValueError("Field 'audio.shape' dimensions must be non-negative.") - - try: - sample_rate = int(sample_rate) - except (TypeError, ValueError) as exc: - raise ValueError("Field 'audio.sample_rate' must be an integer.") from exc - if sample_rate <= 0: - raise ValueError("Field 'audio.sample_rate' must be positive.") - - try: - raw = base64.b64decode(encoded, validate=True) - except (binascii.Error, ValueError) as exc: - raise ValueError("Field 'audio.data' is not valid base64.") from exc - - expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4 - if len(raw) != expected_bytes: - raise ValueError( - f"Field 'audio.data' byte size mismatch: expected {expected_bytes}, got {len(raw)}." - ) - - array = np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple) - waveform = torch.from_numpy(array.copy()) - return { - "waveform": ensure_contiguous(waveform), - "sample_rate": sample_rate, - } + return decode_audio_payload(audio_payload) @server.PromptServer.instance.routes.post("/distributed/prepare_job") diff --git a/api/worker_routes.py b/api/worker_routes.py index ae29363..e3d042f 100644 --- a/api/worker_routes.py +++ b/api/worker_routes.py @@ -22,7 +22,16 @@ from ..workers.detection import ( is_docker_environment, is_runpod_environment, ) -from ..utils.async_helpers import PromptValidationError, queue_prompt_payload +try: + from ..utils.async_helpers import PromptValidationError, queue_prompt_payload +except ImportError: + from ..utils.async_helpers import queue_prompt_payload + + class PromptValidationError(RuntimeError): + def __init__(self, message, validation_error=None, node_errors=None): + super().__init__(str(message)) + self.validation_error = validation_error if isinstance(validation_error, dict) else {} + self.node_errors = node_errors if isinstance(node_errors, dict) else {} @server.PromptServer.instance.routes.get("/distributed/worker_ws") diff --git a/nodes/collector.py b/nodes/collector.py index ff90070..898f153 100644 --- a/nodes/collector.py +++ b/nodes/collector.py @@ -15,6 +15,7 @@ from ..utils.config import get_worker_timeout_seconds, load_config, is_master_de from ..utils.constants import HEARTBEAT_INTERVAL from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous from ..utils.network import build_worker_url, get_client_session, probe_worker +from ..utils.audio_payload import encode_audio_payload from ..utils.async_helpers import run_async_in_server_loop prompt_server = _server.PromptServer.instance @@ -66,7 +67,7 @@ class DistributedCollectorNode: if batch_size == 0: return - encoded_audio = self._encode_audio_payload(audio) + encoded_audio = encode_audio_payload(audio) session = await get_client_session() url = f"{master_url}/distributed/job_complete" @@ -97,30 +98,6 @@ class DistributedCollectorNode: debug_log(f"Worker - Full error details: URL={url}") raise # Re-raise to handle at caller level - def _encode_audio_payload(self, audio): - """Serialize AUDIO dict into a JSON-safe payload for canonical job_complete envelopes.""" - if not isinstance(audio, dict): - return None - - waveform = audio.get("waveform") - if waveform is None or not isinstance(waveform, torch.Tensor) or waveform.numel() == 0: - return None - - sample_rate = audio.get("sample_rate", 44100) - try: - sample_rate = int(sample_rate) - except (TypeError, ValueError): - sample_rate = 44100 - - waveform_cpu = waveform.detach().to(device="cpu", dtype=torch.float32).contiguous() - data_bytes = waveform_cpu.numpy().tobytes() - return { - "sample_rate": sample_rate, - "shape": [int(dim) for dim in waveform_cpu.shape], - "dtype": "float32", - "data": base64.b64encode(data_bytes).decode("ascii"), - } - def _combine_audio(self, master_audio, worker_audio, empty_audio): """Combine audio from master and workers into a single audio output.""" audio_pieces = [] diff --git a/tests/api/test_distributed_queue.py b/tests/api/test_distributed_queue.py index ecf2562..b553624 100644 --- a/tests/api/test_distributed_queue.py +++ b/tests/api/test_distributed_queue.py @@ -110,6 +110,40 @@ def _load_job_routes_module(): image_module.ensure_contiguous = lambda tensor: tensor sys.modules[f"{package_name}.utils.image"] = image_module + audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload") + + def _decode_audio_payload(payload): + if payload is None: + return None + if not isinstance(payload, dict): + raise ValueError("Field 'audio' must be an object when provided.") + + encoded = payload.get("data") + shape = payload.get("shape") + dtype = payload.get("dtype", "float32") + sample_rate = payload.get("sample_rate", 44100) + if not isinstance(encoded, str) or not encoded.strip(): + raise ValueError("Field 'audio.data' must be a non-empty base64 string.") + if not isinstance(shape, list) or len(shape) != 3: + raise ValueError("Field 'audio.shape' must be a 3-item list.") + if dtype != "float32": + raise ValueError("Field 'audio.dtype' must be 'float32'.") + try: + shape_tuple = tuple(int(dim) for dim in shape) + except Exception as exc: + raise ValueError("Field 'audio.shape' must contain integers.") from exc + + raw = base64.b64decode(encoded, validate=True) + expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4 + if len(raw) != expected_bytes: + raise ValueError("Field 'audio.data' byte size mismatch.") + + waveform = torch.from_numpy(np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple).copy()) + return {"waveform": waveform, "sample_rate": int(sample_rate)} + + audio_payload_module.decode_audio_payload = _decode_audio_payload + sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module + network_module = types.ModuleType(f"{package_name}.utils.network") async def _handle_api_error(_request, error, status=500): @@ -276,6 +310,16 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase): with self.assertRaises(ValueError): job_routes._decode_audio_payload(bad) + def test_decode_audio_payload_rejects_bad_dtype(self): + payload = { + "sample_rate": 44100, + "shape": [1, 2, 4], + "dtype": "float16", + "data": base64.b64encode((np.zeros((1, 2, 4), dtype=np.float32)).tobytes()).decode("ascii"), + } + with self.assertRaises(ValueError): + job_routes._decode_audio_payload(payload) + if __name__ == "__main__": unittest.main() diff --git a/tests/api/test_worker_routes.py b/tests/api/test_worker_routes.py index 03cdf4b..34b49f2 100644 --- a/tests/api/test_worker_routes.py +++ b/tests/api/test_worker_routes.py @@ -195,6 +195,14 @@ def _load_worker_routes_module(): return "prompt-id" async_helpers_module.queue_prompt_payload = _queue_prompt_payload + + class _PromptValidationError(RuntimeError): + def __init__(self, message="invalid prompt", validation_error=None, node_errors=None): + super().__init__(message) + self.validation_error = validation_error if isinstance(validation_error, dict) else {} + self.node_errors = node_errors if isinstance(node_errors, dict) else {} + + async_helpers_module.PromptValidationError = _PromptValidationError sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module schemas_module = types.ModuleType(f"{package_name}.api.schemas") diff --git a/utils/audio_payload.py b/utils/audio_payload.py new file mode 100644 index 0000000..ff70225 --- /dev/null +++ b/utils/audio_payload.py @@ -0,0 +1,103 @@ +import base64 +import binascii +import os + +import numpy as np +import torch + +from .image import ensure_contiguous + + +MAX_AUDIO_PAYLOAD_BYTES = int( + os.environ.get("COMFYUI_MAX_AUDIO_PAYLOAD_BYTES", str(256 * 1024 * 1024)) +) + + +def encode_audio_payload(audio_payload): + """Serialize an AUDIO dict into JSON-safe canonical envelope payload.""" + if not isinstance(audio_payload, dict): + return None + + waveform = audio_payload.get("waveform") + if waveform is None or not isinstance(waveform, torch.Tensor) or waveform.numel() == 0: + return None + + sample_rate = audio_payload.get("sample_rate", 44100) + try: + sample_rate = int(sample_rate) + except (TypeError, ValueError): + sample_rate = 44100 + + waveform_cpu = waveform.detach().to(device="cpu", dtype=torch.float32).contiguous() + data_bytes = waveform_cpu.numpy().tobytes() + if len(data_bytes) > MAX_AUDIO_PAYLOAD_BYTES: + raise ValueError( + f"Audio payload too large: {len(data_bytes)} bytes exceeds {MAX_AUDIO_PAYLOAD_BYTES}." + ) + + return { + "sample_rate": sample_rate, + "shape": [int(dim) for dim in waveform_cpu.shape], + "dtype": "float32", + "data": base64.b64encode(data_bytes).decode("ascii"), + } + + +def decode_audio_payload(audio_payload): + """Decode canonical envelope audio payload into an AUDIO dict.""" + if audio_payload is None: + return None + if not isinstance(audio_payload, dict): + raise ValueError("Field 'audio' must be an object when provided.") + + encoded = audio_payload.get("data") + shape = audio_payload.get("shape") + sample_rate = audio_payload.get("sample_rate", 44100) + dtype = audio_payload.get("dtype", "float32") + + if not isinstance(encoded, str) or not encoded.strip(): + raise ValueError("Field 'audio.data' must be a non-empty base64 string.") + if not isinstance(shape, list) or len(shape) != 3: + raise ValueError("Field 'audio.shape' must be a 3-item list [batch, channels, samples].") + if dtype != "float32": + raise ValueError("Field 'audio.dtype' must be 'float32'.") + + try: + shape_tuple = tuple(int(dim) for dim in shape) + except (TypeError, ValueError) as exc: + raise ValueError("Field 'audio.shape' must contain integers.") from exc + + if shape_tuple[0] <= 0 or shape_tuple[1] <= 0 or shape_tuple[2] < 0: + raise ValueError( + "Field 'audio.shape' must be [batch>0, channels>0, samples>=0]." + ) + + try: + sample_rate = int(sample_rate) + except (TypeError, ValueError) as exc: + raise ValueError("Field 'audio.sample_rate' must be an integer.") from exc + if sample_rate <= 0: + raise ValueError("Field 'audio.sample_rate' must be positive.") + + try: + raw = base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError) as exc: + raise ValueError("Field 'audio.data' is not valid base64.") from exc + + if len(raw) > MAX_AUDIO_PAYLOAD_BYTES: + raise ValueError( + f"Field 'audio.data' too large: {len(raw)} bytes exceeds {MAX_AUDIO_PAYLOAD_BYTES}." + ) + + expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4 + if len(raw) != expected_bytes: + raise ValueError( + f"Field 'audio.data' byte size mismatch: expected {expected_bytes}, got {len(raw)}." + ) + + array = np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple) + waveform = torch.from_numpy(array.copy()) + return { + "waveform": ensure_contiguous(waveform), + "sample_rate": sample_rate, + }