diff --git a/tests/pipelines/bria/test_pipeline_bria.py b/tests/pipelines/bria/test_pipeline_bria.py index 40dccaf3b9f1..e4d326e260c2 100644 --- a/tests/pipelines/bria/test_pipeline_bria.py +++ b/tests/pipelines/bria/test_pipeline_bria.py @@ -138,7 +138,7 @@ def test_encode_prompt_works_in_isolation(self): pass def test_bria_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] @@ -151,7 +151,7 @@ def test_bria_different_prompts(self): assert max_diff > 1e-6 def test_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 57)] @@ -165,7 +165,7 @@ def test_image_output_shape(self): assert (output_height, output_width) == (expected_height, expected_width) def test_bria_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() height_width_pairs = [(16, 16), (32, 32), (64, 64)] diff --git a/tests/pipelines/bria_fibo/test_pipeline_bria_fibo.py b/tests/pipelines/bria_fibo/test_pipeline_bria_fibo.py index 0845106854ab..653293b7c100 100644 --- a/tests/pipelines/bria_fibo/test_pipeline_bria_fibo.py +++ b/tests/pipelines/bria_fibo/test_pipeline_bria_fibo.py @@ -27,7 +27,9 @@ from ...testing_utils import assert_tensors_close, torch_device from ..testing_utils import ( BasePipelineTesterConfig, - PipelineOffloadTesterMixin, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, PipelineTesterMixin, ) @@ -69,7 +71,17 @@ def get_dummy_components(self): scheduler = FlowMatchEulerDiscreteScheduler() torch.manual_seed(0) - text_encoder = SmolLM3ForCausalLM(SmolLM3Config(hidden_size=32)) + text_encoder = SmolLM3ForCausalLM( + SmolLM3Config( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + # `vocab_size` stays at the SmolLM3 default: the pipeline hardcodes the beginning-of-text id + # (128000) for empty prompts, so a smaller vocabulary would not be a valid text encoder here. + ) + ) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") return { @@ -105,7 +117,7 @@ def test_inference(self): assert generated_image.shape == self.output_shape # fmt: off - expected_slice = torch.tensor([0.4025, 0.4722, 0.4377, 0.6178, 0.3643, 0.4914, 0.3694, 0.5096, 0.5980, 0.5516, 0.5228, 0.4731, 0.6202, 0.2424, 0.6280, 0.3556]) + expected_slice = torch.tensor([0.3804, 0.4799, 0.5112, 0.5784, 0.3970, 0.4709, 0.4027, 0.4882, 0.5627, 0.4635, 0.4718, 0.3876, 0.5741, 0.2936, 0.5912, 0.4014]) # fmt: on generated_slice = generated_image.flatten() @@ -117,7 +129,7 @@ def test_encode_prompt_works_in_isolation(self): pass def test_bria_fibo_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] @@ -130,7 +142,7 @@ def test_bria_fibo_different_prompts(self): assert max_diff > 1e-6 def test_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (64, 64), (32, 64)] @@ -141,5 +153,22 @@ def test_image_output_shape(self): assert (output_height, output_width) == (height, width) -class TestBriaFiboPipelineMemory(BriaFiboPipelineTesterConfig, PipelineOffloadTesterMixin): - pass +class TestBriaFiboPipelineMemory(BriaFiboPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Bria FIBO pipeline.""" + + +class TestBriaFiboPipelineLoRA(BriaFiboPipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the Bria FIBO pipeline.""" + + @pytest.mark.skip( + "`_load_lora_into_text_encoder` only infers per-module ranks for CLIP-style names " + "(`.q_proj`/`.k_proj`/`.v_proj`/`.out_proj`/`.fc1`/`.fc2`, see `src/diffusers/loaders/lora_base.py`), so the " + "LLaMA-style `.o_proj` on the SmolLM3 text encoder falls back to the default rank and the non-uniform " + "`rank_pattern` this test builds cannot round-trip." + ) + def test_simple_inference_with_partial_text_lora(self): + pass + + +class TestBriaFiboPipelineLoRAMemory(BriaFiboPipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the Bria FIBO pipeline.""" diff --git a/tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py b/tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py index 69829204602f..2f79157a8991 100644 --- a/tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py +++ b/tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py @@ -29,7 +29,9 @@ from ...testing_utils import assert_tensors_close, torch_device from ..testing_utils import ( BasePipelineTesterConfig, - PipelineOffloadTesterMixin, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, PipelineTesterMixin, ) @@ -73,7 +75,17 @@ def get_dummy_components(self): z_dim=16, ) scheduler = FlowMatchEulerDiscreteScheduler() - text_encoder = SmolLM3ForCausalLM(SmolLM3Config(hidden_size=32)) + text_encoder = SmolLM3ForCausalLM( + SmolLM3Config( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + # `vocab_size` stays at the SmolLM3 default: the pipeline hardcodes the beginning-of-text id + # (128000) for empty prompts, so a smaller vocabulary would not be a valid text encoder here. + ) + ) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") return { @@ -111,7 +123,7 @@ def test_inference(self): assert generated_image.shape == self.output_shape # fmt: off - expected_slice = torch.tensor([0.5615, 0.4469, 0.4043, 0.4312, 0.3783, 0.4426, 0.4081, 0.4446, 0.6549, 0.6302, 0.6324, 0.5871, 0.6117, 0.6647, 0.5871, 0.6246]) + expected_slice = torch.tensor([0.5594, 0.4469, 0.4011, 0.4329, 0.3747, 0.4408, 0.4074, 0.4452, 0.6472, 0.6353, 0.6258, 0.5867, 0.6104, 0.6624, 0.5824, 0.6277]) # fmt: on generated_slice = generated_image.flatten() @@ -131,7 +143,7 @@ def test_inference_batch_single_identical(self): pass def test_bria_fibo_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] @@ -144,7 +156,7 @@ def test_bria_fibo_different_prompts(self): assert max_diff > 1e-6 def test_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (64, 64), (32, 64)] @@ -155,7 +167,7 @@ def test_image_output_shape(self): assert (output_height, output_width) == (height, width) def test_bria_fibo_multi_reference_uses_distinct_rope_time_planes(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) references = [ Image.new("RGB", (336, 192), (255, 255, 255)), @@ -180,7 +192,7 @@ def test_bria_fibo_multi_reference_uses_distinct_rope_time_planes(self): assert image.shape == self.output_shape def test_batched_prompts_with_multiple_references(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() inputs.update( prompt=[inputs["prompt"], inputs["prompt"].replace("squirrel", "robot")], @@ -192,7 +204,7 @@ def test_batched_prompts_with_multiple_references(self): assert (images[0] - images[1]).abs().max() > 1e-4 def test_multi_reference_mask_requires_single_reference(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() inputs["image"] = [inputs["image"], Image.new("RGB", (160, 96), (0, 0, 0))] inputs["mask"] = Image.new("L", (336, 192), 255) @@ -200,7 +212,7 @@ def test_multi_reference_mask_requires_single_reference(self): pipe(**inputs) def test_bria_fibo_edit_mask(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() mask = Image.fromarray((np.ones((192, 336)) * 255).astype(np.uint8), mode="L") @@ -211,7 +223,7 @@ def test_bria_fibo_edit_mask(self): assert output.shape == (3, 192, 336) def test_bria_fibo_edit_mask_image_size_mismatch(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() mask = Image.fromarray((np.ones((64, 64)) * 255).astype(np.uint8), mode="L") @@ -221,7 +233,7 @@ def test_bria_fibo_edit_mask_image_size_mismatch(self): pipe(**inputs) def test_bria_fibo_edit_mask_no_image(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() mask = Image.fromarray((np.ones((32, 32)) * 255).astype(np.uint8), mode="L") @@ -233,5 +245,22 @@ def test_bria_fibo_edit_mask_no_image(self): pipe(**inputs) -class TestBriaFiboEditPipelineMemory(BriaFiboEditPipelineTesterConfig, PipelineOffloadTesterMixin): - pass +class TestBriaFiboEditPipelineMemory(BriaFiboEditPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Bria FIBO Edit pipeline.""" + + +class TestBriaFiboEditPipelineLoRA(BriaFiboEditPipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the Bria FIBO Edit pipeline.""" + + @pytest.mark.skip( + "`_load_lora_into_text_encoder` only infers per-module ranks for CLIP-style names " + "(`.q_proj`/`.k_proj`/`.v_proj`/`.out_proj`/`.fc1`/`.fc2`, see `src/diffusers/loaders/lora_base.py`), so the " + "LLaMA-style `.o_proj` on the SmolLM3 text encoder falls back to the default rank and the non-uniform " + "`rank_pattern` this test builds cannot round-trip." + ) + def test_simple_inference_with_partial_text_lora(self): + pass + + +class TestBriaFiboEditPipelineLoRAMemory(BriaFiboEditPipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the Bria FIBO Edit pipeline.""" diff --git a/tests/pipelines/testing_utils/lora.py b/tests/pipelines/testing_utils/lora.py index d6cde97fcef7..a7ba8fcff78c 100644 --- a/tests/pipelines/testing_utils/lora.py +++ b/tests/pipelines/testing_utils/lora.py @@ -62,6 +62,7 @@ def check_module_lora_metadata(parsed_metadata: dict, lora_metadatas: dict, modu # Keyed by `config.model_type` — both `CLIPTextModel` and `CLIPTextModelWithProjection` report `clip_text_model`. TEXT_ENCODER_TARGET_MODULES = { "clip_text_model": ["q_proj", "k_proj", "v_proj", "out_proj"], + "smollm3": ["q_proj", "k_proj", "v_proj", "o_proj"], } @@ -447,9 +448,16 @@ def test_simple_inference_with_text_denoiser_lora_and_scale(self, base_pipe_outp msg="Lora + 0 scale should lead to same result as no LoRA", ) - if self.text_encoder_components: - text_encoder_root = getattr(pipe.text_encoder, "text_model", pipe.text_encoder) - assert text_encoder_root.encoder.layers[0].self_attn.q_proj.scaling["default"] == 1.0, ( + for name in self.text_encoder_components: + # Walk the modules rather than indexing a fixed path: text encoder architectures nest their attention + # layers differently (CLIP under `text_model.encoder.layers`, decoder-only ones under `model.layers`). + scalings = [ + module.scaling["default"] + for module in getattr(pipe, name).modules() + if hasattr(module, "lora_A") and "default" in module.scaling + ] + assert scalings, f"No LoRA layers found on {name}" + assert all(scaling == 1.0 for scaling in scalings), ( "The scaling parameter has not been correctly restored!" )