add fibo-2 LoRA training and loading (#15084)

add fibo-2 LoRA training scripts (text-to-image + image editing)

docs/source/en/api/loaders/lora.md

@@ -40,6 +40,7 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi
 - [`LTX2LoraLoaderMixin`] provides similar functions for [Flux2](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ltx2).
 - [`MiniMaxH3LoraLoaderMixin`] provides similar functions for [MiniMax-H3](https://huggingface.co/docs/diffusers/main/en/api/pipelines/minimax_h3).
 - [`RFD4ProteinaLoraLoaderMixin`] provides similar functions for [RFD4-Proteína](https://huggingface.co/docs/diffusers/main/en/api/pipelines/rfd4_proteina).
+- [`BriaFibo2LoraLoaderMixin`] provides similar functions for [fibo-2](https://huggingface.co/docs/diffusers/main/en/api/pipelines/bria_fibo2).
 - [`LoraBaseMixin`] provides a base class with several utility methods to fuse, unfuse, unload, LoRAs and more.
 
 > [!TIP]
@@ -77,6 +78,10 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi
 
 [[autodoc]] loaders.lora_pipeline.Flux3LoraLoaderMixin
 
+## BriaFibo2LoraLoaderMixin
+
+[[autodoc]] loaders.lora_pipeline.BriaFibo2LoraLoaderMixin
+
 ## ErnieImageLoraLoaderMixin
 
 [[autodoc]] loaders.lora_pipeline.ErnieImageLoraLoaderMixin

examples/dreambooth/README_bria_fibo2.md

@@ -0,0 +1,204 @@
+# DreamBooth LoRA training for fibo-2 (Bria)
+
+[fibo-2](https://huggingface.co/briaai) is Bria's flow-matching text-to-image and image-editing model
+(`BriaFibo2Pipeline`): a `BriaFibo2Transformer2DModel` denoiser conditioned on a Qwen3-VL text encoder, with an
+`AutoencoderKLFlux2` VAE. This directory has two LoRA trainers:
+
+- `train_dreambooth_lora_bria_fibo2.py` — text-to-image.
+- `train_dreambooth_lora_bria_fibo2_img2img.py` — image editing (image-to-image).
+
+Both train a LoRA adapter on the transformer only; the text encoder and VAE are frozen.
+
+> [!NOTE]
+> fibo-2 is trained on **structured JSON captions** and renders them best, but a plain-text short description is
+> accepted too (and used as-is). Keep your training captions in the form you intend to prompt with. The trainers
+> clean captions the way the pipeline does, so a JSON caption and its cleaned, compact form train the same.
+
+## Captions
+
+fibo-2 was trained on **structured JSON captions**, so prompt and train in the same form. The trainers pass each
+caption through the pipeline's own helpers, cleaning and compacting a JSON caption exactly as inference does (empty
+values and the `aesthetic_score`/`preference_score` keys are dropped), so the model sees the same text both times.
+
+**Text-to-image.** A caption is a JSON object; include only the keys you need:
+
+```json
+{
+  "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"
+}
+```
+
+Pass one caption per image through `--caption_column` (each value a JSON string, or a plain sentence used as the
+`short_description`), or a single `--instance_prompt` for the whole set. To place literal text in the image, fibo-2
+uses a `text_render` field; the trainers here do not train that path, but the base model still renders short quoted
+text in a prompt reasonably well.
+
+**Image editing.** The edit instruction (`--caption_column`, with `--instance_prompt` as the fallback) takes either
+form the pipeline accepts, and the two give different text embeddings — train in the one you will prompt with:
+
+- a **plain instruction** — `"Make the background a snowy forest"`; or
+- a **JSON object** with the instruction under `edit_instruction` and a structured caption of the *result*:
+  `{"short_description": "...", "edit_instruction": "Make the background a snowy forest"}` (the caption may instead
+  sit under a `structured_caption` key).
+
+## Running locally
+
+### Installing the dependencies
+
+Install `diffusers` from source and then the trainer requirements:
+
+```bash
+pip install -e .
+cd examples/dreambooth
+pip install -r requirements_bria_fibo2.txt
+```
+
+And initialize an [🤗 Accelerate](https://github.com/huggingface/accelerate/) environment with:
+
+```bash
+accelerate config
+```
+
+Or for a default accelerate configuration without answering questions about your environment:
+
+```bash
+accelerate config default
+```
+
+fibo-2's Qwen3-VL text encoder needs `transformers>=5.20.1`.
+
+## Text-to-image
+
+Train a LoRA on a folder of instance images (one trigger prompt for all of them):
+
+```bash
+export MODEL_NAME="briaai/fibo-2"
+export INSTANCE_DIR="dog"
+export OUTPUT_DIR="trained-fibo2-lora"
+
+accelerate launch train_dreambooth_lora_bria_fibo2.py \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --instance_data_dir=$INSTANCE_DIR \
+  --output_dir=$OUTPUT_DIR \
+  --mixed_precision="bf16" \
+  --instance_prompt="a photo of sks dog" \
+  --resolution=1024 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=4 \
+  --learning_rate=1e-4 \
+  --lr_scheduler="constant" \
+  --lr_warmup_steps=0 \
+  --max_train_steps=1000 \
+  --validation_prompt="A photo of sks dog in a bucket" \
+  --validation_epochs=25 \
+  --seed=0 \
+  --push_to_hub
+```
+
+To train on a dataset with per-image captions, pass `--dataset_name` and `--caption_column` (the captions can be
+structured JSON captions or plain short descriptions) instead of `--instance_data_dir`:
+
+```bash
+accelerate launch train_dreambooth_lora_bria_fibo2.py \
+  --pretrained_model_name_or_path="briaai/fibo-2" \
+  --dataset_name="your-dataset" \
+  --caption_column="text" \
+  --instance_prompt="in the style of TOK" \
+  --output_dir="trained-fibo2-lora" \
+  --mixed_precision="bf16" \
+  --resolution=1024 \
+  --train_batch_size=1 \
+  --learning_rate=1e-4 \
+  --max_train_steps=1000
+```
+
+### Inference
+
+```python
+from diffusers import BriaFibo2Pipeline
+import torch
+
+pipeline = BriaFibo2Pipeline.from_pretrained("briaai/fibo-2", torch_dtype=torch.bfloat16).to("cuda")
+pipeline.load_lora_weights("trained-fibo2-lora", weight_name="pytorch_lora_weights.safetensors")
+image = pipeline("A photo of sks dog in a bucket").images[0]
+image.save("fibo2.png")
+```
+
+## Image editing (image-to-image)
+
+The editing trainer learns from **(conditioning image, target image, edit instruction)** triplets, so it needs a
+dataset (`--dataset_name`) with a source-image column, a target-image column and a caption column holding the
+instruction. The conditioning and target images are given the same geometry so each pair stays aligned; the
+conditioning image is both shown to the Qwen3-VL text encoder (as the image to edit) and VAE-encoded into the
+transformer's clean context tokens, exactly as `BriaFibo2Pipeline` does at inference.
+
+```bash
+accelerate launch train_dreambooth_lora_bria_fibo2_img2img.py \
+  --pretrained_model_name_or_path="briaai/fibo-2" \
+  --dataset_name="your-edit-dataset" \
+  --image_column="target_image" \
+  --cond_image_column="source_image" \
+  --caption_column="instruction" \
+  --instance_prompt="make it a watercolor painting" \
+  --output_dir="trained-fibo2-edit-lora" \
+  --mixed_precision="bf16" \
+  --resolution=1024 \
+  --train_batch_size=1 \
+  --learning_rate=1e-4 \
+  --max_train_steps=2000 \
+  --validation_prompt="make it a watercolor painting" \
+  --validation_image="path/to/source.png" \
+  --validation_epochs=25
+```
+
+`--instance_prompt` is the fallback instruction used when a pair has no caption. `--caption_column` holds the edit
+instruction as plain text.
+
+> [!NOTE]
+> For editing, the conditioning image is shown to Qwen3-VL, and its vision tokens count towards
+> `--max_sequence_length` (a 768px image is up to ~576 tokens). The editing trainer defaults it to 1024 for that
+> reason; raise it further only if you pair large images with long structured captions.
+
+### Inference
+
+```python
+from diffusers import BriaFibo2Pipeline
+from diffusers.utils import load_image
+import torch
+
+pipeline = BriaFibo2Pipeline.from_pretrained("briaai/fibo-2", torch_dtype=torch.bfloat16).to("cuda")
+pipeline.load_lora_weights("trained-fibo2-edit-lora", weight_name="pytorch_lora_weights.safetensors")
+image = load_image("path/to/source.png")
+edited = pipeline("make it a watercolor painting", image=image).images[0]
+edited.save("fibo2_edit.png")
+```
+
+## Notes
+
+- **LoRA layers.** By default the adapter targets **every linear layer** of the transformer (PEFT's `"all-linear"`),
+  matching Bria's reference trainer. Pass `--lora_layers` with a comma-separated list to target a narrower set, e.g.
+  `--lora_layers="to_q,to_k,to_v,to_out.0,w1,w2,w3"` for attention and feed-forward only, or
+  `--lora_layers="to_q,to_k,to_v,to_out.0"` for attention only.
+- **Aspect-ratio bucketing.** Pass `--use_aspect_ratio_buckets` to bucket images by aspect ratio (optionally with an
+  explicit `--aspect_ratio_buckets` list), instead of resizing everything to a square `--resolution`.
+- **Caption dropout.** `--caption_dropout 0.1` trains a fraction of samples unconditionally, which improves
+  classifier-free guidance. For editing, the conditioning image is always kept; only the instruction is dropped.
+- **Memory.** Use `--cache_latents` to pre-encode images, `--gradient_checkpointing`, `--use_8bit_adam`, and
+  `--offload` to move the VAE and text encoder to CPU when idle. `--bnb_quantization_config_path` enables
+  bitsandbytes-quantized training of the transformer, and `--do_fp8_training` enables FP8 (TorchAO) training.
+- **Prodigy.** Pass `--optimizer="prodigy"` (with `--learning_rate=1.0`) to use the Prodigy optimizer.
+
+> [!NOTE]
+> These scripts are provided as a starting point; tune hyperparameters (rank, learning rate, steps) for your data.
+
+## More resources
+
+Bria publishes a single-GPU reference LoRA trainer, the full caption schema, and their recommended settings in the
+model's [`fine_tuning` folder](https://huggingface.co/briaai/Fibo-2.0-dev/tree/main/fine_tuning). Their suggested
+starting point is rank 64 (alpha = rank), learning rate 1e-4 constant, 1024px, gradient accumulation 8, and LoRA on
+every linear layer. The trainers here also default to every linear layer; consider raising `--rank`/`--lora_alpha`
+(the defaults are small) and lowering the target set with `--lora_layers` if you want a lighter adapter.

examples/dreambooth/requirements_bria_fibo2.txt

@@ -0,0 +1,11 @@
+accelerate>=0.31.0
+torchvision
+transformers>=5.20.1
+ftfy
+tensorboard
+Jinja2
+peft>=0.11.1
+sentencepiece
+bitsandbytes
+prodigyopt
+datasets>=3.1.0

examples/dreambooth/test_dreambooth_lora_bria_fibo2.py

@@ -0,0 +1,190 @@
+# 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 logging
+import os
+import sys
+import tempfile
+
+import safetensors
+
+
+sys.path.append("..")
+from test_examples_utils import ExamplesTestsAccelerate, run_command  # noqa: E402
+
+
+logging.basicConfig(level=logging.DEBUG)
+
+logger = logging.getLogger()
+stream_handler = logging.StreamHandler(sys.stdout)
+logger.addHandler(stream_handler)
+
+# Tiny randomly-initialized fibo-2 repo used for fast CI. The env var is an escape hatch for running against a local
+# build. The tiny model scales the components down like tests/pipelines/bria_fibo2 does, keeping a 36-layer Qwen3-VL
+# text encoder so the trainer's real text-encoder wiring works unchanged.
+FIBO2_TINY_REPO = os.environ.get("FIBO2_TINY_REPO", "hf-internal-testing/tiny-bria-fibo2")
+
+
+class DreamBoothLoRABriaFibo2(ExamplesTestsAccelerate):
+    instance_data_dir = "docs/source/en/imgs"
+    instance_prompt = "dog"
+    pretrained_model_name_or_path = FIBO2_TINY_REPO
+    script_path = "examples/dreambooth/train_dreambooth_lora_bria_fibo2.py"
+    transformer_layer_type = "layers.0.attention.to_k"
+
+    def test_dreambooth_lora_bria_fibo2(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_args = f"""
+                {self.script_path}
+                --pretrained_model_name_or_path {self.pretrained_model_name_or_path}
+                --instance_data_dir {self.instance_data_dir}
+                --instance_prompt {self.instance_prompt}
+                --resolution 64
+                --train_batch_size 1
+                --gradient_accumulation_steps 1
+                --max_train_steps 2
+                --learning_rate 5.0e-04
+                --scale_lr
+                --lr_scheduler constant
+                --lr_warmup_steps 0
+                --max_sequence_length 16
+                --output_dir {tmpdir}
+                """.split()
+
+            run_command(self._launch_args + test_args)
+            # save_pretrained smoke test
+            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
+
+            # make sure the state_dict has the correct naming in the parameters.
+            lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
+            is_lora = all("lora" in k for k in lora_state_dict.keys())
+            self.assertTrue(is_lora)
+
+            # when not training the text encoder, all the parameters in the state dict should start
+            # with `"transformer"` in their names.
+            starts_with_transformer = all(key.startswith("transformer") for key in lora_state_dict.keys())
+            self.assertTrue(starts_with_transformer)
+
+    def test_dreambooth_lora_bria_fibo2_latent_caching(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_args = f"""
+                {self.script_path}
+                --pretrained_model_name_or_path {self.pretrained_model_name_or_path}
+                --instance_data_dir {self.instance_data_dir}
+                --instance_prompt {self.instance_prompt}
+                --resolution 64
+                --train_batch_size 1
+                --gradient_accumulation_steps 1
+                --max_train_steps 2
+                --cache_latents
+                --learning_rate 5.0e-04
+                --scale_lr
+                --lr_scheduler constant
+                --lr_warmup_steps 0
+                --max_sequence_length 16
+                --output_dir {tmpdir}
+                """.split()
+
+            run_command(self._launch_args + test_args)
+            # save_pretrained smoke test
+            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
+
+            # make sure the state_dict has the correct naming in the parameters.
+            lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
+            is_lora = all("lora" in k for k in lora_state_dict.keys())
+            self.assertTrue(is_lora)
+
+            starts_with_transformer = all(key.startswith("transformer") for key in lora_state_dict.keys())
+            self.assertTrue(starts_with_transformer)
+
+    def test_dreambooth_lora_bria_fibo2_layers(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_args = f"""
+                {self.script_path}
+                --pretrained_model_name_or_path {self.pretrained_model_name_or_path}
+                --instance_data_dir {self.instance_data_dir}
+                --instance_prompt {self.instance_prompt}
+                --resolution 64
+                --train_batch_size 1
+                --gradient_accumulation_steps 1
+                --max_train_steps 2
+                --cache_latents
+                --learning_rate 5.0e-04
+                --scale_lr
+                --lora_layers {self.transformer_layer_type}
+                --lr_scheduler constant
+                --lr_warmup_steps 0
+                --max_sequence_length 16
+                --output_dir {tmpdir}
+                """.split()
+
+            run_command(self._launch_args + test_args)
+            # save_pretrained smoke test
+            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
+
+            # make sure the state_dict has the correct naming in the parameters.
+            lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
+            is_lora = all("lora" in k for k in lora_state_dict.keys())
+            self.assertTrue(is_lora)
+
+            # only the layers named by `--lora_layers` should be trained, so every key sits under that path
+            starts_with_transformer = all(
+                key.startswith(f"transformer.{self.transformer_layer_type}") for key in lora_state_dict.keys()
+            )
+            self.assertTrue(starts_with_transformer)
+
+    def test_dreambooth_lora_bria_fibo2_checkpointing_checkpoints_total_limit(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_args = f"""
+            {self.script_path}
+            --pretrained_model_name_or_path={self.pretrained_model_name_or_path}
+            --instance_data_dir={self.instance_data_dir}
+            --output_dir={tmpdir}
+            --instance_prompt={self.instance_prompt}
+            --resolution=64
+            --train_batch_size=1
+            --gradient_accumulation_steps=1
+            --max_train_steps=6
+            --checkpoints_total_limit=2
+            --max_sequence_length 16
+            --checkpointing_steps=2
+            """.split()
+
+            run_command(self._launch_args + test_args)
+
+            self.assertEqual(
+                {x for x in os.listdir(tmpdir) if "checkpoint" in x},
+                {"checkpoint-4", "checkpoint-6"},
+            )
+
+    def test_dreambooth_lora_bria_fibo2_checkpointing_checkpoints_total_limit_removes_multiple_checkpoints(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_args = f"""
+            {self.script_path}
+            --pretrained_model_name_or_path={self.pretrained_model_name_or_path}
+            --instance_data_dir={self.instance_data_dir}
+            --output_dir={tmpdir}
+            --instance_prompt={self.instance_prompt}
+            --resolution=64
+            --train_batch_size=1
+            --gradient_accumulation_steps=1
+            --max_train_steps=4
+            --checkpointing_steps=2
+            --max_sequence_length 16
+            """.split()
+
+            run_command(self._launch_args + test_args)
+
+            self.assertEqual({x for x in os.listdir(tmpdir) if "checkpoint" in x}, {"checkpoint-2", "checkpoint-4"})

examples/dreambooth/test_dreambooth_lora_bria_fibo2_img2img.py

@@ -0,0 +1,145 @@
+# 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 logging
+import os
+import sys
+import tempfile
+
+import safetensors
+
+
+sys.path.append("..")
+from test_examples_utils import ExamplesTestsAccelerate, run_command  # noqa: E402
+
+
+logging.basicConfig(level=logging.DEBUG)
+
+logger = logging.getLogger()
+stream_handler = logging.StreamHandler(sys.stdout)
+logger.addHandler(stream_handler)
+
+# Tiny randomly-initialized fibo-2 repo used for fast CI. The env var is an escape hatch for running against a local
+# build. The tiny model scales the components down like tests/pipelines/bria_fibo2 does, keeping a 36-layer Qwen3-VL
+# text encoder so the trainer's real text-encoder wiring works unchanged.
+FIBO2_TINY_REPO = os.environ.get("FIBO2_TINY_REPO", "hf-internal-testing/tiny-bria-fibo2")
+
+
+class DreamBoothLoRABriaFibo2Img2Img(ExamplesTestsAccelerate):
+    # image editing (image-to-image) training needs a paired conditioning image, so it reads a dataset with an
+    # image + conditioning_image column rather than a plain folder
+    dataset_name = "hf-internal-testing/fill10"
+    image_column = "image"
+    cond_image_column = "conditioning_image"
+    caption_column = "text"
+    instance_prompt = "dog"
+    pretrained_model_name_or_path = FIBO2_TINY_REPO
+    script_path = "examples/dreambooth/train_dreambooth_lora_bria_fibo2_img2img.py"
+
+    def test_dreambooth_lora_bria_fibo2_img2img(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_args = f"""
+                {self.script_path}
+                --pretrained_model_name_or_path {self.pretrained_model_name_or_path}
+                --dataset_name {self.dataset_name}
+                --image_column {self.image_column}
+                --cond_image_column {self.cond_image_column}
+                --caption_column {self.caption_column}
+                --instance_prompt {self.instance_prompt}
+                --resolution 64
+                --train_batch_size 1
+                --gradient_accumulation_steps 1
+                --max_train_steps 2
+                --learning_rate 5.0e-04
+                --lr_scheduler constant
+                --lr_warmup_steps 0
+                --max_sequence_length 256
+                --output_dir {tmpdir}
+                """.split()
+
+            run_command(self._launch_args + test_args)
+            # save_pretrained smoke test
+            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
+
+            # make sure the state_dict has the correct naming in the parameters.
+            lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
+            is_lora = all("lora" in k for k in lora_state_dict.keys())
+            self.assertTrue(is_lora)
+
+            # when not training the text encoder, all the parameters in the state dict should start
+            # with `"transformer"` in their names.
+            starts_with_transformer = all(key.startswith("transformer") for key in lora_state_dict.keys())
+            self.assertTrue(starts_with_transformer)
+
+    def test_dreambooth_lora_bria_fibo2_img2img_latent_caching(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_args = f"""
+                {self.script_path}
+                --pretrained_model_name_or_path {self.pretrained_model_name_or_path}
+                --dataset_name {self.dataset_name}
+                --image_column {self.image_column}
+                --cond_image_column {self.cond_image_column}
+                --caption_column {self.caption_column}
+                --instance_prompt {self.instance_prompt}
+                --resolution 64
+                --train_batch_size 1
+                --gradient_accumulation_steps 1
+                --max_train_steps 2
+                --cache_latents
+                --learning_rate 5.0e-04
+                --lr_scheduler constant
+                --lr_warmup_steps 0
+                --max_sequence_length 256
+                --output_dir {tmpdir}
+                """.split()
+
+            run_command(self._launch_args + test_args)
+            # save_pretrained smoke test
+            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
+
+            # make sure the state_dict has the correct naming in the parameters.
+            lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
+            is_lora = all("lora" in k for k in lora_state_dict.keys())
+            self.assertTrue(is_lora)
+
+            starts_with_transformer = all(key.startswith("transformer") for key in lora_state_dict.keys())
+            self.assertTrue(starts_with_transformer)
+
+    def test_dreambooth_lora_bria_fibo2_img2img_checkpointing_checkpoints_total_limit(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_args = f"""
+            {self.script_path}
+            --pretrained_model_name_or_path={self.pretrained_model_name_or_path}
+            --dataset_name={self.dataset_name}
+            --image_column={self.image_column}
+            --cond_image_column={self.cond_image_column}
+            --caption_column={self.caption_column}
+            --output_dir={tmpdir}
+            --instance_prompt={self.instance_prompt}
+            --resolution=64
+            --train_batch_size=1
+            --gradient_accumulation_steps=1
+            --max_train_steps=6
+            --checkpoints_total_limit=2
+            --max_sequence_length 256
+            --checkpointing_steps=2
+            """.split()
+
+            run_command(self._launch_args + test_args)
+
+            self.assertEqual(
+                {x for x in os.listdir(tmpdir) if "checkpoint" in x},
+                {"checkpoint-4", "checkpoint-6"},
+            )

src/diffusers/loaders/__init__.py

@@ -43,6 +43,7 @@ if is_torch_available():
             "ZImageLoraLoaderMixin",
             "Flux2LoraLoaderMixin",
             "Flux3LoraLoaderMixin",
+            "BriaFibo2LoraLoaderMixin",
             "Ideogram4LoraLoaderMixin",
             "ErnieImageLoraLoaderMixin",
             "CosmosLoraLoaderMixin",
@@ -80,6 +81,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
                 AmusedLoraLoaderMixin,
                 AnimaLoraLoaderMixin,
                 AuraFlowLoraLoaderMixin,
+                BriaFibo2LoraLoaderMixin,
                 CogVideoXLoraLoaderMixin,
                 CogView4LoraLoaderMixin,
                 CosmosLoraLoaderMixin,

src/diffusers/loaders/lora_pipeline.py

@@ -6163,6 +6163,195 @@ class Flux3LoraLoaderMixin(LoraBaseMixin):
         super().unfuse_lora(components=components, **kwargs)
 
 
+class BriaFibo2LoraLoaderMixin(LoraBaseMixin):
+    r"""
+    Load LoRA layers into [`BriaFibo2Transformer2DModel`]. Specific to [`BriaFibo2Pipeline`].
+    """
+
+    _lora_loadable_modules = ["transformer"]
+    transformer_name = TRANSFORMER_NAME
+
+    @classmethod
+    @validate_hf_hub_args
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict
+    def lora_state_dict(
+        cls,
+        pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+        **kwargs,
+    ):
+        r"""
+        See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details.
+        """
+        # Load the main state dict first which has the LoRA layers for either of
+        # transformer and text encoder or both.
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        proxies = kwargs.pop("proxies", None)
+        local_files_only = kwargs.pop("local_files_only", None)
+        token = kwargs.pop("token", None)
+        revision = kwargs.pop("revision", None)
+        subfolder = kwargs.pop("subfolder", None)
+        weight_name = kwargs.pop("weight_name", None)
+        use_safetensors = kwargs.pop("use_safetensors", None)
+        return_lora_metadata = kwargs.pop("return_lora_metadata", False)
+
+        allow_pickle = False
+        if use_safetensors is None:
+            use_safetensors = True
+            allow_pickle = True
+
+        user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
+
+        state_dict, metadata = _fetch_state_dict(
+            pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
+            weight_name=weight_name,
+            use_safetensors=use_safetensors,
+            local_files_only=local_files_only,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            proxies=proxies,
+            token=token,
+            revision=revision,
+            subfolder=subfolder,
+            user_agent=user_agent,
+            allow_pickle=allow_pickle,
+        )
+
+        is_dora_scale_present = any("dora_scale" in k for k in state_dict)
+        if is_dora_scale_present:
+            warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
+            logger.warning(warn_msg)
+            state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
+
+        out = (state_dict, metadata) if return_lora_metadata else state_dict
+        return out
+
+    @require_peft_backend
+    # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
+    def load_lora_weights(
+        self,
+        pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+        adapter_name: str | None = None,
+        hotswap: bool = False,
+        **kwargs,
+    ):
+        """
+        See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details.
+        """
+        low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
+
+        # if a dict is passed, copy it instead of modifying it inplace
+        if isinstance(pretrained_model_name_or_path_or_dict, dict):
+            pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
+
+        # First, ensure that the checkpoint is a compatible one and can be successfully loaded.
+        kwargs["return_lora_metadata"] = True
+        state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
+
+        is_correct_format = all("lora" in key for key in state_dict.keys())
+        if not is_correct_format:
+            raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.")
+
+        self.load_lora_into_transformer(
+            state_dict,
+            transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
+            adapter_name=adapter_name,
+            metadata=metadata,
+            _pipeline=self,
+            low_cpu_mem_usage=low_cpu_mem_usage,
+            hotswap=hotswap,
+        )
+
+    @classmethod
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer
+    def load_lora_into_transformer(
+        cls,
+        state_dict,
+        transformer,
+        adapter_name=None,
+        _pipeline=None,
+        low_cpu_mem_usage=False,
+        hotswap: bool = False,
+        metadata=None,
+    ):
+        """
+        See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details.
+        """
+        # Load the layers corresponding to transformer.
+        logger.info(f"Loading {cls.transformer_name}.")
+        transformer.load_lora_adapter(
+            state_dict,
+            network_alphas=None,
+            adapter_name=adapter_name,
+            metadata=metadata,
+            _pipeline=_pipeline,
+            low_cpu_mem_usage=low_cpu_mem_usage,
+            hotswap=hotswap,
+        )
+
+    @classmethod
+    # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
+    def save_lora_weights(
+        cls,
+        save_directory: str | os.PathLike,
+        transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
+        is_main_process: bool = True,
+        weight_name: str = None,
+        save_function: Callable = None,
+        safe_serialization: bool = True,
+        transformer_lora_adapter_metadata: dict | None = None,
+    ):
+        r"""
+        See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
+        """
+        lora_layers = {}
+        lora_metadata = {}
+
+        if transformer_lora_layers:
+            lora_layers[cls.transformer_name] = transformer_lora_layers
+            lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata
+
+        if not lora_layers:
+            raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.")
+
+        cls._save_lora_weights(
+            save_directory=save_directory,
+            lora_layers=lora_layers,
+            lora_metadata=lora_metadata,
+            is_main_process=is_main_process,
+            weight_name=weight_name,
+            save_function=save_function,
+            safe_serialization=safe_serialization,
+        )
+
+    # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
+    def fuse_lora(
+        self,
+        components: list[str] = ["transformer"],
+        lora_scale: float = 1.0,
+        safe_fusing: bool = False,
+        adapter_names: list[str] | None = None,
+        **kwargs,
+    ):
+        r"""
+        See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details.
+        """
+        super().fuse_lora(
+            components=components,
+            lora_scale=lora_scale,
+            safe_fusing=safe_fusing,
+            adapter_names=adapter_names,
+            **kwargs,
+        )
+
+    # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
+    def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
+        r"""
+        See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
+        """
+        super().unfuse_lora(components=components, **kwargs)
+
+
 class Ideogram4LoraLoaderMixin(LoraBaseMixin):
     r"""
     Load LoRA layers into [`Ideogram4Transformer2DModel`]. Specific to [`Ideogram4Pipeline`].

src/diffusers/pipelines/bria_fibo2/pipeline_bria_fibo2.py

@@ -22,6 +22,7 @@ import PIL.Image
 import torch
 from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
 
+from ...loaders import BriaFibo2LoraLoaderMixin
 from ...models import AutoencoderKLFlux2, BriaFibo2Transformer2DModel
 from ...schedulers import FlowMatchEulerDiscreteScheduler
 from ...utils import is_torch_xla_available, logging, replace_example_docstring
@@ -167,7 +168,7 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-class BriaFibo2Pipeline(DiffusionPipeline):
+class BriaFibo2Pipeline(DiffusionPipeline, BriaFibo2LoraLoaderMixin):
     r"""
     Pipeline for text-to-image generation and image editing with Fibo 2.0.