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-turbo is distilled for editing in 4 steps, with guidance_scale=1.0 turning guidance off.
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-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).
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.
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 does the same: flat (128, 128, 128) grey, through the mask made binary at 128.
Fonts for text images: Pillow's built-in font draws most text. For characters it lacks, the pipeline uses Aileron (CC0) and DejaVu Sans (Bitstream Vera license), as training did. It reads them from the fonts folder of briaai/Fibo-2.0-dev only when a text first needs them, so ASCII text never downloads anything. The fonts are in both release repos.
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

sayakpaul · 2026-10-19T13:10:03Z

+        edit["edit_instruction"] = instruction
+        return json.dumps(edit, separators=(",", ":"), ensure_ascii=False)
+
+    def _get_qwen_prompt_embeds(


Is it possible to make it a staticmethod and also pass the processor and text_encoder?

sure. now it takes text_encoder and processor, plus the template, max length and layers (defaulting to module constants), like Flux 2’s _get_mistral_3_small_prompt_embeds.

…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.

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

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

muhammadali-afridi commented Oct 8, 2026 •

edited

Loading

muhammadali-afridi Oct 16, 2026 •

edited

Loading

muhammadali-afridi commented Oct 16, 2026 •

edited

Loading