feat: support audio-only distributed workflows (#87)

Support image, audio, or combined distributed collection with guarded audio-only worker, master, endpoint, preview, and delegate-only paths.

tests/api/test_distributed_queue.py

@@ -302,6 +302,77 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(queued["audio"]["sample_rate"], 44100)
         self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
 
+    async def test_job_complete_accepts_audio_without_image(self):
+        queue = asyncio.Queue()
+        job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
+        job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
+        request = _FakeRequest(
+            {
+                "job_id": "audio-only-job",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "audio": self._encoded_audio_payload(),
+                "is_last": True,
+            }
+        )
+
+        with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
+            response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        decode_image.assert_not_called()
+        queued = await queue.get()
+        self.assertIsNone(queued["tensor"])
+        self.assertEqual(queued["audio"]["sample_rate"], 44100)
+
+    async def test_job_complete_rejects_payload_without_image_or_audio(self):
+        request = _FakeRequest(
+            {
+                "job_id": "job-1",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "is_last": True,
+            }
+        )
+
+        response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("image or audio", response.payload.get("message", "").lower())
+
+    async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
+        request = _FakeRequest(
+            {
+                "job_id": "job-1",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "image": 123,
+                "audio": self._encoded_audio_payload(),
+                "is_last": True,
+            }
+        )
+
+        response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("image", response.payload.get("message", "").lower())
+
+    async def test_job_complete_returns_400_for_malformed_audio(self):
+        request = _FakeRequest(
+            {
+                "job_id": "job-1",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
+                "is_last": True,
+            }
+        )
+
+        response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("audio.shape", response.payload.get("message", "").lower())
+
     def test_decode_audio_payload_rejects_bad_shape(self):
         bad = {
             "sample_rate": 44100,

tests/test_collector_list_inputs.py

@@ -124,6 +124,49 @@ def test_collector_opts_into_comfyui_list_inputs():
     assert collector.INPUT_IS_LIST is True
 
 
+def test_collector_exposes_images_as_optional_input():
+    input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()
+
+    assert "images" not in input_types["required"]
+    assert input_types["optional"]["images"] == ("IMAGE",)
+
+
+def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
+    collector = _load_collector_module().DistributedCollectorNode()
+    audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
+
+    images, returned_audio = collector.run(images=None, audio=[audio])
+
+    assert images is None
+    assert returned_audio is audio
+
+
+def test_collector_rejects_missing_images_and_audio():
+    collector = _load_collector_module().DistributedCollectorNode()
+
+    try:
+        collector.run(images=None, audio=None)
+    except ValueError as exc:
+        assert "image or audio" in str(exc).lower()
+    else:
+        raise AssertionError("Expected collector to reject a run with no media input")
+
+
+def test_delegate_only_master_allows_no_local_media_input():
+    collector = _load_collector_module().DistributedCollectorNode()
+
+    images, audio = collector.run(
+        images=None,
+        audio=None,
+        multi_job_id=["delegate-audio-job"],
+        delegate_only=[True],
+        enabled_worker_ids=["[]"],
+    )
+
+    assert images is None
+    assert tuple(audio["waveform"].shape) == (1, 2, 1)
+
+
 def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
     collector = _load_collector_module().DistributedCollectorNode()
     first = torch.zeros(1, 2, 2, 3)
@@ -203,3 +246,119 @@ def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final
     assert [payload["is_last"] for payload in posted_payloads] == [False, True]
     assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
     assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
+
+
+def test_audio_only_worker_sends_one_completion_without_image():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
+    posted = []
+
+    class _FakeResponse:
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc, tb):
+            return False
+
+        def raise_for_status(self):
+            return None
+
+    class _FakeSession:
+        def post(self, url, json, timeout):
+            posted.append((url, json, timeout.total))
+            return _FakeResponse()
+
+    async def _fake_get_client_session():
+        return _FakeSession()
+
+    module.get_client_session = _fake_get_client_session
+    module.encode_audio_payload = lambda value: {"encoded": value is audio}
+
+    images, returned_audio = collector.run(
+        images=None,
+        audio=[audio],
+        multi_job_id=["audio-job"],
+        is_worker=[True],
+        master_url=["http://master"],
+        worker_id=["worker-a"],
+    )
+
+    assert images is None
+    assert returned_audio is audio
+    assert len(posted) == 1
+    assert posted[0][0] == "http://master/distributed/job_complete"
+    assert posted[0][2] == 600
+    assert posted[0][1] == {
+        "job_id": "audio-job",
+        "worker_id": "worker-a",
+        "batch_idx": 0,
+        "audio": {"encoded": True},
+        "is_last": True,
+    }
+
+
+def test_audio_only_master_combines_local_and_worker_audio():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
+    worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    queue.put_nowait(
+        {
+            "worker_id": "worker-a",
+            "image_index": 0,
+            "tensor": None,
+            "audio": worker_audio,
+            "is_last": True,
+        }
+    )
+    module.prompt_server.distributed_pending_jobs = {"audio-job": queue}
+
+    images, combined_audio = asyncio.run(
+        collector.execute(
+            images=None,
+            audio=master_audio,
+            multi_job_id="audio-job",
+            enabled_worker_ids='["worker-a"]',
+        )
+    )
+
+    assert images is None
+    assert combined_audio["sample_rate"] == 48000
+    assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
+    assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
+    assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])
+
+
+def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    queue.put_nowait(
+        {
+            "worker_id": "worker-a",
+            "image_index": 0,
+            "tensor": None,
+            "audio": worker_audio,
+            "is_last": True,
+        }
+    )
+    module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}
+
+    images, combined_audio = asyncio.run(
+        collector.execute(
+            images=None,
+            audio=None,
+            multi_job_id="delegate-audio-job",
+            enabled_worker_ids='["worker-a"]',
+            delegate_only=True,
+        )
+    )
+
+    assert images is None
+    assert combined_audio["sample_rate"] == 48000
+    assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])

tests/test_prompt_transform.py

@@ -71,6 +71,15 @@ def _collector_only_prompt():
     }
 
 
+def _audio_only_prompt():
+    """1(Audio source) → 2(DistributedCollector) → 3(Audio sink)."""
+    return {
+        "1": {"class_type": "LoadAudio", "inputs": {}},
+        "2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
+        "3": {"class_type": "SaveAudio", "inputs": {"audio": ["2", 1]}},
+    }
+
+
 def _delegate_prompt():
     """1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
     return {
@@ -217,6 +226,35 @@ class PrunePromptForWorkerTests(unittest.TestCase):
         self.assertEqual(len(preview_nodes), 1)
         self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0])
 
+    def test_injects_preview_audio_for_audio_only_collector(self):
+        result = pt.prune_prompt_for_worker(_audio_only_prompt())
+        preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewAudio"]
+        self.assertEqual(len(preview_nodes), 1)
+        self.assertEqual(preview_nodes[0]["inputs"]["audio"], ["2", 1])
+        self.assertFalse(any(n.get("class_type") == "PreviewImage" for n in result.values()))
+
+    def test_prefers_image_preview_when_collector_has_images_and_audio(self):
+        prompt = _linear_prompt()
+        prompt["6"] = {"class_type": "LoadAudio", "inputs": {}}
+        prompt["4"]["inputs"]["audio"] = ["6", 0]
+        result = pt.prune_prompt_for_worker(prompt)
+        self.assertEqual(
+            len([n for n in result.values() if n.get("class_type") == "PreviewImage"]),
+            1,
+        )
+        self.assertFalse(any(n.get("class_type") == "PreviewAudio" for n in result.values()))
+
+    def test_preserves_image_preview_for_distributed_upscale(self):
+        prompt = {
+            "1": {"class_type": "LoadImage", "inputs": {}},
+            "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
+            "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
+        }
+        result = pt.prune_prompt_for_worker(prompt)
+        preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
+        self.assertEqual(len(preview_nodes), 1)
+        self.assertEqual(preview_nodes[0]["inputs"]["images"], ["2", 0])
+
     def test_no_preview_image_when_no_downstream(self):
         result = pt.prune_prompt_for_worker(_collector_only_prompt())
         preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
@@ -242,7 +280,7 @@ class PrunePromptForWorkerTests(unittest.TestCase):
     def test_upscale_node_is_treated_as_distributed(self):
         prompt = {
             "1": {"class_type": "KSampler", "inputs": {}},
-            "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}},
+            "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
             "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
         }
         result = pt.prune_prompt_for_worker(prompt)
@@ -285,12 +323,22 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
         placeholder_id = empty_nodes[0][0]
         self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0])
 
+    def test_audio_only_collector_does_not_get_image_placeholder(self):
+        prompt = _audio_only_prompt()
+        result = pt.prepare_delegate_master_prompt(prompt, ["2"])
+        empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
+        self.assertEqual(empty_nodes, [])
+        self.assertNotIn("images", result["2"].get("inputs", {}))
+        self.assertNotIn("audio", result["2"].get("inputs", {}))
+
     def test_one_placeholder_per_collector(self):
         """Two collectors → two placeholders."""
         prompt = {
-            "1": {"class_type": "DistributedCollector", "inputs": {}},
-            "2": {"class_type": "DistributedCollector", "inputs": {}},
+            "1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0]}},
+            "2": {"class_type": "DistributedCollector", "inputs": {"images": ["11", 0]}},
             "3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
+            "10": {"class_type": "LoadImage", "inputs": {}},
+            "11": {"class_type": "LoadImage", "inputs": {}},
         }
         result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"])
         empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]