diff --git a/docs/source/en/api/loaders/lora.md b/docs/source/en/api/loaders/lora.md
index 4c8eb20b299a..d58e446f3541 100644
--- a/docs/source/en/api/loaders/lora.md
+++ b/docs/source/en/api/loaders/lora.md
@@ -39,6 +39,7 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi
- [`ErnieImageLoraLoaderMixin`] provides similar functions for [Ernie-Image](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ernie_image).
- [`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).
+- [`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]
@@ -72,6 +73,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
diff --git a/examples/dreambooth/README_bria_fibo2.md b/examples/dreambooth/README_bria_fibo2.md
new file mode 100644
index 000000000000..0b334cebbfa4
--- /dev/null
+++ b/examples/dreambooth/README_bria_fibo2.md
@@ -0,0 +1,164 @@
+# 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.
+
+## 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` (the default is 512). Raise it if you edit large images so the image tokens are not
+> truncated.
+
+### 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 attention (`to_q`, `to_k`, `to_v`, `to_out.0`) and the SwiGLU
+ feed-forward projections (`w1`, `w2`, `w3`) across every block. Pass `--lora_layers` with a comma-separated list
+ to target a different set (e.g. `--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.
diff --git a/examples/dreambooth/requirements_bria_fibo2.txt b/examples/dreambooth/requirements_bria_fibo2.txt
new file mode 100644
index 000000000000..9810bcd5b7aa
--- /dev/null
+++ b/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
diff --git a/examples/dreambooth/test_dreambooth_lora_bria_fibo2.py b/examples/dreambooth/test_dreambooth_lora_bria_fibo2.py
new file mode 100644
index 000000000000..184a9bee4b5f
--- /dev/null
+++ b/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"})
diff --git a/examples/dreambooth/test_dreambooth_lora_bria_fibo2_img2img.py b/examples/dreambooth/test_dreambooth_lora_bria_fibo2_img2img.py
new file mode 100644
index 000000000000..27081198c1bd
--- /dev/null
+++ b/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"},
+ )
diff --git a/examples/dreambooth/train_dreambooth_lora_bria_fibo2.py b/examples/dreambooth/train_dreambooth_lora_bria_fibo2.py
new file mode 100644
index 000000000000..b9abc04b6af6
--- /dev/null
+++ b/examples/dreambooth/train_dreambooth_lora_bria_fibo2.py
@@ -0,0 +1,2141 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. 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.
+
+# /// script
+# dependencies = [
+# "diffusers @ git+https://github.com/huggingface/diffusers.git",
+# "torch>=2.0.0",
+# "accelerate>=0.31.0",
+# "transformers>=5.20.1",
+# "ftfy",
+# "tensorboard",
+# "Jinja2",
+# "peft>=0.11.1",
+# "sentencepiece",
+# "torchvision",
+# "datasets",
+# "bitsandbytes",
+# "prodigyopt",
+# ]
+# ///
+
+import argparse
+import copy
+import itertools
+import json
+import logging
+import math
+import os
+import random
+import shutil
+import warnings
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed
+from huggingface_hub import create_repo, upload_folder
+from huggingface_hub.utils import insecure_hashlib
+from peft import LoraConfig, prepare_model_for_kbit_training, set_peft_model_state_dict
+from peft.utils import get_peft_model_state_dict
+from PIL import Image
+from PIL.ImageOps import exif_transpose
+from torch.utils.data import Dataset
+from torch.utils.data.sampler import BatchSampler
+from torchvision import transforms
+from torchvision.transforms import functional as TF
+from tqdm.auto import tqdm
+from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
+
+import diffusers
+from diffusers import (
+ AutoencoderKLFlux2,
+ BitsAndBytesConfig,
+ BriaFibo2Pipeline,
+ BriaFibo2Transformer2DModel,
+ FlowMatchEulerDiscreteScheduler,
+)
+from diffusers.loaders import BriaFibo2LoraLoaderMixin
+from diffusers.optimization import get_scheduler
+from diffusers.training_utils import (
+ _collate_lora_metadata,
+ _to_cpu_contiguous,
+ cast_training_params,
+ compute_density_for_timestep_sampling,
+ compute_loss_weighting_for_sd3,
+ find_nearest_bucket,
+ free_memory,
+ generate_aspect_ratio_buckets,
+ get_fsdp_kwargs_from_accelerator,
+ offload_models,
+ parse_buckets_string,
+ wrap_with_fsdp,
+)
+from diffusers.utils import (
+ check_min_version,
+ convert_unet_state_dict_to_peft,
+ is_wandb_available,
+)
+from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
+from diffusers.utils.import_utils import is_torch_npu_available
+from diffusers.utils.torch_utils import is_compiled_module
+
+
+if getattr(torch, "distributed", None) is not None:
+ import torch.distributed as dist
+
+if is_wandb_available():
+ import wandb
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.40.0.dev0")
+
+logger = get_logger(__name__)
+
+# fibo-2 wraps the prompt in Qwen3-VL's chat template before encoding, as the pipeline does.
+PROMPT_TEMPLATE_ENCODE = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+
+# The Qwen3-VL layers behind the transformer's six text inputs: the Perceiver's bundle first, then one per
+# injection block. Each bundle concatenates three layers, so every text input is 3 * hidden_dim wide. This is model
+# wiring (it must match `BriaFibo2Pipeline.text_encoder_out_layers`), not a tuning knob, so it lives here.
+TEXT_ENCODER_OUT_LAYERS = ((9, 20, 31), (5, 16, 27), (7, 18, 29), (9, 20, 31), (11, 22, 33), (13, 24, 35))
+
+
+# --- fibo-2 latent helpers (identical to the pipeline's) ---
+def patchify_latents(latents):
+ """(B, C, H, W) -> (B, C*4, H//2, W//2), the 2x2 patch order the VAE BatchNorm statistics live in."""
+ b, c, h, w = latents.shape
+ latents = latents.view(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4)
+ return latents.reshape(b, c * 4, h // 2, w // 2)
+
+
+def unpatchify_latents(latents):
+ """(B, C*4, H//2, W//2) -> (B, C, H, W), the inverse of `patchify_latents`."""
+ b, c, h, w = latents.shape
+ latents = latents.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3)
+ return latents.reshape(b, c // 4, h * 2, w * 2)
+
+
+def normalize_latents(latents, bn_mean, bn_std):
+ """Normalize VAE latents with the VAE's BatchNorm statistics, exactly as `BriaFibo2Pipeline` does: the
+ statistics are stored for the 2x2-patchified layout, so patchify, normalize, then unpatchify back to the
+ `in_channels` latents the transformer takes. The result is the clean sample the flow-matching target denoises to
+ (for a conditioning image, the clean context tokens)."""
+ latents = patchify_latents(latents)
+ latents = (latents - bn_mean) / bn_std
+ return unpatchify_latents(latents)
+
+
+@torch.no_grad()
+def get_fibo2_prompt_embeds(text_encoder, processor, prompt, device, dtype, max_sequence_length):
+ """Encode prompts with the frozen Qwen3-VL encoder the way `BriaFibo2Pipeline._get_qwen_prompt_embeds` does.
+
+ Each prompt goes through Qwen3-VL's chat template, the six output bundles are each scaled to unit RMS per token
+ in float32, and every bundle concatenates its three layers. Returns `(B, 6, seq_len, 3 * hidden_dim)` embeddings
+ and their bool token mask. Prompts are encoded one at a time (no padding), so a prompt's real-token embeddings
+ are independent of any batch it is later stacked into -- Qwen3-VL masks padding, so this matches the pipeline.
+ """
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ text = [PROMPT_TEMPLATE_ENCODE.format(p) for p in prompt]
+ text_inputs = processor(
+ text=text,
+ images=None,
+ padding="longest",
+ max_length=max_sequence_length,
+ truncation=True,
+ return_tensors="pt",
+ ).to(device)
+ hidden_states = text_encoder(**text_inputs, output_hidden_states=True).hidden_states
+
+ normalized = {}
+ for layer in sorted({layer for layers in 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 TEXT_ENCODER_OUT_LAYERS]
+ return torch.stack(bundles, dim=1), text_inputs.attention_mask.bool()
+
+
+def pad_and_stack_embeds(embeds_list, mask_list, device):
+ """Right-pad a batch's per-prompt `(1, 6, seq_len, dim)` embeddings to the batch's longest prompt and stack
+ them, returning the batched embeddings and a bool mask (or `None` when nothing is padded). Right-padding and the
+ mask mirror the pipeline's `padding="longest"`; the transformer reads real token lengths from the mask to size
+ its gist tokens, so padding is free of effect."""
+ max_len = max(embeds.shape[2] for embeds in embeds_list)
+ padded_embeds, padded_masks = [], []
+ for embeds, mask in zip(embeds_list, mask_list):
+ seq_len = embeds.shape[2]
+ if seq_len < max_len:
+ embed_pad = embeds.new_zeros(embeds.shape[0], embeds.shape[1], max_len - seq_len, embeds.shape[3])
+ embeds = torch.cat([embeds, embed_pad], dim=2)
+ mask_pad = mask.new_zeros(mask.shape[0], max_len - seq_len)
+ mask = torch.cat([mask, mask_pad], dim=1)
+ padded_embeds.append(embeds)
+ padded_masks.append(mask)
+ embeds = torch.cat(padded_embeds, dim=0).to(device)
+ mask = torch.cat(padded_masks, dim=0).to(device)
+ return embeds, (None if mask.all() else mask)
+
+
+def save_model_card(
+ repo_id: str,
+ images=None,
+ base_model: str = None,
+ instance_prompt=None,
+ validation_prompt=None,
+ repo_folder=None,
+ quant_training=None,
+):
+ widget_dict = []
+ if images is not None:
+ for i, image in enumerate(images):
+ image.save(os.path.join(repo_folder, f"image_{i}.png"))
+ widget_dict.append(
+ {"text": validation_prompt if validation_prompt else " ", "output": {"url": f"image_{i}.png"}}
+ )
+
+ model_description = f"""
+# fibo-2 DreamBooth LoRA - {repo_id}
+
+
+
+## Model description
+
+These are {repo_id} DreamBooth LoRA weights for {base_model}.
+
+The weights were trained using [DreamBooth](https://dreambooth.github.io/) with the [fibo-2 diffusers trainer](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_bria_fibo2.md).
+
+Quant training? {quant_training}
+
+## Trigger words
+
+You should use `{instance_prompt}` to trigger the image generation.
+
+fibo-2 reads structured JSON captions and renders them best; a plain-text short description also works. Keep your
+prompts in the form you trained on.
+
+## Download model
+
+[Download the *.safetensors LoRA]({repo_id}/tree/main) in the Files & versions tab.
+
+## Use it with the [๐งจ diffusers library](https://github.com/huggingface/diffusers)
+
+```py
+from diffusers import BriaFibo2Pipeline
+import torch
+
+pipeline = BriaFibo2Pipeline.from_pretrained("{base_model}", torch_dtype=torch.bfloat16).to("cuda")
+pipeline.load_lora_weights("{repo_id}", weight_name="pytorch_lora_weights.safetensors")
+image = pipeline("{validation_prompt if validation_prompt else instance_prompt}").images[0]
+```
+
+For more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters)
+
+## License
+
+Please adhere to the licensing terms as described in the model repository.
+"""
+ model_card = load_or_create_model_card(
+ repo_id_or_path=repo_id,
+ from_training=True,
+ license="other",
+ base_model=base_model,
+ prompt=instance_prompt,
+ model_description=model_description,
+ widget=widget_dict,
+ )
+ tags = [
+ "text-to-image",
+ "diffusers-training",
+ "diffusers",
+ "lora",
+ "bria",
+ "fibo-2",
+ "template:sd-lora",
+ ]
+
+ model_card = populate_model_card(model_card, tags=tags)
+ model_card.save(os.path.join(repo_folder, "README.md"))
+
+
+def log_validation(
+ pipeline,
+ args,
+ accelerator,
+ pipeline_args,
+ epoch,
+ torch_dtype,
+ is_final_validation=False,
+):
+ args.num_validation_images = args.num_validation_images if args.num_validation_images else 1
+ logger.info(
+ f"Running validation... \n Generating {args.num_validation_images} images with prompt:"
+ f" {args.validation_prompt}."
+ )
+ pipeline = pipeline.to(accelerator.device)
+ pipeline.set_progress_bar_config(disable=True)
+
+ # run inference
+ generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed is not None else None
+ autocast_ctx = torch.autocast(accelerator.device.type) if not is_final_validation else nullcontext()
+
+ images = []
+ for _ in range(args.num_validation_images):
+ with autocast_ctx:
+ image = pipeline(**pipeline_args, generator=generator).images[0]
+ images.append(image)
+
+ for tracker in accelerator.trackers:
+ phase_name = "test" if is_final_validation else "validation"
+ if tracker.name == "tensorboard":
+ np_images = np.stack([np.asarray(img) for img in images])
+ tracker.writer.add_images(phase_name, np_images, epoch, dataformats="NHWC")
+ if tracker.name == "wandb":
+ tracker.log(
+ {
+ phase_name: [
+ wandb.Image(image, caption=f"{i}: {args.validation_prompt}") for i, image in enumerate(images)
+ ]
+ }
+ )
+
+ del pipeline
+ free_memory()
+
+ return images
+
+
+def make_validation_pipeline(args, transformer, weight_dtype, device):
+ # Build a standard BriaFibo2Pipeline for validation that REUSES the in-memory trained transformer, so validation
+ # runs on exactly the weights being trained and the DiT is never loaded a second time. Only the small auxiliary
+ # components are (re)loaded (they are freed after latent caching). This mirrors how flux2_klein passes the
+ # unwrapped transformer into its validation pipeline.
+ vae = AutoencoderKLFlux2.from_pretrained(
+ args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision, torch_dtype=torch.float32
+ )
+ text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
+ args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, torch_dtype=weight_dtype
+ )
+ processor = Qwen3VLProcessor.from_pretrained(
+ args.pretrained_model_name_or_path, subfolder="processor", revision=args.revision
+ )
+ scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+ args.pretrained_model_name_or_path, subfolder="scheduler", revision=args.revision
+ )
+ pipeline = BriaFibo2Pipeline(
+ transformer=transformer,
+ scheduler=scheduler,
+ vae=vae,
+ text_encoder=text_encoder,
+ processor=processor,
+ )
+ pipeline.to(device)
+ return pipeline
+
+
+def module_filter_fn(mod: torch.nn.Module, fqn: str):
+ # don't convert the output projection
+ if fqn.endswith("final_layer.linear"):
+ return False
+ # don't convert linear modules with weight dimensions not divisible by 16
+ if isinstance(mod, torch.nn.Linear):
+ if mod.in_features % 16 != 0 or mod.out_features % 16 != 0:
+ return False
+ return True
+
+
+def parse_args(input_args=None):
+ parser = argparse.ArgumentParser(description="Simple example of a training script.")
+ parser.add_argument(
+ "--pretrained_model_name_or_path",
+ type=str,
+ default=None,
+ required=True,
+ help="Path to pretrained model or model identifier from huggingface.co/models.",
+ )
+ parser.add_argument(
+ "--revision",
+ type=str,
+ default=None,
+ required=False,
+ help="Revision of pretrained model identifier from huggingface.co/models.",
+ )
+ parser.add_argument(
+ "--bnb_quantization_config_path",
+ type=str,
+ default=None,
+ help="Quantization config in a JSON file that will be used to define the bitsandbytes quant config of the DiT.",
+ )
+ parser.add_argument(
+ "--do_fp8_training",
+ action="store_true",
+ help="if we are doing FP8 training.",
+ )
+ parser.add_argument(
+ "--variant",
+ type=str,
+ default=None,
+ help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
+ )
+ parser.add_argument(
+ "--dataset_name",
+ type=str,
+ default=None,
+ help=(
+ "The name of the Dataset (from the HuggingFace hub) containing the training data of instance images (could be your own, possibly private,"
+ " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,"
+ " or to a folder containing files that ๐ค Datasets can understand."
+ ),
+ )
+ parser.add_argument(
+ "--dataset_config_name",
+ type=str,
+ default=None,
+ help="The config of the Dataset, leave as None if there's only one config.",
+ )
+ parser.add_argument(
+ "--instance_data_dir",
+ type=str,
+ default=None,
+ help=("A folder containing the training data. "),
+ )
+
+ parser.add_argument(
+ "--cache_dir",
+ type=str,
+ default=None,
+ help="The directory where the downloaded models and datasets will be stored.",
+ )
+
+ parser.add_argument(
+ "--image_column",
+ type=str,
+ default="image",
+ help="The column of the dataset containing the target image. By "
+ "default, the standard Image Dataset maps out 'file_name' "
+ "to 'image'.",
+ )
+ parser.add_argument(
+ "--caption_column",
+ type=str,
+ default="text",
+ help="The column of the dataset containing the instance prompt for each image. fibo-2 is trained on "
+ "structured JSON captions; a plain-text short description is accepted too, and is used as it is.",
+ )
+
+ parser.add_argument("--repeats", type=int, default=1, help="How many times to repeat the training data.")
+
+ parser.add_argument(
+ "--class_data_dir",
+ type=str,
+ default=None,
+ required=False,
+ help="A folder containing the training data of class images.",
+ )
+ parser.add_argument(
+ "--instance_prompt",
+ type=str,
+ default=None,
+ required=True,
+ help="The prompt with identifier specifying the instance, e.g. 'photo of a TOK dog', 'in the style of TOK'",
+ )
+ parser.add_argument(
+ "--class_prompt",
+ type=str,
+ default=None,
+ help="The prompt to specify images in the same class as provided instance images.",
+ )
+ parser.add_argument(
+ "--max_sequence_length",
+ type=int,
+ default=512,
+ help="Cap on the number of Qwen3-VL prompt tokens; longer prompts are truncated.",
+ )
+ parser.add_argument(
+ "--validation_prompt",
+ type=str,
+ default=None,
+ help="A prompt that is used during validation to verify that the model is learning.",
+ )
+ parser.add_argument(
+ "--skip_final_inference",
+ default=False,
+ action="store_true",
+ help="Whether to skip the final inference step with loaded lora weights upon training completion. This will run intermediate validation inference if `validation_prompt` is provided. Specify to reduce memory.",
+ )
+ parser.add_argument(
+ "--final_validation_prompt",
+ type=str,
+ default=None,
+ help="A prompt that is used during a final validation to verify that the model is learning. Ignored if `--validation_prompt` is provided.",
+ )
+ parser.add_argument(
+ "--num_validation_images",
+ type=int,
+ default=4,
+ help="Number of images that should be generated during validation with `validation_prompt`.",
+ )
+ parser.add_argument(
+ "--validation_epochs",
+ type=int,
+ default=50,
+ help=(
+ "Run dreambooth validation every X epochs. Dreambooth validation consists of running the prompt"
+ " `args.validation_prompt` multiple times: `args.num_validation_images`."
+ ),
+ )
+ parser.add_argument(
+ "--rank",
+ type=int,
+ default=4,
+ help=("The dimension of the LoRA update matrices."),
+ )
+ parser.add_argument(
+ "--lora_alpha",
+ type=int,
+ default=4,
+ help="LoRA alpha to be used for additional scaling.",
+ )
+ parser.add_argument("--lora_dropout", type=float, default=0.0, help="Dropout probability for LoRA layers")
+
+ parser.add_argument(
+ "--with_prior_preservation",
+ default=False,
+ action="store_true",
+ help="Flag to add prior preservation loss.",
+ )
+ parser.add_argument("--prior_loss_weight", type=float, default=1.0, help="The weight of prior preservation loss.")
+ parser.add_argument(
+ "--num_class_images",
+ type=int,
+ default=100,
+ help=(
+ "Minimal class images for prior preservation loss. If there are not enough images already present in"
+ " class_data_dir, additional images will be sampled with class_prompt."
+ ),
+ )
+ parser.add_argument(
+ "--output_dir",
+ type=str,
+ default="fibo2-dreambooth-lora",
+ help="The output directory where the model predictions and checkpoints will be written.",
+ )
+ parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+ parser.add_argument(
+ "--resolution",
+ type=int,
+ default=512,
+ help=(
+ "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+ " resolution"
+ ),
+ )
+ parser.add_argument(
+ "--aspect_ratio_buckets",
+ type=str,
+ default=None,
+ help=(
+ "Aspect ratio buckets to use for training. Define as a string of 'h1,w1;h2,w2;...'. "
+ "e.g. '1024,1024;768,1360;1360,768;880,1168;1168,880;1248,832;832,1248'. "
+ "Requires --use_aspect_ratio_buckets. Images are resized to cover and cropped to the nearest "
+ "listed bucket (smaller images are upscaled). When set, --resolution is ignored."
+ ),
+ )
+ parser.add_argument(
+ "--use_aspect_ratio_buckets",
+ action="store_true",
+ help=(
+ "Enable aspect-ratio bucketing. Without --aspect_ratio_buckets, the buckets are computed on the "
+ "fly from --resolution and capped to each image's own resolution, so smaller images are assigned "
+ "to a smaller bucket instead of being upscaled. Provide --aspect_ratio_buckets to use an explicit list."
+ ),
+ )
+ parser.add_argument(
+ "--center_crop",
+ default=False,
+ action="store_true",
+ help=(
+ "Whether to center crop the input images to the resolution. If not set, the images will be randomly"
+ " cropped. The images will be resized to the resolution first before cropping."
+ ),
+ )
+ parser.add_argument(
+ "--random_flip",
+ action="store_true",
+ help="whether to randomly flip images horizontally",
+ )
+ parser.add_argument(
+ "--caption_dropout",
+ type=float,
+ default=0.0,
+ help=(
+ "Probability of replacing an instance image's caption with an empty string during training, so that"
+ " fraction of samples is trained unconditionally. Improves classifier-free guidance. A common value is"
+ " 0.1. Class/prior-preservation captions are never dropped."
+ ),
+ )
+ parser.add_argument(
+ "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
+ )
+ parser.add_argument(
+ "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images."
+ )
+ parser.add_argument("--num_train_epochs", type=int, default=1)
+ parser.add_argument(
+ "--max_train_steps",
+ type=int,
+ default=None,
+ help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+ )
+ parser.add_argument(
+ "--checkpointing_steps",
+ type=int,
+ default=500,
+ help=(
+ "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+ " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+ " training using `--resume_from_checkpoint`."
+ ),
+ )
+ parser.add_argument(
+ "--checkpoints_total_limit",
+ type=int,
+ default=None,
+ help=("Max number of checkpoints to store."),
+ )
+ parser.add_argument(
+ "--resume_from_checkpoint",
+ type=str,
+ default=None,
+ help=(
+ "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+ ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+ ),
+ )
+ parser.add_argument(
+ "--gradient_accumulation_steps",
+ type=int,
+ default=1,
+ help="Number of updates steps to accumulate before performing a backward/update pass.",
+ )
+ parser.add_argument(
+ "--gradient_checkpointing",
+ action="store_true",
+ help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+ )
+ parser.add_argument(
+ "--learning_rate",
+ type=float,
+ default=1e-4,
+ help="Initial learning rate (after the potential warmup period) to use.",
+ )
+
+ parser.add_argument(
+ "--scale_lr",
+ action="store_true",
+ default=False,
+ help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+ )
+ parser.add_argument(
+ "--lr_scheduler",
+ type=str,
+ default="constant",
+ help=(
+ 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+ ' "constant", "constant_with_warmup"]'
+ ),
+ )
+ parser.add_argument(
+ "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+ )
+ parser.add_argument(
+ "--lr_num_cycles",
+ type=int,
+ default=1,
+ help="Number of hard resets of the lr in cosine_with_restarts scheduler.",
+ )
+ parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.")
+ parser.add_argument(
+ "--dataloader_num_workers",
+ type=int,
+ default=0,
+ help=(
+ "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+ ),
+ )
+ parser.add_argument(
+ "--weighting_scheme",
+ type=str,
+ default="none",
+ choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"],
+ help=('We default to the "none" weighting scheme for uniform sampling and uniform loss'),
+ )
+ parser.add_argument(
+ "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme."
+ )
+ parser.add_argument(
+ "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme."
+ )
+ parser.add_argument(
+ "--mode_scale",
+ type=float,
+ default=1.29,
+ help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
+ )
+ parser.add_argument(
+ "--optimizer",
+ type=str,
+ default="AdamW",
+ help=('The optimizer type to use. Choose between ["AdamW", "prodigy"]'),
+ )
+
+ parser.add_argument(
+ "--use_8bit_adam",
+ action="store_true",
+ help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW",
+ )
+
+ parser.add_argument(
+ "--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers."
+ )
+ parser.add_argument(
+ "--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers."
+ )
+ parser.add_argument(
+ "--prodigy_beta3",
+ type=float,
+ default=None,
+ help="coefficients for computing the Prodigy stepsize using running averages. If set to None, "
+ "uses the value of square root of beta2. Ignored if optimizer is adamW",
+ )
+ parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay")
+ parser.add_argument("--adam_weight_decay", type=float, default=1e-04, help="Weight decay to use for unet params")
+
+ parser.add_argument(
+ "--lora_layers",
+ type=str,
+ default=None,
+ help=(
+ 'The transformer modules to apply LoRA training on. Please specify the layers in a comma separated. E.g. - "to_k,to_q,to_v,to_out.0" will result in lora training of attention layers only'
+ ),
+ )
+
+ parser.add_argument(
+ "--adam_epsilon",
+ type=float,
+ default=1e-08,
+ help="Epsilon value for the Adam optimizer and Prodigy optimizers.",
+ )
+
+ parser.add_argument(
+ "--prodigy_use_bias_correction",
+ type=bool,
+ default=True,
+ help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW",
+ )
+ parser.add_argument(
+ "--prodigy_safeguard_warmup",
+ type=bool,
+ default=True,
+ help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. "
+ "Ignored if optimizer is adamW",
+ )
+ parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+ parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+ parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
+ parser.add_argument(
+ "--hub_model_id",
+ type=str,
+ default=None,
+ help="The name of the repository to keep in sync with the local `output_dir`.",
+ )
+ parser.add_argument(
+ "--logging_dir",
+ type=str,
+ default="logs",
+ help=(
+ "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+ " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+ ),
+ )
+ parser.add_argument(
+ "--allow_tf32",
+ action="store_true",
+ help=(
+ "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+ " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+ ),
+ )
+ parser.add_argument(
+ "--cache_latents",
+ action="store_true",
+ default=False,
+ help="Cache the VAE latents",
+ )
+ parser.add_argument(
+ "--report_to",
+ type=str,
+ default="tensorboard",
+ help=(
+ 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+ ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+ ),
+ )
+ parser.add_argument(
+ "--mixed_precision",
+ type=str,
+ default=None,
+ choices=["no", "fp16", "bf16"],
+ help=(
+ "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+ " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the"
+ " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+ ),
+ )
+ parser.add_argument(
+ "--upcast_before_saving",
+ action="store_true",
+ default=False,
+ help=(
+ "Whether to upcast the trained transformer layers to float32 before saving (at the end of training). "
+ "Defaults to precision dtype used for training to save memory"
+ ),
+ )
+ parser.add_argument(
+ "--offload",
+ action="store_true",
+ help="Whether to offload the VAE and the text encoder to CPU when they are not used.",
+ )
+ parser.add_argument(
+ "--prior_generation_precision",
+ type=str,
+ default=None,
+ choices=["no", "fp32", "fp16", "bf16"],
+ help=(
+ "Choose prior generation precision between fp32, fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+ " 1.10.and an Nvidia Ampere GPU. Default to fp16 if a GPU is available else fp32."
+ ),
+ )
+ parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
+ parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU")
+ parser.add_argument("--fsdp_text_encoder", action="store_true", help="Use FSDP for text encoder")
+
+ if input_args is not None:
+ args = parser.parse_args(input_args)
+ else:
+ args = parser.parse_args()
+
+ if args.dataset_name is None and args.instance_data_dir is None:
+ raise ValueError("Specify either `--dataset_name` or `--instance_data_dir`")
+
+ if args.dataset_name is not None and args.instance_data_dir is not None:
+ raise ValueError("Specify only one of `--dataset_name` or `--instance_data_dir`")
+ if args.do_fp8_training and args.bnb_quantization_config_path:
+ raise ValueError("Both `do_fp8_training` and `bnb_quantization_config_path` cannot be passed.")
+
+ env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
+ if env_local_rank != -1 and env_local_rank != args.local_rank:
+ args.local_rank = env_local_rank
+
+ if args.with_prior_preservation:
+ if args.class_data_dir is None:
+ raise ValueError("You must specify a data directory for class images.")
+ if args.class_prompt is None:
+ raise ValueError("You must specify prompt for class images.")
+ else:
+ # logger is not available yet
+ if args.class_data_dir is not None:
+ warnings.warn("You need not use --class_data_dir without --with_prior_preservation.")
+ if args.class_prompt is not None:
+ warnings.warn("You need not use --class_prompt without --with_prior_preservation.")
+
+ return args
+
+
+class DreamBoothDataset(Dataset):
+ """
+ A dataset to prepare the instance and class images with the prompts for fine-tuning the model.
+ It pre-processes the images.
+ """
+
+ def __init__(
+ self,
+ instance_data_root,
+ instance_prompt,
+ class_prompt,
+ class_data_root=None,
+ class_num=None,
+ size=1024,
+ repeats=1,
+ center_crop=False,
+ buckets=None,
+ use_aspect_ratio_buckets=False,
+ bucket_divisibility=16,
+ bucket_base_resolutions=None,
+ ):
+ self.size = size
+ self.resolution = size
+ self.center_crop = center_crop
+
+ self.instance_prompt = instance_prompt
+ self.custom_instance_prompts = None
+ self.class_prompt = class_prompt
+
+ # Explicit user-provided bucket list (or None). The concrete list of buckets actually used is
+ # built from the data in `self.buckets` during preprocessing below.
+ self._explicit_buckets = buckets
+ self.use_aspect_ratio_buckets = use_aspect_ratio_buckets
+ self.bucket_divisibility = bucket_divisibility
+ self.bucket_base_resolutions = bucket_base_resolutions
+
+ # if --dataset_name is provided or a metadata jsonl file is provided in the local --instance_data directory,
+ # we load the training data using load_dataset
+ if args.dataset_name is not None:
+ try:
+ from datasets import load_dataset
+ except ImportError:
+ raise ImportError(
+ "You are trying to load your data using the datasets library. If you wish to train using custom "
+ "captions please install the datasets library: `pip install datasets`. If you wish to load a "
+ "local folder containing images only, specify --instance_data_dir instead."
+ )
+ # Downloading and loading a dataset from the hub.
+ # See more about loading custom images at
+ # https://huggingface.co/docs/datasets/v2.0.0/en/dataset_script
+ dataset = load_dataset(
+ args.dataset_name,
+ args.dataset_config_name,
+ cache_dir=args.cache_dir,
+ )
+ # Preprocessing the datasets.
+ column_names = dataset["train"].column_names
+
+ # 6. Get the column names for input/target.
+ if args.image_column is None:
+ image_column = column_names[0]
+ logger.info(f"image column defaulting to {image_column}")
+ else:
+ image_column = args.image_column
+ if image_column not in column_names:
+ raise ValueError(
+ f"`--image_column` value '{args.image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
+ )
+ instance_images = dataset["train"][image_column]
+
+ if args.caption_column is None:
+ logger.info(
+ "No caption column provided, defaulting to instance_prompt for all images. If your dataset "
+ "contains captions/prompts for the images, make sure to specify the "
+ "column as --caption_column"
+ )
+ self.custom_instance_prompts = None
+ else:
+ if args.caption_column not in column_names:
+ raise ValueError(
+ f"`--caption_column` value '{args.caption_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
+ )
+ custom_instance_prompts = dataset["train"][args.caption_column]
+ # create final list of captions according to --repeats
+ self.custom_instance_prompts = []
+ for caption in custom_instance_prompts:
+ self.custom_instance_prompts.extend(itertools.repeat(caption, repeats))
+ else:
+ self.instance_data_root = Path(instance_data_root)
+ if not self.instance_data_root.exists():
+ raise ValueError("Instance images root doesn't exists.")
+
+ instance_images = [Image.open(path) for path in list(Path(instance_data_root).iterdir())]
+ self.custom_instance_prompts = None
+
+ self.instance_images = []
+ for img in instance_images:
+ self.instance_images.extend(itertools.repeat(img, repeats))
+
+ self.pixel_values = []
+ self.buckets = []
+ bucket_to_idx = {}
+ for i, image in enumerate(self.instance_images):
+ image = exif_transpose(image)
+ if not image.mode == "RGB":
+ image = image.convert("RGB")
+
+ width, height = image.size
+
+ # Assign the image to a bucket.
+ target = self._bucket_for_image(height, width)
+ if target not in bucket_to_idx:
+ bucket_to_idx[target] = len(self.buckets)
+ self.buckets.append(target)
+ bucket_idx = bucket_to_idx[target]
+
+ # based on the bucket assignment, define the transformations
+ image = self.train_transform(
+ image,
+ size=target,
+ center_crop=args.center_crop,
+ random_flip=args.random_flip,
+ )
+ self.pixel_values.append((image, bucket_idx))
+
+ self.num_instance_images = len(self.instance_images)
+ self._length = self.num_instance_images
+
+ if class_data_root is not None:
+ self.class_data_root = Path(class_data_root)
+ self.class_data_root.mkdir(parents=True, exist_ok=True)
+ self.class_images_path = list(self.class_data_root.iterdir())
+ if class_num is not None:
+ self.num_class_images = min(len(self.class_images_path), class_num)
+ else:
+ self.num_class_images = len(self.class_images_path)
+ self._length = max(self.num_class_images, self.num_instance_images)
+ else:
+ self.class_data_root = None
+
+ def __len__(self):
+ return self._length
+
+ def __getitem__(self, index):
+ example = {}
+ instance_image, bucket_idx = self.pixel_values[index % self.num_instance_images]
+ example["index"] = index
+ example["instance_images"] = instance_image
+ example["bucket_idx"] = bucket_idx
+ if self.custom_instance_prompts:
+ caption = self.custom_instance_prompts[index % self.num_instance_images]
+ if caption:
+ example["instance_prompt"] = caption
+ else:
+ example["instance_prompt"] = self.instance_prompt
+
+ else: # custom prompts were provided, but length does not match size of image dataset
+ example["instance_prompt"] = self.instance_prompt
+
+ if self.class_data_root:
+ class_image = Image.open(self.class_images_path[index % self.num_class_images])
+ class_image = exif_transpose(class_image)
+
+ if not class_image.mode == "RGB":
+ class_image = class_image.convert("RGB")
+ # Match the class image to the paired instance image's bucket so they can be stacked into one batch.
+ example["class_images"] = self.train_transform(
+ class_image, size=self.buckets[bucket_idx], center_crop=self.center_crop
+ )
+ example["class_prompt"] = self.class_prompt
+
+ return example
+
+ def _bucket_for_image(self, height, width):
+ # An explicit bucket list takes priority: pick the nearest, upscaling smaller images to cover it.
+ if self._explicit_buckets is not None:
+ return self._explicit_buckets[find_nearest_bucket(height, width, self._explicit_buckets)]
+ # On-the-fly bucketing: cap the ladder to the image's own resolution so smaller images are
+ # assigned to a smaller bucket rather than being upscaled (mirrors ostris' bucketing).
+ if self.use_aspect_ratio_buckets:
+ resolution = min(self.resolution, round((height * width) ** 0.5))
+ ladder = generate_aspect_ratio_buckets(
+ resolution,
+ divisibility=self.bucket_divisibility,
+ base_resolutions=self.bucket_base_resolutions,
+ )
+ return ladder[find_nearest_bucket(height, width, ladder)]
+ # No bucketing: a single square bucket reproduces the fixed-size resize + crop.
+ return (self.resolution, self.resolution)
+
+ def train_transform(self, image, size, center_crop=False, random_flip=False):
+ # Resize preserving aspect ratio so the image covers the bucket, then crop to the bucket size.
+ target_height, target_width = size
+ width, height = image.size
+ scale = max(target_height / height, target_width / width)
+ new_height, new_width = round(height * scale), round(width * scale)
+ image = TF.resize(image, [new_height, new_width], interpolation=transforms.InterpolationMode.BILINEAR)
+ if center_crop:
+ image = TF.center_crop(image, size)
+ else:
+ i, j, h, w = transforms.RandomCrop.get_params(image, output_size=size)
+ image = TF.crop(image, i, j, h, w)
+ if random_flip and random.random() < 0.5:
+ image = TF.hflip(image)
+ return TF.normalize(TF.to_tensor(image), [0.5], [0.5])
+
+
+def collate_fn(examples, with_prior_preservation=False):
+ indices = [example["index"] for example in examples]
+ pixel_values = [example["instance_images"] for example in examples]
+ # Keep instance_prompts unchanged for prompt cache precompute; prompts may be extended with class prompts below.
+ instance_prompts = [example["instance_prompt"] for example in examples]
+ prompts = [example["instance_prompt"] for example in examples]
+
+ # Concat class and instance examples for prior preservation.
+ # We do this to avoid doing two forward passes.
+ if with_prior_preservation:
+ pixel_values += [example["class_images"] for example in examples]
+ prompts += [example["class_prompt"] for example in examples]
+
+ pixel_values = torch.stack(pixel_values)
+ pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()
+
+ batch = {
+ "indices": indices,
+ "pixel_values": pixel_values,
+ "instance_prompts": instance_prompts,
+ "prompts": prompts,
+ }
+ return batch
+
+
+class BucketBatchSampler(BatchSampler):
+ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, seed: int = None):
+ if not isinstance(batch_size, int) or batch_size <= 0:
+ raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size))
+ if not isinstance(drop_last, bool):
+ raise ValueError("drop_last should be a boolean value, but got drop_last={}".format(drop_last))
+
+ self.dataset = dataset
+ self.batch_size = batch_size
+ self.drop_last = drop_last
+ self.generator = random.Random(seed) if seed is not None else random
+
+ # Group indices by bucket
+ self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))]
+ for idx, (_, bucket_idx) in enumerate(self.dataset.pixel_values):
+ self.bucket_indices[bucket_idx].append(idx)
+
+ self.sampler_len = 0
+ for indices_in_bucket in self.bucket_indices:
+ num_batches, remainder = divmod(len(indices_in_bucket), self.batch_size)
+ self.sampler_len += num_batches
+ if remainder > 0 and not self.drop_last:
+ self.sampler_len += 1
+
+ def __iter__(self):
+ batches = []
+ for indices_in_bucket in self.bucket_indices:
+ shuffled_indices = indices_in_bucket.copy()
+ self.generator.shuffle(shuffled_indices)
+ for i in range(0, len(shuffled_indices), self.batch_size):
+ batch = shuffled_indices[i : i + self.batch_size]
+ if len(batch) < self.batch_size and self.drop_last:
+ continue
+ batches.append(batch)
+
+ self.generator.shuffle(batches)
+ for batch in batches:
+ yield batch
+
+ def __len__(self):
+ return self.sampler_len
+
+
+class PromptDataset(Dataset):
+ "A simple dataset to prepare the prompts to generate class images on multiple GPUs."
+
+ def __init__(self, prompt, num_samples):
+ self.prompt = prompt
+ self.num_samples = num_samples
+
+ def __len__(self):
+ return self.num_samples
+
+ def __getitem__(self, index):
+ example = {}
+ example["prompt"] = self.prompt
+ example["index"] = index
+ return example
+
+
+def main(args):
+ if args.report_to == "wandb" and args.hub_token is not None:
+ raise ValueError(
+ "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token."
+ " Please use `hf auth login` to authenticate with the Hub."
+ )
+
+ if torch.backends.mps.is_available() and args.mixed_precision == "bf16":
+ # due to pytorch#99272, MPS does not yet support bfloat16.
+ raise ValueError(
+ "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
+ )
+ if args.do_fp8_training:
+ from torchao.float8 import Float8LinearConfig, convert_to_float8_training
+
+ logging_dir = Path(args.output_dir, args.logging_dir)
+
+ accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+ kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+ accelerator = Accelerator(
+ gradient_accumulation_steps=args.gradient_accumulation_steps,
+ mixed_precision=args.mixed_precision,
+ log_with=args.report_to,
+ project_config=accelerator_project_config,
+ kwargs_handlers=[kwargs],
+ )
+
+ # Disable AMP for MPS.
+ if torch.backends.mps.is_available():
+ accelerator.native_amp = False
+
+ if args.report_to == "wandb":
+ if not is_wandb_available():
+ raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+
+ # Make one log on every process with the configuration for debugging.
+ logging.basicConfig(
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+ datefmt="%m/%d/%Y %H:%M:%S",
+ level=logging.INFO,
+ )
+ logger.info(accelerator.state, main_process_only=False)
+ if accelerator.is_local_main_process:
+ transformers.utils.logging.set_verbosity_warning()
+ diffusers.utils.logging.set_verbosity_info()
+ else:
+ transformers.utils.logging.set_verbosity_error()
+ diffusers.utils.logging.set_verbosity_error()
+
+ # If passed along, set the training seed now.
+ if args.seed is not None:
+ set_seed(args.seed)
+
+ # Generate class images if prior preservation is enabled.
+ if args.with_prior_preservation:
+ class_images_dir = Path(args.class_data_dir)
+ if not class_images_dir.exists():
+ class_images_dir.mkdir(parents=True)
+ cur_class_images = len(list(class_images_dir.iterdir()))
+
+ if cur_class_images < args.num_class_images:
+ has_supported_fp16_accelerator = torch.cuda.is_available() or torch.backends.mps.is_available()
+ torch_dtype = torch.float16 if has_supported_fp16_accelerator else torch.float32
+ if args.prior_generation_precision == "fp32":
+ torch_dtype = torch.float32
+ elif args.prior_generation_precision == "fp16":
+ torch_dtype = torch.float16
+ elif args.prior_generation_precision == "bf16":
+ torch_dtype = torch.bfloat16
+
+ pipeline = BriaFibo2Pipeline.from_pretrained(
+ args.pretrained_model_name_or_path,
+ torch_dtype=torch_dtype,
+ revision=args.revision,
+ variant=args.variant,
+ )
+
+ num_new_images = args.num_class_images - cur_class_images
+ logger.info(f"Number of class images to sample: {num_new_images}.")
+
+ sample_dataset = PromptDataset(args.class_prompt, num_new_images)
+ sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)
+
+ sample_dataloader = accelerator.prepare(sample_dataloader)
+ pipeline.to(accelerator.device)
+
+ for example in tqdm(
+ sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
+ ):
+ images = pipeline(prompt=example["prompt"]).images
+
+ for i, image in enumerate(images):
+ hash_image = insecure_hashlib.sha1(image.tobytes()).hexdigest()
+ image_filename = class_images_dir / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg"
+ image.save(image_filename)
+
+ del pipeline
+ free_memory()
+
+ # Handle the repository creation
+ if accelerator.is_main_process:
+ if args.output_dir is not None:
+ os.makedirs(args.output_dir, exist_ok=True)
+
+ if args.push_to_hub:
+ repo_id = create_repo(
+ repo_id=args.hub_model_id or Path(args.output_dir).name,
+ exist_ok=True,
+ ).repo_id
+
+ # Load the Qwen3-VL processor (tokenizer + image processor)
+ processor = Qwen3VLProcessor.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="processor",
+ revision=args.revision,
+ )
+
+ # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision
+ # as these weights are only used for inference, keeping weights in full precision is not required.
+ weight_dtype = torch.float32
+ if accelerator.mixed_precision == "fp16":
+ weight_dtype = torch.float16
+ elif accelerator.mixed_precision == "bf16":
+ weight_dtype = torch.bfloat16
+
+ # Load scheduler and models. fibo-2 shares FLUX.2's flow-matching scheduler (dynamic shift), loaded here so the
+ # training noise schedule matches what the model was trained and is sampled with.
+ noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="scheduler",
+ revision=args.revision,
+ )
+ noise_scheduler_copy = copy.deepcopy(noise_scheduler)
+ vae = AutoencoderKLFlux2.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="vae",
+ revision=args.revision,
+ variant=args.variant,
+ )
+ latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(accelerator.device)
+ latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(
+ accelerator.device
+ )
+
+ quantization_config = None
+ if args.bnb_quantization_config_path is not None:
+ with open(args.bnb_quantization_config_path, "r") as f:
+ config_kwargs = json.load(f)
+ if "load_in_4bit" in config_kwargs and config_kwargs["load_in_4bit"]:
+ config_kwargs["bnb_4bit_compute_dtype"] = weight_dtype
+ quantization_config = BitsAndBytesConfig(**config_kwargs)
+
+ transformer = BriaFibo2Transformer2DModel.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="transformer",
+ revision=args.revision,
+ variant=args.variant,
+ quantization_config=quantization_config,
+ torch_dtype=weight_dtype,
+ )
+ if args.bnb_quantization_config_path is not None:
+ transformer = prepare_model_for_kbit_training(transformer, use_gradient_checkpointing=False)
+
+ text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="text_encoder",
+ revision=args.revision,
+ variant=args.variant,
+ torch_dtype=weight_dtype,
+ )
+ text_encoder.requires_grad_(False)
+
+ # We only train the additional adapter LoRA layers
+ transformer.requires_grad_(False)
+ vae.requires_grad_(False)
+
+ if args.enable_npu_flash_attention:
+ if is_torch_npu_available():
+ logger.info("npu flash attention enabled.")
+ transformer.set_attention_backend("_native_npu")
+ else:
+ raise ValueError("npu flash attention requires torch_npu extensions and is supported only on npu device ")
+
+ if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:
+ # due to pytorch#99272, MPS does not yet support bfloat16.
+ raise ValueError(
+ "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
+ )
+
+ to_kwargs = {"dtype": weight_dtype, "device": accelerator.device} if not args.offload else {"dtype": weight_dtype}
+ # flux vae is stable in bf16 so load it in weight_dtype to reduce memory
+ vae.to(**to_kwargs)
+ # we never offload the transformer to CPU, so we can just use the accelerator device
+ transformer_to_kwargs = (
+ {"device": accelerator.device}
+ if args.bnb_quantization_config_path is not None
+ else {"device": accelerator.device, "dtype": weight_dtype}
+ )
+
+ is_fsdp = getattr(accelerator.state, "fsdp_plugin", None) is not None
+ if not is_fsdp:
+ transformer.to(**transformer_to_kwargs)
+
+ if args.do_fp8_training:
+ convert_to_float8_training(
+ transformer, module_filter_fn=module_filter_fn, config=Float8LinearConfig(pad_inner_dim=True)
+ )
+
+ text_encoder.to(**to_kwargs)
+
+ if args.gradient_checkpointing:
+ transformer.enable_gradient_checkpointing()
+
+ if args.lora_layers is not None:
+ target_modules = [layer.strip() for layer in args.lora_layers.split(",")]
+ else:
+ # Attention (self- and cross-attention) + SwiGLU feed-forward projections across every block. The modulation
+ # (`adaLN_modulation.linear`), patch embedding (`x_embedder`) and output projection (`final_layer.linear`)
+ # are left out, matching the attention+FFN set the sibling trainers default to. Pass `--lora_layers` to
+ # target a different set.
+ target_modules = [
+ "to_q",
+ "to_k",
+ "to_v",
+ "to_out.0",
+ "w1",
+ "w2",
+ "w3",
+ ]
+
+ # now we will add new LoRA weights the transformer layers
+ transformer_lora_config = LoraConfig(
+ r=args.rank,
+ lora_alpha=args.lora_alpha,
+ lora_dropout=args.lora_dropout,
+ init_lora_weights="gaussian",
+ target_modules=target_modules,
+ )
+ transformer.add_adapter(transformer_lora_config)
+
+ def unwrap_model(model):
+ model = accelerator.unwrap_model(model)
+ model = model._orig_mod if is_compiled_module(model) else model
+ return model
+
+ # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
+ def save_model_hook(models, weights, output_dir):
+ transformer_cls = type(unwrap_model(transformer))
+
+ # 1) Validate and pick the transformer model
+ modules_to_save: dict[str, Any] = {}
+ transformer_model = None
+
+ for model in models:
+ if isinstance(unwrap_model(model), transformer_cls):
+ transformer_model = model
+ modules_to_save["transformer"] = model
+ else:
+ raise ValueError(f"unexpected save model: {model.__class__}")
+
+ if transformer_model is None:
+ raise ValueError("No transformer model found in 'models'")
+
+ # 2) Optionally gather FSDP state dict once
+ state_dict = accelerator.get_state_dict(model) if is_fsdp else None
+
+ # 3) Only main process materializes the LoRA state dict
+ transformer_lora_layers_to_save = None
+ if accelerator.is_main_process:
+ peft_kwargs = {}
+ if is_fsdp:
+ peft_kwargs["state_dict"] = state_dict
+
+ transformer_lora_layers_to_save = get_peft_model_state_dict(
+ unwrap_model(transformer_model) if is_fsdp else transformer_model,
+ **peft_kwargs,
+ )
+
+ if is_fsdp:
+ transformer_lora_layers_to_save = _to_cpu_contiguous(transformer_lora_layers_to_save)
+
+ # make sure to pop weight so that corresponding model is not saved again
+ if weights:
+ weights.pop()
+
+ BriaFibo2LoraLoaderMixin.save_lora_weights(
+ output_dir,
+ transformer_lora_layers=transformer_lora_layers_to_save,
+ **_collate_lora_metadata(modules_to_save),
+ )
+
+ def load_model_hook(models, input_dir):
+ transformer_ = None
+
+ if not is_fsdp:
+ while len(models) > 0:
+ model = models.pop()
+
+ if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
+ transformer_ = unwrap_model(model)
+ else:
+ raise ValueError(f"unexpected save model: {model.__class__}")
+ else:
+ transformer_ = BriaFibo2Transformer2DModel.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="transformer",
+ )
+ transformer_.add_adapter(transformer_lora_config)
+
+ lora_state_dict = BriaFibo2LoraLoaderMixin.lora_state_dict(input_dir)
+
+ transformer_state_dict = {
+ f"{k.replace('transformer.', '')}": v for k, v in lora_state_dict.items() if k.startswith("transformer.")
+ }
+ transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict)
+ incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default")
+ if incompatible_keys is not None:
+ # check only for unexpected keys
+ unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
+ if unexpected_keys:
+ logger.warning(
+ f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
+ f" {unexpected_keys}. "
+ )
+
+ # Make sure the trainable params are in float32. This is again needed since the base models
+ # are in `weight_dtype`. More details:
+ # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
+ if args.mixed_precision == "fp16":
+ models = [transformer_]
+ # only upcast trainable parameters (LoRA) into fp32
+ cast_training_params(models)
+
+ accelerator.register_save_state_pre_hook(save_model_hook)
+ accelerator.register_load_state_pre_hook(load_model_hook)
+
+ # Enable TF32 for faster training on Ampere GPUs,
+ # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+ if args.allow_tf32 and torch.cuda.is_available():
+ torch.backends.cuda.matmul.allow_tf32 = True
+
+ if args.scale_lr:
+ args.learning_rate = (
+ args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
+ )
+
+ # Make sure the trainable params are in float32.
+ if args.mixed_precision == "fp16":
+ models = [transformer]
+ # only upcast trainable parameters (LoRA) into fp32
+ cast_training_params(models, dtype=torch.float32)
+
+ transformer_lora_parameters = list(filter(lambda p: p.requires_grad, transformer.parameters()))
+
+ # Optimization parameters
+ transformer_parameters_with_lr = {"params": transformer_lora_parameters, "lr": args.learning_rate}
+ params_to_optimize = [transformer_parameters_with_lr]
+
+ # Optimizer creation
+ if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"):
+ logger.warning(
+ f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]."
+ "Defaulting to adamW"
+ )
+ args.optimizer = "adamw"
+
+ if args.use_8bit_adam and not args.optimizer.lower() == "adamw":
+ logger.warning(
+ f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was "
+ f"set to {args.optimizer.lower()}"
+ )
+
+ if args.optimizer.lower() == "adamw":
+ if args.use_8bit_adam:
+ try:
+ import bitsandbytes as bnb
+ except ImportError:
+ raise ImportError(
+ "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+ )
+
+ optimizer_class = bnb.optim.AdamW8bit
+ else:
+ optimizer_class = torch.optim.AdamW
+
+ optimizer = optimizer_class(
+ params_to_optimize,
+ betas=(args.adam_beta1, args.adam_beta2),
+ weight_decay=args.adam_weight_decay,
+ eps=args.adam_epsilon,
+ )
+
+ if args.optimizer.lower() == "prodigy":
+ try:
+ import prodigyopt
+ except ImportError:
+ raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`")
+
+ optimizer_class = prodigyopt.Prodigy
+
+ if args.learning_rate <= 0.1:
+ logger.warning(
+ "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0"
+ )
+
+ optimizer = optimizer_class(
+ params_to_optimize,
+ betas=(args.adam_beta1, args.adam_beta2),
+ beta3=args.prodigy_beta3,
+ weight_decay=args.adam_weight_decay,
+ eps=args.adam_epsilon,
+ decouple=args.prodigy_decouple,
+ use_bias_correction=args.prodigy_use_bias_correction,
+ safeguard_warmup=args.prodigy_safeguard_warmup,
+ )
+
+ # Resolve the bucketing mode. Bucketing must be enabled explicitly with --use_aspect_ratio_buckets;
+ # a bucket list without that flag is an error. With the flag, an explicit --aspect_ratio_buckets list
+ # drives assignment, otherwise buckets are computed on the fly inside the dataset. Without the flag a
+ # single square bucket reproduces the fixed-size resize + crop.
+ if args.aspect_ratio_buckets is not None and not args.use_aspect_ratio_buckets:
+ raise ValueError("--aspect_ratio_buckets requires --use_aspect_ratio_buckets to be set.")
+ if args.aspect_ratio_buckets is not None:
+ buckets = parse_buckets_string(args.aspect_ratio_buckets)
+ use_aspect_ratio_buckets = False
+ logger.info(f"Using explicit aspect ratio buckets: {buckets}")
+ elif args.use_aspect_ratio_buckets:
+ buckets = None
+ use_aspect_ratio_buckets = True
+ logger.info(
+ "No --aspect_ratio_buckets provided; auto-computing aspect ratio buckets on the fly from --resolution."
+ )
+ else:
+ buckets = [(args.resolution, args.resolution)]
+ use_aspect_ratio_buckets = False
+
+ # Dataset and DataLoaders creation:
+ train_dataset = DreamBoothDataset(
+ instance_data_root=args.instance_data_dir,
+ instance_prompt=args.instance_prompt,
+ class_prompt=args.class_prompt,
+ class_data_root=args.class_data_dir if args.with_prior_preservation else None,
+ class_num=args.num_class_images,
+ size=args.resolution,
+ repeats=args.repeats,
+ center_crop=args.center_crop,
+ buckets=buckets,
+ use_aspect_ratio_buckets=use_aspect_ratio_buckets,
+ )
+ precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts
+ batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True, seed=args.seed)
+ train_dataloader = torch.utils.data.DataLoader(
+ train_dataset,
+ batch_sampler=batch_sampler,
+ collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
+ num_workers=args.dataloader_num_workers,
+ )
+
+ def compute_text_embeddings(prompt, text_encoder, processor):
+ # fibo-2 has no classic text-encoding pipeline; encode directly with the frozen Qwen3-VL encoder, exactly as
+ # the pipeline does. Returns `(1, 6, seq_len, dim)` embeddings and their bool mask, per prompt.
+ with torch.no_grad():
+ prompt_embeds, prompt_embeds_mask = get_fibo2_prompt_embeds(
+ text_encoder, processor, prompt, accelerator.device, weight_dtype, args.max_sequence_length
+ )
+ return prompt_embeds, prompt_embeds_mask
+
+ # If no type of tuning is done on the text_encoder and custom instance prompts are NOT
+ # provided (i.e. the --instance_prompt is used for all images), we encode the instance prompt once to avoid
+ # the redundant encoding.
+ if not train_dataset.custom_instance_prompts:
+ with offload_models(text_encoder, device=accelerator.device, offload=args.offload):
+ instance_prompt_hidden_states, instance_prompt_mask = compute_text_embeddings(
+ args.instance_prompt, text_encoder, processor
+ )
+
+ # Handle class prompt for prior-preservation.
+ if args.with_prior_preservation:
+ with offload_models(text_encoder, device=accelerator.device, offload=args.offload):
+ class_prompt_hidden_states, class_prompt_mask = compute_text_embeddings(
+ args.class_prompt, text_encoder, processor
+ )
+
+ # fibo-2 validation runs through the standard pipeline (which encodes the prompt itself), so we only assemble
+ # the generation kwargs here rather than precomputing text embeddings.
+ validation_pipeline_args = {}
+ if args.validation_prompt is not None:
+ validation_pipeline_args = {
+ "prompt": args.validation_prompt,
+ "height": args.resolution,
+ "width": args.resolution,
+ "num_inference_steps": 30,
+ "guidance_scale": 5.0,
+ }
+
+ # Init FSDP for text encoder
+ if args.fsdp_text_encoder:
+ fsdp_kwargs = get_fsdp_kwargs_from_accelerator(accelerator)
+ text_encoder = wrap_with_fsdp(
+ model=text_encoder,
+ device=accelerator.device,
+ offload=args.offload,
+ limit_all_gathers=True,
+ use_orig_params=True,
+ fsdp_kwargs=fsdp_kwargs,
+ )
+ dist.barrier()
+
+ # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images),
+ # pack the statically computed variables appropriately here. This is so that we don't
+ # have to pass them to the dataloader.
+ if not train_dataset.custom_instance_prompts:
+ static_prompt_embeds = instance_prompt_hidden_states
+ static_prompt_mask = instance_prompt_mask
+
+ # if cache_latents is set to True, we encode images to latents and store them.
+ # Similar to pre-encoding in the case of a single instance prompt, if custom prompts are provided
+ # we encode them in advance as well.
+ if args.cache_latents:
+ instance_latents_cache = [None] * train_dataset.num_instance_images
+ class_latents_cache = [None] * train_dataset.num_instance_images if args.with_prior_preservation else None
+ if train_dataset.custom_instance_prompts:
+ prompt_embeds_cache = [None] * train_dataset.num_instance_images
+ prompt_mask_cache = [None] * train_dataset.num_instance_images
+ if precompute_latents:
+ cache_batch_sampler = BucketBatchSampler(
+ train_dataset, batch_size=args.train_batch_size, drop_last=False, seed=args.seed
+ )
+ cache_dataloader = torch.utils.data.DataLoader(
+ train_dataset,
+ batch_sampler=cache_batch_sampler,
+ collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
+ num_workers=args.dataloader_num_workers,
+ )
+ for batch in tqdm(cache_dataloader, desc="Caching latents"):
+ with torch.no_grad():
+ sample_indices = batch["indices"]
+ if args.cache_latents:
+ with offload_models(vae, device=accelerator.device, offload=args.offload):
+ batch["pixel_values"] = batch["pixel_values"].to(
+ accelerator.device, non_blocking=True, dtype=vae.dtype
+ )
+ latents = vae.encode(batch["pixel_values"]).latent_dist.mode()
+ if args.with_prior_preservation:
+ instance_latents, class_latents = torch.chunk(latents, 2, dim=0)
+ else:
+ instance_latents = latents
+ for i, idx in enumerate(sample_indices):
+ instance_latents_cache[idx] = instance_latents[i : i + 1]
+ if args.with_prior_preservation:
+ class_latents_cache[idx] = class_latents[i : i + 1]
+ if train_dataset.custom_instance_prompts:
+ if args.fsdp_text_encoder:
+ prompt_embeds, prompt_mask = compute_text_embeddings(
+ batch["instance_prompts"], text_encoder, processor
+ )
+ else:
+ with offload_models(text_encoder, device=accelerator.device, offload=args.offload):
+ prompt_embeds, prompt_mask = compute_text_embeddings(
+ batch["instance_prompts"], text_encoder, processor
+ )
+ for i, idx in enumerate(sample_indices):
+ prompt_embeds_cache[idx] = prompt_embeds[i : i + 1]
+ prompt_mask_cache[idx] = prompt_mask[i : i + 1]
+
+ if args.cache_latents:
+ assert all(latents is not None for latents in instance_latents_cache), "Latent cache has unfilled entries."
+ if args.with_prior_preservation:
+ assert all(latents is not None for latents in class_latents_cache), (
+ "Class latent cache has unfilled entries."
+ )
+ if train_dataset.custom_instance_prompts:
+ assert all(embeds is not None for embeds in prompt_embeds_cache), (
+ "Prompt embedding cache has unfilled entries."
+ )
+ assert all(mask is not None for mask in prompt_mask_cache), "Prompt mask cache has unfilled entries."
+
+ # Caption dropout swaps the empty-prompt embedding in for a dropped caption. Precompute it once here while the
+ # text encoder is still loaded (it is deleted just below). The batch's embeddings are right-padded to a shared
+ # length at step time, so the empty prompt's own (shorter) window needs no special handling.
+ empty_prompt_embeds = empty_prompt_mask = None
+ if args.caption_dropout > 0:
+ with offload_models(text_encoder, device=accelerator.device, offload=args.offload):
+ empty_prompt_embeds, empty_prompt_mask = compute_text_embeddings("", text_encoder, processor)
+
+ # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624
+ if args.cache_latents:
+ vae = vae.to("cpu")
+ del vae
+
+ # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624
+ text_encoder = text_encoder.to("cpu")
+ del text_encoder
+ free_memory()
+
+ # Scheduler and math around the number of training steps.
+ # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation.
+ num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes
+ if args.max_train_steps is None:
+ len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes)
+ num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps)
+ num_training_steps_for_scheduler = (
+ args.num_train_epochs * accelerator.num_processes * num_update_steps_per_epoch
+ )
+ else:
+ num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes
+
+ lr_scheduler = get_scheduler(
+ args.lr_scheduler,
+ optimizer=optimizer,
+ num_warmup_steps=num_warmup_steps_for_scheduler,
+ num_training_steps=num_training_steps_for_scheduler,
+ num_cycles=args.lr_num_cycles,
+ power=args.lr_power,
+ )
+
+ # Prepare everything with our `accelerator`.
+ transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+ transformer, optimizer, train_dataloader, lr_scheduler
+ )
+
+ # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+ num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+ if args.max_train_steps is None:
+ args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+ if num_training_steps_for_scheduler != args.max_train_steps:
+ logger.warning(
+ f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match "
+ f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. "
+ f"This inconsistency may result in the learning rate scheduler not functioning properly."
+ )
+ # Afterwards we recalculate our number of training epochs
+ args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+ # We need to initialize the trackers we use, and also store our configuration.
+ # The trackers initializes automatically on the main process.
+ if accelerator.is_main_process:
+ tracker_name = "dreambooth-fibo2-lora"
+ accelerator.init_trackers(tracker_name, config=vars(args).copy())
+
+ # Train!
+ total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+ logger.info("***** Running training *****")
+ logger.info(f" Num examples = {len(train_dataset)}")
+ logger.info(f" Num batches each epoch = {len(train_dataloader)}")
+ logger.info(f" Num Epochs = {args.num_train_epochs}")
+ logger.info(f" Instantaneous batch size per device = {args.train_batch_size}")
+ logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+ logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+ logger.info(f" Total optimization steps = {args.max_train_steps}")
+ global_step = 0
+ first_epoch = 0
+
+ # Potentially load in the weights and states from a previous save
+ if args.resume_from_checkpoint:
+ if args.resume_from_checkpoint != "latest":
+ path = os.path.basename(args.resume_from_checkpoint)
+ else:
+ # Get the mos recent checkpoint
+ dirs = os.listdir(args.output_dir)
+ dirs = [d for d in dirs if d.startswith("checkpoint")]
+ dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+ path = dirs[-1] if len(dirs) > 0 else None
+
+ if path is None:
+ accelerator.print(
+ f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+ )
+ args.resume_from_checkpoint = None
+ initial_global_step = 0
+ else:
+ accelerator.print(f"Resuming from checkpoint {path}")
+ accelerator.load_state(os.path.join(args.output_dir, path))
+ global_step = int(path.split("-")[1])
+
+ initial_global_step = global_step
+ first_epoch = global_step // num_update_steps_per_epoch
+
+ else:
+ initial_global_step = 0
+
+ progress_bar = tqdm(
+ range(0, args.max_train_steps),
+ initial=initial_global_step,
+ desc="Steps",
+ # Only show the progress bar once on each machine.
+ disable=not accelerator.is_local_main_process,
+ )
+
+ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
+ sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype)
+ schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device)
+ timesteps = timesteps.to(accelerator.device)
+ step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+ sigma = sigmas[step_indices].flatten()
+ while len(sigma.shape) < n_dim:
+ sigma = sigma.unsqueeze(-1)
+ return sigma
+
+ for epoch in range(first_epoch, args.num_train_epochs):
+ transformer.train()
+
+ for batch in train_dataloader:
+ models_to_accumulate = [transformer]
+ sample_indices = batch["indices"]
+ n_inst = len(sample_indices)
+
+ with accelerator.accumulate(models_to_accumulate):
+ # Assemble the batch's per-prompt embeddings as [instance..., class...], substituting the empty
+ # prompt for dropped captions, then right-pad them to a shared length. Only instance rows are ever
+ # dropped; class/prior-preservation captions are kept.
+ if train_dataset.custom_instance_prompts:
+ embeds_list = [prompt_embeds_cache[idx] for idx in sample_indices]
+ mask_list = [prompt_mask_cache[idx] for idx in sample_indices]
+ else:
+ embeds_list = [static_prompt_embeds] * n_inst
+ mask_list = [static_prompt_mask] * n_inst
+
+ if args.caption_dropout > 0:
+ drop_mask = torch.rand(n_inst) < args.caption_dropout
+ embeds_list = [empty_prompt_embeds if drop_mask[i] else embeds_list[i] for i in range(n_inst)]
+ mask_list = [empty_prompt_mask if drop_mask[i] else mask_list[i] for i in range(n_inst)]
+
+ if args.with_prior_preservation:
+ embeds_list = embeds_list + [class_prompt_hidden_states] * n_inst
+ mask_list = mask_list + [class_prompt_mask] * n_inst
+
+ prompt_embeds, prompt_embeds_mask = pad_and_stack_embeds(embeds_list, mask_list, accelerator.device)
+ prompt_embeds = prompt_embeds.to(dtype=weight_dtype)
+
+ # Convert images to latent space
+ if args.cache_latents:
+ model_input = torch.cat([instance_latents_cache[idx] for idx in sample_indices], dim=0)
+ if args.with_prior_preservation:
+ model_input = torch.cat(
+ [model_input, torch.cat([class_latents_cache[idx] for idx in sample_indices], dim=0)],
+ dim=0,
+ )
+ model_input = model_input.to(accelerator.device)
+ else:
+ with offload_models(vae, device=accelerator.device, offload=args.offload):
+ pixel_values = batch["pixel_values"].to(device=accelerator.device, dtype=vae.dtype)
+ model_input = vae.encode(pixel_values).latent_dist.mode()
+
+ # Normalize the clean latents with the VAE BatchNorm statistics, exactly as the pipeline does, into
+ # the `in_channels` latents the transformer denoises.
+ model_input = normalize_latents(model_input.float(), latents_bn_mean, latents_bn_std).to(weight_dtype)
+
+ # Sample noise that we'll add to the latents
+ noise = torch.randn_like(model_input)
+ bsz = model_input.shape[0]
+
+ # Sample a random timestep for each image
+ # for weighting schemes where we sample timesteps non-uniformly
+ u = compute_density_for_timestep_sampling(
+ weighting_scheme=args.weighting_scheme,
+ batch_size=bsz,
+ logit_mean=args.logit_mean,
+ logit_std=args.logit_std,
+ mode_scale=args.mode_scale,
+ )
+ indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
+ timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device)
+
+ # Add noise according to flow matching.
+ # zt = (1 - texp) * x + texp * z1
+ sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
+ noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise
+
+ # fibo-2 is not guidance-distilled: no guidance argument. The flow timestep is a per-sample scalar
+ # in [0, 1] (the transformer scales it by t_scale internally and gives the image tokens that level,
+ # the text context sits at t = 0). The transformer patchifies the latents and returns only the
+ # generated image tokens, so there is nothing to pack/slice around it.
+ timestep = sigmas.reshape(bsz)
+
+ with accelerator.autocast():
+ model_pred = transformer(
+ hidden_states=noisy_model_input.to(weight_dtype),
+ timestep=timestep,
+ encoder_hidden_states=prompt_embeds,
+ encoder_attention_mask=prompt_embeds_mask,
+ return_dict=False,
+ )[0]
+
+ # these weighting schemes use a uniform timestep sampling
+ # and instead post-weight the loss
+ weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas)
+
+ # flow matching loss (velocity target)
+ target = noise - model_input
+
+ if args.with_prior_preservation:
+ # Chunk the noise and model_pred into two parts and compute the loss on each part separately.
+ model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0)
+ target, target_prior = torch.chunk(target, 2, dim=0)
+ weighting, weighting_prior = torch.chunk(weighting, 2, dim=0)
+
+ # Compute prior loss
+ prior_loss = torch.mean(
+ (weighting_prior.float() * (model_pred_prior.float() - target_prior.float()) ** 2).reshape(
+ target_prior.shape[0], -1
+ ),
+ 1,
+ )
+ prior_loss = prior_loss.mean()
+
+ # Compute regular loss.
+ loss = torch.mean(
+ (weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1),
+ 1,
+ )
+ loss = loss.mean()
+
+ if args.with_prior_preservation:
+ # Add the prior loss to the instance loss.
+ loss = loss + args.prior_loss_weight * prior_loss
+
+ accelerator.backward(loss)
+ if accelerator.sync_gradients:
+ params_to_clip = transformer.parameters()
+ accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+
+ optimizer.step()
+ lr_scheduler.step()
+ optimizer.zero_grad()
+
+ # Checks if the accelerator has performed an optimization step behind the scenes
+ if accelerator.sync_gradients:
+ progress_bar.update(1)
+ global_step += 1
+
+ if accelerator.is_main_process or is_fsdp:
+ if global_step % args.checkpointing_steps == 0:
+ # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
+ if args.checkpoints_total_limit is not None:
+ checkpoints = os.listdir(args.output_dir)
+ checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
+ checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
+
+ # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
+ if len(checkpoints) >= args.checkpoints_total_limit:
+ num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
+ removing_checkpoints = checkpoints[0:num_to_remove]
+
+ logger.info(
+ f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
+ )
+ logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
+
+ for removing_checkpoint in removing_checkpoints:
+ removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
+ shutil.rmtree(removing_checkpoint)
+
+ save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
+ accelerator.save_state(save_path)
+ logger.info(f"Saved state to {save_path}")
+
+ logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
+ progress_bar.set_postfix(**logs)
+ accelerator.log(logs, step=global_step)
+
+ if global_step >= args.max_train_steps:
+ break
+
+ if accelerator.is_main_process:
+ if args.validation_prompt is not None and epoch % args.validation_epochs == 0:
+ transformer_ = unwrap_model(transformer)
+ transformer_.eval()
+ pipeline = make_validation_pipeline(args, transformer_, weight_dtype, accelerator.device)
+ images = log_validation(
+ pipeline=pipeline,
+ args=args,
+ accelerator=accelerator,
+ pipeline_args=validation_pipeline_args,
+ epoch=epoch,
+ torch_dtype=weight_dtype,
+ )
+ transformer_.train()
+ free_memory()
+
+ # Save the lora layers
+ accelerator.wait_for_everyone()
+
+ if is_fsdp:
+ transformer = unwrap_model(transformer)
+ state_dict = accelerator.get_state_dict(transformer)
+ if accelerator.is_main_process:
+ modules_to_save = {}
+ if is_fsdp:
+ if args.bnb_quantization_config_path is None:
+ if args.upcast_before_saving:
+ state_dict = {
+ k: v.to(torch.float32) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
+ }
+ else:
+ state_dict = {
+ k: v.to(weight_dtype) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
+ }
+
+ transformer_lora_layers = get_peft_model_state_dict(
+ transformer,
+ state_dict=state_dict,
+ )
+ transformer_lora_layers = {
+ k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v
+ for k, v in transformer_lora_layers.items()
+ }
+
+ else:
+ transformer = unwrap_model(transformer)
+ if args.bnb_quantization_config_path is None:
+ if args.upcast_before_saving:
+ transformer.to(torch.float32)
+ else:
+ transformer = transformer.to(weight_dtype)
+ transformer_lora_layers = get_peft_model_state_dict(transformer)
+
+ modules_to_save["transformer"] = transformer
+
+ BriaFibo2LoraLoaderMixin.save_lora_weights(
+ save_directory=args.output_dir,
+ transformer_lora_layers=transformer_lora_layers,
+ **_collate_lora_metadata(modules_to_save),
+ )
+
+ images = []
+ run_validation = (args.validation_prompt and args.num_validation_images > 0) or (args.final_validation_prompt)
+ should_run_final_inference = not args.skip_final_inference and run_validation
+ if should_run_final_inference:
+ pipeline = make_validation_pipeline(args, transformer, weight_dtype, accelerator.device)
+
+ # run inference
+ images = []
+ if args.validation_prompt and args.num_validation_images > 0:
+ images = log_validation(
+ pipeline=pipeline,
+ args=args,
+ accelerator=accelerator,
+ pipeline_args=validation_pipeline_args,
+ epoch=epoch,
+ is_final_validation=True,
+ torch_dtype=weight_dtype,
+ )
+ del pipeline
+ free_memory()
+
+ validation_prompt = args.validation_prompt if args.validation_prompt else args.final_validation_prompt
+ quant_training = None
+ if args.do_fp8_training:
+ quant_training = "FP8 TorchAO"
+ elif args.bnb_quantization_config_path:
+ quant_training = "BitsandBytes"
+ save_model_card(
+ (args.hub_model_id or Path(args.output_dir).name) if not args.push_to_hub else repo_id,
+ images=images,
+ base_model=args.pretrained_model_name_or_path,
+ instance_prompt=args.instance_prompt,
+ validation_prompt=validation_prompt,
+ repo_folder=args.output_dir,
+ quant_training=quant_training,
+ )
+
+ if args.push_to_hub:
+ upload_folder(
+ repo_id=repo_id,
+ folder_path=args.output_dir,
+ commit_message="End of training",
+ ignore_patterns=["step_*", "epoch_*"],
+ )
+
+ accelerator.end_training()
+
+
+if __name__ == "__main__":
+ args = parse_args()
+ main(args)
diff --git a/examples/dreambooth/train_dreambooth_lora_bria_fibo2_img2img.py b/examples/dreambooth/train_dreambooth_lora_bria_fibo2_img2img.py
new file mode 100644
index 000000000000..0275e94619ad
--- /dev/null
+++ b/examples/dreambooth/train_dreambooth_lora_bria_fibo2_img2img.py
@@ -0,0 +1,1956 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. 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.
+
+# /// script
+# dependencies = [
+# "diffusers @ git+https://github.com/huggingface/diffusers.git",
+# "torch>=2.0.0",
+# "accelerate>=0.31.0",
+# "transformers>=5.20.1",
+# "ftfy",
+# "tensorboard",
+# "Jinja2",
+# "peft>=0.11.1",
+# "sentencepiece",
+# "torchvision",
+# "datasets",
+# "bitsandbytes",
+# "prodigyopt",
+# ]
+# ///
+
+import argparse
+import copy
+import itertools
+import json
+import logging
+import math
+import os
+import random
+import shutil
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed
+from huggingface_hub import create_repo, upload_folder
+from peft import LoraConfig, prepare_model_for_kbit_training, set_peft_model_state_dict
+from peft.utils import get_peft_model_state_dict
+from PIL.ImageOps import exif_transpose
+from torch.utils.data import Dataset
+from torch.utils.data.sampler import BatchSampler
+from torchvision import transforms
+from torchvision.transforms import functional as TF
+from tqdm.auto import tqdm
+from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
+
+import diffusers
+from diffusers import (
+ AutoencoderKLFlux2,
+ BitsAndBytesConfig,
+ BriaFibo2Pipeline,
+ BriaFibo2Transformer2DModel,
+ FlowMatchEulerDiscreteScheduler,
+)
+from diffusers.loaders import BriaFibo2LoraLoaderMixin
+from diffusers.optimization import get_scheduler
+from diffusers.training_utils import (
+ _collate_lora_metadata,
+ _to_cpu_contiguous,
+ cast_training_params,
+ compute_density_for_timestep_sampling,
+ compute_loss_weighting_for_sd3,
+ find_nearest_bucket,
+ free_memory,
+ generate_aspect_ratio_buckets,
+ get_fsdp_kwargs_from_accelerator,
+ offload_models,
+ parse_buckets_string,
+ wrap_with_fsdp,
+)
+from diffusers.utils import (
+ check_min_version,
+ convert_unet_state_dict_to_peft,
+ is_wandb_available,
+ load_image,
+)
+from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
+from diffusers.utils.import_utils import is_torch_npu_available
+from diffusers.utils.torch_utils import is_compiled_module
+
+
+if getattr(torch, "distributed", None) is not None:
+ import torch.distributed as dist
+
+if is_wandb_available():
+ import wandb
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.40.0.dev0")
+
+logger = get_logger(__name__)
+
+# fibo-2 wraps the prompt in Qwen3-VL's chat template before encoding, as the pipeline does.
+PROMPT_TEMPLATE_ENCODE = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+
+# Where the image to edit stands in a Qwen3-VL prompt.
+VISION_MARKER = "<|vision_start|><|image_pad|><|vision_end|>"
+
+# Qwen3-VL sees the image to edit at most this many pixels a side (`BriaFibo2Pipeline.text_encoder_image_size`).
+TEXT_ENCODER_IMAGE_SIZE = 768
+
+# The Qwen3-VL layers behind the transformer's six text inputs: the Perceiver's bundle first, then one per
+# injection block. Each bundle concatenates three layers, so every text input is 3 * hidden_dim wide. This is model
+# wiring (it must match `BriaFibo2Pipeline.text_encoder_out_layers`), not a tuning knob, so it lives here.
+TEXT_ENCODER_OUT_LAYERS = ((9, 20, 31), (5, 16, 27), (7, 18, 29), (9, 20, 31), (11, 22, 33), (13, 24, 35))
+
+
+# --- fibo-2 latent helpers (identical to the pipeline's) ---
+def patchify_latents(latents):
+ """(B, C, H, W) -> (B, C*4, H//2, W//2), the 2x2 patch order the VAE BatchNorm statistics live in."""
+ b, c, h, w = latents.shape
+ latents = latents.view(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4)
+ return latents.reshape(b, c * 4, h // 2, w // 2)
+
+
+def unpatchify_latents(latents):
+ """(B, C*4, H//2, W//2) -> (B, C, H, W), the inverse of `patchify_latents`."""
+ b, c, h, w = latents.shape
+ latents = latents.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3)
+ return latents.reshape(b, c // 4, h * 2, w * 2)
+
+
+def normalize_latents(latents, bn_mean, bn_std):
+ """Normalize VAE latents with the VAE's BatchNorm statistics, exactly as `BriaFibo2Pipeline` does: the
+ statistics are stored for the 2x2-patchified layout, so patchify, normalize, then unpatchify back to the
+ `in_channels` latents the transformer takes. The result is the clean sample the flow-matching target denoises to
+ (for a conditioning image, the clean context tokens)."""
+ latents = patchify_latents(latents)
+ latents = (latents - bn_mean) / bn_std
+ return unpatchify_latents(latents)
+
+
+def _encode_qwen3vl(text_encoder, processor, text, images, device, dtype, max_sequence_length):
+ text_inputs = processor(
+ text=text,
+ images=images,
+ padding="longest",
+ max_length=max_sequence_length,
+ truncation=True,
+ return_tensors="pt",
+ ).to(device)
+ hidden_states = text_encoder(**text_inputs, output_hidden_states=True).hidden_states
+ normalized = {}
+ for layer in sorted({layer for layers in 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 TEXT_ENCODER_OUT_LAYERS]
+ return torch.stack(bundles, dim=1), text_inputs.attention_mask.bool()
+
+
+@torch.no_grad()
+def get_fibo2_edit_prompt_embeds(text_encoder, processor, instruction, cond_image, device, dtype, max_sequence_length):
+ """Encode a single edit example with the frozen Qwen3-VL encoder the way `BriaFibo2Pipeline` does for one image
+ to edit and a plain-text instruction: the prompt is the JSON object `{"image": , "edit_instruction":
+ }`, and Qwen3-VL sees the image to edit in place of the marker. The six output bundles are each
+ scaled to unit RMS per token in float32 and concatenate their three layers. Returns `(1, 6, seq_len, dim)`
+ embeddings and their bool token mask.
+
+ Encoding is per example (each example has its own image to edit), so the embeddings can be right-padded and
+ stacked into a batch afterwards -- Qwen3-VL masks padding, so this matches the pipeline.
+ """
+ edit = {"image": VISION_MARKER, "edit_instruction": instruction}
+ prompt = json.dumps(edit, separators=(",", ":"), ensure_ascii=False)
+ text = [PROMPT_TEMPLATE_ENCODE.format(prompt)]
+ return _encode_qwen3vl(text_encoder, processor, text, [cond_image], device, dtype, max_sequence_length)
+
+
+def pad_and_stack_embeds(embeds_list, mask_list, device):
+ """Right-pad a batch's per-example `(1, 6, seq_len, dim)` embeddings to the batch's longest prompt and stack
+ them, returning the batched embeddings and a bool mask (or `None` when nothing is padded). Right-padding and the
+ mask mirror the pipeline's `padding="longest"`; the transformer reads real token lengths from the mask to size
+ its gist tokens, so padding is free of effect."""
+ max_len = max(embeds.shape[2] for embeds in embeds_list)
+ padded_embeds, padded_masks = [], []
+ for embeds, mask in zip(embeds_list, mask_list):
+ seq_len = embeds.shape[2]
+ if seq_len < max_len:
+ embed_pad = embeds.new_zeros(embeds.shape[0], embeds.shape[1], max_len - seq_len, embeds.shape[3])
+ embeds = torch.cat([embeds, embed_pad], dim=2)
+ mask_pad = mask.new_zeros(mask.shape[0], max_len - seq_len)
+ mask = torch.cat([mask, mask_pad], dim=1)
+ padded_embeds.append(embeds)
+ padded_masks.append(mask)
+ embeds = torch.cat(padded_embeds, dim=0).to(device)
+ mask = torch.cat(padded_masks, dim=0).to(device)
+ return embeds, (None if mask.all() else mask)
+
+
+def save_model_card(
+ repo_id: str,
+ images=None,
+ base_model: str = None,
+ instance_prompt=None,
+ validation_prompt=None,
+ repo_folder=None,
+ quant_training=None,
+):
+ widget_dict = []
+ if images is not None:
+ for i, image in enumerate(images):
+ image.save(os.path.join(repo_folder, f"image_{i}.png"))
+ widget_dict.append(
+ {"text": validation_prompt if validation_prompt else " ", "output": {"url": f"image_{i}.png"}}
+ )
+
+ model_description = f"""
+# fibo-2 DreamBooth LoRA - {repo_id}
+
+
+
+## Model description
+
+These are {repo_id} DreamBooth LoRA weights for {base_model}.
+
+The weights were trained on **image editing** (image-to-image) using [DreamBooth](https://dreambooth.github.io/) with the [fibo-2 diffusers trainer](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_bria_fibo2.md).
+
+Quant training? {quant_training}
+
+## Trigger words
+
+You should use `{instance_prompt}` as the edit instruction.
+
+## Download model
+
+[Download the *.safetensors LoRA]({repo_id}/tree/main) in the Files & versions tab.
+
+## Use it with the [๐งจ diffusers library](https://github.com/huggingface/diffusers)
+
+```py
+from diffusers import BriaFibo2Pipeline
+from diffusers.utils import load_image
+import torch
+
+pipeline = BriaFibo2Pipeline.from_pretrained("{base_model}", torch_dtype=torch.bfloat16).to("cuda")
+pipeline.load_lora_weights("{repo_id}", weight_name="pytorch_lora_weights.safetensors")
+image = load_image("path/to/image.png")
+edited = pipeline("{validation_prompt if validation_prompt else instance_prompt}", image=image).images[0]
+```
+
+For more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters)
+
+## License
+
+Please adhere to the licensing terms as described in the model repository.
+"""
+ model_card = load_or_create_model_card(
+ repo_id_or_path=repo_id,
+ from_training=True,
+ license="other",
+ base_model=base_model,
+ prompt=instance_prompt,
+ model_description=model_description,
+ widget=widget_dict,
+ )
+ tags = [
+ "image-to-image",
+ "diffusers-training",
+ "diffusers",
+ "lora",
+ "bria",
+ "fibo-2",
+ "template:sd-lora",
+ ]
+
+ model_card = populate_model_card(model_card, tags=tags)
+ model_card.save(os.path.join(repo_folder, "README.md"))
+
+
+def log_validation(
+ pipeline,
+ args,
+ accelerator,
+ pipeline_args,
+ epoch,
+ torch_dtype,
+ is_final_validation=False,
+):
+ args.num_validation_images = args.num_validation_images if args.num_validation_images else 1
+ logger.info(
+ f"Running validation... \n Generating {args.num_validation_images} images with prompt:"
+ f" {args.validation_prompt}."
+ )
+ pipeline = pipeline.to(accelerator.device)
+ pipeline.set_progress_bar_config(disable=True)
+
+ # run inference
+ generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed is not None else None
+ autocast_ctx = torch.autocast(accelerator.device.type) if not is_final_validation else nullcontext()
+
+ images = []
+ for _ in range(args.num_validation_images):
+ with autocast_ctx:
+ image = pipeline(**pipeline_args, generator=generator).images[0]
+ images.append(image)
+
+ for tracker in accelerator.trackers:
+ phase_name = "test" if is_final_validation else "validation"
+ if tracker.name == "tensorboard":
+ np_images = np.stack([np.asarray(img) for img in images])
+ tracker.writer.add_images(phase_name, np_images, epoch, dataformats="NHWC")
+ if tracker.name == "wandb":
+ tracker.log(
+ {
+ phase_name: [
+ wandb.Image(image, caption=f"{i}: {args.validation_prompt}") for i, image in enumerate(images)
+ ]
+ }
+ )
+
+ del pipeline
+ free_memory()
+
+ return images
+
+
+def make_validation_pipeline(args, transformer, weight_dtype, device):
+ # Build a standard BriaFibo2Pipeline for validation that REUSES the in-memory trained transformer, so validation
+ # runs on exactly the weights being trained and the DiT is never loaded a second time. Only the small auxiliary
+ # components are (re)loaded (they are freed after latent caching). This mirrors how flux2_klein passes the
+ # unwrapped transformer into its validation pipeline.
+ vae = AutoencoderKLFlux2.from_pretrained(
+ args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision, torch_dtype=torch.float32
+ )
+ text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
+ args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, torch_dtype=weight_dtype
+ )
+ processor = Qwen3VLProcessor.from_pretrained(
+ args.pretrained_model_name_or_path, subfolder="processor", revision=args.revision
+ )
+ scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+ args.pretrained_model_name_or_path, subfolder="scheduler", revision=args.revision
+ )
+ pipeline = BriaFibo2Pipeline(
+ transformer=transformer,
+ scheduler=scheduler,
+ vae=vae,
+ text_encoder=text_encoder,
+ processor=processor,
+ )
+ pipeline.to(device)
+ return pipeline
+
+
+def module_filter_fn(mod: torch.nn.Module, fqn: str):
+ # don't convert the output projection
+ if fqn.endswith("final_layer.linear"):
+ return False
+ # don't convert linear modules with weight dimensions not divisible by 16
+ if isinstance(mod, torch.nn.Linear):
+ if mod.in_features % 16 != 0 or mod.out_features % 16 != 0:
+ return False
+ return True
+
+
+def parse_args(input_args=None):
+ parser = argparse.ArgumentParser(description="Simple example of a training script.")
+ parser.add_argument(
+ "--pretrained_model_name_or_path",
+ type=str,
+ default=None,
+ required=True,
+ help="Path to pretrained model or model identifier from huggingface.co/models.",
+ )
+ parser.add_argument(
+ "--revision",
+ type=str,
+ default=None,
+ required=False,
+ help="Revision of pretrained model identifier from huggingface.co/models.",
+ )
+ parser.add_argument(
+ "--bnb_quantization_config_path",
+ type=str,
+ default=None,
+ help="Quantization config in a JSON file that will be used to define the bitsandbytes quant config of the DiT.",
+ )
+ parser.add_argument(
+ "--do_fp8_training",
+ action="store_true",
+ help="if we are doing FP8 training.",
+ )
+ parser.add_argument(
+ "--variant",
+ type=str,
+ default=None,
+ help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
+ )
+ parser.add_argument(
+ "--dataset_name",
+ type=str,
+ default=None,
+ help=(
+ "The name of the Dataset (from the HuggingFace hub) containing the training data of instance images (could be your own, possibly private,"
+ " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,"
+ " or to a folder containing files that ๐ค Datasets can understand."
+ ),
+ )
+ parser.add_argument(
+ "--dataset_config_name",
+ type=str,
+ default=None,
+ help="The config of the Dataset, leave as None if there's only one config.",
+ )
+ parser.add_argument(
+ "--instance_data_dir",
+ type=str,
+ default=None,
+ help=("A folder containing the training data. "),
+ )
+
+ parser.add_argument(
+ "--cache_dir",
+ type=str,
+ default=None,
+ help="The directory where the downloaded models and datasets will be stored.",
+ )
+
+ parser.add_argument(
+ "--image_column",
+ type=str,
+ default="image",
+ help="The column of the dataset containing the target (edited) image.",
+ )
+ parser.add_argument(
+ "--cond_image_column",
+ type=str,
+ default=None,
+ help="The column of the dataset containing the conditioning (source) image to edit.",
+ )
+ parser.add_argument(
+ "--caption_column",
+ type=str,
+ default="text",
+ help="The column of the dataset containing the edit instruction for each pair, as plain text.",
+ )
+
+ parser.add_argument("--repeats", type=int, default=1, help="How many times to repeat the training data.")
+
+ parser.add_argument(
+ "--instance_prompt",
+ type=str,
+ default=None,
+ required=True,
+ help="The edit instruction, used for every pair when --caption_column is not provided.",
+ )
+ parser.add_argument(
+ "--max_sequence_length",
+ type=int,
+ default=512,
+ help="Cap on the number of Qwen3-VL prompt tokens; longer prompts are truncated.",
+ )
+ parser.add_argument(
+ "--validation_prompt",
+ type=str,
+ default=None,
+ help="An edit instruction that is used during validation to verify that the model is learning.",
+ )
+ parser.add_argument(
+ "--validation_image",
+ type=str,
+ default=None,
+ help="Path or URL of the image to edit during validation with `--validation_prompt`.",
+ )
+ parser.add_argument(
+ "--skip_final_inference",
+ default=False,
+ action="store_true",
+ help="Whether to skip the final inference step with loaded lora weights upon training completion. This will run intermediate validation inference if `validation_prompt` is provided. Specify to reduce memory.",
+ )
+ parser.add_argument(
+ "--num_validation_images",
+ type=int,
+ default=4,
+ help="Number of images that should be generated during validation with `validation_prompt`.",
+ )
+ parser.add_argument(
+ "--validation_epochs",
+ type=int,
+ default=50,
+ help=(
+ "Run dreambooth validation every X epochs. Dreambooth validation consists of running the prompt"
+ " `args.validation_prompt` multiple times: `args.num_validation_images`."
+ ),
+ )
+ parser.add_argument(
+ "--rank",
+ type=int,
+ default=4,
+ help=("The dimension of the LoRA update matrices."),
+ )
+ parser.add_argument(
+ "--lora_alpha",
+ type=int,
+ default=4,
+ help="LoRA alpha to be used for additional scaling.",
+ )
+ parser.add_argument("--lora_dropout", type=float, default=0.0, help="Dropout probability for LoRA layers")
+
+ parser.add_argument(
+ "--output_dir",
+ type=str,
+ default="fibo2-dreambooth-lora",
+ help="The output directory where the model predictions and checkpoints will be written.",
+ )
+ parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+ parser.add_argument(
+ "--resolution",
+ type=int,
+ default=512,
+ help=(
+ "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+ " resolution"
+ ),
+ )
+ parser.add_argument(
+ "--aspect_ratio_buckets",
+ type=str,
+ default=None,
+ help=(
+ "Aspect ratio buckets to use for training. Define as a string of 'h1,w1;h2,w2;...'. "
+ "e.g. '1024,1024;768,1360;1360,768;880,1168;1168,880;1248,832;832,1248'. "
+ "Requires --use_aspect_ratio_buckets. Images are resized to cover and cropped to the nearest "
+ "listed bucket (smaller images are upscaled). When set, --resolution is ignored."
+ ),
+ )
+ parser.add_argument(
+ "--use_aspect_ratio_buckets",
+ action="store_true",
+ help=(
+ "Enable aspect-ratio bucketing. Without --aspect_ratio_buckets, the buckets are computed on the "
+ "fly from --resolution and capped to each image's own resolution, so smaller images are assigned "
+ "to a smaller bucket instead of being upscaled. Provide --aspect_ratio_buckets to use an explicit list."
+ ),
+ )
+ parser.add_argument(
+ "--center_crop",
+ default=False,
+ action="store_true",
+ help=(
+ "Whether to center crop the input images to the resolution. If not set, the images will be randomly"
+ " cropped. The images will be resized to the resolution first before cropping. The same geometry is"
+ " applied to the conditioning image so each pair stays aligned."
+ ),
+ )
+ parser.add_argument(
+ "--random_flip",
+ action="store_true",
+ help="whether to randomly flip images horizontally (the conditioning image is flipped with its target)",
+ )
+ parser.add_argument(
+ "--caption_dropout",
+ type=float,
+ default=0.0,
+ help=(
+ "Probability of replacing a pair's edit instruction with an empty string during training, so that"
+ " fraction of samples is trained unconditionally. Improves classifier-free guidance. A common value is"
+ " 0.1. The conditioning image is always kept."
+ ),
+ )
+ parser.add_argument(
+ "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
+ )
+ parser.add_argument("--num_train_epochs", type=int, default=1)
+ parser.add_argument(
+ "--max_train_steps",
+ type=int,
+ default=None,
+ help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+ )
+ parser.add_argument(
+ "--checkpointing_steps",
+ type=int,
+ default=500,
+ help=(
+ "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+ " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+ " training using `--resume_from_checkpoint`."
+ ),
+ )
+ parser.add_argument(
+ "--checkpoints_total_limit",
+ type=int,
+ default=None,
+ help=("Max number of checkpoints to store."),
+ )
+ parser.add_argument(
+ "--resume_from_checkpoint",
+ type=str,
+ default=None,
+ help=(
+ "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+ ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+ ),
+ )
+ parser.add_argument(
+ "--gradient_accumulation_steps",
+ type=int,
+ default=1,
+ help="Number of updates steps to accumulate before performing a backward/update pass.",
+ )
+ parser.add_argument(
+ "--gradient_checkpointing",
+ action="store_true",
+ help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+ )
+ parser.add_argument(
+ "--learning_rate",
+ type=float,
+ default=1e-4,
+ help="Initial learning rate (after the potential warmup period) to use.",
+ )
+
+ parser.add_argument(
+ "--scale_lr",
+ action="store_true",
+ default=False,
+ help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+ )
+ parser.add_argument(
+ "--lr_scheduler",
+ type=str,
+ default="constant",
+ help=(
+ 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+ ' "constant", "constant_with_warmup"]'
+ ),
+ )
+ parser.add_argument(
+ "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+ )
+ parser.add_argument(
+ "--lr_num_cycles",
+ type=int,
+ default=1,
+ help="Number of hard resets of the lr in cosine_with_restarts scheduler.",
+ )
+ parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.")
+ parser.add_argument(
+ "--dataloader_num_workers",
+ type=int,
+ default=0,
+ help=(
+ "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+ ),
+ )
+ parser.add_argument(
+ "--weighting_scheme",
+ type=str,
+ default="none",
+ choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"],
+ help=('We default to the "none" weighting scheme for uniform sampling and uniform loss'),
+ )
+ parser.add_argument(
+ "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme."
+ )
+ parser.add_argument(
+ "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme."
+ )
+ parser.add_argument(
+ "--mode_scale",
+ type=float,
+ default=1.29,
+ help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
+ )
+ parser.add_argument(
+ "--optimizer",
+ type=str,
+ default="AdamW",
+ help=('The optimizer type to use. Choose between ["AdamW", "prodigy"]'),
+ )
+
+ parser.add_argument(
+ "--use_8bit_adam",
+ action="store_true",
+ help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW",
+ )
+
+ parser.add_argument(
+ "--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers."
+ )
+ parser.add_argument(
+ "--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers."
+ )
+ parser.add_argument(
+ "--prodigy_beta3",
+ type=float,
+ default=None,
+ help="coefficients for computing the Prodigy stepsize using running averages. If set to None, "
+ "uses the value of square root of beta2. Ignored if optimizer is adamW",
+ )
+ parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay")
+ parser.add_argument("--adam_weight_decay", type=float, default=1e-04, help="Weight decay to use for unet params")
+
+ parser.add_argument(
+ "--lora_layers",
+ type=str,
+ default=None,
+ help=(
+ 'The transformer modules to apply LoRA training on. Please specify the layers in a comma separated. E.g. - "to_k,to_q,to_v,to_out.0" will result in lora training of attention layers only'
+ ),
+ )
+
+ parser.add_argument(
+ "--adam_epsilon",
+ type=float,
+ default=1e-08,
+ help="Epsilon value for the Adam optimizer and Prodigy optimizers.",
+ )
+
+ parser.add_argument(
+ "--prodigy_use_bias_correction",
+ type=bool,
+ default=True,
+ help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW",
+ )
+ parser.add_argument(
+ "--prodigy_safeguard_warmup",
+ type=bool,
+ default=True,
+ help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. "
+ "Ignored if optimizer is adamW",
+ )
+ parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+ parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+ parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
+ parser.add_argument(
+ "--hub_model_id",
+ type=str,
+ default=None,
+ help="The name of the repository to keep in sync with the local `output_dir`.",
+ )
+ parser.add_argument(
+ "--logging_dir",
+ type=str,
+ default="logs",
+ help=(
+ "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+ " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+ ),
+ )
+ parser.add_argument(
+ "--allow_tf32",
+ action="store_true",
+ help=(
+ "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+ " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+ ),
+ )
+ parser.add_argument(
+ "--cache_latents",
+ action="store_true",
+ default=False,
+ help="Cache the VAE latents",
+ )
+ parser.add_argument(
+ "--report_to",
+ type=str,
+ default="tensorboard",
+ help=(
+ 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+ ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+ ),
+ )
+ parser.add_argument(
+ "--mixed_precision",
+ type=str,
+ default=None,
+ choices=["no", "fp16", "bf16"],
+ help=(
+ "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+ " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the"
+ " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+ ),
+ )
+ parser.add_argument(
+ "--upcast_before_saving",
+ action="store_true",
+ default=False,
+ help=(
+ "Whether to upcast the trained transformer layers to float32 before saving (at the end of training). "
+ "Defaults to precision dtype used for training to save memory"
+ ),
+ )
+ parser.add_argument(
+ "--offload",
+ action="store_true",
+ help="Whether to offload the VAE and the text encoder to CPU when they are not used.",
+ )
+ parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
+ parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU")
+ parser.add_argument("--fsdp_text_encoder", action="store_true", help="Use FSDP for text encoder")
+
+ if input_args is not None:
+ args = parser.parse_args(input_args)
+ else:
+ args = parser.parse_args()
+
+ if args.cond_image_column is None:
+ raise ValueError(
+ "You must provide --cond_image_column for image editing (image-to-image) training. For plain "
+ "text-to-image training, see train_dreambooth_lora_bria_fibo2.py instead."
+ )
+
+ if args.dataset_name is None and args.instance_data_dir is None:
+ raise ValueError("Specify either `--dataset_name` or `--instance_data_dir`")
+
+ if args.dataset_name is not None and args.instance_data_dir is not None:
+ raise ValueError("Specify only one of `--dataset_name` or `--instance_data_dir`")
+ if args.do_fp8_training and args.bnb_quantization_config_path:
+ raise ValueError("Both `do_fp8_training` and `bnb_quantization_config_path` cannot be passed.")
+
+ env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
+ if env_local_rank != -1 and env_local_rank != args.local_rank:
+ args.local_rank = env_local_rank
+
+ return args
+
+
+class DreamBoothDataset(Dataset):
+ """
+ A dataset to prepare the (conditioning image, target image, edit instruction) triplets for fine-tuning the
+ model. It pre-processes the images: the conditioning and target images are given the same geometry so each pair
+ stays aligned, and a thumbnail of the conditioning image is kept for the Qwen3-VL text encoder.
+ """
+
+ def __init__(
+ self,
+ instance_data_root,
+ instance_prompt,
+ size=1024,
+ repeats=1,
+ center_crop=False,
+ buckets=None,
+ use_aspect_ratio_buckets=False,
+ bucket_divisibility=16,
+ bucket_base_resolutions=None,
+ ):
+ self.size = size
+ self.resolution = size
+ self.center_crop = center_crop
+
+ self.instance_prompt = instance_prompt
+ self.custom_instance_prompts = None
+
+ # Explicit user-provided bucket list (or None). The concrete list of buckets actually used is
+ # built from the data in `self.buckets` during preprocessing below.
+ self._explicit_buckets = buckets
+ self.use_aspect_ratio_buckets = use_aspect_ratio_buckets
+ self.bucket_divisibility = bucket_divisibility
+ self.bucket_base_resolutions = bucket_base_resolutions
+
+ if args.dataset_name is not None:
+ try:
+ from datasets import load_dataset
+ except ImportError:
+ raise ImportError(
+ "You are trying to load your data using the datasets library. If you wish to train using custom "
+ "captions please install the datasets library: `pip install datasets`. If you wish to load a "
+ "local folder containing images only, specify --instance_data_dir instead."
+ )
+ dataset = load_dataset(
+ args.dataset_name,
+ args.dataset_config_name,
+ cache_dir=args.cache_dir,
+ )
+ column_names = dataset["train"].column_names
+
+ if args.image_column not in column_names:
+ raise ValueError(
+ f"`--image_column` value '{args.image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
+ )
+ if args.cond_image_column not in column_names:
+ raise ValueError(
+ f"`--cond_image_column` value '{args.cond_image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
+ )
+ instance_images = dataset["train"][args.image_column]
+ cond_images = dataset["train"][args.cond_image_column]
+ assert len(instance_images) == len(cond_images)
+
+ if args.caption_column is None or args.caption_column not in column_names:
+ logger.info(
+ "No caption column provided, defaulting to instance_prompt (the edit instruction) for all pairs."
+ )
+ self.custom_instance_prompts = None
+ else:
+ custom_instance_prompts = dataset["train"][args.caption_column]
+ self.custom_instance_prompts = []
+ for caption in custom_instance_prompts:
+ self.custom_instance_prompts.extend(itertools.repeat(caption, repeats))
+ else:
+ raise ValueError(
+ "Image editing training needs paired (conditioning, target) images with instructions, so pass a "
+ "`--dataset_name` with `--image_column`, `--cond_image_column` and `--caption_column`."
+ )
+
+ self.instance_images = []
+ self.cond_images = []
+ for img in instance_images:
+ self.instance_images.extend(itertools.repeat(img, repeats))
+ for img in cond_images:
+ self.cond_images.extend(itertools.repeat(img, repeats))
+
+ self.pixel_values = []
+ self.cond_pixel_values = []
+ self.cond_text_images = []
+ self.buckets = []
+ bucket_to_idx = {}
+ for i, image in enumerate(self.instance_images):
+ image = exif_transpose(image)
+ if not image.mode == "RGB":
+ image = image.convert("RGB")
+ cond_image = exif_transpose(self.cond_images[i])
+ if not cond_image.mode == "RGB":
+ cond_image = cond_image.convert("RGB")
+
+ width, height = image.size
+
+ # Assign the pair to a bucket by the TARGET image's resolution.
+ target = self._bucket_for_image(height, width)
+ if target not in bucket_to_idx:
+ bucket_to_idx[target] = len(self.buckets)
+ self.buckets.append(target)
+ bucket_idx = bucket_to_idx[target]
+
+ image_tensor, cond_tensor, cond_crop = self.paired_transform(
+ image,
+ cond_image,
+ size=target,
+ center_crop=args.center_crop,
+ random_flip=args.random_flip,
+ )
+ # The Qwen3-VL text encoder sees a thumbnail of the (aligned) conditioning crop, at most 768px a side.
+ cond_thumbnail = cond_crop.copy()
+ cond_thumbnail.thumbnail((TEXT_ENCODER_IMAGE_SIZE, TEXT_ENCODER_IMAGE_SIZE))
+ self.pixel_values.append((image_tensor, bucket_idx))
+ self.cond_pixel_values.append(cond_tensor)
+ self.cond_text_images.append(cond_thumbnail)
+
+ self.num_instance_images = len(self.instance_images)
+ self._length = self.num_instance_images
+
+ def __len__(self):
+ return self._length
+
+ def __getitem__(self, index):
+ example = {}
+ instance_image, bucket_idx = self.pixel_values[index % self.num_instance_images]
+ example["index"] = index
+ example["instance_images"] = instance_image
+ example["cond_images"] = self.cond_pixel_values[index % self.num_instance_images]
+ example["cond_text_images"] = self.cond_text_images[index % self.num_instance_images]
+ example["bucket_idx"] = bucket_idx
+ if self.custom_instance_prompts:
+ caption = self.custom_instance_prompts[index % self.num_instance_images]
+ example["instance_prompt"] = caption if caption else self.instance_prompt
+ else:
+ example["instance_prompt"] = self.instance_prompt
+ return example
+
+ def _bucket_for_image(self, height, width):
+ # An explicit bucket list takes priority: pick the nearest, upscaling smaller images to cover it.
+ if self._explicit_buckets is not None:
+ return self._explicit_buckets[find_nearest_bucket(height, width, self._explicit_buckets)]
+ # On-the-fly bucketing: cap the ladder to the image's own resolution so smaller images are
+ # assigned to a smaller bucket rather than being upscaled (mirrors ostris' bucketing).
+ if self.use_aspect_ratio_buckets:
+ resolution = min(self.resolution, round((height * width) ** 0.5))
+ ladder = generate_aspect_ratio_buckets(
+ resolution,
+ divisibility=self.bucket_divisibility,
+ base_resolutions=self.bucket_base_resolutions,
+ )
+ return ladder[find_nearest_bucket(height, width, ladder)]
+ # No bucketing: a single square bucket reproduces the fixed-size resize + crop.
+ return (self.resolution, self.resolution)
+
+ def paired_transform(self, image, cond_image, size, center_crop=False, random_flip=False):
+ # Resize preserving aspect ratio so the image covers the bucket, then crop to the bucket size. The same
+ # geometry is applied to the conditioning image so the pair stays aligned. Returns the two normalized tensors
+ # and the cropped conditioning image as a PIL image (for the Qwen3-VL thumbnail).
+ target_height, target_width = size
+ width, height = image.size
+ scale = max(target_height / height, target_width / width)
+ new_size = [round(height * scale), round(width * scale)]
+ image = TF.resize(image, new_size, interpolation=transforms.InterpolationMode.BILINEAR)
+ cond_image = TF.resize(cond_image, new_size, interpolation=transforms.InterpolationMode.BILINEAR)
+ if center_crop:
+ image = TF.center_crop(image, size)
+ cond_image = TF.center_crop(cond_image, size)
+ else:
+ i, j, h, w = transforms.RandomCrop.get_params(image, output_size=size)
+ image = TF.crop(image, i, j, h, w)
+ cond_image = TF.crop(cond_image, i, j, h, w)
+ if random_flip and random.random() < 0.5:
+ image = TF.hflip(image)
+ cond_image = TF.hflip(cond_image)
+ image_tensor = TF.normalize(TF.to_tensor(image), [0.5], [0.5])
+ cond_tensor = TF.normalize(TF.to_tensor(cond_image), [0.5], [0.5])
+ return image_tensor, cond_tensor, cond_image
+
+
+def collate_fn(examples):
+ indices = [example["index"] for example in examples]
+ pixel_values = [example["instance_images"] for example in examples]
+ cond_pixel_values = [example["cond_images"] for example in examples]
+ cond_text_images = [example["cond_text_images"] for example in examples]
+ prompts = [example["instance_prompt"] for example in examples]
+
+ pixel_values = torch.stack(pixel_values).to(memory_format=torch.contiguous_format).float()
+ cond_pixel_values = torch.stack(cond_pixel_values).to(memory_format=torch.contiguous_format).float()
+
+ batch = {
+ "indices": indices,
+ "pixel_values": pixel_values,
+ "cond_pixel_values": cond_pixel_values,
+ "cond_text_images": cond_text_images,
+ "prompts": prompts,
+ }
+ return batch
+
+
+class BucketBatchSampler(BatchSampler):
+ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, seed: int = None):
+ if not isinstance(batch_size, int) or batch_size <= 0:
+ raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size))
+ if not isinstance(drop_last, bool):
+ raise ValueError("drop_last should be a boolean value, but got drop_last={}".format(drop_last))
+
+ self.dataset = dataset
+ self.batch_size = batch_size
+ self.drop_last = drop_last
+ self.generator = random.Random(seed) if seed is not None else random
+
+ # Group indices by bucket
+ self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))]
+ for idx, (_, bucket_idx) in enumerate(self.dataset.pixel_values):
+ self.bucket_indices[bucket_idx].append(idx)
+
+ self.sampler_len = 0
+ for indices_in_bucket in self.bucket_indices:
+ num_batches, remainder = divmod(len(indices_in_bucket), self.batch_size)
+ self.sampler_len += num_batches
+ if remainder > 0 and not self.drop_last:
+ self.sampler_len += 1
+
+ def __iter__(self):
+ batches = []
+ for indices_in_bucket in self.bucket_indices:
+ shuffled_indices = indices_in_bucket.copy()
+ self.generator.shuffle(shuffled_indices)
+ for i in range(0, len(shuffled_indices), self.batch_size):
+ batch = shuffled_indices[i : i + self.batch_size]
+ if len(batch) < self.batch_size and self.drop_last:
+ continue
+ batches.append(batch)
+
+ self.generator.shuffle(batches)
+ for batch in batches:
+ yield batch
+
+ def __len__(self):
+ return self.sampler_len
+
+
+def main(args):
+ if args.report_to == "wandb" and args.hub_token is not None:
+ raise ValueError(
+ "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token."
+ " Please use `hf auth login` to authenticate with the Hub."
+ )
+
+ if torch.backends.mps.is_available() and args.mixed_precision == "bf16":
+ # due to pytorch#99272, MPS does not yet support bfloat16.
+ raise ValueError(
+ "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
+ )
+ if args.do_fp8_training:
+ from torchao.float8 import Float8LinearConfig, convert_to_float8_training
+
+ logging_dir = Path(args.output_dir, args.logging_dir)
+
+ accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+ kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+ accelerator = Accelerator(
+ gradient_accumulation_steps=args.gradient_accumulation_steps,
+ mixed_precision=args.mixed_precision,
+ log_with=args.report_to,
+ project_config=accelerator_project_config,
+ kwargs_handlers=[kwargs],
+ )
+
+ # Disable AMP for MPS.
+ if torch.backends.mps.is_available():
+ accelerator.native_amp = False
+
+ if args.report_to == "wandb":
+ if not is_wandb_available():
+ raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+
+ # Make one log on every process with the configuration for debugging.
+ logging.basicConfig(
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+ datefmt="%m/%d/%Y %H:%M:%S",
+ level=logging.INFO,
+ )
+ logger.info(accelerator.state, main_process_only=False)
+ if accelerator.is_local_main_process:
+ transformers.utils.logging.set_verbosity_warning()
+ diffusers.utils.logging.set_verbosity_info()
+ else:
+ transformers.utils.logging.set_verbosity_error()
+ diffusers.utils.logging.set_verbosity_error()
+
+ # If passed along, set the training seed now.
+ if args.seed is not None:
+ set_seed(args.seed)
+
+ # Handle the repository creation
+ if accelerator.is_main_process:
+ if args.output_dir is not None:
+ os.makedirs(args.output_dir, exist_ok=True)
+
+ if args.push_to_hub:
+ repo_id = create_repo(
+ repo_id=args.hub_model_id or Path(args.output_dir).name,
+ exist_ok=True,
+ ).repo_id
+
+ # Load the Qwen3-VL processor (tokenizer + image processor)
+ processor = Qwen3VLProcessor.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="processor",
+ revision=args.revision,
+ )
+
+ # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision
+ # as these weights are only used for inference, keeping weights in full precision is not required.
+ weight_dtype = torch.float32
+ if accelerator.mixed_precision == "fp16":
+ weight_dtype = torch.float16
+ elif accelerator.mixed_precision == "bf16":
+ weight_dtype = torch.bfloat16
+
+ # Load scheduler and models. fibo-2 shares FLUX.2's flow-matching scheduler (dynamic shift), loaded here so the
+ # training noise schedule matches what the model was trained and is sampled with.
+ noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="scheduler",
+ revision=args.revision,
+ )
+ noise_scheduler_copy = copy.deepcopy(noise_scheduler)
+ vae = AutoencoderKLFlux2.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="vae",
+ revision=args.revision,
+ variant=args.variant,
+ )
+ latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(accelerator.device)
+ latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(
+ accelerator.device
+ )
+
+ quantization_config = None
+ if args.bnb_quantization_config_path is not None:
+ with open(args.bnb_quantization_config_path, "r") as f:
+ config_kwargs = json.load(f)
+ if "load_in_4bit" in config_kwargs and config_kwargs["load_in_4bit"]:
+ config_kwargs["bnb_4bit_compute_dtype"] = weight_dtype
+ quantization_config = BitsAndBytesConfig(**config_kwargs)
+
+ transformer = BriaFibo2Transformer2DModel.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="transformer",
+ revision=args.revision,
+ variant=args.variant,
+ quantization_config=quantization_config,
+ torch_dtype=weight_dtype,
+ )
+ if args.bnb_quantization_config_path is not None:
+ transformer = prepare_model_for_kbit_training(transformer, use_gradient_checkpointing=False)
+
+ text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="text_encoder",
+ revision=args.revision,
+ variant=args.variant,
+ torch_dtype=weight_dtype,
+ )
+ text_encoder.requires_grad_(False)
+
+ # We only train the additional adapter LoRA layers
+ transformer.requires_grad_(False)
+ vae.requires_grad_(False)
+
+ if args.enable_npu_flash_attention:
+ if is_torch_npu_available():
+ logger.info("npu flash attention enabled.")
+ transformer.set_attention_backend("_native_npu")
+ else:
+ raise ValueError("npu flash attention requires torch_npu extensions and is supported only on npu device ")
+
+ if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:
+ # due to pytorch#99272, MPS does not yet support bfloat16.
+ raise ValueError(
+ "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
+ )
+
+ to_kwargs = {"dtype": weight_dtype, "device": accelerator.device} if not args.offload else {"dtype": weight_dtype}
+ # flux vae is stable in bf16 so load it in weight_dtype to reduce memory
+ vae.to(**to_kwargs)
+ # we never offload the transformer to CPU, so we can just use the accelerator device
+ transformer_to_kwargs = (
+ {"device": accelerator.device}
+ if args.bnb_quantization_config_path is not None
+ else {"device": accelerator.device, "dtype": weight_dtype}
+ )
+
+ is_fsdp = getattr(accelerator.state, "fsdp_plugin", None) is not None
+ if not is_fsdp:
+ transformer.to(**transformer_to_kwargs)
+
+ if args.do_fp8_training:
+ convert_to_float8_training(
+ transformer, module_filter_fn=module_filter_fn, config=Float8LinearConfig(pad_inner_dim=True)
+ )
+
+ text_encoder.to(**to_kwargs)
+
+ if args.gradient_checkpointing:
+ transformer.enable_gradient_checkpointing()
+
+ if args.lora_layers is not None:
+ target_modules = [layer.strip() for layer in args.lora_layers.split(",")]
+ else:
+ # Attention (self- and cross-attention) + SwiGLU feed-forward projections across every block. The modulation
+ # (`adaLN_modulation.linear`), patch embedding (`x_embedder`) and output projection (`final_layer.linear`)
+ # are left out, matching the attention+FFN set the sibling trainers default to. Pass `--lora_layers` to
+ # target a different set.
+ target_modules = [
+ "to_q",
+ "to_k",
+ "to_v",
+ "to_out.0",
+ "w1",
+ "w2",
+ "w3",
+ ]
+
+ # now we will add new LoRA weights the transformer layers
+ transformer_lora_config = LoraConfig(
+ r=args.rank,
+ lora_alpha=args.lora_alpha,
+ lora_dropout=args.lora_dropout,
+ init_lora_weights="gaussian",
+ target_modules=target_modules,
+ )
+ transformer.add_adapter(transformer_lora_config)
+
+ def unwrap_model(model):
+ model = accelerator.unwrap_model(model)
+ model = model._orig_mod if is_compiled_module(model) else model
+ return model
+
+ # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
+ def save_model_hook(models, weights, output_dir):
+ transformer_cls = type(unwrap_model(transformer))
+
+ # 1) Validate and pick the transformer model
+ modules_to_save: dict[str, Any] = {}
+ transformer_model = None
+
+ for model in models:
+ if isinstance(unwrap_model(model), transformer_cls):
+ transformer_model = model
+ modules_to_save["transformer"] = model
+ else:
+ raise ValueError(f"unexpected save model: {model.__class__}")
+
+ if transformer_model is None:
+ raise ValueError("No transformer model found in 'models'")
+
+ # 2) Optionally gather FSDP state dict once
+ state_dict = accelerator.get_state_dict(model) if is_fsdp else None
+
+ # 3) Only main process materializes the LoRA state dict
+ transformer_lora_layers_to_save = None
+ if accelerator.is_main_process:
+ peft_kwargs = {}
+ if is_fsdp:
+ peft_kwargs["state_dict"] = state_dict
+
+ transformer_lora_layers_to_save = get_peft_model_state_dict(
+ unwrap_model(transformer_model) if is_fsdp else transformer_model,
+ **peft_kwargs,
+ )
+
+ if is_fsdp:
+ transformer_lora_layers_to_save = _to_cpu_contiguous(transformer_lora_layers_to_save)
+
+ # make sure to pop weight so that corresponding model is not saved again
+ if weights:
+ weights.pop()
+
+ BriaFibo2LoraLoaderMixin.save_lora_weights(
+ output_dir,
+ transformer_lora_layers=transformer_lora_layers_to_save,
+ **_collate_lora_metadata(modules_to_save),
+ )
+
+ def load_model_hook(models, input_dir):
+ transformer_ = None
+
+ if not is_fsdp:
+ while len(models) > 0:
+ model = models.pop()
+
+ if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
+ transformer_ = unwrap_model(model)
+ else:
+ raise ValueError(f"unexpected save model: {model.__class__}")
+ else:
+ transformer_ = BriaFibo2Transformer2DModel.from_pretrained(
+ args.pretrained_model_name_or_path,
+ subfolder="transformer",
+ )
+ transformer_.add_adapter(transformer_lora_config)
+
+ lora_state_dict = BriaFibo2LoraLoaderMixin.lora_state_dict(input_dir)
+
+ transformer_state_dict = {
+ f"{k.replace('transformer.', '')}": v for k, v in lora_state_dict.items() if k.startswith("transformer.")
+ }
+ transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict)
+ incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default")
+ if incompatible_keys is not None:
+ # check only for unexpected keys
+ unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
+ if unexpected_keys:
+ logger.warning(
+ f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
+ f" {unexpected_keys}. "
+ )
+
+ # Make sure the trainable params are in float32. This is again needed since the base models
+ # are in `weight_dtype`. More details:
+ # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
+ if args.mixed_precision == "fp16":
+ models = [transformer_]
+ # only upcast trainable parameters (LoRA) into fp32
+ cast_training_params(models)
+
+ accelerator.register_save_state_pre_hook(save_model_hook)
+ accelerator.register_load_state_pre_hook(load_model_hook)
+
+ # Enable TF32 for faster training on Ampere GPUs,
+ # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+ if args.allow_tf32 and torch.cuda.is_available():
+ torch.backends.cuda.matmul.allow_tf32 = True
+
+ if args.scale_lr:
+ args.learning_rate = (
+ args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
+ )
+
+ # Make sure the trainable params are in float32.
+ if args.mixed_precision == "fp16":
+ models = [transformer]
+ # only upcast trainable parameters (LoRA) into fp32
+ cast_training_params(models, dtype=torch.float32)
+
+ transformer_lora_parameters = list(filter(lambda p: p.requires_grad, transformer.parameters()))
+
+ # Optimization parameters
+ transformer_parameters_with_lr = {"params": transformer_lora_parameters, "lr": args.learning_rate}
+ params_to_optimize = [transformer_parameters_with_lr]
+
+ # Optimizer creation
+ if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"):
+ logger.warning(
+ f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]."
+ "Defaulting to adamW"
+ )
+ args.optimizer = "adamw"
+
+ if args.use_8bit_adam and not args.optimizer.lower() == "adamw":
+ logger.warning(
+ f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was "
+ f"set to {args.optimizer.lower()}"
+ )
+
+ if args.optimizer.lower() == "adamw":
+ if args.use_8bit_adam:
+ try:
+ import bitsandbytes as bnb
+ except ImportError:
+ raise ImportError(
+ "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+ )
+
+ optimizer_class = bnb.optim.AdamW8bit
+ else:
+ optimizer_class = torch.optim.AdamW
+
+ optimizer = optimizer_class(
+ params_to_optimize,
+ betas=(args.adam_beta1, args.adam_beta2),
+ weight_decay=args.adam_weight_decay,
+ eps=args.adam_epsilon,
+ )
+
+ if args.optimizer.lower() == "prodigy":
+ try:
+ import prodigyopt
+ except ImportError:
+ raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`")
+
+ optimizer_class = prodigyopt.Prodigy
+
+ if args.learning_rate <= 0.1:
+ logger.warning(
+ "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0"
+ )
+
+ optimizer = optimizer_class(
+ params_to_optimize,
+ betas=(args.adam_beta1, args.adam_beta2),
+ beta3=args.prodigy_beta3,
+ weight_decay=args.adam_weight_decay,
+ eps=args.adam_epsilon,
+ decouple=args.prodigy_decouple,
+ use_bias_correction=args.prodigy_use_bias_correction,
+ safeguard_warmup=args.prodigy_safeguard_warmup,
+ )
+
+ # Resolve the bucketing mode. Bucketing must be enabled explicitly with --use_aspect_ratio_buckets;
+ # a bucket list without that flag is an error. With the flag, an explicit --aspect_ratio_buckets list
+ # drives assignment, otherwise buckets are computed on the fly inside the dataset. Without the flag a
+ # single square bucket reproduces the fixed-size resize + crop.
+ if args.aspect_ratio_buckets is not None and not args.use_aspect_ratio_buckets:
+ raise ValueError("--aspect_ratio_buckets requires --use_aspect_ratio_buckets to be set.")
+ if args.aspect_ratio_buckets is not None:
+ buckets = parse_buckets_string(args.aspect_ratio_buckets)
+ use_aspect_ratio_buckets = False
+ logger.info(f"Using explicit aspect ratio buckets: {buckets}")
+ elif args.use_aspect_ratio_buckets:
+ buckets = None
+ use_aspect_ratio_buckets = True
+ logger.info(
+ "No --aspect_ratio_buckets provided; auto-computing aspect ratio buckets on the fly from --resolution."
+ )
+ else:
+ buckets = [(args.resolution, args.resolution)]
+ use_aspect_ratio_buckets = False
+
+ # Dataset and DataLoaders creation:
+ train_dataset = DreamBoothDataset(
+ instance_data_root=args.instance_data_dir,
+ instance_prompt=args.instance_prompt,
+ size=args.resolution,
+ repeats=args.repeats,
+ center_crop=args.center_crop,
+ buckets=buckets,
+ use_aspect_ratio_buckets=use_aspect_ratio_buckets,
+ )
+ batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True, seed=args.seed)
+ train_dataloader = torch.utils.data.DataLoader(
+ train_dataset,
+ batch_sampler=batch_sampler,
+ collate_fn=collate_fn,
+ num_workers=args.dataloader_num_workers,
+ )
+
+ # fibo-2 validation runs through the standard pipeline (which encodes the prompt and the image to edit itself),
+ # so we only assemble the generation kwargs here rather than precomputing embeddings.
+ validation_pipeline_args = {}
+ if args.validation_prompt is not None and args.validation_image is not None:
+ validation_pipeline_args = {
+ "prompt": args.validation_prompt,
+ "image": load_image(args.validation_image).convert("RGB"),
+ "num_inference_steps": 30,
+ "guidance_scale": 5.0,
+ }
+
+ # Init FSDP for text encoder
+ if args.fsdp_text_encoder:
+ fsdp_kwargs = get_fsdp_kwargs_from_accelerator(accelerator)
+ text_encoder = wrap_with_fsdp(
+ model=text_encoder,
+ device=accelerator.device,
+ offload=args.offload,
+ limit_all_gathers=True,
+ use_orig_params=True,
+ fsdp_kwargs=fsdp_kwargs,
+ )
+ dist.barrier()
+
+ # The edit instruction's text embeddings depend on the conditioning image (Qwen3-VL sees it), so they are
+ # computed per example and cached here, together with the conditioning/target VAE latents when --cache_latents is
+ # set. When --caption_dropout is on, the empty-instruction embeddings (still with the conditioning image) are
+ # cached too, so a dropped sample can be swapped to them.
+ instance_latents_cache = [None] * train_dataset.num_instance_images if args.cache_latents else None
+ cond_latents_cache = [None] * train_dataset.num_instance_images if args.cache_latents else None
+ prompt_embeds_cache = [None] * train_dataset.num_instance_images
+ prompt_mask_cache = [None] * train_dataset.num_instance_images
+ empty_prompt_embeds_cache = [None] * train_dataset.num_instance_images if args.caption_dropout > 0 else None
+ empty_prompt_mask_cache = [None] * train_dataset.num_instance_images if args.caption_dropout > 0 else None
+
+ cache_batch_sampler = BucketBatchSampler(
+ train_dataset, batch_size=args.train_batch_size, drop_last=False, seed=args.seed
+ )
+ cache_dataloader = torch.utils.data.DataLoader(
+ train_dataset,
+ batch_sampler=cache_batch_sampler,
+ collate_fn=collate_fn,
+ num_workers=args.dataloader_num_workers,
+ )
+ for batch in tqdm(cache_dataloader, desc="Caching latents and prompt embeddings"):
+ with torch.no_grad():
+ sample_indices = batch["indices"]
+ if args.cache_latents:
+ with offload_models(vae, device=accelerator.device, offload=args.offload):
+ pixel_values = batch["pixel_values"].to(accelerator.device, non_blocking=True, dtype=vae.dtype)
+ cond_pixel_values = batch["cond_pixel_values"].to(
+ accelerator.device, non_blocking=True, dtype=vae.dtype
+ )
+ instance_latents = vae.encode(pixel_values).latent_dist.mode()
+ cond_latents = vae.encode(cond_pixel_values).latent_dist.mode()
+ for i, idx in enumerate(sample_indices):
+ instance_latents_cache[idx] = instance_latents[i : i + 1]
+ cond_latents_cache[idx] = cond_latents[i : i + 1]
+
+ with offload_models(text_encoder, device=accelerator.device, offload=args.offload):
+ for i, idx in enumerate(sample_indices):
+ instruction = batch["prompts"][i]
+ cond_image = batch["cond_text_images"][i]
+ prompt_embeds, prompt_mask = get_fibo2_edit_prompt_embeds(
+ text_encoder,
+ processor,
+ instruction,
+ cond_image,
+ accelerator.device,
+ weight_dtype,
+ args.max_sequence_length,
+ )
+ prompt_embeds_cache[idx] = prompt_embeds
+ prompt_mask_cache[idx] = prompt_mask
+ if args.caption_dropout > 0:
+ empty_embeds, empty_mask = get_fibo2_edit_prompt_embeds(
+ text_encoder,
+ processor,
+ "",
+ cond_image,
+ accelerator.device,
+ weight_dtype,
+ args.max_sequence_length,
+ )
+ empty_prompt_embeds_cache[idx] = empty_embeds
+ empty_prompt_mask_cache[idx] = empty_mask
+
+ assert all(embeds is not None for embeds in prompt_embeds_cache), "Prompt embedding cache has unfilled entries."
+ if args.cache_latents:
+ assert all(latents is not None for latents in instance_latents_cache), "Latent cache has unfilled entries."
+ assert all(latents is not None for latents in cond_latents_cache), (
+ "Conditioning latent cache has unfilled entries."
+ )
+
+ # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624
+ if args.cache_latents:
+ vae = vae.to("cpu")
+ del vae
+
+ # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624
+ text_encoder = text_encoder.to("cpu")
+ del text_encoder
+ free_memory()
+
+ # Scheduler and math around the number of training steps.
+ # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation.
+ num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes
+ if args.max_train_steps is None:
+ len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes)
+ num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps)
+ num_training_steps_for_scheduler = (
+ args.num_train_epochs * accelerator.num_processes * num_update_steps_per_epoch
+ )
+ else:
+ num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes
+
+ lr_scheduler = get_scheduler(
+ args.lr_scheduler,
+ optimizer=optimizer,
+ num_warmup_steps=num_warmup_steps_for_scheduler,
+ num_training_steps=num_training_steps_for_scheduler,
+ num_cycles=args.lr_num_cycles,
+ power=args.lr_power,
+ )
+
+ # Prepare everything with our `accelerator`.
+ transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+ transformer, optimizer, train_dataloader, lr_scheduler
+ )
+
+ # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+ num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+ if args.max_train_steps is None:
+ args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+ if num_training_steps_for_scheduler != args.max_train_steps:
+ logger.warning(
+ f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match "
+ f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. "
+ f"This inconsistency may result in the learning rate scheduler not functioning properly."
+ )
+ # Afterwards we recalculate our number of training epochs
+ args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+ # We need to initialize the trackers we use, and also store our configuration.
+ # The trackers initializes automatically on the main process.
+ if accelerator.is_main_process:
+ tracker_name = "dreambooth-fibo2-img2img-lora"
+ accelerator.init_trackers(tracker_name, config=vars(args).copy())
+
+ # Train!
+ total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+ logger.info("***** Running training *****")
+ logger.info(f" Num examples = {len(train_dataset)}")
+ logger.info(f" Num batches each epoch = {len(train_dataloader)}")
+ logger.info(f" Num Epochs = {args.num_train_epochs}")
+ logger.info(f" Instantaneous batch size per device = {args.train_batch_size}")
+ logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+ logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+ logger.info(f" Total optimization steps = {args.max_train_steps}")
+ global_step = 0
+ first_epoch = 0
+
+ # Potentially load in the weights and states from a previous save
+ if args.resume_from_checkpoint:
+ if args.resume_from_checkpoint != "latest":
+ path = os.path.basename(args.resume_from_checkpoint)
+ else:
+ # Get the mos recent checkpoint
+ dirs = os.listdir(args.output_dir)
+ dirs = [d for d in dirs if d.startswith("checkpoint")]
+ dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+ path = dirs[-1] if len(dirs) > 0 else None
+
+ if path is None:
+ accelerator.print(
+ f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+ )
+ args.resume_from_checkpoint = None
+ initial_global_step = 0
+ else:
+ accelerator.print(f"Resuming from checkpoint {path}")
+ accelerator.load_state(os.path.join(args.output_dir, path))
+ global_step = int(path.split("-")[1])
+
+ initial_global_step = global_step
+ first_epoch = global_step // num_update_steps_per_epoch
+
+ else:
+ initial_global_step = 0
+
+ progress_bar = tqdm(
+ range(0, args.max_train_steps),
+ initial=initial_global_step,
+ desc="Steps",
+ # Only show the progress bar once on each machine.
+ disable=not accelerator.is_local_main_process,
+ )
+
+ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
+ sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype)
+ schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device)
+ timesteps = timesteps.to(accelerator.device)
+ step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+ sigma = sigmas[step_indices].flatten()
+ while len(sigma.shape) < n_dim:
+ sigma = sigma.unsqueeze(-1)
+ return sigma
+
+ for epoch in range(first_epoch, args.num_train_epochs):
+ transformer.train()
+
+ for batch in train_dataloader:
+ models_to_accumulate = [transformer]
+ sample_indices = batch["indices"]
+ n_inst = len(sample_indices)
+
+ with accelerator.accumulate(models_to_accumulate):
+ # Assemble the batch's per-example edit-instruction embeddings, substituting the empty-instruction
+ # embedding (with the conditioning image kept) for dropped samples, then right-pad them to a shared
+ # length.
+ embeds_list = [prompt_embeds_cache[idx] for idx in sample_indices]
+ mask_list = [prompt_mask_cache[idx] for idx in sample_indices]
+ if args.caption_dropout > 0:
+ drop_mask = torch.rand(n_inst) < args.caption_dropout
+ embeds_list = [
+ empty_prompt_embeds_cache[sample_indices[i]] if drop_mask[i] else embeds_list[i]
+ for i in range(n_inst)
+ ]
+ mask_list = [
+ empty_prompt_mask_cache[sample_indices[i]] if drop_mask[i] else mask_list[i]
+ for i in range(n_inst)
+ ]
+ prompt_embeds, prompt_embeds_mask = pad_and_stack_embeds(embeds_list, mask_list, accelerator.device)
+ prompt_embeds = prompt_embeds.to(dtype=weight_dtype)
+
+ # Convert target + conditioning images to latent space
+ if args.cache_latents:
+ model_input = torch.cat([instance_latents_cache[idx] for idx in sample_indices], dim=0)
+ cond_model_input = torch.cat([cond_latents_cache[idx] for idx in sample_indices], dim=0)
+ model_input = model_input.to(accelerator.device)
+ cond_model_input = cond_model_input.to(accelerator.device)
+ else:
+ with offload_models(vae, device=accelerator.device, offload=args.offload):
+ pixel_values = batch["pixel_values"].to(device=accelerator.device, dtype=vae.dtype)
+ cond_pixel_values = batch["cond_pixel_values"].to(device=accelerator.device, dtype=vae.dtype)
+ model_input = vae.encode(pixel_values).latent_dist.mode()
+ cond_model_input = vae.encode(cond_pixel_values).latent_dist.mode()
+
+ # Normalize the clean latents with the VAE BatchNorm statistics, exactly as the pipeline does. The
+ # target becomes the sample the flow-matching target denoises to; the conditioning latents become the
+ # clean context tokens the transformer reads but never predicts.
+ model_input = normalize_latents(model_input.float(), latents_bn_mean, latents_bn_std).to(weight_dtype)
+ cond_model_input = normalize_latents(cond_model_input.float(), latents_bn_mean, latents_bn_std).to(
+ weight_dtype
+ )
+
+ # Sample noise that we'll add to the latents
+ noise = torch.randn_like(model_input)
+ bsz = model_input.shape[0]
+
+ # Sample a random timestep for each image
+ # for weighting schemes where we sample timesteps non-uniformly
+ u = compute_density_for_timestep_sampling(
+ weighting_scheme=args.weighting_scheme,
+ batch_size=bsz,
+ logit_mean=args.logit_mean,
+ logit_std=args.logit_std,
+ mode_scale=args.mode_scale,
+ )
+ indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
+ timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device)
+
+ # Add noise according to flow matching.
+ # zt = (1 - texp) * x + texp * z1
+ sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
+ noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise
+
+ # fibo-2 is not guidance-distilled: no guidance argument. Only the (noisy) target tokens get the flow
+ # timestep; the conditioning image's tokens are passed as clean context (`context_latents`) at t = 0,
+ # and the transformer returns only the target tokens.
+ timestep = sigmas.reshape(bsz)
+
+ with accelerator.autocast():
+ model_pred = transformer(
+ hidden_states=noisy_model_input.to(weight_dtype),
+ timestep=timestep,
+ encoder_hidden_states=prompt_embeds,
+ encoder_attention_mask=prompt_embeds_mask,
+ context_latents=[cond_model_input.to(weight_dtype)],
+ return_dict=False,
+ )[0]
+
+ # these weighting schemes use a uniform timestep sampling
+ # and instead post-weight the loss
+ weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas)
+
+ # flow matching loss (velocity target)
+ target = noise - model_input
+
+ # Compute regular loss.
+ loss = torch.mean(
+ (weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1),
+ 1,
+ )
+ loss = loss.mean()
+
+ accelerator.backward(loss)
+ if accelerator.sync_gradients:
+ params_to_clip = transformer.parameters()
+ accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+
+ optimizer.step()
+ lr_scheduler.step()
+ optimizer.zero_grad()
+
+ # Checks if the accelerator has performed an optimization step behind the scenes
+ if accelerator.sync_gradients:
+ progress_bar.update(1)
+ global_step += 1
+
+ if accelerator.is_main_process or is_fsdp:
+ if global_step % args.checkpointing_steps == 0:
+ # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
+ if args.checkpoints_total_limit is not None:
+ checkpoints = os.listdir(args.output_dir)
+ checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
+ checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
+
+ # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
+ if len(checkpoints) >= args.checkpoints_total_limit:
+ num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
+ removing_checkpoints = checkpoints[0:num_to_remove]
+
+ logger.info(
+ f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
+ )
+ logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
+
+ for removing_checkpoint in removing_checkpoints:
+ removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
+ shutil.rmtree(removing_checkpoint)
+
+ save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
+ accelerator.save_state(save_path)
+ logger.info(f"Saved state to {save_path}")
+
+ logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
+ progress_bar.set_postfix(**logs)
+ accelerator.log(logs, step=global_step)
+
+ if global_step >= args.max_train_steps:
+ break
+
+ if accelerator.is_main_process:
+ if (
+ args.validation_prompt is not None
+ and args.validation_image is not None
+ and (epoch % args.validation_epochs == 0)
+ ):
+ transformer_ = unwrap_model(transformer)
+ transformer_.eval()
+ pipeline = make_validation_pipeline(args, transformer_, weight_dtype, accelerator.device)
+ images = log_validation(
+ pipeline=pipeline,
+ args=args,
+ accelerator=accelerator,
+ pipeline_args=validation_pipeline_args,
+ epoch=epoch,
+ torch_dtype=weight_dtype,
+ )
+ transformer_.train()
+ free_memory()
+
+ # Save the lora layers
+ accelerator.wait_for_everyone()
+
+ if is_fsdp:
+ transformer = unwrap_model(transformer)
+ state_dict = accelerator.get_state_dict(transformer)
+ if accelerator.is_main_process:
+ modules_to_save = {}
+ if is_fsdp:
+ if args.bnb_quantization_config_path is None:
+ if args.upcast_before_saving:
+ state_dict = {
+ k: v.to(torch.float32) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
+ }
+ else:
+ state_dict = {
+ k: v.to(weight_dtype) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
+ }
+
+ transformer_lora_layers = get_peft_model_state_dict(
+ transformer,
+ state_dict=state_dict,
+ )
+ transformer_lora_layers = {
+ k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v
+ for k, v in transformer_lora_layers.items()
+ }
+
+ else:
+ transformer = unwrap_model(transformer)
+ if args.bnb_quantization_config_path is None:
+ if args.upcast_before_saving:
+ transformer.to(torch.float32)
+ else:
+ transformer = transformer.to(weight_dtype)
+ transformer_lora_layers = get_peft_model_state_dict(transformer)
+
+ modules_to_save["transformer"] = transformer
+
+ BriaFibo2LoraLoaderMixin.save_lora_weights(
+ save_directory=args.output_dir,
+ transformer_lora_layers=transformer_lora_layers,
+ **_collate_lora_metadata(modules_to_save),
+ )
+
+ images = []
+ run_validation = args.validation_prompt and args.validation_image and args.num_validation_images > 0
+ should_run_final_inference = not args.skip_final_inference and run_validation
+ if should_run_final_inference:
+ pipeline = make_validation_pipeline(args, transformer, weight_dtype, accelerator.device)
+ images = log_validation(
+ pipeline=pipeline,
+ args=args,
+ accelerator=accelerator,
+ pipeline_args=validation_pipeline_args,
+ epoch=epoch,
+ is_final_validation=True,
+ torch_dtype=weight_dtype,
+ )
+ del pipeline
+ free_memory()
+
+ quant_training = None
+ if args.do_fp8_training:
+ quant_training = "FP8 TorchAO"
+ elif args.bnb_quantization_config_path:
+ quant_training = "BitsandBytes"
+ save_model_card(
+ (args.hub_model_id or Path(args.output_dir).name) if not args.push_to_hub else repo_id,
+ images=images,
+ base_model=args.pretrained_model_name_or_path,
+ instance_prompt=args.instance_prompt,
+ validation_prompt=args.validation_prompt,
+ repo_folder=args.output_dir,
+ quant_training=quant_training,
+ )
+
+ if args.push_to_hub:
+ upload_folder(
+ repo_id=repo_id,
+ folder_path=args.output_dir,
+ commit_message="End of training",
+ ignore_patterns=["step_*", "epoch_*"],
+ )
+
+ accelerator.end_training()
+
+
+if __name__ == "__main__":
+ args = parse_args()
+ main(args)
diff --git a/src/diffusers/loaders/__init__.py b/src/diffusers/loaders/__init__.py
index 72a24ba3803f..76d049710732 100644
--- a/src/diffusers/loaders/__init__.py
+++ b/src/diffusers/loaders/__init__.py
@@ -42,6 +42,7 @@
"ZImageLoraLoaderMixin",
"Flux2LoraLoaderMixin",
"Flux3LoraLoaderMixin",
+ "BriaFibo2LoraLoaderMixin",
"Ideogram4LoraLoaderMixin",
"ErnieImageLoraLoaderMixin",
"CosmosLoraLoaderMixin",
@@ -78,6 +79,7 @@
AmusedLoraLoaderMixin,
AnimaLoraLoaderMixin,
AuraFlowLoraLoaderMixin,
+ BriaFibo2LoraLoaderMixin,
CogVideoXLoraLoaderMixin,
CogView4LoraLoaderMixin,
CosmosLoraLoaderMixin,
diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py
index 22d596c6c682..8e2e341b81cc 100644
--- a/src/diffusers/loaders/lora_pipeline.py
+++ b/src/diffusers/loaders/lora_pipeline.py
@@ -6163,6 +6163,192 @@ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
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
+ 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.
+ """
+ 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`].
diff --git a/src/diffusers/pipelines/bria_fibo2/pipeline_bria_fibo2.py b/src/diffusers/pipelines/bria_fibo2/pipeline_bria_fibo2.py
index 4ef562923b7e..43d0c8ba379f 100644
--- a/src/diffusers/pipelines/bria_fibo2/pipeline_bria_fibo2.py
+++ b/src/diffusers/pipelines/bria_fibo2/pipeline_bria_fibo2.py
@@ -26,6 +26,7 @@
from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
from ...image_processor import VaeImageProcessor
+from ...loaders import BriaFibo2LoraLoaderMixin
from ...models import AutoencoderKLFlux2, BriaFibo2Transformer2DModel
from ...schedulers import FlowMatchEulerDiscreteScheduler
from ...utils import is_torch_xla_available, logging, replace_example_docstring
@@ -164,7 +165,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.