Skip to content

refactor bria fibo edit pipeline tests to the new mixin structure #14577

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 1 commit into from
Aug 24, 2026
+61 −56
Merged

refactor bria fibo edit pipeline tests to the new mixin structure #14577

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
117 changes: 61 additions & 56 deletions tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest

import numpy as np
import pytest
import torch
from PIL import Image
from transformers import AutoTokenizer
Expand All @@ -26,24 +25,20 @@
FlowMatchEulerDiscreteScheduler,
)
from diffusers.models.transformers.transformer_bria_fibo import BriaFiboTransformer2DModel
from tests.pipelines.test_pipelines_common import PipelineTesterMixin

from ...testing_utils import (
enable_full_determinism,
torch_device,
from ...testing_utils import assert_tensors_close, torch_device
from ..testing_utils import (
BasePipelineTesterConfig,
PipelineOffloadTesterMixin,
PipelineTesterMixin,
)


enable_full_determinism()


class BriaFiboPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class BriaFiboEditPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = BriaFiboEditPipeline
params = frozenset(["prompt", "height", "width", "guidance_scale"])
batch_params = frozenset(["prompt"])
test_xformers_attention = False
test_layerwise_casting = False
test_group_offloading = False
required_input_params_in_call_signature = frozenset(["prompt", "height", "width", "guidance_scale"])
batch_input_params = frozenset(["prompt"])
output_shape = (3, 192, 336)

def get_dummy_components(self):
torch.manual_seed(0)
Expand Down Expand Up @@ -81,111 +76,121 @@ def get_dummy_components(self):
text_encoder = SmolLM3ForCausalLM(SmolLM3Config(hidden_size=32))
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")

components = {
return {
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
"transformer": transformer,
"vae": vae,
}
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)
def get_dummy_inputs(self):
inputs = {
"prompt": '{"text": "A painting of a squirrel eating a burger","edit_instruction": "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": 192,
"width": 336,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
image = Image.new("RGB", (336, 192), (255, 255, 255))
inputs["image"] = image
inputs["image"] = Image.new("RGB", (336, 192), (255, 255, 255))
return inputs

@unittest.skip(reason="will not be supported due to dim-fusion")

class TestBriaFiboEditPipeline(BriaFiboEditPipelineTesterConfig, 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.5611, 0.4472, 0.4037, 0.4324, 0.3775, 0.4414, 0.4093, 0.4446, 0.6525, 0.6329, 0.6308, 0.5895, 0.6117, 0.6657, 0.5850, 0.6254])
# 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("will not be supported due to dim-fusion")
def test_encode_prompt_works_in_isolation(self):
pass

@unittest.skip(reason="Batching is not supported yet")
@pytest.mark.skip("Batching is not supported yet")
def test_num_images_per_prompt(self):
pass

@unittest.skip(reason="Batching is not supported yet")
@pytest.mark.skip("Batching is not supported yet")
def test_inference_batch_consistent(self):
pass

@unittest.skip(reason="Batching is not supported yet")
@pytest.mark.skip("Batching is not supported yet")
def test_inference_batch_single_identical(self):
pass

def test_bria_fibo_different_prompts(self):
pipe = self.pipeline_class(**self.get_dummy_components())
pipe = pipe.to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
pipe = self.pipeline_class(**self.get_dummy_components()).to(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"] = {"edit_instruction": "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())
pipe = pipe.to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
inputs = self.get_dummy_inputs()

height_width_pairs = [(32, 32), (64, 64), (32, 64)]
for height, width in height_width_pairs:
expected_height = height
expected_width = width

inputs.update({"height": height, "width": width})
image = pipe(**inputs).images[0]
output_height, output_width, _ = image.shape
assert (output_height, output_width) == (expected_height, expected_width)
_, output_height, output_width = image.shape
assert (output_height, output_width) == (height, width)

def test_bria_fibo_edit_mask(self):
pipe = self.pipeline_class(**self.get_dummy_components())
pipe = pipe.to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
inputs = self.get_dummy_inputs()

mask = Image.fromarray((np.ones((192, 336)) * 255).astype(np.uint8), mode="L")

inputs.update({"mask": mask})
output = pipe(**inputs).images[0]

assert output.shape == (192, 336, 3)
assert output.shape == (3, 192, 336)

def test_bria_fibo_edit_mask_image_size_mismatch(self):
pipe = self.pipeline_class(**self.get_dummy_components())
pipe = pipe.to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
inputs = self.get_dummy_inputs()

mask = Image.fromarray((np.ones((64, 64)) * 255).astype(np.uint8), mode="L")

inputs.update({"mask": mask})
with self.assertRaisesRegex(ValueError, "Mask and image must have the same size"):
with pytest.raises(ValueError, match="Mask and image must have the same size"):
pipe(**inputs)

def test_bria_fibo_edit_mask_no_image(self):
pipe = self.pipeline_class(**self.get_dummy_components())
pipe = pipe.to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
inputs = self.get_dummy_inputs()

mask = Image.fromarray((np.ones((32, 32)) * 255).astype(np.uint8), mode="L")

# Remove image from inputs if it's there (it shouldn't be by default from get_dummy_inputs)
inputs.pop("image", None)
inputs.update({"mask": mask})

with self.assertRaisesRegex(ValueError, "If mask is provided, image must also be provided"):
with pytest.raises(ValueError, match="If mask is provided, image must also be provided"):
pipe(**inputs)


class TestBriaFiboEditPipelineMemory(BriaFiboEditPipelineTesterConfig, PipelineOffloadTesterMixin):
pass
Loading