Repository navigation
[modular] conditional block triggers should be workflow inputs #15100
Description
A bit of history on conditional blocks and workflows
Conditional blocks (AutoPipelineBlocks, ConditionalPipelineBlocks) are convenient to use, but they come at a cost of clarity: it is harder to figure out what is actually running. We first addressed that with get_execution_blocks, which works on any blockset with sub-blocks: give it a set of inputs and it returns the blocks that would run for them. Later, #13028 added workflows on a higher level, on the final assembled blockset only: it declares its workflows in _workflow_map, get_workflow(name) returns the blocks of one, and from_pretrained(workflow=...) and load_components(workflow=...) use that to prune blocks and components.
This leaves a gap for a conditional block whose trigger input is not also a workflow input. Let's look at a dummy example:
Three blocks: prompt_enhancer is optional, keyed on enhance_prompt; block1 and block2 each pick a text-to-image or image-to-image branch on image. Two declared workflows, t2i and i2i.
simple_get_workflow.py
from diffusers.modular_pipelines import (
AutoPipelineBlocks,
ConditionalPipelineBlocks,
ModularPipelineBlocks,
SequentialPipelineBlocks,
)
from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam
def leaf(name, inputs, outputs):
"""A leaf block that only declares inputs/outputs; it does nothing when run."""
class Block(ModularPipelineBlocks):
model_name = "dummy"
@property
def description(self):
return name
@property
def inputs(self):
return [InputParam(n) for n in inputs]
@property
def intermediate_outputs(self):
return [OutputParam(n) for n in outputs]
def __call__(self, components, state):
return components, state
Block.__name__ = Block.__qualname__ = name
return Block
PromptEnhancer = leaf("prompt_enhancer", ["prompt", "enhance_prompt"], ["prompt"])
Block1T2I = leaf("block1_t2i", ["prompt"], ["latents"])
Block1I2I = leaf("block1_i2i", ["prompt", "image"], ["latents"])
Block2T2I = leaf("block2_t2i", ["latents"], ["images"])
Block2I2I = leaf("block2_i2i", ["latents", "image"], ["images"])
class AutoPromptEnhancer(ConditionalPipelineBlocks):
block_classes = [PromptEnhancer]
block_names = ["enhance"]
block_trigger_inputs = ["enhance_prompt"]
def select_block(self, enhance_prompt=None):
return "enhance" if enhance_prompt else None
class AutoBlock1(AutoPipelineBlocks):
block_classes = [Block1I2I, Block1T2I]
block_names = ["i2i", "t2i"]
block_trigger_inputs = ["image", None]
class AutoBlock2(AutoPipelineBlocks):
block_classes = [Block2I2I, Block2T2I]
block_names = ["i2i", "t2i"]
block_trigger_inputs = ["image", None]
class Blocks(SequentialPipelineBlocks):
model_name = "dummy"
block_classes = [AutoPromptEnhancer, AutoBlock1, AutoBlock2]
block_names = ["prompt_enhancer", "block1", "block2"]
_workflow_map = {"t2i": {"prompt": True}, "i2i": {"prompt": True, "image": True}}
blocks = Blocks()
for name in blocks.available_workflows:
wf = blocks.get_workflow(name)
print(f"get_workflow({name!r}) -> one {type(wf).__name__}: {[type(b).__name__ for b in wf.sub_blocks.values()]}")
print()
print("for comparison, get_execution_blocks with enhance_prompt=True:")
for kwargs in ({"prompt": True}, {"prompt": True, "image": True}):
resolved = blocks.get_execution_blocks(enhance_prompt=True, **kwargs)
print(f" {kwargs} -> {[type(b).__name__ for b in resolved.sub_blocks.values()]}")get_workflow('t2i') -> one SequentialPipelineBlocks: ['block1_t2i', 'block2_t2i']
get_workflow('i2i') -> one SequentialPipelineBlocks: ['block1_i2i', 'block2_i2i']
for comparison, get_execution_blocks with enhance_prompt=True:
{'prompt': True} -> ['prompt_enhancer', 'block1_t2i', 'block2_t2i']
{'prompt': True, 'image': True} -> ['prompt_enhancer', 'block1_i2i', 'block2_i2i']
block1 and block2 are conditioned on image, a workflow input, so each workflow resolves them to the right branch. prompt_enhancer is conditioned on enhance_prompt, which no workflow entry mentions, so get_workflow resolves it as if the input had not been passed and drops it from both. In other words, t2i actually has two possible paths, with and without prompt_enhancer, and get_workflow returns one of them with no sign of the other. A pipeline loaded with workflow="t2i" therefore has no prompt_enhancer, none of its components, and no enhance_prompt input. But you should be able to enhance the prompt in the text-to-image workflow.
We could technically make get_workflow return both paths, but I think a simpler solution is to not let it happen: do not create conditional blocks on inputs that do not decide a workflow: a block can read the input and branch in its own code, or it can be an optional block that users insert into or pop out of their workflow.
So "workflow" is where I think we should draw the line, as a trade between convenience and clarity.
Proposed rule of thumb: a conditional block's triggers must be workflow inputs, or intermediates derived from them.
With that rule, every workflow entry decides every conditional block in the tree, so each workflow has exactly one path and get_workflow returns the whole truth.
Where we do this today
A scan of the 28 released blocksets for conditional triggers that are neither workflow inputs nor intermediates (scan_triggers.py below) finds these:
- LTX-2 / LTX-2.5
enable_prompt_enhancementand ernieuse_pe: opt-in prompt enhancers.get_workflow("text2video")on LTX-2 drops the block and itsprompt_enhancerandprocessorcomponents. - LTX-2 / LTX-2.5
conditioning_attention_strengthandconditioning_attention_mask: the in-context attention-mask step runs when the strength is below 1 or a mask is given.get_workflow("in_context")drops it, so a pipeline loaded by workflow cannot use either input. - SDXL
padding_mask_crop: the decode step picks the inpaint decoder (which pastes the crop back) onpadding_mask_croprather than onmask_image.get_workflow("inpainting")resolves to the plain decoder. - Cosmos 3
control_videos: thetransferworkflow is deliberately unavailable, so this one is known.
Under the rule the prompt enhancers become plain blocks that users pop out when they do not want them, the SDXL decode is keyed on mask_image and reads padding_mask_crop as a regular input, and the LTX-2 attention mask becomes a regular input of the in-context branch.
scan_triggers.py
"""List conditional-block triggers that are neither workflow inputs nor intermediates, for every released blockset."""
import importlib
import inspect
import pkgutil
import diffusers.modular_pipelines as mp
from diffusers.modular_pipelines.modular_pipeline import ConditionalPipelineBlocks, SequentialPipelineBlocks
def all_blocks(block):
yield block
for sub in getattr(block, "sub_blocks", {}).values():
yield from all_blocks(sub)
seen = set()
for modinfo in pkgutil.walk_packages(mp.__path__, mp.__name__ + "."):
if not modinfo.name.rsplit(".", 1)[-1].startswith("modular_blocks"):
continue
mod = importlib.import_module(modinfo.name)
for name, cls in inspect.getmembers(mod, inspect.isclass):
if cls.__module__ != mod.__name__ or not issubclass(cls, SequentialPipelineBlocks):
continue
if getattr(cls, "_workflow_map", None) is None or cls in seen:
continue
seen.add(cls)
blocks = cls()
keys = set()
for entry in blocks._workflow_map.values():
for d in entry if isinstance(entry, tuple) else (entry,):
keys |= set(d)
produced = {o.name for b in all_blocks(blocks) for o in getattr(b, "intermediate_outputs", [])}
triggers = {
t
for b in all_blocks(blocks)
if isinstance(b, ConditionalPipelineBlocks)
for t in b.block_trigger_inputs
if t
}
bad = sorted(triggers - keys - produced)
if bad:
print(f"{name}: {bad}")
print(f"checked {len(seen)} blocksets")Follow-ups
- say so in the modular guide and in the
ConditionalPipelineBlocksdocstring; - a test over the released blocksets that every trigger of a conditional block is a workflow-map key or an intermediate output of a block in the tree, which is what the scan above does;
- migrate the cases listed above.
Posted by Claude Code on behalf of @yiyixuxu, who reviewed and approved this text.
Activity