From e522883e44feb78c6e20d5558ab212451034efa3 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 9 Oct 2026 06:04:34 +0000 Subject: [PATCH] set overwrite_modular_index=True in tests --- .../anima/test_modular_pipeline_anima.py | 2 +- .../cosmos/test_modular_pipeline_cosmos3.py | 2 +- .../test_modular_pipeline_cosmos3_distilled.py | 2 +- .../flux/test_modular_pipeline_flux.py | 4 ++-- .../test_modular_pipeline_minimax_music3.py | 2 +- .../test_modular_pipeline_stable_diffusion_3.py | 8 ++++---- .../test_modular_pipeline_loading.py | 17 ++++++++++++++--- .../test_modular_pipelines_custom_blocks.py | 2 +- .../modular_pipelines/testing_utils/loading.py | 4 ++-- 9 files changed, 27 insertions(+), 16 deletions(-) diff --git a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py index c3ae0d9f2e0d..5f2e239af706 100644 --- a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py +++ b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py @@ -146,7 +146,7 @@ def test_save_load_components(self): pipe = self.get_pipeline() with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, safe_serialization=True) + pipe.save_pretrained(tmpdir, safe_serialization=True, overwrite_modular_index=True) pipe = self.pipeline_class.from_pretrained(tmpdir) pipe.load_components() diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index c3bcb267758f..2ed5e4cf8d57 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -514,7 +514,7 @@ def test_set_timesteps_native_flow_schedule(self): class TestCosmos3OmniModularPipelineLoading(Cosmos3OmniModularPipelineTesterConfig, ModularLoadingTesterMixin): def test_save_from_pretrained(self, tmp_path): base_pipe = self.get_pipeline().to(torch_device) - base_pipe.save_pretrained(str(tmp_path)) + base_pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) loaded_pipe = ModularPipeline.from_pretrained(str(tmp_path)) loaded_pipe.load_components(dtype=torch.float32) diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py index 86bc6ad4d6f2..c41906bd7a81 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -205,7 +205,7 @@ class TestCosmos3DistilledModularPipelineLoading( ): def test_save_from_pretrained(self, tmp_path): base_pipe = self.get_pipeline().to(torch_device) - base_pipe.save_pretrained(str(tmp_path)) + base_pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) loaded_pipe = ModularPipeline.from_pretrained(str(tmp_path)) loaded_pipe.load_components(torch_dtype=torch.float32) diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 502df3f447bf..deda31afac6e 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -162,7 +162,7 @@ def test_float16_inference(self): class TestFluxImg2ImgModularPipelineLoading(FluxImg2ImgModularPipelineTesterConfig, ModularLoadingTesterMixin): def test_save_from_pretrained(self, tmp_path, base_pipe_output): base_pipe = self.get_pipeline().to(torch_device) - base_pipe.save_pretrained(str(tmp_path)) + base_pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) pipe = ModularPipeline.from_pretrained(tmp_path) pipe.load_components(dtype=torch.float32) @@ -249,7 +249,7 @@ def test_float16_inference(self): class TestFluxKontextModularPipelineLoading(FluxKontextModularPipelineTesterConfig, ModularLoadingTesterMixin): def test_save_from_pretrained(self, tmp_path, base_pipe_output): base_pipe = self.get_pipeline().to(torch_device) - base_pipe.save_pretrained(str(tmp_path)) + base_pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) pipe = ModularPipeline.from_pretrained(tmp_path) pipe.load_components(dtype=torch.float32) diff --git a/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py b/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py index 923c4704dd7b..58a9c155aaa3 100644 --- a/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py +++ b/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py @@ -139,7 +139,7 @@ class TestMiniMaxMusic3ModularPipelineLoading(MiniMaxMusic3ModularPipelineTester def test_save_from_pretrained(self, tmp_path, base_pipe_output): # the common implementation indexes 4-D image outputs; compare the audio waveform directly base_pipe = self.get_pipeline().to(torch_device) - base_pipe.save_pretrained(str(tmp_path)) + base_pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) pipe = ModularPipeline.from_pretrained(str(tmp_path)) pipe.load_components(dtype=torch.float32) diff --git a/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py b/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py index 03d0c5687794..45161539ecf2 100644 --- a/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py +++ b/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py @@ -86,7 +86,7 @@ def test_save_from_pretrained(self, tmp_path): base_pipe = self.get_pipeline().to(torch_device) pipes.append(base_pipe) - base_pipe.save_pretrained(str(tmp_path)) + base_pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) pipe = self.pipeline_class.from_pretrained(tmp_path).to(torch_device) pipe.load_components(dtype=torch.float32) pipe.to(torch_device) @@ -102,7 +102,7 @@ def test_save_from_pretrained(self, tmp_path): def test_load_expected_components_from_save_pretrained(self, tmp_path): base_pipe = self.get_pipeline() - base_pipe.save_pretrained(str(tmp_path)) + base_pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) pipe = self.pipeline_class.from_pretrained(tmp_path) pipe.load_components(dtype=torch.float32) @@ -191,7 +191,7 @@ def test_save_from_pretrained(self, tmp_path): base_pipe = self.get_pipeline().to(torch_device) pipes.append(base_pipe) - base_pipe.save_pretrained(str(tmp_path)) + base_pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) pipe = self.pipeline_class.from_pretrained(tmp_path).to(torch_device) pipe.load_components(dtype=torch.float32) pipe.to(torch_device) @@ -208,7 +208,7 @@ def test_save_from_pretrained(self, tmp_path): def test_load_expected_components_from_save_pretrained(self, tmp_path): base_pipe = self.get_pipeline() - base_pipe.save_pretrained(str(tmp_path)) + base_pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) pipe = self.pipeline_class.from_pretrained(tmp_path) pipe.load_components(dtype=torch.float32) diff --git a/tests/modular_pipelines/test_modular_pipeline_loading.py b/tests/modular_pipelines/test_modular_pipeline_loading.py index 4871eab4e67d..06da7e2ad16c 100644 --- a/tests/modular_pipelines/test_modular_pipeline_loading.py +++ b/tests/modular_pipelines/test_modular_pipeline_loading.py @@ -16,6 +16,7 @@ import json import os import shutil +import warnings import pytest import torch @@ -156,7 +157,7 @@ def test_save_pretrained_roundtrip_with_local_model(self, tmp_path): original_state_dict = pipe.unet.state_dict() save_dir = str(tmp_path / "my-pipeline") - pipe.save_pretrained(save_dir) + pipe.save_pretrained(save_dir, overwrite_modular_index=True) loaded_pipe = ModularPipeline.from_pretrained(save_dir) loaded_pipe.load_components(dtype=torch.float32) @@ -204,7 +205,8 @@ def test_save_pretrained_default_writes_self_contained_local_copy(self, tmp_path pipe.load_components(names=["unet"], dtype=torch.float32) save_dir = str(tmp_path / "my-pipeline") - pipe.save_pretrained(save_dir) + with pytest.warns(FutureWarning, match="overwrite_modular_index"): + pipe.save_pretrained(save_dir) with open(os.path.join(save_dir, "modular_model_index.json")) as f: index = json.load(f) @@ -221,6 +223,15 @@ def test_save_pretrained_default_writes_self_contained_local_copy(self, tmp_path loaded_pipe.load_components(names=["unet"], dtype=torch.float32, local_files_only=True) assert loaded_pipe.unet is not None + @pytest.mark.parametrize("overwrite_modular_index", [True, False]) + def test_save_pretrained_explicit_overwrite_modular_index_does_not_warn(self, tmp_path, overwrite_modular_index): + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + pipe.load_components(names=["unet"], dtype=torch.float32) + + with warnings.catch_warnings(): + warnings.filterwarnings("error", message=".*overwrite_modular_index", category=FutureWarning) + pipe.save_pretrained(str(tmp_path / "my-pipeline"), overwrite_modular_index=overwrite_modular_index) + def test_save_pretrained_overwrite_modular_index(self, tmp_path): """With overwrite_modular_index=True, all component references should point to the save directory.""" pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") @@ -266,7 +277,7 @@ def test_init_fallback_when_blocks_class_name_is_base_class(self, tmp_path): # 3. Save and reload — the saved config will have _blocks_class_name="SequentialPipelineBlocks" save_dir = str(tmp_path / "pipeline") - t2i_pipe.save_pretrained(save_dir) + t2i_pipe.save_pretrained(save_dir, overwrite_modular_index=True) loaded_pipe = ModularPipeline.from_pretrained(save_dir) # 4. Verify it fell back to default_blocks_name and has correct blocks diff --git a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py index 20f65740d511..4d7e85cec25c 100644 --- a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py +++ b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py @@ -364,7 +364,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block = ModularPipelineBlocks.from_pretrained(pipeline_repo_dir, trust_remote_code=True) pipe = block.init_pipeline() - pipe.save_pretrained(pipeline_repo_dir) + pipe.save_pretrained(pipeline_repo_dir, overwrite_modular_index=True) # Step 3: Load the pipeline from the saved directory. loaded_pipe = ModularPipeline.from_pretrained(pipeline_repo_dir, trust_remote_code=True) diff --git a/tests/modular_pipelines/testing_utils/loading.py b/tests/modular_pipelines/testing_utils/loading.py index 63b1872f07eb..b9a0555f8ae9 100644 --- a/tests/modular_pipelines/testing_utils/loading.py +++ b/tests/modular_pipelines/testing_utils/loading.py @@ -38,7 +38,7 @@ class ModularLoadingTesterMixin(BaseModularPipelineOutputMixin): def test_save_from_pretrained(self, tmp_path, base_pipe_output): base_pipe = self.get_pipeline().to(torch_device) - base_pipe.save_pretrained(str(tmp_path)) + base_pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) pipe = ModularPipeline.from_pretrained(tmp_path) pipe.load_components(dtype=torch.float32) @@ -67,7 +67,7 @@ def test_load_expected_components_from_pretrained(self, tmp_path): 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) + pipe.save_pretrained(save_dir, overwrite_modular_index=True) expected = get_specified_components(save_dir) loaded_pipe = ModularPipeline.from_pretrained(save_dir)