update

docs/source/en/optimization/attention_backends.md

@@ -143,7 +143,6 @@ Refer to the table below for a complete list of available attention backends and
 | `flash_varlen` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | Variable length FlashAttention |
 | `flash_varlen_hub` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | Variable length FlashAttention from kernels |
 | `aiter` | [AI Tensor Engine for ROCm](https://github.com/ROCm/aiter) | FlashAttention for AMD ROCm |
-| `flash_4_hub` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-4 |
 | `_flash_3` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-3 |
 | `_flash_varlen_3` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | Variable length FlashAttention-3 |
 | `_flash_3_hub` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-3 from kernels |

src/diffusers/loaders/lora_conversion_utils.py

@@ -2443,191 +2443,6 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
     return converted_state_dict
 
 
-def _convert_kohya_flux2_lora_to_diffusers(state_dict):
-    def _convert_to_ai_toolkit(sds_sd, ait_sd, sds_key, ait_key):
-        if sds_key + ".lora_down.weight" not in sds_sd:
-            return
-        down_weight = sds_sd.pop(sds_key + ".lora_down.weight")
-
-        # scale weight by alpha and dim
-        rank = down_weight.shape[0]
-        default_alpha = torch.tensor(rank, dtype=down_weight.dtype, device=down_weight.device, requires_grad=False)
-        alpha = sds_sd.pop(sds_key + ".alpha", default_alpha).item()
-        scale = alpha / rank
-
-        scale_down = scale
-        scale_up = 1.0
-        while scale_down * 2 < scale_up:
-            scale_down *= 2
-            scale_up /= 2
-
-        ait_sd[ait_key + ".lora_A.weight"] = down_weight * scale_down
-        ait_sd[ait_key + ".lora_B.weight"] = sds_sd.pop(sds_key + ".lora_up.weight") * scale_up
-
-    def _convert_to_ai_toolkit_cat(sds_sd, ait_sd, sds_key, ait_keys, dims=None):
-        if sds_key + ".lora_down.weight" not in sds_sd:
-            return
-        down_weight = sds_sd.pop(sds_key + ".lora_down.weight")
-        up_weight = sds_sd.pop(sds_key + ".lora_up.weight")
-        sd_lora_rank = down_weight.shape[0]
-
-        default_alpha = torch.tensor(
-            sd_lora_rank, dtype=down_weight.dtype, device=down_weight.device, requires_grad=False
-        )
-        alpha = sds_sd.pop(sds_key + ".alpha", default_alpha)
-        scale = alpha / sd_lora_rank
-
-        scale_down = scale
-        scale_up = 1.0
-        while scale_down * 2 < scale_up:
-            scale_down *= 2
-            scale_up /= 2
-
-        down_weight = down_weight * scale_down
-        up_weight = up_weight * scale_up
-
-        num_splits = len(ait_keys)
-        if dims is None:
-            dims = [up_weight.shape[0] // num_splits] * num_splits
-        else:
-            assert sum(dims) == up_weight.shape[0]
-
-        # check if upweight is sparse
-        is_sparse = False
-        if sd_lora_rank % num_splits == 0:
-            ait_rank = sd_lora_rank // num_splits
-            is_sparse = True
-            i = 0
-            for j in range(len(dims)):
-                for k in range(len(dims)):
-                    if j == k:
-                        continue
-                    is_sparse = is_sparse and torch.all(
-                        up_weight[i : i + dims[j], k * ait_rank : (k + 1) * ait_rank] == 0
-                    )
-                i += dims[j]
-            if is_sparse:
-                logger.info(f"weight is sparse: {sds_key}")
-
-        ait_down_keys = [k + ".lora_A.weight" for k in ait_keys]
-        ait_up_keys = [k + ".lora_B.weight" for k in ait_keys]
-        if not is_sparse:
-            ait_sd.update(dict.fromkeys(ait_down_keys, down_weight))
-            ait_sd.update({k: v for k, v in zip(ait_up_keys, torch.split(up_weight, dims, dim=0))})  # noqa: C416
-        else:
-            ait_sd.update({k: v for k, v in zip(ait_down_keys, torch.chunk(down_weight, num_splits, dim=0))})  # noqa: C416
-            i = 0
-            for j in range(len(dims)):
-                ait_sd[ait_up_keys[j]] = up_weight[i : i + dims[j], j * ait_rank : (j + 1) * ait_rank].contiguous()
-                i += dims[j]
-
-    # Detect number of blocks from keys
-    num_double_layers = 0
-    num_single_layers = 0
-    for key in state_dict.keys():
-        if key.startswith("lora_unet_double_blocks_"):
-            block_idx = int(key.split("_")[4])
-            num_double_layers = max(num_double_layers, block_idx + 1)
-        elif key.startswith("lora_unet_single_blocks_"):
-            block_idx = int(key.split("_")[4])
-            num_single_layers = max(num_single_layers, block_idx + 1)
-
-    ait_sd = {}
-
-    for i in range(num_double_layers):
-        # Attention projections
-        _convert_to_ai_toolkit(
-            state_dict,
-            ait_sd,
-            f"lora_unet_double_blocks_{i}_img_attn_proj",
-            f"transformer.transformer_blocks.{i}.attn.to_out.0",
-        )
-        _convert_to_ai_toolkit_cat(
-            state_dict,
-            ait_sd,
-            f"lora_unet_double_blocks_{i}_img_attn_qkv",
-            [
-                f"transformer.transformer_blocks.{i}.attn.to_q",
-                f"transformer.transformer_blocks.{i}.attn.to_k",
-                f"transformer.transformer_blocks.{i}.attn.to_v",
-            ],
-        )
-        _convert_to_ai_toolkit(
-            state_dict,
-            ait_sd,
-            f"lora_unet_double_blocks_{i}_txt_attn_proj",
-            f"transformer.transformer_blocks.{i}.attn.to_add_out",
-        )
-        _convert_to_ai_toolkit_cat(
-            state_dict,
-            ait_sd,
-            f"lora_unet_double_blocks_{i}_txt_attn_qkv",
-            [
-                f"transformer.transformer_blocks.{i}.attn.add_q_proj",
-                f"transformer.transformer_blocks.{i}.attn.add_k_proj",
-                f"transformer.transformer_blocks.{i}.attn.add_v_proj",
-            ],
-        )
-        # MLP layers (Flux2 uses ff.linear_in/linear_out)
-        _convert_to_ai_toolkit(
-            state_dict,
-            ait_sd,
-            f"lora_unet_double_blocks_{i}_img_mlp_0",
-            f"transformer.transformer_blocks.{i}.ff.linear_in",
-        )
-        _convert_to_ai_toolkit(
-            state_dict,
-            ait_sd,
-            f"lora_unet_double_blocks_{i}_img_mlp_2",
-            f"transformer.transformer_blocks.{i}.ff.linear_out",
-        )
-        _convert_to_ai_toolkit(
-            state_dict,
-            ait_sd,
-            f"lora_unet_double_blocks_{i}_txt_mlp_0",
-            f"transformer.transformer_blocks.{i}.ff_context.linear_in",
-        )
-        _convert_to_ai_toolkit(
-            state_dict,
-            ait_sd,
-            f"lora_unet_double_blocks_{i}_txt_mlp_2",
-            f"transformer.transformer_blocks.{i}.ff_context.linear_out",
-        )
-
-    for i in range(num_single_layers):
-        # Single blocks: linear1 -> attn.to_qkv_mlp_proj (fused, no split needed)
-        _convert_to_ai_toolkit(
-            state_dict,
-            ait_sd,
-            f"lora_unet_single_blocks_{i}_linear1",
-            f"transformer.single_transformer_blocks.{i}.attn.to_qkv_mlp_proj",
-        )
-        # Single blocks: linear2 -> attn.to_out
-        _convert_to_ai_toolkit(
-            state_dict,
-            ait_sd,
-            f"lora_unet_single_blocks_{i}_linear2",
-            f"transformer.single_transformer_blocks.{i}.attn.to_out",
-        )
-
-    # Handle optional extra keys
-    extra_mappings = {
-        "lora_unet_img_in": "transformer.x_embedder",
-        "lora_unet_txt_in": "transformer.context_embedder",
-        "lora_unet_time_in_in_layer": "transformer.time_guidance_embed.timestep_embedder.linear_1",
-        "lora_unet_time_in_out_layer": "transformer.time_guidance_embed.timestep_embedder.linear_2",
-        "lora_unet_final_layer_linear": "transformer.proj_out",
-    }
-    for sds_key, ait_key in extra_mappings.items():
-        _convert_to_ai_toolkit(state_dict, ait_sd, sds_key, ait_key)
-
-    remaining_keys = list(state_dict.keys())
-    if remaining_keys:
-        logger.warning(f"Unsupported keys for Kohya Flux2 LoRA conversion: {remaining_keys}")
-
-    return ait_sd
-
-
 def _convert_non_diffusers_z_image_lora_to_diffusers(state_dict):
     """
     Convert non-diffusers ZImage LoRA state dict to diffusers format.

src/diffusers/loaders/lora_pipeline.py

@@ -43,7 +43,6 @@ from .lora_conversion_utils import (
     _convert_bfl_flux_control_lora_to_diffusers,
     _convert_fal_kontext_lora_to_diffusers,
     _convert_hunyuan_video_lora_to_diffusers,
-    _convert_kohya_flux2_lora_to_diffusers,
     _convert_kohya_flux_lora_to_diffusers,
     _convert_musubi_wan_lora_to_diffusers,
     _convert_non_diffusers_flux2_lora_to_diffusers,
@@ -5674,13 +5673,6 @@ class Flux2LoraLoaderMixin(LoraBaseMixin):
             logger.warning(warn_msg)
             state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
 
-        is_kohya = any(".lora_down.weight" in k for k in state_dict)
-        if is_kohya:
-            state_dict = _convert_kohya_flux2_lora_to_diffusers(state_dict)
-            # Kohya already takes care of scaling the LoRA parameters with alpha.
-            out = (state_dict, metadata) if return_lora_metadata else state_dict
-            return out
-
         is_peft_format = any(k.startswith("base_model.model.") for k in state_dict)
         if is_peft_format:
             state_dict = {k.replace("base_model.model.", "diffusion_model."): v for k, v in state_dict.items()}

src/diffusers/models/attention_dispatch.py

@@ -229,7 +229,6 @@ class AttentionBackendName(str, Enum):
     FLASH_HUB = "flash_hub"
     FLASH_VARLEN = "flash_varlen"
     FLASH_VARLEN_HUB = "flash_varlen_hub"
-    FLASH_4_HUB = "flash_4_hub"
     _FLASH_3 = "_flash_3"
     _FLASH_VARLEN_3 = "_flash_varlen_3"
     _FLASH_3_HUB = "_flash_3_hub"
@@ -359,11 +358,6 @@ _HUB_KERNELS_REGISTRY: dict["AttentionBackendName", _HubKernelConfig] = {
         function_attr="sageattn",
         version=1,
     ),
-    AttentionBackendName.FLASH_4_HUB: _HubKernelConfig(
-        repo_id="kernels-staging/flash-attn4",
-        function_attr="flash_attn_func",
-        version=0,
-    ),
 }
 
 
@@ -527,7 +521,6 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None
         AttentionBackendName._FLASH_3_HUB,
         AttentionBackendName._FLASH_3_VARLEN_HUB,
         AttentionBackendName.SAGE_HUB,
-        AttentionBackendName.FLASH_4_HUB,
     ]:
         if not is_kernels_available():
             raise RuntimeError(
@@ -538,11 +531,6 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None
                 f"Backend '{backend.value}' needs to be used with a `kernels` version of at least 0.12. Please update with `pip install -U kernels`."
             )
 
-        if backend == AttentionBackendName.FLASH_4_HUB and not is_kernels_available(">=", "0.12.3"):
-            raise RuntimeError(
-                f"Backend '{backend.value}' needs to be used with a `kernels` version of at least 0.12.3. Please update with `pip install -U kernels`."
-            )
-
     elif backend == AttentionBackendName.AITER:
         if not _CAN_USE_AITER_ATTN:
             raise RuntimeError(
@@ -2688,37 +2676,6 @@ def _flash_attention_3_varlen_hub(
     return (out, lse) if return_lse else out
 
 
-@_AttentionBackendRegistry.register(
-    AttentionBackendName.FLASH_4_HUB,
-    constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
-    supports_context_parallel=False,
-)
-def _flash_attention_4_hub(
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    attn_mask: torch.Tensor | None = None,
-    scale: float | None = None,
-    is_causal: bool = False,
-    return_lse: bool = False,
-    _parallel_config: "ParallelConfig" | None = None,
-) -> torch.Tensor:
-    if attn_mask is not None:
-        raise ValueError("`attn_mask` is not supported for flash-attn 4.")
-
-    func = _HUB_KERNELS_REGISTRY[AttentionBackendName.FLASH_4_HUB].kernel_fn
-    out = func(
-        q=query,
-        k=key,
-        v=value,
-        softmax_scale=scale,
-        causal=is_causal,
-    )
-    if isinstance(out, tuple):
-        return (out[0], out[1]) if return_lse else out[0]
-    return out
-
-
 @_AttentionBackendRegistry.register(
     AttentionBackendName._FLASH_VARLEN_3,
     constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],

src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py

@@ -324,18 +324,17 @@ class AudioLDM2Pipeline(DiffusionPipeline):
             `inputs_embeds (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`):
                 The sequence of generated hidden-states.
         """
+        cache_position_kwargs = {}
+        if is_transformers_version("<", "4.52.1"):
+            cache_position_kwargs["input_ids"] = inputs_embeds
+        else:
+            cache_position_kwargs["seq_length"] = inputs_embeds.shape[0]
+            cache_position_kwargs["device"] = (
+                self.language_model.device if getattr(self, "language_model", None) is not None else self.device
+            )
+        cache_position_kwargs["model_kwargs"] = model_kwargs
         max_new_tokens = max_new_tokens if max_new_tokens is not None else self.language_model.config.max_new_tokens
-        if hasattr(self.language_model, "_get_initial_cache_position"):
-            cache_position_kwargs = {}
-            if is_transformers_version("<", "4.52.1"):
-                cache_position_kwargs["input_ids"] = inputs_embeds
-            else:
-                cache_position_kwargs["seq_length"] = inputs_embeds.shape[0]
-                cache_position_kwargs["device"] = (
-                    self.language_model.device if getattr(self, "language_model", None) is not None else self.device
-                )
-            cache_position_kwargs["model_kwargs"] = model_kwargs
-            model_kwargs = self.language_model._get_initial_cache_position(**cache_position_kwargs)
+        model_kwargs = self.language_model._get_initial_cache_position(**cache_position_kwargs)
 
         for _ in range(max_new_tokens):
             # prepare model inputs

tests/lora/test_lora_layers_wanvace.py

@@ -28,6 +28,7 @@ from diffusers.utils.import_utils import is_peft_available
 
 from ..testing_utils import (
     floats_tensor,
+    is_flaky,
     require_peft_backend,
     require_peft_version_greater,
     skip_mps,
@@ -45,6 +46,7 @@ from .utils import PeftLoraLoaderMixinTests  # noqa: E402
 
 @require_peft_backend
 @skip_mps
+@is_flaky(max_attempts=10, description="very flaky class")
 class WanVACELoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):
     pipeline_class = WanVACEPipeline
     scheduler_cls = FlowMatchEulerDiscreteScheduler
@@ -71,8 +73,8 @@ class WanVACELoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):
         "base_dim": 3,
         "z_dim": 4,
         "dim_mult": [1, 1, 1, 1],
-        "latents_mean": [-0.7571, -0.7089, -0.9113, -0.7245],
-        "latents_std": [2.8184, 1.4541, 2.3275, 2.6558],
+        "latents_mean": torch.randn(4).numpy().tolist(),
+        "latents_std": torch.randn(4).numpy().tolist(),
         "num_res_blocks": 1,
         "temperal_downsample": [False, True, True],
     }

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -5,7 +5,6 @@ from typing import Callable
 
 import pytest
 import torch
-from huggingface_hub import hf_hub_download
 
 import diffusers
 from diffusers import AutoModel, ComponentsManager, ModularPipeline, ModularPipelineBlocks
@@ -33,33 +32,6 @@ from ..testing_utils import (
 )
 
 
-def _get_specified_components(path_or_repo_id, cache_dir=None):
-    if os.path.isdir(path_or_repo_id):
-        config_path = os.path.join(path_or_repo_id, "modular_model_index.json")
-    else:
-        try:
-            config_path = hf_hub_download(
-                repo_id=path_or_repo_id,
-                filename="modular_model_index.json",
-                local_dir=cache_dir,
-            )
-        except Exception:
-            return None
-
-    with open(config_path) as f:
-        config = json.load(f)
-
-    components = set()
-    for k, v in config.items():
-        if isinstance(v, (str, int, float, bool)):
-            continue
-        for entry in v:
-            if isinstance(entry, dict) and (entry.get("repo") or entry.get("pretrained_model_name_or_path")):
-                components.add(k)
-                break
-    return components
-
-
 class ModularPipelineTesterMixin:
     """
     It provides a set of common tests for each modular pipeline,
@@ -388,39 +360,6 @@ class ModularPipelineTesterMixin:
 
         assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
 
-    def test_load_expected_components_from_pretrained(self, tmp_path):
-        pipe = self.get_pipeline()
-        expected = _get_specified_components(self.pretrained_model_name_or_path, cache_dir=tmp_path)
-        if not expected:
-            pytest.skip("Skipping test as we couldn't fetch the expected components.")
-
-        actual = {
-            name
-            for name in pipe.components
-            if getattr(pipe, name, None) is not None
-            and getattr(getattr(pipe, name), "_diffusers_load_id", None) not in (None, "null")
-        }
-        assert expected == actual, f"Component mismatch: missing={expected - actual}, unexpected={actual - expected}"
-
-    def test_load_expected_components_from_save_pretrained(self, tmp_path):
-        pipe = self.get_pipeline()
-        save_dir = str(tmp_path / "saved-pipeline")
-        pipe.save_pretrained(save_dir)
-
-        expected = _get_specified_components(save_dir)
-        loaded_pipe = ModularPipeline.from_pretrained(save_dir)
-        loaded_pipe.load_components(torch_dtype=torch.float32)
-
-        actual = {
-            name
-            for name in loaded_pipe.components
-            if getattr(loaded_pipe, name, None) is not None
-            and getattr(getattr(loaded_pipe, name), "_diffusers_load_id", None) not in (None, "null")
-        }
-        assert expected == actual, (
-            f"Component mismatch after save/load: missing={expected - actual}, unexpected={actual - expected}"
-        )
-
     def test_modular_index_consistency(self, tmp_path):
         pipe = self.get_pipeline()
         components_spec = pipe._component_specs