AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-27 21:33:02 +03:00
Code Activity

refactor bria pipeline tests to the new mixin structure (#14547)

Browse Source 
* refactor bria pipeline tests to the new mixin structure

* add lora tester class

* remove lora tester class

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
This commit is contained in:
Akshan Krithick
2026-08-21 20:19:52 -07:00
committed by GitHub
parent 145dba5349
commit 398106a40f
1 changed files with 62 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files

212
tests/pipelines/bria/test_pipeline_bria.py
View File

@@ -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 import (
)
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 @@ class BriaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
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 @@ class BriaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
"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 @@ class BriaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
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 @@ class BriaPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
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 @@ class BriaPipelineSlowTests(unittest.TestCase):
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}"