diff --git a/docs/source/en/api/pipelines/bria_fibo_edit.md b/docs/source/en/api/pipelines/bria_fibo_edit.md index b46dd78cdb90..c9672fa46706 100644 --- a/docs/source/en/api/pipelines/bria_fibo_edit.md +++ b/docs/source/en/api/pipelines/bria_fibo_edit.md @@ -16,8 +16,11 @@ Fibo Edit is an 8B parameter image-to-image model that introduces a new paradigm Featuring native masking for granular precision, it moves beyond simple prompt-based diffusion to offer explicit, interpretable control optimized for production environments. Its lightweight architecture is designed for deep customization, empowering researchers to build specialized "Edit" models for domain-specific tasks while delivering top-tier aesthetic quality +Refer to the Bria Fibo Edit Hugging Face [page](https://huggingface.co/briaai/Fibo-Edit-1.5-base) to learn more. A distilled checkpoint is available at [Fibo-Edit-1.5-turbo](https://huggingface.co/briaai/Fibo-Edit-1.5-turbo). + ## Usage -_As the model is gated, before using it with diffusers you first need to go to the [Bria Fibo Hugging Face page](https://huggingface.co/briaai/Fibo-Edit), fill in the form and accept the gate. Once you are in, you need to login so that your system knows you’ve accepted the gate._ + +_As the model is gated, before using it with diffusers you first need to go to the [Bria Fibo Edit Hugging Face page](https://huggingface.co/briaai/Fibo-Edit-1.5-base), fill in the form and accept the gate. Once you are in, you need to login so that your system knows you’ve accepted the gate._ Use the command below to log in: @@ -25,9 +28,8 @@ Use the command below to log in: hf auth login ``` - ## BriaFiboEditPipeline [[autodoc]] BriaFiboEditPipeline - all - - __call__ \ No newline at end of file + - __call__ diff --git a/src/diffusers/models/transformers/transformer_bria_fibo.py b/src/diffusers/models/transformers/transformer_bria_fibo.py index a02f59461a1e..9ec0ea1647a6 100644 --- a/src/diffusers/models/transformers/transformer_bria_fibo.py +++ b/src/diffusers/models/transformers/transformer_bria_fibo.py @@ -441,6 +441,7 @@ class BriaFiboTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, From """ _supports_gradient_checkpointing = True + _repeated_blocks = ["BriaFiboTransformerBlock", "BriaFiboSingleTransformerBlock"] @register_to_config def __init__( diff --git a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py index 2aef63215f08..398758294bcc 100644 --- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py @@ -432,16 +432,6 @@ def prepare_latents( return latents, latent_image_ids - @staticmethod - def _prepare_attention_mask(attention_mask): - attention_matrix = torch.einsum("bi,bj->bij", attention_mask, attention_mask) - - # convert to 0 - keep, -inf ignore - attention_matrix = torch.where( - attention_matrix == 1, 0.0, -torch.inf - ) # Apply -inf to ignored tokens for nulling softmax score - return attention_matrix - @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( @@ -615,13 +605,15 @@ def __call__( if guidance_scale > 1: latent_attention_mask = latent_attention_mask.repeat(2, 1) - attention_mask = torch.cat([prompt_attention_mask, latent_attention_mask], dim=1) - attention_mask = self._prepare_attention_mask(attention_mask) # batch, seq => batch, seq, seq - attention_mask = attention_mask.unsqueeze(dim=1).to(dtype=self.transformer.dtype) # for head broadcasting + attention_mask = torch.cat([prompt_attention_mask, latent_attention_mask], dim=1).bool() if self._joint_attention_kwargs is None: self._joint_attention_kwargs = {} - self._joint_attention_kwargs["attention_mask"] = attention_mask + if not attention_mask.all(): + # Bool key-padding mask (batch, 1, 1, seq): every real query attends to the same keys as + # with a full (seq, seq) matrix, and varlen backends require bool. When nothing is padded + # the mask is a no-op, and omitting it keeps backends without mask support usable. + self._joint_attention_kwargs["attention_mask"] = attention_mask[:, None, None, :] # Adapt scheduler to dynamic shifting (resolution dependent) @@ -630,7 +622,7 @@ def __call__( else: seq_len = (height // self.vae_scale_factor) * (width // self.vae_scale_factor) - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + sigmas = None if timesteps is not None else np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) mu = calculate_shift( seq_len, @@ -646,7 +638,7 @@ def __call__( self.scheduler, num_inference_steps=num_inference_steps, device=device, - timesteps=None, + timesteps=timesteps, sigmas=sigmas, mu=mu, ) diff --git a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py index 664ee7d090d2..28858322fb96 100644 --- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py @@ -50,38 +50,47 @@ torch.FloatTensor, Image.Image, List[Image.Image], List[torch.FloatTensor], np.ndarray, List[np.ndarray] ] -# TODO: Update example docstring EXAMPLE_DOC_STRING = """ Example: ```python import torch + from PIL import Image + from diffusers import BriaFiboEditPipeline - from diffusers.modular_pipelines import ModularPipeline + from diffusers.modular_pipelines import ModularPipelineBlocks - torch.set_grad_enabled(False) - vlm_pipe = ModularPipelineBlocks.from_pretrained("briaai/FIBO-VLM-prompt-to-JSON", trust_remote_code=True) + # This prompt-to-JSON block calls Gemini and needs GEMINI_API_KEY in the environment. + vlm_pipe = ModularPipelineBlocks.from_pretrained("briaai/FIBO-edit-gemini-prompt-to-JSON", trust_remote_code=True) vlm_pipe = vlm_pipe.init_pipeline() pipe = BriaFiboEditPipeline.from_pretrained( - "briaai/fibo-edit", + "briaai/Fibo-Edit-1.5-base", torch_dtype=torch.bfloat16, ) pipe.to("cuda") - output = vlm_pipe( - prompt="A hyper-detailed, ultra-fluffy owl sitting in the trees at night, looking directly at the camera with wide, adorable, expressive eyes. Its feathers are soft and voluminous, catching the cool moonlight with subtle silver highlights. The owl's gaze is curious and full of charm, giving it a whimsical, storybook-like personality." + image = Image.open("owl.png") + json_prompt = vlm_pipe(image=image, prompt="Make the owl into a cat").values["json_prompt"] + + result = pipe(prompt=json_prompt, image=image, num_inference_steps=30, guidance_scale=5) + + # Multiple reference images: pass a list. Each reference conditions the edit at its + # own aspect ratio; the output resolution follows the first reference. + owl, forest = Image.open("owl.png"), Image.open("forest.png") + json_prompt = vlm_pipe( + image=[owl, forest], prompt="Place the owl from the first image in the forest from the second image" + ).values["json_prompt"] + result = pipe( + prompt=json_prompt, + image=[owl, forest], + num_inference_steps=30, + guidance_scale=5, ) - json_prompt_generate = json.loads(output.values["json_prompt"]) - - image = Image.open("image_generate.png") - - edit_prompt = "Make the owl to be a cat" - json_prompt_generate["edit_instruction"] = edit_prompt - - results_generate = pipe( - prompt=json_prompt_generate, num_inference_steps=50, guidance_scale=3.5, image=image, output_type="np" - ) + # The distilled Turbo checkpoint edits in 4 steps without classifier-free guidance. + pipe = BriaFiboEditPipeline.from_pretrained("briaai/Fibo-Edit-1.5-turbo", torch_dtype=torch.bfloat16) + pipe.to("cuda") + result = pipe(prompt=json_prompt, image=[owl, forest], num_inference_steps=4, guidance_scale=1) ``` """ @@ -173,18 +182,15 @@ def get_mask_size(mask: PipelineMaskInput): return None -def get_image_size(image: PipelineImageInput): - """ - Get the size of the image. +def _vae_safe_dims(width, height, base_resolution=1024, multiple=16): + """Return VAE-safe ``(width, height)`` for a reference image. + + Dimensions are rounded to a multiple of 16 (required by the VAE) and large references are capped at + ``base_resolution`` square pixels, so each reference keeps its own aspect ratio without producing an unbounded + context sequence. """ - if isinstance(image, torch.Tensor): - return image.shape[-2:] - elif isinstance(image, Image.Image): - return image.size[::-1] # (height, width) - elif isinstance(image, list): - return [get_image_size(i) for i in image] - else: - return None + scale = min(1.0, ((base_resolution * base_resolution) / float(width * height)) ** 0.5) + return tuple(max(multiple, int(round(side * scale / multiple)) * multiple) for side in (width, height)) def paste_mask_on_image(mask: PipelineMaskInput, image: PipelineImageInput): @@ -590,22 +596,12 @@ def prepare_latents( return latents, latent_image_ids - @staticmethod - def _prepare_attention_mask(attention_mask): - attention_matrix = torch.einsum("bi,bj->bij", attention_mask, attention_mask) - - # convert to 0 - keep, -inf ignore - attention_matrix = torch.where( - attention_matrix == 1, 0.0, -torch.inf - ) # Apply -inf to ignored tokens for nulling softmax score - return attention_matrix - @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, prompt: Union[str, List[str]] = None, - image: Optional[PipelineImageInput] = None, + image: Optional[Union[Image.Image, List[Image.Image]]] = None, mask: Optional[PipelineMaskInput] = None, height: int | None = None, width: int | None = None, @@ -632,9 +628,10 @@ def __call__( Args: prompt (`str` or `List[str]`): The prompt or prompts to guide the image generation. - image (`PIL.Image.Image` or `torch.FloatTensor`, *optional*): - The image to guide the image generation. If not defined, the pipeline will generate an image from - scratch. + image (`PIL.Image.Image` or `List[PIL.Image.Image]`, *optional*): + One or more reference images to guide the image generation. A list is interpreted as multiple + references (not a batch): each reference is VAE-encoded at its own aspect ratio and placed on its own + RoPE time plane 1, 2, ... . If not defined, the pipeline generates an image from scratch. mask (`PipelineMaskInput`, *optional*): Optional mask defining the region of `image` to be edited. Pixels covered by the mask are regenerated while the rest of the image is preserved. @@ -642,7 +639,7 @@ def __call__( The height in pixels of the generated image. This is set to 1024 by default for the best results. width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): The width in pixels of the generated image. This is set to 1024 by default for the best results. - num_inference_steps (`int`, *optional*, defaults to 50): + num_inference_steps (`int`, *optional*, defaults to 30): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. seed (`int`, *optional*): @@ -691,7 +688,8 @@ def __call__( max_sequence_length (`int` defaults to 3000): Maximum sequence length to use with the `prompt`. do_patching (`bool`, *optional*, defaults to `False`): Whether to use patching. _auto_resize (`bool`, *optional*, defaults to `True`): - Whether to automatically resize the input image to the preferred resolutions. + Whether to snap the default output resolution (taken from the first reference image) to the preferred + resolutions. Examples: Returns: [`~pipelines.flux.BriaFiboPipelineOutput`] or `tuple`: [`~pipelines.flux.BriaFiboPipelineOutput`] if @@ -699,9 +697,18 @@ def __call__( generated images. """ + # A list is interpreted as multiple references, not as a batch; Fibo Edit has no batched-edit API. + references = [] if image is None else image if isinstance(image, list) else [image] + if image is not None and ( + not references or not all(isinstance(reference, Image.Image) for reference in references) + ): + raise ValueError("`image` must be a `PIL.Image.Image` or a non-empty list of them.") + if height is None or width is None: - if image is not None: - image_height, image_width = self.image_processor.get_default_height_width(image) + if references: + # Output resolution follows the first reference image; the other + # references only condition generation at their own size. + image_width, image_height = _vae_safe_dims(*references[0].size) if _auto_resize: image_width, image_height = min( PREFERRED_RESOLUTION[1024 * 1024], @@ -714,7 +721,7 @@ def __call__( # 1. Check inputs. Raise error if not correct self.check_inputs( seed=seed, - image=image, + references=references, mask=mask, prompt=prompt, height=height, @@ -723,8 +730,8 @@ def __call__( max_sequence_length=max_sequence_length, ) - if mask is not None and image is not None: - image = paste_mask_on_image(mask, image) + if mask is not None: + references[0] = paste_mask_on_image(mask, references[0]) self._guidance_scale = guidance_scale self._joint_attention_kwargs = joint_attention_kwargs @@ -732,8 +739,10 @@ def __call__( # 2. Define call parameters - if prompt is not None and is_valid_edit_json(prompt): + if isinstance(prompt, dict): prompt = json.dumps(prompt) + elif isinstance(prompt, list): + prompt = [json.dumps(p) if isinstance(p, dict) else p for p in prompt] if isinstance(prompt, str): batch_size = 1 else: @@ -742,6 +751,7 @@ def __call__( device = self._execution_device if generator is None and seed is not None: generator = torch.Generator(device=device).manual_seed(seed) + lora_scale = ( self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None ) @@ -782,12 +792,8 @@ def __call__( # duplicate last layer prompt_layers = prompt_layers + [prompt_layers[-1]] * (total_num_layers_transformer - len(prompt_layers)) - # Preprocess image - if image is not None: - image = self.image_processor.resize(image, height, width) - image = self.image_processor.preprocess(image, height, width) - - # 5. Prepare latent variables + # 5. Prepare generated and reference latent variables. Every reference is + # VAE-encoded at its own size and receives RoPE time id 1, 2, ... . num_channels_latents = self.transformer.config.in_channels if do_patching: num_channels_latents = int(num_channels_latents / 4) @@ -804,18 +810,21 @@ def __call__( do_patching, ) - if image is not None: - image_latents, image_ids = self.prepare_image_latents( - image=image, - batch_size=batch_size * num_images_per_prompt, + reference_latents, reference_ids = [], [] + for reference_index, reference in enumerate(references, start=1): + packed, ids = self.prepare_reference_latents( + image=reference, num_channels_latents=num_channels_latents, - height=height, - width=width, dtype=prompt_embeds.dtype, device=device, - generator=generator, + do_patching=do_patching, + reference_index=reference_index, ) - latent_image_ids = torch.cat([latent_image_ids, image_ids], dim=0) # dim 0 is sequence dimension + reference_latents.append(packed) + reference_ids.append(ids) + if reference_latents: + image_latents = torch.cat(reference_latents, dim=1).repeat(prompt_batch_size, 1, 1) + latent_image_ids = torch.cat([latent_image_ids, *reference_ids], dim=0) else: image_latents = None @@ -839,12 +848,14 @@ def __call__( [prompt_attention_mask, latent_attention_mask, image_latent_attention_mask], dim=1 ) - attention_mask = self.create_attention_matrix(attention_mask) # batch, seq => batch, seq, seq - attention_mask = attention_mask.unsqueeze(dim=1).to(dtype=self.transformer.dtype) # for head broadcasting - if self._joint_attention_kwargs is None: self._joint_attention_kwargs = {} - self._joint_attention_kwargs["attention_mask"] = attention_mask + attention_mask = attention_mask.bool() + if not attention_mask.all(): + # Bool key-padding mask (batch, 1, 1, seq): every real query attends to the same keys as + # with a full (seq, seq) matrix, and varlen backends require bool. When nothing is padded + # the mask is a no-op, and omitting it keeps backends without mask support usable. + self._joint_attention_kwargs["attention_mask"] = attention_mask[:, None, None, :] # Adapt scheduler to dynamic shifting (resolution dependent) @@ -853,7 +864,7 @@ def __call__( else: seq_len = (height // self.vae_scale_factor) * (width // self.vae_scale_factor) - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + sigmas = None if timesteps is not None else np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) mu = calculate_shift( seq_len, @@ -869,7 +880,7 @@ def __call__( self.scheduler, num_inference_steps=num_inference_steps, device=device, - timesteps=None, + timesteps=timesteps, sigmas=sigmas, mu=mu, ) @@ -978,56 +989,47 @@ def __call__( return BriaFiboPipelineOutput(images=image) - def prepare_image_latents( + def prepare_reference_latents( self, - image: torch.Tensor, - batch_size: int, + image: Image.Image, num_channels_latents: int, - height: int, - width: int, dtype: torch.dtype, device: torch.device, - generator: torch.Generator | list[torch.Generator] | None = None, + do_patching: bool = False, + reference_index: int = 1, ): - image = image.to(device=device, dtype=dtype) - - height = int(height) // self.vae_scale_factor - width = int(width) // self.vae_scale_factor - - # scaling - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1).to(device, dtype) - ) - latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( - device, dtype - ) - - image_latents_cthw = self.vae.encode(image.unsqueeze(2)).latent_dist.mean - latents_scaled = [(latent - latents_mean) * latents_std for latent in image_latents_cthw] - image_latents_cthw = torch.concat(latents_scaled, dim=0) - image_latents_bchw = image_latents_cthw[:, :, 0, :, :] - - image_latent_height, image_latent_width = image_latents_bchw.shape[2:] - image_latents_bsd = self._pack_latents_no_patch( - latents=image_latents_bchw, - batch_size=batch_size, - num_channels_latents=num_channels_latents, - height=image_latent_height, - width=image_latent_width, - ) - # breakpoint() - image_ids = self._prepare_latent_image_ids( - batch_size=batch_size, height=image_latent_height, width=image_latent_width, device=device, dtype=dtype - ) - # image ids are the same as latent ids with the first dimension set to 1 instead of 0 - image_ids[..., 0] = 1 - return image_latents_bsd, image_ids + """VAE-encode one PIL reference at its own size and pack it as an edit-context token stream.""" + vae_dtype = next(self.vae.parameters()).dtype + reference_width, reference_height = _vae_safe_dims(*image.size) + encoded_input = self.image_processor.preprocess( + image.convert("RGB"), height=reference_height, width=reference_width + ).to(device=device) + + latent = self.vae.encode(encoded_input.to(dtype=vae_dtype).unsqueeze(2)).latent_dist.mean[:, :, 0, :, :] + latents_mean = torch.tensor(self.vae.config.latents_mean, device=device, dtype=latent.dtype).view(1, -1, 1, 1) + latents_std = torch.tensor(self.vae.config.latents_std, device=device, dtype=latent.dtype).view(1, -1, 1, 1) + latent = (latent - latents_mean) / latents_std + + latent_height, latent_width = latent.shape[-2:] + if do_patching: + if latent_height % 2 or latent_width % 2: + raise ValueError("Patched reference latents require even height and width.") + packed = self._pack_latents(latent, 1, num_channels_latents, latent_height, latent_width) + id_height, id_width = latent_height // 2, latent_width // 2 + else: + packed = self._pack_latents_no_patch(latent, 1, num_channels_latents, latent_height, latent_width) + id_height, id_width = latent_height, latent_width + image_ids = self._prepare_latent_image_ids(1, id_height, id_width, device, dtype) + # Fibo Edit conditions on references by placing each one on a distinct RoPE + # time plane. Generated image tokens remain on plane zero. + image_ids[..., 0] = reference_index + return packed.to(dtype=dtype), image_ids def check_inputs( self, prompt, seed, - image, + references, mask, height, width, @@ -1036,15 +1038,15 @@ def check_inputs( ): if seed is not None and not isinstance(seed, int): raise ValueError("Seed must be an integer") - if image is not None and not isinstance(image, (torch.Tensor, Image.Image, list)): - raise ValueError("Image must be a valid image") - if image is None and mask is not None: + if not references and mask is not None: raise ValueError("If mask is provided, image must also be provided") + if mask is not None and len(references) != 1: + raise ValueError("Masks are supported only with exactly one reference image.") if mask is not None and not is_valid_mask(mask): raise ValueError("Mask must be a valid mask") - if mask is not None and image is not None and not (get_mask_size(mask) == get_image_size(image)): + if mask is not None and tuple(get_mask_size(mask)) != references[0].size[::-1]: raise ValueError("Mask and image must have the same size") if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: @@ -1061,17 +1063,9 @@ def check_inputs( if prompt is None: raise ValueError("`prompt` must be provided.") - elif not is_valid_edit_json(prompt): - raise ValueError(f"`prompt` has to be a valid JSON string or dict but is {type(prompt)}") + prompts = prompt if isinstance(prompt, list) else [prompt] + if not prompts or not all(is_valid_edit_json(p) for p in prompts): + raise ValueError(f"`prompt` has to be a valid edit JSON string/dict or a list of them but is {prompt!r}") if max_sequence_length is not None and max_sequence_length > 3000: raise ValueError(f"`max_sequence_length` cannot be greater than 3000 but is {max_sequence_length}") - - def create_attention_matrix(self, attention_mask): - attention_matrix = torch.einsum("bi,bj->bij", attention_mask, attention_mask) - - # convert to 0 - keep, -inf ignore - attention_matrix = torch.where( - attention_matrix == 1, 0.0, -torch.inf - ) # Apply -inf to ignored tokens for nulling softmax score - return attention_matrix 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 3a0aab41030c..69829204602f 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 @@ -111,7 +111,7 @@ def test_inference(self): 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]) + expected_slice = torch.tensor([0.5615, 0.4469, 0.4043, 0.4312, 0.3783, 0.4426, 0.4081, 0.4446, 0.6549, 0.6302, 0.6324, 0.5871, 0.6117, 0.6647, 0.5871, 0.6246]) # fmt: on generated_slice = generated_image.flatten() @@ -122,10 +122,6 @@ def test_inference(self): def test_encode_prompt_works_in_isolation(self): pass - @pytest.mark.skip("Batching is not supported yet") - def test_num_images_per_prompt(self): - pass - @pytest.mark.skip("Batching is not supported yet") def test_inference_batch_consistent(self): pass @@ -158,6 +154,51 @@ def test_image_output_shape(self): _, output_height, output_width = image.shape assert (output_height, output_width) == (height, width) + def test_bria_fibo_multi_reference_uses_distinct_rope_time_planes(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + + references = [ + Image.new("RGB", (336, 192), (255, 255, 255)), + Image.new("RGB", (160, 96), (0, 0, 0)), + ] + num_channels_latents = pipe.transformer.config.in_channels + for reference_index, reference in enumerate(references, start=1): + packed, ids = pipe.prepare_reference_latents( + image=reference, + num_channels_latents=num_channels_latents, + dtype=torch.float32, + device=torch_device, + reference_index=reference_index, + ) + expected_tokens = (reference.height // 16) * (reference.width // 16) + assert packed.shape[:2] == (1, expected_tokens) + assert (ids[:, 0] == reference_index).all() + + inputs = self.get_dummy_inputs() + inputs.update(image=references, num_inference_steps=1) + image = pipe(**inputs).images[0] + assert image.shape == self.output_shape + + def test_batched_prompts_with_multiple_references(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs() + inputs.update( + prompt=[inputs["prompt"], inputs["prompt"].replace("squirrel", "robot")], + image=[inputs["image"], Image.new("RGB", (160, 96), (0, 0, 0))], + num_inference_steps=2, + ) + images = pipe(**inputs).images + assert images.shape == (2, *self.output_shape) + assert (images[0] - images[1]).abs().max() > 1e-4 + + def test_multi_reference_mask_requires_single_reference(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs() + inputs["image"] = [inputs["image"], Image.new("RGB", (160, 96), (0, 0, 0))] + inputs["mask"] = Image.new("L", (336, 192), 255) + with pytest.raises(ValueError, match="exactly one reference"): + pipe(**inputs) + def test_bria_fibo_edit_mask(self): pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) inputs = self.get_dummy_inputs()