"""Exercise the actual loader, schemas and executor in a fresh CPU process. No HTTP listener, workers, model downloads or live installation changes. The original contracts were captured from 32ac027 using the same core. """ import asyncio import inspect import json import os from pathlib import Path import sys import tempfile from unittest.mock import patch from PIL import Image ROOT = Path(__file__).resolve().parents[2] COMFY_ROOT = Path(sys.argv[1]).resolve() sys.path.insert(0, str(COMFY_ROOT)) os.environ['COMFYUI_IS_WORKER'] = '1' import comfy.cli_args comfy.cli_args.args.cpu = True comfy.cli_args.args.disable_assets = True from app.assets.manager import default_asset_manager from comfy_api.v0_0_2 import io import execution import nodes import server import torch BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text()) # Inputs intentionally added after the V1 baseline was captured. They must be # optional, appended after the baseline inputs, and keep legacy behaviour as # their default so saved workflows are unaffected. ADDED_OPTIONAL_INPUTS = { 'DistributedCollector': { 'transfer_format': [ ['legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1'], {'default': 'legacy_png'}, ], 'transfer_quality': ['INT', {'default': 90, 'min': 1, 'max': 100, 'step': 1}], }, } # Nodes added after the V1 baseline; not part of saved-workflow contracts. # DistributedWorkerSink is internal (dev-only) and only inserted by the master. ADDED_NODES = {'DistributedWorkerSink'} def normalized_input(value): kind = value[0] opts = dict(value[1]) if len(value) > 1 else {} if kind == 'STRING': # V3 serializes the same default single-line widget explicitly. opts.setdefault('multiline', False) if kind == 'COMBO': kind = opts['options'] opts = {key: val for key, val in opts.items() if key != 'options'} if isinstance(kind, list): # Single-selection is the V1 dropdown default as well. opts.setdefault('multiselect', False) return [kind, opts] async def map_node(cls, values, extra=None): """Use the actual input parser and sanitized V3 executor clones.""" prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe', extra_data=extra or {}) assert not missing, missing # Use a worker thread for synchronous nodes: collector/upscale bridge back # to PromptServer's running loop, just as ComfyUI's prompt worker does. def run(): return asyncio.run(execution._async_map_node_over_list( 'v3-acceptance', 'probe', cls, prepared, cls.FUNCTION, v3_data=hidden)) return await asyncio.to_thread(run) def check_saved_workflows(mapping): checked = 0 occurrences = 0 for path in sorted((ROOT / 'workflows').glob('*.json')): workflow = json.loads(path.read_text()) local_nodes = {node['id']: node for node in workflow['nodes']} for saved in workflow['nodes']: if saved['type'] not in mapping: continue occurrences += 1 info = mapping[saved['type']].GET_NODE_INFO_V1() inputs = {**info['input'].get('required', {}), **info['input'].get('optional', {})} for socket in saved.get('inputs', []): assert socket['name'] in inputs, (path.name, saved['id'], socket) assert socket['type'] == inputs[socket['name']][0], (path.name, socket) for index, output in enumerate(saved.get('outputs', [])): assert output['type'] == info['output'][index], (path.name, index) for link in workflow['links']: if link[1] == saved['id']: assert link[2] < len(info['output']), (path.name, link) target = local_nodes[link[3]]['inputs'][link[4]] assert target['type'] == info['output'][link[2]], (path.name, link) checked += 1 assert checked == 5 and occurrences == 10, (checked, occurrences) print('SAVED_WORKFLOW_CONTRACTS_OK', checked, occurrences) async def check_execution(mapping, module, prompt_server, asset_manager): result = await map_node(mapping['DistributedSeed'], {'seed': 123, 'is_worker': True, 'worker_id': 'worker_2'}) assert result[0].result == (126,) values = json.dumps({'_type': 'INT', '3': '17'}) result = await map_node(mapping['DistributedValue'], {'default_value': '4', 'worker_values': values, 'is_worker': True, 'worker_id': 'worker_2'}) assert result[0].result == (17,) metadata = {'workflow': {'nodes': [{'id': 'probe', 'widgets_values': []}]}} result = await map_node(mapping['DistributedModelName'], {'text': 'model.ckpt'}, {'extra_pnginfo': metadata}) assert result[0].result == ('model.ckpt',) assert result[0].ui == {'text': ['model.ckpt']} assert metadata['workflow']['nodes'][0]['widgets_values'] == [['model.ckpt']] result = await map_node(mapping['DistributedEmptyImage'], {'height': 8, 'width': 8, 'channels': 3}) empty = result[0].result[0] assert empty.shape == (0, 8, 8, 3) and empty.numel() == 0 images = torch.arange(10 * 8 * 8 * 3, dtype=torch.float32).reshape(10, 8, 8, 3) for count in (1, 3, 10): result = await map_node(mapping['ImageBatchDivider'], {'images': images, 'divide_by': count}) assert len(result[0].result) == 10 assert torch.equal(torch.cat(result[0].result[:count]), images) audio = {'waveform': torch.arange(33, dtype=torch.float32).reshape(1, 1, 33), 'sample_rate': 24000} result = await map_node(mapping['AudioBatchDivider'], {'audio': audio, 'divide_by': 3}) assert len(result[0].result) == 10 assert torch.equal(torch.cat([item['waveform'] for item in result[0].result[:3]], dim=-1), audio['waveform']) # INPUT_IS_LIST must preserve all images/audio but unwrap transport scalars. collector = mapping['DistributedCollector'] prepared, missing, hidden = execution.get_input_data( {'load_balance': False, 'multi_job_id': '', 'pass_through': True}, collector, 'collector-list') assert not missing prepared['images'] = [images[:2], images[2:]] prepared['audio'] = [audio, audio] result = await execution._async_map_node_over_list( 'v3-acceptance', 'collector-list', collector, prepared, collector.FUNCTION, v3_data=hidden) assert torch.equal(result[0].result[0], images) assert result[0].result[1]['waveform'].shape[-1] == 66 # Exercise real async aggregation without an HTTP endpoint or worker. queue = asyncio.Queue() await queue.put({'worker_id': 'worker_2', 'tensor': images[2:3], 'image_index': 0, 'is_last': True}) prompt_server.distributed_pending_jobs['v3-aggregate'] = queue result = await map_node(collector, {'images': images[:2], 'load_balance': False, 'multi_job_id': 'v3-aggregate', 'enabled_worker_ids': '["worker_2"]'}) assert torch.equal(result[0].result[0], images[:3]) assert 'v3-aggregate' not in prompt_server.distributed_pending_jobs # The registered node must forward the transfer options to the runtime # helper, and default to the legacy transfer when they are not set. collector_runtime = sys.modules[module.__name__ + '.nodes.collector'].DistributedCollectorNode forwarded = [] def fake_run(self, **kwargs): forwarded.append(kwargs) return (kwargs['images'], kwargs['audio']) with patch.object(collector_runtime, 'run', fake_run): await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True, 'transfer_format': 'av1', 'transfer_quality': 70}) await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True}) # Set values arrive list-wrapped (is_input_list); unset ones use the # execute() defaults. The runtime helper unwraps both forms. assert [(kw['transfer_format'], kw['transfer_quality']) for kw in forwarded] == [ (['av1'], [70]), ('legacy_png', 90)], forwarded # Prove V3's per-execution class clones do not leak helper instance state; # replace only the GPU/model boundary, not the input parser or executor. upscale_runtime = sys.modules[module.__name__ + '.nodes.distributed_upscale'].UltimateSDUpscaleDistributed seen = [] def fake_upscale(self, *args): assert not hasattr(self, 'acceptance_marker') self.acceptance_marker = True seen.append(args) return (args[0],) inputs = {'upscaled_image': images, 'model': object(), 'positive': [], 'negative': [], 'vae': object(), 'seed': 1, 'steps': 1, 'cfg': 1.0, 'sampler_name': 'euler', 'scheduler': 'normal', 'denoise': 0.5, 'tile_width': 64, 'tile_height': 64, 'padding': 0, 'mask_blur': 0, 'force_uniform_tiles': True, 'tiled_decode': False, 'multi_job_id': 'tile-job', 'is_worker': True, 'master_url': 'http://master.invalid', 'enabled_worker_ids': '["worker_2"]', 'worker_id': 'worker_2', 'tile_indices': '[2,3]', 'dynamic_threshold': 9} with patch.object(upscale_runtime, 'run', fake_upscale): for _ in range(2): result = await map_node(mapping['UltimateSDUpscaleDistributed'], inputs) assert result[0].result[0] is images assert len(seen) == 2 and seen[0][-7:] == tuple(inputs[name] for name in ( 'multi_job_id', 'is_worker', 'master_url', 'enabled_worker_ids', 'worker_id', 'tile_indices', 'dynamic_threshold')), seen[0] import math assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id='tile-job')) assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id='')) graph = { '1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 10, 'color': 0}}, '2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False, 'pass_through': True}}, '3': {'class_type': 'ImageBatchDivider', 'inputs': {'images': ['2', 0], 'divide_by': 10}}, '4': {'class_type': 'PreviewImage', 'inputs': {'images': ['3', 9]}}, } valid = await execution.validate_prompt('v3-graph', graph, None) assert valid[0], valid invalid = {'1': {'class_type': 'UltimateSDUpscaleDistributed', 'inputs': {**inputs, 'sampler_name': 'INVALID_ENUM', 'scheduler': 'INVALID_ENUM'}}, '2': {'class_type': 'PreviewImage', 'inputs': {'images': ['1', 0]}}} rejected = await execution.validate_prompt('v3-invalid', invalid, None) assert not rejected[0] errors = [error for entry in rejected[3].values() for error in entry['errors']] enum_errors = [error['extra_info']['input_name'] for error in errors if error['type'] == 'value_not_in_list'] assert {'sampler_name', 'scheduler'} <= set(enum_errors), errors import folder_paths scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch')) with tempfile.TemporaryDirectory(prefix='v3-preview-', dir=scratch) as temp: with patch.object(folder_paths, 'temp_directory', temp): executor = execution.PromptExecutor( prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager) await asyncio.to_thread(executor.execute, graph, 'v3-graph', {}, valid[2]) assert executor.success, executor.status_messages history = executor.history_result record = history['outputs']['4']['images'][0] preview = Path(temp) / record.get('subfolder', '') / record['filename'] with Image.open(preview) as image: assert image.size == (8, 8) and image.mode == 'RGB' assert image.getextrema() == ((0, 0), (0, 0), (0, 0)) # A worker prompt terminated by the sink validates and runs through the # real executor without writing anything to the temp directory. sink_graph = { '1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 4, 'color': 0}}, '2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False, 'pass_through': True}}, '3': {'class_type': 'DistributedWorkerSink', 'inputs': {'value': ['2', 0]}}, } valid = await execution.validate_prompt('v3-sink', sink_graph, None) assert valid[0], valid assert valid[2] == ['3'], valid[2] collected = [] def record_run(self, **kwargs): collected.append(kwargs['images']) return (kwargs['images'][0], None) with tempfile.TemporaryDirectory(prefix='v3-sink-', dir=scratch) as temp: with patch.object(folder_paths, 'temp_directory', temp), patch.object(collector_runtime, 'run', record_run): executor = execution.PromptExecutor( prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager) await asyncio.to_thread(executor.execute, sink_graph, 'v3-sink', {}, valid[2]) assert executor.success, executor.status_messages assert len(collected) == 1 and collected[0][0].shape == (4, 8, 8, 3), collected assert list(Path(temp).rglob('*')) == [], list(Path(temp).rglob('*')) assert executor.history_result['outputs'].get('3', {}) == {}, executor.history_result print('V3_EXECUTION_OK nine nodes; upscale GPU boundary mocked; preview artifact verified; sink writes nothing') async def main(): asset_manager = default_asset_manager() prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager) assert not (ROOT / 'distributed.py').exists(), 'obsolete root bootstrap remains' assert await nodes.load_custom_node(str(ROOT)), 'ComfyUI loader rejected the pack' module = sys.modules[str(ROOT).replace('.', '_x_')] assert not hasattr(module, 'NODE_CLASS_MAPPINGS'), 'V1 map shadows V3 entrypoint' extension = await module.comfy_entrypoint() classes = await extension.get_node_list() mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes} assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES) == 9 assert set(mapping) == set(BASELINE) | ADDED_NODES for node_id in ADDED_NODES: assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id] sink_info = mapping['DistributedWorkerSink'].GET_NODE_INFO_V1() assert sink_info['output_node'] and sink_info['dev_only'], sink_info for node_id, cls in mapping.items(): if node_id in ADDED_NODES: continue assert issubclass(cls, io.ComfyNode) assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls old = dict(BASELINE[node_id]) if node_id in ('ImageBatchDivider', 'AudioBatchDivider'): # V1's ByPassTypeTuple advertises '*' when indexed, while its # underlying tuple and existing frontend declare IMAGE/AUDIO. # Native V3 declares all ten existing typed sockets explicitly. assert old['output'] == ['*'] * 10 old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10 new = json.loads(json.dumps(cls.GET_NODE_INFO_V1())) for group in ('required', 'optional'): old_inputs = old['input'].get(group, {}) new_inputs = dict(new['input'].get(group, {})) if group == 'optional': for name, expected in ADDED_OPTIONAL_INPUTS.get(node_id, {}).items(): assert name in new_inputs, (node_id, name, 'added input missing') actual = normalized_input(new_inputs.pop(name)) kind, opts = normalized_input(expected) assert actual[0] == kind, (node_id, name, actual) for key, value in opts.items(): assert actual[1].get(key) == value, (node_id, name, key, actual) assert list(old_inputs) == list(new_inputs), (node_id, group, 'input order') for name, original in old_inputs.items(): assert normalized_input(original) == normalized_input(new_inputs[name]), (node_id, name, original, new_inputs[name]) for key in ('output', 'output_name', 'output_is_list', 'is_input_list', 'output_node', 'category', 'display_name'): assert old[key] == new[key], (node_id, key, old[key], new[key]) # Standard context lives in cls.hidden; orchestrator metadata remains # accepted by its original kwarg name, without creating new widgets. signature = inspect.signature(cls.execute) for name, field in old['input'].get('hidden', {}).items(): if isinstance(field, list): assert cls.GET_SCHEMA().accept_all_inputs, node_id assert name in signature.parameters, (node_id, name) assert signature.parameters[name].default == field[1]['default'], (node_id, name) else: assert name in new['input']['hidden'], (node_id, name) expected_hidden = [] if node_id == 'DistributedModelName': expected_hidden.extend(['unique_id', 'extra_pnginfo']) if old['output_node']: expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden) assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden')) print('SCHEMA_PARITY_OK', len(mapping)) check_saved_workflows(mapping) await check_execution(mapping, module, prompt_server, asset_manager) print('V3_ACCEPTANCE_OK') asyncio.run(main())