Thanks to visit codestin.com
Credit goes to github.com

Skip to content

[tests] refactor UNet model tests to align with the new pattern#13153

Merged
sayakpaul merged 30 commits into
mainfrom
unet-model-tests-refactor
Jun 10, 2026
Merged

[tests] refactor UNet model tests to align with the new pattern#13153
sayakpaul merged 30 commits into
mainfrom
unet-model-tests-refactor

Conversation

@sayakpaul

Copy link
Copy Markdown
Member

What does this PR do?

Some comments in-line. I have run the tests locally and all of them pass.

@sayakpaul sayakpaul requested a review from DN6 February 16, 2026 10:16
sayakpaul and others added 5 commits February 16, 2026 16:05
Refactor UNet1D model tests to follow the modern testing pattern using
BaseModelTesterConfig and focused mixin classes (ModelTesterMixin,
MemoryTesterMixin, TrainingTesterMixin, LoraTesterMixin).

Both UNet1D standard and RL variants now have separate config classes
and dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <[email protected]>
Refactor UNet2D model tests (standard, LDM, NCSN++) to follow the
modern testing pattern. Each variant gets its own config class and
dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <[email protected]>
…ing mixins

Refactor UNet3DConditionModel tests to follow the modern testing pattern
with separate classes for core, attention, memory, training, and LoRA.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
…ing mixins

Refactor UNetControlNetXSModel tests to follow the modern testing
pattern with separate classes for core, memory, training, and LoRA.
Specialized tests (from_unet, freeze_unet, forward_no_control,
time_embedding_mixing) remain in the core test class.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
…sting mixins

Refactored the spatiotemporal UNet test file to follow the modern modular testing
pattern with BaseModelTesterConfig and focused test classes:

- UNetSpatioTemporalTesterConfig: Base configuration with model setup
- TestUNetSpatioTemporal: Core model tests (ModelTesterMixin, UNetTesterMixin)
- TestUNetSpatioTemporalAttention: Attention-related tests (AttentionTesterMixin)
- TestUNetSpatioTemporalMemory: Memory/offloading tests (MemoryTesterMixin)
- TestUNetSpatioTemporalTraining: Training tests (TrainingTesterMixin)
- TestUNetSpatioTemporalLoRA: LoRA adapter tests (LoraTesterMixin)

Co-Authored-By: Claude Opus 4.6 <[email protected]>
Comment on lines -290 to +292
image = model(**self.get_dummy_inputs(), return_dict=False)[0]
new_image = new_model(**self.get_dummy_inputs(), return_dict=False)[0]
inputs_dict = self.get_dummy_inputs()
image = model(**inputs_dict, return_dict=False)[0]
new_image = new_model(**inputs_dict, return_dict=False)[0]

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To ensure reproducibility.

Comment on lines -95 to -97
if self.model_class.__name__ == "UNet2DConditionModel":
recompile_limit = 2

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not needed as we pass recompile_limit explicitly now.

with pytest.raises(RuntimeError, match=msg):
model.enable_lora_hotswap(target_rank=32)

def test_enable_lora_hotswap_called_after_adapter_added_warning(self, caplog):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It was needed because caplog doesn't capture these warnings properly.

@sayakpaul

Copy link
Copy Markdown
Member Author

@DN6 a gentle ping.

@github-actions github-actions Bot added tests size/L PR with diff > 200 LOC labels Apr 12, 2026
@sayakpaul sayakpaul requested a review from dg845 May 27, 2026 02:54
@sayakpaul

Copy link
Copy Markdown
Member Author

@dg845 a gentle ping.

@dg845 dg845 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR! Can you fix the merge conflicts?

@sayakpaul sayakpaul merged commit 1c7f759 into main Jun 10, 2026
29 of 31 checks passed
@sayakpaul sayakpaul deleted the unet-model-tests-refactor branch June 10, 2026 07:03
DN6 pushed a commit that referenced this pull request Jun 10, 2026
* refactor unet2d condition model tests.

* fix tests

* up

* fix

* Revert "fix"

This reverts commit 46d44b7.

* up

* recompile limit

* [tests] refactor test_models_unet_1d.py to use modular testing mixins

Refactor UNet1D model tests to follow the modern testing pattern using
BaseModelTesterConfig and focused mixin classes (ModelTesterMixin,
MemoryTesterMixin, TrainingTesterMixin, LoraTesterMixin).

Both UNet1D standard and RL variants now have separate config classes
and dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_2d.py to use modular testing mixins

Refactor UNet2D model tests (standard, LDM, NCSN++) to follow the
modern testing pattern. Each variant gets its own config class and
dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_3d_condition.py to use modular testing mixins

Refactor UNet3DConditionModel tests to follow the modern testing pattern
with separate classes for core, attention, memory, training, and LoRA.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_controlnetxs.py to use modular testing mixins

Refactor UNetControlNetXSModel tests to follow the modern testing
pattern with separate classes for core, memory, training, and LoRA.
Specialized tests (from_unet, freeze_unet, forward_no_control,
time_embedding_mixing) remain in the core test class.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_spatiotemporal.py to use modular testing mixins

Refactored the spatiotemporal UNet test file to follow the modern modular testing
pattern with BaseModelTesterConfig and focused test classes:

- UNetSpatioTemporalTesterConfig: Base configuration with model setup
- TestUNetSpatioTemporal: Core model tests (ModelTesterMixin, UNetTesterMixin)
- TestUNetSpatioTemporalAttention: Attention-related tests (AttentionTesterMixin)
- TestUNetSpatioTemporalMemory: Memory/offloading tests (MemoryTesterMixin)
- TestUNetSpatioTemporalTraining: Training tests (TrainingTesterMixin)
- TestUNetSpatioTemporalLoRA: LoRA adapter tests (LoraTesterMixin)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* remove test suites that are passed.

* fix consistencydecodervae tests

* Revert "fix consistencydecodervae tests"

This reverts commit 41b036b.

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: dg845 <[email protected]>
sayakpaul added a commit that referenced this pull request Jun 16, 2026
* update

* update

* update

* update

* [CI] Refactor SD3 Transformer Test (#13340)

* update

* update

---------

Co-authored-by: Sayak Paul <[email protected]>

* refactor unet tests (3d_condition, motion, controlnetxs) (#13897)

* refactor unet_3d_condition tests

* refactor unet_motion tests

* refactor unet_controlnetxs tests

* refactor unet_1d tests (#13898)

* refactor unet_1d tests

* use per-sample output_shape for unet_1d tests

---------

Co-authored-by: Sayak Paul <[email protected]>

* refactor unet_2d tests (#13901)

Co-authored-by: Sayak Paul <[email protected]>

* [chore] log quant config to the user_agent (#13850)

log quant config to the user_agent

* Integrate AutoRound into Diffusers (#13552)

* support auto_round

Signed-off-by: Xin He <[email protected]>

* add document and unit tests

Signed-off-by: Xin He <[email protected]>

* fix CI

Signed-off-by: Xin He <[email protected]>

* Apply suggestions from code review

Co-authored-by: Steven Liu <[email protected]>

* update document and overwrite the default quantization_config with specified backend.

Signed-off-by: Xin He <[email protected]>

* add UT and fix bug

Signed-off-by: Xin He <[email protected]>

* update per comments

Signed-off-by: Xin He <[email protected]>

* update per comments

Signed-off-by: Xin He <[email protected]>

* fix compile error in doc

Signed-off-by: Xin He <[email protected]>

* Apply style fixes

* small nits

* Add auto_round dependency to the versions table

Signed-off-by: Xin He <[email protected]>

* fix make deps_table_check_updated

Signed-off-by: Xin He <[email protected]>

* fix CI

Signed-off-by: Xin He <[email protected]>

---------

Signed-off-by: Xin He <[email protected]>
Co-authored-by: Steven Liu <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Sayak Paul <[email protected]>

* [tests] refactor UNet model tests to align with the new pattern (#13153)

* refactor unet2d condition model tests.

* fix tests

* up

* fix

* Revert "fix"

This reverts commit 46d44b7.

* up

* recompile limit

* [tests] refactor test_models_unet_1d.py to use modular testing mixins

Refactor UNet1D model tests to follow the modern testing pattern using
BaseModelTesterConfig and focused mixin classes (ModelTesterMixin,
MemoryTesterMixin, TrainingTesterMixin, LoraTesterMixin).

Both UNet1D standard and RL variants now have separate config classes
and dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_2d.py to use modular testing mixins

Refactor UNet2D model tests (standard, LDM, NCSN++) to follow the
modern testing pattern. Each variant gets its own config class and
dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_3d_condition.py to use modular testing mixins

Refactor UNet3DConditionModel tests to follow the modern testing pattern
with separate classes for core, attention, memory, training, and LoRA.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_controlnetxs.py to use modular testing mixins

Refactor UNetControlNetXSModel tests to follow the modern testing
pattern with separate classes for core, memory, training, and LoRA.
Specialized tests (from_unet, freeze_unet, forward_no_control,
time_embedding_mixing) remain in the core test class.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_spatiotemporal.py to use modular testing mixins

Refactored the spatiotemporal UNet test file to follow the modern modular testing
pattern with BaseModelTesterConfig and focused test classes:

- UNetSpatioTemporalTesterConfig: Base configuration with model setup
- TestUNetSpatioTemporal: Core model tests (ModelTesterMixin, UNetTesterMixin)
- TestUNetSpatioTemporalAttention: Attention-related tests (AttentionTesterMixin)
- TestUNetSpatioTemporalMemory: Memory/offloading tests (MemoryTesterMixin)
- TestUNetSpatioTemporalTraining: Training tests (TrainingTesterMixin)
- TestUNetSpatioTemporalLoRA: LoRA adapter tests (LoraTesterMixin)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* remove test suites that are passed.

* fix consistencydecodervae tests

* Revert "fix consistencydecodervae tests"

This reverts commit 41b036b.

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: dg845 <[email protected]>

* [tests] fix vidtok tests (#13894)

* fix vidtok tests

* style

* Update tests/models/autoencoders/test_models_autoencoder_vidtok.py

Co-authored-by: dg845 <[email protected]>

* Apply style fixes

---------

Co-authored-by: dg845 <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

* clean up

---------

Signed-off-by: Xin He <[email protected]>
Co-authored-by: Sayak Paul <[email protected]>
Co-authored-by: Akshan Krithick <[email protected]>
Co-authored-by: Xin He <[email protected]>
Co-authored-by: Steven Liu <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: dg845 <[email protected]>
DN6 pushed a commit that referenced this pull request Jul 1, 2026
* refactor unet2d condition model tests.

* fix tests

* up

* fix

* Revert "fix"

This reverts commit 46d44b7.

* up

* recompile limit

* [tests] refactor test_models_unet_1d.py to use modular testing mixins

Refactor UNet1D model tests to follow the modern testing pattern using
BaseModelTesterConfig and focused mixin classes (ModelTesterMixin,
MemoryTesterMixin, TrainingTesterMixin, LoraTesterMixin).

Both UNet1D standard and RL variants now have separate config classes
and dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_2d.py to use modular testing mixins

Refactor UNet2D model tests (standard, LDM, NCSN++) to follow the
modern testing pattern. Each variant gets its own config class and
dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_3d_condition.py to use modular testing mixins

Refactor UNet3DConditionModel tests to follow the modern testing pattern
with separate classes for core, attention, memory, training, and LoRA.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_controlnetxs.py to use modular testing mixins

Refactor UNetControlNetXSModel tests to follow the modern testing
pattern with separate classes for core, memory, training, and LoRA.
Specialized tests (from_unet, freeze_unet, forward_no_control,
time_embedding_mixing) remain in the core test class.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_spatiotemporal.py to use modular testing mixins

Refactored the spatiotemporal UNet test file to follow the modern modular testing
pattern with BaseModelTesterConfig and focused test classes:

- UNetSpatioTemporalTesterConfig: Base configuration with model setup
- TestUNetSpatioTemporal: Core model tests (ModelTesterMixin, UNetTesterMixin)
- TestUNetSpatioTemporalAttention: Attention-related tests (AttentionTesterMixin)
- TestUNetSpatioTemporalMemory: Memory/offloading tests (MemoryTesterMixin)
- TestUNetSpatioTemporalTraining: Training tests (TrainingTesterMixin)
- TestUNetSpatioTemporalLoRA: LoRA adapter tests (LoraTesterMixin)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* remove test suites that are passed.

* fix consistencydecodervae tests

* Revert "fix consistencydecodervae tests"

This reverts commit 41b036b.

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: dg845 <[email protected]>
DN6 added a commit that referenced this pull request Jul 1, 2026
* update

* update

* update

* update

* [CI] Refactor SD3 Transformer Test (#13340)

* update

* update

---------

Co-authored-by: Sayak Paul <[email protected]>

* refactor unet tests (3d_condition, motion, controlnetxs) (#13897)

* refactor unet_3d_condition tests

* refactor unet_motion tests

* refactor unet_controlnetxs tests

* refactor unet_1d tests (#13898)

* refactor unet_1d tests

* use per-sample output_shape for unet_1d tests

---------

Co-authored-by: Sayak Paul <[email protected]>

* refactor unet_2d tests (#13901)

Co-authored-by: Sayak Paul <[email protected]>

* [chore] log quant config to the user_agent (#13850)

log quant config to the user_agent

* Integrate AutoRound into Diffusers (#13552)

* support auto_round

Signed-off-by: Xin He <[email protected]>

* add document and unit tests

Signed-off-by: Xin He <[email protected]>

* fix CI

Signed-off-by: Xin He <[email protected]>

* Apply suggestions from code review

Co-authored-by: Steven Liu <[email protected]>

* update document and overwrite the default quantization_config with specified backend.

Signed-off-by: Xin He <[email protected]>

* add UT and fix bug

Signed-off-by: Xin He <[email protected]>

* update per comments

Signed-off-by: Xin He <[email protected]>

* update per comments

Signed-off-by: Xin He <[email protected]>

* fix compile error in doc

Signed-off-by: Xin He <[email protected]>

* Apply style fixes

* small nits

* Add auto_round dependency to the versions table

Signed-off-by: Xin He <[email protected]>

* fix make deps_table_check_updated

Signed-off-by: Xin He <[email protected]>

* fix CI

Signed-off-by: Xin He <[email protected]>

---------

Signed-off-by: Xin He <[email protected]>
Co-authored-by: Steven Liu <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Sayak Paul <[email protected]>

* [tests] refactor UNet model tests to align with the new pattern (#13153)

* refactor unet2d condition model tests.

* fix tests

* up

* fix

* Revert "fix"

This reverts commit 46d44b7.

* up

* recompile limit

* [tests] refactor test_models_unet_1d.py to use modular testing mixins

Refactor UNet1D model tests to follow the modern testing pattern using
BaseModelTesterConfig and focused mixin classes (ModelTesterMixin,
MemoryTesterMixin, TrainingTesterMixin, LoraTesterMixin).

Both UNet1D standard and RL variants now have separate config classes
and dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_2d.py to use modular testing mixins

Refactor UNet2D model tests (standard, LDM, NCSN++) to follow the
modern testing pattern. Each variant gets its own config class and
dedicated test classes organized by concern (core, memory, training,
LoRA, hub loading).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_3d_condition.py to use modular testing mixins

Refactor UNet3DConditionModel tests to follow the modern testing pattern
with separate classes for core, attention, memory, training, and LoRA.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_controlnetxs.py to use modular testing mixins

Refactor UNetControlNetXSModel tests to follow the modern testing
pattern with separate classes for core, memory, training, and LoRA.
Specialized tests (from_unet, freeze_unet, forward_no_control,
time_embedding_mixing) remain in the core test class.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* [tests] refactor test_models_unet_spatiotemporal.py to use modular testing mixins

Refactored the spatiotemporal UNet test file to follow the modern modular testing
pattern with BaseModelTesterConfig and focused test classes:

- UNetSpatioTemporalTesterConfig: Base configuration with model setup
- TestUNetSpatioTemporal: Core model tests (ModelTesterMixin, UNetTesterMixin)
- TestUNetSpatioTemporalAttention: Attention-related tests (AttentionTesterMixin)
- TestUNetSpatioTemporalMemory: Memory/offloading tests (MemoryTesterMixin)
- TestUNetSpatioTemporalTraining: Training tests (TrainingTesterMixin)
- TestUNetSpatioTemporalLoRA: LoRA adapter tests (LoraTesterMixin)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* remove test suites that are passed.

* fix consistencydecodervae tests

* Revert "fix consistencydecodervae tests"

This reverts commit 41b036b.

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: dg845 <[email protected]>

* [tests] fix vidtok tests (#13894)

* fix vidtok tests

* style

* Update tests/models/autoencoders/test_models_autoencoder_vidtok.py

Co-authored-by: dg845 <[email protected]>

* Apply style fixes

---------

Co-authored-by: dg845 <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

* clean up

---------

Signed-off-by: Xin He <[email protected]>
Co-authored-by: Sayak Paul <[email protected]>
Co-authored-by: Akshan Krithick <[email protected]>
Co-authored-by: Xin He <[email protected]>
Co-authored-by: Steven Liu <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: dg845 <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size/L PR with diff > 200 LOC tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants