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 fibo edit pipeline tests to the new mixin structure (#14577)

Browse Source
This commit is contained in:
Akshan Krithick
2026-08-23 20:08:40 -07:00
committed by GitHub
parent 58eb52c080
commit ad90df4847
1 changed files with 61 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

117
tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py
View File

@@ -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
@@ -26,24 +25,20 @@ from diffusers import (
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)
@@ -81,111 +76,121 @@ class BriaFiboPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
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