From 189491a4f21999c07ce398583647bcf9f82a89b0 Mon Sep 17 00:00:00 2001 From: DN6 Date: Thu, 19 Mar 2026 13:05:08 +0530 Subject: [PATCH 1/3] update --- .../test_modular_pipelines_custom_blocks.py | 64 ++++++++++++++++++- 1 file changed, 63 insertions(+), 1 deletion(-) diff --git a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py index 7c6e97a36eb7..49ff5008d53d 100644 --- a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py +++ b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py @@ -31,7 +31,7 @@ WanModularPipeline, ) -from ..testing_utils import nightly, require_torch, slow +from ..testing_utils import nightly, require_torch, require_torch_accelerator, slow, torch_device class DummyCustomBlockSimple(ModularPipelineBlocks): @@ -341,6 +341,68 @@ def __call__(self, components, state: PipelineState) -> PipelineState: loaded_pipe.update_components(custom_model=custom_model) assert getattr(loaded_pipe, "custom_model", None) is not None + def _create_tiny_model_dir(self, model_dir): + TINY_MODEL_CODE = ( + "import torch\n" + "from diffusers import ModelMixin, ConfigMixin\n" + "from diffusers.configuration_utils import register_to_config\n" + "\n" + "class TinyModel(ModelMixin, ConfigMixin):\n" + " @register_to_config\n" + " def __init__(self, hidden_size=4):\n" + " super().__init__()\n" + " self.linear = torch.nn.Linear(hidden_size, hidden_size)\n" + "\n" + " def forward(self, x):\n" + " return self.linear(x)\n" + ) + + with open(os.path.join(model_dir, "modeling.py"), "w") as f: + f.write(TINY_MODEL_CODE) + + config = { + "_class_name": "TinyModel", + "_diffusers_version": "0.0.0", + "auto_map": {"AutoModel": "modeling.TinyModel"}, + "hidden_size": 4, + } + with open(os.path.join(model_dir, "config.json"), "w") as f: + json.dump(config, f) + + torch.save( + {"linear.weight": torch.randn(4, 4), "linear.bias": torch.randn(4)}, + os.path.join(model_dir, "diffusion_pytorch_model.bin"), + ) + + def test_automodel_type_hint_preserves_torch_dtype(self, tmp_path): + """Regression test for #13271: torch_dtype was incorrectly removed when type_hint is AutoModel.""" + from diffusers import AutoModel + + self._create_tiny_model_dir(tmp_path) + + spec = ComponentSpec( + name="model", + type_hint=AutoModel, + pretrained_model_name_or_path=str(tmp_path), + ) + loaded = spec.load(torch_dtype=torch.float16, trust_remote_code=True) + assert loaded.dtype == torch.float16 + + @require_torch_accelerator + def test_automodel_type_hint_preserves_device(self, tmp_path): + """Test that ComponentSpec with AutoModel type_hint correctly passes device_map.""" + from diffusers import AutoModel + + self._create_tiny_model_dir(tmp_path) + + spec = ComponentSpec( + name="model", + type_hint=AutoModel, + pretrained_model_name_or_path=str(tmp_path), + ) + loaded = spec.load(device_map=torch_device, trust_remote_code=True) + assert loaded.device.type == torch_device + def test_custom_block_loads_from_hub(self): repo_id = "hf-internal-testing/tiny-modular-diffusers-block" block = ModularPipelineBlocks.from_pretrained(repo_id, trust_remote_code=True) From 5219182752561610303bfd1e015340f7630d15bb Mon Sep 17 00:00:00 2001 From: DN6 Date: Thu, 19 Mar 2026 13:08:43 +0530 Subject: [PATCH 2/3] update --- .../test_modular_pipelines_custom_blocks.py | 71 ++++++++++--------- 1 file changed, 36 insertions(+), 35 deletions(-) diff --git a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py index 49ff5008d53d..a070086f4e84 100644 --- a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py +++ b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py @@ -34,6 +34,40 @@ from ..testing_utils import nightly, require_torch, require_torch_accelerator, slow, torch_device +def _create_tiny_model_dir(model_dir): + TINY_MODEL_CODE = ( + "import torch\n" + "from diffusers import ModelMixin, ConfigMixin\n" + "from diffusers.configuration_utils import register_to_config\n" + "\n" + "class TinyModel(ModelMixin, ConfigMixin):\n" + " @register_to_config\n" + " def __init__(self, hidden_size=4):\n" + " super().__init__()\n" + " self.linear = torch.nn.Linear(hidden_size, hidden_size)\n" + "\n" + " def forward(self, x):\n" + " return self.linear(x)\n" + ) + + with open(os.path.join(model_dir, "modeling.py"), "w") as f: + f.write(TINY_MODEL_CODE) + + config = { + "_class_name": "TinyModel", + "_diffusers_version": "0.0.0", + "auto_map": {"AutoModel": "modeling.TinyModel"}, + "hidden_size": 4, + } + with open(os.path.join(model_dir, "config.json"), "w") as f: + json.dump(config, f) + + torch.save( + {"linear.weight": torch.randn(4, 4), "linear.bias": torch.randn(4)}, + os.path.join(model_dir, "diffusion_pytorch_model.bin"), + ) + + class DummyCustomBlockSimple(ModularPipelineBlocks): def __init__(self, use_dummy_model_component=False): self.use_dummy_model_component = use_dummy_model_component @@ -341,44 +375,11 @@ def __call__(self, components, state: PipelineState) -> PipelineState: loaded_pipe.update_components(custom_model=custom_model) assert getattr(loaded_pipe, "custom_model", None) is not None - def _create_tiny_model_dir(self, model_dir): - TINY_MODEL_CODE = ( - "import torch\n" - "from diffusers import ModelMixin, ConfigMixin\n" - "from diffusers.configuration_utils import register_to_config\n" - "\n" - "class TinyModel(ModelMixin, ConfigMixin):\n" - " @register_to_config\n" - " def __init__(self, hidden_size=4):\n" - " super().__init__()\n" - " self.linear = torch.nn.Linear(hidden_size, hidden_size)\n" - "\n" - " def forward(self, x):\n" - " return self.linear(x)\n" - ) - - with open(os.path.join(model_dir, "modeling.py"), "w") as f: - f.write(TINY_MODEL_CODE) - - config = { - "_class_name": "TinyModel", - "_diffusers_version": "0.0.0", - "auto_map": {"AutoModel": "modeling.TinyModel"}, - "hidden_size": 4, - } - with open(os.path.join(model_dir, "config.json"), "w") as f: - json.dump(config, f) - - torch.save( - {"linear.weight": torch.randn(4, 4), "linear.bias": torch.randn(4)}, - os.path.join(model_dir, "diffusion_pytorch_model.bin"), - ) - def test_automodel_type_hint_preserves_torch_dtype(self, tmp_path): """Regression test for #13271: torch_dtype was incorrectly removed when type_hint is AutoModel.""" from diffusers import AutoModel - self._create_tiny_model_dir(tmp_path) + _create_tiny_model_dir(tmp_path) spec = ComponentSpec( name="model", @@ -393,7 +394,7 @@ def test_automodel_type_hint_preserves_device(self, tmp_path): """Test that ComponentSpec with AutoModel type_hint correctly passes device_map.""" from diffusers import AutoModel - self._create_tiny_model_dir(tmp_path) + _create_tiny_model_dir(tmp_path) spec = ComponentSpec( name="model", From d90a4dfe57d30b58060b14b7af1d7c796c66c5f0 Mon Sep 17 00:00:00 2001 From: DN6 Date: Thu, 19 Mar 2026 13:16:40 +0530 Subject: [PATCH 3/3] update --- .../test_modular_pipelines_custom_blocks.py | 78 +++++++++++++++---- 1 file changed, 62 insertions(+), 16 deletions(-) diff --git a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py index a070086f4e84..59d6a3e75f55 100644 --- a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py +++ b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py @@ -379,30 +379,76 @@ def test_automodel_type_hint_preserves_torch_dtype(self, tmp_path): """Regression test for #13271: torch_dtype was incorrectly removed when type_hint is AutoModel.""" from diffusers import AutoModel - _create_tiny_model_dir(tmp_path) + model_dir = str(tmp_path / "model") + os.makedirs(model_dir) + _create_tiny_model_dir(model_dir) + + class DtypeTestBlock(ModularPipelineBlocks): + @property + def expected_components(self): + return [ComponentSpec("model", AutoModel, pretrained_model_name_or_path=model_dir)] + + @property + def inputs(self) -> List[InputParam]: + return [InputParam("prompt", type_hint=str, required=True)] + + @property + def intermediate_inputs(self) -> List[InputParam]: + return [] + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [OutputParam("output", type_hint=str)] + + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.output = "test" + self.set_block_state(state, block_state) + return components, state + + block = DtypeTestBlock() + pipe = block.init_pipeline() + pipe.load_components(torch_dtype=torch.float16, trust_remote_code=True) - spec = ComponentSpec( - name="model", - type_hint=AutoModel, - pretrained_model_name_or_path=str(tmp_path), - ) - loaded = spec.load(torch_dtype=torch.float16, trust_remote_code=True) - assert loaded.dtype == torch.float16 + assert pipe.model.dtype == torch.float16 @require_torch_accelerator def test_automodel_type_hint_preserves_device(self, tmp_path): """Test that ComponentSpec with AutoModel type_hint correctly passes device_map.""" from diffusers import AutoModel - _create_tiny_model_dir(tmp_path) + model_dir = str(tmp_path / "model") + os.makedirs(model_dir) + _create_tiny_model_dir(model_dir) - spec = ComponentSpec( - name="model", - type_hint=AutoModel, - pretrained_model_name_or_path=str(tmp_path), - ) - loaded = spec.load(device_map=torch_device, trust_remote_code=True) - assert loaded.device.type == torch_device + class DeviceTestBlock(ModularPipelineBlocks): + @property + def expected_components(self): + return [ComponentSpec("model", AutoModel, pretrained_model_name_or_path=model_dir)] + + @property + def inputs(self) -> List[InputParam]: + return [InputParam("prompt", type_hint=str, required=True)] + + @property + def intermediate_inputs(self) -> List[InputParam]: + return [] + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [OutputParam("output", type_hint=str)] + + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.output = "test" + self.set_block_state(state, block_state) + return components, state + + block = DeviceTestBlock() + pipe = block.init_pipeline() + pipe.load_components(device_map=torch_device, trust_remote_code=True) + + assert pipe.model.device.type == torch_device def test_custom_block_loads_from_hub(self): repo_id = "hf-internal-testing/tiny-modular-diffusers-block"