Add Bria fibo-2 by muhammadali-afridi · Pull Request #14991 · huggingface/diffusers

muhammadali-afridi · 2026-10-08T16:15:42Z

What does this PR do?

Adds fibo-2, Bria's text-to-image and image-editing model.

BriaFibo2Transformer2DModel: an 8.1B flow-matching transformer with a single token stream, in the style of Z-Image. Its text encoder is Qwen3-VL-4B.
- A Perceiver condenses Qwen3-VL's hidden states into gist tokens, which join the image tokens.
- Five blocks also read Qwen3-VL through gated cross-attention.
- Images to edit come in as context_latents. Each is patchified by the shared embedder and gets its own RoPE plane.
BriaFibo2Pipeline:
- text-to-image from structured JSON captions;
- editing of up to 5 images, with an optional mask;
- captions cleaned and serialized the way training did, for generation and edits, in every caption form fibo-2 was trained with: a full caption, a short description (plain text, or "structured_caption": "..." in an edit), or a short description with its text_render list;
- images of the caption's text, which Qwen3-VL sees beside the prompt. The pipeline draws them the way training did, by default (render_text_images=False turns them off), or the caller passes them as text_images.
briaai/Fibo-2.0-dev, the full model, runs 30 steps with guidance 5, the pipeline's defaults. briaai/Fibo-2.0-edit-turbo is distilled for editing in 4 steps, with guidance_scale=1.0 turning guidance off.
BriaFibo2ImageProcessor: prepares the images to edit and the mask the way training did, and gives the size fibo-2 was trained at for an aspect ratio.
scripts/convert_bria_fibo2_to_diffusers.py converts the original checkpoints.
Docs and tests.

Weights: briaai/Fibo-2.0-dev and briaai/Fibo-2.0-edit-turbo, in diffusers format. They're private until the release, and we're happy to give reviewers early access.

Verification

I compared the pipeline against Bria's reference inference code, using the original checkpoints, the same seeds and bf16:

A100, SDPA:
- Text-to-image is pixel-identical, for the full checkpoint and a 4-step distilled one.
- Editing is pixel-identical when the text encoder gets the same images as the reference.
H100: the reference code runs FlashAttention-3 there. With pipe.transformer.set_attention_backend("_flash_3_varlen_hub"):
- text-to-image is pixel-identical for both checkpoints;
- editing is pixel-identical too (2 full-checkpoint cases and 4 distilled ones).
Prompts and text images: compared with Bria's training and evaluation code, every prompt is byte-identical:
- the 215 captions of a text-to-image benchmark, in the full and short forms, compact and pretty-printed (860 prompts), with the prompts that mark and blank their text images, and the images;
- the 115 captions of an edit benchmark, in the full and short forms, with and without text images (460 prompts).
With the fonts training used passed as fonts, the drawn images are also byte-identical on 17 test strings, including non-Latin scripts, wrapped lines and characters no font has.
Text-encoder images: by default the pipeline thumbnails each image straight to 768 px for Qwen3-VL, as in training. The reference code upsamples 4× first, which changes that input slightly.
Review changes: the refactor for review (one prompt parser, fonts as an argument, BriaFibo2ImageProcessor) changes no output. On 153 generation and edit cases through the tiny test pipeline, the text and images Qwen3-VL receives, the VAE's inputs and the outputs are byte-identical before and after. At full scale, so are the trained sizes (75 aspect ratios), the crops, the masking and the prepared images to edit.
Batching: in float32, a batch of prompts of different lengths matches the separate single-prompt calls. This includes the varlen attention backends; see the self-review notes.
Checks (rebased on main):
- make style, make quality and make fix-copies pass.
- check_copies, check_dummies, check_support_list, check_forward_call_docstrings and check_return_annotations pass.
New tests: they pass on a Mac (CPU/MPS), except the group offloading and layerwise casting memory tests. Those fail the same way for Z-Image and Flux 2 on that machine, because MPS has no CUDA streams and no float8.

Notes

mask: fibo-2 was trained on edits whose image has the region to regenerate greyed out. The pipeline's image processor does the same: flat (128, 128, 128) grey, through the mask made binary at 128, at the image's own size before it's resized.
Fonts for text images: Pillow's built-in font draws most text. For characters it lacks, training used Aileron (CC0) and DejaVu Sans (Bitstream Vera license), which are in the fonts folder of both release repos. The pipeline doesn't download them: __call__ takes them as fonts, and the docs show how to get them. Without them, it warns when a text needs them.
Caption cleaning: no training caption had empty values or the aesthetic_score and preference_score fields. The pipeline drops them from text-to-image and edit captions alike and serializes compactly. The negative prompt is used as given.
Standard vs modular: this is a standard pipeline. Happy to add a modular one if that's preferred.

Self-review notes

I ran the self-review skill on the diff.

Fixed:

The transformer forward docstring had no Returns: section, which check_forward_call_docstrings flags.
Varlen backends kept the wrong keys.
- With prompts of different lengths, the padded gist tokens sat between the real gist tokens and the image tokens.
- The varlen backends (flash_varlen, _flash_3_varlen_hub, sage_varlen) keep the first mask.sum() keys, so they kept the wrong ones.
- Padded gist tokens now go last. test_padded_batch_matches_single_prompts covers this.
Removed dead code: a PyTorch 2.0 check in the attention processor, and casts that changed nothing.
__call__ read the first image's size before check_inputs had validated it.

Not fixed, on purpose:

Full-graph compile tests are skipped. RopeEmbedder, copied from Z-Image, builds its frequency table under a device context on first use. Z-Image skips these tests for the same reason.
Gist counts break the graph under torch.compile(transformer).
- With a text mask, the gist count per prompt comes from the text lengths through .tolist().
- Regional compilation (compile_repeated_blocks) is not affected.

Before submitting

Who can review?

@yiyixuxu @sayakpaul @DN6

🤖 Generated with Claude Code

fibo-2 is Bria's 8.1B flow-matching DiT, built on the Z-Image block and conditioned on Qwen3-VL hidden states two ways: a Perceiver resampler turns them into gist tokens that share the stream with the image tokens, and five blocks add gated cross-attention to their own bundle of Qwen3-VL layers. It uses the Flux 2 VAE. One pipeline class covers both of the model's jobs and both checkpoints: text-to-image generation, instruction editing of up to five images with an optional mask, the turbo checkpoint (4 steps, no guidance) and the guided one. - BriaFibo2Transformer2DModel, with images to edit passed as context latents - BriaFibo2Pipeline, including: - captions cleaned and serialized the way training did, for generation and edits, in every caption form fibo-2 was trained with: a full caption, a short description, or a short description with its text_render list - images of the caption's text, which Qwen3-VL sees beside the prompt: drawn the way training drew them, by default (`render_text_images`), or passed as `text_images` - masked edits, the region greyed out through a binary mask as in training - a conversion script from the original checkpoint layout - model and pipeline tests, and docs Verified against Bria's reference implementation on an A100 and an H100.

sayakpaul

Looks pretty good actually! Can we see some examples, please?

Also, we can remove tests where we're calling private methods.

sayakpaul · 2026-10-16T09:35:29Z

+        def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+            with torch.amp.autocast("cuda", enabled=False):
+                x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2))
+                freqs_cis = freqs_cis.unsqueeze(2)
+                x_out = torch.view_as_real(x * freqs_cis).flatten(3)
+                return x_out.type_as(x_in)


We usually avoid autocast blocks within the code. Possible to explicitly handle the casting?

Sure. The same block is in TimestepEmbedder, which is # copied from Z-Image’s. will change it there too if you’d like.

sayakpaul · 2026-10-16T09:38:42Z

+    @pytest.mark.skip(
+        "RopeEmbedder, copied from Z-Image, builds its frequency table under a `torch.device` context on first use, "
+        "which Dynamo cannot trace in a full graph. Z-Image skips this test too."
+    )
+    def test_torch_compile_recompilation_and_graph_break(self):
+        pass


Can we xfail instead?

sayakpaul · 2026-10-16T09:38:51Z

+    @pytest.mark.skip(
+        "RopeEmbedder, copied from Z-Image, cannot be traced in a full graph (see above). Z-Image skips this test too."
+    )
+    def test_compile_on_different_shapes(self):
+        pass


sayakpaul · 2026-10-16T09:39:29Z

+    def test_edit_prompt(self):
+        # The prompt shape fibo-2 was trained to edit with: a vision marker per image, the cleaned caption, then the
+        # instruction
+        pipe = self.get_pipeline()
+        marker = "<|vision_start|><|image_pad|><|vision_end|>"
+        caption = {
+            "short_description": "a blue bicycle",
+            "objects": [],
+            "aesthetics": {"mood": "calm", "aesthetic_score": "high"},
+            "edit_instruction": "Paint <image_1> like <image_2>",
+        }
+
+        assert pipe._get_edit_prompt("Make it blue", 1) == f'{{"image":"{marker}","edit_instruction":"Make it blue"}}'
+        assert pipe._get_edit_prompt(json.dumps(caption), 2) == (
+            f'{{"image_1":"{marker}","image_2":"{marker}","structured_caption":{{"short_description":"a blue bicycle",'
+            f'"aesthetics":{{"mood":"calm"}}}},"edit_instruction":"Paint <image_1> like <image_2>"}}'
+        )


We can remove tests for methods are private.

The rotation already runs in float32 and casts back to the input's dtype, and none of its ops are autocast ops, so the autocast block changed nothing.

…pping them Both fail with torch._dynamo.exc.Unsupported, on the torch.device context in which RopeEmbedder, copied from Z-Image, builds its frequency table. Strict, so they show up once Dynamo can trace it.

Removes the tests of private prompt and text-image helpers. The caption cleaning test now checks that a caption renders like its trained form, and the text images test draws its strips with the public draw_text_image.

muhammadali-afridi · 2026-10-16T16:41:35Z

@sayakpaul, regarding the example, let me confirm with my team about the best way to share weights with you since we don't plan to make it public before the official release

sayakpaul · 2026-10-17T16:56:43Z

Cool, we can merge once the date is up.

HuggingFaceDocBuilderDev · 2026-10-19T13:11:36Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

sayakpaul

Excited to see this model!

sayakpaul · 2026-10-19T13:08:04Z

+        encoder_attention_mask: torch.Tensor | None = None,
+        context_latents: list[torch.Tensor] | None = None,
+        return_dict: bool = True,
+    ) -> torch.Tensor | Transformer2DModelOutput:


Looks like this should be a tuple[torch.Tensor]? I know that this persists across different models under src/diffusers/models (which we can fix in a different PR).

you are right. fixed

…er and processor The chat template, the sequence length and the Qwen3-VL layers become arguments, defaulting to module constants; encode_prompt passes the pipeline's own. The edit and text prompt helpers become classmethods. Embeddings are unchanged.

…-2.0-turbo The docs and the example load Fibo-2.0-dev and show the turbo checkpoint's 4-step edits without guidance. Text-image fonts download from Fibo-2.0-dev, and the conversion script's usage converts a local original checkpoint.

yiyixuxu

thanks for the PR!
sorry i'm bit late on review. left some comments

yiyixuxu · 2026-10-19T21:22:23Z

+        ]
+
+    @classmethod
+    def _mark_text_images(cls, caption: dict, marker: str) -> dict:


can you inline this?

Kept it since it is used in 3 places: the edit prompt, the generation prompt, and a generation’s first steps (with empty markers). Happy to inline it if you still prefer.

yiyixuxu · 2026-10-19T21:27:01Z

+    def _get_text_image_fonts(self) -> list:
+        # The fallback fonts, read on first need from the `fonts` folder of the folder the pipeline was loaded from, or
+        # downloaded from fibo-2's repo when it has none
+        if self._text_image_fonts is None:
+            folder = os.path.join(self.name_or_path, "fonts") if self.name_or_path else None
+            paths = [
+                os.path.join(folder, name)
+                if folder is not None and os.path.isdir(folder)
+                else hf_hub_download(FONTS_REPO_ID, name, subfolder="fonts")
+                for name in FALLBACK_FONT_FILES
+            ]
+            self._text_image_fonts = [load_font(path) for path in paths]
+        return self._text_image_fonts
+


Suggested change

def _get_text_image_fonts(self) -> list:

# The fallback fonts, read on first need from the `fonts` folder of the folder the pipeline was loaded from, or

# downloaded from fibo-2's repo when it has none

if self._text_image_fonts is None:

folder = os.path.join(self.name_or_path, "fonts") if self.name_or_path else None

paths = [

os.path.join(folder, name)

if folder is not None and os.path.isdir(folder)

else hf_hub_download(FONTS_REPO_ID, name, subfolder="fonts")

for name in FALLBACK_FONT_FILES

]

self._text_image_fonts = [load_font(path) for path in paths]

return self._text_image_fonts

can we remove this? pipeline should not resolve or download fonts at call time. we can add a fonts argument to the pipeline so that draw_tet_image can take it directly; you can also move this code to the doc/model card to show user how to get the font outside of the pipeline

sure. __call__ now takes fonts and passes them to draw_text_image, and the docs show how to get the ones training used from the model repo

yiyixuxu · 2026-10-19T21:30:49Z

+        latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2)
+        return latents
+
+    def _get_image_size(self, aspect_ratio: float) -> tuple[int, int]:


can you move image sizing into an image processr? if you need a custom one, you can take a look at flux3 https://github.com/huggingface/diffusers/blob/main/src/diffusers/modular_pipelines/flux3/image_processor.py

done. moved into BriaFibo2ImageProcessor

yiyixuxu · 2026-10-19T21:35:02Z

+        if image is not None:
+            if mask is not None:
+                # As in training: flat (128, 128, 128) grey, through a binary mask
+                grey = PIL.Image.new("RGB", image[0].size, (128, 128, 128))


is it possible for both image and mask to go through an image processor like the other inpainting pipelines? https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux/pipeline_flux_fill.py#L922

sure. I made both an image and a mask to go through an image processor, but we have one difference from Flux Fill. It greys out the region at the image’s own size before resizing, because that’s the order training used, so the region’s edges look the way Fibo 2 saw them

…mask images in BriaFibo2ImageProcessor - `_parse_prompt` replaces `_clean_caption`, `_normalize_caption`, `_split_edit_prompt`, `_get_generation_prompt` and `_get_caption`. `_text_render_entries`, `_get_text_prompt` and `_draw_text_images` are inlined. - The pipeline no longer fetches fonts: `__call__` takes `fonts`, and the docs show how to get the ones training used from briaai/Fibo-2.0-dev. A text that needs them warns without them. - `BriaFibo2ImageProcessor` holds the trained sizes, the cover-and-crop and the masking. The masking stays before the crop, at the image's own size, where training greyed the images. The prompts and images Qwen3-VL sees, the VAE's inputs and the outputs are byte-identical to before.

…eprocess, as Flux3's image processor does

…ibo 2.0, as its model cards do

sayakpaul · 2026-10-20T15:24:43Z

/diffusers-bot pytest tests/pipelines/bria_fibo2/test_pipeline_bria_fibo2.py tests/models/transformers/test_models_transformer_bria_fibo2.py

github-actions · 2026-10-20T15:25:02Z

✅ pytest tests/pipelines/bria_fibo2/test_pipeline_bria_fibo2.py tests/models/transformers/test_models_transformer_bria_fibo2.py passed on GPU — view logs.

sayakpaul

Very cool!

github-actions Bot added documentation Improvements or additions to documentation models tests utils pipelines size/L PR with diff > 200 LOC labels Oct 8, 2026

muhammadali-afridi force-pushed the ali/fibo2 branch 2 times, most recently from e3fdebc to 2b7b792 Compare October 15, 2026 08:47

muhammadali-afridi marked this pull request as ready for review October 15, 2026 08:51

muhammadali-afridi force-pushed the ali/fibo2 branch 2 times, most recently from 553af68 to 5c38cb1 Compare October 15, 2026 09:36

muhammadali-afridi force-pushed the ali/fibo2 branch from 5c38cb1 to 4f7c850 Compare October 15, 2026 10:55

sayakpaul reviewed Oct 16, 2026

View reviewed changes

muhammadali-afridi added 3 commits October 16, 2026 08:17

Apply RoPE with explicit casting instead of an autocast block

0ddf8b8

The rotation already runs in float32 and casts back to the input's dtype, and none of its ops are autocast ops, so the autocast block changed nothing.

Mark the full-graph compile tests as expected failures instead of ski…

eee5acf

…pping them Both fail with torch._dynamo.exc.Unsupported, on the torch.device context in which RopeEmbedder, copied from Z-Image, builds its frequency table. Strict, so they show up once Dynamo can trace it.

Test the pipeline through its public call only

60f9527

Removes the tests of private prompt and text-image helpers. The caption cleaning test now checks that a caption renders like its trained form, and the text images test draws its strips with the public draw_text_image.

Merge branch 'main' into ali/fibo2

d44fa85

sayakpaul approved these changes Oct 19, 2026

View reviewed changes

linoytsaban mentioned this pull request Oct 19, 2026

add fibo-2 LoRA training and loading #15084

Open

muhammadali-afridi added 3 commits October 19, 2026 08:55

Annotate the transformer's tuple return as tuple[torch.Tensor]

ae1b33a

yiyixuxu reviewed Oct 19, 2026

View reviewed changes

muhammadali-afridi added 3 commits October 19, 2026 15:29

Prepare the images to edit and the mask in BriaFibo2ImageProcessor.pr…

531c87b

…eprocess, as Flux3's image processor does

Point to the turbo checkpoint's new name, briaai/Fibo-2.0-edit-turbo

f5bbd61

sayakpaul reviewed Oct 20, 2026

View reviewed changes

muhammadali-afridi and others added 2 commits October 19, 2026 22:44

Show a working example for each use in the docs, and call the model F…

e158bde

…ibo 2.0, as its model cards do

Merge branch 'main' into ali/fibo2

ab96e2b

sayakpaul approved these changes Oct 20, 2026

View reviewed changes

sayakpaul merged commit de988b7 into huggingface:main Oct 20, 2026
12 checks passed

kijai mentioned this pull request Oct 20, 2026

feat: Add Fibo 2.0 support Comfy-Org/ComfyUI#17240

Merged

Conversation

muhammadali-afridi commented Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Verification

Notes

Self-review notes

Before submitting

Who can review?

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

muhammadali-afridi Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

muhammadali-afridi commented Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sayakpaul commented Oct 17, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Oct 19, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sayakpaul commented Oct 20, 2026

Uh oh!

github-actions Bot commented Oct 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

muhammadali-afridi commented Oct 8, 2026 •

edited

Loading

muhammadali-afridi Oct 16, 2026 •

edited

Loading

muhammadali-afridi commented Oct 16, 2026 •

edited

Loading

github-actions Bot commented Oct 20, 2026 •

edited

Loading