diff --git a/tests/lora/test_lora_layers_ace_step.py b/tests/lora/test_lora_layers_ace_step.py deleted file mode 100644 index 753c1f32f073..000000000000 --- a/tests/lora/test_lora_layers_ace_step.py +++ /dev/null @@ -1,179 +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 sys -import unittest - -import torch -from transformers import AutoTokenizer, Qwen3Config, Qwen3Model - -from diffusers import AutoencoderOobleck, FlowMatchEulerDiscreteScheduler -from diffusers.models.transformers.ace_step_transformer import AceStepTransformer1DModel -from diffusers.pipelines.ace_step import AceStepConditionEncoder, AceStepPipeline -from diffusers.utils.import_utils import is_peft_available - -from ..testing_utils import require_peft_backend, skip_mps - - -if is_peft_available(): - from peft import LoraConfig - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -@skip_mps -class AceStepLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = AceStepPipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {"num_train_timesteps": 1, "shift": 1.0} - - transformer_cls = AceStepTransformer1DModel - transformer_kwargs = { - "hidden_size": 32, - "intermediate_size": 64, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 8, - "in_channels": 24, - "audio_acoustic_hidden_dim": 8, - "patch_size": 2, - "rope_theta": 10000.0, - "sliding_window": 16, - } - - vae_cls = AutoencoderOobleck - vae_kwargs = { - "encoder_hidden_size": 6, - "downsampling_ratios": [1, 2], - "decoder_channels": 3, - "decoder_input_channels": 8, - "audio_channels": 2, - "channel_multiples": [2, 4], - "sampling_rate": 4, - } - - tokenizer_cls, tokenizer_id = AutoTokenizer, "Qwen/Qwen3-Embedding-0.6B" - text_encoder_cls, text_encoder_id = None, None - - text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"] - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 2, 1) - - def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): - scheduler_cls = scheduler_cls or self.scheduler_cls - rank = 4 - lora_alpha = rank if lora_alpha is None else lora_alpha - - torch.manual_seed(0) - transformer = self.transformer_cls(**self.transformer_kwargs) - - scheduler = scheduler_cls(**self.scheduler_kwargs) - - torch.manual_seed(0) - vae = self.vae_cls(**self.vae_kwargs) - - torch.manual_seed(0) - qwen3_config = Qwen3Config( - hidden_size=32, - intermediate_size=64, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - head_dim=8, - vocab_size=151936, - max_position_embeddings=256, - ) - text_encoder = Qwen3Model(qwen3_config) - tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_id) - - torch.manual_seed(0) - condition_encoder = AceStepConditionEncoder( - hidden_size=32, - intermediate_size=64, - text_hidden_dim=32, - timbre_hidden_dim=8, - num_lyric_encoder_hidden_layers=2, - num_timbre_encoder_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - head_dim=8, - rope_theta=10000.0, - sliding_window=16, - ) - - text_lora_config = LoraConfig( - r=rank, - lora_alpha=lora_alpha, - target_modules=self.text_encoder_target_modules, - init_lora_weights=False, - use_dora=use_dora, - ) - - denoiser_lora_config = LoraConfig( - r=rank, - lora_alpha=lora_alpha, - target_modules=self.denoiser_target_modules, - init_lora_weights=False, - use_dora=use_dora, - ) - - pipeline_components = { - "scheduler": scheduler, - "vae": vae, - "text_encoder": text_encoder, - "tokenizer": tokenizer, - "transformer": transformer, - "condition_encoder": condition_encoder, - "audio_tokenizer": None, - "audio_token_detokenizer": None, - } - - return pipeline_components, text_lora_config, denoiser_lora_config - - def get_dummy_inputs(self, with_generator=True): - generator = torch.manual_seed(0) - noise = torch.randn(1, 4, 8) - input_ids = torch.randint(1, 10, size=(1, 10), generator=generator) - - pipeline_inputs = { - "prompt": "A beautiful piano piece", - "lyrics": "[verse]\nSoft notes", - "audio_duration": 0.4, - "num_inference_steps": 2, - "max_text_length": 32, - "output_type": "np", - } - if with_generator: - pipeline_inputs["generator"] = generator - - return noise, input_ids, pipeline_inputs - - @unittest.skip("Not supported in AceStep.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in AceStep.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - - @unittest.skip("Not supported in AceStep.") - def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): - pass diff --git a/tests/lora/test_lora_layers_auraflow.py b/tests/lora/test_lora_layers_auraflow.py deleted file mode 100644 index c3ac8bd08d8a..000000000000 --- a/tests/lora/test_lora_layers_auraflow.py +++ /dev/null @@ -1,114 +0,0 @@ -# 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 sys -import unittest - -import torch -from transformers import AutoTokenizer, UMT5EncoderModel - -from diffusers import ( - AuraFlowPipeline, - AuraFlowTransformer2DModel, - FlowMatchEulerDiscreteScheduler, -) - -from ..testing_utils import ( - floats_tensor, - is_peft_available, - require_peft_backend, -) - - -if is_peft_available(): - pass - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -class AuraFlowLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = AuraFlowPipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - - transformer_kwargs = { - "sample_size": 64, - "patch_size": 1, - "in_channels": 4, - "num_mmdit_layers": 1, - "num_single_dit_layers": 1, - "attention_head_dim": 16, - "num_attention_heads": 2, - "joint_attention_dim": 32, - "caption_projection_dim": 32, - "pos_embed_max_size": 64, - } - transformer_cls = AuraFlowTransformer2DModel - vae_kwargs = { - "sample_size": 32, - "in_channels": 3, - "out_channels": 3, - "block_out_channels": (4,), - "layers_per_block": 1, - "latent_channels": 4, - "norm_num_groups": 1, - "use_quant_conv": False, - "use_post_quant_conv": False, - "shift_factor": 0.0609, - "scaling_factor": 1.5035, - } - tokenizer_cls, tokenizer_id = AutoTokenizer, "hf-internal-testing/tiny-random-t5" - text_encoder_cls, text_encoder_id = UMT5EncoderModel, "hf-internal-testing/tiny-random-umt5" - text_encoder_target_modules = ["q", "k", "v", "o"] - denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0", "linear_1"] - - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 8, 8, 3) - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 10 - num_channels = 4 - sizes = (32, 32) - - generator = torch.manual_seed(0) - noise = floats_tensor((batch_size, num_channels) + sizes) - input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) - - pipeline_inputs = { - "prompt": "A painting of a squirrel eating a burger", - "num_inference_steps": 4, - "guidance_scale": 0.0, - "height": 8, - "width": 8, - "output_type": "np", - } - if with_generator: - pipeline_inputs.update({"generator": generator}) - - return noise, input_ids, pipeline_inputs - - @unittest.skip("Not supported in AuraFlow.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in AuraFlow.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass diff --git a/tests/lora/test_lora_layers_cogview4.py b/tests/lora/test_lora_layers_cogview4.py deleted file mode 100644 index dd1db0e64d92..000000000000 --- a/tests/lora/test_lora_layers_cogview4.py +++ /dev/null @@ -1,162 +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 sys -import tempfile -import unittest - -import numpy as np -import torch -from parameterized import parameterized -from transformers import AutoTokenizer, GlmModel - -from diffusers import AutoencoderKL, CogView4Pipeline, CogView4Transformer2DModel, FlowMatchEulerDiscreteScheduler - -from ..testing_utils import ( - floats_tensor, - require_peft_backend, - require_torch_accelerator, - skip_mps, - torch_device, -) - - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -class TokenizerWrapper: - @staticmethod - def from_pretrained(*args, **kwargs): - return AutoTokenizer.from_pretrained( - "hf-internal-testing/tiny-random-cogview4", subfolder="tokenizer", trust_remote_code=True - ) - - -@require_peft_backend -@skip_mps -class CogView4LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = CogView4Pipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - - transformer_kwargs = { - "patch_size": 2, - "in_channels": 4, - "num_layers": 2, - "attention_head_dim": 4, - "num_attention_heads": 4, - "out_channels": 4, - "text_embed_dim": 32, - "time_embed_dim": 8, - "condition_dim": 4, - } - transformer_cls = CogView4Transformer2DModel - vae_kwargs = { - "block_out_channels": [32, 64], - "in_channels": 3, - "out_channels": 3, - "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"], - "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"], - "latent_channels": 4, - "sample_size": 128, - } - vae_cls = AutoencoderKL - tokenizer_cls, tokenizer_id, tokenizer_subfolder = ( - TokenizerWrapper, - "hf-internal-testing/tiny-random-cogview4", - "tokenizer", - ) - text_encoder_cls, text_encoder_id, text_encoder_subfolder = ( - GlmModel, - "hf-internal-testing/tiny-random-cogview4", - "text_encoder", - ) - - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 32, 32, 3) - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 16 - num_channels = 4 - sizes = (4, 4) - - generator = torch.manual_seed(0) - noise = floats_tensor((batch_size, num_channels) + sizes) - input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) - - pipeline_inputs = { - "prompt": "", - "num_inference_steps": 1, - "guidance_scale": 6.0, - "height": 32, - "width": 32, - "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) - - def test_simple_inference_save_pretrained(self): - """ - Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained - """ - components, _, _ = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdirname: - pipe.save_pretrained(tmpdirname) - - pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname) - pipe_from_pretrained.to(torch_device) - - images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) - - @parameterized.expand([("block_level", True), ("leaf_level", False)]) - @require_torch_accelerator - def test_group_offloading_inference_denoiser(self, offload_type, use_stream): - # TODO: We don't run the (leaf_level, True) test here that is enabled for other models. - # The reason for this can be found here: https://github.com/huggingface/diffusers/pull/11804#issuecomment-3013325338 - super()._test_group_offloading_inference_denoiser(offload_type, use_stream) - - @unittest.skip("Not supported in CogView4.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in CogView4.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass diff --git a/tests/lora/test_lora_layers_helios.py b/tests/lora/test_lora_layers_helios.py deleted file mode 100644 index e88ca06a69ef..000000000000 --- a/tests/lora/test_lora_layers_helios.py +++ /dev/null @@ -1,116 +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 sys -import unittest - -import torch -from transformers import AutoTokenizer, T5EncoderModel - -from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, HeliosPipeline, HeliosTransformer3DModel - -from ..testing_utils import floats_tensor, require_peft_backend, skip_mps - - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -@skip_mps -class HeliosLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = HeliosPipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - - transformer_kwargs = { - "patch_size": (1, 2, 2), - "num_attention_heads": 2, - "attention_head_dim": 12, - "in_channels": 16, - "out_channels": 16, - "text_dim": 32, - "freq_dim": 256, - "ffn_dim": 32, - "num_layers": 2, - "cross_attn_norm": True, - "qk_norm": "rms_norm_across_heads", - "rope_dim": (4, 4, 4), - "has_multi_term_memory_patch": True, - "guidance_cross_attn": True, - "zero_history_timestep": True, - "is_amplify_history": False, - } - transformer_cls = HeliosTransformer3DModel - vae_kwargs = { - "base_dim": 3, - "z_dim": 16, - "dim_mult": [1, 1, 1, 1], - "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, 33, 32, 32, 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) - - 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) - - pipeline_inputs = { - "prompt": "", - "num_frames": num_frames, - "num_inference_steps": 1, - "guidance_scale": 6.0, - "height": 32, - "width": 32, - "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 Helios.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Helios.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass diff --git a/tests/pipelines/ace_step/test_ace_step.py b/tests/pipelines/ace_step/test_ace_step.py index e049840b159b..5ced68691272 100644 --- a/tests/pipelines/ace_step/test_ace_step.py +++ b/tests/pipelines/ace_step/test_ace_step.py @@ -15,8 +15,8 @@ import math -import unittest +import pytest import torch from transformers import AutoTokenizer, Qwen3Config, Qwen3Model @@ -30,13 +30,19 @@ ) from ...testing_utils import enable_full_determinism -from ..test_pipelines_common import PipelineTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) enable_full_determinism() -class AceStepConditionEncoderTests(unittest.TestCase): +class TestAceStepConditionEncoder: """Fast tests for the AceStepConditionEncoder.""" def get_tiny_config(self): @@ -90,55 +96,34 @@ def test_forward_shape(self): ) # Output should be packed: batch_size x (lyric + timbre + text seq_len) x hidden_size - self.assertEqual(enc_hidden.shape[0], batch_size) - self.assertEqual(enc_hidden.shape[2], config["hidden_size"]) - self.assertEqual(enc_mask.shape[0], batch_size) - self.assertEqual(enc_mask.shape[1], enc_hidden.shape[1]) + assert enc_hidden.shape[0] == batch_size + assert enc_hidden.shape[2] == config["hidden_size"] + assert enc_mask.shape[0] == batch_size + assert enc_mask.shape[1] == enc_hidden.shape[1] - def test_save_load_config(self): + def test_save_load_config(self, tmp_path): """Test that the condition encoder config can be saved and loaded.""" - import tempfile - config = self.get_tiny_config() encoder = AceStepConditionEncoder(**config) - with tempfile.TemporaryDirectory() as tmpdir: - encoder.save_config(tmpdir) - loaded = AceStepConditionEncoder.from_config(tmpdir) + encoder.save_config(tmp_path) + loaded = AceStepConditionEncoder.from_config(tmp_path) - self.assertEqual(encoder.config.hidden_size, loaded.config.hidden_size) - self.assertEqual(encoder.config.text_hidden_dim, loaded.config.text_hidden_dim) - self.assertEqual(encoder.config.timbre_hidden_dim, loaded.config.timbre_hidden_dim) + assert encoder.config.hidden_size == loaded.config.hidden_size + assert encoder.config.text_hidden_dim == loaded.config.text_hidden_dim + assert encoder.config.timbre_hidden_dim == loaded.config.timbre_hidden_dim -class AceStepPipelineFastTests(PipelineTesterMixin, unittest.TestCase): - """Fast end-to-end tests for AceStepPipeline with tiny models.""" - +class AceStepPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = AceStepPipeline - params = frozenset( - [ - "prompt", - "lyrics", - "audio_duration", - "vocal_language", - "guidance_scale", - "shift", - ] + required_input_params_in_call_signature = frozenset( + ["prompt", "lyrics", "audio_duration", "vocal_language", "guidance_scale", "shift"] ) - batch_params = frozenset(["prompt", "lyrics"]) - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "output_type", - "return_dict", - ] - ) - - # ACE-Step uses custom attention, not standard diffusers attention processors - test_attention_slicing = False - test_xformers_attention = False + batch_input_params = frozenset(["prompt", "lyrics"]) + # ACE-Step generates audio, so there is no `num_images_per_prompt`. + optional_input_params = frozenset(["num_inference_steps", "generator", "latents", "output_type", "return_dict"]) + # `(channels, samples)` for the short `audio_duration` used by the dummy inputs. + output_shape = (2, 7) def get_dummy_components(self): torch.manual_seed(0) @@ -231,7 +216,7 @@ def get_dummy_components(self): scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1, shift=1.0) - components = { + return { "transformer": transformer, "condition_encoder": condition_encoder, "vae": vae, @@ -241,113 +226,54 @@ def get_dummy_components(self): "audio_tokenizer": audio_tokenizer, "audio_token_detokenizer": audio_token_detokenizer, } - return components - - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { + + def get_dummy_inputs(self): + return { "prompt": "A beautiful piano piece", "lyrics": "[verse]\nSoft notes in the morning", - "audio_duration": 0.4, # Very short for fast test (10 latent frames at 25Hz) + # Short for a fast test, but long enough that the decoded waveform carries enough samples for the + # output comparisons the common tests make (the tiny VAE here runs at `latents_per_second == 2`). + "audio_duration": 2.0, "num_inference_steps": 2, - "generator": generator, + "generator": self.get_generator(0), "max_text_length": 32, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + "output_type": "pt", } - return inputs - def test_ace_step_basic(self): - """Test basic text-to-music generation.""" - device = "cpu" - components = self.get_dummy_components() - pipe = AceStepPipeline(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) - generator = torch.Generator(device=device).manual_seed(0) - output = pipe( - prompt="A beautiful piano piece", - lyrics="[verse]\nSoft notes in the morning", - audio_duration=0.4, - num_inference_steps=2, - generator=generator, - max_text_length=32, - ) - audio = output.audios - self.assertIsNotNone(audio) - self.assertEqual(audio.ndim, 3) # [batch, channels, samples] +class TestAceStepPipeline(AceStepPipelineTesterConfig, PipelineTesterMixin): + """Fast end-to-end tests for AceStepPipeline with tiny models.""" def test_ace_step_batch(self): """Test batch generation.""" - device = "cpu" - components = self.get_dummy_components() - pipe = AceStepPipeline(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() - generator = torch.Generator(device=device).manual_seed(42) - output = pipe( - prompt=["Piano piece", "Guitar solo"], - lyrics=["[verse]\nHello", "[chorus]\nWorld"], - audio_duration=0.4, - num_inference_steps=2, - generator=generator, - max_text_length=32, + audio = self.run_pipe( + pipe, prompt=["Piano piece", "Guitar solo"], lyrics=["[verse]\nHello", "[chorus]\nWorld"] ) - audio = output.audios - self.assertIsNotNone(audio) - self.assertEqual(audio.shape[0], 2) # batch size = 2 + assert audio.shape[0] == 2 # batch size = 2 def test_ace_step_latent_output(self): """Test that output_type='latent' returns latents.""" - device = "cpu" - components = self.get_dummy_components() - pipe = AceStepPipeline(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() - generator = torch.Generator(device=device).manual_seed(0) - output = pipe( - prompt="A test prompt", - lyrics="", - audio_duration=0.4, - num_inference_steps=2, - generator=generator, - output_type="latent", - max_text_length=32, - ) - latents = output.audios - self.assertIsNotNone(latents) + latents = self.run_pipe(pipe, lyrics="", output_type="latent") # Latent shape: [batch, latent_length, acoustic_dim] - self.assertEqual(latents.ndim, 3) - self.assertEqual(latents.shape[0], 1) + assert latents.ndim == 3 + assert latents.shape[0] == 1 def test_ace_step_return_dict_false(self): """Test that return_dict=False returns a tuple.""" - device = "cpu" - components = self.get_dummy_components() - pipe = AceStepPipeline(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() - generator = torch.Generator(device=device).manual_seed(0) - output = pipe( - prompt="A test prompt", - lyrics="", - audio_duration=0.4, - num_inference_steps=2, - generator=generator, - return_dict=False, - max_text_length=32, - ) - self.assertIsInstance(output, tuple) - self.assertEqual(len(output), 1) + inputs = self.get_dummy_inputs() + output = pipe(**inputs, return_dict=False) + assert isinstance(output, tuple) + assert len(output) == 1 def test_audio_codes_cover_path(self): - components = self.get_dummy_components() - pipe = AceStepPipeline(**components) + pipe = self.get_pipeline() output = pipe( prompt="A test prompt", @@ -358,15 +284,15 @@ def test_audio_codes_cover_path(self): max_text_length=32, ) - self.assertEqual(output.audios.shape[1], 4) + assert output.audios.shape[1] == 4 - def test_save_load_local(self, expected_max_difference=7e-3): + def test_save_load_local(self, tmp_path, base_pipe_output, expected_max_difference=7e-3): # increase tolerance to account for large composite model - super().test_save_load_local(expected_max_difference=expected_max_difference) + super().test_save_load_local(tmp_path, base_pipe_output, expected_max_difference=expected_max_difference) - def test_save_load_optional_components(self, expected_max_difference=7e-3): + def test_save_load_optional_components(self, tmp_path, expected_max_difference=7e-3): # increase tolerance to account for large composite model - super().test_save_load_optional_components(expected_max_difference=expected_max_difference) + super().test_save_load_optional_components(tmp_path, expected_max_difference=expected_max_difference) def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=7e-3): # increase tolerance for audio pipeline @@ -378,94 +304,61 @@ def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_d expected_slice=expected_slice, expected_max_difference=expected_max_difference ) - # ACE-Step does not use num_images_per_prompt - def test_num_images_per_prompt(self): - pass - - # ACE-Step does not use standard schedulers - @unittest.skip("ACE-Step uses built-in flow matching schedule, not diffusers schedulers") - def test_karras_schedulers_shape(self): - pass - - # ACE-Step does not support prompt_embeds directly - @unittest.skip("ACE-Step does not support prompt_embeds / negative_prompt_embeds") - def test_cfg(self): - pass - - def test_float16_inference(self, expected_max_diff=5e-2): - super().test_float16_inference(expected_max_diff=expected_max_diff) - - @unittest.skip( + @pytest.mark.skip( "ACE-Step __call__ does not accept prompt_embeds, so encode_prompt isolation test is not applicable" ) def test_encode_prompt_works_in_isolation(self): pass - @unittest.skip("Sequential CPU offloading produces NaN with tiny random models") - def test_sequential_cpu_offload_forward_pass(self): - pass - - @unittest.skip("Sequential CPU offloading produces NaN with tiny random models") - def test_sequential_offload_forward_pass_twice(self): - pass - def test_encode_prompt(self): """Test that encode_prompt returns correct shapes.""" - device = "cpu" - components = self.get_dummy_components() - pipe = AceStepPipeline(**components) - pipe = pipe.to(device) + pipe = self.get_pipeline() text_hidden, text_mask, lyric_hidden, lyric_mask = pipe.encode_prompt( prompt="A test prompt", lyrics="[verse]\nHello world", - device=device, + device="cpu", max_text_length=32, max_lyric_length=64, ) - self.assertEqual(text_hidden.ndim, 3) # [batch, seq_len, hidden_dim] - self.assertEqual(text_mask.ndim, 2) # [batch, seq_len] - self.assertEqual(lyric_hidden.ndim, 3) - self.assertEqual(lyric_mask.ndim, 2) - self.assertEqual(text_hidden.shape[0], 1) - self.assertEqual(lyric_hidden.shape[0], 1) + assert text_hidden.ndim == 3 # [batch, seq_len, hidden_dim] + assert text_mask.ndim == 2 # [batch, seq_len] + assert lyric_hidden.ndim == 3 + assert lyric_mask.ndim == 2 + assert text_hidden.shape[0] == 1 + assert lyric_hidden.shape[0] == 1 def test_prepare_latents(self): """Test that prepare_latents returns correct shapes.""" - device = "cpu" - components = self.get_dummy_components() - pipe = AceStepPipeline(**components) - pipe = pipe.to(device) + pipe = self.get_pipeline() latents = pipe.prepare_latents( batch_size=2, audio_duration=1.0, dtype=torch.float32, - device=device, + device="cpu", ) expected_length = math.ceil(1.0 * pipe.latents_per_second) - self.assertEqual(latents.shape, (2, expected_length, 8)) + assert latents.shape == (2, expected_length, 8) def test_timestep_schedule(self): """Test that the timestep schedule is generated correctly.""" - components = self.get_dummy_components() - pipe = AceStepPipeline(**components) + pipe = self.get_pipeline() # Test standard schedule schedule = pipe._get_timestep_schedule(num_inference_steps=8, shift=3.0) - self.assertEqual(len(schedule), 8) - self.assertAlmostEqual(schedule[0].item(), 1.0, places=5) + assert len(schedule) == 8 + assert schedule[0].item() == pytest.approx(1.0, abs=1e-5) # Test truncated schedule schedule = pipe._get_timestep_schedule(num_inference_steps=4, shift=3.0) - self.assertEqual(len(schedule), 4) + assert len(schedule) == 4 def test_format_prompt(self): """Test that prompt formatting works correctly.""" - components = self.get_dummy_components() - pipe = AceStepPipeline(**components) + pipe = self.get_pipeline() text, lyrics = pipe._format_prompt( prompt="A piano piece", @@ -474,12 +367,28 @@ def test_format_prompt(self): audio_duration=30.0, ) - self.assertIn("A piano piece", text) - self.assertIn("30 seconds", text) - self.assertIn("[verse]", lyrics) - self.assertIn("Hello", lyrics) - self.assertIn("en", lyrics) + assert "A piano piece" in text + assert "30 seconds" in text + assert "[verse]" in lyrics + assert "Hello" in lyrics + assert "en" in lyrics + + +class TestAceStepPipelineMemory(AceStepPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the ACE-Step pipeline.""" + + @pytest.mark.skip("Sequential CPU offloading produces NaN with tiny random models") + def test_sequential_cpu_offload_forward_pass(self): + pass + + @pytest.mark.skip("Sequential CPU offloading produces NaN with tiny random models") + def test_sequential_offload_forward_pass_twice(self): + pass + + +class TestAceStepPipelineLoRA(AceStepPipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the ACE-Step pipeline.""" -if __name__ == "__main__": - unittest.main() +class TestAceStepPipelineLoRAMemory(AceStepPipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the ACE-Step pipeline.""" diff --git a/tests/pipelines/aura_flow/test_pipeline_aura_flow.py b/tests/pipelines/aura_flow/test_pipeline_aura_flow.py index 1eb9d1035c33..a45f65ebb89b 100644 --- a/tests/pipelines/aura_flow/test_pipeline_aura_flow.py +++ b/tests/pipelines/aura_flow/test_pipeline_aura_flow.py @@ -1,34 +1,28 @@ -import unittest - -import numpy as np import torch from transformers import AutoTokenizer, UMT5EncoderModel from diffusers import AuraFlowPipeline, AuraFlowTransformer2DModel, AutoencoderKL, FlowMatchEulerDiscreteScheduler -from ..test_pipelines_common import ( +from ...testing_utils import assert_tensors_close +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, PipelineTesterMixin, check_qkv_fusion_matches_attn_procs_length, check_qkv_fusion_processors_exist, ) -class AuraFlowPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class AuraFlowPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = AuraFlowPipeline - params = frozenset( - [ - "prompt", - "height", - "width", - "guidance_scale", - "negative_prompt", - "prompt_embeds", - "negative_prompt_embeds", - ] + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "negative_prompt", "prompt_embeds", "negative_prompt_embeds"] ) - batch_params = frozenset(["prompt", "negative_prompt"]) - test_layerwise_casting = True - test_group_offloading = True + batch_input_params = frozenset(["prompt", "negative_prompt"]) + # `height` / `width` default to `transformer.config.sample_size * vae_scale_factor` (32 * 2). + output_shape = (3, 64, 64) def get_dummy_components(self): torch.manual_seed(0) @@ -70,38 +64,31 @@ def get_dummy_components(self): "vae": vae, } - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) - - inputs = { + def get_dummy_inputs(self): + return { "prompt": "A painting of a squirrel eating a burger", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, - "output_type": "np", "height": None, "width": None, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } - return inputs - def test_attention_slicing_forward_pass(self): - # Attention slicing needs to implemented differently for this because how single DiT and MMDiT - # blocks interfere with each other. - return + +class TestAuraFlowPipeline(AuraFlowPipelineTesterConfig, PipelineTesterMixin): + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-3): + # AuraFlow pads the prompt embeddings to a common length, so batched and single runs diverge slightly more. + super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) def test_fused_qkv_projections(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + # Run on CPU to keep the device-dependent `torch.Generator` deterministic. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - original_image_slice = image[0, -3:, -3:, -1] + image = self.run_pipe(pipe) + original_image_slice = image[0, -1, -3:, -3:] # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added # to the pipeline level. @@ -113,25 +100,50 @@ def test_fused_qkv_projections(self): pipe.transformer, pipe.transformer.original_attn_processors ), "Something wrong with the attention processors concerning the fused QKV projections." - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice_fused = image[0, -3:, -3:, -1] + image = self.run_pipe(pipe) + image_slice_fused = image[0, -1, -3:, -3:] pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice_disabled = image[0, -3:, -3:, -1] - - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), ( - "Fusion of QKV projections shouldn't affect the outputs." + image = self.run_pipe(pipe) + image_slice_disabled = image[0, -1, -3:, -3:] + + assert_tensors_close( + original_image_slice, + image_slice_fused, + atol=1e-3, + rtol=1e-3, + msg="Fusion of QKV projections shouldn't affect the outputs.", ) - assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." + assert_tensors_close( + image_slice_fused, + image_slice_disabled, + atol=1e-3, + rtol=1e-3, + msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", ) - assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Original outputs should match when fused QKV projections are disabled." + assert_tensors_close( + original_image_slice, + image_slice_disabled, + atol=1e-2, + rtol=1e-2, + msg="Original outputs should match when fused QKV projections are disabled.", ) - @unittest.skip("xformers attention processor does not exist for AuraFlow") - def test_xformers_attention_forwardGenerator_pass(self): - pass + +class TestAuraFlowPipelineMemory(AuraFlowPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the AuraFlow pipeline.""" + + +class TestAuraFlowPipelineLoRA(AuraFlowPipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the AuraFlow pipeline.""" + + # Adapting the attention projections alone barely moves the output of this tiny AuraFlow (max diff ~3e-5, below + # the tolerances the tests assert against), so the feed-forward `linear_1` layers are adapted as well. + denoiser_target_modules = {"transformer": ["to_q", "to_k", "to_v", "to_out.0", "linear_1"]} + + +class TestAuraFlowPipelineLoRAMemory(AuraFlowPipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the AuraFlow pipeline.""" + + # See `TestAuraFlowPipelineLoRA`. + denoiser_target_modules = {"transformer": ["to_q", "to_k", "to_v", "to_out.0", "linear_1"]} diff --git a/tests/pipelines/cogview4/test_cogview4.py b/tests/pipelines/cogview4/test_cogview4.py index 6a8521d27761..7eb047a30337 100644 --- a/tests/pipelines/cogview4/test_cogview4.py +++ b/tests/pipelines/cogview4/test_cogview4.py @@ -12,42 +12,32 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import unittest - -import numpy as np +import pytest import torch from transformers import AutoTokenizer, GlmConfig, GlmForCausalLM from diffusers import AutoencoderKL, CogView4Pipeline, CogView4Transformer2DModel, FlowMatchEulerDiscreteScheduler -from ...testing_utils import enable_full_determinism, torch_device -from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import PipelineTesterMixin, to_np +from ...testing_utils import enable_full_determinism, require_torch_accelerator +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) enable_full_determinism() -class CogView4PipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class CogView4PipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = CogView4Pipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "negative_prompt", "prompt_embeds", "negative_prompt_embeds"] ) - - test_xformers_attention = False - test_layerwise_casting = True + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (3, 16, 16) def get_dummy_components(self): torch.manual_seed(0) @@ -91,143 +81,48 @@ def get_dummy_components(self): # TODO(aryan): change this to THUDM/CogView4 once released tokenizer = AutoTokenizer.from_pretrained("THUDM/glm-4-9b-chat", trust_remote_code=True) - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, "text_encoder": text_encoder, "tokenizer": tokenizer, } - return components - - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { + + def get_dummy_inputs(self): + return { "prompt": "dance monkey", "negative_prompt": "bad", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "height": 16, "width": 16, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "output_type": "pt", } - return inputs - - def test_inference(self): - device = "cpu" - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs)[0] - generated_image = image[0] - - self.assertEqual(generated_image.shape, (3, 16, 16)) - expected_image = torch.randn(3, 16, 16) - max_diff = np.abs(generated_image - expected_image).max() - self.assertLessEqual(max_diff, 1e10) - - def test_callback_inputs(self): - sig = inspect.signature(self.pipeline_class.__call__) - has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters - has_callback_step_end = "callback_on_step_end" in sig.parameters - - if not (has_callback_tensor_inputs and has_callback_step_end): - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - self.assertTrue( - hasattr(pipe, "_callback_tensor_inputs"), - f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs", - ) - def callback_inputs_subset(pipe, i, t, callback_kwargs): - # iterate over callback args - for tensor_name, tensor_value in callback_kwargs.items(): - # check that we're only passing in allowed tensor inputs - assert tensor_name in pipe._callback_tensor_inputs - - return callback_kwargs - - def callback_inputs_all(pipe, i, t, callback_kwargs): - for tensor_name in pipe._callback_tensor_inputs: - assert tensor_name in callback_kwargs - - # iterate over callback args - for tensor_name, tensor_value in callback_kwargs.items(): - # check that we're only passing in allowed tensor inputs - assert tensor_name in pipe._callback_tensor_inputs - - return callback_kwargs - - inputs = self.get_dummy_inputs(torch_device) - - # Test passing in a subset - inputs["callback_on_step_end"] = callback_inputs_subset - inputs["callback_on_step_end_tensor_inputs"] = ["latents"] - output = pipe(**inputs)[0] - - # Test passing in a everything - inputs["callback_on_step_end"] = callback_inputs_all - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - output = pipe(**inputs)[0] - - def callback_inputs_change_tensor(pipe, i, t, callback_kwargs): - is_last = i == (pipe.num_timesteps - 1) - if is_last: - callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"]) - return callback_kwargs - - inputs["callback_on_step_end"] = callback_inputs_change_tensor - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - output = pipe(**inputs)[0] - assert output.abs().sum() < 1e10 - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-3) - - def test_attention_slicing_forward_pass( - self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 - ): - if not self.test_attention_slicing: - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - generator_device = "cpu" - inputs = self.get_dummy_inputs(generator_device) - output_without_slicing = pipe(**inputs)[0] - - pipe.enable_attention_slicing(slice_size=1) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing1 = pipe(**inputs)[0] - - pipe.enable_attention_slicing(slice_size=2) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing2 = pipe(**inputs)[0] - - if test_max_difference: - max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max() - max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max() - self.assertLess( - max(max_diff1, max_diff2), - expected_max_diff, - "Attention slicing should not affect the inference results", - ) + +class TestCogView4Pipeline(CogView4PipelineTesterConfig, PipelineTesterMixin): + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-3): + super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) + + +class TestCogView4PipelineMemory(CogView4PipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the CogView4 pipeline.""" + + +class TestCogView4PipelineLoRA(CogView4PipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the CogView4 pipeline.""" + + +class TestCogView4PipelineLoRAMemory(CogView4PipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the CogView4 pipeline.""" + + @pytest.mark.parametrize("offload_type,use_stream", [("block_level", True), ("leaf_level", False)]) + @require_torch_accelerator + def test_group_offloading_inference_denoiser(self, tmp_path, offload_type, use_stream): + # TODO: We don't run the (leaf_level, True) case that is enabled for other models. + # The reason for this can be found here: https://github.com/huggingface/diffusers/pull/11804#issuecomment-3013325338 + super().test_group_offloading_inference_denoiser(tmp_path, offload_type, use_stream) diff --git a/tests/pipelines/helios/test_helios.py b/tests/pipelines/helios/test_helios.py index 5de50b350923..0bbb166e88bc 100644 --- a/tests/pipelines/helios/test_helios.py +++ b/tests/pipelines/helios/test_helios.py @@ -13,44 +13,44 @@ # limitations under the License. import gc -import unittest +import pytest import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLWan, HeliosPipeline, HeliosScheduler, HeliosTransformer3DModel from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, enable_full_determinism, require_torch_accelerator, slow, torch_device, ) -from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import PipelineTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) enable_full_determinism() -class HeliosPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class HeliosPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = HeliosPipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "negative_prompt", "prompt_embeds", "negative_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (33, 3, 16, 16) + # Helios is a video pipeline: it exposes `num_videos_per_prompt`, not the base default `num_images_per_prompt`. + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"] ) - test_xformers_attention = False def get_dummy_components(self): torch.manual_seed(0) @@ -88,46 +88,39 @@ def get_dummy_components(self): is_amplify_history=False, ) - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, "text_encoder": text_encoder, "tokenizer": tokenizer, } - return components - - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { + + def get_dummy_inputs(self): + return { "prompt": "dance monkey", "negative_prompt": "negative", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 1.0, "height": 16, "width": 16, "num_frames": 9, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "output_type": "pt", } - return inputs - def test_inference(self): - device = "cpu" - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) +class TestHeliosPipeline(HeliosPipelineTesterConfig, PipelineTesterMixin): + def test_inference(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (33, 3, 16, 16)) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4529, 0.4527, 0.4499, 0.4542, 0.4528, 0.4524, 0.4531, 0.4534, 0.5328, @@ -136,36 +129,42 @@ def test_inference(self): generated_slice = generated_video.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3)) + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) - @unittest.skip("Helios uses a lot of mixed precision internally, which is not suitable for this test case") + @pytest.mark.skip("Helios uses a lot of mixed precision internally, which is not suitable for this test case") def test_save_load_float16(self): pass - @unittest.skip("Test not supported") - def test_attention_slicing_forward_pass(self): - pass - - @unittest.skip("Optional components not applicable for Helios") + @pytest.mark.skip("Optional components not applicable for Helios") def test_save_load_optional_components(self): pass +class TestHeliosPipelineMemory(HeliosPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Helios pipeline.""" + + +class TestHeliosPipelineLoRA(HeliosPipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the Helios pipeline.""" + + +class TestHeliosPipelineLoRAMemory(HeliosPipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the Helios pipeline.""" + + @slow @require_torch_accelerator -class HeliosPipelineIntegrationTests(unittest.TestCase): +class TestHeliosPipelineIntegration: prompt = "A painting of a squirrel eating a burger." - def setUp(self): - super().setUp() + @pytest.fixture(autouse=True) + def cleanup(self): gc.collect() backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() + yield gc.collect() backend_empty_cache(torch_device) - @unittest.skip("TODO: test needs to be implemented") + @pytest.mark.skip("TODO: test needs to be implemented") def test_helios(self): pass