From efafaf644e23d9fd7e32170acac431d0b6cece42 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Thu, 27 Aug 2026 09:20:30 -0700 Subject: [PATCH] migrate ideogram4 lora tests to the pipeline-level mixins --- tests/lora/test_lora_layers_ideogram4.py | 188 ------------------ .../ideogram4/test_pipeline_ideogram4.py | 10 + 2 files changed, 10 insertions(+), 188 deletions(-) delete mode 100644 tests/lora/test_lora_layers_ideogram4.py diff --git a/tests/lora/test_lora_layers_ideogram4.py b/tests/lora/test_lora_layers_ideogram4.py deleted file mode 100644 index 3675a6effe69..000000000000 --- a/tests/lora/test_lora_layers_ideogram4.py +++ /dev/null @@ -1,188 +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 ( - AutoencoderKLFlux2, - FlowMatchEulerDiscreteScheduler, - Ideogram4Pipeline, - Ideogram4Transformer2DModel, -) -from diffusers.pipelines.ideogram4.pipeline_ideogram4 import QWEN3_VL_ACTIVATION_LAYERS - -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 - - -# The text conditioning concatenates the hidden states of these Qwen3-VL decoder layers, so the dummy text -# encoder must be deep enough to expose the last tapped layer, and `llm_features_dim` must match the product. -_TEXT_HIDDEN_SIZE = 8 -_NUM_TEXT_LAYERS = max(QWEN3_VL_ACTIVATION_LAYERS) + 1 -_LLM_FEATURES_DIM = len(QWEN3_VL_ACTIVATION_LAYERS) * _TEXT_HIDDEN_SIZE - - -@require_peft_backend -class Ideogram4LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = Ideogram4Pipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - - transformer_kwargs = { - "in_channels": 16, - "num_layers": 2, - "attention_head_dim": 8, - "num_attention_heads": 4, - "intermediate_size": 32, - "adaln_dim": 16, - "llm_features_dim": _LLM_FEATURES_DIM, - "rope_theta": 10_000, - "mrope_section": (2, 1, 1), - "norm_eps": 1e-5, - } - transformer_cls = Ideogram4Transformer2DModel - - vae_kwargs = { - "in_channels": 3, - "out_channels": 3, - "down_block_types": ("DownEncoderBlock2D",), - "up_block_types": ("UpDecoderBlock2D",), - "block_out_channels": (8,), - "layers_per_block": 1, - "latent_channels": 4, - "norm_num_groups": 1, - "sample_size": 32, - "patch_size": (2, 2), - "use_quant_conv": False, - "use_post_quant_conv": False, - } - vae_cls = AutoencoderKLFlux2 - - tokenizer_cls, tokenizer_id = Qwen2Tokenizer, "hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration" - - # Ideogram4'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 Ideogram4 loader. - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 16, 16, 3) - - def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): - # The Ideogram4 pipeline takes a second (unconditional) transformer and 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) - unconditional_transformer = self.transformer_cls(**self.transformer_kwargs) - - torch.manual_seed(0) - vae = self.vae_cls(**self.vae_kwargs) - - torch.manual_seed(0) - text_config = { - "hidden_size": _TEXT_HIDDEN_SIZE, - "num_hidden_layers": _NUM_TEXT_LAYERS, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "intermediate_size": 16, - "head_dim": 8, - "vocab_size": 151936, - "max_position_embeddings": 256, - "rope_theta": 10_000.0, - } - vision_config = { - "hidden_size": 8, - "depth": 2, - "num_heads": 2, - "intermediate_size": 16, - "out_hidden_size": _TEXT_HIDDEN_SIZE, - "patch_size": 14, - } - text_encoder = Qwen3VLModel(Qwen3VLConfig(text_config=text_config, vision_config=vision_config)) - tokenizer = self.tokenizer_cls.from_pretrained(self.tokenizer_id) - - scheduler = scheduler_cls(**self.scheduler_kwargs) - - 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, - "unconditional_transformer": unconditional_transformer, - } - - return pipeline_components, text_lora_config, denoiser_lora_config - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 32 - num_channels = 4 - sizes = (16, 16) - - 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_schedule": [1.0, 1.0], - "height": 16, - "width": 16, - "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 Ideogram4.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Ideogram4.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass diff --git a/tests/pipelines/ideogram4/test_pipeline_ideogram4.py b/tests/pipelines/ideogram4/test_pipeline_ideogram4.py index 29eb6e46447d..f125e8f5e97c 100644 --- a/tests/pipelines/ideogram4/test_pipeline_ideogram4.py +++ b/tests/pipelines/ideogram4/test_pipeline_ideogram4.py @@ -29,6 +29,8 @@ from ...testing_utils import assert_tensors_close, torch_device from ..testing_utils import ( BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, MemoryTesterMixin, PipelineTesterMixin, ) @@ -338,3 +340,11 @@ def test_sequential_offload_forward_pass_twice(self, expected_max_diff=2e-4): ) def test_group_offloading_inference(self): pass + + +class TestIdeogram4PipelineLoRA(Ideogram4PipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the Ideogram4 pipeline.""" + + +class TestIdeogram4PipelineLoRAMemory(Ideogram4PipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA offloading tests for the Ideogram4 pipeline."""