fix

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

@@ -29,31 +29,8 @@ 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
@@ -63,39 +40,7 @@ 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/loaders/lora_conversion_utils.py

@@ -2321,14 +2321,6 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
     prefix = "diffusion_model."
     original_state_dict = {k[len(prefix) :]: v for k, v in state_dict.items()}
 
-    has_lora_down_up = any("lora_down" in k or "lora_up" in k for k in original_state_dict.keys())
-    if has_lora_down_up:
-        temp_state_dict = {}
-        for k, v in original_state_dict.items():
-            new_key = k.replace("lora_down", "lora_A").replace("lora_up", "lora_B")
-            temp_state_dict[new_key] = v
-        original_state_dict = temp_state_dict
-
     num_double_layers = 0
     num_single_layers = 0
     for key in original_state_dict.keys():
@@ -2345,15 +2337,13 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
         attn_prefix = f"single_transformer_blocks.{sl}.attn"
 
         for lora_key in lora_keys:
-            linear1_key = f"{single_block_prefix}.linear1.{lora_key}.weight"
-            if linear1_key in original_state_dict:
-                converted_state_dict[f"{attn_prefix}.to_qkv_mlp_proj.{lora_key}.weight"] = original_state_dict.pop(
-                    linear1_key
-                )
+            converted_state_dict[f"{attn_prefix}.to_qkv_mlp_proj.{lora_key}.weight"] = original_state_dict.pop(
+                f"{single_block_prefix}.linear1.{lora_key}.weight"
+            )
 
-            linear2_key = f"{single_block_prefix}.linear2.{lora_key}.weight"
-            if linear2_key in original_state_dict:
-                converted_state_dict[f"{attn_prefix}.to_out.{lora_key}.weight"] = original_state_dict.pop(linear2_key)
+            converted_state_dict[f"{attn_prefix}.to_out.{lora_key}.weight"] = original_state_dict.pop(
+                f"{single_block_prefix}.linear2.{lora_key}.weight"
+            )
 
     for dl in range(num_double_layers):
         transformer_block_prefix = f"transformer_blocks.{dl}"
@@ -2362,10 +2352,6 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
             for attn_type in attn_types:
                 attn_prefix = f"{transformer_block_prefix}.attn"
                 qkv_key = f"double_blocks.{dl}.{attn_type}.qkv.{lora_key}.weight"
-
-                if qkv_key not in original_state_dict:
-                    continue
-
                 fused_qkv_weight = original_state_dict.pop(qkv_key)
 
                 if lora_key == "lora_A":
@@ -2397,9 +2383,8 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
         for org_proj, diff_proj in proj_mappings:
             for lora_key in lora_keys:
                 original_key = f"double_blocks.{dl}.{org_proj}.{lora_key}.weight"
-                if original_key in original_state_dict:
-                    diffusers_key = f"{transformer_block_prefix}.{diff_proj}.{lora_key}.weight"
-                    converted_state_dict[diffusers_key] = original_state_dict.pop(original_key)
+                diffusers_key = f"{transformer_block_prefix}.{diff_proj}.{lora_key}.weight"
+                converted_state_dict[diffusers_key] = original_state_dict.pop(original_key)
 
         mlp_mappings = [
             ("img_mlp.0", "ff.linear_in"),
@@ -2410,27 +2395,8 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
         for org_mlp, diff_mlp in mlp_mappings:
             for lora_key in lora_keys:
                 original_key = f"double_blocks.{dl}.{org_mlp}.{lora_key}.weight"
-                if original_key in original_state_dict:
-                    diffusers_key = f"{transformer_block_prefix}.{diff_mlp}.{lora_key}.weight"
-                    converted_state_dict[diffusers_key] = original_state_dict.pop(original_key)
-
-    extra_mappings = {
-        "img_in": "x_embedder",
-        "txt_in": "context_embedder",
-        "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1",
-        "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2",
-        "final_layer.linear": "proj_out",
-        "final_layer.adaLN_modulation.1": "norm_out.linear",
-        "single_stream_modulation.lin": "single_stream_modulation.linear",
-        "double_stream_modulation_img.lin": "double_stream_modulation_img.linear",
-        "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear",
-    }
-
-    for org_key, diff_key in extra_mappings.items():
-        for lora_key in lora_keys:
-            original_key = f"{org_key}.{lora_key}.weight"
-            if original_key in original_state_dict:
-                converted_state_dict[f"{diff_key}.{lora_key}.weight"] = original_state_dict.pop(original_key)
+                diffusers_key = f"{transformer_block_prefix}.{diff_mlp}.{lora_key}.weight"
+                converted_state_dict[diffusers_key] = original_state_dict.pop(original_key)
 
     if len(original_state_dict) > 0:
         raise ValueError(f"`original_state_dict` should be empty at this point but has {original_state_dict.keys()=}.")

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -6,7 +6,7 @@ import pytest
 import torch
 
 import diffusers
-from diffusers import AutoModel, ComponentsManager, ModularPipeline, ModularPipelineBlocks
+from diffusers import ComponentsManager, ModularPipeline, ModularPipelineBlocks
 from diffusers.guiders import ClassifierFreeGuidance
 from diffusers.modular_pipelines.modular_pipeline_utils import (
     ComponentSpec,
@@ -598,69 +598,3 @@ 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