refactor: centralize audio payload codec and harden worker dispatch compatibility

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")

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")

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 = []

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()

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")

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,
+    }