Support workflows on conditional pipeline blocks

.ai/modular.md

@@ -181,6 +181,8 @@ class AutoDenoise(ConditionalPipelineBlocks):
     default_block_name = "text2video"
 ```
 
+Conditional blocks can also declare `_workflow_map` (workflow name → trigger inputs), like the final blockset does — then `block.get_workflow("image2video")` resolves the branch for that workflow. Opt-in per block; when set, use the same workflow names as the final blockset (`test_sub_block_workflow_map` in the tester mixin checks the two levels agree), and map each name to the triggers *this* block switches on — which may be intermediates rather than user inputs (e.g. flux2's core denoise switches on `image_latents`, not `image`). A workflow where the block is skipped maps to `{}`.
+
 ## Key pattern: Checkpoint variants
 
 A different checkpoint (distilled / turbo / a variant with its own schedule) can have its own blockset mapped to it: give the variant a `ModularPipeline` subclass carrying its `default_blocks_name`, and checkpoints route to it automatically — via `_class_name` in `modular_model_index.json`, or, for repos that only ship a standard `model_index.json`, a config-keyed map fn in `MODULAR_PIPELINE_MAPPING` (see `_flux2_klein_map_fn`).

src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py

@@ -104,6 +104,10 @@ class Flux2AutoVaeEncoderStep(AutoPipelineBlocks):
     block_classes = [Flux2VaeEncoderSequentialStep]
     block_names = ["img_conditioning"]
     block_trigger_inputs = ["image"]
+    _workflow_map = {
+        "text2image": {},
+        "image_conditioned": {"image": True},
+    }
 
     @property
     def description(self):
@@ -260,6 +264,10 @@ class Flux2AutoCoreDenoiseStep(AutoPipelineBlocks):
     block_classes = [Flux2ImageConditionedCoreDenoiseStep, Flux2CoreDenoiseStep]
     block_names = ["image_conditioned", "text2image"]
     block_trigger_inputs = ["image_latents", None]
+    _workflow_map = {
+        "text2image": {},
+        "image_conditioned": {"image_latents": True},
+    }
 
     @property
     def description(self):

src/diffusers/modular_pipelines/modular_pipeline.py

@@ -385,23 +385,41 @@ def get_execution_blocks(self, **kwargs):
         """
         raise NotImplementedError(f"`get_execution_blocks` is not implemented for {self.__class__.__name__}")
 
-    # currently only SequentialPipelineBlocks support workflows
     @property
     def available_workflows(self):
         """
-        Returns a list of available workflow names. Must be implemented by subclasses that define `_workflow_map`.
+        Returns a list of available workflow names, as defined by `_workflow_map`.
         """
-        raise NotImplementedError(f"`available_workflows` is not implemented for {self.__class__.__name__}")
+        if self._workflow_map is None:
+            raise NotImplementedError(
+                f"workflows is not supported because _workflow_map is not set for {self.__class__.__name__}"
+            )
+
+        return list(self._workflow_map.keys())
 
     def get_workflow(self, workflow_name: str):
         """
-        Get the execution blocks for a specific workflow. Must be implemented by subclasses that define
-        `_workflow_map`.
+        Get the execution blocks for a specific workflow, resolved by passing the workflow's trigger inputs to
+        `get_execution_blocks`.
 
         Args:
             workflow_name: Name of the workflow to retrieve.
         """
-        raise NotImplementedError(f"`get_workflow` is not implemented for {self.__class__.__name__}")
+        if self._workflow_map is None:
+            raise NotImplementedError(
+                f"workflows is not supported because _workflow_map is not set for {self.__class__.__name__}"
+            )
+
+        if workflow_name not in self._workflow_map:
+            raise ValueError(
+                f"Workflow {workflow_name!r} not found in {self.__class__.__name__}; "
+                f"available workflows: {self.available_workflows}"
+            )
+
+        trigger_inputs = self._workflow_map[workflow_name]
+        workflow_blocks = self.get_execution_blocks(**trigger_inputs)
+
+        return workflow_blocks
 
     @classmethod
     def from_pretrained(
@@ -984,29 +1002,6 @@ def expected_configs(self):
                     expected_configs.append(config)
         return expected_configs
 
-    @property
-    def available_workflows(self):
-        if self._workflow_map is None:
-            raise NotImplementedError(
-                f"workflows is not supported because _workflow_map is not set for {self.__class__.__name__}"
-            )
-
-        return list(self._workflow_map.keys())
-
-    def get_workflow(self, workflow_name: str):
-        if self._workflow_map is None:
-            raise NotImplementedError(
-                f"workflows is not supported because _workflow_map is not set for {self.__class__.__name__}"
-            )
-
-        if workflow_name not in self._workflow_map:
-            raise ValueError(f"Workflow {workflow_name} not found in {self.__class__.__name__}")
-
-        trigger_inputs = self._workflow_map[workflow_name]
-        workflow_blocks = self.get_execution_blocks(**trigger_inputs)
-
-        return workflow_blocks
-
     @classmethod
     def from_blocks_dict(
         cls, blocks_dict: dict[str, Any], description: str | None = None
@@ -1616,6 +1611,7 @@ def __init__(
         collection: str | None = None,
         modular_config_dict: dict[str, Any] | None = None,
         config_dict: dict[str, Any] | None = None,
+        workflow: str | None = None,
         **kwargs,
     ):
         """
@@ -1640,6 +1636,8 @@ def __init__(
                 Optional ComponentsManager for managing multiple component cross different pipelines and apply
                 offloading strategies.
             collection: Optional collection name for organizing components in the ComponentsManager.
+            workflow: Optional workflow name (one of `blocks.available_workflows`); the pipeline is built from that
+                workflow's execution blocks only, so it only expects the components that workflow uses.
             **kwargs: Additional arguments passed to `load_config()` when loading pretrained configuration.
 
         Examples:
@@ -1708,6 +1706,11 @@ def __init__(
             else:
                 logger.warning(f"`blocks` is `None`, no default blocks class found for {self.__class__.__name__}")
 
+        if workflow is not None:
+            if blocks is None:
+                raise ValueError(f"`workflow={workflow!r}` was passed but no pipeline blocks could be resolved.")
+            blocks = blocks.get_workflow(workflow)
+
         self._blocks = blocks
         self._components_manager = components_manager
         self._collection = collection
@@ -1832,6 +1835,10 @@ def from_pretrained(
                 offloading strategies.
             collection (`str`, optional):`
                 Collection name for organizing components in the ComponentsManager.
+            workflow (`str`, optional):
+                Build the pipeline from a single workflow of the repo's blocks (one of
+                `blocks.available_workflows`, e.g. `"text2image"`) instead of the full blocks. The pipeline then
+                only expects — and `load_components` only loads — the components that workflow uses.
         """
         from ..pipelines.pipeline_loading_utils import _get_pipeline_class
 
@@ -2338,24 +2345,37 @@ def update_components(self, **kwargs):
             config_to_register[name] = new_value
         self.register_to_config(**config_to_register)
 
-    def load_components(self, names: list[str] | str | None = None, **kwargs):
+    def load_components(self, names: list[str] | str | None = None, workflow: str | None = None, **kwargs):
         """
         Load selected components from specs.
 
         Args:
             names: list of component names to load. If None, will load all components with
                    default_creation_method == "from_pretrained". If provided as a list or string, will load only the
                    specified components.
+            workflow: name of a workflow (one of `blocks.available_workflows`); loads the components that workflow
+                   uses and that are not loaded yet. Call again with another workflow later to load just the missing
+                   pieces. Cannot be combined with `names`.
             **kwargs: additional kwargs to be passed to `from_pretrained()`.Can be:
              - a single value to be applied to all components to be loaded, e.g. torch_dtype=torch.bfloat16
              - a dict, e.g. torch_dtype={"unet": torch.bfloat16, "default": torch.float32}
              - if potentially override ComponentSpec if passed a different loading field in kwargs, e.g.
                `pretrained_model_name_or_path`, `variant`, `revision`, etc.
-             - if potentially override ComponentSpec if passed a different loading field in kwargs, e.g.
-               `pretrained_model_name_or_path`, `variant`, `revision`, etc.
         """
+        if workflow is not None and names is not None:
+            raise ValueError("`names` and `workflow` cannot be passed together.")
 
-        if names is None:
+        if workflow is not None:
+            workflow_components = {c.name for c in self.blocks.get_workflow(workflow).expected_components}
+            names = [
+                name
+                for name in self._component_specs.keys()
+                if name in workflow_components
+                and self._component_specs[name].default_creation_method == "from_pretrained"
+                and self._component_specs[name].pretrained_model_name_or_path is not None
+                and getattr(self, name, None) is None
+            ]
+        elif names is None:
             names = [
                 name
                 for name in self._component_specs.keys()

tests/modular_pipelines/test_conditional_pipeline_blocks.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 
+import pytest
+
 from diffusers.modular_pipelines import (
     AutoPipelineBlocks,
     ConditionalPipelineBlocks,
@@ -127,6 +129,11 @@ class AutoImageBlocks(AutoPipelineBlocks):
     block_classes = [InpaintBlock, ImageToImageBlock, TextToImageBlock]
     block_names = ["inpaint", "img2img", "text2img"]
     block_trigger_inputs = ["mask", "image", None]
+    _workflow_map = {
+        "inpaint": {"mask": True},
+        "img2img": {"image": True},
+        "text2img": {},
+    }
 
     @property
     def description(self):
@@ -224,6 +231,27 @@ def test_auto_image_workflow(self):
         assert isinstance(execution, ImageToImageBlock)
 
 
+class TestConditionalPipelineBlocksWorkflows:
+    def test_available_workflows(self):
+        assert AutoImageBlocks().available_workflows == ["inpaint", "img2img", "text2img"]
+
+    def test_get_workflow_resolves_branch(self):
+        blocks = AutoImageBlocks()
+        assert isinstance(blocks.get_workflow("inpaint"), InpaintBlock)
+        assert isinstance(blocks.get_workflow("img2img"), ImageToImageBlock)
+        assert isinstance(blocks.get_workflow("text2img"), TextToImageBlock)
+
+    def test_get_workflow_unknown_name_raises(self):
+        with pytest.raises(ValueError, match="not found"):
+            AutoImageBlocks().get_workflow("video")
+
+    def test_workflow_map_not_set_raises(self):
+        with pytest.raises(NotImplementedError, match="_workflow_map is not set"):
+            ConditionalImageBlocks().available_workflows
+        with pytest.raises(NotImplementedError, match="_workflow_map is not set"):
+            ConditionalImageBlocks().get_workflow("inpaint")
+
+
 class TestConditionalPipelineBlocksStructure:
     def test_block_names_accessible(self):
         blocks = ConditionalImageBlocks()

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -10,6 +10,7 @@
 import diffusers
 from diffusers import AutoModel, ComponentsManager, ControlNetModel, ModularPipeline, ModularPipelineBlocks
 from diffusers.guiders import ClassifierFreeGuidance
+from diffusers.modular_pipelines import LoopSequentialPipelineBlocks
 from diffusers.modular_pipelines.modular_pipeline_utils import (
     ComponentSpec,
     ConfigSpec,
@@ -478,6 +479,44 @@ def test_workflow_map(self):
                     f"{actual_block.__class__.__name__}, expected {expected_class_name}"
                 )
 
+    def test_sub_block_workflow_map(self):
+        blocks = self.pipeline_blocks_class()
+        if blocks._workflow_map is None:
+            pytest.skip("Skipping test as _workflow_map is not set")
+
+        sub_blocks_with_map = {
+            name: sub_block for name, sub_block in blocks.sub_blocks.items() if sub_block._workflow_map is not None
+        }
+        if not sub_blocks_with_map:
+            pytest.skip("Skipping test as no sub-block sets _workflow_map")
+
+        for workflow_name, expected_blocks in self.expected_workflow_blocks.items():
+            for sub_name, sub_block in sub_blocks_with_map.items():
+                # sub-blocks must use the same workflow names as the top-level blocks
+                assert workflow_name in sub_block._workflow_map, (
+                    f"Workflow '{workflow_name}' is not defined in {sub_block.__class__.__name__}._workflow_map"
+                )
+
+                # the expected entries under this sub-block, with the sub-block prefix stripped
+                expected_under = [
+                    (name[len(sub_name) + 1 :] if name != sub_name else "", expected_class_name)
+                    for name, expected_class_name in expected_blocks
+                    if name == sub_name or name.startswith(f"{sub_name}.")
+                ]
+
+                resolved = sub_block.get_workflow(workflow_name)
+                if resolved is None:
+                    actual = []
+                elif resolved.sub_blocks and not isinstance(resolved, LoopSequentialPipelineBlocks):
+                    actual = [(name, block.__class__.__name__) for name, block in resolved.sub_blocks.items()]
+                else:
+                    actual = [("", resolved.__class__.__name__)]
+
+                assert actual == expected_under, (
+                    f"Workflow '{workflow_name}': sub-block '{sub_name}' resolves to {actual}, "
+                    f"expected {expected_under} from expected_workflow_blocks"
+                )
+
 
 class ModularGuiderTesterMixin:
     def test_guider_cfg(self, expected_max_diff=1e-2):