Merge branch 'main' into klein-lora-fixes

refactor: use defaultdict for _SET_ADAPTER_SCALE_FN_MAPPING

Co-authored-by: Alexkkir <alexkkir@gmail.coom>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

src/diffusers/loaders/lora_conversion_utils.py

@@ -2443,6 +2443,191 @@ 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,6 +43,7 @@ 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,
@@ -5673,6 +5674,13 @@ 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/loaders/peft.py

@@ -15,6 +15,7 @@
 import inspect
 import json
 import os
+from collections import defaultdict
 from functools import partial
 from pathlib import Path
 from typing import Literal
@@ -44,33 +45,13 @@ from .unet_loader_utils import _maybe_expand_lora_scales
 
 logger = logging.get_logger(__name__)
 
-_SET_ADAPTER_SCALE_FN_MAPPING = {
-    "UNet2DConditionModel": _maybe_expand_lora_scales,
-    "UNetMotionModel": _maybe_expand_lora_scales,
-    "SD3Transformer2DModel": lambda model_cls, weights: weights,
-    "FluxTransformer2DModel": lambda model_cls, weights: weights,
-    "CogVideoXTransformer3DModel": lambda model_cls, weights: weights,
-    "ConsisIDTransformer3DModel": lambda model_cls, weights: weights,
-    "HeliosTransformer3DModel": lambda model_cls, weights: weights,
-    "MochiTransformer3DModel": lambda model_cls, weights: weights,
-    "HunyuanVideoTransformer3DModel": lambda model_cls, weights: weights,
-    "LTXVideoTransformer3DModel": lambda model_cls, weights: weights,
-    "SanaTransformer2DModel": lambda model_cls, weights: weights,
-    "AuraFlowTransformer2DModel": lambda model_cls, weights: weights,
-    "Lumina2Transformer2DModel": lambda model_cls, weights: weights,
-    "WanTransformer3DModel": lambda model_cls, weights: weights,
-    "CogView4Transformer2DModel": lambda model_cls, weights: weights,
-    "HiDreamImageTransformer2DModel": lambda model_cls, weights: weights,
-    "HunyuanVideoFramepackTransformer3DModel": lambda model_cls, weights: weights,
-    "WanVACETransformer3DModel": lambda model_cls, weights: weights,
-    "ChromaTransformer2DModel": lambda model_cls, weights: weights,
-    "ChronoEditTransformer3DModel": lambda model_cls, weights: weights,
-    "QwenImageTransformer2DModel": lambda model_cls, weights: weights,
-    "Flux2Transformer2DModel": lambda model_cls, weights: weights,
-    "ZImageTransformer2DModel": lambda model_cls, weights: weights,
-    "LTX2VideoTransformer3DModel": lambda model_cls, weights: weights,
-    "LTX2TextConnectors": lambda model_cls, weights: weights,
-}
+_SET_ADAPTER_SCALE_FN_MAPPING = defaultdict(
+    lambda: (lambda model_cls, weights: weights),
+    {
+        "UNet2DConditionModel": _maybe_expand_lora_scales,
+        "UNetMotionModel": _maybe_expand_lora_scales,
+    },
+)
 
 
 class PeftAdapterMixin:

src/diffusers/models/transformers/transformer_z_image.py

@@ -788,9 +788,12 @@ class ZImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOr
         freqs_cis = pad_sequence(freqs_cis, batch_first=True, padding_value=0.0)[:, : feats.shape[1]]
 
         # Attention mask
-        attn_mask = torch.zeros((bsz, max_seqlen), dtype=torch.bool, device=device)
-        for i, seq_len in enumerate(item_seqlens):
-            attn_mask[i, :seq_len] = 1
+        if all(seq == max_seqlen for seq in item_seqlens):
+            attn_mask = None
+        else:
+            attn_mask = torch.zeros((bsz, max_seqlen), dtype=torch.bool, device=device)
+            for i, seq_len in enumerate(item_seqlens):
+                attn_mask[i, :seq_len] = 1
 
         # Noise mask
         noise_mask_tensor = None
@@ -871,9 +874,12 @@ class ZImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOr
         unified_freqs = pad_sequence(unified_freqs, batch_first=True, padding_value=0.0)
 
         # Attention mask
-        attn_mask = torch.zeros((bsz, max_seqlen), dtype=torch.bool, device=device)
-        for i, seq_len in enumerate(unified_seqlens):
-            attn_mask[i, :seq_len] = 1
+        if all(seq == max_seqlen for seq in unified_seqlens):
+            attn_mask = None
+        else:
+            attn_mask = torch.zeros((bsz, max_seqlen), dtype=torch.bool, device=device)
+            for i, seq_len in enumerate(unified_seqlens):
+                attn_mask[i, :seq_len] = 1
 
         # Noise mask
         noise_mask_tensor = None

tests/modular_pipelines/test_conditional_pipeline_blocks.py

@@ -1,242 +0,0 @@
-# Copyright 2025 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.
-
-
-from diffusers.modular_pipelines import (
-    AutoPipelineBlocks,
-    ConditionalPipelineBlocks,
-    InputParam,
-    ModularPipelineBlocks,
-)
-
-
-class TextToImageBlock(ModularPipelineBlocks):
-    model_name = "text2img"
-
-    @property
-    def inputs(self):
-        return [InputParam(name="prompt")]
-
-    @property
-    def intermediate_outputs(self):
-        return []
-
-    @property
-    def description(self):
-        return "text-to-image workflow"
-
-    def __call__(self, components, state):
-        block_state = self.get_block_state(state)
-        block_state.workflow = "text2img"
-        self.set_block_state(state, block_state)
-        return components, state
-
-
-class ImageToImageBlock(ModularPipelineBlocks):
-    model_name = "img2img"
-
-    @property
-    def inputs(self):
-        return [InputParam(name="prompt"), InputParam(name="image")]
-
-    @property
-    def intermediate_outputs(self):
-        return []
-
-    @property
-    def description(self):
-        return "image-to-image workflow"
-
-    def __call__(self, components, state):
-        block_state = self.get_block_state(state)
-        block_state.workflow = "img2img"
-        self.set_block_state(state, block_state)
-        return components, state
-
-
-class InpaintBlock(ModularPipelineBlocks):
-    model_name = "inpaint"
-
-    @property
-    def inputs(self):
-        return [InputParam(name="prompt"), InputParam(name="image"), InputParam(name="mask")]
-
-    @property
-    def intermediate_outputs(self):
-        return []
-
-    @property
-    def description(self):
-        return "inpaint workflow"
-
-    def __call__(self, components, state):
-        block_state = self.get_block_state(state)
-        block_state.workflow = "inpaint"
-        self.set_block_state(state, block_state)
-        return components, state
-
-
-class ConditionalImageBlocks(ConditionalPipelineBlocks):
-    block_classes = [InpaintBlock, ImageToImageBlock, TextToImageBlock]
-    block_names = ["inpaint", "img2img", "text2img"]
-    block_trigger_inputs = ["mask", "image"]
-    default_block_name = "text2img"
-
-    @property
-    def description(self):
-        return "Conditional image blocks for testing"
-
-    def select_block(self, mask=None, image=None) -> str | None:
-        if mask is not None:
-            return "inpaint"
-        if image is not None:
-            return "img2img"
-        return None  # falls back to default_block_name
-
-
-class OptionalConditionalBlocks(ConditionalPipelineBlocks):
-    block_classes = [InpaintBlock, ImageToImageBlock]
-    block_names = ["inpaint", "img2img"]
-    block_trigger_inputs = ["mask", "image"]
-    default_block_name = None  # no default; block can be skipped
-
-    @property
-    def description(self):
-        return "Optional conditional blocks (skippable)"
-
-    def select_block(self, mask=None, image=None) -> str | None:
-        if mask is not None:
-            return "inpaint"
-        if image is not None:
-            return "img2img"
-        return None
-
-
-class AutoImageBlocks(AutoPipelineBlocks):
-    block_classes = [InpaintBlock, ImageToImageBlock, TextToImageBlock]
-    block_names = ["inpaint", "img2img", "text2img"]
-    block_trigger_inputs = ["mask", "image", None]
-
-    @property
-    def description(self):
-        return "Auto image blocks for testing"
-
-
-class TestConditionalPipelineBlocksSelectBlock:
-    def test_select_block_with_mask(self):
-        blocks = ConditionalImageBlocks()
-        assert blocks.select_block(mask="something") == "inpaint"
-
-    def test_select_block_with_image(self):
-        blocks = ConditionalImageBlocks()
-        assert blocks.select_block(image="something") == "img2img"
-
-    def test_select_block_with_mask_and_image(self):
-        blocks = ConditionalImageBlocks()
-        assert blocks.select_block(mask="m", image="i") == "inpaint"
-
-    def test_select_block_no_triggers_returns_none(self):
-        blocks = ConditionalImageBlocks()
-        assert blocks.select_block() is None
-
-    def test_select_block_explicit_none_values(self):
-        blocks = ConditionalImageBlocks()
-        assert blocks.select_block(mask=None, image=None) is None
-
-
-class TestConditionalPipelineBlocksWorkflowSelection:
-    def test_default_workflow_when_no_triggers(self):
-        blocks = ConditionalImageBlocks()
-        execution = blocks.get_execution_blocks()
-        assert execution is not None
-        assert isinstance(execution, TextToImageBlock)
-
-    def test_mask_trigger_selects_inpaint(self):
-        blocks = ConditionalImageBlocks()
-        execution = blocks.get_execution_blocks(mask=True)
-        assert isinstance(execution, InpaintBlock)
-
-    def test_image_trigger_selects_img2img(self):
-        blocks = ConditionalImageBlocks()
-        execution = blocks.get_execution_blocks(image=True)
-        assert isinstance(execution, ImageToImageBlock)
-
-    def test_mask_and_image_selects_inpaint(self):
-        blocks = ConditionalImageBlocks()
-        execution = blocks.get_execution_blocks(mask=True, image=True)
-        assert isinstance(execution, InpaintBlock)
-
-    def test_skippable_block_returns_none(self):
-        blocks = OptionalConditionalBlocks()
-        execution = blocks.get_execution_blocks()
-        assert execution is None
-
-    def test_skippable_block_still_selects_when_triggered(self):
-        blocks = OptionalConditionalBlocks()
-        execution = blocks.get_execution_blocks(image=True)
-        assert isinstance(execution, ImageToImageBlock)
-
-
-class TestAutoPipelineBlocksSelectBlock:
-    def test_auto_select_mask(self):
-        blocks = AutoImageBlocks()
-        assert blocks.select_block(mask="m") == "inpaint"
-
-    def test_auto_select_image(self):
-        blocks = AutoImageBlocks()
-        assert blocks.select_block(image="i") == "img2img"
-
-    def test_auto_select_default(self):
-        blocks = AutoImageBlocks()
-        # No trigger -> returns None -> falls back to default (text2img)
-        assert blocks.select_block() is None
-
-    def test_auto_select_priority_order(self):
-        blocks = AutoImageBlocks()
-        assert blocks.select_block(mask="m", image="i") == "inpaint"
-
-
-class TestAutoPipelineBlocksWorkflowSelection:
-    def test_auto_default_workflow(self):
-        blocks = AutoImageBlocks()
-        execution = blocks.get_execution_blocks()
-        assert isinstance(execution, TextToImageBlock)
-
-    def test_auto_mask_workflow(self):
-        blocks = AutoImageBlocks()
-        execution = blocks.get_execution_blocks(mask=True)
-        assert isinstance(execution, InpaintBlock)
-
-    def test_auto_image_workflow(self):
-        blocks = AutoImageBlocks()
-        execution = blocks.get_execution_blocks(image=True)
-        assert isinstance(execution, ImageToImageBlock)
-
-
-class TestConditionalPipelineBlocksStructure:
-    def test_block_names_accessible(self):
-        blocks = ConditionalImageBlocks()
-        sub = dict(blocks.sub_blocks)
-        assert set(sub.keys()) == {"inpaint", "img2img", "text2img"}
-
-    def test_sub_block_types(self):
-        blocks = ConditionalImageBlocks()
-        sub = dict(blocks.sub_blocks)
-        assert isinstance(sub["inpaint"], InpaintBlock)
-        assert isinstance(sub["img2img"], ImageToImageBlock)
-        assert isinstance(sub["text2img"], TextToImageBlock)
-
-    def test_description(self):
-        blocks = ConditionalImageBlocks()
-        assert "Conditional" in blocks.description

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -10,6 +10,11 @@ from huggingface_hub import hf_hub_download
 import diffusers
 from diffusers import AutoModel, ComponentsManager, ModularPipeline, ModularPipelineBlocks
 from diffusers.guiders import ClassifierFreeGuidance
+from diffusers.modular_pipelines import (
+    ConditionalPipelineBlocks,
+    LoopSequentialPipelineBlocks,
+    SequentialPipelineBlocks,
+)
 from diffusers.modular_pipelines.modular_pipeline_utils import (
     ComponentSpec,
     ConfigSpec,
@@ -20,6 +25,7 @@ from diffusers.modular_pipelines.modular_pipeline_utils import (
 from diffusers.utils import logging
 
 from ..testing_utils import (
+    CaptureLogger,
     backend_empty_cache,
     numpy_cosine_similarity_distance,
     require_accelerator,
@@ -492,6 +498,117 @@ class ModularGuiderTesterMixin:
         assert max_diff > expected_max_diff, "Output with CFG must be different from normal inference"
 
 
+class TestCustomBlockRequirements:
+    def get_dummy_block_pipe(self):
+        class DummyBlockOne:
+            # keep two arbitrary deps so that we can test warnings.
+            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
+
+        class DummyBlockTwo:
+            # keep two dependencies that will be available during testing.
+            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
+
+        pipe = SequentialPipelineBlocks.from_blocks_dict(
+            {"dummy_block_one": DummyBlockOne, "dummy_block_two": DummyBlockTwo}
+        )
+        return pipe
+
+    def get_dummy_conditional_block_pipe(self):
+        class DummyBlockOne:
+            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
+
+        class DummyBlockTwo:
+            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
+
+        class DummyConditionalBlocks(ConditionalPipelineBlocks):
+            block_classes = [DummyBlockOne, DummyBlockTwo]
+            block_names = ["block_one", "block_two"]
+            block_trigger_inputs = []
+
+            def select_block(self, **kwargs):
+                return "block_one"
+
+        return DummyConditionalBlocks()
+
+    def get_dummy_loop_block_pipe(self):
+        class DummyBlockOne:
+            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
+
+        class DummyBlockTwo:
+            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
+
+        return LoopSequentialPipelineBlocks.from_blocks_dict({"block_one": DummyBlockOne, "block_two": DummyBlockTwo})
+
+    def test_sequential_block_requirements_save_load(self, tmp_path):
+        pipe = self.get_dummy_block_pipe()
+        pipe.save_pretrained(str(tmp_path))
+
+        config_path = tmp_path / "modular_config.json"
+
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        assert "requirements" in config
+        requirements = config["requirements"]
+
+        expected_requirements = {
+            "xyz": ">=0.8.0",
+            "abc": ">=10.0.0",
+            "transformers": ">=4.44.0",
+            "diffusers": ">=0.2.0",
+        }
+        assert expected_requirements == requirements
+
+    def test_sequential_block_requirements_warnings(self, tmp_path):
+        pipe = self.get_dummy_block_pipe()
+
+        logger = logging.get_logger("diffusers.modular_pipelines.modular_pipeline_utils")
+        logger.setLevel(30)
+
+        with CaptureLogger(logger) as cap_logger:
+            pipe.save_pretrained(str(tmp_path))
+
+        template = "{req} was specified in the requirements but wasn't found in the current environment"
+        msg_xyz = template.format(req="xyz")
+        msg_abc = template.format(req="abc")
+        assert msg_xyz in str(cap_logger.out)
+        assert msg_abc in str(cap_logger.out)
+
+    def test_conditional_block_requirements_save_load(self, tmp_path):
+        pipe = self.get_dummy_conditional_block_pipe()
+        pipe.save_pretrained(str(tmp_path))
+
+        config_path = tmp_path / "modular_config.json"
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        assert "requirements" in config
+        expected_requirements = {
+            "xyz": ">=0.8.0",
+            "abc": ">=10.0.0",
+            "transformers": ">=4.44.0",
+            "diffusers": ">=0.2.0",
+        }
+        assert expected_requirements == config["requirements"]
+
+    def test_loop_block_requirements_save_load(self, tmp_path):
+        pipe = self.get_dummy_loop_block_pipe()
+        pipe.save_pretrained(str(tmp_path))
+
+        config_path = tmp_path / "modular_config.json"
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        assert "requirements" in config
+        expected_requirements = {
+            "xyz": ">=0.8.0",
+            "abc": ">=10.0.0",
+            "transformers": ">=4.44.0",
+            "diffusers": ">=0.2.0",
+        }
+        assert expected_requirements == config["requirements"]
+
+
 class TestModularModelCardContent:
     def create_mock_block(self, name="TestBlock", description="Test block description"):
         class MockBlock:

tests/modular_pipelines/test_modular_pipelines_custom_blocks.py

@@ -24,18 +24,14 @@ import torch
 from diffusers import FluxTransformer2DModel
 from diffusers.modular_pipelines import (
     ComponentSpec,
-    ConditionalPipelineBlocks,
     InputParam,
-    LoopSequentialPipelineBlocks,
     ModularPipelineBlocks,
     OutputParam,
     PipelineState,
-    SequentialPipelineBlocks,
     WanModularPipeline,
 )
-from diffusers.utils import logging
 
-from ..testing_utils import CaptureLogger, nightly, require_torch, require_torch_accelerator, slow, torch_device
+from ..testing_utils import nightly, require_torch, require_torch_accelerator, slow, torch_device
 
 
 def _create_tiny_model_dir(model_dir):
@@ -467,117 +463,6 @@ class TestModularCustomBlocks:
         assert output_prompt.startswith("Modular diffusers + ")
 
 
-class TestCustomBlockRequirements:
-    def get_dummy_block_pipe(self):
-        class DummyBlockOne:
-            # keep two arbitrary deps so that we can test warnings.
-            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
-
-        class DummyBlockTwo:
-            # keep two dependencies that will be available during testing.
-            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
-
-        pipe = SequentialPipelineBlocks.from_blocks_dict(
-            {"dummy_block_one": DummyBlockOne, "dummy_block_two": DummyBlockTwo}
-        )
-        return pipe
-
-    def get_dummy_conditional_block_pipe(self):
-        class DummyBlockOne:
-            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
-
-        class DummyBlockTwo:
-            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
-
-        class DummyConditionalBlocks(ConditionalPipelineBlocks):
-            block_classes = [DummyBlockOne, DummyBlockTwo]
-            block_names = ["block_one", "block_two"]
-            block_trigger_inputs = []
-
-            def select_block(self, **kwargs):
-                return "block_one"
-
-        return DummyConditionalBlocks()
-
-    def get_dummy_loop_block_pipe(self):
-        class DummyBlockOne:
-            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
-
-        class DummyBlockTwo:
-            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
-
-        return LoopSequentialPipelineBlocks.from_blocks_dict({"block_one": DummyBlockOne, "block_two": DummyBlockTwo})
-
-    def test_sequential_block_requirements_save_load(self, tmp_path):
-        pipe = self.get_dummy_block_pipe()
-        pipe.save_pretrained(str(tmp_path))
-
-        config_path = tmp_path / "modular_config.json"
-
-        with open(config_path, "r") as f:
-            config = json.load(f)
-
-        assert "requirements" in config
-        requirements = config["requirements"]
-
-        expected_requirements = {
-            "xyz": ">=0.8.0",
-            "abc": ">=10.0.0",
-            "transformers": ">=4.44.0",
-            "diffusers": ">=0.2.0",
-        }
-        assert expected_requirements == requirements
-
-    def test_sequential_block_requirements_warnings(self, tmp_path):
-        pipe = self.get_dummy_block_pipe()
-
-        logger = logging.get_logger("diffusers.modular_pipelines.modular_pipeline_utils")
-        logger.setLevel(30)
-
-        with CaptureLogger(logger) as cap_logger:
-            pipe.save_pretrained(str(tmp_path))
-
-        template = "{req} was specified in the requirements but wasn't found in the current environment"
-        msg_xyz = template.format(req="xyz")
-        msg_abc = template.format(req="abc")
-        assert msg_xyz in str(cap_logger.out)
-        assert msg_abc in str(cap_logger.out)
-
-    def test_conditional_block_requirements_save_load(self, tmp_path):
-        pipe = self.get_dummy_conditional_block_pipe()
-        pipe.save_pretrained(str(tmp_path))
-
-        config_path = tmp_path / "modular_config.json"
-        with open(config_path, "r") as f:
-            config = json.load(f)
-
-        assert "requirements" in config
-        expected_requirements = {
-            "xyz": ">=0.8.0",
-            "abc": ">=10.0.0",
-            "transformers": ">=4.44.0",
-            "diffusers": ">=0.2.0",
-        }
-        assert expected_requirements == config["requirements"]
-
-    def test_loop_block_requirements_save_load(self, tmp_path):
-        pipe = self.get_dummy_loop_block_pipe()
-        pipe.save_pretrained(str(tmp_path))
-
-        config_path = tmp_path / "modular_config.json"
-        with open(config_path, "r") as f:
-            config = json.load(f)
-
-        assert "requirements" in config
-        expected_requirements = {
-            "xyz": ">=0.8.0",
-            "abc": ">=10.0.0",
-            "transformers": ">=4.44.0",
-            "diffusers": ">=0.2.0",
-        }
-        assert expected_requirements == config["requirements"]
-
-
 @slow
 @nightly
 @require_torch

utils/fetch_torch_cuda_pipeline_test_matrix.py

@@ -43,7 +43,7 @@ def filter_pipelines(usage_dict, usage_cutoff=10000):
 
 
 def fetch_pipeline_objects():
-    models = api.list_models(library="diffusers")
+    models = api.list_models(filter="diffusers")
     downloads = defaultdict(int)
 
     for model in models: