diff --git a/docs/source/en/modular_diffusers/custom_blocks.md b/docs/source/en/modular_diffusers/custom_blocks.md index b412e0e58abc..66e1de172b34 100644 --- a/docs/source/en/modular_diffusers/custom_blocks.md +++ b/docs/source/en/modular_diffusers/custom_blocks.md @@ -332,4 +332,49 @@ Make your custom block work with Mellon's visual interface. See the [Mellon Cust Browse the [Modular Diffusers Custom Blocks](https://huggingface.co/collections/diffusers/modular-diffusers-custom-blocks) collection for inspiration and ready-to-use blocks. - \ No newline at end of file + + +## Dependencies + +Declaring package dependencies in custom blocks prevents runtime import errors later on. Diffusers validates the dependencies and returns a warning if a package is missing or incompatible. + +Set a `_requirements` attribute in your block class, mapping package names to version specifiers. + +```py +from diffusers.modular_pipelines import PipelineBlock + +class MyCustomBlock(PipelineBlock): + _requirements = { + "transformers": ">=4.44.0", + "sentencepiece": ">=0.2.0" + } +``` + +When there are blocks with different requirements, Diffusers merges their requirements. + +```py +from diffusers.modular_pipelines import SequentialPipelineBlocks + +class BlockA(PipelineBlock): + _requirements = {"transformers": ">=4.44.0"} + # ... + +class BlockB(PipelineBlock): + _requirements = {"sentencepiece": ">=0.2.0"} + # ... + +pipe = SequentialPipelineBlocks.from_blocks_dict({ + "block_a": BlockA, + "block_b": BlockB, +}) +``` + +When this block is saved with [`~ModularPipeline.save_pretrained`], the requirements are saved to the `modular_config.json` file. When this block is loaded, Diffusers checks each requirement against the current environment. If there is a mismatch or a package isn't found, Diffusers returns the following warning. + +```md +# missing package +xyz-package was specified in the requirements but wasn't found in the current environment. + +# version mismatch +xyz requirement 'specific-version' is not satisfied by the installed version 'actual-version'. Things might work unexpected. +``` diff --git a/src/diffusers/commands/custom_blocks.py b/src/diffusers/commands/custom_blocks.py index 43d9ea88577a..953240c5a2c3 100644 --- a/src/diffusers/commands/custom_blocks.py +++ b/src/diffusers/commands/custom_blocks.py @@ -89,8 +89,6 @@ def run(self): # automap = self._create_automap(parent_class=parent_class, child_class=child_class) # with open(CONFIG, "w") as f: # json.dump(automap, f) - with open("requirements.txt", "w") as f: - f.write("") def _choose_block(self, candidates, chosen=None): for cls, base in candidates: diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 8d662080124c..a563d2aa99eb 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -47,6 +47,7 @@ InputParam, InsertableDict, OutputParam, + _validate_requirements, combine_inputs, combine_outputs, format_components, @@ -297,6 +298,7 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): config_name = "modular_config.json" model_name = None + _requirements: dict[str, str] | None = None _workflow_map = None @classmethod @@ -411,6 +413,9 @@ def from_pretrained( "Selected model repository does not happear to have any custom code or does not have a valid `config.json` file." ) + if "requirements" in config and config["requirements"] is not None: + _ = _validate_requirements(config["requirements"]) + class_ref = config["auto_map"][cls.__name__] module_file, class_name = class_ref.split(".") module_file = module_file + ".py" @@ -435,8 +440,13 @@ def save_pretrained(self, save_directory, push_to_hub=False, **kwargs): module = full_mod.rsplit(".", 1)[-1].replace("__dynamic__", "") parent_module = self.save_pretrained.__func__.__qualname__.split(".", 1)[0] auto_map = {f"{parent_module}": f"{module}.{cls_name}"} - self.register_to_config(auto_map=auto_map) + + # resolve requirements + requirements = _validate_requirements(getattr(self, "_requirements", None)) + if requirements: + self.register_to_config(requirements=requirements) + self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs) config = dict(self.config) self._internal_dict = FrozenDict(config) @@ -658,6 +668,15 @@ def outputs(self) -> list[str]: combined_outputs = combine_outputs(*named_outputs) return combined_outputs + @property + # Copied from diffusers.modular_pipelines.modular_pipeline.SequentialPipelineBlocks._requirements + def _requirements(self) -> dict[str, str]: + requirements = {} + for block_name, block in self.sub_blocks.items(): + if getattr(block, "_requirements", None): + requirements[block_name] = block._requirements + return requirements + # used for `__repr__` def _get_trigger_inputs(self) -> set: """ @@ -1247,6 +1266,14 @@ def doc(self): expected_configs=self.expected_configs, ) + @property + def _requirements(self) -> dict[str, str]: + requirements = {} + for block_name, block in self.sub_blocks.items(): + if getattr(block, "_requirements", None): + requirements[block_name] = block._requirements + return requirements + class LoopSequentialPipelineBlocks(ModularPipelineBlocks): """ @@ -1385,6 +1412,15 @@ def intermediate_outputs(self) -> list[str]: def outputs(self) -> list[str]: return next(reversed(self.sub_blocks.values())).intermediate_outputs + @property + # Copied from diffusers.modular_pipelines.modular_pipeline.SequentialPipelineBlocks._requirements + def _requirements(self) -> dict[str, str]: + requirements = {} + for block_name, block in self.sub_blocks.items(): + if getattr(block, "_requirements", None): + requirements[block_name] = block._requirements + return requirements + def __init__(self): sub_blocks = InsertableDict() for block_name, block in zip(self.block_names, self.block_classes): diff --git a/src/diffusers/modular_pipelines/modular_pipeline_utils.py b/src/diffusers/modular_pipelines/modular_pipeline_utils.py index 68bb1fe2fd0c..fa82f17a9108 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline_utils.py +++ b/src/diffusers/modular_pipelines/modular_pipeline_utils.py @@ -22,10 +22,12 @@ import PIL.Image import torch +from packaging.specifiers import InvalidSpecifier, SpecifierSet from ..configuration_utils import ConfigMixin, FrozenDict from ..loaders.single_file_utils import _is_single_file_path_or_url from ..utils import DIFFUSERS_LOAD_ID_FIELDS, is_torch_available, logging +from ..utils.import_utils import _is_package_available if is_torch_available(): @@ -1020,6 +1022,89 @@ def make_doc_string( return output +def _validate_requirements(reqs): + if reqs is None: + normalized_reqs = {} + else: + if not isinstance(reqs, dict): + raise ValueError( + "Requirements must be provided as a dictionary mapping package names to version specifiers." + ) + normalized_reqs = _normalize_requirements(reqs) + + if not normalized_reqs: + return {} + + final: dict[str, str] = {} + for req, specified_ver in normalized_reqs.items(): + req_available, req_actual_ver = _is_package_available(req) + if not req_available: + logger.warning(f"{req} was specified in the requirements but wasn't found in the current environment.") + + if specified_ver: + try: + specifier = SpecifierSet(specified_ver) + except InvalidSpecifier as err: + raise ValueError(f"Requirement specifier '{specified_ver}' for {req} is invalid.") from err + + if req_actual_ver == "N/A": + logger.warning( + f"Version of {req} could not be determined to validate requirement '{specified_ver}'. Things might work unexpected." + ) + elif not specifier.contains(req_actual_ver, prereleases=True): + logger.warning( + f"{req} requirement '{specified_ver}' is not satisfied by the installed version {req_actual_ver}. Things might work unexpected." + ) + + final[req] = specified_ver + + return final + + +def _normalize_requirements(reqs): + if not reqs: + return {} + + normalized: "OrderedDict[str, str]" = OrderedDict() + + def _accumulate(mapping: dict[str, Any]): + for pkg, spec in mapping.items(): + if isinstance(spec, dict): + # This is recursive because blocks are composable. This way, we can merge requirements + # from multiple blocks. + _accumulate(spec) + continue + + pkg_name = str(pkg).strip() + if not pkg_name: + raise ValueError("Requirement package name cannot be empty.") + + spec_str = "" if spec is None else str(spec).strip() + if spec_str and not spec_str.startswith(("<", ">", "=", "!", "~")): + spec_str = f"=={spec_str}" + + existing_spec = normalized.get(pkg_name) + if existing_spec is not None: + if not existing_spec and spec_str: + normalized[pkg_name] = spec_str + elif existing_spec and spec_str and existing_spec != spec_str: + try: + combined_spec = SpecifierSet(",".join(filter(None, [existing_spec, spec_str]))) + except InvalidSpecifier: + logger.warning( + f"Conflicting requirements for '{pkg_name}' detected: '{existing_spec}' vs '{spec_str}'. Keeping '{existing_spec}'." + ) + else: + normalized[pkg_name] = str(combined_spec) + continue + + normalized[pkg_name] = spec_str + + _accumulate(reqs) + + return normalized + + def combine_inputs(*named_input_lists: list[tuple[str, list[InputParam]]]) -> list[InputParam]: """ Combines multiple lists of InputParam objects from different blocks. For duplicate inputs, updates only if current diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index 589698ffc73b..c1a402a2fd8f 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -10,6 +10,11 @@ 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,7 +24,13 @@ ) from diffusers.utils import logging -from ..testing_utils import backend_empty_cache, numpy_cosine_similarity_distance, require_accelerator, torch_device +from ..testing_utils import ( + CaptureLogger, + backend_empty_cache, + numpy_cosine_similarity_distance, + require_accelerator, + torch_device, +) class ModularPipelineTesterMixin: @@ -429,6 +440,117 @@ def test_guider_cfg(self, expected_max_diff=1e-2): 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(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(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(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(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: