Robert Wojciechowski
cedc8d4591
Support image, audio, or combined distributed collection with guarded audio-only worker, master, endpoint, preview, and delegate-only paths.
365 lines
12 KiB
Python
365 lines
12 KiB
Python
import importlib.util
|
|
import sys
|
|
import types
|
|
import asyncio
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
def _load_collector_module():
|
|
module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
|
|
package_name = "dist_collector_list_testpkg"
|
|
|
|
for mod_name in list(sys.modules):
|
|
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
|
del sys.modules[mod_name]
|
|
|
|
root_pkg = types.ModuleType(package_name)
|
|
root_pkg.__path__ = []
|
|
sys.modules[package_name] = root_pkg
|
|
|
|
nodes_pkg = types.ModuleType(f"{package_name}.nodes")
|
|
nodes_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.nodes"] = nodes_pkg
|
|
|
|
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
|
utils_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.utils"] = utils_pkg
|
|
|
|
class _Routes:
|
|
def post(self, _path):
|
|
return lambda fn: fn
|
|
|
|
def get(self, _path):
|
|
return lambda fn: fn
|
|
|
|
prompt_server = types.SimpleNamespace(
|
|
routes=_Routes(),
|
|
distributed_jobs_lock=None,
|
|
distributed_pending_jobs={},
|
|
)
|
|
server_module = types.ModuleType("server")
|
|
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
|
sys.modules["server"] = server_module
|
|
|
|
comfy_module = types.ModuleType("comfy")
|
|
model_management = types.ModuleType("comfy.model_management")
|
|
|
|
class InterruptProcessingException(Exception):
|
|
pass
|
|
|
|
model_management.InterruptProcessingException = InterruptProcessingException
|
|
model_management.throw_exception_if_processing_interrupted = lambda: None
|
|
comfy_module.model_management = model_management
|
|
|
|
comfy_utils = types.ModuleType("comfy.utils")
|
|
|
|
class ProgressBar:
|
|
def __init__(self, _total):
|
|
self.total = _total
|
|
self.updates = []
|
|
|
|
def update(self, value):
|
|
self.updates.append(value)
|
|
|
|
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
|
|
|
|
aiohttp_module = types.ModuleType("aiohttp")
|
|
aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
|
|
sys.modules["aiohttp"] = aiohttp_module
|
|
|
|
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
|
logging_module.debug_log = lambda *_args, **_kwargs: None
|
|
logging_module.log = lambda *_args, **_kwargs: None
|
|
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
|
|
|
config_module = types.ModuleType(f"{package_name}.utils.config")
|
|
config_module.get_worker_timeout_seconds = lambda: 0.1
|
|
config_module.load_config = lambda: {"workers": []}
|
|
config_module.is_master_delegate_only = lambda: False
|
|
sys.modules[f"{package_name}.utils.config"] = config_module
|
|
|
|
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
|
constants_module.HEARTBEAT_INTERVAL = 1.0
|
|
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
|
|
|
image_module = types.ModuleType(f"{package_name}.utils.image")
|
|
def _ensure_contiguous(tensor):
|
|
return tensor.contiguous() if hasattr(tensor, "contiguous") else tensor
|
|
|
|
image_module.ensure_contiguous = _ensure_contiguous
|
|
image_module.tensor_to_pil = lambda *_args, **_kwargs: None
|
|
image_module.pil_to_tensor = lambda value: value
|
|
sys.modules[f"{package_name}.utils.image"] = image_module
|
|
|
|
network_module = types.ModuleType(f"{package_name}.utils.network")
|
|
network_module.build_worker_url = lambda worker: "http://worker"
|
|
network_module.get_client_session = lambda: None
|
|
network_module.probe_worker = lambda *_args, **_kwargs: None
|
|
sys.modules[f"{package_name}.utils.network"] = network_module
|
|
|
|
audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
|
|
audio_payload_module.encode_audio_payload = lambda audio: audio
|
|
sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
|
|
|
|
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
|
async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
|
|
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
|
|
|
spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.collector", module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec is not None and spec.loader is not None
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def test_collector_opts_into_comfyui_list_inputs():
|
|
collector = _load_collector_module().DistributedCollectorNode
|
|
|
|
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)
|
|
second = torch.ones(1, 2, 2, 3)
|
|
|
|
images, audio = collector.run(
|
|
images=[first, second],
|
|
load_balance=[False],
|
|
audio=[None],
|
|
multi_job_id=[""],
|
|
is_worker=[False],
|
|
master_url=[""],
|
|
enabled_worker_ids=["[]"],
|
|
worker_batch_size=[1],
|
|
worker_id=[""],
|
|
pass_through=[False],
|
|
delegate_only=[False],
|
|
)
|
|
|
|
assert tuple(images.shape) == (2, 2, 2, 3)
|
|
assert torch.equal(images[0:1], first)
|
|
assert torch.equal(images[1:2], second)
|
|
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
|
|
|
|
|
def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final_item():
|
|
module = _load_collector_module()
|
|
collector = module.DistributedCollectorNode()
|
|
first = torch.zeros(1, 2, 2, 3)
|
|
second = torch.ones(1, 2, 2, 3)
|
|
posted_payloads = []
|
|
|
|
class _FakeImage:
|
|
def save(self, fp, format=None, compress_level=None):
|
|
fp.write(b"png-bytes")
|
|
|
|
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_payloads.append(json)
|
|
return _FakeResponse()
|
|
|
|
async def _fake_get_client_session():
|
|
return _FakeSession()
|
|
|
|
module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage()
|
|
module.get_client_session = _fake_get_client_session
|
|
module.encode_audio_payload = lambda _audio: None
|
|
|
|
images, audio = collector.run(
|
|
images=[first, second],
|
|
load_balance=[False],
|
|
audio=[None],
|
|
multi_job_id=["job-list-1"],
|
|
is_worker=[True],
|
|
master_url=["http://master"],
|
|
enabled_worker_ids=["[]"],
|
|
worker_batch_size=[1],
|
|
worker_id=["worker-a"],
|
|
pass_through=[False],
|
|
delegate_only=[False],
|
|
)
|
|
|
|
assert tuple(images.shape) == (2, 2, 2, 3)
|
|
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
|
assert len(posted_payloads) == 2
|
|
assert [payload["batch_idx"] for payload in posted_payloads] == [0, 1]
|
|
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"])
|