# coding=utf-8 # Copyright 2026 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 copy import os import torch from diffusers import MotionAdapter, UNet2DConditionModel, UNetMotionModel from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin, ) enable_full_determinism() class UNetMotionModelTesterConfig(BaseModelTesterConfig): @property def model_class(self): return UNetMotionModel @property def main_input_name(self) -> str: return "sample" @property def output_shape(self) -> tuple: return (4, 4, 16, 16) @property def generator(self): return torch.Generator("cpu").manual_seed(0) def get_init_dict(self) -> dict: return { "block_out_channels": (16, 32), "norm_num_groups": 16, "down_block_types": ("CrossAttnDownBlockMotion", "DownBlockMotion"), "up_block_types": ("UpBlockMotion", "CrossAttnUpBlockMotion"), "cross_attention_dim": 16, "num_attention_heads": 2, "out_channels": 4, "in_channels": 4, "layers_per_block": 1, "sample_size": 16, } def get_dummy_inputs(self) -> dict: batch_size = 4 num_channels = 4 num_frames = 4 sizes = (16, 16) noise = randn_tensor( (batch_size, num_channels, num_frames, *sizes), generator=self.generator, device=torch_device ) timestep = torch.tensor([10], device=torch_device) encoder_hidden_states = randn_tensor( (batch_size * num_frames, 4, 16), generator=self.generator, device=torch_device ) return {"sample": noise, "timestep": timestep, "encoder_hidden_states": encoder_hidden_states} class TestUNetMotionModel(UNetMotionModelTesterConfig, ModelTesterMixin): def test_from_unet2d(self): torch.manual_seed(0) unet2d = UNet2DConditionModel() torch.manual_seed(1) model = self.model_class.from_unet2d(unet2d) model_state_dict = model.state_dict() for param_name, param_value in unet2d.named_parameters(): assert torch.equal(model_state_dict[param_name], param_value) def test_freeze_unet2d(self): model = self.model_class(**self.get_init_dict()) model.freeze_unet2d_params() for param_name, param_value in model.named_parameters(): if "motion_modules" not in param_name: assert not param_value.requires_grad else: assert param_value.requires_grad def test_loading_motion_adapter(self): model = self.model_class() adapter = MotionAdapter() model.load_motion_modules(adapter) for idx, down_block in enumerate(model.down_blocks): adapter_state_dict = adapter.down_blocks[idx].motion_modules.state_dict() for param_name, param_value in down_block.motion_modules.named_parameters(): assert torch.equal(adapter_state_dict[param_name], param_value) for idx, up_block in enumerate(model.up_blocks): adapter_state_dict = adapter.up_blocks[idx].motion_modules.state_dict() for param_name, param_value in up_block.motion_modules.named_parameters(): assert torch.equal(adapter_state_dict[param_name], param_value) mid_block_adapter_state_dict = adapter.mid_block.motion_modules.state_dict() for param_name, param_value in model.mid_block.motion_modules.named_parameters(): assert torch.equal(mid_block_adapter_state_dict[param_name], param_value) def test_saving_motion_modules(self, tmp_path): torch.manual_seed(0) init_dict = self.get_init_dict() model = self.model_class(**init_dict).to(torch_device) model.save_motion_modules(tmp_path) assert os.path.isfile(os.path.join(tmp_path, "diffusion_pytorch_model.safetensors")) adapter_loaded = MotionAdapter.from_pretrained(tmp_path) torch.manual_seed(0) model_loaded = self.model_class(**init_dict) model_loaded.load_motion_modules(adapter_loaded) model_loaded.to(torch_device) with torch.no_grad(): output = model(**self.get_dummy_inputs())[0] output_loaded = model_loaded(**self.get_dummy_inputs())[0] assert (output - output_loaded).abs().max().item() <= 1e-4, "Models give different forward passes" def test_feed_forward_chunking(self): init_dict = self.get_init_dict() init_dict["block_out_channels"] = (32, 64) init_dict["norm_num_groups"] = 32 model = self.model_class(**init_dict).to(torch_device).eval() with torch.no_grad(): output = model(**self.get_dummy_inputs())[0] model.enable_forward_chunking() with torch.no_grad(): output_2 = model(**self.get_dummy_inputs())[0] assert output.shape == output_2.shape, "Shape doesn't match" assert (output - output_2).abs().max() < 1e-2 def test_pickle(self): model = self.model_class(**self.get_init_dict()).to(torch_device) with torch.no_grad(): sample = model(**self.get_dummy_inputs()).sample sample_copy = copy.copy(sample) assert (sample - sample_copy).abs().max() < 1e-4 def test_forward_with_norm_groups(self): init_dict = self.get_init_dict() init_dict["norm_num_groups"] = 16 init_dict["block_out_channels"] = (16, 32) model = self.model_class(**init_dict).to(torch_device).eval() with torch.no_grad(): output = model(**self.get_dummy_inputs()).sample assert output.shape == self.get_dummy_inputs()["sample"].shape, "Input and output shapes do not match" def test_asymmetric_motion_model(self): init_dict = self.get_init_dict() init_dict["layers_per_block"] = (2, 3) init_dict["transformer_layers_per_block"] = ((1, 2), (3, 4, 5)) init_dict["reverse_transformer_layers_per_block"] = ((7, 6, 7, 4), (4, 2, 2)) init_dict["temporal_transformer_layers_per_block"] = ((2, 5), (2, 3, 5)) init_dict["reverse_temporal_transformer_layers_per_block"] = ((5, 4, 3, 4), (3, 2, 2)) init_dict["num_attention_heads"] = (2, 4) init_dict["motion_num_attention_heads"] = (4, 4) init_dict["reverse_motion_num_attention_heads"] = (2, 2) init_dict["use_motion_mid_block"] = True init_dict["mid_block_layers"] = 2 init_dict["transformer_layers_per_mid_block"] = (1, 5) init_dict["temporal_transformer_layers_per_mid_block"] = (2, 4) model = self.model_class(**init_dict).to(torch_device).eval() with torch.no_grad(): output = model(**self.get_dummy_inputs()).sample assert output.shape == self.get_dummy_inputs()["sample"].shape, "Input and output shapes do not match" class TestUNetMotionModelTraining(UNetMotionModelTesterConfig, TrainingTesterMixin): """Training tests for UNetMotionModel.""" def test_gradient_checkpointing_is_applied(self): expected_set = { "CrossAttnUpBlockMotion", "CrossAttnDownBlockMotion", "UNetMidBlockCrossAttnMotion", "UpBlockMotion", "Transformer2DModel", "DownBlockMotion", } super().test_gradient_checkpointing_is_applied(expected_set=expected_set) class TestUNetMotionModelMemory(UNetMotionModelTesterConfig, MemoryTesterMixin): """Memory optimization tests for UNetMotionModel.""" class TestUNetMotionModelAttention(UNetMotionModelTesterConfig, AttentionTesterMixin): """Attention processor tests for UNetMotionModel."""