Skip to content

refactor bria pipeline tests to the new mixin structure #14547

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sayakpaul merged 6 commits into from
Aug 22, 2026
Merged

refactor bria pipeline tests to the new mixin structure #14547

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
3c71abe
refactor bria pipeline tests to the new mixin structure
akshan-main Aug 20, 2026
13712b5
Merge branch 'main' into refactor-bria-pipeline-tests
sayakpaul Aug 21, 2026
1fc8d98
add lora tester class
akshan-main Aug 21, 2026
e79f9f4
remove lora tester class
akshan-main Aug 21, 2026
ee7e6ef
Merge branch 'main' into refactor-bria-pipeline-tests
sayakpaul Aug 22, 2026
8729f90
Merge branch 'main' into refactor-bria-pipeline-tests
sayakpaul Aug 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
212 changes: 62 additions & 150 deletions tests/pipelines/bria/test_pipeline_bria.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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)

Expand Down Expand Up @@ -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}"
Loading