From 3c71abee73530b56eff2a82f98cb9d9db4cc92f0 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Thu, 20 Aug 2026 16:26:53 -0700 Subject: [PATCH 1/3] refactor bria pipeline tests to the new mixin structure --- tests/pipelines/bria/test_pipeline_bria.py | 212 ++++++--------------- 1 file changed, 62 insertions(+), 150 deletions(-) diff --git a/tests/pipelines/bria/test_pipeline_bria.py b/tests/pipelines/bria/test_pipeline_bria.py index 96a5f557d5fe..40dccaf3b9f1 100644 --- a/tests/pipelines/bria/test_pipeline_bria.py +++ b/tests/pipelines/bria/test_pipeline_bria.py @@ -13,10 +13,9 @@ # limitations under the License. import gc -import tempfile -import unittest import numpy as np +import pytest import torch from huggingface_hub import hf_hub_download from transformers import AutoConfig, T5EncoderModel, T5TokenizerFast @@ -28,32 +27,28 @@ ) from diffusers.pipelines.bria import BriaPipeline -# from ..test_pipelines_common import PipelineTesterMixin, check_qkv_fused_layers_exist -from tests.pipelines.test_pipelines_common import PipelineTesterMixin, to_np - from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, - enable_full_determinism, numpy_cosine_similarity_distance, require_torch_accelerator, slow, torch_device, ) +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, +) -enable_full_determinism() - - -class BriaPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class BriaPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = BriaPipeline - params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds"]) - batch_params = frozenset(["prompt"]) - test_xformers_attention = False - - # there is no xformers processor for Flux - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 16, 16) def get_dummy_components(self): torch.manual_seed(0) @@ -93,7 +88,7 @@ def get_dummy_components(self): text_encoder = T5EncoderModel(config) tokenizer = T5TokenizerFast.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "scheduler": scheduler, "text_encoder": text_encoder, "tokenizer": tokenizer, @@ -102,43 +97,62 @@ def get_dummy_components(self): "image_encoder": None, "feature_extractor": None, } - 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", "negative_prompt": "bad, ugly", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "height": 16, "width": 16, "max_sequence_length": 48, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + "output_type": "pt", } - return inputs + +class TestBriaPipeline(BriaPipelineTesterConfig, 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.4820, 0.3644, 0.4211, 0.5917, 0.7707, 0.4361, 0.4391, 0.6226, 0.3278, 0.6498, 0.3452, 0.4684, 0.4824, 0.5849, 0.3553, 0.3937]) + # 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) + + @pytest.mark.skip( + "BriaPipeline.__init__ dereferences the vae config, so it cannot be built with text-only components" + ) 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) - inputs = self.get_dummy_inputs(torch_device) + + inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] - inputs = self.get_dummy_inputs(torch_device) + + inputs = self.get_dummy_inputs() inputs["prompt"] = "a different prompt" output_different_prompts = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + + max_diff = (output_same_prompt - output_different_prompts).abs().max() assert max_diff > 1e-6 def test_image_output_shape(self): pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 57)] for height, width in height_width_pairs: @@ -147,55 +161,12 @@ def test_image_output_shape(self): inputs.update({"height": height, "width": width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape + _, output_height, output_width = image.shape assert (output_height, output_width) == (expected_height, expected_width) - @unittest.skipIf(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU") - @require_torch_accelerator - def test_save_load_float16(self, expected_max_diff=1e-2): - components = self.get_dummy_components() - for name, module in components.items(): - if hasattr(module, "half"): - components[name] = module.to(torch_device).half() - - 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) - - inputs = self.get_dummy_inputs(torch_device) - output = pipe(**inputs)[0] - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir) - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, dtype=torch.float16) - for component in pipe_loaded.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) - - for name, component in pipe_loaded.components.items(): - if name == "vae": - continue - if hasattr(component, "dtype"): - self.assertTrue( - component.dtype == torch.float16, - f"`{name}.dtype` switched from `float16` to {component.dtype} after loading.", - ) - - inputs = self.get_dummy_inputs(torch_device) - output_loaded = pipe_loaded(**inputs)[0] - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess( - max_diff, expected_max_diff, "The output of the fp16 pipeline changed after saving and loading." - ) - def test_bria_image_output_shape(self): pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(16, 16), (32, 32), (64, 64)] for height, width in height_width_pairs: @@ -204,53 +175,25 @@ def test_bria_image_output_shape(self): inputs.update({"height": height, "width": width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape + _, output_height, output_width = image.shape assert (output_height, output_width) == (expected_height, expected_width) - def test_to_dtype(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - - model_dtypes = [component.dtype for component in components.values() if hasattr(component, "dtype")] - self.assertTrue([dtype == torch.float32 for dtype in model_dtypes] == [True, True, True]) - def test_dtype_dict(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - with tempfile.TemporaryDirectory() as tmpdirname: - pipe.save_pretrained(tmpdirname) - dtype_dict = {"transformer": torch.bfloat16, "default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained(tmpdirname, dtype=dtype_dict) - - self.assertEqual(loaded_pipe.transformer.dtype, torch.bfloat16) - self.assertEqual(loaded_pipe.text_encoder.dtype, torch.float16) - self.assertEqual(loaded_pipe.vae.dtype, torch.float16) - - with tempfile.TemporaryDirectory() as tmpdirname: - pipe.save_pretrained(tmpdirname) - dtype_dict = {"default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained(tmpdirname, dtype=dtype_dict) - - self.assertEqual(loaded_pipe.transformer.dtype, torch.float16) - self.assertEqual(loaded_pipe.text_encoder.dtype, torch.float16) - self.assertEqual(loaded_pipe.vae.dtype, torch.float16) +class TestBriaPipelineMemory(BriaPipelineTesterConfig, MemoryTesterMixin): + pass @slow @require_torch_accelerator -class BriaPipelineSlowTests(unittest.TestCase): +class TestBriaPipelineSlow: pipeline_class = BriaPipeline repo_id = "briaai/BRIA-3.2" - 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) @@ -281,40 +224,9 @@ def test_bria_inference_bf16(self): image = pipe(**inputs).images[0] image_slice = image[0, :10, :10].flatten() - expected_slice = np.array( - [ - 0.59729785, - 0.6153719, - 0.595112, - 0.5884763, - 0.59366125, - 0.5795311, - 0.58325, - 0.58449626, - 0.57737637, - 0.58432233, - 0.5867875, - 0.57824117, - 0.5819089, - 0.5830988, - 0.57730293, - 0.57647324, - 0.5769151, - 0.57312685, - 0.57926565, - 0.5823928, - 0.57783926, - 0.57162863, - 0.575649, - 0.5745547, - 0.5740556, - 0.5799735, - 0.57799566, - 0.5715559, - 0.5771242, - 0.5773058, - ], - dtype=np.float32, - ) + # fmt: off + expected_slice = np.array([0.59729785, 0.6153719, 0.595112, 0.5884763, 0.59366125, 0.5795311, 0.58325, 0.58449626, 0.57737637, 0.58432233, 0.5867875, 0.57824117, 0.5819089, 0.5830988, 0.57730293, 0.57647324, 0.5769151, 0.57312685, 0.57926565, 0.5823928, 0.57783926, 0.57162863, 0.575649, 0.5745547, 0.5740556, 0.5799735, 0.57799566, 0.5715559, 0.5771242, 0.5773058], dtype=np.float32) + # fmt: on + max_diff = numpy_cosine_similarity_distance(expected_slice, image_slice) - self.assertLess(max_diff, 1e-4, f"Image slice is different from expected slice: {max_diff:.4f}") + assert max_diff < 1e-4, f"Image slice is different from expected slice: {max_diff:.4f}" From 1fc8d98b1766e1c8e99cb87993928d0c5ea45775 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Thu, 20 Aug 2026 20:05:03 -0700 Subject: [PATCH 2/3] add lora tester class --- tests/pipelines/bria/test_pipeline_bria.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/pipelines/bria/test_pipeline_bria.py b/tests/pipelines/bria/test_pipeline_bria.py index 40dccaf3b9f1..debdc24baa49 100644 --- a/tests/pipelines/bria/test_pipeline_bria.py +++ b/tests/pipelines/bria/test_pipeline_bria.py @@ -37,6 +37,7 @@ ) from ..testing_utils import ( BasePipelineTesterConfig, + LoraTesterMixin, MemoryTesterMixin, PipelineTesterMixin, ) @@ -183,6 +184,10 @@ class TestBriaPipelineMemory(BriaPipelineTesterConfig, MemoryTesterMixin): pass +class TestBriaPipelineLoRA(BriaPipelineTesterConfig, LoraTesterMixin): + pass + + @slow @require_torch_accelerator class TestBriaPipelineSlow: From e79f9f44be11f8f808c3a193e3b801ddb1989cb7 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Fri, 21 Aug 2026 15:21:24 -0700 Subject: [PATCH 3/3] remove lora tester class --- tests/pipelines/bria/test_pipeline_bria.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/pipelines/bria/test_pipeline_bria.py b/tests/pipelines/bria/test_pipeline_bria.py index debdc24baa49..40dccaf3b9f1 100644 --- a/tests/pipelines/bria/test_pipeline_bria.py +++ b/tests/pipelines/bria/test_pipeline_bria.py @@ -37,7 +37,6 @@ ) from ..testing_utils import ( BasePipelineTesterConfig, - LoraTesterMixin, MemoryTesterMixin, PipelineTesterMixin, ) @@ -184,10 +183,6 @@ class TestBriaPipelineMemory(BriaPipelineTesterConfig, MemoryTesterMixin): pass -class TestBriaPipelineLoRA(BriaPipelineTesterConfig, LoraTesterMixin): - pass - - @slow @require_torch_accelerator class TestBriaPipelineSlow: