From c14a1e80ecb395bfd74fb8faa16ec90fcc113283 Mon Sep 17 00:00:00 2001 From: Akshan Krithick <97239696+akshan-main@users.noreply.github.com> Date: Wed, 26 Aug 2026 21:41:34 -0700 Subject: [PATCH] migrate wan vace lora tests to the pipeline-level mixins (#14614) Co-authored-by: Sayak Paul --- tests/lora/test_lora_layers_wanvace.py | 188 ------------------------- tests/pipelines/wan/test_wan_vace.py | 63 ++++++++- 2 files changed, 61 insertions(+), 190 deletions(-) delete mode 100644 tests/lora/test_lora_layers_wanvace.py diff --git a/tests/lora/test_lora_layers_wanvace.py b/tests/lora/test_lora_layers_wanvace.py deleted file mode 100644 index c17d0d07a0..0000000000 --- a/tests/lora/test_lora_layers_wanvace.py +++ /dev/null @@ -1,188 +0,0 @@ -# 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 os -import sys -import tempfile -import unittest - -import numpy as np -import safetensors.torch -import torch -from PIL import Image -from transformers import AutoTokenizer, T5EncoderModel - -from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanVACEPipeline, WanVACETransformer3DModel -from diffusers.utils.import_utils import is_peft_available - -from ..testing_utils import ( - floats_tensor, - require_peft_backend, - require_peft_version_greater, - skip_mps, - torch_device, -) - - -if is_peft_available(): - from peft.utils import get_peft_model_state_dict - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -@skip_mps -class WanVACELoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = WanVACEPipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - - transformer_kwargs = { - "patch_size": (1, 2, 2), - "num_attention_heads": 2, - "attention_head_dim": 8, - "in_channels": 4, - "out_channels": 4, - "text_dim": 32, - "freq_dim": 16, - "ffn_dim": 16, - "num_layers": 2, - "cross_attn_norm": True, - "qk_norm": "rms_norm_across_heads", - "rope_max_seq_len": 16, - "vace_layers": [0], - "vace_in_channels": 72, - } - transformer_cls = WanVACETransformer3DModel - vae_kwargs = { - "base_dim": 3, - "z_dim": 4, - "dim_mult": [1, 1, 1, 1], - "latents_mean": [-0.7571, -0.7089, -0.9113, -0.7245], - "latents_std": [2.8184, 1.4541, 2.3275, 2.6558], - "num_res_blocks": 1, - "temperal_downsample": [False, True, True], - } - vae_cls = AutoencoderKLWan - has_two_text_encoders = True - tokenizer_cls, tokenizer_id = AutoTokenizer, "hf-internal-testing/tiny-random-t5" - text_encoder_cls, text_encoder_id = T5EncoderModel, "hf-internal-testing/tiny-random-t5" - - text_encoder_target_modules = ["q", "k", "v", "o"] - - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 9, 16, 16, 3) - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 16 - num_channels = 4 - num_frames = 9 - num_latent_frames = 3 # (num_frames - 1) // temporal_compression_ratio + 1 - sizes = (4, 4) - height, width = 16, 16 - - generator = torch.manual_seed(0) - noise = floats_tensor((batch_size, num_latent_frames, num_channels) + sizes) - input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) - video = [Image.new("RGB", (height, width))] * num_frames - mask = [Image.new("L", (height, width), 0)] * num_frames - - pipeline_inputs = { - "video": video, - "mask": mask, - "prompt": "", - "num_frames": num_frames, - "num_inference_steps": 1, - "guidance_scale": 6.0, - "height": height, - "width": height, - "max_sequence_length": sequence_length, - "output_type": "np", - } - if with_generator: - pipeline_inputs.update({"generator": generator}) - - return noise, input_ids, pipeline_inputs - - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) - - @unittest.skip("Not supported in Wan VACE.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Wan VACE.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - - @require_peft_version_greater("0.13.2") - def test_lora_exclude_modules_wanvace(self): - exclude_module_name = "vace_blocks.0.proj_out" - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(torch_device) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - self.assertTrue(output_no_lora.shape == self.output_shape) - - # only supported for `denoiser` now - denoiser_lora_config.target_modules = ["proj_out"] - denoiser_lora_config.exclude_modules = [exclude_module_name] - pipe, _ = self.add_adapters_to_pipeline( - pipe, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config - ) - # The state dict shouldn't contain the modules to be excluded from LoRA. - state_dict_from_model = get_peft_model_state_dict(pipe.transformer, adapter_name="default") - self.assertTrue(not any(exclude_module_name in k for k in state_dict_from_model)) - self.assertTrue(any("proj_out" in k for k in state_dict_from_model)) - output_lora_exclude_modules = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdir: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights(save_directory=tmpdir, **lora_state_dicts) - pipe.unload_lora_weights() - - # Check in the loaded state dict. - loaded_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")) - self.assertTrue(not any(exclude_module_name in k for k in loaded_state_dict)) - self.assertTrue(any("proj_out" in k for k in loaded_state_dict)) - - # Check in the state dict obtained after loading LoRA. - pipe.load_lora_weights(tmpdir) - state_dict_from_model = get_peft_model_state_dict(pipe.transformer, adapter_name="default_0") - self.assertTrue(not any(exclude_module_name in k for k in state_dict_from_model)) - self.assertTrue(any("proj_out" in k for k in state_dict_from_model)) - - output_lora_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_no_lora, output_lora_exclude_modules, atol=1e-3, rtol=1e-3), - "LoRA should change outputs.", - ) - self.assertTrue( - np.allclose(output_lora_exclude_modules, output_lora_pretrained, atol=1e-3, rtol=1e-3), - "Lora outputs should match.", - ) - - def test_simple_inference_with_text_denoiser_lora_and_scale(self): - super().test_simple_inference_with_text_denoiser_lora_and_scale() diff --git a/tests/pipelines/wan/test_wan_vace.py b/tests/pipelines/wan/test_wan_vace.py index 6c4e825211..de070fa359 100644 --- a/tests/pipelines/wan/test_wan_vace.py +++ b/tests/pipelines/wan/test_wan_vace.py @@ -13,7 +13,10 @@ # limitations under the License. +import os + import pytest +import safetensors.torch import torch from PIL import Image from transformers import AutoConfig, AutoTokenizer, T5EncoderModel @@ -25,9 +28,20 @@ from diffusers import ( WanVACEPipeline, WanVACETransformer3DModel, ) +from diffusers.utils.import_utils import is_peft_available -from ...testing_utils import assert_tensors_close, torch_device -from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from ...testing_utils import assert_tensors_close, require_peft_version_greater, torch_device +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) + + +if is_peft_available(): + from peft.utils import get_peft_model_state_dict class WanVACEPipelineTesterConfig(BasePipelineTesterConfig): @@ -242,3 +256,48 @@ class TestWanVACEPipeline(WanVACEPipelineTesterConfig, PipelineTesterMixin): class TestWanVACEPipelineMemory(WanVACEPipelineTesterConfig, MemoryTesterMixin): pass + + +class TestWanVACEPipelineLoRA(WanVACEPipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the Wan VACE pipeline.""" + + @require_peft_version_greater("0.13.2") + def test_lora_exclude_modules(self, tmp_path, base_pipe_output): + exclude_module_name = "vace_blocks.0.proj_out" + pipe = self.get_pipeline().to(torch_device) + + self.add_adapters_to_pipeline( + pipe, components=["transformer"], target_modules=["proj_out"], exclude_modules=[exclude_module_name] + ) + # The state dict should not contain the modules excluded from LoRA. + state_dict_from_model = get_peft_model_state_dict(pipe.transformer, adapter_name="default") + assert not any(exclude_module_name in k for k in state_dict_from_model) + assert any("proj_out" in k for k in state_dict_from_model) + output_lora_exclude_modules = self.run_pipe(pipe) + + denoiser_state_dict = get_peft_model_state_dict(pipe.transformer) + self.pipeline_class.save_lora_weights(tmp_path, transformer_lora_layers=denoiser_state_dict) + pipe.unload_lora_weights() + + # Check in the saved state dict. + loaded_state_dict = safetensors.torch.load_file(os.path.join(tmp_path, "pytorch_lora_weights.safetensors")) + assert not any(exclude_module_name in k for k in loaded_state_dict) + assert any("proj_out" in k for k in loaded_state_dict) + + # Check in the state dict obtained after loading LoRA. + pipe.load_lora_weights(tmp_path) + state_dict_from_model = get_peft_model_state_dict(pipe.transformer, adapter_name="default_0") + assert not any(exclude_module_name in k for k in state_dict_from_model) + assert any("proj_out" in k for k in state_dict_from_model) + + output_lora_pretrained = self.run_pipe(pipe) + assert not torch.allclose(base_pipe_output, output_lora_exclude_modules, atol=1e-3, rtol=1e-3), ( + "LoRA should change outputs." + ) + assert torch.allclose(output_lora_exclude_modules, output_lora_pretrained, atol=1e-3, rtol=1e-3), ( + "Lora outputs should match." + ) + + +class TestWanVACEPipelineLoRAMemory(WanVACEPipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA offloading tests for the Wan VACE pipeline."""