From 6ebd990336c36006c5d72491c5674d6b987030c1 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 27 Feb 2026 14:59:58 +0530 Subject: [PATCH 1/2] add a test to check modular index consistency --- .../test_modular_pipelines_common.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index e97b543ff85d..a025221d8010 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -1,4 +1,6 @@ import gc +import json +import os import tempfile from typing import Callable @@ -349,6 +351,29 @@ def test_save_from_pretrained(self): assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 + def test_modular_index_consistency(self): + pipe = self.get_pipeline() + components_spec = pipe._component_specs + components = sorted(components_spec.keys()) + + with tempfile.TemporaryDirectory() as tmpdir: + pipe.save_pretrained(tmpdir) + index_file = os.path.join(tmpdir, "modular_model_index.json") + assert os.path.exists(index_file) + + with open(index_file) as f: + index_contents = json.load(f) + + to_check_attrs = {"pretrained_model_name_or_path", "revision", "subfolder"} + for component in components: + spec = components_spec[component] + for attr in to_check_attrs: + if getattr(spec, "pretrained_model_name_or_path", None) is not None: + for attr in to_check_attrs: + assert component in index_contents, f"{component} should be present in index but isn't." + attr_value_from_index = index_contents[component][2][attr] + assert getattr(spec, attr) == attr_value_from_index + def test_workflow_map(self): blocks = self.pipeline_blocks_class() if blocks._workflow_map is None: From 94457fd6b15df98260d2444a16b3d4aaec845740 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 27 Feb 2026 15:02:17 +0530 Subject: [PATCH 2/2] check for compulsory keys. --- tests/modular_pipelines/test_modular_pipelines_common.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index a025221d8010..c94f41935938 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -364,6 +364,10 @@ def test_modular_index_consistency(self): with open(index_file) as f: index_contents = json.load(f) + compulsory_keys = {"_blocks_class_name", "_class_name", "_diffusers_version"} + for k in compulsory_keys: + assert k in index_contents + to_check_attrs = {"pretrained_model_name_or_path", "revision", "subfolder"} for component in components: spec = components_spec[component]