agentdevsl · cursor · Jun 24, 2026 · Jun 24, 2026
diff --git a/src/diffusers/guiders/auto_guidance.py b/src/diffusers/guiders/auto_guidance.py
@@ -104,7 +104,7 @@ def __init__(
                     f"Expected `auto_guidance_layers` to be an int or a list of ints, but got {type(auto_guidance_layers)}."
                 )
             auto_guidance_config = [
-                LayerSkipConfig(layer, fqn="auto", dropout=dropout) for layer in auto_guidance_layers
+                LayerSkipConfig(indices=[layer], fqn="auto", dropout=dropout) for layer in auto_guidance_layers
             ]
 
         if isinstance(auto_guidance_config, dict):

diff --git a/src/diffusers/guiders/skip_layer_guidance.py b/src/diffusers/guiders/skip_layer_guidance.py
@@ -136,7 +136,7 @@ def __init__(
                 raise ValueError(
                     f"Expected `skip_layer_guidance_layers` to be an int or a list of ints, but got {type(skip_layer_guidance_layers)}."
                 )
-            skip_layer_config = [LayerSkipConfig(layer, fqn="auto") for layer in skip_layer_guidance_layers]
+            skip_layer_config = [LayerSkipConfig(indices=[layer], fqn="auto") for layer in skip_layer_guidance_layers]
 
         if isinstance(skip_layer_config, dict):
             skip_layer_config = LayerSkipConfig.from_dict(skip_layer_config)

diff --git a/src/diffusers/hooks/layer_skip.py b/src/diffusers/hooks/layer_skip.py
@@ -166,7 +166,13 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
                 original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
                     "encoder_hidden_states", args, kwargs
                 )
-                output = (original_hidden_states, original_encoder_hidden_states)
+                max_idx = max(
+                    self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index
+                )
+                ret_list = [None] * (max_idx + 1)
+                ret_list[self._metadata.return_hidden_states_index] = original_hidden_states
+                ret_list[self._metadata.return_encoder_hidden_states_index] = original_encoder_hidden_states
+                output = tuple(ret_list)
         else:
             output = self.fn_ref.original_forward(*args, **kwargs)
             output = torch.nn.functional.dropout(output, p=self.dropout)

diff --git a/src/diffusers/models/model_loading_utils.py b/src/diffusers/models/model_loading_utils.py
@@ -18,6 +18,7 @@
 import importlib
 import inspect
 import os
+import threading
 from array import array
 from collections import OrderedDict, defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -51,6 +52,8 @@
 
 logger = logging.get_logger(__name__)
 
+_parallel_load_lock = threading.Lock()
+
 _CLASS_REMAPPING_DICT = {
     "Transformer2DModel": {
         "ada_norm_zero": "DiTTransformer2DModel",
@@ -360,31 +363,32 @@ def _load_shard_file(
     if hf_quantizer is not None:
         state_dict = hf_quantizer.maybe_update_state_dict(state_dict)
 
-    mismatched_keys = _find_mismatched_keys(
-        state_dict,
-        model_state_dict,
-        loaded_keys,
-        ignore_mismatched_sizes,
-    )
-    error_msgs = []
-    if low_cpu_mem_usage:
-        offload_index, state_dict_index = load_model_dict_into_meta(
-            model,
+    with _parallel_load_lock:
+        mismatched_keys = _find_mismatched_keys(
             state_dict,
-            device_map=device_map,
-            dtype=dtype,
-            hf_quantizer=hf_quantizer,
-            keep_in_fp32_modules=keep_in_fp32_modules,
-            unexpected_keys=unexpected_keys,
-            offload_folder=offload_folder,
-            offload_index=offload_index,
-            state_dict_index=state_dict_index,
-            state_dict_folder=state_dict_folder,
+            model_state_dict,
+            loaded_keys,
+            ignore_mismatched_sizes,
         )
-    else:
-        assign_to_params_buffers = check_support_param_buffer_assignment(model, state_dict)
+        error_msgs = []
+        if low_cpu_mem_usage:
+            offload_index, state_dict_index = load_model_dict_into_meta(
+                model,
+                state_dict,
+                device_map=device_map,
+                dtype=dtype,
+                hf_quantizer=hf_quantizer,
+                keep_in_fp32_modules=keep_in_fp32_modules,
+                unexpected_keys=unexpected_keys,
+                offload_folder=offload_folder,
+                offload_index=offload_index,
+                state_dict_index=state_dict_index,
+                state_dict_folder=state_dict_folder,
+            )
+        else:
+            assign_to_params_buffers = check_support_param_buffer_assignment(model, state_dict)
 
-        error_msgs += _load_state_dict_into_model(model, state_dict, assign_to_params_buffers)
+            error_msgs += _load_state_dict_into_model(model, state_dict, assign_to_params_buffers)
     return offload_index, state_dict_index, mismatched_keys, error_msgs
 
 

diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py
@@ -990,7 +990,9 @@ def _identify_model_variants(folder: str, variant: str, config: dict) -> dict:
         for sub_folder in os.listdir(folder):
             folder_path = os.path.join(folder, sub_folder)
             is_folder = os.path.isdir(folder_path) and sub_folder in config
-            variant_exists = is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path))
+            variant_exists = is_folder and any(
+                "." in p and p.split(".", 1)[1].startswith(variant) for p in os.listdir(folder_path)
+            )
             if variant_exists:
                 model_variants[sub_folder] = variant
     return model_variants

diff --git a/tests/guiders/test_skip_layer_guidance.py b/tests/guiders/test_skip_layer_guidance.py
@@ -0,0 +1,34 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+from diffusers.guiders.auto_guidance import AutoGuidance
+from diffusers.guiders.skip_layer_guidance import SkipLayerGuidance
+
+
+class SkipLayerGuidanceConfigTest(unittest.TestCase):
+    def test_shorthand_layers_wrap_each_index_in_a_list(self):
+        guider = SkipLayerGuidance(skip_layer_guidance_layers=[7, 8, 9])
+        self.assertEqual([config.indices for config in guider.skip_layer_config], [[7], [8], [9]])
+
+    def test_single_int_layer_shorthand(self):
+        guider = SkipLayerGuidance(skip_layer_guidance_layers=7)
+        self.assertEqual([config.indices for config in guider.skip_layer_config], [[7]])
+
+
+class AutoGuidanceConfigTest(unittest.TestCase):
+    def test_shorthand_layers_wrap_each_index_in_a_list(self):
+        guider = AutoGuidance(auto_guidance_layers=[3, 4], dropout=1.0)
+        self.assertEqual([config.indices for config in guider.auto_guidance_config], [[3], [4]])
diff --git a/tests/hooks/test_layer_skip.py b/tests/hooks/test_layer_skip.py
@@ -0,0 +1,56 @@
+# Copyright 2025 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry
+from diffusers.hooks.layer_skip import LayerSkipConfig, apply_layer_skip
+from diffusers.models import ModelMixin
+
+
+class FluxLikeBlock(torch.nn.Module):
+    def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
+        return encoder_hidden_states + 1.0, hidden_states + 2.0
+
+
+class FluxLikeTransformer(ModelMixin):
+    def __init__(self):
+        super().__init__()
+        self.transformer_blocks = torch.nn.ModuleList([FluxLikeBlock(), FluxLikeBlock()])
+
+    def forward(self, hidden_states, encoder_hidden_states=None):
+        for block in self.transformer_blocks:
+            encoder_hidden_states, hidden_states = block(
+                hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states
+            )
+        return encoder_hidden_states, hidden_states
+
+
+def test_transformer_block_skip_hook_respects_return_order():
+    TransformerBlockRegistry.register(
+        FluxLikeBlock,
+        TransformerBlockMetadata(return_hidden_states_index=1, return_encoder_hidden_states_index=0),
+    )
+
+    model = FluxLikeTransformer()
+    hidden_states = torch.zeros(2, 3)
+    encoder_hidden_states = torch.ones(2, 3)
+
+    apply_layer_skip(model, LayerSkipConfig(indices=[0], fqn="transformer_blocks"))
+
+    out_encoder, out_hidden = model(hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states)
+
+    # Block 0 is skipped (identity inputs), block 1 still runs (+1 / +2).
+    torch.testing.assert_close(out_encoder, encoder_hidden_states + 1.0)
+    torch.testing.assert_close(out_hidden, hidden_states + 2.0)
diff --git a/tests/others/test_remote_utils.py b/tests/others/test_remote_utils.py
@@ -0,0 +1,103 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import io
+import json
+import unittest
+from unittest.mock import Mock
+
+import torch
+from PIL import Image
+
+from diffusers.utils.remote_utils import (
+    check_inputs_decode,
+    detect_image_type,
+    postprocess_decode,
+    prepare_decode,
+    prepare_encode,
+)
+
+
+class RemoteUtilsTest(unittest.TestCase):
+    def test_detect_image_type(self):
+        self.assertEqual(detect_image_type(b"\xff\xd8\xff"), "jpeg")
+        self.assertEqual(detect_image_type(b"\x89PNG\r\n\x1a\n"), "png")
+        self.assertEqual(detect_image_type(b"GIF89a"), "gif")
+        self.assertEqual(detect_image_type(b"BM"), "bmp")
+        self.assertEqual(detect_image_type(b"unknown"), "unknown")
+
+    def test_check_inputs_decode_packed_latents_requires_hw(self):
+        tensor = torch.randn(4, 8, 8)
+        with self.assertRaises(ValueError):
+            check_inputs_decode("http://example.com", tensor)
+
+    def test_check_inputs_decode_processor_required(self):
+        tensor = torch.randn(1, 4, 8, 8)
+        with self.assertRaises(ValueError):
+            check_inputs_decode(
+                "http://example.com",
+                tensor,
+                processor=None,
+                output_type="pt",
+                return_type="pil",
+                partial_postprocess=False,
+            )
+
+    def test_prepare_decode_sets_accept_header_for_jpeg(self):
+        tensor = torch.randn(1, 4, 8, 8, dtype=torch.float16)
+        payload = prepare_decode(tensor, output_type="pil", image_format="jpg")
+        self.assertEqual(payload["headers"]["Accept"], "image/jpeg")
+        self.assertEqual(payload["params"]["output_type"], "pil")
+        self.assertEqual(payload["params"]["shape"], list(tensor.shape))
+
+    def test_prepare_encode_tensor_includes_shape_and_dtype(self):
+        tensor = torch.randn(1, 3, 8, 8, dtype=torch.float16)
+        payload = prepare_encode(tensor, scaling_factor=0.18215)
+        self.assertEqual(payload["params"]["shape"], list(tensor.shape))
+        self.assertEqual(payload["params"]["dtype"], "float16")
+        self.assertEqual(payload["params"]["scaling_factor"], 0.18215)
+
+    def test_prepare_encode_pil_image(self):
+        image = Image.new("RGB", (8, 8), color="red")
+        payload = prepare_encode(image)
+        self.assertIn(b"PNG", payload["data"][:8])
+
+    def test_postprocess_decode_pil_without_processor(self):
+        buffer = io.BytesIO()
+        Image.new("RGB", (4, 4), color="blue").save(buffer, format="PNG")
+        response = Mock()
+        response.content = buffer.getvalue()
+
+        output = postprocess_decode(response, processor=None, output_type="pil", return_type="pil")
+        self.assertIsInstance(output, Image.Image)
+        self.assertEqual(output.size, (4, 4))
+        self.assertEqual(output.format, "png")
+
+    def test_postprocess_decode_pt_tensor(self):
+        tensor = torch.arange(16, dtype=torch.float32).reshape(1, 4, 2, 2)
+        response = Mock()
+        response.content = tensor.numpy().tobytes()
+        response.headers = {
+            "shape": json.dumps(list(tensor.shape)),
+            "dtype": "float32",
+        }
+
+        output = postprocess_decode(
+            response,
+            processor=None,
+            output_type="pt",
+            return_type="pt",
+            partial_postprocess=False,
+        )
+        torch.testing.assert_close(output, tensor)