From 4f7c85069c0ea7e86e6ec6ffb7db0ac3c0682403 Mon Sep 17 00:00:00 2001 From: muhammadali-afridi Date: Thu, 8 Oct 2026 09:03:20 -0700 Subject: [PATCH 1/4] Add Bria fibo-2: BriaFibo2Pipeline and BriaFibo2Transformer2DModel 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. --- docs/source/en/_toctree.yml | 4 + .../en/api/models/transformer_bria_fibo2.md | 19 + docs/source/en/api/pipelines/bria_fibo2.md | 93 ++ scripts/convert_bria_fibo2_to_diffusers.py | 167 ++++ src/diffusers/__init__.py | 4 + src/diffusers/models/__init__.py | 2 + src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_bria_fibo2.py | 656 +++++++++++++ src/diffusers/pipelines/__init__.py | 2 + .../pipelines/bria_fibo2/__init__.py | 49 + .../bria_fibo2/pipeline_bria_fibo2.py | 927 ++++++++++++++++++ .../pipelines/bria_fibo2/pipeline_output.py | 20 + .../pipelines/bria_fibo2/text_images.py | 140 +++ src/diffusers/utils/dummy_pt_objects.py | 15 + .../dummy_torch_and_transformers_objects.py | 15 + .../test_models_transformer_bria_fibo2.py | 183 ++++ tests/pipelines/bria_fibo2/__init__.py | 0 .../bria_fibo2/test_pipeline_bria_fibo2.py | 458 +++++++++ 18 files changed, 2755 insertions(+) create mode 100644 docs/source/en/api/models/transformer_bria_fibo2.md create mode 100644 docs/source/en/api/pipelines/bria_fibo2.md create mode 100644 scripts/convert_bria_fibo2_to_diffusers.py create mode 100644 src/diffusers/models/transformers/transformer_bria_fibo2.py create mode 100644 src/diffusers/pipelines/bria_fibo2/__init__.py create mode 100644 src/diffusers/pipelines/bria_fibo2/pipeline_bria_fibo2.py create mode 100644 src/diffusers/pipelines/bria_fibo2/pipeline_output.py create mode 100644 src/diffusers/pipelines/bria_fibo2/text_images.py create mode 100644 tests/models/transformers/test_models_transformer_bria_fibo2.py create mode 100644 tests/pipelines/bria_fibo2/__init__.py create mode 100644 tests/pipelines/bria_fibo2/test_pipeline_bria_fibo2.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index ab98a5243059..fd25522a4b85 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -321,6 +321,8 @@ title: AnyFlowTransformer3DModel - local: api/models/aura_flow_transformer2d title: AuraFlowTransformer2DModel + - local: api/models/transformer_bria_fibo2 + title: BriaFibo2Transformer2DModel - local: api/models/transformer_bria_fibo title: BriaFiboTransformer2DModel - local: api/models/bria_transformer @@ -543,6 +545,8 @@ title: Bria 3.2 - local: api/pipelines/bria_fibo title: Bria Fibo + - local: api/pipelines/bria_fibo2 + title: Bria Fibo 2 - local: api/pipelines/bria_fibo_edit title: Bria Fibo Edit - local: api/pipelines/chroma diff --git a/docs/source/en/api/models/transformer_bria_fibo2.md b/docs/source/en/api/models/transformer_bria_fibo2.md new file mode 100644 index 000000000000..737777680a7a --- /dev/null +++ b/docs/source/en/api/models/transformer_bria_fibo2.md @@ -0,0 +1,19 @@ + + +# BriaFibo2Transformer2DModel + +The single-stream transformer of [Bria](https://huggingface.co/briaai)'s fibo-2, built on the Z-Image block. Its text conditioning comes from Qwen3-VL in two ways: a Perceiver resampler turns the text encoder's hidden states into gist tokens that share the stream with the image tokens, and five blocks add a gated cross-attention to their own bundle of Qwen3-VL layers. To edit, the latents of the images to edit join the stream as clean tokens, each image on its own RoPE plane. + +## BriaFibo2Transformer2DModel + +[[autodoc]] BriaFibo2Transformer2DModel diff --git a/docs/source/en/api/pipelines/bria_fibo2.md b/docs/source/en/api/pipelines/bria_fibo2.md new file mode 100644 index 000000000000..3bea6d2fee3e --- /dev/null +++ b/docs/source/en/api/pipelines/bria_fibo2.md @@ -0,0 +1,93 @@ + + +# Bria Fibo 2 + +fibo-2 is an 8.1B parameter flow-matching transformer from [Bria](https://huggingface.co/briaai) that generates images from structured JSON captions. Its text encoder is [Qwen3-VL](https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct): a Perceiver condenses Qwen3-VL's hidden states into gist tokens that join the image tokens in one stream, and five of the transformer's blocks also read Qwen3-VL directly through gated cross-attention. + +The checkpoint, [`briaai/fibo-2`](https://huggingface.co/briaai/fibo-2), runs 30 steps with guidance 5, the pipeline's defaults. With guidance on, the default negative prompt is the null caption the model was trained with: the structured caption with every value left empty. + +fibo-2 is trained on structured JSON captions and will not work well with freeform text. Pass the caption serialized as JSON: + +```py +import json + +prompt = json.dumps(caption, separators=(",", ":"), ensure_ascii=False) +``` + +The pipeline cleans the caption the way training did: it drops empty values and the `aesthetic_score` and `preference_score` fields, which no training caption had, and serializes the caption compactly. Two short forms were trained too and are kept as they are: the caption's `short_description` alone, as plain text, and a caption of exactly a `"short_description"` and a `"text_render"` list, which keeps the list even when it's empty to say the image has no text. + +To turn freeform text into a structured caption, see the prompt-to-JSON models on the [Bria Fibo](https://huggingface.co/briaai/FIBO) page. + +## Editing + +The same pipeline edits images. Pass the image as `image` and the edit instruction as the prompt: + +```py +import torch +from diffusers import BriaFibo2Pipeline +from diffusers.utils import load_image + +pipe = BriaFibo2Pipeline.from_pretrained("briaai/fibo-2", dtype=torch.bfloat16).to("cuda") +image = load_image( + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yarn-art-pikachu.png" +) +edited = pipe("Make the background a snowy forest", image=image).images[0] +``` + +- The prompt can also be a JSON object with an `"edit_instruction"` key and a caption of the result, in any of the forms fibo-2 was trained to edit with. The pipeline turns it into the prompt fibo-2 was trained to edit with. + - **A full caption:** the caption's keys beside `"edit_instruction"`, as Bria's prompt-to-JSON models write for edits. The pipeline drops empty values and the `aesthetic_score` and `preference_score` fields, which training removed from every caption. + - **A short description:** a string in a `"structured_caption"` key, for example `{"structured_caption": "A woman in a red coat on a snowy street.", "edit_instruction": "Make the coat green."}`. + - **A short description and its text:** exactly a `"short_description"` and a `"text_render"` list. The list is kept even when it's empty, which tells fibo-2 the result has no text. + + A `"structured_caption"` key can hold either JSON form too. +- Up to 5 images can be edited together, for example to put an object from one into another. The instruction refers to them in order as ``, `` and so on. +- Unless `height` and `width` are given, the result takes the size fibo-2 was trained at for the aspect ratio of the first image, about one megapixel. A single image is cropped to the size of the result. Several images are each cropped to the trained size for their own aspect ratio. +- `mask` marks the region of a single image to regenerate, white on black. fibo-2 was trained on edits whose image has that region greyed out, and the pipeline greys it out the same way before the image is encoded: flat (128, 128, 128) grey, through the mask made binary at 128. + +## Text in images + +When the caption asks for text, in its `text_render` entries, fibo-2 tends to render it better if Qwen3-VL also sees each text drawn on its own. The pipeline draws these images by default, the way training did, and they can be turned off: + +```py +image = pipe(prompt, render_text_images=False).images[0] +``` + +Training showed fibo-2 these images for some captions with text and not for others, so it works either way. + +- Each `text_render` entry that has `text` gets an image: the text in one line, black on white, in Pillow's built-in font at size 20. A text longer than 256 characters is split into lines of 256. +- Characters that font lacks are drawn in Aileron, then DejaVu Sans. The pipeline reads these fonts when a text first needs them, from the `fonts` folder of the folder it was loaded from, or else from [`briaai/fibo-2`](https://huggingface.co/briaai/fibo-2). +- The pipeline marks each entry with an `image_tokens` field for its image, and Qwen3-VL sees the images after any images to edit. +- A generation runs its first steps, while the noise level is above `pipe.text_images_max_noise_level` (0.9), on the caption without the images. An edit uses them at every step. +- A caption without text gets no images, nor does a short description, as in training. + +To draw the images yourself, pass them as `text_images`, one per `text_render` entry that has `text`, in order, for a single prompt. They take precedence over the pipeline's own. Images drawn differently from training's are out of distribution. + +## Attention backends + +The transformer runs on any of the diffusers [attention backends](../../optimization/attention_backends), for example FlashAttention-3 on Hopper GPUs such as the H100: + +```py +pipe.transformer.set_attention_backend("_flash_3_varlen_hub") +``` + +`_flash_3_hub` also works for a single prompt, but it doesn't take masks, so it can't run a batch of prompts of different lengths (see below). + +## Batching + +Several prompts can be generated in one call. Prompts of different lengths are padded, and the padding is masked out of every attention layer, so such a batch needs an attention backend that takes a mask: the default one, or a varlen backend such as `flash_varlen`. In float32, a batch reproduces the images of the separate calls; in bfloat16 the attention kernels differ between padded and unpadded inputs, so batched images can differ from single calls in fine details. + +## BriaFibo2Pipeline + +[[autodoc]] BriaFibo2Pipeline + - all + - __call__ diff --git a/scripts/convert_bria_fibo2_to_diffusers.py b/scripts/convert_bria_fibo2_to_diffusers.py new file mode 100644 index 000000000000..e2fb02a54894 --- /dev/null +++ b/scripts/convert_bria_fibo2_to_diffusers.py @@ -0,0 +1,167 @@ +import argparse +import json +import os + +import safetensors.torch +import torch +from accelerate import init_empty_weights +from huggingface_hub import snapshot_download +from transformers import AutoImageProcessor, Qwen3VLForConditionalGeneration, Qwen3VLProcessor + +from diffusers import ( + AutoencoderKLFlux2, + BriaFibo2Pipeline, + BriaFibo2Transformer2DModel, + FlowMatchEulerDiscreteScheduler, +) +from diffusers.configuration_utils import FrozenDict + + +""" +# Transformer only + +python scripts/convert_bria_fibo2_to_diffusers.py \ + --original_repo_id briaai/fibo-2 \ + --output_path fibo-2-diffusers + +# Full pipeline + +python scripts/convert_bria_fibo2_to_diffusers.py \ + --original_repo_id briaai/fibo-2 \ + --output_path fibo-2-diffusers \ + --full_pipe +""" + +parser = argparse.ArgumentParser() +parser.add_argument("--original_repo_id", default=None, type=str) +parser.add_argument("--checkpoint_path", default=None, type=str) +parser.add_argument("--output_path", type=str, required=True) +parser.add_argument("--full_pipe", action="store_true") + +args = parser.parse_args() + + +TRANSFORMER_KEYS_RENAME_DICT = { + # one embedder and one final layer instead of Z-Image's one per patch size + "all_x_embedder.2-1.": "x_embedder.", + "all_final_layer.2-1.": "final_layer.", + # injection blocks and Perceiver layers: cross-attention to the text encoder + "cross_q.": "cross_attention.to_q.", + "cross_k.": "cross_attention.to_k.", + "cross_v.": "cross_attention.to_v.", + "cross_out.": "cross_attention.to_out.0.", + "cross_norm_q.": "cross_attention.norm_q.", + "cross_norm_k.": "cross_attention.norm_k.", + # Perceiver layers: self-attention among the gist tokens + "self_q.": "self_attention.to_q.", + "self_k.": "self_attention.to_k.", + "self_v.": "self_attention.to_v.", + "self_out.": "self_attention.to_out.0.", + "self_norm_q.": "self_attention.norm_q.", + "self_norm_k.": "self_attention.norm_k.", +} + +# BriaFibo2Pipeline.text_encoder_out_layers: the Qwen3-VL layers behind the Perceiver, then each injection block +PIPELINE_TEXT_ENCODER_LAYERS = [[9, 20, 31], [5, 16, 27], [7, 18, 29], [9, 20, 31], [11, 22, 33], [13, 24, 35]] + + +def get_checkpoint_path(args): + if args.original_repo_id is not None: + allow_patterns = None if args.full_pipe else ["transformer/*"] + return snapshot_download(args.original_repo_id, allow_patterns=allow_patterns) + elif args.checkpoint_path is not None: + return args.checkpoint_path + else: + raise ValueError("please provide either `original_repo_id` or a local `checkpoint_path`") + + +def load_original_checkpoint(checkpoint_path): + transformer_path = os.path.join(checkpoint_path, "transformer") + with open(os.path.join(transformer_path, "config.json")) as f: + original_config = json.load(f) + with open(os.path.join(transformer_path, "diffusion_pytorch_model.safetensors.index.json")) as f: + shard_names = sorted(set(json.load(f)["weight_map"].values())) + + original_state_dict = {} + for shard_name in shard_names: + original_state_dict.update(safetensors.torch.load_file(os.path.join(transformer_path, shard_name))) + return original_config, original_state_dict + + +def get_transformer_config(original_config): + if original_config["perceiver_num_heads"] not in (None, original_config["n_heads"]): + raise ValueError("The Perceiver must use as many attention heads as the main blocks.") + # BriaFibo2Pipeline feeds the transformer fixed Qwen3-VL layers: the Perceiver's, then three per injection block + text_encoder_layers = [original_config["perceiver_llm_layers"], *original_config["injection_llm_layers"]] + if text_encoder_layers != PIPELINE_TEXT_ENCODER_LAYERS: + raise ValueError( + f"The checkpoint reads Qwen3-VL layers {text_encoder_layers}, but BriaFibo2Pipeline feeds it" + f" {PIPELINE_TEXT_ENCODER_LAYERS}." + ) + + return { + "in_channels": original_config["in_channels"], + "patch_size": original_config["patch_size"], + "dim": original_config["dim"], + "n_layers": original_config["n_layers"], + "n_refiner_layers": original_config["n_refiner_layers"], + "n_heads": original_config["n_heads"], + "cap_feat_dim": original_config["cond_llm_dim"], + "injection_layer_ids": original_config["injection_layer_ids"], + "perceiver_num_layers": original_config["perceiver_num_layers"], + "min_num_gist_tokens": original_config["min_num_gist_tokens"], + "max_num_gist_tokens": original_config["max_num_gist_tokens"], + "gist_step": original_config["gist_step"], + "gist_min_text_len": original_config["gist_min_text_len"], + "gist_max_text_len": original_config["gist_max_text_len"], + } + + +def convert_bria_fibo2_transformer_to_diffusers(original_config, original_state_dict): + with init_empty_weights(): + transformer = BriaFibo2Transformer2DModel.from_config(get_transformer_config(original_config)) + + converted_state_dict = {} + for key, value in original_state_dict.items(): + new_key = key + for old, new in TRANSFORMER_KEYS_RENAME_DICT.items(): + new_key = new_key.replace(old, new) + converted_state_dict[new_key] = value + + transformer.load_state_dict(converted_state_dict, strict=True, assign=True) + return transformer + + +def main(args): + checkpoint_path = get_checkpoint_path(args) + original_config, original_state_dict = load_original_checkpoint(checkpoint_path) + transformer = convert_bria_fibo2_transformer_to_diffusers(original_config, original_state_dict) + + if not args.full_pipe: + transformer.save_pretrained(os.path.join(args.output_path, "transformer")) + return + + vae = AutoencoderKLFlux2.from_pretrained(checkpoint_path, subfolder="vae") + # the original config carries `force_upcast`, which the Flux 2 VAE doesn't take; left in, every load warns about it + vae_config = dict(vae.config) + vae_config.pop("force_upcast", None) + vae._internal_dict = FrozenDict(vae_config) + text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( + checkpoint_path, subfolder="text_encoder", dtype=torch.bfloat16 + ) + # fibo-2's limits on an image's size in pixels, not Qwen3-VL's defaults; they decide how many tokens an image to + # edit becomes + processor = Qwen3VLProcessor.from_pretrained(checkpoint_path, subfolder="text_encoder") + processor.image_processor = AutoImageProcessor.from_pretrained( + checkpoint_path, subfolder="text_encoder", size={"shortest_edge": 4 * 28 * 28, "longest_edge": 2048 * 28 * 28} + ) + scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(checkpoint_path, subfolder="scheduler") + + pipe = BriaFibo2Pipeline( + transformer=transformer, scheduler=scheduler, vae=vae, text_encoder=text_encoder, processor=processor + ) + pipe.save_pretrained(args.output_path) + + +if __name__ == "__main__": + main(args) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 36cd6d283884..7dc9f9fd611a 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -264,6 +264,7 @@ "AutoencoderTiny", "AutoencoderVidTok", "AutoModel", + "BriaFibo2Transformer2DModel", "BriaFiboTransformer2DModel", "BriaTransformer2DModel", "CacheMixin", @@ -626,6 +627,7 @@ "AuraFlowPipeline", "BlipDiffusionControlNetPipeline", "BlipDiffusionPipeline", + "BriaFibo2Pipeline", "BriaFiboEditPipeline", "BriaFiboPipeline", "BriaPipeline", @@ -1162,6 +1164,7 @@ AutoencoderTiny, AutoencoderVidTok, AutoModel, + BriaFibo2Transformer2DModel, BriaFiboTransformer2DModel, BriaTransformer2DModel, CacheMixin, @@ -1499,6 +1502,7 @@ AudioLDM2UNet2DConditionModel, AudioLDMPipeline, AuraFlowPipeline, + BriaFibo2Pipeline, BriaFiboEditPipeline, BriaFiboPipeline, BriaPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index cbe283b73954..e112f2a36eb5 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -113,6 +113,7 @@ _import_structure["transformers.transformer_anyflow_far"] = ["AnyFlowFARTransformer3DModel"] _import_structure["transformers.transformer_bria"] = ["BriaTransformer2DModel"] _import_structure["transformers.transformer_bria_fibo"] = ["BriaFiboTransformer2DModel"] + _import_structure["transformers.transformer_bria_fibo2"] = ["BriaFibo2Transformer2DModel"] _import_structure["transformers.transformer_chroma"] = ["ChromaTransformer2DModel"] _import_structure["transformers.transformer_chronoedit"] = ["ChronoEditTransformer3DModel"] _import_structure["transformers.transformer_cogview3plus"] = ["CogView3PlusTransformer2DModel"] @@ -254,6 +255,7 @@ AnyFlowFARTransformer3DModel, AnyFlowTransformer3DModel, AuraFlowTransformer2DModel, + BriaFibo2Transformer2DModel, BriaFiboTransformer2DModel, BriaTransformer2DModel, ChromaTransformer2DModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 079c4a37a236..692d1869f768 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -24,6 +24,7 @@ from .transformer_anyflow_far import AnyFlowFARTransformer3DModel from .transformer_bria import BriaTransformer2DModel from .transformer_bria_fibo import BriaFiboTransformer2DModel + from .transformer_bria_fibo2 import BriaFibo2Transformer2DModel from .transformer_chroma import ChromaTransformer2DModel from .transformer_chronoedit import ChronoEditTransformer3DModel from .transformer_cogview3plus import CogView3PlusTransformer2DModel diff --git a/src/diffusers/models/transformers/transformer_bria_fibo2.py b/src/diffusers/models/transformers/transformer_bria_fibo2.py new file mode 100644 index 000000000000..7f5bec18aead --- /dev/null +++ b/src/diffusers/models/transformers/transformer_bria_fibo2.py @@ -0,0 +1,656 @@ +# Copyright 2026 Bria AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ...models.modeling_utils import ModelMixin +from ...models.normalization import RMSNorm +from ...utils.torch_utils import maybe_allow_in_graph +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_outputs import Transformer2DModelOutput + + +ADALN_EMBED_DIM = 256 + + +# Copied from diffusers.models.transformers.transformer_z_image.TimestepEmbedder +class TimestepEmbedder(nn.Module): + def __init__(self, out_size, mid_size=None, frequency_embedding_size=256): + super().__init__() + if mid_size is None: + mid_size = out_size + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, mid_size, bias=True), + nn.SiLU(), + nn.Linear(mid_size, out_size, bias=True), + ) + + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + with torch.amp.autocast("cuda", enabled=False): + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t) -> torch.Tensor: + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + compute_dtype = getattr(self.mlp[0], "compute_dtype", None) + if weight_dtype.is_floating_point: + t_freq = t_freq.to(weight_dtype) + elif compute_dtype is not None: + t_freq = t_freq.to(compute_dtype) + t_emb = self.mlp(t_freq) + return t_emb + + +# Copied from diffusers.models.transformers.transformer_z_image.select_per_token +def select_per_token( + value_noisy: torch.Tensor, + value_clean: torch.Tensor, + noise_mask: torch.Tensor, + seq_len: int, +) -> torch.Tensor: + noise_mask_expanded = noise_mask.unsqueeze(-1) # (batch, seq_len, 1) + return torch.where( + noise_mask_expanded == 1, + value_noisy.unsqueeze(1).expand(-1, seq_len, -1), + value_clean.unsqueeze(1).expand(-1, seq_len, -1), + ) + + +# Copied from diffusers.models.transformers.transformer_z_image.FeedForward +class FeedForward(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + self.w2 = nn.Linear(hidden_dim, dim, bias=False) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + + def _forward_silu_gating(self, x1, x3): + return F.silu(x1) * x3 + + def forward(self, x) -> torch.Tensor: + return self.w2(self._forward_silu_gating(self.w1(x), self.w3(x))) + + +# Copied from diffusers.models.transformers.transformer_z_image.FinalLayer +class FinalLayer(nn.Module): + def __init__(self, hidden_size, out_channels): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, out_channels, bias=True) + + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True), + ) + + def forward(self, x, c=None, noise_mask=None, c_noisy=None, c_clean=None) -> torch.Tensor: + seq_len = x.shape[1] + + if noise_mask is not None: + # Per-token modulation + scale_noisy = 1.0 + self.adaLN_modulation(c_noisy) + scale_clean = 1.0 + self.adaLN_modulation(c_clean) + scale = select_per_token(scale_noisy, scale_clean, noise_mask, seq_len) + else: + # Original global modulation + assert c is not None, "Either c or (c_noisy, c_clean) must be provided" + scale = 1.0 + self.adaLN_modulation(c) + scale = scale.unsqueeze(1) + + x = self.norm_final(x) * scale + x = self.linear(x) + return x + + +# Copied from diffusers.models.transformers.transformer_z_image.RopeEmbedder +class RopeEmbedder: + def __init__( + self, + theta: float = 256.0, + axes_dims: list[int] = (16, 56, 56), + axes_lens: list[int] = (64, 128, 128), + ): + self.theta = theta + self.axes_dims = axes_dims + self.axes_lens = axes_lens + assert len(axes_dims) == len(axes_lens), "axes_dims and axes_lens must have the same length" + self.freqs_cis = None + + @staticmethod + def precompute_freqs_cis(dim: list[int], end: list[int], theta: float = 256.0): + with torch.device("cpu"): + freqs_cis = [] + for i, (d, e) in enumerate(zip(dim, end)): + freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d)) + timestep = torch.arange(e, device=freqs.device, dtype=torch.float64) + freqs = torch.outer(timestep, freqs).float() + freqs_cis_i = torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64) # complex64 + freqs_cis.append(freqs_cis_i) + + return freqs_cis + + def __call__(self, ids: torch.Tensor): + assert ids.ndim == 2 + assert ids.shape[-1] == len(self.axes_dims) + device = ids.device + + if self.freqs_cis is None: + self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, self.axes_lens, theta=self.theta) + self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + else: + # Ensure freqs_cis are on the same device as ids + if self.freqs_cis[0].device != device: + self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + + result = [] + for i in range(len(self.axes_dims)): + index = ids[:, i] + result.append(self.freqs_cis[i][index]) + return torch.cat(result, dim=-1) + + +def compute_gist_count( + text_lengths: list[int], min_gist: int, max_gist: int, gist_step: int, min_text_len: int, max_text_len: int +) -> list[int]: + """Number of gist tokens for each prompt, from its text length in tokens. + + Lengths between `min_text_len` and `max_text_len` map linearly onto `[min_gist, max_gist]`, rounded up to a + multiple of `gist_step`. + """ + span = max_text_len - min_text_len + counts = [] + for n in text_lengths: + ratio = min(1.0, max(0.0, (float(n) - min_text_len) / span)) + raw = min_gist + ratio * (max_gist - min_gist) + quantized = int(math.ceil(raw / gist_step)) * gist_step + counts.append(max(min_gist, min(max_gist, quantized))) + return counts + + +class BriaFibo2AttnProcessor: + """Self-attention (with RoPE when `freqs_cis` is given) or cross-attention to `encoder_hidden_states`.""" + + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "BriaFibo2Attention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + freqs_cis: torch.Tensor | None = None, + ) -> torch.Tensor: + query = attn.to_q(hidden_states) + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) + + # Apply Norms + query = attn.norm_q(query) + key = attn.norm_k(key) + + # Apply RoPE + 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) + + if freqs_cis is not None: + query = apply_rotary_emb(query, freqs_cis) + key = apply_rotary_emb(key, freqs_cis) + + # From [batch, seq_len] to [batch, 1, 1, seq_len] -> broadcast to [batch, heads, seq_len, seq_len] + if attention_mask is not None and attention_mask.ndim == 2: + attention_mask = attention_mask[:, None, None, :] + + # Compute joint attention + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + + # Reshape back + hidden_states = hidden_states.flatten(2, 3) + + return attn.to_out[0](hidden_states) + + +class BriaFibo2Attention(nn.Module, AttentionModuleMixin): + _default_processor_cls = BriaFibo2AttnProcessor + _available_processors = [BriaFibo2AttnProcessor] + # The processor reads the separate to_q, to_k and to_v projections + _supports_qkv_fusion = False + + def __init__(self, query_dim: int, heads: int, dim_head: int, kv_dim: int | None = None, eps: float = 1e-5): + super().__init__() + self.heads = heads + kv_dim = kv_dim or query_dim + self.to_q = nn.Linear(query_dim, heads * dim_head, bias=False) + self.to_k = nn.Linear(kv_dim, heads * dim_head, bias=False) + self.to_v = nn.Linear(kv_dim, heads * dim_head, bias=False) + self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, query_dim, bias=False)]) + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + self.set_processor(self._default_processor_cls()) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + freqs_cis: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, freqs_cis) + + +@maybe_allow_in_graph +class BriaFibo2TransformerBlock(nn.Module): + """Z-Image's block. With `cross_attention_dim` it becomes an injection block that also reads the text encoder.""" + + def __init__(self, dim: int, n_heads: int, norm_eps: float, cross_attention_dim: int | None = None): + super().__init__() + self.attention = BriaFibo2Attention(dim, n_heads, dim // n_heads, eps=norm_eps) + self.feed_forward = FeedForward(dim=dim, hidden_dim=(dim * 8 // 3) // 64 * 64) + + self.attention_norm1 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm1 = RMSNorm(dim, eps=norm_eps) + + self.attention_norm2 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) + + self.adaLN_modulation = nn.Sequential(nn.Linear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True)) + + if cross_attention_dim is not None: + self.cross_attn_norm1 = RMSNorm(dim, eps=norm_eps) + self.cross_attention = BriaFibo2Attention( + dim, n_heads, dim // n_heads, kv_dim=cross_attention_dim, eps=norm_eps + ) + self.cross_attn_norm2 = RMSNorm(dim, eps=norm_eps) + self.cross_attn_gate = nn.Parameter(torch.zeros(1)) + else: + self.cross_attention = None + + def forward( + self, + x: torch.Tensor, + attn_mask: torch.Tensor | None, + freqs_cis: torch.Tensor, + adaln_input: torch.Tensor | None = None, + noise_mask: torch.Tensor | None = None, + adaln_noisy: torch.Tensor | None = None, + adaln_clean: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + seq_len = x.shape[1] + + if noise_mask is not None: + # Per-token modulation: different modulation for noisy/clean tokens + mod_noisy = self.adaLN_modulation(adaln_noisy) + mod_clean = self.adaLN_modulation(adaln_clean) + + scale_msa_noisy, gate_msa_noisy, scale_mlp_noisy, gate_mlp_noisy = mod_noisy.chunk(4, dim=1) + scale_msa_clean, gate_msa_clean, scale_mlp_clean, gate_mlp_clean = mod_clean.chunk(4, dim=1) + + gate_msa_noisy, gate_mlp_noisy = gate_msa_noisy.tanh(), gate_mlp_noisy.tanh() + gate_msa_clean, gate_mlp_clean = gate_msa_clean.tanh(), gate_mlp_clean.tanh() + + scale_msa_noisy, scale_mlp_noisy = 1.0 + scale_msa_noisy, 1.0 + scale_mlp_noisy + scale_msa_clean, scale_mlp_clean = 1.0 + scale_msa_clean, 1.0 + scale_mlp_clean + + scale_msa = select_per_token(scale_msa_noisy, scale_msa_clean, noise_mask, seq_len) + scale_mlp = select_per_token(scale_mlp_noisy, scale_mlp_clean, noise_mask, seq_len) + gate_msa = select_per_token(gate_msa_noisy, gate_msa_clean, noise_mask, seq_len) + gate_mlp = select_per_token(gate_mlp_noisy, gate_mlp_clean, noise_mask, seq_len) + else: + # Global modulation: same modulation for all tokens (avoid double select) + mod = self.adaLN_modulation(adaln_input) + scale_msa, gate_msa, scale_mlp, gate_mlp = mod.unsqueeze(1).chunk(4, dim=2) + gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + + # Attention block + attn_out = self.attention(self.attention_norm1(x) * scale_msa, attention_mask=attn_mask, freqs_cis=freqs_cis) + x = x + gate_msa * self.attention_norm2(attn_out) + + # Gated cross-attention to the text encoder (injection blocks only) + if self.cross_attention is not None: + cross_out = self.cross_attention( + self.cross_attn_norm1(x), + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + ) + x = x + self.cross_attn_gate.tanh() * self.cross_attn_norm2(cross_out) + + # FFN block + x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp)) + + return x + + +class BriaFibo2PerceiverLayer(nn.Module): + """Gist tokens read the text encoder (cross-attention), then each other (self-attention), then an FFN.""" + + def __init__(self, dim: int, n_heads: int, cross_attention_dim: int, norm_eps: float): + super().__init__() + self.cross_attn_norm = RMSNorm(dim, eps=norm_eps) + self.cross_attention = BriaFibo2Attention( + dim, n_heads, dim // n_heads, kv_dim=cross_attention_dim, eps=norm_eps + ) + self.self_attn_norm = RMSNorm(dim, eps=norm_eps) + self.self_attention = BriaFibo2Attention(dim, n_heads, dim // n_heads, eps=norm_eps) + self.ffn_norm = RMSNorm(dim, eps=norm_eps) + self.ffn = FeedForward(dim=dim, hidden_dim=(dim * 8 // 3) // 64 * 64) + + def forward( + self, + latents: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + latents = latents + self.cross_attention( + self.cross_attn_norm(latents), + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + ) + latents = latents + self.self_attention(self.self_attn_norm(latents), attention_mask=attention_mask) + latents = latents + self.ffn(self.ffn_norm(latents)) + return latents + + +class BriaFibo2PerceiverResampler(nn.Module): + """Squeezes the text encoder's hidden states into `num_gist_tokens` gist tokens.""" + + def __init__( + self, dim: int, n_heads: int, cross_attention_dim: int, num_queries: int, num_layers: int, norm_eps: float + ): + super().__init__() + self.latent_queries = nn.Parameter(torch.randn(num_queries, dim) * 0.02) + self.layers = nn.ModuleList( + [BriaFibo2PerceiverLayer(dim, n_heads, cross_attention_dim, norm_eps) for _ in range(num_layers)] + ) + self.norm_out = RMSNorm(dim, eps=norm_eps) + + def forward( + self, + encoder_hidden_states: torch.Tensor, + num_gist_tokens: int, + encoder_attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + # `attention_mask` marks each prompt's real gist tokens when a batch pads them to `num_gist_tokens` + batch_size = encoder_hidden_states.shape[0] + latents = self.latent_queries[:num_gist_tokens].unsqueeze(0).expand(batch_size, -1, -1) + for layer in self.layers: + latents = layer(latents, encoder_hidden_states, encoder_attention_mask, attention_mask) + return self.norm_out(latents) + + +class BriaFibo2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin): + _supports_gradient_checkpointing = True + _no_split_modules = ["BriaFibo2TransformerBlock", "BriaFibo2PerceiverLayer"] + _repeated_blocks = ["BriaFibo2TransformerBlock"] + _skip_layerwise_casting_patterns = ["t_embedder"] + + @register_to_config + def __init__( + self, + in_channels: int = 32, + patch_size: int = 2, + dim: int = 3584, + n_layers: int = 44, + n_refiner_layers: int = 2, + n_heads: int = 28, + norm_eps: float = 1e-5, + cap_feat_dim: int = 7680, + injection_layer_ids: tuple[int, ...] = (6, 14, 22, 30, 38), + perceiver_num_layers: int = 2, + min_num_gist_tokens: int = 256, + max_num_gist_tokens: int = 512, + gist_step: int = 64, + gist_min_text_len: int = 256, + gist_max_text_len: int = 2048, + rope_theta: float = 256.0, + t_scale: float = 1000.0, + axes_dims: tuple[int, ...] = (32, 48, 48), + axes_lens: tuple[int, ...] = (2048, 512, 512), + ): + super().__init__() + self.gradient_checkpointing = False + + self.x_embedder = nn.Linear(patch_size * patch_size * in_channels, dim, bias=True) + self.t_embedder = TimestepEmbedder(min(dim, ADALN_EMBED_DIM), mid_size=1024) + self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=axes_dims, axes_lens=axes_lens) + + self.noise_refiner = nn.ModuleList( + [BriaFibo2TransformerBlock(dim, n_heads, norm_eps) for _ in range(n_refiner_layers)] + ) + self.perceiver = BriaFibo2PerceiverResampler( + dim, n_heads, cap_feat_dim, max_num_gist_tokens, perceiver_num_layers, norm_eps + ) + self.layers = nn.ModuleList( + [ + BriaFibo2TransformerBlock( + dim, n_heads, norm_eps, cross_attention_dim=cap_feat_dim if i in injection_layer_ids else None + ) + for i in range(n_layers) + ] + ) + self.final_layer = FinalLayer(dim, patch_size * patch_size * in_channels) + + @staticmethod + # Copied from diffusers.models.transformers.transformer_z_image.ZImageTransformer2DModel.create_coordinate_grid + def create_coordinate_grid(size, start=None, device=None): + if start is None: + start = (0 for _ in size) + axes = [torch.arange(x0, x0 + span, dtype=torch.int32, device=device) for x0, span in zip(start, size)] + grids = torch.meshgrid(axes, indexing="ij") + return torch.stack(grids, dim=-1) + + def _patchify(self, latents: torch.Tensor) -> torch.Tensor: + # 2x2 latent patches become tokens, features ordered (channel, dy, dx) as in training + batch_size, channels, height, width = latents.shape + patch_size = self.config.patch_size + latents = latents.view(batch_size, channels, height // patch_size, patch_size, width // patch_size, patch_size) + return latents.permute(0, 2, 4, 1, 3, 5).reshape( + batch_size, (height // patch_size) * (width // patch_size), -1 + ) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor | None = None, + context_latents: list[torch.Tensor] | None = None, + return_dict: bool = True, + ) -> torch.Tensor | Transformer2DModelOutput: + """ + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, in_channels, height, width)`): + Normalized VAE latents of the image being generated. + timestep (`torch.Tensor` of shape `(batch_size,)`): + Noise level in `[0, 1]`: 1 is pure noise, 0 is a clean image. + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, 1 + len(injection_layer_ids), text_len, cap_feat_dim)`): + Text encoder features: the Perceiver's bundle first, then one bundle per injection block. + encoder_attention_mask (`torch.Tensor` of shape `(batch_size, text_len)`, *optional*): + Bool mask of the real text tokens, for a batch of prompts padded to the same length. Leave it as `None` + when nothing is padded. + context_latents (`list[torch.Tensor]`, *optional*): + Normalized VAE latents of the images to edit, one tensor of shape `(batch_size, in_channels, height, + width)` per image, in order. Each image can have its own size. Their tokens condition the generation as + clean tokens and are not predicted. + return_dict (`bool`, defaults to `True`): + Whether to return a [`~models.modeling_outputs.Transformer2DModelOutput`] instead of a tuple. + + Returns: + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: The predicted velocity, with the shape of + `hidden_states`. If `return_dict` is `False`, a tuple whose first element is that tensor. + """ + batch_size, channels, height, width = hidden_states.shape + patch_size = self.config.patch_size + height_tokens, width_tokens = height // patch_size, width // patch_size + num_image_tokens = height_tokens * width_tokens + device = hidden_states.device + + # 1. Timestep embeddings: image tokens get the real timestep, clean tokens (the gist and the context) get t = 0 + adaln_noisy = self.t_embedder(timestep * self.config.t_scale) + adaln_clean = self.t_embedder(torch.zeros_like(timestep)) + + # 2. Patchify the image and the context images, which share the patch embedding. The image sits on RoPE plane + # `max_num_gist_tokens + 1` whatever the gist count, and context image k on the k-th plane after it with its + # own rows and columns + image_plane = self.config.max_num_gist_tokens + 1 + tokens = [self._patchify(hidden_states)] + ids = [self.create_coordinate_grid((1, height_tokens, width_tokens), (image_plane, 0, 0), device)] + for k, latents in enumerate(context_latents or []): + tokens.append(self._patchify(latents)) + context_size = (1, latents.shape[2] // patch_size, latents.shape[3] // patch_size) + ids.append(self.create_coordinate_grid(context_size, (image_plane + k + 1, 0, 0), device)) + hidden_states = self.x_embedder(torch.cat(tokens, dim=1)) + image_freqs = self.rope_embedder(torch.cat([grid.flatten(0, 2) for grid in ids])).unsqueeze(0) + + # 3. Noise refiner. With context images, the noise mask marks the image tokens as the only noisy ones + refiner_noise_mask = None + if context_latents: + refiner_noise_mask = torch.zeros(batch_size, hidden_states.shape[1], dtype=torch.long, device=device) + refiner_noise_mask[:, :num_image_tokens] = 1 + for block in self.noise_refiner: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, None, image_freqs, adaln_noisy, refiner_noise_mask, adaln_noisy, adaln_clean + ) + else: + hidden_states = block( + hidden_states, + None, + image_freqs, + adaln_input=adaln_noisy, + noise_mask=refiner_noise_mask, + adaln_noisy=adaln_noisy, + adaln_clean=adaln_clean, + ) + + # 4. Gist count per prompt from its text length, then the Perceiver turns the first bundle into gist tokens. + # In a batch, prompts with fewer gist tokens are padded to the largest count and masked out + if encoder_attention_mask is not None: + text_lengths = encoder_attention_mask.sum(dim=1).tolist() + else: + text_lengths = [encoder_hidden_states.shape[2]] * batch_size + gist_counts = compute_gist_count( + text_lengths, + min_gist=self.config.min_num_gist_tokens, + max_gist=self.config.max_num_gist_tokens, + gist_step=self.config.gist_step, + min_text_len=self.config.gist_min_text_len, + max_text_len=self.config.gist_max_text_len, + ) + num_gist_tokens = max(gist_counts) + gist_mask = None + if min(gist_counts) < num_gist_tokens: + gist_mask = ( + torch.arange(num_gist_tokens, device=device) < torch.tensor(gist_counts, device=device)[:, None] + ) + gist = self.perceiver(encoder_hidden_states[:, 0], num_gist_tokens, encoder_attention_mask, gist_mask) + gist_ids = self.create_coordinate_grid((num_gist_tokens, 1, 1), (1, 0, 0), device).flatten(0, 2) + + # 5. One stream: gist tokens first, then the image tokens, then the context tokens, the order fibo-2 was trained + # with. When the batch pads gist tokens, the gist goes last instead, so the padding ends every row as varlen + # attention backends expect. Positions come from RoPE, so the order changes nothing but float rounding. The + # noise mask marks which tokens are noisy + num_context_tokens = hidden_states.shape[1] - num_image_tokens + gist_freqs = self.rope_embedder(gist_ids).unsqueeze(0) + attention_mask = None + if gist_mask is None: + image_start = num_gist_tokens + hidden_states = torch.cat([gist, hidden_states], dim=1) + freqs_cis = torch.cat([gist_freqs, image_freqs], dim=1) + else: + image_start = 0 + hidden_states = torch.cat([hidden_states, gist], dim=1) + freqs_cis = torch.cat([image_freqs, gist_freqs], dim=1) + image_mask = torch.ones(batch_size, num_image_tokens + num_context_tokens, dtype=torch.bool, device=device) + attention_mask = torch.cat([image_mask, gist_mask], dim=1) + noise_mask = torch.zeros(batch_size, hidden_states.shape[1], dtype=torch.long, device=device) + noise_mask[:, image_start : image_start + num_image_tokens] = 1 + + # 6. The main blocks. Injection blocks also read their own bundle, in `injection_layer_ids` order + bundle_index = {layer_id: index + 1 for index, layer_id in enumerate(self.config.injection_layer_ids)} + for layer_id, block in enumerate(self.layers): + text = encoder_hidden_states[:, bundle_index[layer_id]] if layer_id in bundle_index else None + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + attention_mask, + freqs_cis, + None, + noise_mask, + adaln_noisy, + adaln_clean, + text, + encoder_attention_mask, + ) + else: + hidden_states = block( + hidden_states, + attention_mask, + freqs_cis, + noise_mask=noise_mask, + adaln_noisy=adaln_noisy, + adaln_clean=adaln_clean, + encoder_hidden_states=text, + encoder_attention_mask=encoder_attention_mask, + ) + + # 7. Final layer on the image tokens only, then unpatchify back to latents + hidden_states = self.final_layer(hidden_states[:, image_start : image_start + num_image_tokens], c=adaln_noisy) + hidden_states = hidden_states.view(batch_size, height_tokens, width_tokens, channels, patch_size, patch_size) + output = hidden_states.permute(0, 3, 1, 4, 2, 5).reshape(batch_size, channels, height, width) + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 4f1f600e4faf..ff8772b73fc6 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -170,6 +170,7 @@ ] _import_structure["bria"] = ["BriaPipeline"] _import_structure["bria_fibo"] = ["BriaFiboPipeline", "BriaFiboEditPipeline"] + _import_structure["bria_fibo2"] = ["BriaFibo2Pipeline"] _import_structure["flux2"] = [ "Flux2Pipeline", "Flux2KleinPipeline", @@ -622,6 +623,7 @@ from .aura_flow import AuraFlowPipeline from .bria import BriaPipeline from .bria_fibo import BriaFiboEditPipeline, BriaFiboPipeline + from .bria_fibo2 import BriaFibo2Pipeline from .chroma import ChromaImg2ImgPipeline, ChromaInpaintPipeline, ChromaPipeline from .chronoedit import ChronoEditPipeline from .cogvideo import ( diff --git a/src/diffusers/pipelines/bria_fibo2/__init__.py b/src/diffusers/pipelines/bria_fibo2/__init__.py new file mode 100644 index 000000000000..568946ee8642 --- /dev/null +++ b/src/diffusers/pipelines/bria_fibo2/__init__.py @@ -0,0 +1,49 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_bria_fibo2"] = ["BriaFibo2Pipeline"] + + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_bria_fibo2 import BriaFibo2Pipeline + +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/bria_fibo2/pipeline_bria_fibo2.py b/src/diffusers/pipelines/bria_fibo2/pipeline_bria_fibo2.py new file mode 100644 index 000000000000..4ef562923b7e --- /dev/null +++ b/src/diffusers/pipelines/bria_fibo2/pipeline_bria_fibo2.py @@ -0,0 +1,927 @@ +# Copyright 2026 Bria AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import inspect +import json +import math +import os +from typing import Callable + +import numpy as np +import PIL.Image +import torch +from huggingface_hub import hf_hub_download +from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor + +from ...image_processor import VaeImageProcessor +from ...models import AutoencoderKLFlux2, BriaFibo2Transformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import BriaFibo2PipelineOutput +from .text_images import FALLBACK_FONT_FILES, FONTS_REPO_ID, draw_text_image, load_font + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +# The null caption fibo-2 was trained with for classifier-free guidance: every key of the structured caption, every +# value empty +DEFAULT_NEGATIVE_PROMPT = ( + '{"short_description":"","objects":[],"background_setting":"","lighting":{},' + '"aesthetics":{},"photographic_characteristics":{},"style_medium":"",' + '"context":"","artistic_style":""}' +) + +# Where an image stands in a Qwen3-VL prompt +VISION_MARKER = "<|vision_start|><|image_pad|><|vision_end|>" + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import json + + >>> import torch + >>> from diffusers import BriaFibo2Pipeline + >>> from diffusers.utils import load_image + + >>> # fibo-2 reads structured JSON captions, serialized compactly the way it was trained + >>> caption = { + ... "short_description": "A red bicycle leaning against a white brick wall on a sunny morning.", + ... "objects": [{"description": "a red city bicycle with a wicker basket", "location": "center"}], + ... "background_setting": "a white brick wall with a green door", + ... "lighting": {"conditions": "bright morning sun", "direction": "from the right"}, + ... "style_medium": "photograph", + ... } + >>> prompt = json.dumps(caption, separators=(",", ":"), ensure_ascii=False) + + >>> # 30 steps with guidance 5, the defaults + >>> pipe = BriaFibo2Pipeline.from_pretrained("briaai/fibo-2", dtype=torch.bfloat16).to("cuda") + >>> image = pipe(prompt).images[0] + >>> image.save("fibo2.png") + + >>> # editing: pass the image, and the instruction as the prompt + >>> image = load_image( + ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yarn-art-pikachu.png" + ... ) + >>> edited = pipe("Make the background a snowy forest", image=image).images[0] + >>> edited.save("fibo2_edit.png") + ``` +""" + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class BriaFibo2Pipeline(DiffusionPipeline): + r""" + Pipeline for text-to-image generation and image editing with fibo-2. + + Args: + transformer ([`BriaFibo2Transformer2DModel`]): + The transformer that denoises the image latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + Scheduler used with `transformer` to denoise the latents. + vae ([`AutoencoderKLFlux2`]): + Variational auto-encoder that maps images to latents and back. + text_encoder (`Qwen3VLForConditionalGeneration`): + Qwen3-VL. Its hidden states condition the transformer. + processor (`Qwen3VLProcessor`): + Turns the prompt, and the images to edit, into Qwen3-VL inputs. + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents"] + + def __init__( + self, + transformer: BriaFibo2Transformer2DModel, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLFlux2, + text_encoder: Qwen3VLForConditionalGeneration, + processor: Qwen3VLProcessor, + ): + super().__init__() + + self.register_modules( + transformer=transformer, + scheduler=scheduler, + vae=vae, + text_encoder=text_encoder, + processor=processor, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + self.default_sample_size = 128 + self.tokenizer_max_length = 4096 + self.prompt_template_encode = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + # Qwen3-VL layers behind the transformer's six text inputs: the Perceiver's, then one per injection block + self.text_encoder_out_layers = ((9, 20, 31), (5, 16, 27), (7, 18, 29), (9, 20, 31), (11, 22, 33), (13, 24, 35)) + # Editing: fibo-2 edits up to 5 images at once, and Qwen3-VL sees each one at most 768 pixels a side + self.max_num_images = 5 + self.text_encoder_image_size = 768 + # Images of the caption's text: a generation shows them to Qwen3-VL only once the noise level is at or below + # this, as research renders; an edit shows them at every step + self.text_images_max_noise_level = 0.9 + # The fallback fonts the pipeline draws those images with, loaded when a text first needs them + self._text_image_fonts = None + + @staticmethod + def _clean_caption(caption: dict) -> dict: + # Training captions hold no empty values (None, "", {}, [] or NaN), and no aesthetic_score or preference_score: + # training removes them from the caption and from its sections + def empty(value): + if isinstance(value, float): + return math.isnan(value) + return value is None or (isinstance(value, (str, dict, list)) and not value) + + def clean(value): + if isinstance(value, dict): + value = {k: clean(v) for k, v in value.items()} + return {k: v for k, v in value.items() if not empty(v)} + if isinstance(value, list): + return [v for v in map(clean, value) if not empty(v)] + return value + + scores = ("aesthetic_score", "preference_score") + caption = { + key: {k: v for k, v in value.items() if k not in scores} if isinstance(value, dict) else value + for key, value in clean(caption).items() + if key not in scores + } + return {key: value for key, value in caption.items() if value != {}} + + @classmethod + def _normalize_caption(cls, caption: dict) -> dict: + # A caption in the form fibo-2 was trained with. One of exactly a "short_description" and a "text_render" list + # keeps those keys in that order, and the list even when it's empty: it then says the image has no text. Any + # other caption is cleaned + if caption.keys() == {"short_description", "text_render"}: + caption = cls._clean_caption(caption) + return { + "short_description": caption.get("short_description") or "", + "text_render": caption.get("text_render") or [], + } + return cls._clean_caption(caption) + + @classmethod + def _get_generation_prompt(cls, prompt: str) -> str: + # A generation prompt as fibo-2 was trained on it: a JSON caption in its trained form, serialized compactly. A + # prompt that isn't a JSON caption, such as a short description, stays as it is + try: + caption = json.loads(prompt) + except json.JSONDecodeError: + return prompt + if not isinstance(caption, dict): + return prompt + return json.dumps(cls._normalize_caption(caption), separators=(",", ":"), ensure_ascii=False) + + @staticmethod + def _text_render_entries(caption: dict) -> list[dict]: + # The caption's text_render entries that have text: each one gets an image of its text + return [ + entry + for entry in caption.get("text_render") or [] + if isinstance(entry, dict) and isinstance(entry.get("text"), str) and entry["text"] + ] + + @classmethod + def _mark_text_images(cls, caption: dict, marker: str) -> dict: + # A copy of the caption with an "image_tokens" entry on each text_render entry that has text: the vision marker + # of its image, or empty where the image isn't shown + caption = copy.deepcopy(caption) + for entry in cls._text_render_entries(caption): + entry["image_tokens"] = marker + return caption + + def _get_text_prompt(self, prompt: str, marker: str) -> str: + # A generation caption with its text entries marked, serialized compactly the way fibo-2 was trained + caption = self._mark_text_images(json.loads(prompt), marker) + return json.dumps(caption, separators=(",", ":"), ensure_ascii=False) + + @classmethod + def _split_edit_prompt(cls, prompt: str) -> tuple[dict | str | None, str]: + # An edit prompt's caption of the result, if it has one, and its instruction. The prompt is the instruction + # alone, or a JSON object with an "edit_instruction" key and the caption: the object's other keys, or its + # "structured_caption" key. A short description, as a string, is kept as it is, and a JSON caption takes its + # trained form + try: + edit = json.loads(prompt) + except json.JSONDecodeError: + edit = None + if not isinstance(edit, dict): + return None, prompt + edit = dict(edit) + instruction = edit.pop("edit_instruction") + caption = edit.pop("structured_caption") if "structured_caption" in edit else edit + if isinstance(caption, str): + return caption or None, instruction + return cls._normalize_caption(caption) or None, instruction + + def _get_caption(self, prompt: str, edit: bool) -> dict | None: + # The JSON caption whose text_render entries get images: the prompt for a generation, the caption of the result + # for an edit. None for a prompt without one, or with a short description + if edit: + caption = self._split_edit_prompt(prompt)[0] + else: + try: + caption = json.loads(prompt) + except json.JSONDecodeError: + caption = None + return caption if isinstance(caption, dict) else None + + 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 + + def _draw_text_images(self, caption: dict) -> list[PIL.Image.Image]: + # An image of each text_render entry's text, as fibo-2's training drew them + return [ + draw_text_image(entry["text"], self._get_text_image_fonts) for entry in self._text_render_entries(caption) + ] + + def _get_edit_prompt(self, prompt: str, num_images: int, text_images: bool = False) -> str: + # An edit prompt is one JSON object: a vision marker per image, the caption of the result when the prompt has + # one, then the edit instruction. With images of the caption's text, each text entry gets the marker of its + # image, which follows the images to edit + caption, instruction = self._split_edit_prompt(prompt) + if text_images and isinstance(caption, dict): + caption = self._mark_text_images(caption, VISION_MARKER) + + marker = VISION_MARKER + if num_images == 1: + edit = {"image": marker} + else: + edit = {f"image_{k}": marker for k in range(1, num_images + 1)} + if caption: + edit["structured_caption"] = caption + edit["edit_instruction"] = instruction + return json.dumps(edit, separators=(",", ":"), ensure_ascii=False) + + def _get_qwen_prompt_embeds( + self, + prompt: str | list[str], + image: list[PIL.Image.Image] | None, + device: torch.device, + text_images: list[list[PIL.Image.Image]] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + prompt = [prompt] if isinstance(prompt, str) else prompt + text_images = text_images or [[]] * len(prompt) + if len(text_images) != len(prompt): + raise ValueError( + f"`text_images` holds images for {len(text_images)} prompts, but there are {len(prompt)}." + ) + # Qwen3-VL takes the batch's images in the order of their markers: each prompt's images to edit, which every + # prompt edits, then the images of its caption's text + text, images = [], [] + for p, own_text_images in zip(prompt, text_images): + if image is not None: + p = self._get_edit_prompt(p, len(image), text_images=bool(own_text_images)) + images += image + elif own_text_images: + p = self._get_text_prompt(p, VISION_MARKER) + text.append(self.prompt_template_encode.format(p)) + images += own_text_images + text_inputs = self.processor( + text=text, + images=images or None, + padding="longest", + max_length=self.tokenizer_max_length, + truncation=True, + return_tensors="pt", + ).to(device) + hidden_states = self.text_encoder(**text_inputs, output_hidden_states=True).hidden_states + + # Each layer is scaled to unit RMS per token in float32, then every bundle concatenates its three layers + dtype = self.text_encoder.dtype + normalized = {} + for layer in sorted({layer for layers in self.text_encoder_out_layers for layer in layers}): + states = hidden_states[layer].float() + normalized[layer] = (states * torch.rsqrt(states.pow(2).mean(dim=-1, keepdim=True) + 1e-6)).to(dtype) + bundles = [ + torch.cat([normalized[layer] for layer in layers], dim=-1) for layers in self.text_encoder_out_layers + ] + return torch.stack(bundles, dim=1), text_inputs.attention_mask.bool() + + def encode_prompt( + self, + prompt: str | list[str], + image: list[PIL.Image.Image] | None = None, + device: torch.device | None = None, + num_images_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + text_images: list[PIL.Image.Image] | list[list[PIL.Image.Image]] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + r""" + Args: + prompt (`str` or `list[str]`): + The prompt or prompts to encode. With `image`, each prompt is an edit instruction, or a JSON object + with an `"edit_instruction"` key and the caption of the result: its other keys, or a + `"structured_caption"` key holding the caption or a short description. A generation prompt is encoded + as given: `__call__` first puts a JSON caption in the form fibo-2 was trained with. + image (`list[PIL.Image.Image]`, *optional*): + The images to edit, as Qwen3-VL sees them. Every prompt edits all of them. + device (`torch.device`, *optional*): + The device to put the embeddings on. + num_images_per_prompt (`int`, defaults to 1): + Number of images to generate per prompt. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-computed text embeddings of shape `(batch_size, 6, text_len, 7680)`, used instead of `prompt`. + prompt_embeds_mask (`torch.Tensor`, *optional*): + Bool mask of the real tokens in `prompt_embeds`, of shape `(batch_size, text_len)`. + text_images (`list[PIL.Image.Image]` or `list[list[PIL.Image.Image]]`, *optional*): + Images of the caption's text, one per `text_render` entry that has `text`, in order, or a list of them + per prompt. Each entry gets the vision marker of its image, and Qwen3-VL sees the images after any + images to edit. + + Returns: + The text embeddings, and their mask, or `None` when no prompt was padded. + """ + device = device or self._execution_device + if text_images is not None and all(isinstance(img, PIL.Image.Image) for img in text_images): + # The images of a single prompt's text + text_images = [text_images] + if prompt_embeds is None: + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image, device, text_images) + + prompt_embeds = prompt_embeds.to(device).repeat_interleave(num_images_per_prompt, dim=0) + if prompt_embeds_mask is not None: + prompt_embeds_mask = prompt_embeds_mask.to(device).repeat_interleave(num_images_per_prompt, dim=0) + if prompt_embeds_mask.all(): + # Nothing is padded: attention runs without a mask, which every attention backend supports + prompt_embeds_mask = None + return prompt_embeds, prompt_embeds_mask + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._patchify_latents + def _patchify_latents(latents): + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 3, 5, 2, 4) + latents = latents.reshape(batch_size, num_channels_latents * 4, height // 2, width // 2) + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._unpatchify_latents + def _unpatchify_latents(latents): + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width) + latents = latents.permute(0, 1, 4, 2, 5, 3) + 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]: + # The (width, height) fibo-2 was trained at for an aspect ratio: of the sizes divisible by 16, at least 64 + # pixels a side and within 10% of the default 1024x1024 area, the one that crops the least from a picture with + # that aspect ratio, then the one closest to the default area + multiple = self.vae_scale_factor * 2 + min_side = 4 * multiple + area = (self.default_sample_size * self.vae_scale_factor) ** 2 + sizes = [ + (width, height) + for width in range(min_side, 2 * area // min_side, multiple) + for height in range(min_side, 2 * area // width, multiple) + if 0.9 * area <= width * height <= 1.1 * area + ] + + def crop(size): + ratio = size[0] / size[1] + return 1 - min(ratio / aspect_ratio, aspect_ratio / ratio) + + return min(sizes, key=lambda size: (round(crop(size), 9), abs(size[0] * size[1] - area), -size[0])) + + @staticmethod + def _resize_and_crop(image: PIL.Image.Image, width: int, height: int) -> PIL.Image.Image: + # Scale the image to cover (width, height), keeping its aspect ratio, then crop the center, as in training + image = image.convert("RGB") + scale = max(height / image.height, width / image.width) + resized_width, resized_height = math.ceil(image.width * scale), math.ceil(image.height * scale) + image = image.resize((resized_width, resized_height), resample=PIL.Image.LANCZOS) + left, top = (resized_width - width) // 2, (resized_height - height) // 2 + return image.crop((left, top, left + width, top + height)) + + def prepare_image_latents(self, image, batch_size, device, dtype): + # The images to edit are normalized with the VAE's BatchNorm statistics, like the image being generated. The + # pixels are made contiguous so the VAE's convolutions run the same kernels as in the reference implementation + image_latents = [] + for img in image: + img = self.image_processor.preprocess(img).contiguous().to(device=device, dtype=self.vae.dtype) + latents = self._patchify_latents(self.vae.encode(img).latent_dist.mode().float()) + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + latents_bn_scale = (1 / torch.sqrt(self.vae.bn.running_var + self.vae.config.batch_norm_eps)).view( + 1, -1, 1, 1 + ) + latents = (latents - latents_bn_mean) * latents_bn_scale.to(latents.device, latents.dtype) + latents = self._unpatchify_latents(latents) + image_latents.append(latents.repeat(batch_size, 1, 1, 1).to(dtype)) + return image_latents + + def check_inputs( + self, + prompt, + height, + width, + image=None, + mask=None, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + text_images=None, + ): + if any(size is not None and size % (self.vae_scale_factor * 2) != 0 for size in (height, width)): + raise ValueError( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}." + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + "Cannot forward both `prompt` and `prompt_embeds`. Please make sure to only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError("Provide either `prompt` or `prompt_embeds`.") + elif prompt is not None and not isinstance(prompt, (str, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + "Cannot forward both `negative_prompt` and `negative_prompt_embeds`. Please make sure to only forward" + " one of the two." + ) + + if isinstance(prompt, list) and isinstance(negative_prompt, list) and len(prompt) != len(negative_prompt): + raise ValueError( + f"`negative_prompt` has {len(negative_prompt)} prompts but `prompt` has {len(prompt)}. Pass one" + " negative prompt per prompt, or a single string for all of them." + ) + + if image is not None: + if not all(isinstance(img, PIL.Image.Image) for img in image): + raise ValueError("`image` has to be a `PIL.Image.Image` or a list of them.") + if not 1 <= len(image) <= self.max_num_images: + raise ValueError(f"`image` can hold 1 to {self.max_num_images} images but holds {len(image)}.") + for p in [] if prompt is None else [prompt] if isinstance(prompt, str) else prompt: + try: + edit = json.loads(p) + except json.JSONDecodeError: + continue + if not isinstance(edit, dict): + continue + if "edit_instruction" not in edit: + raise ValueError( + "To edit an image, `prompt` has to be the edit instruction, or a JSON caption of the result with" + f' an "edit_instruction" key, but is {p}' + ) + if "structured_caption" in edit and ( + not isinstance(edit["structured_caption"], (dict, str)) or len(edit) > 2 + ): + raise ValueError( + 'An edit prompt with a "structured_caption" key has to hold the caption of the result there, as' + f' a JSON object or a string, and nothing else but the "edit_instruction", but is {p}' + ) + + if mask is not None: + if image is None or len(image) != 1: + raise ValueError("`mask` needs exactly one image in `image`.") + if not isinstance(mask, PIL.Image.Image): + raise ValueError(f"`mask` has to be a `PIL.Image.Image` but is {type(mask)}") + if mask.size != image[0].size: + raise ValueError(f"`mask` has size {mask.size} but `image` has size {image[0].size}.") + + if text_images is not None: + if not isinstance(prompt, str): + raise ValueError(f"`text_images` belong to a single caption, but `prompt` is {prompt}") + if not all(isinstance(img, PIL.Image.Image) for img in text_images): + raise ValueError("`text_images` has to be a list of `PIL.Image.Image`.") + caption = self._get_caption(prompt, edit=image is not None) + if caption is None: + raise ValueError(f"`text_images` needs a JSON caption in `prompt`, but it is {prompt}") + num_texts = len(self._text_render_entries(caption)) + if len(text_images) != num_texts: + raise ValueError( + f"`text_images` holds {len(text_images)} images, but the caption has {num_texts} `text_render`" + " entries with text: pass one image per entry, in order." + ) + + def prepare_latents(self, batch_size, num_channels_latents, height, width, device, generator, latents=None): + if latents is not None: + return latents.to(device=device, dtype=torch.float32) + + shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + return randn_tensor(shape, generator=generator, device=device, dtype=torch.float32) + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: str | list[str] | None = None, + image: PIL.Image.Image | list[PIL.Image.Image] | None = None, + mask: PIL.Image.Image | None = None, + text_images: list[PIL.Image.Image] | None = None, + render_text_images: bool = True, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 30, + guidance_scale: float = 5.0, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str = "pil", + return_dict: bool = True, + callback_on_step_end: Callable[[int, int, dict], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + ) -> BriaFibo2PipelineOutput | tuple: + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `list[str]`, *optional*): + The prompt to guide the image generation, usually a structured JSON prompt, which is cleaned as + fibo-2's training cleaned captions: empty values and the `aesthetic_score` and `preference_score` + fields are dropped, and it is serialized compactly. With `image`, the edit instruction, or a JSON + object with an `"edit_instruction"` key and the caption of the result: its other keys, or a + `"structured_caption"` key holding the caption or a short description. + image (`PIL.Image.Image` or `list[PIL.Image.Image]`, *optional*): + The image to edit, or up to 5 images to edit together, which the instruction can refer to as + ``, `` and so on. Without `image`, the pipeline generates an image from the prompt. + mask (`PIL.Image.Image`, *optional*): + The region of `image` to regenerate, white on black, when editing a single image. As in training, the + region is greyed out before the image is encoded, through the mask made binary at 128. + text_images (`list[PIL.Image.Image]`, *optional*): + Images of the text the caption asks for, one per `text_render` entry that has `text`, in order: each + entry's text drawn on its own, as `render_text_images` draws them. Qwen3-VL sees them with the prompt, + each marked in its caption entry, which helps fibo-2 render the text. Needs `prompt` as a single JSON + caption, and takes precedence over `render_text_images`. A generation runs its first steps, while the + noise level is above `text_images_max_noise_level` (0.9), without them; an edit uses them at every + step. + render_text_images (`bool`, *optional*, defaults to `True`): + Whether to draw the `text_images` of each caption, the way fibo-2's training drew them: each text in + one line, black on white, in Pillow's built-in font at size 20. Characters that font lacks are drawn in + Aileron, then DejaVu Sans, read on first need from the `fonts` folder of the folder the pipeline was + loaded from, or downloaded from `briaai/fibo-2`. A caption without text, and a prompt without a JSON + caption, such as a short description, get no images, as in training. + height (`int`, *optional*): + The height in pixels of the generated image. Must be divisible by 16. Defaults to 1024, or when + editing, to the size fibo-2 was trained at for the aspect ratio of the first image, about one + megapixel. + width (`int`, *optional*): + The width in pixels of the generated image. Must be divisible by 16. Defaults like `height`. + num_inference_steps (`int`, *optional*, defaults to 30): + The number of denoising steps. + guidance_scale (`float`, *optional*, defaults to 5.0): + Classifier-free guidance scale. Values above 1 run the negative prompt as a second pass at every step, + and 1.0 turns guidance off. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt to guide away from when guidance is on. Defaults to the null caption fibo-2 was trained + with: the structured caption with every value empty. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Generator(s) used to draw the starting noise, for deterministic generation. + latents (`torch.Tensor`, *optional*): + Pre-generated starting noise of shape `(batch_size, 32, height // 8, width // 8)`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-computed text embeddings from `encode_prompt`, used instead of `prompt`. + prompt_embeds_mask (`torch.Tensor`, *optional*): + The mask `encode_prompt` returned with `prompt_embeds`. Needed when they are padded. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-computed text embeddings of the negative prompt, used instead of `negative_prompt`. + negative_prompt_embeds_mask (`torch.Tensor`, *optional*): + The mask `encode_prompt` returned with `negative_prompt_embeds`. Needed when they are padded. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format: `"pil"`, `"np"`, `"pt"` or `"latent"`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`BriaFibo2PipelineOutput`] instead of a plain tuple. + callback_on_step_end (`Callable`, *optional*): + A function called at the end of each denoising step as `callback_on_step_end(self, step, timestep, + callback_kwargs)`. + callback_on_step_end_tensor_inputs (`list[str]`, *optional*, defaults to `["latents"]`): + The tensors passed to `callback_on_step_end` in `callback_kwargs`. + + Examples: + + Returns: + [`BriaFibo2PipelineOutput`] or `tuple`: the generated images. + """ + if isinstance(image, PIL.Image.Image): + image = [image] + + # 1. Check inputs + self.check_inputs( + prompt, + height, + width, + image, + mask, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + callback_on_step_end_tensor_inputs, + text_images, + ) + if prompt is not None and image is None: + # A generation caption in its trained form. An edit prompt's caption takes it when the prompt is encoded + prompt = ( + self._get_generation_prompt(prompt) + if isinstance(prompt, str) + else [self._get_generation_prompt(p) for p in prompt] + ) + # The images of each prompt's text: the caller's, or drawn from its caption + if text_images is not None: + text_images = [list(text_images)] + elif render_text_images and prompt is not None: + prompts = [prompt] if isinstance(prompt, str) else prompt + captions = [self._get_caption(p, edit=image is not None) for p in prompts] + text_images = [[] if caption is None else self._draw_text_images(caption) for caption in captions] + if not any(text_images or []): + text_images = None + if image is not None and height is None and width is None: + width, height = self._get_image_size(image[0].width / image[0].height) + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + self._guidance_scale = guidance_scale + self._interrupt = False + device = self._execution_device + + # 2. Prepare the images to edit. A single image is cropped to the size of the result, several images each to + # the size for their own aspect ratio. The VAE encodes the crops, and Qwen3-VL sees them at most 768 pixels a + # side + vae_images = text_encoder_images = None + 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)) + binary_mask = mask.convert("L").point(lambda value: 255 if value >= 128 else 0) + image = [PIL.Image.composite(grey, image[0].convert("RGB"), binary_mask)] + if len(image) == 1: + image_sizes = [(width, height)] + else: + image_sizes = [self._get_image_size(img.width / img.height) for img in image] + vae_images = [self._resize_and_crop(img, *size) for img, size in zip(image, image_sizes)] + text_encoder_images = [img.copy() for img in vae_images] + for img in text_encoder_images: + img.thumbnail((self.text_encoder_image_size, self.text_encoder_image_size)) + + # 3. Encode the prompt with the images to edit, and the negative prompt, without them, when guidance is on + if prompt is not None: + batch_size = 1 if isinstance(prompt, str) else len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + prompt_embeds, prompt_embeds_mask = self.encode_prompt( + prompt, + text_encoder_images, + device, + num_images_per_prompt, + prompt_embeds, + prompt_embeds_mask, + text_images=text_images, + ) + # A generation with images of its text runs its first steps on the captions without them: each text entry's + # marker left empty + base_prompt_embeds = base_prompt_embeds_mask = None + if text_images is not None and image is None: + prompts = [prompt] if isinstance(prompt, str) else prompt + base_prompts = [self._get_text_prompt(p, "") if own else p for p, own in zip(prompts, text_images)] + base_prompt_embeds, base_prompt_embeds_mask = self.encode_prompt( + base_prompts, device=device, num_images_per_prompt=num_images_per_prompt + ) + if self.do_classifier_free_guidance: + if negative_prompt is None and negative_prompt_embeds is None: + negative_prompt = DEFAULT_NEGATIVE_PROMPT + if isinstance(negative_prompt, str): + negative_prompt = [negative_prompt] * batch_size + negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( + negative_prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + prompt_embeds=negative_prompt_embeds, + prompt_embeds_mask=negative_prompt_embeds_mask, + ) + + # 4. Prepare the starting noise, kept in float32 between steps so the updates aren't rounded to the + # transformer's dtype, and the latents of the images to edit + latents = self.prepare_latents( + prompt_embeds.shape[0], self.transformer.config.in_channels, height, width, device, generator, latents + ) + image_latents = None + if vae_images is not None: + image_latents = self.prepare_image_latents(vae_images, latents.shape[0], device, self.transformer.dtype) + + # 5. Prepare timesteps: sigmas from 1 down to 1 / steps, shifted by the image size + patch_size = self.transformer.config.patch_size + image_seq_len = (latents.shape[2] // patch_size) * (latents.shape[3] // patch_size) + mu = calculate_shift( + image_seq_len, + self.scheduler.config.base_image_seq_len, + self.scheduler.config.max_image_seq_len, + self.scheduler.config.base_shift, + self.scheduler.config.max_shift, + ) + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu + ) + self._num_timesteps = len(timesteps) + + # 6. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # the transformer takes the noise level in [0, 1] + timestep = t.expand(latents.shape[0]) / 1000 + latent_model_input = latents.to(self.transformer.dtype) + use_text_images = base_prompt_embeds is None or float(t) / 1000 <= self.text_images_max_noise_level + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds if use_text_images else base_prompt_embeds, + encoder_attention_mask=prompt_embeds_mask if use_text_images else base_prompt_embeds_mask, + context_latents=image_latents, + return_dict=False, + )[0].float() + + # The negative prompt runs as a second pass, as in the reference implementation: batching it with the + # prompt would pad both to the longer of the two. The images to edit condition both passes + if self.do_classifier_free_guidance: + noise_pred_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + encoder_attention_mask=negative_prompt_embeds_mask, + context_latents=image_latents, + return_dict=False, + )[0].float() + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred - noise_pred_uncond) + + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + latents = callback_outputs.pop("latents", latents) + + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + # 7. Decode: undo the latent normalization with the VAE's BatchNorm statistics, then run the VAE decoder + if output_type == "latent": + image = latents + else: + latents = self._patchify_latents(latents) + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + latents_bn_scale = (1 / torch.sqrt(self.vae.bn.running_var + self.vae.config.batch_norm_eps)).view( + 1, -1, 1, 1 + ) + latents = latents / latents_bn_scale.to(latents.device, latents.dtype) + latents_bn_mean + latents = self._unpatchify_latents(latents) + image = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return BriaFibo2PipelineOutput(images=image) diff --git a/src/diffusers/pipelines/bria_fibo2/pipeline_output.py b/src/diffusers/pipelines/bria_fibo2/pipeline_output.py new file mode 100644 index 000000000000..56f658cc9d49 --- /dev/null +++ b/src/diffusers/pipelines/bria_fibo2/pipeline_output.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass + +import numpy as np +import PIL.Image + +from ...utils import BaseOutput + + +@dataclass +class BriaFibo2PipelineOutput(BaseOutput): + """ + Output class for BriaFibo2 pipelines. + + Args: + images (`list[PIL.Image.Image]` or `np.ndarray`) + List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. + """ + + images: list[PIL.Image.Image, np.ndarray] diff --git a/src/diffusers/pipelines/bria_fibo2/text_images.py b/src/diffusers/pipelines/bria_fibo2/text_images.py new file mode 100644 index 000000000000..d74b9eaa0dd9 --- /dev/null +++ b/src/diffusers/pipelines/bria_fibo2/text_images.py @@ -0,0 +1,140 @@ +# Copyright 2026 Bria AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Images of the text a fibo-2 caption asks for, drawn the way fibo-2's training drew them: each text on its own, in one +line, black on white. Pillow's built-in font draws the text. A character it lacks is drawn with the first of the +fallback fonts that has it: Aileron, the full version of the built-in font, then DejaVu Sans. +""" + +import functools +from typing import Callable, Sequence + +import PIL.Image +from PIL import ImageDraw, ImageFont + + +# The size the text is drawn at, and the number of characters longer texts are split into lines of +FONT_SIZE = 20 +MAX_LINE_CHARS = 256 +# The fallback fonts, in order, from the `fonts` folder of fibo-2's repo: Aileron (CC0) and DejaVu Sans (Bitstream +# Vera license) +FONTS_REPO_ID = "briaai/fibo-2" +FALLBACK_FONT_FILES = ("Aileron-Regular.otf", "DejaVuSans.ttf") +_PADDING = 2 +_LINE_SPACING = 4 + + +def load_font(path: str) -> ImageFont.FreeTypeFont: + # With the basic layout, as Pillow's built-in font, so the text doesn't set differently where libraqm is installed + return ImageFont.truetype(path, FONT_SIZE, layout_engine=ImageFont.Layout.BASIC) + + +def draw_text_image(text: str, get_fallback_fonts: Callable[[], Sequence[ImageFont.FreeTypeFont]]) -> PIL.Image.Image: + """ + Draws a text the way fibo-2's training drew each `text_render` entry's text. + + The text is drawn in one line, its whitespace collapsed, black on white with 2 pixels of padding. A text longer + than 256 characters is split into lines of 256. `get_fallback_fonts` is only called for a text with a character + Pillow's built-in font lacks, and returns the fonts to try for it, in order. + """ + font = _get_default_font() + line = " ".join(text.split()) + lines = [line[i : i + MAX_LINE_CHARS] for i in range(0, len(line), MAX_LINE_CHARS)] or [""] + if all(ord(character) < 128 or _has_glyph(font, character) for character in line): + return _draw_one_font(lines, font) + return _draw_mixed(lines, [font, *get_fallback_fonts()]) + + +@functools.lru_cache(maxsize=None) +def _get_default_font() -> ImageFont.FreeTypeFont: + try: + font = ImageFont.load_default(FONT_SIZE) + except TypeError: # Pillow before 10.1 can't size its built-in font + font = None + if not isinstance(font, ImageFont.FreeTypeFont): + raise ImportError("Drawing the images of a caption's text needs Pillow 10.1 or later, built with FreeType.") + return font + + +@functools.lru_cache(maxsize=4096) +def _has_glyph(font: ImageFont.FreeTypeFont, character: str) -> bool: + # A font without a glyph for the character draws its "notdef" box, as it does for U+FFFE, which no font has + try: + notdef = font.getmask("￾") + mask = font.getmask(character) + except Exception: # as in training, a font that can't say is taken to have the glyph + return True + return not (mask.size == notdef.size and bytes(mask) == bytes(notdef)) + + +def _font_runs(text: str, fonts: Sequence[ImageFont.FreeTypeFont]) -> list[tuple[str, ImageFont.FreeTypeFont]]: + # The text split into runs of characters drawn with the same font: the first font with a glyph for the character, + # or the built-in font, the first, when no font has one + runs = [] + for character in text: + font = fonts[0] + if not (ord(character) < 128 or _has_glyph(fonts[0], character)): + font = next((candidate for candidate in fonts[1:] if _has_glyph(candidate, character)), fonts[0]) + if runs and runs[-1][1] is font: + runs[-1][0] += character + else: + runs.append([character, font]) + return [(run, font) for run, font in runs] + + +def _draw_one_font(lines: list[str], font: ImageFont.FreeTypeFont) -> PIL.Image.Image: + # Each line is placed by its bounding box, and is at least as tall as "Ag" + draw = ImageDraw.Draw(PIL.Image.new("RGB", (1, 1), "white")) + min_box = draw.textbbox((0, 0), "Ag", font=font) + min_line_height = max(1, min_box[3] - min_box[1]) + line_sizes = [] + max_width = 1 + total_height = 0 + for line in lines: + box = draw.textbbox((0, 0), line or " ", font=font) + line_height = max(min_line_height, box[3] - box[1]) + line_sizes.append((box, line_height)) + max_width = max(max_width, box[2] - box[0]) + total_height += line_height + total_height += _LINE_SPACING * (len(lines) - 1) + + image = PIL.Image.new("RGB", (max_width + _PADDING * 2, total_height + _PADDING * 2), "white") + draw = ImageDraw.Draw(image) + y = _PADDING + for line, (box, line_height) in zip(lines, line_sizes): + draw.text((_PADDING - box[0], y - box[1]), line, font=font, fill=(0, 0, 0)) + y += line_height + _LINE_SPACING + return image + + +def _draw_mixed(lines: list[str], fonts: Sequence[ImageFont.FreeTypeFont]) -> PIL.Image.Image: + # Each run is drawn with its own font, on the line's baseline, and each line is as tall as the built-in font's + # ascent and descent + ascent, descent = fonts[0].getmetrics() + line_height = max(1, ascent + descent) + runs_per_line = [_font_runs(line or " ", fonts) for line in lines] + width = max(max(1, round(sum(font.getlength(run) for run, font in runs))) for runs in runs_per_line) + height = line_height * len(lines) + _LINE_SPACING * (len(lines) - 1) + + image = PIL.Image.new("RGB", (width + _PADDING * 2, height + _PADDING * 2), "white") + draw = ImageDraw.Draw(image) + y = _PADDING + for runs in runs_per_line: + x = float(_PADDING) + for run, font in runs: + draw.text((x, y + ascent), run, font=font, fill=(0, 0, 0), anchor="ls") + x += font.getlength(run) + y += line_height + _LINE_SPACING + return image diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 9e1f5f027ce0..dbd800992343 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -994,6 +994,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class BriaFibo2Transformer2DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class BriaFiboTransformer2DModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index a068ef8672e8..0bfba5d5ee72 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -1442,6 +1442,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class BriaFibo2Pipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class BriaFiboEditPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/transformers/test_models_transformer_bria_fibo2.py b/tests/models/transformers/test_models_transformer_bria_fibo2.py new file mode 100644 index 000000000000..cfb6998dee40 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_bria_fibo2.py @@ -0,0 +1,183 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import BriaFibo2Transformer2DModel +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class BriaFibo2TransformerTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return BriaFibo2Transformer2DModel + + @property + def output_shape(self) -> tuple[int, ...]: + return (4, 8, 8) + + @property + def input_shape(self) -> tuple[int, ...]: + return (4, 8, 8) + + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | tuple[int, ...]]: + return { + "in_channels": 4, + "patch_size": 2, + "dim": 32, + "n_layers": 2, + "n_refiner_layers": 1, + "n_heads": 2, + "cap_feat_dim": 24, + "injection_layer_ids": (1,), + "perceiver_num_layers": 1, + "min_num_gist_tokens": 4, + "max_num_gist_tokens": 8, + "gist_step": 4, + "gist_min_text_len": 4, + "gist_max_text_len": 16, + "axes_dims": (4, 6, 6), + "axes_lens": (16, 16, 16), + } + + def get_dummy_inputs(self, height: int = 8, width: int = 8) -> dict[str, torch.Tensor]: + batch_size = 1 + text_len = 6 + # one text bundle for the Perceiver, plus one per injection block + num_bundles = 1 + len(self.get_init_dict()["injection_layer_ids"]) + + return { + "hidden_states": randn_tensor( + (batch_size, 4, height, width), generator=self.generator, device=torch_device + ), + "timestep": torch.tensor([0.5], device=torch_device).expand(batch_size), + "encoder_hidden_states": randn_tensor( + (batch_size, num_bundles, text_len, 24), generator=self.generator, device=torch_device + ), + } + + +class TestBriaFibo2TransformerModel(BriaFibo2TransformerTesterConfig, ModelTesterMixin): + def test_context_latents(self): + # Images to edit, each of its own size, condition the output without being predicted + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + context_latents = [ + randn_tensor((1, 4, 6, 10), generator=self.generator, device=torch_device), + randn_tensor((1, 4, 8, 8), generator=self.generator, device=torch_device), + ] + with torch.no_grad(): + output = model(**inputs).sample + output_with_context = model(**inputs, context_latents=context_latents).sample + + assert output_with_context.shape == output.shape + assert not torch.allclose(output_with_context, output) + + def test_padded_batch_matches_single_prompts(self): + # Prompts of 4 and 16 text tokens get 4 and 8 gist tokens, so the batch pads the first prompt's gist. The + # padding has to end every row of the attention mask, which varlen attention backends assume + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + hidden_states = randn_tensor((2, 4, 8, 8), generator=self.generator, device=torch_device) + timestep = torch.tensor([0.5, 0.5], device=torch_device) + encoder_hidden_states = randn_tensor((2, 2, 16, 24), generator=self.generator, device=torch_device) + encoder_attention_mask = torch.ones(2, 16, dtype=torch.bool, device=torch_device) + encoder_attention_mask[0, 4:] = False + text_lengths = (4, 16) + + for context_latents in (None, [randn_tensor((2, 4, 6, 10), generator=self.generator, device=torch_device)]): + attention_masks = [] + hook = model.layers[0].register_forward_pre_hook(lambda module, args: attention_masks.append(args[1])) + with torch.no_grad(): + batch = model( + hidden_states, + timestep, + encoder_hidden_states, + encoder_attention_mask, + context_latents=context_latents, + ).sample + hook.remove() + with torch.no_grad(): + single = [ + model( + hidden_states[i : i + 1], + timestep[i : i + 1], + encoder_hidden_states[i : i + 1, :, : text_lengths[i]], + context_latents=None if context_latents is None else [c[i : i + 1] for c in context_latents], + ).sample + for i in range(2) + ] + + attention_mask = attention_masks[0] + assert not attention_mask.all() + assert torch.equal(attention_mask, attention_mask.long().cumprod(dim=1).bool()) + assert torch.allclose(batch, torch.cat(single), atol=1e-5, rtol=1e-5) + + +class TestBriaFibo2TransformerMemory(BriaFibo2TransformerTesterConfig, MemoryTesterMixin): + pass + + +class TestBriaFibo2TransformerTorchCompile(BriaFibo2TransformerTesterConfig, TorchCompileTesterMixin): + @property + def different_shapes_for_compilation(self): + return [(4, 4), (4, 8), (8, 8)] + + @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 + + def test_torch_compile_repeated_blocks(self): + # BriaFibo2TransformerBlock runs three ways: in the noise refiner (one timestep for every token), in the main + # blocks (per-token timesteps) and in the injection blocks (plus cross-attention), so it compiles three times + super().test_torch_compile_repeated_blocks(recompile_limit=3) + + @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 + + +class TestBriaFibo2TransformerTraining(BriaFibo2TransformerTesterConfig, TrainingTesterMixin): + pass + + +class TestBriaFibo2TransformerAttention(BriaFibo2TransformerTesterConfig, AttentionTesterMixin): + pass diff --git a/tests/pipelines/bria_fibo2/__init__.py b/tests/pipelines/bria_fibo2/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/bria_fibo2/test_pipeline_bria_fibo2.py b/tests/pipelines/bria_fibo2/test_pipeline_bria_fibo2.py new file mode 100644 index 000000000000..df28d3583434 --- /dev/null +++ b/tests/pipelines/bria_fibo2/test_pipeline_bria_fibo2.py @@ -0,0 +1,458 @@ +import json + +import numpy as np +import PIL.Image +import pytest +import torch +from transformers import AutoConfig, Qwen3VLForConditionalGeneration, Qwen3VLProcessor + +from diffusers import ( + AutoencoderKLFlux2, + BriaFibo2Pipeline, + BriaFibo2Transformer2DModel, + FlowMatchEulerDiscreteScheduler, +) + +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin + + +class BriaFibo2PipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = BriaFibo2Pipeline + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "negative_prompt", "prompt_embeds", "negative_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 8, 8) + # `encode_prompt` runs the Qwen3-VL processor as well as the text encoder + text_stack_component_names = ("text", "processor") + + def get_dummy_components(self): + tiny_ckpt_id = "huangfeice/tiny-random-Qwen3VLForConditionalGeneration" + + torch.manual_seed(0) + transformer = BriaFibo2Transformer2DModel( + in_channels=1, + patch_size=2, + dim=32, + n_layers=2, + n_refiner_layers=1, + n_heads=2, + cap_feat_dim=48, # three Qwen3-VL layers of width 16, side by side + injection_layer_ids=(1,), + perceiver_num_layers=1, + min_num_gist_tokens=4, + max_num_gist_tokens=8, + gist_step=4, + gist_min_text_len=4, + gist_max_text_len=16, + axes_dims=(4, 6, 6), + axes_lens=(16, 16, 16), + ) + + # The pipeline reads Qwen3-VL's hidden states up to layer 35, so the tiny text encoder keeps 36 layers + config = AutoConfig.from_pretrained(tiny_ckpt_id) + config.text_config.num_hidden_layers = 36 + torch.manual_seed(0) + text_encoder = Qwen3VLForConditionalGeneration(config) + processor = Qwen3VLProcessor.from_pretrained(tiny_ckpt_id) + text_encoder.resize_token_embeddings(len(processor.tokenizer)) + + torch.manual_seed(0) + vae = AutoencoderKLFlux2( + sample_size=32, + in_channels=3, + out_channels=3, + down_block_types=("DownEncoderBlock2D",), + up_block_types=("UpDecoderBlock2D",), + block_out_channels=(4,), + layers_per_block=1, + latent_channels=1, + norm_num_groups=1, + use_quant_conv=False, + use_post_quant_conv=False, + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + return { + "transformer": transformer, + "scheduler": scheduler, + "vae": vae, + "text_encoder": text_encoder, + "processor": processor, + } + + def get_dummy_inputs(self): + inputs = { + "prompt": '{"short_description":"a red bicycle leaning against a white brick wall"}', + "generator": self.get_generator(0), + "num_inference_steps": 2, + "guidance_scale": 1.0, + "height": 8, + "width": 8, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", + } + return inputs + + +class TestBriaFibo2Pipeline(BriaFibo2PipelineTesterConfig, PipelineTesterMixin): + def test_default_negative_prompt_with_guidance(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs["guidance_scale"] = 5.0 + default_negative = pipe(**inputs).images + + inputs = self.get_dummy_inputs() + inputs["guidance_scale"] = 5.0 + inputs["negative_prompt"] = '{"short_description":"blurry"}' + other_negative = pipe(**inputs).images + + assert default_negative.shape == (1, *self.output_shape) + assert not torch.allclose(default_negative, other_negative) + + def get_edit_image(self, width=16, height=12, seed=0): + pixels = np.random.default_rng(seed).integers(0, 256, (height, width, 3), dtype=np.uint8) + return PIL.Image.fromarray(pixels) + + def test_edit(self): + pipe = self.get_pipeline() + generated = pipe(**self.get_dummy_inputs()).images + + inputs = self.get_dummy_inputs() + inputs["prompt"] = "Make the bicycle blue" + inputs["image"] = self.get_edit_image() + edited = pipe(**inputs).images + + assert edited.shape == (1, *self.output_shape) + assert not torch.allclose(edited, generated) + + def test_edit_with_mask(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs["prompt"] = "Remove the bicycle" + inputs["image"] = self.get_edit_image() + edited = pipe(**inputs).images + + inputs = self.get_dummy_inputs() + inputs["prompt"] = "Remove the bicycle" + inputs["image"] = self.get_edit_image() + inputs["mask"] = PIL.Image.new("L", (16, 12)) + inputs["mask"].paste(255, (4, 2, 12, 10)) + masked = pipe(**inputs).images + + assert masked.shape == (1, *self.output_shape) + assert not torch.allclose(masked, edited) + + def test_mask_should_be_made_binary(self): + # Training greyed out the region through a binary mask: from 128 up a pixel is greyed out fully, below not at all + pipe = self.get_pipeline() + + def edit(value): + mask = PIL.Image.new("L", (16, 12)) + mask.paste(value, (4, 2, 12, 10)) + inputs = {**self.get_dummy_inputs(), "prompt": "Remove the bicycle", "image": self.get_edit_image()} + return pipe(**inputs, mask=mask).images + + assert torch.equal(edit(128), edit(255)) + assert torch.equal(edit(127), edit(0)) + assert not torch.allclose(edit(255), edit(0)) + + def test_edit_several_images(self): + pipe = self.get_pipeline() + # Several images are each encoded at the default area, which the tiny model keeps small + pipe.default_sample_size = 16 + inputs = self.get_dummy_inputs() + inputs["prompt"] = '{"short_description":"a blue bicycle","edit_instruction":"Paint like "}' + inputs["image"] = [self.get_edit_image(16, 12, seed=0), self.get_edit_image(10, 20, seed=1)] + edited = pipe(**inputs).images + + assert edited.shape == (1, *self.output_shape) + + 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 like ", + } + + 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 like "}}' + ) + + def test_edit_prompt_caption_forms(self): + # fibo-2 was also trained to edit with a short description, as a string, and with a short description and the + # caption's text_render list, in that order, the list kept even when empty + pipe = self.get_pipeline() + marker = "<|vision_start|><|image_pad|><|vision_end|>" + description = "A woman in a red coat on a snowy street." + instruction = "Make the coat green." + short_json = {"short_description": description, "text_render": []} + + short = json.dumps({"structured_caption": description, "edit_instruction": instruction}) + assert pipe._get_edit_prompt(short, 1) == ( + f'{{"image":"{marker}","structured_caption":"{description}","edit_instruction":"{instruction}"}}' + ) + for prompt in [ + {**short_json, "edit_instruction": instruction}, + {"text_render": [], "short_description": description, "edit_instruction": instruction}, + {"structured_caption": short_json, "edit_instruction": instruction}, + ]: + assert pipe._get_edit_prompt(json.dumps(prompt), 1) == ( + f'{{"image":"{marker}","structured_caption":{{"short_description":"{description}","text_render":[]}},' + f'"edit_instruction":"{instruction}"}}' + ) + # An empty description is no caption, as in training + no_caption = json.dumps({"structured_caption": "", "edit_instruction": instruction}) + assert pipe._get_edit_prompt(no_caption, 1) == f'{{"image":"{marker}","edit_instruction":"{instruction}"}}' + + def get_text_image(self, width=24, height=8): + return PIL.Image.new("RGB", (width, height), "white") + + def record_text_lengths(self, pipe): + # The length of the text the transformer is conditioned on, at each call + lengths = [] + pipe.transformer.register_forward_pre_hook( + lambda module, args, kwargs: lengths.append(kwargs["encoder_hidden_states"].shape[2]), with_kwargs=True + ) + return lengths + + def test_text_prompt(self): + # Each text_render entry with text gets the marker of its image, or an empty one where the image isn't shown + pipe = self.get_pipeline() + marker = "<|vision_start|><|image_pad|><|vision_end|>" + prompt = json.dumps( + { + "short_description": "a shop sign", + "text_render": [{"text": "OPEN"}, {"text": ""}, {"location": "door"}, {"text": "7 AM"}], + } + ) + entries = '[{{"text":"OPEN","image_tokens":"{0}"}},{{"text":""}},{{"location":"door"}},{{"text":"7 AM","image_tokens":"{0}"}}]' + + assert pipe._get_text_prompt(prompt, marker) == ( + f'{{"short_description":"a shop sign","text_render":{entries.format(marker)}}}' + ) + assert ( + pipe._get_text_prompt(prompt, "") + == f'{{"short_description":"a shop sign","text_render":{entries.format("")}}}' + ) + + def test_edit_prompt_with_text_images(self): + # The caption is cleaned first, then each text entry is marked: its image follows the image to edit + pipe = self.get_pipeline() + marker = "<|vision_start|><|image_pad|><|vision_end|>" + caption = { + "short_description": "a cafe window", + "text_render": [{"text": "OPEN DAILY", "font": ""}], + "edit_instruction": "Turn it into a flat illustration", + } + + assert pipe._get_edit_prompt(json.dumps(caption), 1, text_images=True) == ( + f'{{"image":"{marker}","structured_caption":{{"short_description":"a cafe window","text_render":' + f'[{{"text":"OPEN DAILY","image_tokens":"{marker}"}}]}},"edit_instruction":"Turn it into a flat illustration"}}' + ) + + def test_generation_prompt_should_take_the_trained_form(self): + # Empty values, NaN and the score fields dropped, serialized compactly. A short description and a caption of a + # short description and its text_render list keep their own trained forms + pipe = self.get_pipeline() + pretty = ( + '{\n "short_description": "a red bicycle",\n "objects": [],\n "context": NaN,\n' + ' "aesthetics": {"mood": "calm", "aesthetic_score": "high"},\n "preference_score": "very high"\n}' + ) + clean = '{"short_description":"a red bicycle","aesthetics":{"mood":"calm"}}' + + assert pipe._get_generation_prompt(pretty) == clean + assert pipe._get_generation_prompt("A red bicycle by a door.") == "A red bicycle by a door." + assert ( + pipe._get_generation_prompt('{"text_render": [], "short_description": "a sign"}') + == '{"short_description":"a sign","text_render":[]}' + ) + from_pretty = pipe(**{**self.get_dummy_inputs(), "prompt": pretty}).images + from_clean = pipe(**{**self.get_dummy_inputs(), "prompt": clean}).images + assert torch.equal(from_pretty, from_clean) + + def test_generation_with_text_images(self): + pipe = self.get_pipeline() + prompt = '{"short_description":"a shop sign","text_render":[{"text":"OPEN"}]}' + without = pipe(**{**self.get_dummy_inputs(), "prompt": prompt, "render_text_images": False}).images + with_images = pipe( + **{**self.get_dummy_inputs(), "prompt": prompt, "text_images": [self.get_text_image()]} + ).images + + assert with_images.shape == (1, *self.output_shape) + assert not torch.allclose(with_images, without) + + def test_generation_shows_text_images_only_at_low_noise(self): + pipe = self.get_pipeline() + lengths = self.record_text_lengths(pipe) + inputs = { + **self.get_dummy_inputs(), + "num_inference_steps": 4, + "prompt": '{"short_description":"a shop sign","text_render":[{"text":"OPEN"}]}', + "text_images": [self.get_text_image()], + } + pipe(**inputs) + + shown = [level <= pipe.text_images_max_noise_level for level in (pipe.scheduler.timesteps / 1000).tolist()] + assert True in shown and False in shown + assert [length == max(lengths) for length in lengths] == shown + + def test_edit_shows_text_images_at_every_step(self): + pipe = self.get_pipeline() + lengths = self.record_text_lengths(pipe) + inputs = { + **self.get_dummy_inputs(), + "num_inference_steps": 4, + "prompt": '{"text_render":[{"text":"OPEN"}],"edit_instruction":"Make the sign blue"}', + "image": self.get_edit_image(), + "text_images": [self.get_text_image()], + } + pipe(**inputs) + + assert len(lengths) == 4 and len(set(lengths)) == 1 + + def test_text_images_are_checked(self): + pipe = self.get_pipeline() + caption = '{"short_description":"a shop sign","text_render":[{"text":"OPEN"}]}' + for wrong_inputs in [ + {"prompt": caption, "text_images": [self.get_text_image()] * 2}, + {"prompt": "Write OPEN on the sign", "text_images": [self.get_text_image()]}, + {"prompt": [caption, caption], "text_images": [self.get_text_image()]}, + {"prompt": caption, "text_images": ["not an image"]}, + { + "prompt": '{"structured_caption":"a shop sign","edit_instruction":"Make the sign blue"}', + "image": self.get_edit_image(), + "text_images": [self.get_text_image()], + }, + ]: + with pytest.raises(ValueError): + pipe(**{**self.get_dummy_inputs(), **wrong_inputs}) + + def test_draw_text_images(self): + # One image per text_render entry with text: the text in one line, black on white, split into lines of 256 + # characters + pipe = self.get_pipeline() + caption = {"text_render": [{"text": "OPEN\n DAILY"}, {"text": ""}, {"location": "door"}, {"text": "x" * 257}]} + images = pipe._draw_text_images(caption) + one_line = pipe._draw_text_images({"text_render": [{"text": "OPEN DAILY"}, {"text": "x" * 256}]}) + + assert len(images) == 2 + assert images[0].mode == "RGB" and images[0].getpixel((0, 0)) == (255, 255, 255) + assert np.asarray(images[0]).min() == 0 + assert np.array_equal(np.asarray(images[0]), np.asarray(one_line[0])) + # Lines of 256 and 1 characters, each as tall as the line of 256, and 4 pixels apart + assert images[1].size == (one_line[1].width, 2 * one_line[1].height) + + def test_text_images_read_fonts_only_when_needed(self): + # Pillow's built-in font draws ASCII on its own: only a character it lacks needs the fallback fonts + pipe = self.get_pipeline() + calls = [] + pipe._get_text_image_fonts = lambda: calls.append(True) or [] + + pipe._draw_text_images({"text_render": [{"text": "OPEN 7AM-4PM"}]}) + assert calls == [] + pipe._draw_text_images({"text_render": [{"text": "OPEN 7AM–4PM"}]}) + assert calls == [True] + + @pytest.mark.parametrize("edit", [False, True], ids=["generation", "edit"]) + def test_render_text_images(self, edit): + # By default the pipeline draws the images it would otherwise be given as text_images + pipe = self.get_pipeline() + caption = {"short_description": "a shop sign", "text_render": [{"text": "OPEN"}, {"text": "7 AM"}]} + inputs = {**self.get_dummy_inputs(), "prompt": json.dumps(caption)} + if edit: + inputs["prompt"] = json.dumps({**caption, "edit_instruction": "Make the sign blue"}) + inputs["image"] = self.get_edit_image() + + rendered = pipe(**inputs).images + inputs["generator"] = self.get_generator(0) + given = pipe(**inputs, text_images=pipe._draw_text_images(caption)).images + inputs["generator"] = self.get_generator(0) + without = pipe(**inputs, render_text_images=False).images + + assert torch.equal(rendered, given) + assert not torch.allclose(rendered, without) + + def test_render_text_images_without_text(self): + # A caption without text, a prompt without a caption and a short description get no images of text + pipe = self.get_pipeline() + for prompt_inputs in [ + {"prompt": '{"short_description":"a shop sign","text_render":[]}'}, + {"prompt": "a shop sign"}, + { + "prompt": '{"structured_caption":"a shop sign","edit_instruction":"Make the sign blue"}', + "image": self.get_edit_image(), + }, + ]: + inputs = {**self.get_dummy_inputs(), **prompt_inputs} + rendered = pipe(**inputs).images + inputs["generator"] = self.get_generator(0) + without = pipe(**inputs, render_text_images=False).images + + assert torch.equal(rendered, without) + + def test_given_text_images_should_take_precedence_over_drawn_ones(self): + pipe = self.get_pipeline() + inputs = { + **self.get_dummy_inputs(), + "prompt": '{"short_description":"a shop sign","text_render":[{"text":"OPEN"}]}', + } + given = pipe(**inputs, text_images=[self.get_text_image()]).images + inputs["generator"] = self.get_generator(0) + only_given = pipe(**inputs, text_images=[self.get_text_image()], render_text_images=False).images + inputs["generator"] = self.get_generator(0) + drawn = pipe(**inputs).images + + assert torch.equal(given, only_given) + assert not torch.allclose(given, drawn) + + def test_a_batch_should_draw_each_captions_own_text_images(self): + # Each prompt's images follow its own markers, so a batch renders as its prompts do one at a time + pipe = self.get_pipeline() + captions = [ + '{"short_description":"a shop sign","text_render":[{"text":"OPEN"}]}', + '{"short_description":"a red bicycle"}', + ] + generators = [self.get_generator(0), self.get_generator(1)] + batch = pipe(**{**self.get_dummy_inputs(), "prompt": captions, "generator": generators}).images + singles = [ + pipe(**{**self.get_dummy_inputs(), "prompt": caption, "generator": self.get_generator(seed)}).images + for seed, caption in enumerate(captions) + ] + + assert torch.allclose(batch, torch.cat(singles), atol=1e-4) + + def test_edit_inputs_are_checked(self): + pipe = self.get_pipeline() + mask = PIL.Image.new("L", (16, 12)) + for wrong_inputs in [ + {"image": [self.get_edit_image()] * 6}, + {"mask": mask}, + {"image": [self.get_edit_image()] * 2, "mask": mask}, + {"image": self.get_edit_image(), "mask": PIL.Image.new("L", (12, 16))}, + {"image": self.get_edit_image(), "prompt": '{"short_description":"no edit instruction"}'}, + { + "image": self.get_edit_image(), + "prompt": '{"structured_caption":["a sign"],"edit_instruction":"Make it blue"}', + }, + { + "image": self.get_edit_image(), + "prompt": '{"structured_caption":"a sign","objects":[],"edit_instruction":"Make it blue"}', + }, + ]: + inputs = {**self.get_dummy_inputs(), "prompt": "Make it blue", **wrong_inputs} + with pytest.raises(ValueError): + pipe(**inputs) + + +class TestBriaFibo2PipelineMemory(BriaFibo2PipelineTesterConfig, MemoryTesterMixin): + pass From 0ddf8b86dd0d5b83b4793ea01ca83417ed908cd9 Mon Sep 17 00:00:00 2001 From: muhammadali-afridi Date: Fri, 16 Oct 2026 08:17:39 -0700 Subject: [PATCH 2/4] Apply RoPE with explicit casting instead of an autocast block 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. --- .../models/transformers/transformer_bria_fibo2.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_bria_fibo2.py b/src/diffusers/models/transformers/transformer_bria_fibo2.py index 7f5bec18aead..4c93a96231d1 100644 --- a/src/diffusers/models/transformers/transformer_bria_fibo2.py +++ b/src/diffusers/models/transformers/transformer_bria_fibo2.py @@ -223,13 +223,11 @@ def __call__( query = attn.norm_q(query) key = attn.norm_k(key) - # Apply RoPE + # Apply RoPE in float32, then cast back to the input's dtype 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) + x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2)) + x_out = torch.view_as_real(x * freqs_cis.unsqueeze(2)).flatten(3) + return x_out.type_as(x_in) if freqs_cis is not None: query = apply_rotary_emb(query, freqs_cis) From eee5acf0db9153728fee257abfceaccc0aa31eb9 Mon Sep 17 00:00:00 2001 From: muhammadali-afridi Date: Fri, 16 Oct 2026 08:19:09 -0700 Subject: [PATCH 3/4] Mark the full-graph compile tests as expected failures instead of skipping 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_models_transformer_bria_fibo2.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/models/transformers/test_models_transformer_bria_fibo2.py b/tests/models/transformers/test_models_transformer_bria_fibo2.py index cfb6998dee40..1dd343fd5c0f 100644 --- a/tests/models/transformers/test_models_transformer_bria_fibo2.py +++ b/tests/models/transformers/test_models_transformer_bria_fibo2.py @@ -156,23 +156,29 @@ class TestBriaFibo2TransformerTorchCompile(BriaFibo2TransformerTesterConfig, Tor def different_shapes_for_compilation(self): return [(4, 4), (4, 8), (8, 8)] - @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." + @pytest.mark.xfail( + reason=( + "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." + ), + raises=torch._dynamo.exc.Unsupported, + strict=True, ) def test_torch_compile_recompilation_and_graph_break(self): - pass + super().test_torch_compile_recompilation_and_graph_break() def test_torch_compile_repeated_blocks(self): # BriaFibo2TransformerBlock runs three ways: in the noise refiner (one timestep for every token), in the main # blocks (per-token timesteps) and in the injection blocks (plus cross-attention), so it compiles three times super().test_torch_compile_repeated_blocks(recompile_limit=3) - @pytest.mark.skip( - "RopeEmbedder, copied from Z-Image, cannot be traced in a full graph (see above). Z-Image skips this test too." + @pytest.mark.xfail( + reason="RopeEmbedder, copied from Z-Image, cannot be traced in a full graph (see above).", + raises=torch._dynamo.exc.Unsupported, + strict=True, ) def test_compile_on_different_shapes(self): - pass + super().test_compile_on_different_shapes() class TestBriaFibo2TransformerTraining(BriaFibo2TransformerTesterConfig, TrainingTesterMixin): From 60f95274bf07ed705d675eb09d99b6ae125b601e Mon Sep 17 00:00:00 2001 From: muhammadali-afridi Date: Fri, 16 Oct 2026 08:20:48 -0700 Subject: [PATCH 4/4] Test the pipeline through its public call only 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. --- .../bria_fibo2/test_pipeline_bria_fibo2.py | 131 ++---------------- 1 file changed, 14 insertions(+), 117 deletions(-) diff --git a/tests/pipelines/bria_fibo2/test_pipeline_bria_fibo2.py b/tests/pipelines/bria_fibo2/test_pipeline_bria_fibo2.py index df28d3583434..d08a32bb3961 100644 --- a/tests/pipelines/bria_fibo2/test_pipeline_bria_fibo2.py +++ b/tests/pipelines/bria_fibo2/test_pipeline_bria_fibo2.py @@ -12,6 +12,7 @@ BriaFibo2Transformer2DModel, FlowMatchEulerDiscreteScheduler, ) +from diffusers.pipelines.bria_fibo2.text_images import draw_text_image from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin @@ -170,50 +171,6 @@ def test_edit_several_images(self): assert edited.shape == (1, *self.output_shape) - 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 like ", - } - - 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 like "}}' - ) - - def test_edit_prompt_caption_forms(self): - # fibo-2 was also trained to edit with a short description, as a string, and with a short description and the - # caption's text_render list, in that order, the list kept even when empty - pipe = self.get_pipeline() - marker = "<|vision_start|><|image_pad|><|vision_end|>" - description = "A woman in a red coat on a snowy street." - instruction = "Make the coat green." - short_json = {"short_description": description, "text_render": []} - - short = json.dumps({"structured_caption": description, "edit_instruction": instruction}) - assert pipe._get_edit_prompt(short, 1) == ( - f'{{"image":"{marker}","structured_caption":"{description}","edit_instruction":"{instruction}"}}' - ) - for prompt in [ - {**short_json, "edit_instruction": instruction}, - {"text_render": [], "short_description": description, "edit_instruction": instruction}, - {"structured_caption": short_json, "edit_instruction": instruction}, - ]: - assert pipe._get_edit_prompt(json.dumps(prompt), 1) == ( - f'{{"image":"{marker}","structured_caption":{{"short_description":"{description}","text_render":[]}},' - f'"edit_instruction":"{instruction}"}}' - ) - # An empty description is no caption, as in training - no_caption = json.dumps({"structured_caption": "", "edit_instruction": instruction}) - assert pipe._get_edit_prompt(no_caption, 1) == f'{{"image":"{marker}","edit_instruction":"{instruction}"}}' - def get_text_image(self, width=24, height=8): return PIL.Image.new("RGB", (width, height), "white") @@ -225,44 +182,10 @@ def record_text_lengths(self, pipe): ) return lengths - def test_text_prompt(self): - # Each text_render entry with text gets the marker of its image, or an empty one where the image isn't shown - pipe = self.get_pipeline() - marker = "<|vision_start|><|image_pad|><|vision_end|>" - prompt = json.dumps( - { - "short_description": "a shop sign", - "text_render": [{"text": "OPEN"}, {"text": ""}, {"location": "door"}, {"text": "7 AM"}], - } - ) - entries = '[{{"text":"OPEN","image_tokens":"{0}"}},{{"text":""}},{{"location":"door"}},{{"text":"7 AM","image_tokens":"{0}"}}]' - - assert pipe._get_text_prompt(prompt, marker) == ( - f'{{"short_description":"a shop sign","text_render":{entries.format(marker)}}}' - ) - assert ( - pipe._get_text_prompt(prompt, "") - == f'{{"short_description":"a shop sign","text_render":{entries.format("")}}}' - ) - - def test_edit_prompt_with_text_images(self): - # The caption is cleaned first, then each text entry is marked: its image follows the image to edit - pipe = self.get_pipeline() - marker = "<|vision_start|><|image_pad|><|vision_end|>" - caption = { - "short_description": "a cafe window", - "text_render": [{"text": "OPEN DAILY", "font": ""}], - "edit_instruction": "Turn it into a flat illustration", - } - - assert pipe._get_edit_prompt(json.dumps(caption), 1, text_images=True) == ( - f'{{"image":"{marker}","structured_caption":{{"short_description":"a cafe window","text_render":' - f'[{{"text":"OPEN DAILY","image_tokens":"{marker}"}}]}},"edit_instruction":"Turn it into a flat illustration"}}' - ) - def test_generation_prompt_should_take_the_trained_form(self): - # Empty values, NaN and the score fields dropped, serialized compactly. A short description and a caption of a - # short description and its text_render list keep their own trained forms + # Empty values, NaN and the score fields are dropped and the caption serialized compactly, so a pretty-printed + # caption with them renders as its clean compact form. A short description and its text_render list keep + # their trained key order pipe = self.get_pipeline() pretty = ( '{\n "short_description": "a red bicycle",\n "objects": [],\n "context": NaN,\n' @@ -270,15 +193,13 @@ def test_generation_prompt_should_take_the_trained_form(self): ) clean = '{"short_description":"a red bicycle","aesthetics":{"mood":"calm"}}' - assert pipe._get_generation_prompt(pretty) == clean - assert pipe._get_generation_prompt("A red bicycle by a door.") == "A red bicycle by a door." - assert ( - pipe._get_generation_prompt('{"text_render": [], "short_description": "a sign"}') - == '{"short_description":"a sign","text_render":[]}' - ) - from_pretty = pipe(**{**self.get_dummy_inputs(), "prompt": pretty}).images - from_clean = pipe(**{**self.get_dummy_inputs(), "prompt": clean}).images - assert torch.equal(from_pretty, from_clean) + for given, trained in [ + (pretty, clean), + ('{"text_render": [], "short_description": "a sign"}', '{"short_description":"a sign","text_render":[]}'), + ]: + from_given = pipe(**{**self.get_dummy_inputs(), "prompt": given}).images + from_trained = pipe(**{**self.get_dummy_inputs(), "prompt": trained}).images + assert torch.equal(from_given, from_trained) def test_generation_with_text_images(self): pipe = self.get_pipeline() @@ -337,32 +258,6 @@ def test_text_images_are_checked(self): with pytest.raises(ValueError): pipe(**{**self.get_dummy_inputs(), **wrong_inputs}) - def test_draw_text_images(self): - # One image per text_render entry with text: the text in one line, black on white, split into lines of 256 - # characters - pipe = self.get_pipeline() - caption = {"text_render": [{"text": "OPEN\n DAILY"}, {"text": ""}, {"location": "door"}, {"text": "x" * 257}]} - images = pipe._draw_text_images(caption) - one_line = pipe._draw_text_images({"text_render": [{"text": "OPEN DAILY"}, {"text": "x" * 256}]}) - - assert len(images) == 2 - assert images[0].mode == "RGB" and images[0].getpixel((0, 0)) == (255, 255, 255) - assert np.asarray(images[0]).min() == 0 - assert np.array_equal(np.asarray(images[0]), np.asarray(one_line[0])) - # Lines of 256 and 1 characters, each as tall as the line of 256, and 4 pixels apart - assert images[1].size == (one_line[1].width, 2 * one_line[1].height) - - def test_text_images_read_fonts_only_when_needed(self): - # Pillow's built-in font draws ASCII on its own: only a character it lacks needs the fallback fonts - pipe = self.get_pipeline() - calls = [] - pipe._get_text_image_fonts = lambda: calls.append(True) or [] - - pipe._draw_text_images({"text_render": [{"text": "OPEN 7AM-4PM"}]}) - assert calls == [] - pipe._draw_text_images({"text_render": [{"text": "OPEN 7AM–4PM"}]}) - assert calls == [True] - @pytest.mark.parametrize("edit", [False, True], ids=["generation", "edit"]) def test_render_text_images(self, edit): # By default the pipeline draws the images it would otherwise be given as text_images @@ -375,7 +270,9 @@ def test_render_text_images(self, edit): rendered = pipe(**inputs).images inputs["generator"] = self.get_generator(0) - given = pipe(**inputs, text_images=pipe._draw_text_images(caption)).images + # The text is ASCII, which Pillow's built-in font draws without the fallback fonts + drawn = [draw_text_image(entry["text"], lambda: []) for entry in caption["text_render"]] + given = pipe(**inputs, text_images=drawn).images inputs["generator"] = self.get_generator(0) without = pipe(**inputs, render_text_images=False).images