diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 6b0fd74c108e..a5695736581f 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -34,6 +34,7 @@ from ..utils.hub_utils import load_or_create_model_card, populate_model_card from .components_manager import ComponentsManager from .modular_pipeline_utils import ( + MODULAR_MODEL_CARD_TEMPLATE, ComponentSpec, ConfigSpec, InputParam, @@ -41,6 +42,7 @@ OutputParam, format_components, format_configs, + generate_modular_model_card_content, make_doc_string, ) @@ -1753,9 +1755,19 @@ def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id + # Generate modular pipeline card content + card_content = generate_modular_model_card_content(self.blocks) + # Create a new empty model card and eventually tag it - model_card = load_or_create_model_card(repo_id, token=token, is_pipeline=True) - model_card = populate_model_card(model_card) + model_card = load_or_create_model_card( + repo_id, + token=token, + is_pipeline=True, + model_description=MODULAR_MODEL_CARD_TEMPLATE.format(**card_content), + is_modular=True, + ) + model_card = populate_model_card(model_card, tags=card_content["tags"]) + model_card.save(os.path.join(save_directory, "README.md")) # YiYi TODO: maybe order the json file to make it more readable: configs first, then components diff --git a/src/diffusers/modular_pipelines/modular_pipeline_utils.py b/src/diffusers/modular_pipelines/modular_pipeline_utils.py index cab1704df63a..9e11fb7ef79b 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline_utils.py +++ b/src/diffusers/modular_pipelines/modular_pipeline_utils.py @@ -31,6 +31,30 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +# Template for modular pipeline model card description with placeholders +MODULAR_MODEL_CARD_TEMPLATE = """{model_description} + +## Example Usage + +[TODO] + +## Pipeline Architecture + +This modular pipeline is composed of the following blocks: + +{blocks_description} {trigger_inputs_section} + +## Model Components + +{components_description} {configs_section} + +## Input/Output Specification + +### Inputs {inputs_description} + +### Outputs {outputs_description} +""" + class InsertableDict(OrderedDict): def insert(self, key, value, index): @@ -916,3 +940,178 @@ def make_doc_string( output += format_output_params(outputs, indent_level=2) return output + + +def generate_modular_model_card_content(blocks) -> Dict[str, Any]: + """ + Generate model card content for a modular pipeline. + + This function creates a comprehensive model card with descriptions of the pipeline's architecture, components, + configurations, inputs, and outputs. + + Args: + blocks: The pipeline's blocks object containing all pipeline specifications + + Returns: + Dict[str, Any]: A dictionary containing formatted content sections: + - pipeline_name: Name of the pipeline + - model_description: Overall description with pipeline type + - 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 + - trigger_inputs_section: Conditional execution information + - tags: List of relevant tags for the model card + """ + blocks_class_name = blocks.__class__.__name__ + pipeline_name = blocks_class_name.replace("Blocks", " Pipeline") + description = getattr(blocks, "description", "A modular diffusion pipeline.") + + # generate blocks architecture description + blocks_desc_parts = [] + sub_blocks = getattr(blocks, "sub_blocks", None) or {} + if sub_blocks: + for i, (name, block) in enumerate(sub_blocks.items()): + block_class = block.__class__.__name__ + block_desc = block.description.split("\n")[0] if getattr(block, "description", "") else "" + blocks_desc_parts.append(f"{i + 1}. **{name}** (`{block_class}`)") + 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", []) + if components: + components_str = format_components(components, indent_level=0, add_empty_lines=False) + # remove the "Components:" header since template has its own + components_description = components_str.replace("Components:\n", "").strip() + if components_description: + # Convert to enumerated list + lines = [line.strip() for line in components_description.split("\n") if line.strip()] + enumerated_lines = [f"{i + 1}. {line}" for i, line in enumerate(lines)] + components_description = "\n".join(enumerated_lines) + else: + components_description = "No specific components required." + else: + components_description = "No specific components required. Components can be loaded dynamically." + + configs = getattr(blocks, "expected_configs", []) + configs_section = "" + if configs: + configs_str = format_configs(configs, indent_level=0, add_empty_lines=False) + configs_description = configs_str.replace("Configs:\n", "").strip() + 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}") + + 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""" +### Conditional Execution + +This pipeline contains blocks that are selected at runtime based on inputs: +- **Trigger Inputs**: {trigger_inputs_str} +""" + + # generate tags based on pipeline characteristics + tags = ["modular-diffusers", "diffusers"] + + if hasattr(blocks, "model_name") and blocks.model_name: + tags.append(blocks.model_name) + + if 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") + if any(t in triggers for t in ["image", "image_latents"]): + tags.append("image-to-image") + if any(t in triggers for t in ["control_image", "controlnet_cond"]): + tags.append("controlnet") + if not any(t in triggers for t in ["image", "mask", "image_latents", "mask_image"]): + tags.append("text-to-image") + else: + tags.append("text-to-image") + + block_count = len(blocks.sub_blocks) + model_description = f"""This is a modular diffusion pipeline built with 🧨 Diffusers' modular pipeline framework. + +**Pipeline Type**: {blocks_class_name} + +**Description**: {description} + +This pipeline uses a {block_count}-block architecture that can be customized and extended.""" + + return { + "pipeline_name": pipeline_name, + "model_description": model_description, + "blocks_description": blocks_description, + "components_description": components_description, + "configs_section": configs_section, + "inputs_description": inputs_description, + "outputs_description": outputs_description, + "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 d0b05c7d9541..58695bae1e9d 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( license: Optional[str] = None, widget: Optional[List[dict]] = None, inference: Optional[bool] = None, + is_modular: bool = False, ) -> ModelCard: """ Loads or creates a model card. @@ -131,6 +132,8 @@ def load_or_create_model_card( widget (`List[dict]`, *optional*): Widget to accompany a gallery template. inference: (`bool`, optional): Whether to turn on inference widget. Helpful when using `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. """ if not is_jinja_available(): raise ValueError( @@ -159,10 +162,14 @@ def load_or_create_model_card( ) else: card_data = ModelCardData() - component = "pipeline" if is_pipeline else "model" - if model_description is None: - model_description = f"This is the model card of a 🧨 diffusers {component} that has been pushed on the Hub. This model card has been automatically generated." - model_card = ModelCard.from_template(card_data, model_description=model_description) + if is_modular and model_description is not None: + model_card = ModelCard(model_description) + model_card.data = card_data + else: + component = "pipeline" if is_pipeline else "model" + if model_description is None: + model_description = f"This is the model card of a 🧨 diffusers {component} that has been pushed on the Hub. This model card has been automatically generated." + model_card = ModelCard.from_template(card_data, model_description=model_description) return model_card diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index 661fcc253795..a08ca2fb759c 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -8,6 +8,13 @@ import diffusers from diffusers import ComponentsManager, ModularPipeline, ModularPipelineBlocks from diffusers.guiders import ClassifierFreeGuidance +from diffusers.modular_pipelines.modular_pipeline_utils import ( + ComponentSpec, + ConfigSpec, + InputParam, + OutputParam, + generate_modular_model_card_content, +) from diffusers.utils import logging from ..testing_utils import backend_empty_cache, numpy_cosine_similarity_distance, require_accelerator, torch_device @@ -335,3 +342,239 @@ def test_guider_cfg(self, expected_max_diff=1e-2): assert out_cfg.shape == out_no_cfg.shape max_diff = torch.abs(out_cfg - out_no_cfg).max() assert max_diff > expected_max_diff, "Output with CFG must be different from normal inference" + + +class TestModularModelCardContent: + def create_mock_block(self, name="TestBlock", description="Test block description"): + class MockBlock: + def __init__(self, name, description): + self.__class__.__name__ = name + self.description = description + self.sub_blocks = {} + + return MockBlock(name, description) + + def create_mock_blocks( + self, + class_name="TestBlocks", + description="Test pipeline description", + num_blocks=2, + components=None, + configs=None, + inputs=None, + outputs=None, + trigger_inputs=None, + model_name=None, + ): + class MockBlocks: + def __init__(self): + self.__class__.__name__ = class_name + self.description = description + self.sub_blocks = {} + self.expected_components = components or [] + self.expected_configs = configs or [] + self.inputs = inputs or [] + self.outputs = outputs or [] + self.trigger_inputs = trigger_inputs + self.model_name = model_name + + blocks = MockBlocks() + + # Add mock sub-blocks + for i in range(num_blocks): + block_name = f"block_{i}" + blocks.sub_blocks[block_name] = self.create_mock_block(f"Block{i}", f"Description for block {i}") + + return blocks + + def test_basic_model_card_content_structure(self): + """Test that all expected keys are present in the output.""" + blocks = self.create_mock_blocks() + content = generate_modular_model_card_content(blocks) + + expected_keys = [ + "pipeline_name", + "model_description", + "blocks_description", + "components_description", + "configs_section", + "inputs_description", + "outputs_description", + "trigger_inputs_section", + "tags", + ] + + for key in expected_keys: + assert key in content, f"Expected key '{key}' not found in model card content" + + assert isinstance(content["tags"], list), "Tags should be a list" + + def test_pipeline_name_generation(self): + """Test that pipeline name is correctly generated from blocks class name.""" + blocks = self.create_mock_blocks(class_name="StableDiffusionBlocks") + content = generate_modular_model_card_content(blocks) + + assert content["pipeline_name"] == "StableDiffusion Pipeline" + + def test_tags_generation_text_to_image(self): + """Test that text-to-image tags are correctly generated.""" + blocks = self.create_mock_blocks(trigger_inputs=None) + content = generate_modular_model_card_content(blocks) + + assert "modular-diffusers" in content["tags"] + assert "diffusers" in content["tags"] + assert "text-to-image" in content["tags"] + + def test_tags_generation_with_trigger_inputs(self): + """Test that tags are correctly generated based on trigger inputs.""" + # Test inpainting + blocks = self.create_mock_blocks(trigger_inputs=["mask", "prompt"]) + content = generate_modular_model_card_content(blocks) + assert "inpainting" in content["tags"] + + # Test image-to-image + blocks = self.create_mock_blocks(trigger_inputs=["image", "prompt"]) + content = generate_modular_model_card_content(blocks) + assert "image-to-image" in content["tags"] + + # Test controlnet + blocks = self.create_mock_blocks(trigger_inputs=["control_image", "prompt"]) + content = generate_modular_model_card_content(blocks) + assert "controlnet" in content["tags"] + + def test_tags_with_model_name(self): + """Test that model name is included in tags when present.""" + blocks = self.create_mock_blocks(model_name="stable-diffusion-xl") + content = generate_modular_model_card_content(blocks) + + assert "stable-diffusion-xl" in content["tags"] + + def test_components_description_formatting(self): + """Test that components are correctly formatted.""" + components = [ + ComponentSpec(name="vae", description="VAE component"), + ComponentSpec(name="text_encoder", description="Text encoder component"), + ] + blocks = self.create_mock_blocks(components=components) + content = generate_modular_model_card_content(blocks) + + assert "vae" in content["components_description"] + assert "text_encoder" in content["components_description"] + # Should be enumerated + assert "1." in content["components_description"] + + def test_components_description_empty(self): + """Test handling of pipelines without components.""" + blocks = self.create_mock_blocks(components=None) + content = generate_modular_model_card_content(blocks) + + assert "No specific components required" in content["components_description"] + + def test_configs_section_with_configs(self): + """Test that configs section is generated when configs are present.""" + configs = [ + ConfigSpec(name="num_train_timesteps", default=1000, description="Number of training timesteps"), + ] + blocks = self.create_mock_blocks(configs=configs) + content = generate_modular_model_card_content(blocks) + + assert "## Configuration Parameters" in content["configs_section"] + + def test_configs_section_empty(self): + """Test that configs section is empty when no configs are present.""" + blocks = self.create_mock_blocks(configs=None) + content = generate_modular_model_card_content(blocks) + + assert content["configs_section"] == "" + + def test_inputs_description_required_and_optional(self): + """Test that required and optional inputs are correctly formatted.""" + inputs = [ + InputParam(name="prompt", type_hint=str, required=True, description="The input prompt"), + InputParam(name="num_steps", type_hint=int, required=False, default=50, description="Number of steps"), + ] + 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"] + + 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"] + + def test_outputs_description_formatting(self): + """Test that outputs are correctly formatted.""" + outputs = [ + OutputParam(name="images", type_hint=torch.Tensor, description="Generated images"), + ] + 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"] + + 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"] + + def test_trigger_inputs_section_with_triggers(self): + """Test that trigger inputs section is generated when present.""" + blocks = self.create_mock_blocks(trigger_inputs=["mask", "image"]) + content = generate_modular_model_card_content(blocks) + + assert "### Conditional Execution" in content["trigger_inputs_section"] + assert "`mask`" in content["trigger_inputs_section"] + assert "`image`" in content["trigger_inputs_section"] + + def test_trigger_inputs_section_empty(self): + """Test that trigger inputs section is empty when not present.""" + blocks = self.create_mock_blocks(trigger_inputs=None) + content = generate_modular_model_card_content(blocks) + + assert content["trigger_inputs_section"] == "" + + def test_blocks_description_with_sub_blocks(self): + """Test that blocks with sub-blocks are correctly described.""" + + class MockBlockWithSubBlocks: + def __init__(self): + self.__class__.__name__ = "ParentBlock" + self.description = "Parent block" + self.sub_blocks = { + "child1": self.create_child_block("ChildBlock1", "Child 1 description"), + "child2": self.create_child_block("ChildBlock2", "Child 2 description"), + } + + def create_child_block(self, name, desc): + class ChildBlock: + def __init__(self): + self.__class__.__name__ = name + self.description = desc + + return ChildBlock() + + blocks = self.create_mock_blocks() + blocks.sub_blocks["parent"] = MockBlockWithSubBlocks() + + content = generate_modular_model_card_content(blocks) + + assert "parent" in content["blocks_description"] + assert "child1" in content["blocks_description"] + assert "child2" in content["blocks_description"] + + def test_model_description_includes_block_count(self): + """Test that model description includes the number of blocks.""" + blocks = self.create_mock_blocks(num_blocks=5) + content = generate_modular_model_card_content(blocks) + + assert "5-block architecture" in content["model_description"]