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 DN6February 16, 2026 10:16
sayakpauland 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
MemberAuthor

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
MemberAuthor

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
MemberAuthor

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
MemberAuthor

@DN6 a gentle ping.

@github-actionsgithub-actionsBot added tests size/L PR with diff > 200 LOC labels Apr 12, 2026
@sayakpaul
sayakpaul requested a review from dg845May 27, 2026 02:54
@akshan-mainakshan-main mentioned this pull request May 28, 2026
6 tasks
@sayakpaul

Copy link
Copy Markdown
MemberAuthor

@dg845 a gentle ping.

@dg845dg845 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 mainJun 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 freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size/LPR with diff > 200 LOCtests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@sayakpaul@dg845