diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 76a850b63c4e..074ffe743975 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -1836,6 +1836,7 @@ def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = create_pr = kwargs.pop("create_pr", False) token = kwargs.pop("token", None) repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + update_model_card = kwargs.pop("update_model_card", False) repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id # Generate modular pipeline card content @@ -1848,6 +1849,7 @@ def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = is_pipeline=True, model_description=MODULAR_MODEL_CARD_TEMPLATE.format(**card_content), is_modular=True, + update_model_card=update_model_card, ) model_card = populate_model_card(model_card, tags=card_content["tags"]) diff --git a/src/diffusers/modular_pipelines/modular_pipeline_utils.py b/src/diffusers/modular_pipelines/modular_pipeline_utils.py index cab17c2aed5c..90e3307c61c2 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline_utils.py +++ b/src/diffusers/modular_pipelines/modular_pipeline_utils.py @@ -50,11 +50,7 @@ {components_description} {configs_section} -## Input/Output Specification - -### Inputs {inputs_description} - -### Outputs {outputs_description} +{io_specification_section} """ @@ -799,6 +795,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"] if header else [] + 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 +1091,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 """ @@ -1075,15 +1110,6 @@ def generate_modular_model_card_content(blocks) -> dict[str, Any]: if block_desc: blocks_desc_parts.append(f" - {block_desc}") - # add sub-blocks if any - if hasattr(block, "sub_blocks") and block.sub_blocks: - for sub_name, sub_block in block.sub_blocks.items(): - sub_class = sub_block.__class__.__name__ - sub_desc = sub_block.description.split("\n")[0] if getattr(sub_block, "description", "") else "" - blocks_desc_parts.append(f" - *{sub_name}*: `{sub_class}`") - if sub_desc: - blocks_desc_parts.append(f" - {sub_desc}") - blocks_description = "\n".join(blocks_desc_parts) if blocks_desc_parts else "No blocks defined." components = getattr(blocks, "expected_components", []) @@ -1109,63 +1135,76 @@ 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 - - # 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 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}") + # Branch on whether workflows are defined + has_workflows = getattr(blocks, "_workflow_map", None) is not 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}") - - inputs_description = "\n".join(inputs_parts) if inputs_parts else "No specific inputs defined." - - # 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}") - - outputs_description = "\n".join(outputs_parts) if outputs_parts else "Standard pipeline outputs." - - 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""" + if has_workflows: + workflow_map = blocks._workflow_map + parts = [] + + # 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 + ) + + parts.append("## Workflow Input Specification\n") + + # Per-workflow details: show trigger inputs with full param descriptions + 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"
\n{wf_name}\n") + parts.append("*Could not resolve workflow blocks.*\n") + parts.append("
\n") + continue + + wf_inputs = workflow_blocks.inputs + # Show only trigger inputs with full parameter descriptions + trigger_params = [p for p in wf_inputs if p.name in trigger_input_names] + + parts.append(f"
\n{wf_name}\n") + + inputs_str = format_params_markdown(trigger_params, header=None) + parts.append(inputs_str if inputs_str else "No additional inputs required.") + parts.append("") + + parts.append("
\n") + + # Common Inputs & Outputs section (like non-workflow pipelines) + all_inputs = blocks.inputs + all_outputs = shared_outputs if shared_outputs is not None else blocks.outputs + + inputs_str = format_params_markdown(all_inputs, "Inputs") + outputs_str = format_params_markdown(all_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." + + parts.append(f"\n## Input/Output Specification\n\n{inputs_description}\n\n{outputs_description}") + + 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 +1217,18 @@ def generate_modular_model_card_content(blocks) -> dict[str, Any]: 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 +1256,7 @@ def generate_modular_model_card_content(blocks) -> dict[str, Any]: "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, } diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index ad1ce988870c..b5eb9ab2e17f 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -107,6 +107,7 @@ def load_or_create_model_card( widget: list[dict] | None = None, inference: bool | None = None, is_modular: bool = False, + update_model_card: bool = False, ) -> ModelCard: """ Loads or creates a model card. @@ -133,6 +134,9 @@ def load_or_create_model_card( `load_or_create_model_card` from a training script. is_modular: (`bool`, optional): Boolean flag to denote if the model card is for a modular pipeline. When True, uses model_description as-is without additional template formatting. + update_model_card: (`bool`, optional): When True, regenerates the model card content even if one + already exists on the remote repo. Existing card metadata (tags, license, etc.) is preserved. Only + supported for modular pipelines (i.e., `is_modular=True`). """ if not is_jinja_available(): raise ValueError( @@ -141,9 +145,17 @@ def load_or_create_model_card( " To install it, please run `pip install Jinja2`." ) + if update_model_card and not is_modular: + raise ValueError("`update_model_card=True` is only supported for modular pipelines (`is_modular=True`).") + try: # Check if the model card is present on the remote repo model_card = ModelCard.load(repo_id_or_path, token=token) + # For modular pipelines, regenerate card content when requested (preserve existing metadata) + if update_model_card and is_modular and model_description is not None: + existing_data = model_card.data + model_card = ModelCard(model_description) + model_card.data = existing_data except (EntryNotFoundError, RepositoryNotFoundError): # Otherwise create a model card from template if from_training: diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index e97b543ff85d..4ba347c17b97 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -454,8 +454,7 @@ def test_basic_model_card_content_structure(self): "blocks_description", "components_description", "configs_section", - "inputs_description", - "outputs_description", + "io_specification_section", "trigger_inputs_section", "tags", ] @@ -552,18 +551,19 @@ def test_inputs_description_required_and_optional(self): 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.""" @@ -573,15 +573,16 @@ def test_outputs_description_formatting(self): 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."""