diff --git a/tests/lora/test_lora_layers_krea2.py b/tests/lora/test_lora_layers_krea2.py deleted file mode 100644 index eaf2a1c5450c..000000000000 --- a/tests/lora/test_lora_layers_krea2.py +++ /dev/null @@ -1,186 +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 unittest - -import torch -from transformers import Qwen2Tokenizer, Qwen3VLConfig, Qwen3VLModel - -from diffusers import ( - AutoencoderKLQwenImage, - FlowMatchEulerDiscreteScheduler, - Krea2Pipeline, - Krea2Transformer2DModel, -) - -from ..testing_utils import floats_tensor, is_peft_available, require_peft_backend - - -if is_peft_available(): - from peft import LoraConfig - - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -class Krea2LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = Krea2Pipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = { - "use_dynamic_shifting": True, - "base_shift": 0.5, - "max_shift": 1.15, - "base_image_seq_len": 256, - "max_image_seq_len": 6400, - } - - transformer_cls = Krea2Transformer2DModel - transformer_kwargs = { - "in_channels": 16, - "num_layers": 2, - "attention_head_dim": 8, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "intermediate_size": 32, - "timestep_embed_dim": 8, - "text_hidden_dim": 16, - "num_text_layers": 3, - "text_num_attention_heads": 2, - "text_num_key_value_heads": 1, - "text_intermediate_size": 16, - "num_layerwise_text_blocks": 1, - "num_refiner_text_blocks": 1, - "axes_dims_rope": (4, 2, 2), - "rope_theta": 1000.0, - } - - z_dim = 4 - vae_cls = AutoencoderKLQwenImage - vae_kwargs = { - "base_dim": z_dim * 6, - "z_dim": z_dim, - "dim_mult": [1, 2, 4], - "num_res_blocks": 1, - "temperal_downsample": [False, True], - "latents_mean": [0.0] * 4, - "latents_std": [1.0] * 4, - } - - tokenizer_cls, tokenizer_id = Qwen2Tokenizer, "hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration" - - # Krea2's attention uses split q/k/v/out projections in the diffusers transformer. - denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"] - # The text encoder (Qwen3-VL) is frozen and not LoRA-adapted by the Krea2 loader. - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 32, 32, 3) - - def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): - # The Krea2 pipeline uses a Qwen3-VL text encoder for which there is no tiny pretrained checkpoint, - # so build the components inline rather than relying on the base implementation. - scheduler_cls = self.scheduler_cls if scheduler_cls is None else 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) - - torch.manual_seed(0) - vae = self.vae_cls(**self.vae_kwargs) - - torch.manual_seed(0) - scheduler = scheduler_cls(**self.scheduler_kwargs) - - torch.manual_seed(0) - config = Qwen3VLConfig( - text_config={ - "hidden_size": 16, - "intermediate_size": 16, - "num_hidden_layers": 2, - "num_attention_heads": 2, - "num_key_value_heads": 2, - "head_dim": 8, - }, - vision_config={ - "depth": 2, - "hidden_size": 16, - "intermediate_size": 16, - "num_heads": 2, - "out_hidden_size": 16, - }, - vocab_size=152064, - ) - text_encoder = Qwen3VLModel(config).eval() - tokenizer = self.tokenizer_cls.from_pretrained(self.tokenizer_id) - - text_lora_config = LoraConfig( - r=rank, - lora_alpha=lora_alpha, - target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], - 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, - "text_encoder_select_layers": (0, 1, 2), - } - - return pipeline_components, text_lora_config, denoiser_lora_config - - 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 dog is dancing", - "num_inference_steps": 2, - "guidance_scale": 3.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 - - @unittest.skip("Not supported in Krea2.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Krea2.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass diff --git a/tests/lora/test_lora_layers_z_image.py b/tests/lora/test_lora_layers_z_image.py deleted file mode 100644 index 741e334d9603..000000000000 --- a/tests/lora/test_lora_layers_z_image.py +++ /dev/null @@ -1,181 +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 Qwen2Tokenizer, Qwen3Config, Qwen3Model - -from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, ZImagePipeline, ZImageTransformer2DModel - -from ..testing_utils import floats_tensor, is_peft_available, require_peft_backend - - -if is_peft_available(): - from peft import LoraConfig - - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -class ZImageLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = ZImagePipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - - transformer_kwargs = { - "all_patch_size": (2,), - "all_f_patch_size": (1,), - "in_channels": 16, - "dim": 32, - "n_layers": 2, - "n_refiner_layers": 1, - "n_heads": 2, - "n_kv_heads": 2, - "norm_eps": 1e-5, - "qk_norm": True, - "cap_feat_dim": 16, - "rope_theta": 256.0, - "t_scale": 1000.0, - "axes_dims": [8, 4, 4], - "axes_lens": [256, 32, 32], - } - transformer_cls = ZImageTransformer2DModel - vae_kwargs = { - "in_channels": 3, - "out_channels": 3, - "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"], - "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"], - "block_out_channels": [32, 64], - "layers_per_block": 1, - "latent_channels": 16, - "norm_num_groups": 32, - "sample_size": 32, - "scaling_factor": 0.3611, - "shift_factor": 0.1159, - } - vae_cls = AutoencoderKL - tokenizer_cls, tokenizer_id = Qwen2Tokenizer, "hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration" - text_encoder_cls, text_encoder_id = Qwen3Model, None # Will be created inline - denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"] - - 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 = 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": 32, - "width": 32, - "max_sequence_length": 16, - "output_type": "np", - } - if with_generator: - pipeline_inputs.update({"generator": generator}) - - return noise, input_ids, pipeline_inputs - - def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): - # Override to create Qwen3Model inline since it doesn't have a pretrained tiny model - torch.manual_seed(0) - config = Qwen3Config( - hidden_size=16, - intermediate_size=16, - num_hidden_layers=2, - num_attention_heads=2, - num_key_value_heads=2, - vocab_size=151936, - max_position_embeddings=512, - ) - text_encoder = Qwen3Model(config) - tokenizer = Qwen2Tokenizer.from_pretrained(self.tokenizer_id) - - transformer = self.transformer_cls(**self.transformer_kwargs) - # `x_pad_token` and `cap_pad_token` are initialized with `torch.empty`. - # This can cause NaN data values in our testing environment. Fixating them - # helps prevent that issue. - with torch.no_grad(): - transformer.x_pad_token.copy_(torch.ones_like(transformer.x_pad_token.data)) - transformer.cap_pad_token.copy_(torch.ones_like(transformer.cap_pad_token.data)) - vae = self.vae_cls(**self.vae_kwargs) - - if scheduler_cls is None: - scheduler_cls = self.scheduler_cls - scheduler = scheduler_cls(**self.scheduler_kwargs) - - rank = 4 - lora_alpha = rank if lora_alpha is None else lora_alpha - - text_lora_config = LoraConfig( - r=rank, - lora_alpha=lora_alpha, - target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], - 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 = { - "transformer": transformer, - "vae": vae, - "scheduler": scheduler, - "text_encoder": text_encoder, - "tokenizer": tokenizer, - } - - return pipeline_components, text_lora_config, denoiser_lora_config - - def test_lora_scale_kwargs_match_fusion(self): - super().test_lora_scale_kwargs_match_fusion(5e-2, 5e-2) - - @unittest.skip("Needs to be debugged.") - def test_set_adapters_match_attention_kwargs(self): - super().test_set_adapters_match_attention_kwargs() - - @unittest.skip("Needs to be debugged.") - def test_simple_inference_with_text_denoiser_lora_and_scale(self): - super().test_simple_inference_with_text_denoiser_lora_and_scale() - - @unittest.skip("Not supported in ZImage.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in ZImage.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass diff --git a/tests/pipelines/krea2/test_krea2.py b/tests/pipelines/krea2/test_krea2.py index 57b50a49f8d8..b21f8464262e 100644 --- a/tests/pipelines/krea2/test_krea2.py +++ b/tests/pipelines/krea2/test_krea2.py @@ -12,9 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - -import numpy as np import torch from transformers import Qwen2Tokenizer, Qwen3VLConfig, Qwen3VLModel @@ -25,33 +22,26 @@ Krea2Transformer2DModel, ) -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 assert_tensors_close, enable_full_determinism +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) enable_full_determinism() -class Krea2PipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class Krea2PipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = Krea2Pipeline - 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", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -119,7 +109,7 @@ def get_dummy_components(self): text_encoder = Qwen3VLModel(config).eval() tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -127,40 +117,31 @@ def get_dummy_components(self): "tokenizer": tokenizer, "text_encoder_select_layers": (0, 1, 2), } - 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 quality", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 3.0, "height": 32, "width": 32, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "output_type": "pt", } - return inputs +class TestKrea2Pipeline(Krea2PipelineTesterConfig, PipelineTesterMixin): def test_inference(self): - device = "cpu" + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - 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) + inputs = self.get_dummy_inputs() image = pipe(**inputs).images generated_image = image[0] - self.assertEqual(generated_image.shape, (3, 32, 32)) + assert generated_image.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.5649, 0.6510, 0.5885, 0.4954, 0.5551, 0.5973, 0.6043, 0.6009, 0.4307, 0.4733, 0.6145, 0.5121, 0.4431, 0.5144, 0.4427, 0.5011]) @@ -168,10 +149,10 @@ def test_inference(self): generated_slice = generated_image.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-3)) + assert_tensors_close(generated_slice, expected_slice, atol=5e-3) - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-1) + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-1): + super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) def test_components_function(self): # Same as the common test, but `text_encoder_select_layers` is a config value (a tuple), not a module, so it @@ -179,61 +160,38 @@ def test_components_function(self): init_components = self.get_dummy_components() init_components = {k: v for k, v in init_components.items() if not isinstance(v, (str, int, float, tuple))} - pipe = self.pipeline_class(**init_components) + pipe = self.get_pipeline(**init_components) - self.assertTrue(hasattr(pipe, "components")) - self.assertTrue(set(pipe.components.keys()) == set(init_components.keys())) + assert hasattr(pipe, "components") + assert set(pipe.components.keys()) == set(init_components.keys()) - def test_encode_prompt_works_in_isolation(self): + def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-4, rtol=1e-4): # Krea 2 enables classifier-free guidance whenever `guidance_scale > 0` and then encodes the (default empty) # negative prompt, which needs the tokenizer. The isolation pipeline carries no tokenizer, so run without # guidance; the common test already forwards only the positive `encode_prompt` outputs. original_get_dummy_inputs = self.get_dummy_inputs - def get_dummy_inputs_without_guidance(device, seed=0): - inputs = original_get_dummy_inputs(device, seed) + def get_dummy_inputs_without_guidance(): + inputs = original_get_dummy_inputs() inputs["guidance_scale"] = 0.0 return inputs self.get_dummy_inputs = get_dummy_inputs_without_guidance try: - super().test_encode_prompt_works_in_isolation() + super().test_encode_prompt_works_in_isolation( + extra_required_param_value_dict=extra_required_param_value_dict, atol=atol, rtol=rtol + ) finally: self.get_dummy_inputs = original_get_dummy_inputs - def test_attention_slicing_forward_pass( - self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 - ): - # Same as the qwenimage override: the common helper assumes channel-last outputs for the mean-pixel check, - # which does not hold for `output_type="pt"`; compare max difference only. - 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 TestKrea2PipelineMemory(Krea2PipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Krea 2 pipeline.""" + + +class TestKrea2PipelineLoRA(Krea2PipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the Krea 2 pipeline.""" + + +class TestKrea2PipelineLoRAMemory(Krea2PipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the Krea 2 pipeline.""" diff --git a/tests/pipelines/testing_utils/lora.py b/tests/pipelines/testing_utils/lora.py index 15af20a9abdb..34f796580dbc 100644 --- a/tests/pipelines/testing_utils/lora.py +++ b/tests/pipelines/testing_utils/lora.py @@ -752,7 +752,9 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): assert pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}" @pytest.mark.parametrize("lora_scale", [1.0, 0.8]) - def test_lora_scale_kwargs_match_fusion(self, base_pipe_output, lora_scale): + def test_lora_scale_kwargs_match_fusion( + self, base_pipe_output, lora_scale, expected_atol=1e-3, expected_rtol=1e-3 + ): attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) pipe = self.get_pipeline().to(torch_device) @@ -772,8 +774,8 @@ def test_lora_scale_kwargs_match_fusion(self, base_pipe_output, lora_scale): assert_tensors_close( outputs_lora_1_fused, outputs_lora_1, - atol=1e-3, - rtol=1e-3, + atol=expected_atol, + rtol=expected_rtol, msg="Fused lora should not change the output", ) assert not torch.allclose(base_pipe_output, outputs_lora_1, atol=1e-3, rtol=1e-3), ( diff --git a/tests/pipelines/z_image/test_z_image.py b/tests/pipelines/z_image/test_z_image.py index d3602bdeff76..470d52614c3e 100644 --- a/tests/pipelines/z_image/test_z_image.py +++ b/tests/pipelines/z_image/test_z_image.py @@ -12,19 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc import os -import unittest -import numpy as np +import pytest import torch from transformers import Qwen2Tokenizer, Qwen3Config, Qwen3Model from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, ZImagePipeline, ZImageTransformer2DModel -from ...testing_utils import 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 assert_tensors_close, torch_device +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) # Z-Image requires torch.use_deterministic_algorithms(False) due to complex64 RoPE operations @@ -37,49 +40,18 @@ if hasattr(torch.backends, "cuda"): torch.backends.cuda.matmul.allow_tf32 = False -# Note: Some tests (test_float16_inference, test_save_load_float16) may fail in full suite +# Note: Some tests (test_half_precision_inference_no_nan, test_save_load_float16) may fail in full suite # due to RopeEmbedder cache state pollution between tests. They pass when run individually. # This is a known test isolation issue, not a functional bug. -class ZImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class ZImagePipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = ZImagePipeline - 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 - test_group_offloading = True - - def setUp(self): - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.synchronize() - torch.manual_seed(0) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(0) - - def tearDown(self): - super().tearDown() - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.synchronize() - torch.manual_seed(0) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(0) + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -138,25 +110,19 @@ def get_dummy_components(self): text_encoder = Qwen3Model(config) tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") - 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 quality", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 3.0, "cfg_normalization": False, @@ -164,23 +130,20 @@ def get_dummy_inputs(self, device, seed=0): "height": 32, "width": 32, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "output_type": "pt", } - return inputs +class TestZImagePipeline(ZImagePipelineTesterConfig, PipelineTesterMixin): def test_inference(self): - device = "cpu" + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - 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) + inputs = self.get_dummy_inputs() image = pipe(**inputs).images generated_image = image[0] - self.assertEqual(generated_image.shape, (3, 32, 32)) + assert generated_image.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4622, 0.4532, 0.4714, 0.5087, 0.5371, 0.5405, 0.4492, 0.4479, 0.2984, 0.2783, 0.5409, 0.6577, 0.3952, 0.5524, 0.5262, 0.453]) @@ -188,119 +151,63 @@ def test_inference(self): generated_slice = generated_image.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-2)) - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-1) - - def test_num_images_per_prompt(self): - import inspect - - sig = inspect.signature(self.pipeline_class.__call__) - - if "num_images_per_prompt" not in sig.parameters: - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - batch_sizes = [1, 2] - num_images_per_prompts = [1, 2] - - for batch_size in batch_sizes: - for num_images_per_prompt in num_images_per_prompts: - inputs = self.get_dummy_inputs(torch_device) - - for key in inputs.keys(): - if key in self.batch_params: - inputs[key] = batch_size * [inputs[key]] - - images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt)[0] - - assert images.shape[0] == batch_size * num_images_per_prompt - - del pipe - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.synchronize() - - 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) + assert_tensors_close(generated_slice, expected_slice, atol=5e-2) - generator_device = "cpu" - inputs = self.get_dummy_inputs(generator_device) - output_without_slicing = pipe(**inputs)[0] + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-1): + # Z-Image pads the batch to a common sequence 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) - 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", - ) - - def test_vae_tiling(self, expected_diff_max: float = 0.2): - generator_device = "cpu" - components = self.get_dummy_components() - - pipe = self.pipeline_class(**components) - pipe.to("cpu") - pipe.set_progress_bar_config(disable=None) + def test_vae_tiling(self, expected_diff_max: float = 0.3): + pipe = self.get_pipeline().to(torch_device) # Without tiling - inputs = self.get_dummy_inputs(generator_device) + inputs = self.get_dummy_inputs() inputs["height"] = inputs["width"] = 128 output_without_tiling = pipe(**inputs)[0] # With tiling (standard AutoencoderKL doesn't accept parameters) pipe.vae.enable_tiling() - inputs = self.get_dummy_inputs(generator_device) + inputs = self.get_dummy_inputs() inputs["height"] = inputs["width"] = 128 output_with_tiling = pipe(**inputs)[0] - self.assertLess( - (to_np(output_without_tiling) - to_np(output_with_tiling)).max(), - expected_diff_max, - "VAE tiling should not affect the inference results", + assert (output_without_tiling - output_with_tiling).abs().max() < expected_diff_max, ( + "VAE tiling should not affect the inference results." ) - def test_pipeline_with_accelerator_device_map(self, expected_max_difference=5e-4): - # Z-Image RoPE embeddings (complex64) have slightly higher numerical tolerance - super().test_pipeline_with_accelerator_device_map(expected_max_difference=expected_max_difference) + +class TestZImagePipelineMemory(ZImagePipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Z-Image pipeline.""" def test_group_offloading_inference(self): # Block-level offloading conflicts with RoPE cache. Pipeline-level offloading (tested separately) works fine. - self.skipTest("Using test_pipeline_level_group_offloading_inference instead") + pytest.skip("Using test_pipeline_level_group_offloading_inference instead") - def test_save_load_float16(self, expected_max_diff=1e-2): - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.synchronize() - torch.manual_seed(0) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(0) - super().test_save_load_float16(expected_max_diff=expected_max_diff) + def test_pipeline_with_accelerator_device_map(self, tmp_path, base_pipe_output, expected_max_difference=5e-4): + # Z-Image RoPE embeddings (complex64) have slightly higher numerical tolerance + super().test_pipeline_with_accelerator_device_map( + tmp_path, base_pipe_output, expected_max_difference=expected_max_difference + ) + + +class TestZImagePipelineLoRA(ZImagePipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the Z-Image pipeline.""" + + @pytest.mark.parametrize("lora_scale", [1.0, 0.8]) + def test_lora_scale_kwargs_match_fusion(self, base_pipe_output, lora_scale): + # Fusing a scaled LoRA drifts a bit more than the default tolerance on Z-Image. + super().test_lora_scale_kwargs_match_fusion( + base_pipe_output, lora_scale, expected_atol=5e-2, expected_rtol=5e-2 + ) + + @pytest.mark.skip("Needs to be debugged.") + def test_set_adapters_match_attention_kwargs(self, tmp_path, base_pipe_output): + pass + + @pytest.mark.skip("Needs to be debugged.") + def test_simple_inference_with_text_denoiser_lora_and_scale(self, base_pipe_output): + pass + + +class TestZImagePipelineLoRAMemory(ZImagePipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the Z-Image pipeline."""