From edbf0e7c15921e7a83df7d831ed18409b7345d3d Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 2 Mar 2026 23:24:31 +0100 Subject: [PATCH 1/8] add --- src/diffusers/modular_pipelines/modular_pipeline.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index c1ac7f3aab4c..0795b3fae123 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -1707,6 +1707,8 @@ def __init__( _blocks_class_name=self._blocks.__class__.__name__ if self._blocks is not None else None ) + self._pretrained_model_name_or_path = pretrained_model_name_or_path + @property def default_call_parameters(self) -> dict[str, Any]: """ @@ -2323,6 +2325,16 @@ def load_components(self, names: list[str] | str | None = None, **kwargs): elif "default" in value: # check if the default is specified component_load_kwargs[key] = value["default"] + # Only pass trust_remote_code to components from the same repo as the pipeline. + # When a user passes trust_remote_code=True, they intend to trust code from the + # pipeline's repo, not from external repos referenced in modular_model_index.json. + if ( + "trust_remote_code" in component_load_kwargs + and self._pretrained_model_name_or_path is not None + and spec.pretrained_model_name_or_path != self._pretrained_model_name_or_path + ): + component_load_kwargs.pop("trust_remote_code") + try: components_to_register[name] = spec.load(**component_load_kwargs) except Exception: From 7178fc6bdc071d1ebf4ffe02899339b129419504 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 3 Mar 2026 06:04:03 +0100 Subject: [PATCH 2/8] update warn --- .../modular_pipelines/modular_pipeline.py | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 0795b3fae123..775f4c16ca18 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -2328,24 +2328,42 @@ def load_components(self, names: list[str] | str | None = None, **kwargs): # Only pass trust_remote_code to components from the same repo as the pipeline. # When a user passes trust_remote_code=True, they intend to trust code from the # pipeline's repo, not from external repos referenced in modular_model_index.json. + trust_remote_code_stripped = False if ( "trust_remote_code" in component_load_kwargs and self._pretrained_model_name_or_path is not None and spec.pretrained_model_name_or_path != self._pretrained_model_name_or_path ): component_load_kwargs.pop("trust_remote_code") + trust_remote_code_stripped = True + + if not spec.pretrained_model_name_or_path: + logger.info(f"Skipping component `{name}`: no pretrained model path specified.") + continue try: components_to_register[name] = spec.load(**component_load_kwargs) except Exception: - logger.warning( + tb = traceback.format_exc() + warning_msg = ( f"\nFailed to create component {name}:\n" f"- Component spec: {spec}\n" f"- load() called with kwargs: {component_load_kwargs}\n" "If this component is not required for your workflow you can safely ignore this message.\n\n" "Traceback:\n" - f"{traceback.format_exc()}" + f"{tb}" ) + if trust_remote_code_stripped and "trust_remote_code" in tb: + warning_msg += ( + f"\nNote: `trust_remote_code=True` was not passed to `{name}` because it comes from " + f"an external repository (`{spec.pretrained_model_name_or_path}`). For safety, `trust_remote_code` is only forwarded to " + f"components from the same repository as the pipeline.\n\n" + f"To load this component manually:\n" + f" from diffusers import AutoModel\n" + f' {name} = AutoModel.from_pretrained("{spec.pretrained_model_name_or_path}", trust_remote_code=True)\n' + f" pipe.update_components({name} = {name})\n" + ) + logger.warning(warning_msg) # Register all components at once self.register_components(**components_to_register) From 812365b26cdaebd7ac9b3e87ada9a12cf7c6ad98 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 3 Mar 2026 06:04:13 +0100 Subject: [PATCH 3/8] add a test --- .../test_modular_pipelines_custom_blocks.py | 150 ++++++++++++++++++ 1 file changed, 150 insertions(+) diff --git a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py index 9c5fd5be326d..766ca0c16f86 100644 --- a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py +++ b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py @@ -192,6 +192,156 @@ def test_custom_block_supported_components(self): assert len(pipe.components) == 1 assert pipe.component_names[0] == "transformer" + def test_trust_remote_code_not_propagated_to_external_repo(self): + """When a modular pipeline repo references a component from an external repo that has custom + code (auto_map in config), calling load_components(trust_remote_code=True) should NOT + propagate trust_remote_code to that external component. The external component should fail + to load.""" + + from diffusers import ModularPipeline + + CUSTOM_MODEL_CODE = ( + "import torch\n" + "from diffusers import ModelMixin, ConfigMixin\n" + "from diffusers.configuration_utils import register_to_config\n" + "\n" + "class CustomModel(ModelMixin, ConfigMixin):\n" + " @register_to_config\n" + " def __init__(self, hidden_size=8):\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 tempfile.TemporaryDirectory() as external_repo_dir, tempfile.TemporaryDirectory() as pipeline_repo_dir: + # Step 1: Create an external model repo with custom code (requires trust_remote_code) + with open(os.path.join(external_repo_dir, "modeling.py"), "w") as f: + f.write(CUSTOM_MODEL_CODE) + + config = { + "_class_name": "CustomModel", + "_diffusers_version": "0.0.0", + "auto_map": {"AutoModel": "modeling.CustomModel"}, + "hidden_size": 8, + } + with open(os.path.join(external_repo_dir, "config.json"), "w") as f: + json.dump(config, f) + + torch.save({}, os.path.join(external_repo_dir, "diffusion_pytorch_model.bin")) + + # Step 2: Create a custom block that references the external repo. + # Define both the class (for direct use) and its code string (for block.py). + class ExternalRefBlock(ModularPipelineBlocks): + @property + def expected_components(self): + return [ + ComponentSpec( + "custom_model", + AutoModel, + pretrained_model_name_or_path=external_repo_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 + + EXTERNAL_REF_BLOCK_CODE_STR = ( + "from typing import List\n" + "from diffusers import AutoModel\n" + "from diffusers.modular_pipelines import (\n" + " ComponentSpec,\n" + " InputParam,\n" + " ModularPipelineBlocks,\n" + " OutputParam,\n" + " PipelineState,\n" + ")\n" + "\n" + "class ExternalRefBlock(ModularPipelineBlocks):\n" + " @property\n" + " def expected_components(self):\n" + " return [\n" + " ComponentSpec(\n" + ' "custom_model",\n' + " AutoModel,\n" + f' pretrained_model_name_or_path="{external_repo_dir}",\n' + " )\n" + " ]\n" + "\n" + " @property\n" + " def inputs(self) -> List[InputParam]:\n" + ' return [InputParam("prompt", type_hint=str, required=True)]\n' + "\n" + " @property\n" + " def intermediate_inputs(self) -> List[InputParam]:\n" + " return []\n" + "\n" + " @property\n" + " def intermediate_outputs(self) -> List[OutputParam]:\n" + ' return [OutputParam("output", type_hint=str)]\n' + "\n" + " def __call__(self, components, state: PipelineState) -> PipelineState:\n" + " block_state = self.get_block_state(state)\n" + ' block_state.output = "test"\n' + " self.set_block_state(state, block_state)\n" + " return components, state\n" + ) + + # Save the block config, write block.py, then load back via from_pretrained + block = ExternalRefBlock() + block.save_pretrained(pipeline_repo_dir) + + # auto_map will reference the module name derived from ExternalRefBlock.__module__, + # which is "test_modular_pipelines_custom_blocks". Write the code file with that name. + code_path = os.path.join(pipeline_repo_dir, "test_modular_pipelines_custom_blocks.py") + with open(code_path, "w") as f: + f.write(EXTERNAL_REF_BLOCK_CODE_STR) + + block = ModularPipelineBlocks.from_pretrained(pipeline_repo_dir, trust_remote_code=True) + pipe = block.init_pipeline() + pipe.save_pretrained(pipeline_repo_dir) + + # Step 3: Load the pipeline from the saved directory. + loaded_pipe = ModularPipeline.from_pretrained(pipeline_repo_dir, trust_remote_code=True) + + assert loaded_pipe._pretrained_model_name_or_path == pipeline_repo_dir + assert loaded_pipe._component_specs["custom_model"].pretrained_model_name_or_path == external_repo_dir + assert getattr(loaded_pipe, "custom_model", None) is None + + # Step 4a: load_components WITHOUT trust_remote_code. + # It should still fail + loaded_pipe.load_components() + assert getattr(loaded_pipe, "custom_model", None) is None + + # Step 4b: load_components with trust_remote_code=True. + # trust_remote_code should be stripped for the external component, so it fails. + # The warning should contain guidance about manually loading with trust_remote_code. + loaded_pipe.load_components(trust_remote_code=True) + assert getattr(loaded_pipe, "custom_model", None) is None + + # Step 4c: Manually load with AutoModel and update_components — this should work. + from diffusers import AutoModel + + custom_model = AutoModel.from_pretrained(external_repo_dir, trust_remote_code=True) + loaded_pipe.update_components(custom_model=custom_model) + assert getattr(loaded_pipe, "custom_model", None) is not None + 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 a605b2a887946f6882be07fbbacd03a208ca8d52 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 3 Mar 2026 08:25:06 +0100 Subject: [PATCH 4/8] updaqte --- .../modular_pipelines/modular_pipeline.py | 34 +++++++++++-------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index ebb0ca5c1ec8..684bf7a02bbb 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -2347,23 +2347,27 @@ def load_components(self, names: list[str] | str | None = None, **kwargs): components_to_register[name] = spec.load(**component_load_kwargs) except Exception: tb = traceback.format_exc() - warning_msg = ( - f"\nFailed to create component {name}:\n" - f"- Component spec: {spec}\n" - f"- load() called with kwargs: {component_load_kwargs}\n" - "If this component is not required for your workflow you can safely ignore this message.\n\n" - "Traceback:\n" - f"{tb}" - ) if trust_remote_code_stripped and "trust_remote_code" in tb: - warning_msg += ( - f"\nNote: `trust_remote_code=True` was not passed to `{name}` because it comes from " - f"an external repository (`{spec.pretrained_model_name_or_path}`). For safety, `trust_remote_code` is only forwarded to " - f"components from the same repository as the pipeline.\n\n" - f"To load this component manually:\n" - f" from diffusers import AutoModel\n" + warning_msg = ( + f"Failed to load component `{name}` from external repository " + f"`{spec.pretrained_model_name_or_path}`.\n\n" + f"`trust_remote_code=True` was not forwarded to `{name}` because it comes from " + f"a different repository than the pipeline (`{self._pretrained_model_name_or_path}`). " + f"For safety, `trust_remote_code` is only forwarded to components from the same " + f"repository as the pipeline.\n\n" + f"You need to load this component manually with `trust_remote_code=True` and pass it " + f"to the pipeline via `pipe.update_components()`. For example, if it is a custom model:\n\n" f' {name} = AutoModel.from_pretrained("{spec.pretrained_model_name_or_path}", trust_remote_code=True)\n' - f" pipe.update_components({name} = {name})\n" + f" pipe.update_components({name}={name})\n" + ) + else: + warning_msg = ( + f"Failed to create component {name}:\n" + f"- Component spec: {spec}\n" + f"- load() called with kwargs: {component_load_kwargs}\n" + "If this component is not required for your workflow you can safely ignore this message.\n\n" + "Traceback:\n" + f"{tb}" ) logger.warning(warning_msg) From 81162568dccd706b5e9635cdc79253d3ee4c4d32 Mon Sep 17 00:00:00 2001 From: "yiyi@huggingface.co" Date: Tue, 3 Mar 2026 09:19:14 +0000 Subject: [PATCH 5/8] update_component with custom model --- src/diffusers/modular_pipelines/modular_pipeline.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 684bf7a02bbb..37e8ebf638d1 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -2256,6 +2256,10 @@ def update_components(self, **kwargs): new_component_spec = current_component_spec if hasattr(self, name) and getattr(self, name) is not None: logger.warning(f"ModularPipeline.update_components: setting {name} to None (spec unchanged)") + elif current_component_spec.default_creation_method == "from_pretrained" and not ( + hasattr(component, "_diffusers_load_id") and component._diffusers_load_id is not None + ): + new_component_spec = ComponentSpec(name=name, type_hint=type(component)) else: new_component_spec = ComponentSpec.from_component(name, component) From fa5141500e47027c9f2b2f4c9fa74938c0885923 Mon Sep 17 00:00:00 2001 From: "yiyi@huggingface.co" Date: Tue, 3 Mar 2026 09:19:26 +0000 Subject: [PATCH 6/8] add more tests --- .../test_modular_pipelines_common.py | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index 486b2c3f4166..589698ffc73b 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -687,6 +687,18 @@ def test_load_components_selective_loading(self): assert pipe.unet is not None assert getattr(pipe, "vae", None) is None + def test_load_components_selective_loading_incremental(self): + """Loading a subset of components should not affect already-loaded components.""" + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + + pipe.load_components(names="unet", torch_dtype=torch.float32) + pipe.load_components(names="text_encoder", torch_dtype=torch.float32) + + assert hasattr(pipe, "unet") + assert pipe.unet is not None + assert hasattr(pipe, "text_encoder") + assert pipe.text_encoder is not None + def test_load_components_skips_invalid_pretrained_path(self): pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") @@ -749,6 +761,36 @@ def test_save_pretrained_roundtrip_with_local_model(self, tmp_path): for key in original_state_dict: assert torch.equal(original_state_dict[key], loaded_state_dict[key]), f"Mismatch in {key}" + def test_save_pretrained_updates_index_for_model_with_no_load_id(self, tmp_path): + """testing the workflow of update the pipeline with a custom model and save the pipeline, + the modular_model_index.json should point to the save directory.""" + import json + + from diffusers import UNet2DConditionModel + + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + pipe.load_components(torch_dtype=torch.float32) + + unet = UNet2DConditionModel.from_pretrained( + "hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet" + ) + assert not hasattr(unet, "_diffusers_load_id") + + pipe.update_components(unet=unet) + + save_dir = str(tmp_path / "my-pipeline") + pipe.save_pretrained(save_dir) + + with open(os.path.join(save_dir, "modular_model_index.json")) as f: + index = json.load(f) + + _library, _cls, unet_spec = index["unet"] + assert unet_spec["pretrained_model_name_or_path"] == save_dir + assert unet_spec["subfolder"] == "unet" + + _library, _cls, vae_spec = index["vae"] + assert vae_spec["pretrained_model_name_or_path"] == "hf-internal-testing/tiny-stable-diffusion-xl-pipe" + def test_save_pretrained_overwrite_modular_index(self, tmp_path): """With overwrite_modular_index=True, all component references should point to the save directory.""" import json From acd21875363d7600bb7827592090e6e2f7395336 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Tue, 3 Mar 2026 00:33:30 -1000 Subject: [PATCH 7/8] Apply suggestion from @DN6 Co-authored-by: Dhruv Nair --- src/diffusers/modular_pipelines/modular_pipeline.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 37e8ebf638d1..6ad5fcebb4d7 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -2256,9 +2256,7 @@ def update_components(self, **kwargs): new_component_spec = current_component_spec if hasattr(self, name) and getattr(self, name) is not None: logger.warning(f"ModularPipeline.update_components: setting {name} to None (spec unchanged)") - elif current_component_spec.default_creation_method == "from_pretrained" and not ( - hasattr(component, "_diffusers_load_id") and component._diffusers_load_id is not None - ): + elif current_component_spec.default_creation_method == "from_pretrained" and getattr(component, "_diffusers_load_id", None) is None new_component_spec = ComponentSpec(name=name, type_hint=type(component)) else: new_component_spec = ComponentSpec.from_component(name, component) From 8be13458479de1f06dafda7cae691fdae880e557 Mon Sep 17 00:00:00 2001 From: "yiyi@huggingface.co" Date: Tue, 3 Mar 2026 10:36:58 +0000 Subject: [PATCH 8/8] up --- src/diffusers/modular_pipelines/modular_pipeline.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 6ad5fcebb4d7..8d662080124c 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -2256,7 +2256,10 @@ def update_components(self, **kwargs): new_component_spec = current_component_spec if hasattr(self, name) and getattr(self, name) is not None: logger.warning(f"ModularPipeline.update_components: setting {name} to None (spec unchanged)") - elif current_component_spec.default_creation_method == "from_pretrained" and getattr(component, "_diffusers_load_id", None) is None + elif ( + current_component_spec.default_creation_method == "from_pretrained" + and getattr(component, "_diffusers_load_id", None) is None + ): new_component_spec = ComponentSpec(name=name, type_hint=type(component)) else: new_component_spec = ComponentSpec.from_component(name, component)