Merge branch 'main' into tests-load-components

tests/modular_pipelines/test_conditional_pipeline_blocks.py

@@ -1,242 +0,0 @@
-# Copyright 2025 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-from diffusers.modular_pipelines import (
-    AutoPipelineBlocks,
-    ConditionalPipelineBlocks,
-    InputParam,
-    ModularPipelineBlocks,
-)
-
-
-class TextToImageBlock(ModularPipelineBlocks):
-    model_name = "text2img"
-
-    @property
-    def inputs(self):
-        return [InputParam(name="prompt")]
-
-    @property
-    def intermediate_outputs(self):
-        return []
-
-    @property
-    def description(self):
-        return "text-to-image workflow"
-
-    def __call__(self, components, state):
-        block_state = self.get_block_state(state)
-        block_state.workflow = "text2img"
-        self.set_block_state(state, block_state)
-        return components, state
-
-
-class ImageToImageBlock(ModularPipelineBlocks):
-    model_name = "img2img"
-
-    @property
-    def inputs(self):
-        return [InputParam(name="prompt"), InputParam(name="image")]
-
-    @property
-    def intermediate_outputs(self):
-        return []
-
-    @property
-    def description(self):
-        return "image-to-image workflow"
-
-    def __call__(self, components, state):
-        block_state = self.get_block_state(state)
-        block_state.workflow = "img2img"
-        self.set_block_state(state, block_state)
-        return components, state
-
-
-class InpaintBlock(ModularPipelineBlocks):
-    model_name = "inpaint"
-
-    @property
-    def inputs(self):
-        return [InputParam(name="prompt"), InputParam(name="image"), InputParam(name="mask")]
-
-    @property
-    def intermediate_outputs(self):
-        return []
-
-    @property
-    def description(self):
-        return "inpaint workflow"
-
-    def __call__(self, components, state):
-        block_state = self.get_block_state(state)
-        block_state.workflow = "inpaint"
-        self.set_block_state(state, block_state)
-        return components, state
-
-
-class ConditionalImageBlocks(ConditionalPipelineBlocks):
-    block_classes = [InpaintBlock, ImageToImageBlock, TextToImageBlock]
-    block_names = ["inpaint", "img2img", "text2img"]
-    block_trigger_inputs = ["mask", "image"]
-    default_block_name = "text2img"
-
-    @property
-    def description(self):
-        return "Conditional image blocks for testing"
-
-    def select_block(self, mask=None, image=None) -> str | None:
-        if mask is not None:
-            return "inpaint"
-        if image is not None:
-            return "img2img"
-        return None  # falls back to default_block_name
-
-
-class OptionalConditionalBlocks(ConditionalPipelineBlocks):
-    block_classes = [InpaintBlock, ImageToImageBlock]
-    block_names = ["inpaint", "img2img"]
-    block_trigger_inputs = ["mask", "image"]
-    default_block_name = None  # no default; block can be skipped
-
-    @property
-    def description(self):
-        return "Optional conditional blocks (skippable)"
-
-    def select_block(self, mask=None, image=None) -> str | None:
-        if mask is not None:
-            return "inpaint"
-        if image is not None:
-            return "img2img"
-        return None
-
-
-class AutoImageBlocks(AutoPipelineBlocks):
-    block_classes = [InpaintBlock, ImageToImageBlock, TextToImageBlock]
-    block_names = ["inpaint", "img2img", "text2img"]
-    block_trigger_inputs = ["mask", "image", None]
-
-    @property
-    def description(self):
-        return "Auto image blocks for testing"
-
-
-class TestConditionalPipelineBlocksSelectBlock:
-    def test_select_block_with_mask(self):
-        blocks = ConditionalImageBlocks()
-        assert blocks.select_block(mask="something") == "inpaint"
-
-    def test_select_block_with_image(self):
-        blocks = ConditionalImageBlocks()
-        assert blocks.select_block(image="something") == "img2img"
-
-    def test_select_block_with_mask_and_image(self):
-        blocks = ConditionalImageBlocks()
-        assert blocks.select_block(mask="m", image="i") == "inpaint"
-
-    def test_select_block_no_triggers_returns_none(self):
-        blocks = ConditionalImageBlocks()
-        assert blocks.select_block() is None
-
-    def test_select_block_explicit_none_values(self):
-        blocks = ConditionalImageBlocks()
-        assert blocks.select_block(mask=None, image=None) is None
-
-
-class TestConditionalPipelineBlocksWorkflowSelection:
-    def test_default_workflow_when_no_triggers(self):
-        blocks = ConditionalImageBlocks()
-        execution = blocks.get_execution_blocks()
-        assert execution is not None
-        assert isinstance(execution, TextToImageBlock)
-
-    def test_mask_trigger_selects_inpaint(self):
-        blocks = ConditionalImageBlocks()
-        execution = blocks.get_execution_blocks(mask=True)
-        assert isinstance(execution, InpaintBlock)
-
-    def test_image_trigger_selects_img2img(self):
-        blocks = ConditionalImageBlocks()
-        execution = blocks.get_execution_blocks(image=True)
-        assert isinstance(execution, ImageToImageBlock)
-
-    def test_mask_and_image_selects_inpaint(self):
-        blocks = ConditionalImageBlocks()
-        execution = blocks.get_execution_blocks(mask=True, image=True)
-        assert isinstance(execution, InpaintBlock)
-
-    def test_skippable_block_returns_none(self):
-        blocks = OptionalConditionalBlocks()
-        execution = blocks.get_execution_blocks()
-        assert execution is None
-
-    def test_skippable_block_still_selects_when_triggered(self):
-        blocks = OptionalConditionalBlocks()
-        execution = blocks.get_execution_blocks(image=True)
-        assert isinstance(execution, ImageToImageBlock)
-
-
-class TestAutoPipelineBlocksSelectBlock:
-    def test_auto_select_mask(self):
-        blocks = AutoImageBlocks()
-        assert blocks.select_block(mask="m") == "inpaint"
-
-    def test_auto_select_image(self):
-        blocks = AutoImageBlocks()
-        assert blocks.select_block(image="i") == "img2img"
-
-    def test_auto_select_default(self):
-        blocks = AutoImageBlocks()
-        # No trigger -> returns None -> falls back to default (text2img)
-        assert blocks.select_block() is None
-
-    def test_auto_select_priority_order(self):
-        blocks = AutoImageBlocks()
-        assert blocks.select_block(mask="m", image="i") == "inpaint"
-
-
-class TestAutoPipelineBlocksWorkflowSelection:
-    def test_auto_default_workflow(self):
-        blocks = AutoImageBlocks()
-        execution = blocks.get_execution_blocks()
-        assert isinstance(execution, TextToImageBlock)
-
-    def test_auto_mask_workflow(self):
-        blocks = AutoImageBlocks()
-        execution = blocks.get_execution_blocks(mask=True)
-        assert isinstance(execution, InpaintBlock)
-
-    def test_auto_image_workflow(self):
-        blocks = AutoImageBlocks()
-        execution = blocks.get_execution_blocks(image=True)
-        assert isinstance(execution, ImageToImageBlock)
-
-
-class TestConditionalPipelineBlocksStructure:
-    def test_block_names_accessible(self):
-        blocks = ConditionalImageBlocks()
-        sub = dict(blocks.sub_blocks)
-        assert set(sub.keys()) == {"inpaint", "img2img", "text2img"}
-
-    def test_sub_block_types(self):
-        blocks = ConditionalImageBlocks()
-        sub = dict(blocks.sub_blocks)
-        assert isinstance(sub["inpaint"], InpaintBlock)
-        assert isinstance(sub["img2img"], ImageToImageBlock)
-        assert isinstance(sub["text2img"], TextToImageBlock)
-
-    def test_description(self):
-        blocks = ConditionalImageBlocks()
-        assert "Conditional" in blocks.description

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -5,10 +5,16 @@ from typing import Callable
 
 import pytest
 import torch
+from huggingface_hub import hf_hub_download
 
 import diffusers
 from diffusers import AutoModel, ComponentsManager, ModularPipeline, ModularPipelineBlocks
 from diffusers.guiders import ClassifierFreeGuidance
+from diffusers.modular_pipelines import (
+    ConditionalPipelineBlocks,
+    LoopSequentialPipelineBlocks,
+    SequentialPipelineBlocks,
+)
 from diffusers.modular_pipelines.modular_pipeline_utils import (
     ComponentSpec,
     ConfigSpec,
@@ -19,6 +25,7 @@ from diffusers.modular_pipelines.modular_pipeline_utils import (
 from diffusers.utils import logging
 
 from ..testing_utils import (
+    CaptureLogger,
     backend_empty_cache,
     numpy_cosine_similarity_distance,
     require_accelerator,
@@ -26,6 +33,33 @@ from ..testing_utils import (
 )
 
 
+def _get_specified_components(path_or_repo_id, cache_dir=None):
+    if os.path.isdir(path_or_repo_id):
+        config_path = os.path.join(path_or_repo_id, "modular_model_index.json")
+    else:
+        try:
+            config_path = hf_hub_download(
+                repo_id=path_or_repo_id,
+                filename="modular_model_index.json",
+                local_dir=cache_dir,
+            )
+        except Exception:
+            return None
+
+    with open(config_path) as f:
+        config = json.load(f)
+
+    components = set()
+    for k, v in config.items():
+        if isinstance(v, (str, int, float, bool)):
+            continue
+        for entry in v:
+            if isinstance(entry, dict) and (entry.get("repo") or entry.get("pretrained_model_name_or_path")):
+                components.add(k)
+                break
+    return components
+
+
 class ModularPipelineTesterMixin:
     """
     It provides a set of common tests for each modular pipeline,
@@ -354,6 +388,39 @@ class ModularPipelineTesterMixin:
 
         assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
 
+    def test_load_expected_components_from_pretrained(self, tmp_path):
+        pipe = self.get_pipeline()
+        expected = _get_specified_components(self.pretrained_model_name_or_path, cache_dir=tmp_path)
+        if not expected:
+            pytest.skip("Skipping test as we couldn't fetch the expected components.")
+
+        actual = {
+            name
+            for name in pipe.components
+            if getattr(pipe, name, None) is not None
+            and getattr(getattr(pipe, name), "_diffusers_load_id", None) not in (None, "null")
+        }
+        assert expected == actual, f"Component mismatch: missing={expected - actual}, unexpected={actual - expected}"
+
+    def test_load_expected_components_from_save_pretrained(self, tmp_path):
+        pipe = self.get_pipeline()
+        save_dir = str(tmp_path / "saved-pipeline")
+        pipe.save_pretrained(save_dir)
+
+        expected = _get_specified_components(save_dir)
+        loaded_pipe = ModularPipeline.from_pretrained(save_dir)
+        loaded_pipe.load_components(torch_dtype=torch.float32)
+
+        actual = {
+            name
+            for name in loaded_pipe.components
+            if getattr(loaded_pipe, name, None) is not None
+            and getattr(getattr(loaded_pipe, name), "_diffusers_load_id", None) not in (None, "null")
+        }
+        assert expected == actual, (
+            f"Component mismatch after save/load: missing={expected - actual}, unexpected={actual - expected}"
+        )
+
     def test_modular_index_consistency(self, tmp_path):
         pipe = self.get_pipeline()
         components_spec = pipe._component_specs
@@ -431,6 +498,117 @@ class ModularGuiderTesterMixin:
         assert max_diff > expected_max_diff, "Output with CFG must be different from normal inference"
 
 
+class TestCustomBlockRequirements:
+    def get_dummy_block_pipe(self):
+        class DummyBlockOne:
+            # keep two arbitrary deps so that we can test warnings.
+            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
+
+        class DummyBlockTwo:
+            # keep two dependencies that will be available during testing.
+            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
+
+        pipe = SequentialPipelineBlocks.from_blocks_dict(
+            {"dummy_block_one": DummyBlockOne, "dummy_block_two": DummyBlockTwo}
+        )
+        return pipe
+
+    def get_dummy_conditional_block_pipe(self):
+        class DummyBlockOne:
+            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
+
+        class DummyBlockTwo:
+            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
+
+        class DummyConditionalBlocks(ConditionalPipelineBlocks):
+            block_classes = [DummyBlockOne, DummyBlockTwo]
+            block_names = ["block_one", "block_two"]
+            block_trigger_inputs = []
+
+            def select_block(self, **kwargs):
+                return "block_one"
+
+        return DummyConditionalBlocks()
+
+    def get_dummy_loop_block_pipe(self):
+        class DummyBlockOne:
+            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
+
+        class DummyBlockTwo:
+            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
+
+        return LoopSequentialPipelineBlocks.from_blocks_dict({"block_one": DummyBlockOne, "block_two": DummyBlockTwo})
+
+    def test_sequential_block_requirements_save_load(self, tmp_path):
+        pipe = self.get_dummy_block_pipe()
+        pipe.save_pretrained(str(tmp_path))
+
+        config_path = tmp_path / "modular_config.json"
+
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        assert "requirements" in config
+        requirements = config["requirements"]
+
+        expected_requirements = {
+            "xyz": ">=0.8.0",
+            "abc": ">=10.0.0",
+            "transformers": ">=4.44.0",
+            "diffusers": ">=0.2.0",
+        }
+        assert expected_requirements == requirements
+
+    def test_sequential_block_requirements_warnings(self, tmp_path):
+        pipe = self.get_dummy_block_pipe()
+
+        logger = logging.get_logger("diffusers.modular_pipelines.modular_pipeline_utils")
+        logger.setLevel(30)
+
+        with CaptureLogger(logger) as cap_logger:
+            pipe.save_pretrained(str(tmp_path))
+
+        template = "{req} was specified in the requirements but wasn't found in the current environment"
+        msg_xyz = template.format(req="xyz")
+        msg_abc = template.format(req="abc")
+        assert msg_xyz in str(cap_logger.out)
+        assert msg_abc in str(cap_logger.out)
+
+    def test_conditional_block_requirements_save_load(self, tmp_path):
+        pipe = self.get_dummy_conditional_block_pipe()
+        pipe.save_pretrained(str(tmp_path))
+
+        config_path = tmp_path / "modular_config.json"
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        assert "requirements" in config
+        expected_requirements = {
+            "xyz": ">=0.8.0",
+            "abc": ">=10.0.0",
+            "transformers": ">=4.44.0",
+            "diffusers": ">=0.2.0",
+        }
+        assert expected_requirements == config["requirements"]
+
+    def test_loop_block_requirements_save_load(self, tmp_path):
+        pipe = self.get_dummy_loop_block_pipe()
+        pipe.save_pretrained(str(tmp_path))
+
+        config_path = tmp_path / "modular_config.json"
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        assert "requirements" in config
+        expected_requirements = {
+            "xyz": ">=0.8.0",
+            "abc": ">=10.0.0",
+            "transformers": ">=4.44.0",
+            "diffusers": ">=0.2.0",
+        }
+        assert expected_requirements == config["requirements"]
+
+
 class TestModularModelCardContent:
     def create_mock_block(self, name="TestBlock", description="Test block description"):
         class MockBlock:

tests/modular_pipelines/test_modular_pipelines_custom_blocks.py

@@ -24,18 +24,14 @@ import torch
 from diffusers import FluxTransformer2DModel
 from diffusers.modular_pipelines import (
     ComponentSpec,
-    ConditionalPipelineBlocks,
     InputParam,
-    LoopSequentialPipelineBlocks,
     ModularPipelineBlocks,
     OutputParam,
     PipelineState,
-    SequentialPipelineBlocks,
     WanModularPipeline,
 )
-from diffusers.utils import logging
 
-from ..testing_utils import CaptureLogger, nightly, require_torch, slow
+from ..testing_utils import nightly, require_torch, slow
 
 
 class DummyCustomBlockSimple(ModularPipelineBlocks):
@@ -358,117 +354,6 @@ class TestModularCustomBlocks:
         assert output_prompt.startswith("Modular diffusers + ")
 
 
-class TestCustomBlockRequirements:
-    def get_dummy_block_pipe(self):
-        class DummyBlockOne:
-            # keep two arbitrary deps so that we can test warnings.
-            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
-
-        class DummyBlockTwo:
-            # keep two dependencies that will be available during testing.
-            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
-
-        pipe = SequentialPipelineBlocks.from_blocks_dict(
-            {"dummy_block_one": DummyBlockOne, "dummy_block_two": DummyBlockTwo}
-        )
-        return pipe
-
-    def get_dummy_conditional_block_pipe(self):
-        class DummyBlockOne:
-            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
-
-        class DummyBlockTwo:
-            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
-
-        class DummyConditionalBlocks(ConditionalPipelineBlocks):
-            block_classes = [DummyBlockOne, DummyBlockTwo]
-            block_names = ["block_one", "block_two"]
-            block_trigger_inputs = []
-
-            def select_block(self, **kwargs):
-                return "block_one"
-
-        return DummyConditionalBlocks()
-
-    def get_dummy_loop_block_pipe(self):
-        class DummyBlockOne:
-            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
-
-        class DummyBlockTwo:
-            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
-
-        return LoopSequentialPipelineBlocks.from_blocks_dict({"block_one": DummyBlockOne, "block_two": DummyBlockTwo})
-
-    def test_sequential_block_requirements_save_load(self, tmp_path):
-        pipe = self.get_dummy_block_pipe()
-        pipe.save_pretrained(str(tmp_path))
-
-        config_path = tmp_path / "modular_config.json"
-
-        with open(config_path, "r") as f:
-            config = json.load(f)
-
-        assert "requirements" in config
-        requirements = config["requirements"]
-
-        expected_requirements = {
-            "xyz": ">=0.8.0",
-            "abc": ">=10.0.0",
-            "transformers": ">=4.44.0",
-            "diffusers": ">=0.2.0",
-        }
-        assert expected_requirements == requirements
-
-    def test_sequential_block_requirements_warnings(self, tmp_path):
-        pipe = self.get_dummy_block_pipe()
-
-        logger = logging.get_logger("diffusers.modular_pipelines.modular_pipeline_utils")
-        logger.setLevel(30)
-
-        with CaptureLogger(logger) as cap_logger:
-            pipe.save_pretrained(str(tmp_path))
-
-        template = "{req} was specified in the requirements but wasn't found in the current environment"
-        msg_xyz = template.format(req="xyz")
-        msg_abc = template.format(req="abc")
-        assert msg_xyz in str(cap_logger.out)
-        assert msg_abc in str(cap_logger.out)
-
-    def test_conditional_block_requirements_save_load(self, tmp_path):
-        pipe = self.get_dummy_conditional_block_pipe()
-        pipe.save_pretrained(str(tmp_path))
-
-        config_path = tmp_path / "modular_config.json"
-        with open(config_path, "r") as f:
-            config = json.load(f)
-
-        assert "requirements" in config
-        expected_requirements = {
-            "xyz": ">=0.8.0",
-            "abc": ">=10.0.0",
-            "transformers": ">=4.44.0",
-            "diffusers": ">=0.2.0",
-        }
-        assert expected_requirements == config["requirements"]
-
-    def test_loop_block_requirements_save_load(self, tmp_path):
-        pipe = self.get_dummy_loop_block_pipe()
-        pipe.save_pretrained(str(tmp_path))
-
-        config_path = tmp_path / "modular_config.json"
-        with open(config_path, "r") as f:
-            config = json.load(f)
-
-        assert "requirements" in config
-        expected_requirements = {
-            "xyz": ">=0.8.0",
-            "abc": ">=10.0.0",
-            "transformers": ">=4.44.0",
-            "diffusers": ">=0.2.0",
-        }
-        assert expected_requirements == config["requirements"]
-
-
 @slow
 @nightly
 @require_torch