diff --git a/__init__.py b/__init__.py index 2029bdc..ae616f6 100644 --- a/__init__.py +++ b/__init__.py @@ -1,29 +1,22 @@ -# Import everything needed from the main module -from .distributed import ( - NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS, - NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS -) +"""ComfyUI-Distributed's native V3 extension entrypoint.""" +from comfy_api.v0_0_2 import ComfyExtension, io -# Import utilities -from .utils.config import ensure_config_exists, CONFIG_FILE -from .utils.logging import debug_log +from .nodes.v3 import NODES +from .runtime.bootstrap import initialize -# Import distributed upscale nodes -from .nodes.distributed_upscale import ( - NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS, - NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS -) +WEB_DIRECTORY = './web' -WEB_DIRECTORY = "./web" -ensure_config_exists() +class DistributedExtension(ComfyExtension): + async def on_load(self) -> None: + initialize() -# Merge node mappings -NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS} -NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS} + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return list(NODES) -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] -debug_log("Loaded Distributed nodes.") -debug_log(f"Config file: {CONFIG_FILE}") -debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}") +async def comfy_entrypoint() -> DistributedExtension: + return DistributedExtension() + + +__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY'] diff --git a/conftest.py b/conftest.py index db98903..cbe5fcf 100644 --- a/conftest.py +++ b/conftest.py @@ -1,7 +1,7 @@ # conftest.py — project-level pytest configuration. # # Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports -# (from .distributed import ...) that fail when pytest tries to import it as a +# (from .nodes.v3 import ...) that fail when pytest tries to import it as a # standalone module during Package.setup() for the root package node. # # Fix: patch Package.setup() to skip the root-package's __init__.py import. diff --git a/distributed.py b/distributed.py deleted file mode 100644 index f44bc7c..0000000 --- a/distributed.py +++ /dev/null @@ -1,51 +0,0 @@ -""" -ComfyUI-Distributed: thin entry point. -All implementation lives in workers/, nodes/, api/. -""" -import atexit -import os - -import server - -from .utils.config import ensure_config_exists -from .utils.logging import debug_log -from .utils.network import cleanup_client_session -from .workers import get_worker_manager -from .workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup -from .upscale.job_store import ensure_tile_jobs_initialized -from .nodes import ( - NODE_CLASS_MAPPINGS, - NODE_DISPLAY_NAME_MAPPINGS, - ImageBatchDivider, - DistributedCollectorNode, - DistributedSeed, - DistributedModelName, - DistributedValue, - AudioBatchDivider, - DistributedEmptyImage, - AnyType, - ByPassTypeTuple, - any_type, -) -from . import api # noqa: F401 - triggers all @routes.* registrations -from .api.queue_orchestration import ensure_distributed_state - -ensure_config_exists() - -# Aiohttp session cleanup -async def _cleanup_session(): - await cleanup_client_session() - - -atexit.register(lambda: None) # placeholder; real cleanup in sync_cleanup - -# Initialize distributed job state on prompt_server -prompt_server = server.PromptServer.instance -ensure_distributed_state(prompt_server) -ensure_tile_jobs_initialized() - -# Worker startup -if not os.environ.get('COMFYUI_IS_WORKER'): - atexit.register(sync_cleanup) - delayed_auto_launch() - register_async_signals() diff --git a/docs/v3-migration.md b/docs/v3-migration.md new file mode 100644 index 0000000..6537b8f --- /dev/null +++ b/docs/v3-migration.md @@ -0,0 +1,53 @@ +# Native V3 node API + +The package registers one `ComfyExtension` through `comfy_entrypoint()` and +imports the versioned `comfy_api.v0_0_2` API. Use a ComfyUI version that provides +that API; there is no V1 registration fallback. + +## Compatibility + +- All eight node IDs, display names, categories, visible input order, defaults + and output order are retained. Existing execution algorithms remain in the + private collector, utilities and upscale modules. +- The image/audio dividers explicitly declare ten typed outputs. Their existing + frontend extensions still show only the selected number of outputs. This + replaces the V1 `ByPassTypeTuple` indexing workaround without changing saved + IMAGE/AUDIO socket indices or the ten returned values. +- Standard hidden context uses `cls.hidden`. Worker/orchestration metadata keeps + its existing prompt-input names and Python defaults via `accept_all_inputs`; + it does not become visible widgets. +- The collector retains list-input handling. Upscale retains its always-changing + fingerprint and creates a private runtime helper for each V3 execution, so + mutable helper state is not attached to sanitized V3 class clones. +- Routes, distributed state and the existing worker startup/shutdown hooks are + initialized by `runtime/bootstrap.py` during `ComfyExtension.on_load()`. + Worker mode still suppresses automatic worker launch. `distributed.py` is no + longer an entrypoint. + +## Verification + +Run ordinary unit tests from this repository: + +```bash +python -m pytest tests -q -o addopts= +``` + +Opt into real-framework acceptance using a ComfyUI checkout and its interpreter: + +```bash +COMFYUI_SOURCE_ROOT=/path/to/ComfyUI \ + /path/to/ComfyUI/.venv/bin/python -m pytest tests -q -o addopts= +``` + +The acceptance subprocess uses CPU mode and the real ComfyUI loader, input +parser, V3 class preparation, prompt validation and `PromptExecutor`. It compares +all node schemas with the V1 fixture, checks all five bundled workflows' node +IDs and socket/link contracts, exercises injected worker metadata, image/audio +lists, collector aggregation and divider outputs, rejects invalid upscale enums, +and decodes an actual preview PNG referenced by executor history. + +The upscale GPU/model boundary is mocked to check argument forwarding and +per-execution helper isolation. Full checkpoint inference, browser canvas +acceptance and multi-host HTTP transport are not exercised. No HTTP listener, +workers or model downloads are started; preview files use a temporary scratch +directory. The test does not install this branch into a live custom-node folder. diff --git a/nodes/__init__.py b/nodes/__init__.py index b7de07f..6a310c3 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -1,31 +1 @@ -from .utilities import ( - DistributedSeed, - DistributedModelName, - DistributedValue, - ImageBatchDivider, - AudioBatchDivider, - DistributedEmptyImage, - AnyType, - ByPassTypeTuple, - any_type, -) -from .collector import DistributedCollectorNode - -NODE_CLASS_MAPPINGS = { - "DistributedCollector": DistributedCollectorNode, - "DistributedSeed": DistributedSeed, - "DistributedModelName": DistributedModelName, - "DistributedValue": DistributedValue, - "ImageBatchDivider": ImageBatchDivider, - "AudioBatchDivider": AudioBatchDivider, - "DistributedEmptyImage": DistributedEmptyImage, -} -NODE_DISPLAY_NAME_MAPPINGS = { - "DistributedCollector": "Distributed Collector", - "DistributedSeed": "Distributed Seed", - "DistributedModelName": "Distributed Model Name", - "DistributedValue": "Distributed Value", - "ImageBatchDivider": "Image Batch Divider", - "AudioBatchDivider": "Audio Segment Divider", - "DistributedEmptyImage": "Distributed Empty Image", -} +"""Private execution helpers; public registration lives in nodes.v3.""" diff --git a/nodes/distributed_upscale.py b/nodes/distributed_upscale.py index 2f9b288..ef7ffd9 100644 --- a/nodes/distributed_upscale.py +++ b/nodes/distributed_upscale.py @@ -268,12 +268,3 @@ class UltimateSDUpscaleDistributed( # Ensure initialization before registering routes ensure_tile_jobs_initialized() - -# Node registration -NODE_CLASS_MAPPINGS = { - "UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)", -} diff --git a/nodes/v3.py b/nodes/v3.py new file mode 100644 index 0000000..0bce348 --- /dev/null +++ b/nodes/v3.py @@ -0,0 +1,181 @@ +"""Native V3 schemas with the existing execution algorithms kept intact. + +Runtime objects are private, per-execution helpers, not registered V1 nodes. +This avoids sharing mutable instance state through V3's sanitized class clones. +""" +import comfy.samplers +from comfy_api.v0_0_2 import io + +from . import utilities as _utilities +from .collector import DistributedCollectorNode as _CollectorRuntime +from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime + + +class DistributedSeed(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id='DistributedSeed', display_name='Distributed Seed', category='utils', + inputs=[io.Int.Input('seed', default=1125899906842, min=0, + max=1125899906842624, force_input=False)], + outputs=[io.Int.Output(display_name='seed')], + accept_all_inputs=True, + ) + + @classmethod + def execute(cls, seed, is_worker=False, worker_id=''): + return io.NodeOutput(*_utilities.DistributedSeed().distribute(seed, is_worker, worker_id)) + + +class DistributedValue(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id='DistributedValue', display_name='Distributed Value', category='utils', + inputs=[io.String.Input('default_value', default=''), + io.String.Input('worker_values', default='{}')], + outputs=[io.AnyType.Output(display_name='value')], + accept_all_inputs=True, + ) + + @classmethod + def execute(cls, default_value, worker_values='{}', is_worker=False, worker_id=''): + return io.NodeOutput(*_utilities.DistributedValue().distribute( + default_value, worker_values, is_worker, worker_id)) + + +class DistributedModelName(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id='DistributedModelName', display_name='Distributed Model Name', category='utils', + inputs=[io.String.Input('text', default='')], + outputs=[io.AnyType.Output(display_name='output')], + hidden=[io.Hidden.unique_id, io.Hidden.extra_pnginfo], is_output_node=True, + ) + + @classmethod + def execute(cls, text): + result = _utilities.DistributedModelName().log_input( + text, unique_id=cls.hidden.unique_id, extra_pnginfo=cls.hidden.extra_pnginfo) + return io.NodeOutput(*result['result'], ui=result['ui']) + + +class ImageBatchDivider(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id='ImageBatchDivider', display_name='Image Batch Divider', category='image', + inputs=[io.Image.Input('images'), + io.Int.Input('divide_by', default=2, min=1, max=10, step=1, + display_mode=io.NumberDisplay.number, + tooltip='Number of parts to divide the batch into')], + # The existing frontend still displays only divide_by sockets. + outputs=[io.Image.Output(display_name=f'batch_{index + 1}') for index in range(10)], + is_output_node=True, + ) + + @classmethod + def execute(cls, images, divide_by): + return io.NodeOutput(*_utilities.ImageBatchDivider().divide_batch(images, divide_by)) + + +class AudioBatchDivider(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id='AudioBatchDivider', display_name='Audio Segment Divider', category='audio', + inputs=[io.Audio.Input('audio'), + io.Int.Input('divide_by', default=2, min=1, max=10, step=1, + display_mode=io.NumberDisplay.number, + tooltip='Number of sequential time segments to create')], + outputs=[io.Audio.Output(display_name=f'audio_{index + 1}') for index in range(10)], + is_output_node=True, + ) + + @classmethod + def execute(cls, audio, divide_by): + return io.NodeOutput(*_utilities.AudioBatchDivider().divide_audio(audio, divide_by)) + + +class DistributedEmptyImage(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id='DistributedEmptyImage', display_name='Distributed Empty Image', category='image', + inputs=[io.Int.Input('height', default=64, min=1, max=4096, step=1), + io.Int.Input('width', default=64, min=1, max=4096, step=1), + io.Int.Input('channels', default=3, min=1, max=4, step=1)], + outputs=[io.Image.Output()], + ) + + @classmethod + def execute(cls, height, width, channels): + return io.NodeOutput(*_utilities.DistributedEmptyImage().create(height, width, channels)) + + +class DistributedCollector(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id='DistributedCollector', display_name='Distributed Collector', category='image', + inputs=[io.Boolean.Input('load_balance', default=False, + tooltip='Run this workflow on one least-busy participant (master included when participating).'), + io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True)], + outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')], + is_input_list=True, accept_all_inputs=True, + ) + + @classmethod + def execute(cls, images=None, 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): + return io.NodeOutput(*_CollectorRuntime().run( + images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id, + is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids, + worker_batch_size=worker_batch_size, worker_id=worker_id, + pass_through=pass_through, delegate_only=delegate_only)) + + +class UltimateSDUpscaleDistributed(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id='UltimateSDUpscaleDistributed', + display_name='Ultimate SD Upscale Distributed (No Upscale)', category='image/upscaling', + inputs=[ + io.Image.Input('upscaled_image'), io.Model.Input('model'), + io.Conditioning.Input('positive'), io.Conditioning.Input('negative'), io.Vae.Input('vae'), + io.Int.Input('seed', default=0, min=0, max=0xffffffffffffffff), + io.Int.Input('steps', default=20, min=1, max=10000), + io.Float.Input('cfg', default=8.0, min=0.0, max=100.0), + io.Combo.Input('sampler_name', options=comfy.samplers.KSampler.SAMPLERS), + io.Combo.Input('scheduler', options=comfy.samplers.KSampler.SCHEDULERS), + io.Float.Input('denoise', default=0.5, min=0.0, max=1.0, step=0.01), + io.Int.Input('tile_width', default=512, min=64, max=2048, step=8), + io.Int.Input('tile_height', default=512, min=64, max=2048, step=8), + io.Int.Input('padding', default=32, min=0, max=256, step=8), + io.Int.Input('mask_blur', default=8, min=0, max=256), + io.Boolean.Input('force_uniform_tiles', default=True), + io.Boolean.Input('tiled_decode', default=False), + ], outputs=[io.Image.Output()], accept_all_inputs=True, + ) + + @classmethod + def fingerprint_inputs(cls, **kwargs): + return _UpscaleRuntime.IS_CHANGED(**kwargs) + + @classmethod + def execute(cls, upscaled_image, model, positive, negative, vae, seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, padding, + mask_blur, force_uniform_tiles, tiled_decode, multi_job_id='', is_worker=False, + master_url='', enabled_worker_ids='[]', worker_id='', tile_indices='', dynamic_threshold=8): + return io.NodeOutput(*_UpscaleRuntime().run( + upscaled_image, model, positive, negative, vae, seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, padding, + mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, is_worker, + master_url, enabled_worker_ids, worker_id, tile_indices, dynamic_threshold)) + + +NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue, + ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed] diff --git a/runtime/__init__.py b/runtime/__init__.py new file mode 100644 index 0000000..0852c34 --- /dev/null +++ b/runtime/__init__.py @@ -0,0 +1 @@ +"""Internal extension lifecycle support (not a node provider).""" diff --git a/runtime/bootstrap.py b/runtime/bootstrap.py new file mode 100644 index 0000000..630b021 --- /dev/null +++ b/runtime/bootstrap.py @@ -0,0 +1,35 @@ +"""Initialize routes, distributed state and the existing worker lifecycle.""" +import atexit +import os + +import server + +from ..utils.config import CONFIG_FILE, ensure_config_exists +from ..utils.logging import debug_log +from ..workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup +from ..upscale.job_store import ensure_tile_jobs_initialized + +_initialized = False + + +def initialize(): + """Called by ComfyExtension.on_load; initialize once per loaded package.""" + global _initialized + if _initialized: + return + + ensure_config_exists() + from .. import api # noqa: F401 - registers the existing @routes.* handlers + from ..api.queue_orchestration import ensure_distributed_state + + ensure_distributed_state(server.PromptServer.instance) + ensure_tile_jobs_initialized() + + if not os.environ.get('COMFYUI_IS_WORKER'): + atexit.register(sync_cleanup) + delayed_auto_launch() + register_async_signals() + + _initialized = True + debug_log('Loaded Distributed nodes.') + debug_log(f'Config file: {CONFIG_FILE}') diff --git a/tests/fixtures/v1_node_contracts.json b/tests/fixtures/v1_node_contracts.json new file mode 100644 index 0000000..dfb6056 --- /dev/null +++ b/tests/fixtures/v1_node_contracts.json @@ -0,0 +1,691 @@ +{ + "DistributedCollector": { + "input": { + "required": { + "load_balance": [ + "BOOLEAN", + { + "default": false, + "tooltip": "Run this workflow on one least-busy participant (master included when participating)." + } + ] + }, + "optional": { + "images": [ + "IMAGE" + ], + "audio": [ + "AUDIO" + ] + }, + "hidden": { + "multi_job_id": [ + "STRING", + { + "default": "" + } + ], + "is_worker": [ + "BOOLEAN", + { + "default": false + } + ], + "master_url": [ + "STRING", + { + "default": "" + } + ], + "enabled_worker_ids": [ + "STRING", + { + "default": "[]" + } + ], + "worker_batch_size": [ + "INT", + { + "default": 1, + "min": 1, + "max": 1024 + } + ], + "worker_id": [ + "STRING", + { + "default": "" + } + ], + "pass_through": [ + "BOOLEAN", + { + "default": false + } + ], + "delegate_only": [ + "BOOLEAN", + { + "default": false + } + ] + } + }, + "input_order": { + "required": [ + "load_balance" + ], + "optional": [ + "images", + "audio" + ], + "hidden": [ + "multi_job_id", + "is_worker", + "master_url", + "enabled_worker_ids", + "worker_batch_size", + "worker_id", + "pass_through", + "delegate_only" + ] + }, + "output": [ + "IMAGE", + "AUDIO" + ], + "output_name": [ + "images", + "audio" + ], + "output_is_list": [ + false, + false + ], + "is_input_list": true, + "output_node": false, + "category": "image", + "display_name": "Distributed Collector" + }, + "DistributedSeed": { + "input": { + "required": { + "seed": [ + "INT", + { + "default": 1125899906842, + "min": 0, + "max": 1125899906842624, + "forceInput": false + } + ] + }, + "hidden": { + "is_worker": [ + "BOOLEAN", + { + "default": false + } + ], + "worker_id": [ + "STRING", + { + "default": "" + } + ] + } + }, + "input_order": { + "required": [ + "seed" + ], + "hidden": [ + "is_worker", + "worker_id" + ] + }, + "output": [ + "INT" + ], + "output_name": [ + "seed" + ], + "output_is_list": [ + false + ], + "is_input_list": false, + "output_node": false, + "category": "utils", + "display_name": "Distributed Seed" + }, + "DistributedModelName": { + "input": { + "required": { + "text": [ + "STRING", + { + "default": "" + } + ] + }, + "hidden": { + "unique_id": "UNIQUE_ID", + "extra_pnginfo": "EXTRA_PNGINFO" + } + }, + "input_order": { + "required": [ + "text" + ], + "hidden": [ + "unique_id", + "extra_pnginfo" + ] + }, + "output": [ + "*" + ], + "output_name": [ + "output" + ], + "output_is_list": [ + false + ], + "is_input_list": false, + "output_node": true, + "category": "utils", + "display_name": "Distributed Model Name" + }, + "DistributedValue": { + "input": { + "required": { + "default_value": [ + "STRING", + { + "default": "" + } + ], + "worker_values": [ + "STRING", + { + "default": "{}" + } + ] + }, + "hidden": { + "is_worker": [ + "BOOLEAN", + { + "default": false + } + ], + "worker_id": [ + "STRING", + { + "default": "" + } + ] + } + }, + "input_order": { + "required": [ + "default_value", + "worker_values" + ], + "hidden": [ + "is_worker", + "worker_id" + ] + }, + "output": [ + "*" + ], + "output_name": [ + "value" + ], + "output_is_list": [ + false + ], + "is_input_list": false, + "output_node": false, + "category": "utils", + "display_name": "Distributed Value" + }, + "ImageBatchDivider": { + "input": { + "required": { + "images": [ + "IMAGE" + ], + "divide_by": [ + "INT", + { + "default": 2, + "min": 1, + "max": 10, + "step": 1, + "display": "number", + "tooltip": "Number of parts to divide the batch into" + } + ] + } + }, + "input_order": { + "required": [ + "images", + "divide_by" + ] + }, + "output": [ + "*", + "*", + "*", + "*", + "*", + "*", + "*", + "*", + "*", + "*" + ], + "output_name": [ + "batch_1", + "batch_2", + "batch_3", + "batch_4", + "batch_5", + "batch_6", + "batch_7", + "batch_8", + "batch_9", + "batch_10" + ], + "output_is_list": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false + ], + "is_input_list": false, + "output_node": true, + "category": "image", + "display_name": "Image Batch Divider" + }, + "AudioBatchDivider": { + "input": { + "required": { + "audio": [ + "AUDIO" + ], + "divide_by": [ + "INT", + { + "default": 2, + "min": 1, + "max": 10, + "step": 1, + "display": "number", + "tooltip": "Number of sequential time segments to create" + } + ] + } + }, + "input_order": { + "required": [ + "audio", + "divide_by" + ] + }, + "output": [ + "*", + "*", + "*", + "*", + "*", + "*", + "*", + "*", + "*", + "*" + ], + "output_name": [ + "audio_1", + "audio_2", + "audio_3", + "audio_4", + "audio_5", + "audio_6", + "audio_7", + "audio_8", + "audio_9", + "audio_10" + ], + "output_is_list": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false + ], + "is_input_list": false, + "output_node": true, + "category": "audio", + "display_name": "Audio Segment Divider" + }, + "DistributedEmptyImage": { + "input": { + "required": { + "height": [ + "INT", + { + "default": 64, + "min": 1, + "max": 4096, + "step": 1 + } + ], + "width": [ + "INT", + { + "default": 64, + "min": 1, + "max": 4096, + "step": 1 + } + ], + "channels": [ + "INT", + { + "default": 3, + "min": 1, + "max": 4, + "step": 1 + } + ] + } + }, + "input_order": { + "required": [ + "height", + "width", + "channels" + ] + }, + "output": [ + "IMAGE" + ], + "output_name": [ + "IMAGE" + ], + "output_is_list": [ + false + ], + "is_input_list": false, + "output_node": false, + "category": "image", + "display_name": "Distributed Empty Image" + }, + "UltimateSDUpscaleDistributed": { + "input": { + "required": { + "upscaled_image": [ + "IMAGE" + ], + "model": [ + "MODEL" + ], + "positive": [ + "CONDITIONING" + ], + "negative": [ + "CONDITIONING" + ], + "vae": [ + "VAE" + ], + "seed": [ + "INT", + { + "default": 0, + "min": 0, + "max": 18446744073709551615 + } + ], + "steps": [ + "INT", + { + "default": 20, + "min": 1, + "max": 10000 + } + ], + "cfg": [ + "FLOAT", + { + "default": 8.0, + "min": 0.0, + "max": 100.0 + } + ], + "sampler_name": [ + [ + "euler", + "euler_cfg_pp", + "euler_ancestral", + "euler_ancestral_cfg_pp", + "heun", + "heunpp2", + "exp_heun_2_x0", + "exp_heun_2_x0_sde", + "dpm_2", + "dpm_2_ancestral", + "lms", + "dpm_fast", + "dpm_adaptive", + "dpmpp_2s_ancestral", + "dpmpp_2s_ancestral_cfg_pp", + "dpmpp_sde", + "dpmpp_sde_gpu", + "dpmpp_2m", + "dpmpp_2m_cfg_pp", + "dpmpp_2m_sde", + "dpmpp_2m_sde_gpu", + "dpmpp_2m_sde_heun", + "dpmpp_2m_sde_heun_gpu", + "dpmpp_3m_sde", + "dpmpp_3m_sde_gpu", + "ddpm", + "lcm", + "ipndm", + "ipndm_v", + "deis", + "cfgpp_ud10_ab", + "res_multistep", + "res_multistep_cfg_pp", + "res_multistep_ancestral", + "res_multistep_ancestral_cfg_pp", + "gradient_estimation", + "gradient_estimation_cfg_pp", + "er_sde", + "seeds_2", + "seeds_3", + "sa_solver", + "sa_solver_pece", + "ddim", + "uni_pc", + "uni_pc_bh2" + ] + ], + "scheduler": [ + [ + "simple", + "sgm_uniform", + "karras", + "exponential", + "ddim_uniform", + "beta", + "normal", + "linear_quadratic", + "kl_optimal" + ] + ], + "denoise": [ + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01 + } + ], + "tile_width": [ + "INT", + { + "default": 512, + "min": 64, + "max": 2048, + "step": 8 + } + ], + "tile_height": [ + "INT", + { + "default": 512, + "min": 64, + "max": 2048, + "step": 8 + } + ], + "padding": [ + "INT", + { + "default": 32, + "min": 0, + "max": 256, + "step": 8 + } + ], + "mask_blur": [ + "INT", + { + "default": 8, + "min": 0, + "max": 256 + } + ], + "force_uniform_tiles": [ + "BOOLEAN", + { + "default": true + } + ], + "tiled_decode": [ + "BOOLEAN", + { + "default": false + } + ] + }, + "hidden": { + "multi_job_id": [ + "STRING", + { + "default": "" + } + ], + "is_worker": [ + "BOOLEAN", + { + "default": false + } + ], + "master_url": [ + "STRING", + { + "default": "" + } + ], + "enabled_worker_ids": [ + "STRING", + { + "default": "[]" + } + ], + "worker_id": [ + "STRING", + { + "default": "" + } + ], + "tile_indices": [ + "STRING", + { + "default": "" + } + ], + "dynamic_threshold": [ + "INT", + { + "default": 8, + "min": 1, + "max": 64 + } + ] + } + }, + "input_order": { + "required": [ + "upscaled_image", + "model", + "positive", + "negative", + "vae", + "seed", + "steps", + "cfg", + "sampler_name", + "scheduler", + "denoise", + "tile_width", + "tile_height", + "padding", + "mask_blur", + "force_uniform_tiles", + "tiled_decode" + ], + "hidden": [ + "multi_job_id", + "is_worker", + "master_url", + "enabled_worker_ids", + "worker_id", + "tile_indices", + "dynamic_threshold" + ] + }, + "output": [ + "IMAGE" + ], + "output_name": [ + "IMAGE" + ], + "output_is_list": [ + false + ], + "is_input_list": false, + "output_node": false, + "category": "image/upscaling", + "display_name": "Ultimate SD Upscale Distributed (No Upscale)" + } +} diff --git a/tests/helpers/v3_runtime_check.py b/tests/helpers/v3_runtime_check.py new file mode 100644 index 0000000..dc148f2 --- /dev/null +++ b/tests/helpers/v3_runtime_check.py @@ -0,0 +1,252 @@ +"""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()) + + +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 + + # 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)) + print('V3_EXECUTION_OK eight nodes; upscale GPU boundary mocked; preview artifact verified') + + +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) == 8 + assert set(mapping) == set(BASELINE) + for node_id, cls in mapping.items(): + 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 = new['input'].get(group, {}) + 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()) diff --git a/tests/test_v3_migration.py b/tests/test_v3_migration.py new file mode 100644 index 0000000..f9ef1d7 --- /dev/null +++ b/tests/test_v3_migration.py @@ -0,0 +1,21 @@ +"""Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT.""" +import os +from pathlib import Path +import subprocess +import sys + +import pytest + + +def test_native_v3_runtime(): + comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT') + if not comfy_root: + pytest.skip('Set COMFYUI_SOURCE_ROOT to run native ComfyUI V3 acceptance') + helper = Path(__file__).parent / 'helpers' / 'v3_runtime_check.py' + result = subprocess.run( + [sys.executable, str(helper), str(Path(comfy_root).resolve())], + capture_output=True, text=True, timeout=120, + env={**os.environ, 'COMFYUI_IS_WORKER': '1'}, + ) + assert result.returncode == 0, result.stdout + '\n' + result.stderr + assert 'V3_ACCEPTANCE_OK' in result.stdout