83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
"""
|
|
ComfyUI Workflow Iterator - A/B testing and batch parameter comparison nodes.
|
|
|
|
This package provides nodes for systematically testing different parameter
|
|
combinations in ComfyUI workflows. Define parameters to vary (CFG, sampler,
|
|
steps, prompts, etc.) and the system automatically runs all combinations,
|
|
producing individually-tagged images plus labeled comparison grids.
|
|
"""
|
|
|
|
import logging
|
|
|
|
logger = logging.getLogger("WorkflowIterator")
|
|
|
|
# Import node classes
|
|
from .nodes.workflow_iterator import WorkflowIteratorNode
|
|
from .nodes.parameter_input import (
|
|
WIParameterInt,
|
|
WIParameterFloat,
|
|
WIParameterString,
|
|
WIParameterCombo,
|
|
WIParameterSeed,
|
|
)
|
|
from .nodes.save_image import WISaveImage
|
|
from .nodes.grid_compositor import WIGridCompositor
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"WorkflowIterator": WorkflowIteratorNode,
|
|
"WIParameterInt": WIParameterInt,
|
|
"WIParameterFloat": WIParameterFloat,
|
|
"WIParameterString": WIParameterString,
|
|
"WIParameterCombo": WIParameterCombo,
|
|
"WIParameterSeed": WIParameterSeed,
|
|
"WISaveImage": WISaveImage,
|
|
"WIGridCompositor": WIGridCompositor,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"WorkflowIterator": "WI Workflow Iterator",
|
|
"WIParameterInt": "WI Parameter (Int)",
|
|
"WIParameterFloat": "WI Parameter (Float)",
|
|
"WIParameterString": "WI Parameter (String)",
|
|
"WIParameterCombo": "WI Parameter (Combo)",
|
|
"WIParameterSeed": "WI Parameter (Seed)",
|
|
"WISaveImage": "WI Save Image",
|
|
"WIGridCompositor": "WI Grid Compositor",
|
|
}
|
|
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
|
|
|
# Register the onprompt handler for multi-queue iteration
|
|
try:
|
|
from server import PromptServer
|
|
from .core.iteration_state import IterationStateManager
|
|
|
|
state_manager = IterationStateManager.instance()
|
|
|
|
def wi_onprompt(json_data):
|
|
"""Intercept prompt submissions to handle parameter iteration."""
|
|
try:
|
|
return state_manager.handle_prompt(json_data)
|
|
except Exception as e:
|
|
logger.error(f"Workflow Iterator onprompt error: {e}", exc_info=True)
|
|
return json_data
|
|
|
|
PromptServer.instance.add_on_prompt_handler(wi_onprompt)
|
|
|
|
# Register custom HTTP routes for batch status and control
|
|
state_manager.register_routes(PromptServer.instance.routes)
|
|
|
|
logger.info("Workflow Iterator: initialized successfully")
|
|
|
|
except ImportError:
|
|
# Running outside ComfyUI (e.g., during testing)
|
|
logger.info(
|
|
"Workflow Iterator: PromptServer not available "
|
|
"(running outside ComfyUI)"
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Workflow Iterator initialization error: {e}", exc_info=True)
|