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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
…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 <noreply@anthropic.com>
…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 <noreply@anthropic.com>
…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 <noreply@anthropic.com>
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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* remove test suites that are passed.

* fix consistencydecodervae tests

* Revert "fix consistencydecodervae tests"

This reverts commit 41b036b.

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
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 <spsayakpaul@gmail.com>

* 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 <spsayakpaul@gmail.com>

* refactor unet_2d tests (#13901)

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

* [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 <xin3.he@intel.com>

* add document and unit tests

Signed-off-by: Xin He <xin3.he@intel.com>

* fix CI

Signed-off-by: Xin He <xin3.he@intel.com>

* Apply suggestions from code review

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

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

Signed-off-by: Xin He <xin3.he@intel.com>

* add UT and fix bug

Signed-off-by: Xin He <xin3.he@intel.com>

* update per comments

Signed-off-by: Xin He <xin3.he@intel.com>

* update per comments

Signed-off-by: Xin He <xin3.he@intel.com>

* fix compile error in doc

Signed-off-by: Xin He <xin3.he@intel.com>

* Apply style fixes

* small nits

* Add auto_round dependency to the versions table

Signed-off-by: Xin He <xin3.he@intel.com>

* fix make deps_table_check_updated

Signed-off-by: Xin He <xin3.he@intel.com>

* fix CI

Signed-off-by: Xin He <xin3.he@intel.com>

---------

Signed-off-by: Xin He <xin3.he@intel.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* remove test suites that are passed.

* fix consistencydecodervae tests

* Revert "fix consistencydecodervae tests"

This reverts commit 41b036b.

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* [tests] fix vidtok tests (#13894)

* fix vidtok tests

* style

* Update tests/models/autoencoders/test_models_autoencoder_vidtok.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Apply style fixes

---------

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

* clean up

---------

Signed-off-by: Xin He <xin3.he@intel.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: Akshan Krithick <97239696+akshan-main@users.noreply.github.com>
Co-authored-by: Xin He <xin3.he@intel.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* remove test suites that are passed.

* fix consistencydecodervae tests

* Revert "fix consistencydecodervae tests"

This reverts commit 41b036b.

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
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 <spsayakpaul@gmail.com>

* 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 <spsayakpaul@gmail.com>

* refactor unet_2d tests (#13901)

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

* [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 <xin3.he@intel.com>

* add document and unit tests

Signed-off-by: Xin He <xin3.he@intel.com>

* fix CI

Signed-off-by: Xin He <xin3.he@intel.com>

* Apply suggestions from code review

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

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

Signed-off-by: Xin He <xin3.he@intel.com>

* add UT and fix bug

Signed-off-by: Xin He <xin3.he@intel.com>

* update per comments

Signed-off-by: Xin He <xin3.he@intel.com>

* update per comments

Signed-off-by: Xin He <xin3.he@intel.com>

* fix compile error in doc

Signed-off-by: Xin He <xin3.he@intel.com>

* Apply style fixes

* small nits

* Add auto_round dependency to the versions table

Signed-off-by: Xin He <xin3.he@intel.com>

* fix make deps_table_check_updated

Signed-off-by: Xin He <xin3.he@intel.com>

* fix CI

Signed-off-by: Xin He <xin3.he@intel.com>

---------

Signed-off-by: Xin He <xin3.he@intel.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* [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 <noreply@anthropic.com>

* remove test suites that are passed.

* fix consistencydecodervae tests

* Revert "fix consistencydecodervae tests"

This reverts commit 41b036b.

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* [tests] fix vidtok tests (#13894)

* fix vidtok tests

* style

* Update tests/models/autoencoders/test_models_autoencoder_vidtok.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Apply style fixes

---------

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

* clean up

---------

Signed-off-by: Xin He <xin3.he@intel.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: Akshan Krithick <97239696+akshan-main@users.noreply.github.com>
Co-authored-by: Xin He <xin3.he@intel.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
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