[WIP] Refactor Model Tests

[DN6]

What does this PR do?

Following the plan outlined for Diffusers 1.0.0, this PR introduces changes to our model testing approach in order to reduce the overhead involved in adding comprehensive tests for new models and standardize tests across all models.

Changes include

1. Introducing feature specific tester Mixins and marks for models (breaking up the very large ModelTesterMixin class)
2. Introduce new test file structure using Config + Mixin pattern
3. New markers for selective test execution
4. Adds a ulility script (generate_model_tests.py) to automatically generate tests based on the model file. Also provide a flag that allows us to include any optional features to test e.g. (we can turn this into a bot down the line)

I've only made changes to Flux to make this PR easy to review. I'll open follow ups in phases for the other models once this is approved.

python utils/generate_model_tests.py src/diffusers/models/transformers/transformer_qwenimage.py 

Will now generate a template test file that can be populated with the necessary config information

# coding=utf-8
# 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 import QwenImageTransformer2DModel
from diffusers.utils.torch_utils import randn_tensor

from ...testing_utils import enable_full_determinism, torch_device
from ..test_modeling_common import LoraHotSwappingForModelTesterMixin
from ..testing_utils import (
    AttentionTesterMixin,
    ContextParallelTesterMixin,
    LoraTesterMixin,
    MemoryTesterMixin,
    ModelTesterMixin,
    TorchCompileTesterMixin,
    TrainingTesterMixin,
)


enable_full_determinism()


class QwenImageTransformerTesterConfig:
    model_class = QwenImageTransformer2DModel
    pretrained_model_name_or_path = ""
    pretrained_model_kwargs = {"subfolder": "transformer"}

    @property
    def generator(self):
        return torch.Generator("cpu").manual_seed(0)

    def get_init_dict(self) -> dict[str, int | list[int]]:
        # __init__ parameters:
        #   patch_size: int = 2
        #   in_channels: int = 64
        #   out_channels: Optional[int] = 16
        #   num_layers: int = 60
        #   attention_head_dim: int = 128
        #   num_attention_heads: int = 24
        #   joint_attention_dim: int = 3584
        #   guidance_embeds: bool = False
        #   axes_dims_rope: Tuple[int, int, int] = <complex>
        return {}

    def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
        # forward() parameters:
        #   hidden_states: torch.Tensor
        #   encoder_hidden_states: torch.Tensor
        #   encoder_hidden_states_mask: torch.Tensor
        #   timestep: torch.LongTensor
        #   img_shapes: Optional[List[Tuple[int, int, int]]]
        #   txt_seq_lens: Optional[List[int]]
        #   guidance: torch.Tensor
        #   attention_kwargs: Optional[Dict[str, Any]]
        #   controlnet_block_samples
        #   return_dict: bool = True
        # TODO: Fill in dummy inputs
        return {}

    @property
    def input_shape(self) -> tuple[int, ...]:
        return (1, 1)

    @property
    def output_shape(self) -> tuple[int, ...]:
        return (1, 1)


class TestQwenImageTransformerModel(QwenImageTransformerTesterConfig, ModelTesterMixin):
    pass


class TestQwenImageTransformerMemory(QwenImageTransformerTesterConfig, MemoryTesterMixin):
    pass


class TestQwenImageTransformerAttention(QwenImageTransformerTesterConfig, AttentionTesterMixin):
    pass


class TestQwenImageTransformerTorchCompile(QwenImageTransformerTesterConfig, TorchCompileTesterMixin):
    different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)]

    def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:
        # TODO: Implement dynamic input generation
        return {}


class TestQwenImageTransformerLora(QwenImageTransformerTesterConfig, LoraTesterMixin):
    pass


class TestQwenImageTransformerContextParallel(QwenImageTransformerTesterConfig, ContextParallelTesterMixin):
    pass


class TestQwenImageTransformerTraining(QwenImageTransformerTesterConfig, TrainingTesterMixin):
    pass


class TestQwenImageTransformerLoraHotSwappingForModel(QwenImageTransformerTesterConfig, LoraHotSwappingForModelTesterMixin):
    different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)]

    def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:
        # TODO: Implement dynamic input generation
        return {}

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[sayakpaul]

Excellent stuff! S

ome general comments:

• Normalize the model outputs to a common format before they go to torch.allclose().
• Initialize the input dict newly before passing to a new initialization of the model with torch.manual_seed(0). This is because some autoencoder models take a generator input.
• Use fixtures wherever possible to reduce boilerplate and take advantage of pytest features.
  • One particular session-level fixture could be base_output. It should help reduce test time quite a bit.
• Use pytest.mark.parametrize where possible.

Okay for me to do in a future PR but:

• Should also account for the attention backends.
• Should we also do a cross between CP and attention backends?
• How about the caching mixins?

Some nits:

• Use torch.no_grad() as an entire decorator as opposed to using it inside the functions.

[sayakpaul]

Do we want to club this with attention backends?

[DN6]

We can. But I think it's a bit out of scope for this PR because we need to create a container with the relevant backends available.

[sayakpaul]

We don't have to do it in this PR then.

But just as a note, most of the attention backends that have complex installation processes, we should actually encourage users to rely on their kernels-variants (FA2, FA3, SAGE). This way, we won't have to build any containers.

[sayakpaul]

How do we implement it in an individual model testing class? For example, say we want to skip it for model X where its attention class doesn't inherit from AttentionModuleMixin?

[DN6]

Ideally, any model using attention also uses AttentionModuleMixin. The options here

1. Do not add tests from attention Mixin to a module file.
2. Add a decorator that makes AttentionModuleMixin a requirement for running attention tests.

[sayakpaul]

But there are important classes like Autoencoders that don't use the Attention mixins.

Let's do this?

Add a decorator that makes AttentionModuleMixin a requirement for running attention tests.

[DN6]

But there are important classes like Autoencoders that don't use the Attention mixins.
The check is for AttentionModuleMixin not AttentionMixin and Autoencoders do use it

diffusers/src/diffusers/models/autoencoders/autoencoder_kl.py

Line 39 in 5851928

ModelMixin, AttentionMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin

[sayakpaul]

Oh, I see your point!

So, as long as there is an Attention module this class should apply.

So, maybe for each of the tests, at the beginning, we could check if their attention classes inherit from

diffusers/tests/models/testing_utils/attention.py

Line 76 in d9b73ff

if isinstance(module, AttentionModuleMixin):

and if that's not the case, we skip.

Otherwise, I think it could be cumbersome to check which model tests should and shouldn't use this class because attention is a common component.

[sayakpaul]

Very cool rewrite!

[sayakpaul]

There should also be a corresponding training test similar to:

diffusers/tests/quantization/bnb/test_4bit.py

Line 418 in 6708f5c

def test_training(self):

[sayakpaul]

dequantize() tests are missing from bitsandbytes but they exist:

diffusers/tests/quantization/bnb/test_4bit.py

Line 490 in 6708f5c

def test_generate_quality_dequantize(self):

[sayakpaul]

Where do we include:

diffusers/tests/quantization/test_torch_compile_utils.py

Line 27 in 6708f5c

class QuantCompileTests:

[sayakpaul]

I guess it cannot currently generate the dummy input and init dicts? I would understand if so because inferring those is quite non-trivial.

[sayakpaul]

Should we include a training test too?

[sayakpaul]

One thing I think we should do is get a coverage report for tests/models with `main and with this PR and confirm we are not skipping anything truly critical.

If we are, then we should likely be able to explain why that's the case.