From 6ebd990336c36006c5d72491c5674d6b987030c1 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 27 Feb 2026 14:59:58 +0530 Subject: [PATCH 1/4] 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/4] 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] From dc9190545e25731335847b1dc3adfb67394da954 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 27 Feb 2026 15:29:41 +0530 Subject: [PATCH 3/4] use fixture for tmp_path in modular tests. --- .../flux/test_modular_pipeline_flux.py | 29 ++++------ .../test_modular_pipelines_common.py | 55 +++++++++---------- .../test_modular_pipelines_custom_blocks.py | 28 +++++----- 3 files changed, 51 insertions(+), 61 deletions(-) diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 9a6b4b9b6fb4..f4fee82ecba5 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -14,7 +14,6 @@ # limitations under the License. import random -import tempfile import numpy as np import PIL @@ -129,18 +128,16 @@ def get_dummy_inputs(self, seed=0): return inputs - def test_save_from_pretrained(self): + def test_save_from_pretrained(self, tmp_path): pipes = [] base_pipe = self.get_pipeline().to(torch_device) pipes.append(base_pipe) - with tempfile.TemporaryDirectory() as tmpdirname: - base_pipe.save_pretrained(tmpdirname) - - pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device) - pipe.load_components(torch_dtype=torch.float32) - pipe.to(torch_device) - pipe.image_processor = VaeImageProcessor(vae_scale_factor=2) + base_pipe.save_pretrained(tmp_path) + pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device) + pipe.load_components(torch_dtype=torch.float32) + pipe.to(torch_device) + pipe.image_processor = VaeImageProcessor(vae_scale_factor=2) pipes.append(pipe) @@ -212,18 +209,16 @@ def get_dummy_inputs(self, seed=0): return inputs - def test_save_from_pretrained(self): + def test_save_from_pretrained(self, tmp_path): pipes = [] base_pipe = self.get_pipeline().to(torch_device) pipes.append(base_pipe) - with tempfile.TemporaryDirectory() as tmpdirname: - base_pipe.save_pretrained(tmpdirname) - - pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device) - pipe.load_components(torch_dtype=torch.float32) - pipe.to(torch_device) - pipe.image_processor = VaeImageProcessor(vae_scale_factor=2) + base_pipe.save_pretrained(tmp_path) + pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device) + pipe.load_components(torch_dtype=torch.float32) + pipe.to(torch_device) + pipe.image_processor = VaeImageProcessor(vae_scale_factor=2) pipes.append(pipe) diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index c94f41935938..cb355fea698d 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -1,7 +1,6 @@ import gc import json import os -import tempfile from typing import Callable import pytest @@ -330,16 +329,15 @@ def test_components_auto_cpu_offload_inference_consistent(self): assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 - def test_save_from_pretrained(self): + def test_save_from_pretrained(self, tmp_path): pipes = [] base_pipe = self.get_pipeline().to(torch_device) pipes.append(base_pipe) - with tempfile.TemporaryDirectory() as tmpdirname: - base_pipe.save_pretrained(tmpdirname) - pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device) - pipe.load_components(torch_dtype=torch.float32) - pipe.to(torch_device) + base_pipe.save_pretrained(tmp_path) + pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device) + pipe.load_components(torch_dtype=torch.float32) + pipe.to(torch_device) pipes.append(pipe) @@ -351,32 +349,31 @@ def test_save_from_pretrained(self): assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 - def test_modular_index_consistency(self): + def test_modular_index_consistency(self, tmp_path): 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) - - 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] - 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 + pipe.save_pretrained(tmp_path) + index_file = os.path.join(tmp_path, "modular_model_index.json") + assert os.path.exists(index_file) + + 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] + 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() diff --git a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py index 9c5fd5be326d..422fa3323e21 100644 --- a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py +++ b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py @@ -14,7 +14,6 @@ import json import os -import tempfile from collections import deque from typing import List @@ -153,25 +152,24 @@ def test_custom_block_output(self): output_prompt = output.values["output_prompt"] assert output_prompt.startswith("Modular diffusers + ") - def test_custom_block_saving_loading(self): + def test_custom_block_saving_loading(self, tmp_path): custom_block = DummyCustomBlockSimple() - with tempfile.TemporaryDirectory() as tmpdir: - custom_block.save_pretrained(tmpdir) - assert any("modular_config.json" in k for k in os.listdir(tmpdir)) + custom_block.save_pretrained(tmp_path) + assert any("modular_config.json" in k for k in os.listdir(tmp_path)) - with open(os.path.join(tmpdir, "modular_config.json"), "r") as f: - config = json.load(f) - auto_map = config["auto_map"] - assert auto_map == {"ModularPipelineBlocks": "test_modular_pipelines_custom_blocks.DummyCustomBlockSimple"} + with open(os.path.join(tmp_path, "modular_config.json"), "r") as f: + config = json.load(f) + auto_map = config["auto_map"] + assert auto_map == {"ModularPipelineBlocks": "test_modular_pipelines_custom_blocks.DummyCustomBlockSimple"} - # For now, the Python script that implements the custom block has to be manually pushed to the Hub. - # This is why, we have to separately save the Python script here. - code_path = os.path.join(tmpdir, "test_modular_pipelines_custom_blocks.py") - with open(code_path, "w") as f: - f.write(CODE_STR) + # For now, the Python script that implements the custom block has to be manually pushed to the Hub. + # This is why, we have to separately save the Python script here. + code_path = os.path.join(tmp_path, "test_modular_pipelines_custom_blocks.py") + with open(code_path, "w") as f: + f.write(CODE_STR) - loaded_custom_block = ModularPipelineBlocks.from_pretrained(tmpdir, trust_remote_code=True) + loaded_custom_block = ModularPipelineBlocks.from_pretrained(tmp_path, trust_remote_code=True) pipe = loaded_custom_block.init_pipeline() prompt = "Diffusers is nice" From de5878117f8bd503b54841b26e61c36b1fa7bffe Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 27 Feb 2026 15:30:23 +0530 Subject: [PATCH 4/4] remove unneeded test. --- .../test_modular_pipelines_common.py | 28 ------------------- 1 file changed, 28 deletions(-) diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index cb355fea698d..41cd5f344779 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -1,6 +1,4 @@ import gc -import json -import os from typing import Callable import pytest @@ -349,32 +347,6 @@ def test_save_from_pretrained(self, tmp_path): assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 - def test_modular_index_consistency(self, tmp_path): - pipe = self.get_pipeline() - components_spec = pipe._component_specs - components = sorted(components_spec.keys()) - - pipe.save_pretrained(tmp_path) - index_file = os.path.join(tmp_path, "modular_model_index.json") - assert os.path.exists(index_file) - - 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] - 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: