diff --git a/tests/lora/test_lora_layers_lumina2.py b/tests/lora/test_lora_layers_lumina2.py deleted file mode 100644 index abe6e941a5d5..000000000000 --- a/tests/lora/test_lora_layers_lumina2.py +++ /dev/null @@ -1,110 +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, GemmaForCausalLM - -from diffusers import ( - AutoencoderKL, - FlowMatchEulerDiscreteScheduler, - Lumina2Pipeline, - Lumina2Transformer2DModel, -) - -from ..testing_utils import floats_tensor, require_peft_backend - - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -class Lumina2LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = Lumina2Pipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - - transformer_kwargs = { - "sample_size": 4, - "patch_size": 2, - "in_channels": 4, - "hidden_size": 8, - "num_layers": 2, - "num_attention_heads": 1, - "num_kv_heads": 1, - "multiple_of": 16, - "ffn_dim_multiplier": None, - "norm_eps": 1e-5, - "scaling_factor": 1.0, - "axes_dim_rope": [4, 2, 2], - "cap_feat_dim": 8, - } - transformer_cls = Lumina2Transformer2DModel - 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, - } - vae_cls = AutoencoderKL - tokenizer_cls, tokenizer_id = AutoTokenizer, "hf-internal-testing/dummy-gemma" - text_encoder_cls, text_encoder_id = GemmaForCausalLM, "hf-internal-testing/dummy-gemma-diffusers" - - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 4, 4, 3) - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 16 - 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": 2, - "guidance_scale": 5.0, - "height": 32, - "width": 32, - "output_type": "np", - } - if with_generator: - pipeline_inputs.update({"generator": generator}) - - return noise, input_ids, pipeline_inputs - - @unittest.skip("Not supported in Lumina2.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Lumina2.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass diff --git a/tests/pipelines/lumina2/test_pipeline_lumina2.py b/tests/pipelines/lumina2/test_pipeline_lumina2.py index abb607635f49..c5f0d8012041 100644 --- a/tests/pipelines/lumina2/test_pipeline_lumina2.py +++ b/tests/pipelines/lumina2/test_pipeline_lumina2.py @@ -1,5 +1,3 @@ -import unittest - import torch from transformers import AutoTokenizer, Gemma2Config, Gemma2Model @@ -10,12 +8,19 @@ Lumina2Transformer2DModel, ) -from ..test_pipelines_common import PipelineTesterMixin +from ...testing_utils import assert_tensors_close +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) -class Lumina2PipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class Lumina2PipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = Lumina2Pipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( [ "prompt", "height", @@ -26,20 +31,9 @@ class Lumina2PipelineFastTests(unittest.TestCase, PipelineTesterMixin): "negative_prompt_embeds", ] ) - batch_params = frozenset(["prompt", "negative_prompt"]) - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] - ) - - test_xformers_attention = False - test_layerwise_casting = True + batch_input_params = frozenset(["prompt", "negative_prompt"]) + # The dummy one-block VAE decodes the 4x4 latents at scale 1, so requested 32x32 comes out 4x4 + output_shape = (3, 4, 4) def get_dummy_components(self): torch.manual_seed(0) @@ -89,28 +83,53 @@ def get_dummy_components(self): ) text_encoder = Gemma2Model(config) - components = { + return { "transformer": transformer, - "vae": vae.eval(), + "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="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, "height": 32, "width": 32, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + "output_type": "pt", } - return inputs + + +class TestLumina2Pipeline(Lumina2PipelineTesterConfig, PipelineTesterMixin): + def test_inference(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() + + inputs = self.get_dummy_inputs() + image = pipe(**inputs).images + generated_image = image[0] + assert generated_image.shape == self.output_shape + + # fmt: off + expected_slice = torch.tensor([0.4409, 0.6402, 0.1740, 0.4674, 0.4631, 0.3840, 0.5556, 0.4289, 0.4979, 0.4755, 0.5825, 0.6095, 0.7116, 0.5101, 0.6170, 0.6536]) + # fmt: on + + generated_slice = generated_image.flatten() + generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) + + +class TestLumina2PipelineMemory(Lumina2PipelineTesterConfig, MemoryTesterMixin): + pass + + +class TestLumina2PipelineLoRA(Lumina2PipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the Lumina2 pipeline.""" + + +class TestLumina2PipelineLoRAMemory(Lumina2PipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA offloading tests for the Lumina2 pipeline."""