diff --git a/tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py b/tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py index e69cc486b197..3a0aab41030c 100644 --- a/tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py +++ b/tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py @@ -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 @@ 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 @@ 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