add tests for robust model loading.

fix

docs/source/en/using-diffusers/automodel.md

@@ -29,8 +29,31 @@ text_encoder = AutoModel.from_pretrained(
 )
 ```
 
+## Custom models
+
 [`AutoModel`] also loads models from the [Hub](https://huggingface.co/models) that aren't included in Diffusers. Set `trust_remote_code=True` in [`AutoModel.from_pretrained`] to load custom models.
 
+A custom model repository needs a Python module with the model class, and a `config.json` with an `auto_map` entry that maps `"AutoModel"` to `"module_file.ClassName"`.
+
+```
+custom/custom-transformer-model/
+├── config.json
+├── my_model.py
+└── diffusion_pytorch_model.safetensors
+```
+
+The `config.json` includes the `auto_map` field pointing to the custom class.
+
+```json
+{
+  "auto_map": {
+    "AutoModel": "my_model.MyCustomModel"
+  }
+}
+```
+
+Then load it with `trust_remote_code=True`.
+
 ```py
 import torch
 from diffusers import AutoModel
@@ -40,7 +63,39 @@ transformer = AutoModel.from_pretrained(
 )
 ```
 
+For a real-world example, [Overworld/Waypoint-1-Small](https://huggingface.co/Overworld/Waypoint-1-Small/tree/main/transformer) hosts a custom `WorldModel` class across several modules in its `transformer` subfolder.
+
+```
+transformer/
+├── config.json          # auto_map: "model.WorldModel"
+├── model.py
+├── attn.py
+├── nn.py
+├── cache.py
+├── quantize.py
+├── __init__.py
+└── diffusion_pytorch_model.safetensors
+```
+
+```py
+import torch
+from diffusers import AutoModel
+
+transformer = AutoModel.from_pretrained(
+    "Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="cuda"
+)
+```
+
 If the custom model inherits from the [`ModelMixin`] class, it gets access to the same features as Diffusers model classes, like [regional compilation](../optimization/fp16#regional-compilation) and [group offloading](../optimization/memory#group-offloading).
 
+> [!WARNING]
+> As a precaution with `trust_remote_code=True`, pass a commit hash to the `revision` argument in [`AutoModel.from_pretrained`] to make sure the code hasn't been updated with new malicious code (unless you fully trust the model owners).
+>
+> ```py
+> transformer = AutoModel.from_pretrained(
+>     "Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, revision="a3d8cb2"
+> )
+> ```
+
 > [!NOTE]
 > Learn more about implementing custom models in the [Community components](../using-diffusers/custom_pipeline_overview#community-components) guide.

src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py

@@ -18,7 +18,6 @@ import re
 from copy import deepcopy
 from typing import Any, Callable, Dict, List, Optional, Union
 
-import ftfy
 import torch
 from transformers import AutoTokenizer, UMT5EncoderModel
 

src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py

@@ -18,7 +18,6 @@ import re
 from copy import deepcopy
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
-import ftfy
 import PIL
 import torch
 from transformers import AutoTokenizer, UMT5EncoderModel

src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py

@@ -19,7 +19,6 @@ import re
 from copy import deepcopy
 from typing import Any, Callable, Dict, List, Optional, Union
 
-import ftfy
 import torch
 from PIL import Image
 from transformers import AutoTokenizer, UMT5EncoderModel

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -6,7 +6,7 @@ import pytest
 import torch
 
 import diffusers
-from diffusers import ComponentsManager, ModularPipeline, ModularPipelineBlocks
+from diffusers import AutoModel, ComponentsManager, ModularPipeline, ModularPipelineBlocks
 from diffusers.guiders import ClassifierFreeGuidance
 from diffusers.modular_pipelines.modular_pipeline_utils import (
     ComponentSpec,
@@ -598,3 +598,69 @@ class TestModularModelCardContent:
         content = generate_modular_model_card_content(blocks)
 
         assert "5-block architecture" in content["model_description"]
+
+
+class TestAutoModelLoadIdTagging:
+    def test_automodel_tags_load_id(self):
+        model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet")
+
+        assert hasattr(model, "_diffusers_load_id"), "Model should have _diffusers_load_id attribute"
+        assert model._diffusers_load_id != "null", "_diffusers_load_id should not be 'null'"
+
+        # Verify load_id contains the expected fields
+        load_id = model._diffusers_load_id
+        assert "hf-internal-testing/tiny-stable-diffusion-xl-pipe" in load_id
+        assert "unet" in load_id
+
+    def test_automodel_update_components(self):
+        pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
+        pipe.load_components(torch_dtype=torch.float32)
+
+        auto_model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet")
+
+        pipe.update_components(unet=auto_model)
+
+        assert pipe.unet is auto_model
+
+        assert "unet" in pipe._component_specs
+        spec = pipe._component_specs["unet"]
+        assert spec.pretrained_model_name_or_path == "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
+        assert spec.subfolder == "unet"
+
+
+class TestLoadComponentsSkipBehavior:
+    def test_load_components_skips_already_loaded(self):
+        pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
+        pipe.load_components(torch_dtype=torch.float32)
+
+        original_unet = pipe.unet
+
+        pipe.load_components()
+
+        # Verify that the unet is the same object (not reloaded)
+        assert pipe.unet is original_unet, "load_components should skip already loaded components"
+
+    def test_load_components_selective_loading(self):
+        pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
+
+        pipe.load_components(names="unet", torch_dtype=torch.float32)
+
+        # Verify only requested component was loaded.
+        assert hasattr(pipe, "unet")
+        assert pipe.unet is not None
+        if "vae" in pipe._component_specs:
+            assert getattr(pipe, "vae", None) is None
+
+    def test_load_components_skips_invalid_pretrained_path(self):
+        pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
+
+        pipe._component_specs["test_component"] = ComponentSpec(
+            name="test_component",
+            type_hint=torch.nn.Module,
+            pretrained_model_name_or_path=None,
+            default_creation_method="from_pretrained",
+        )
+        pipe.load_components(torch_dtype=torch.float32)
+
+        # Verify test_component was not loaded
+        assert not hasattr(pipe, "test_component") or pipe.test_component is None