up

src/diffusers/modular_pipelines/modular_pipeline_utils.py

@@ -48,13 +48,7 @@ This modular pipeline is composed of the following blocks:
 
 ## Model Components
 
-{components_description} {configs_section}
-
-## Input/Output Specification
-
-### Inputs {inputs_description}
-
-### Outputs {outputs_description}
+{components_description} {configs_section} {io_specification_section}
 """
 
 
@@ -799,6 +793,46 @@ def format_output_params(output_params, indent_level=4, max_line_length=115):
     return format_params(output_params, "Outputs", indent_level, max_line_length)
 
 
+def format_params_markdown(params, header="Inputs"):
+    """Format a list of InputParam or OutputParam objects as a markdown bullet-point list.
+
+    Suitable for model cards rendered on Hugging Face Hub.
+
+    Args:
+        params: list of InputParam or OutputParam objects to format
+        header: Header text (e.g. "Inputs" or "Outputs")
+
+    Returns:
+        A formatted markdown string, or empty string if params is empty.
+    """
+    if not params:
+        return ""
+
+    def get_type_str(type_hint):
+        if isinstance(type_hint, UnionType) or get_origin(type_hint) is Union:
+            type_strs = [t.__name__ if hasattr(t, "__name__") else str(t) for t in get_args(type_hint)]
+            return " | ".join(type_strs)
+        return type_hint.__name__ if hasattr(type_hint, "__name__") else str(type_hint)
+
+    lines = [f"**{header}:**\n"]
+    for param in params:
+        type_str = get_type_str(param.type_hint) if param.type_hint != Any else ""
+        name = f"**{param.kwargs_type}" if param.name is None and param.kwargs_type is not None else param.name
+        param_str = f"- `{name}` (`{type_str}`"
+
+        if hasattr(param, "required") and not param.required:
+            param_str += ", *optional*"
+            if param.default is not None:
+                param_str += f", defaults to `{param.default}`"
+        param_str += ")"
+
+        desc = param.description if param.description else "No description provided"
+        param_str += f": {desc}"
+        lines.append(param_str)
+
+    return "\n".join(lines)
+
+
 def format_components(components, indent_level=4, max_line_length=115, add_empty_lines=True):
     """Format a list of ComponentSpec objects into a readable string representation.
 
@@ -1055,8 +1089,7 @@ def generate_modular_model_card_content(blocks) -> dict[str, Any]:
             - blocks_description: Detailed architecture of blocks
             - components_description: List of required components
             - configs_section: Configuration parameters section
-            - inputs_description: Input parameters specification
-            - outputs_description: Output parameters specification
+            - io_specification_section: Input/Output specification (per-workflow or unified)
             - trigger_inputs_section: Conditional execution information
             - tags: List of relevant tags for the model card
     """
@@ -1109,63 +1142,74 @@ def generate_modular_model_card_content(blocks) -> dict[str, Any]:
         if configs_description:
             configs_section = f"\n\n## Configuration Parameters\n\n{configs_description}"
 
-    inputs = blocks.inputs
-    outputs = blocks.outputs
+    # Branch on whether workflows are defined
+    has_workflows = getattr(blocks, "_workflow_map", None) is not None
 
-    # format inputs as markdown list
-    inputs_parts = []
-    required_inputs = [inp for inp in inputs if inp.required]
-    optional_inputs = [inp for inp in inputs if not inp.required]
+    if has_workflows:
+        # Per-workflow I/O sections
+        workflow_map = blocks._workflow_map
+        parts = []
 
-    if required_inputs:
-        inputs_parts.append("**Required:**\n")
-        for inp in required_inputs:
-            if hasattr(inp.type_hint, "__name__"):
-                type_str = inp.type_hint.__name__
-            elif inp.type_hint is not None:
-                type_str = str(inp.type_hint).replace("typing.", "")
-            else:
-                type_str = "Any"
-            desc = inp.description or "No description provided"
-            inputs_parts.append(f"- `{inp.name}` (`{type_str}`): {desc}")
+        # If blocks overrides outputs (e.g. to return just "images" instead of all intermediates),
+        # use that as the shared output for all workflows
+        blocks_outputs = blocks.outputs
+        blocks_intermediate = getattr(blocks, "intermediate_outputs", None)
+        shared_outputs = (
+            blocks_outputs if blocks_intermediate is not None and blocks_outputs != blocks_intermediate else None
+        )
 
-    if optional_inputs:
-        if required_inputs:
-            inputs_parts.append("")
-        inputs_parts.append("**Optional:**\n")
-        for inp in optional_inputs:
-            if hasattr(inp.type_hint, "__name__"):
-                type_str = inp.type_hint.__name__
-            elif inp.type_hint is not None:
-                type_str = str(inp.type_hint).replace("typing.", "")
-            else:
-                type_str = "Any"
-            desc = inp.description or "No description provided"
-            default_str = f", default: `{inp.default}`" if inp.default is not None else ""
-            inputs_parts.append(f"- `{inp.name}` (`{type_str}`){default_str}: {desc}")
+        # Summary section using existing format_workflow
+        parts.append("## Supported Workflows\n")
+        parts.append(format_workflow(workflow_map))
+        parts.append("")
 
-    inputs_description = "\n".join(inputs_parts) if inputs_parts else "No specific inputs defined."
+        # Per-workflow details
+        for wf_name, trigger_inputs in workflow_map.items():
+            trigger_input_names = set(trigger_inputs.keys())
+            try:
+                workflow_blocks = blocks.get_workflow(wf_name)
+            except Exception:
+                parts.append(f"<details>\n<summary><strong>{wf_name}</strong></summary>\n")
+                parts.append(f"> **Trigger inputs**: {', '.join(f'`{t}`' for t in trigger_input_names)}\n")
+                parts.append("*Could not resolve workflow blocks.*\n")
+                parts.append("</details>\n")
+                continue
 
-    # format outputs as markdown list
-    outputs_parts = []
-    for out in outputs:
-        if hasattr(out.type_hint, "__name__"):
-            type_str = out.type_hint.__name__
-        elif out.type_hint is not None:
-            type_str = str(out.type_hint).replace("typing.", "")
-        else:
-            type_str = "Any"
-        desc = out.description or "No description provided"
-        outputs_parts.append(f"- `{out.name}` (`{type_str}`): {desc}")
+            wf_inputs = workflow_blocks.inputs
+            wf_outputs = shared_outputs if shared_outputs is not None else workflow_blocks.outputs
 
-    outputs_description = "\n".join(outputs_parts) if outputs_parts else "Standard pipeline outputs."
+            parts.append(f"<details>\n<summary><strong>{wf_name}</strong></summary>\n")
+            parts.append(f"> **Trigger inputs**: {', '.join(f'`{t}`' for t in trigger_input_names)}\n")
 
-    trigger_inputs_section = ""
-    if hasattr(blocks, "trigger_inputs") and blocks.trigger_inputs:
-        trigger_inputs_list = sorted([t for t in blocks.trigger_inputs if t is not None])
-        if trigger_inputs_list:
-            trigger_inputs_str = ", ".join(f"`{t}`" for t in trigger_inputs_list)
-            trigger_inputs_section = f"""
+            inputs_str = format_params_markdown(wf_inputs, "Inputs")
+            parts.append(inputs_str if inputs_str else "No specific inputs defined.")
+            parts.append("")
+
+            outputs_str = format_params_markdown(wf_outputs, "Outputs")
+            parts.append(outputs_str if outputs_str else "No specific outputs defined.")
+            parts.append("")
+
+            parts.append("</details>\n")
+
+        io_specification_section = "\n".join(parts)
+        # Suppress trigger_inputs_section when workflows are shown (it's redundant)
+        trigger_inputs_section = ""
+    else:
+        # Unified I/O section (original behavior)
+        inputs = blocks.inputs
+        outputs = blocks.outputs
+        inputs_str = format_params_markdown(inputs, "Inputs")
+        outputs_str = format_params_markdown(outputs, "Outputs")
+        inputs_description = inputs_str if inputs_str else "No specific inputs defined."
+        outputs_description = outputs_str if outputs_str else "Standard pipeline outputs."
+        io_specification_section = f"## Input/Output Specification\n\n{inputs_description}\n\n{outputs_description}"
+
+        trigger_inputs_section = ""
+        if hasattr(blocks, "trigger_inputs") and blocks.trigger_inputs:
+            trigger_inputs_list = sorted([t for t in blocks.trigger_inputs if t is not None])
+            if trigger_inputs_list:
+                trigger_inputs_str = ", ".join(f"`{t}`" for t in trigger_inputs_list)
+                trigger_inputs_section = f"""
 ### Conditional Execution
 
 This pipeline contains blocks that are selected at runtime based on inputs:
@@ -1178,7 +1222,18 @@ This pipeline contains blocks that are selected at runtime based on inputs:
     if hasattr(blocks, "model_name") and blocks.model_name:
         tags.append(blocks.model_name)
 
-    if hasattr(blocks, "trigger_inputs") and blocks.trigger_inputs:
+    if has_workflows:
+        # Derive tags from workflow names
+        workflow_names = set(blocks._workflow_map.keys())
+        if any("inpainting" in wf for wf in workflow_names):
+            tags.append("inpainting")
+        if any("image2image" in wf for wf in workflow_names):
+            tags.append("image-to-image")
+        if any("controlnet" in wf for wf in workflow_names):
+            tags.append("controlnet")
+        if any("text2image" in wf for wf in workflow_names):
+            tags.append("text-to-image")
+    elif hasattr(blocks, "trigger_inputs") and blocks.trigger_inputs:
         triggers = blocks.trigger_inputs
         if any(t in triggers for t in ["mask", "mask_image"]):
             tags.append("inpainting")
@@ -1206,8 +1261,7 @@ This pipeline uses a {block_count}-block architecture that can be customized and
         "blocks_description": blocks_description,
         "components_description": components_description,
         "configs_section": configs_section,
-        "inputs_description": inputs_description,
-        "outputs_description": outputs_description,
+        "io_specification_section": io_specification_section,
         "trigger_inputs_section": trigger_inputs_section,
         "tags": tags,
     }

tests/modular_pipelines/flux/test_modular_pipeline_flux.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import random
+import tempfile
 
 import numpy as np
 import PIL
@@ -128,16 +129,18 @@ class TestFluxImg2ImgModularPipelineFast(ModularPipelineTesterMixin):
 
         return inputs
 
-    def test_save_from_pretrained(self, tmp_path):
+    def test_save_from_pretrained(self):
         pipes = []
         base_pipe = self.get_pipeline().to(torch_device)
         pipes.append(base_pipe)
 
-        base_pipe.save_pretrained(tmp_path)
-        pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
-        pipe.load_components(torch_dtype=torch.float32)
-        pipe.to(torch_device)
-        pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            base_pipe.save_pretrained(tmpdirname)
+
+            pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
+            pipe.load_components(torch_dtype=torch.float32)
+            pipe.to(torch_device)
+            pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
 
         pipes.append(pipe)
 
@@ -209,16 +212,18 @@ class TestFluxKontextModularPipelineFast(ModularPipelineTesterMixin):
 
         return inputs
 
-    def test_save_from_pretrained(self, tmp_path):
+    def test_save_from_pretrained(self):
         pipes = []
         base_pipe = self.get_pipeline().to(torch_device)
         pipes.append(base_pipe)
 
-        base_pipe.save_pretrained(tmp_path)
-        pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
-        pipe.load_components(torch_dtype=torch.float32)
-        pipe.to(torch_device)
-        pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            base_pipe.save_pretrained(tmpdirname)
+
+            pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
+            pipe.load_components(torch_dtype=torch.float32)
+            pipe.to(torch_device)
+            pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
 
         pipes.append(pipe)
 

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -1,4 +1,5 @@
 import gc
+import tempfile
 from typing import Callable
 
 import pytest
@@ -327,15 +328,16 @@ class ModularPipelineTesterMixin:
 
         assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
 
-    def test_save_from_pretrained(self, tmp_path):
+    def test_save_from_pretrained(self):
         pipes = []
         base_pipe = self.get_pipeline().to(torch_device)
         pipes.append(base_pipe)
 
-        base_pipe.save_pretrained(tmp_path)
-        pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
-        pipe.load_components(torch_dtype=torch.float32)
-        pipe.to(torch_device)
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            base_pipe.save_pretrained(tmpdirname)
+            pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
+            pipe.load_components(torch_dtype=torch.float32)
+            pipe.to(torch_device)
 
         pipes.append(pipe)
 
@@ -452,8 +454,7 @@ class TestModularModelCardContent:
             "blocks_description",
             "components_description",
             "configs_section",
-            "inputs_description",
-            "outputs_description",
+            "io_specification_section",
             "trigger_inputs_section",
             "tags",
         ]
@@ -550,18 +551,19 @@ class TestModularModelCardContent:
         blocks = self.create_mock_blocks(inputs=inputs)
         content = generate_modular_model_card_content(blocks)
 
-        assert "**Required:**" in content["inputs_description"]
-        assert "**Optional:**" in content["inputs_description"]
-        assert "prompt" in content["inputs_description"]
-        assert "num_steps" in content["inputs_description"]
-        assert "default: `50`" in content["inputs_description"]
+        io_section = content["io_specification_section"]
+        assert "**Inputs:**" in io_section
+        assert "prompt" in io_section
+        assert "num_steps" in io_section
+        assert "*optional*" in io_section
+        assert "defaults to `50`" in io_section
 
     def test_inputs_description_empty(self):
         """Test handling of pipelines without specific inputs."""
         blocks = self.create_mock_blocks(inputs=[])
         content = generate_modular_model_card_content(blocks)
 
-        assert "No specific inputs defined" in content["inputs_description"]
+        assert "No specific inputs defined" in content["io_specification_section"]
 
     def test_outputs_description_formatting(self):
         """Test that outputs are correctly formatted."""
@@ -571,15 +573,16 @@ class TestModularModelCardContent:
         blocks = self.create_mock_blocks(outputs=outputs)
         content = generate_modular_model_card_content(blocks)
 
-        assert "images" in content["outputs_description"]
-        assert "Generated images" in content["outputs_description"]
+        io_section = content["io_specification_section"]
+        assert "images" in io_section
+        assert "Generated images" in io_section
 
     def test_outputs_description_empty(self):
         """Test handling of pipelines without specific outputs."""
         blocks = self.create_mock_blocks(outputs=[])
         content = generate_modular_model_card_content(blocks)
 
-        assert "Standard pipeline outputs" in content["outputs_description"]
+        assert "Standard pipeline outputs" in content["io_specification_section"]
 
     def test_trigger_inputs_section_with_triggers(self):
         """Test that trigger inputs section is generated when present."""

tests/modular_pipelines/test_modular_pipelines_custom_blocks.py

@@ -14,6 +14,7 @@
 
 import json
 import os
+import tempfile
 from collections import deque
 from typing import List
 
@@ -152,24 +153,25 @@ class TestModularCustomBlocks:
         output_prompt = output.values["output_prompt"]
         assert output_prompt.startswith("Modular diffusers + ")
 
-    def test_custom_block_saving_loading(self, tmp_path):
+    def test_custom_block_saving_loading(self):
         custom_block = DummyCustomBlockSimple()
 
-        custom_block.save_pretrained(tmp_path)
-        assert any("modular_config.json" in k for k in os.listdir(tmp_path))
+        with tempfile.TemporaryDirectory() as tmpdir:
+            custom_block.save_pretrained(tmpdir)
+            assert any("modular_config.json" in k for k in os.listdir(tmpdir))
 
-        with open(os.path.join(tmp_path, "modular_config.json"), "r") as f:
-            config = json.load(f)
-        auto_map = config["auto_map"]
-        assert auto_map == {"ModularPipelineBlocks": "test_modular_pipelines_custom_blocks.DummyCustomBlockSimple"}
+            with open(os.path.join(tmpdir, "modular_config.json"), "r") as f:
+                config = json.load(f)
+            auto_map = config["auto_map"]
+            assert auto_map == {"ModularPipelineBlocks": "test_modular_pipelines_custom_blocks.DummyCustomBlockSimple"}
 
-        # For now, the Python script that implements the custom block has to be manually pushed to the Hub.
-        # This is why, we have to separately save the Python script here.
-        code_path = os.path.join(tmp_path, "test_modular_pipelines_custom_blocks.py")
-        with open(code_path, "w") as f:
-            f.write(CODE_STR)
+            # For now, the Python script that implements the custom block has to be manually pushed to the Hub.
+            # This is why, we have to separately save the Python script here.
+            code_path = os.path.join(tmpdir, "test_modular_pipelines_custom_blocks.py")
+            with open(code_path, "w") as f:
+                f.write(CODE_STR)
 
-        loaded_custom_block = ModularPipelineBlocks.from_pretrained(tmp_path, trust_remote_code=True)
+            loaded_custom_block = ModularPipelineBlocks.from_pretrained(tmpdir, trust_remote_code=True)
 
         pipe = loaded_custom_block.init_pipeline()
         prompt = "Diffusers is nice"