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,38 @@ 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} 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_pretrained(
@@ -984,29 +999,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

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):