206 lines
7.2 KiB
Python
206 lines
7.2 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_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"}
|