From 479af16888d4f6678efca8e8ada149415a019c30 Mon Sep 17 00:00:00 2001 From: Shivam Shrirao Date: Sat, 22 Aug 2026 19:57:26 +0000 Subject: [PATCH 1/6] Fibo Edit: multi-reference conditioning, batching, timesteps --- .../source/en/api/pipelines/bria_fibo_edit.md | 24 ++ .../transformers/transformer_bria_fibo.py | 1 + .../pipelines/bria_fibo/pipeline_bria_fibo.py | 27 +- .../bria_fibo/pipeline_bria_fibo_edit.py | 230 +++++++++--------- .../test_pipeline_bria_fibo_edit.py | 49 +++- 5 files changed, 201 insertions(+), 130 deletions(-) diff --git a/docs/source/en/api/pipelines/bria_fibo_edit.md b/docs/source/en/api/pipelines/bria_fibo_edit.md index b46dd78cdb90..45166771e426 100644 --- a/docs/source/en/api/pipelines/bria_fibo_edit.md +++ b/docs/source/en/api/pipelines/bria_fibo_edit.md @@ -25,6 +25,30 @@ Use the command below to log in: hf auth login ``` +## Multiple reference images + +`image` accepts a `PIL.Image.Image` or a list of them. A list is interpreted as multiple references, not a batch: each reference is VAE-encoded at its own aspect ratio (capped at 1024x1024 square pixels) and conditions the edit on its own RoPE time plane, while the output resolution follows the first reference. Masks are supported only with a single reference. + +```python +import json + +import torch +from PIL import Image + +from diffusers import BriaFiboEditPipeline + +pipe = BriaFiboEditPipeline.from_pretrained("briaai/Fibo-Edit", torch_dtype=torch.bfloat16) +pipe.to("cuda") + +prompt = {"edit_instruction": "Place the product from the first image in the scene from the second image"} +result = pipe( + prompt=json.dumps(prompt), + image=[Image.open("product.png"), Image.open("scene.png")], + num_inference_steps=50, + guidance_scale=3.5, +) +result.images[0].save("edit.png") +``` ## BriaFiboEditPipeline 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..edd47f34ad67 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,18 @@ 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 attention_mask.all(): + # Nothing is padded, so the mask is a no-op; skipping it keeps backends without + # mask support (e.g. flash-attn 2/3) usable. + self._joint_attention_kwargs.pop("attention_mask", None) + else: + # 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. + self._joint_attention_kwargs["attention_mask"] = attention_mask[:, None, None, :] # Adapt scheduler to dynamic shifting (resolution dependent) @@ -630,7 +625,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 +641,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..dc3d48829e3a 100644 --- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py @@ -50,13 +50,16 @@ 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 json + 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) @@ -71,16 +74,22 @@ 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." ) - json_prompt_generate = json.loads(output.values["json_prompt"]) + json_prompt = json.loads(output.values["json_prompt"]) image = Image.open("image_generate.png") - edit_prompt = "Make the owl to be a cat" + json_prompt["edit_instruction"] = "Make the owl to be a cat" - json_prompt_generate["edit_instruction"] = edit_prompt + result = pipe(prompt=json_prompt, num_inference_steps=50, guidance_scale=3.5, image=image, output_type="np") - results_generate = pipe( - prompt=json_prompt_generate, num_inference_steps=50, guidance_scale=3.5, image=image, output_type="np" + # Multiple reference images: pass a list. Each reference conditions the edit at its + # own aspect ratio; the output resolution follows the first reference. + json_prompt["edit_instruction"] = "Place the owl from the first image in the forest from the second image" + result = pipe( + prompt=json_prompt, + image=[Image.open("owl.png"), Image.open("forest.png")], + num_inference_steps=50, + guidance_scale=3.5, ) ``` """ @@ -173,18 +182,34 @@ def get_mask_size(mask: PipelineMaskInput): return None -def get_image_size(image: PipelineImageInput): +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. """ - Get the size of the image. + 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 _vae_safe_size(image, base_resolution=1024, multiple=16): + """Resize a PIL reference to VAE-safe dimensions.""" + target = _vae_safe_dims(*image.size, base_resolution, multiple) + return image if target == image.size else image.resize(target, Image.LANCZOS) + + +def _as_reference_images(image): + """Normalize the edit input to a non-empty list of PIL reference images. + + A list is interpreted as multiple references, not as a batch; Fibo Edit has no batched-edit API. """ - 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 + if image is None: + return [] + references = image if isinstance(image, list) else [image] + if 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.") + return references def paste_mask_on_image(mask: PipelineMaskInput, image: PipelineImageInput): @@ -590,22 +615,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 +647,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. @@ -691,7 +707,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 +716,12 @@ def __call__( generated images. """ + references = _as_reference_images(image) 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], @@ -723,8 +743,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 +752,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 +764,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 +805,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 +823,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 +861,17 @@ 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 attention_mask.all(): + # Nothing is padded, so the mask is a no-op; skipping it keeps backends without + # mask support (e.g. flash-attn 2/3) usable. + self._joint_attention_kwargs.pop("attention_mask", None) + else: + # 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. + self._joint_attention_kwargs["attention_mask"] = attention_mask[:, None, None, :] # Adapt scheduler to dynamic shifting (resolution dependent) @@ -853,7 +880,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 +896,7 @@ def __call__( self.scheduler, num_inference_steps=num_inference_steps, device=device, - timesteps=None, + timesteps=timesteps, sigmas=sigmas, mu=mu, ) @@ -978,50 +1005,40 @@ 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 + image = _vae_safe_size(image.convert("RGB")) + pixels = torch.from_numpy(np.array(image)).permute(2, 0, 1).unsqueeze(0) + encoded_input = pixels.to(device=device, dtype=torch.float32) / 255.0 * 2.0 - 1.0 + + 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, @@ -1036,15 +1053,16 @@ 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") + references = _as_reference_images(image) if image is None 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 +1079,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 e69cc486b197..71af0705260e 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 @@ -113,10 +113,6 @@ def get_dummy_inputs(self, device, seed=0): def test_encode_prompt_works_in_isolation(self): pass - @unittest.skip(reason="Batching is not supported yet") - def test_num_images_per_prompt(self): - pass - @unittest.skip(reason="Batching is not supported yet") def test_inference_batch_consistent(self): pass @@ -153,6 +149,51 @@ def test_image_output_shape(self): output_height, output_width, _ = image.shape assert (output_height, output_width) == (expected_height, expected_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) + self.assertEqual(packed.shape[:2], (1, expected_tokens)) + self.assertTrue((ids[:, 0] == reference_index).all()) + + inputs = self.get_dummy_inputs(torch_device) + inputs.update(image=references, num_inference_steps=1) + image = pipe(**inputs).images[0] + self.assertEqual(image.shape, (192, 336, 3)) + + def test_batched_prompts_with_multiple_references(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs(torch_device) + 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 + self.assertEqual(images.shape, (2, 192, 336, 3)) + self.assertGreater(np.abs(images[0] - images[1]).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(torch_device) + inputs["image"] = [inputs["image"], Image.new("RGB", (160, 96), (0, 0, 0))] + inputs["mask"] = Image.new("L", (336, 192), 255) + with self.assertRaisesRegex(ValueError, "exactly one reference"): + pipe(**inputs) + def test_bria_fibo_edit_mask(self): pipe = self.pipeline_class(**self.get_dummy_components()) pipe = pipe.to(torch_device) From 91796cd9cb3728140d61948b67a3a5190c42eb69 Mon Sep 17 00:00:00 2001 From: Shivam Shrirao Date: Sun, 23 Aug 2026 14:32:21 +0000 Subject: [PATCH 2/6] docs: simplify Fibo Edit page, fix EXAMPLE_DOC_STRING to the VLM prompt flow --- .../source/en/api/pipelines/bria_fibo_edit.md | 32 +++------------ .../bria_fibo/pipeline_bria_fibo_edit.py | 39 ++++++++++--------- 2 files changed, 25 insertions(+), 46 deletions(-) diff --git a/docs/source/en/api/pipelines/bria_fibo_edit.md b/docs/source/en/api/pipelines/bria_fibo_edit.md index 45166771e426..6f9698cd2099 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) 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), 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,33 +28,8 @@ Use the command below to log in: hf auth login ``` -## Multiple reference images - -`image` accepts a `PIL.Image.Image` or a list of them. A list is interpreted as multiple references, not a batch: each reference is VAE-encoded at its own aspect ratio (capped at 1024x1024 square pixels) and conditions the edit on its own RoPE time plane, while the output resolution follows the first reference. Masks are supported only with a single reference. - -```python -import json - -import torch -from PIL import Image - -from diffusers import BriaFiboEditPipeline - -pipe = BriaFiboEditPipeline.from_pretrained("briaai/Fibo-Edit", torch_dtype=torch.bfloat16) -pipe.to("cuda") - -prompt = {"edit_instruction": "Place the product from the first image in the scene from the second image"} -result = pipe( - prompt=json.dumps(prompt), - image=[Image.open("product.png"), Image.open("scene.png")], - num_inference_steps=50, - guidance_scale=3.5, -) -result.images[0].save("edit.png") -``` - ## BriaFiboEditPipeline [[autodoc]] BriaFiboEditPipeline - all - - __call__ \ No newline at end of file + - __call__ 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 dc3d48829e3a..202507595714 100644 --- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py @@ -53,44 +53,45 @@ EXAMPLE_DOC_STRING = """ Example: ```python - import json - import torch from PIL import Image from diffusers import BriaFiboEditPipeline 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) + # Fibo Edit takes a full structured VGL JSON prompt; build it from the source image + # and an instruction with the edit prompt-to-JSON modular block. + vlm_pipe = ModularPipelineBlocks.from_pretrained("briaai/FIBO-edit-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", 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." - ) - json_prompt = json.loads(output.values["json_prompt"]) - - image = Image.open("image_generate.png") + image = Image.open("owl.png") + json_prompt = vlm_pipe(image=image, prompt="Make the owl into a cat").values["json_prompt"] - json_prompt["edit_instruction"] = "Make the owl to be a cat" - - result = pipe(prompt=json_prompt, num_inference_steps=50, guidance_scale=3.5, image=image, output_type="np") + 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. - json_prompt["edit_instruction"] = "Place the owl from the first image in the forest from the second image" + 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=[Image.open("owl.png"), Image.open("forest.png")], - num_inference_steps=50, - guidance_scale=3.5, + image=[owl, forest], + num_inference_steps=30, + guidance_scale=5, ) + + # 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) ``` """ @@ -658,7 +659,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*): From dcd5bfcf077076b29356a1fe5dbaa6f2fe58c30b Mon Sep 17 00:00:00 2001 From: Shivam Shrirao Date: Mon, 24 Aug 2026 17:43:34 +0000 Subject: [PATCH 3/6] docs: point Fibo Edit at the Gemini prompt-to-JSON block and the 1.5 base/turbo ids briaai/FIBO-edit-prompt-to-JSON is retired in favor of briaai/FIBO-edit-gemini-prompt-to-JSON, which handles multiple reference images and masks. Also corrects the checkpoint ids to their canonical briaai/Fibo-Edit-1.5-base and briaai/Fibo-Edit-1.5-turbo spellings. --- docs/source/en/api/pipelines/bria_fibo_edit.md | 4 ++-- .../pipelines/bria_fibo/pipeline_bria_fibo_edit.py | 11 ++++++----- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/docs/source/en/api/pipelines/bria_fibo_edit.md b/docs/source/en/api/pipelines/bria_fibo_edit.md index 6f9698cd2099..c9672fa46706 100644 --- a/docs/source/en/api/pipelines/bria_fibo_edit.md +++ b/docs/source/en/api/pipelines/bria_fibo_edit.md @@ -16,11 +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) to learn more. A distilled checkpoint is available at [Fibo-Edit-1.5-Turbo](https://huggingface.co/briaai/Fibo-Edit-1.5-Turbo). +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 Edit Hugging Face page](https://huggingface.co/briaai/Fibo-Edit-1.5), 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: 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 202507595714..d7ebf9b5a3bd 100644 --- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py @@ -59,13 +59,14 @@ from diffusers import BriaFiboEditPipeline from diffusers.modular_pipelines import ModularPipelineBlocks - # Fibo Edit takes a full structured VGL JSON prompt; build it from the source image - # and an instruction with the edit prompt-to-JSON modular block. - vlm_pipe = ModularPipelineBlocks.from_pretrained("briaai/FIBO-edit-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-1.5", + "briaai/Fibo-Edit-1.5-base", torch_dtype=torch.bfloat16, ) pipe.to("cuda") @@ -89,7 +90,7 @@ ) # 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 = 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) ``` From 3eb9ba411b6bfc3fd434c440b5ca3290459a02e6 Mon Sep 17 00:00:00 2001 From: Shivam Shrirao Date: Mon, 24 Aug 2026 19:44:51 +0000 Subject: [PATCH 4/6] apply review feedback on reference handling and the attention mask - collapse the mask branch to `if not attention_mask.all():` in both Fibo pipelines - encode references via `self.image_processor.preprocess(...)`, dropping the manual numpy normalization and the `_vae_safe_size` helper - normalize `image` into a reference list once in `__call__`, pass it to `check_inputs` --- .../pipelines/bria_fibo/pipeline_bria_fibo.py | 11 ++-- .../bria_fibo/pipeline_bria_fibo_edit.py | 50 ++++++------------- 2 files changed, 20 insertions(+), 41 deletions(-) diff --git a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py index edd47f34ad67..398758294bcc 100644 --- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py @@ -609,13 +609,10 @@ def __call__( if self._joint_attention_kwargs is None: self._joint_attention_kwargs = {} - if attention_mask.all(): - # Nothing is padded, so the mask is a no-op; skipping it keeps backends without - # mask support (e.g. flash-attn 2/3) usable. - self._joint_attention_kwargs.pop("attention_mask", None) - else: - # 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. + 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) 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 d7ebf9b5a3bd..bebdcea851f8 100644 --- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py @@ -195,25 +195,6 @@ def _vae_safe_dims(width, height, base_resolution=1024, multiple=16): return tuple(max(multiple, int(round(side * scale / multiple)) * multiple) for side in (width, height)) -def _vae_safe_size(image, base_resolution=1024, multiple=16): - """Resize a PIL reference to VAE-safe dimensions.""" - target = _vae_safe_dims(*image.size, base_resolution, multiple) - return image if target == image.size else image.resize(target, Image.LANCZOS) - - -def _as_reference_images(image): - """Normalize the edit input to a non-empty list of PIL reference images. - - A list is interpreted as multiple references, not as a batch; Fibo Edit has no batched-edit API. - """ - if image is None: - return [] - references = image if isinstance(image, list) else [image] - if 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.") - return references - - def paste_mask_on_image(mask: PipelineMaskInput, image: PipelineImageInput): """convert mask and image to PIL Images and paste the mask on the image""" if isinstance(mask, torch.Tensor): @@ -718,7 +699,13 @@ def __call__( generated images. """ - references = _as_reference_images(image) + # 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 references: # Output resolution follows the first reference image; the other @@ -736,7 +723,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, @@ -866,13 +853,10 @@ def __call__( if self._joint_attention_kwargs is None: self._joint_attention_kwargs = {} attention_mask = attention_mask.bool() - if attention_mask.all(): - # Nothing is padded, so the mask is a no-op; skipping it keeps backends without - # mask support (e.g. flash-attn 2/3) usable. - self._joint_attention_kwargs.pop("attention_mask", None) - else: - # 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. + 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) @@ -1018,9 +1002,8 @@ def prepare_reference_latents( ): """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 - image = _vae_safe_size(image.convert("RGB")) - pixels = torch.from_numpy(np.array(image)).permute(2, 0, 1).unsqueeze(0) - encoded_input = pixels.to(device=device, dtype=torch.float32) / 255.0 * 2.0 - 1.0 + 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) @@ -1046,7 +1029,7 @@ def check_inputs( self, prompt, seed, - image, + references, mask, height, width, @@ -1055,8 +1038,7 @@ def check_inputs( ): if seed is not None and not isinstance(seed, int): raise ValueError("Seed must be an integer") - references = _as_reference_images(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.") From a447dfd57c9971deef05a16d068de634643a9179 Mon Sep 17 00:00:00 2001 From: Shivam Shrirao Date: Mon, 24 Aug 2026 20:27:45 +0000 Subject: [PATCH 5/6] make style --- src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 bebdcea851f8..3c3c98eb412f 100644 --- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py @@ -1003,7 +1003,9 @@ def prepare_reference_latents( """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) + 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) From cf6a8cd6d3ee6606e4a19a6c8006e100b631268f Mon Sep 17 00:00:00 2001 From: Shivam Shrirao Date: Mon, 24 Aug 2026 20:30:54 +0000 Subject: [PATCH 6/6] make style --- src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 3c3c98eb412f..28858322fb96 100644 --- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py @@ -60,9 +60,7 @@ from diffusers.modular_pipelines import ModularPipelineBlocks # 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 = ModularPipelineBlocks.from_pretrained("briaai/FIBO-edit-gemini-prompt-to-JSON", trust_remote_code=True) vlm_pipe = vlm_pipe.init_pipeline() pipe = BriaFiboEditPipeline.from_pretrained(