# Fibo 2.0 LoRA fine-tuning

`train_lora.py` trains a LoRA on Fibo 2.0 Dev for image editing or text-to-image. It runs on one
GPU and uses the diffusers `BriaFibo2Pipeline` to encode its training data exactly as the pipeline
encodes prompts and images at inference. The module docstring has the full details.

## Install

`BriaFibo2Pipeline` is being added to diffusers in
[PR #14991](https://github.com/huggingface/diffusers/pull/14991). Until it is merged, install the
branch the model card uses:

```bash
pip install "torch==2.10.*" torchvision "transformers>=5.19" peft safetensors pillow numpy click \
    "diffusers @ git+https://github.com/muhammadali-afridi/diffusers@ali/fibo2"
```

Then get the script. The model repo is gated, so accept its terms on Hugging Face and log in first:

```bash
hf auth login
hf download briaai/Fibo-2.0-dev --include "fine_tuning/*" --local-dir .
```

## Base model

`--base` defaults to `briaai/Fibo-2.0-dev`, which is downloaded on first use. It also takes a local
folder in the same diffusers layout.

## Data

The trainer picks the mode from the data folder. If `<data>/edits.jsonl` exists it trains edits,
otherwise text-to-image. `--mode` overrides this.

### Edit

`<data>/edits.jsonl` holds one row per training pair. Image paths are relative to `<data>`.

- `sources`: 1 to 4 input images. With several, the instruction can name them `<image_1>`,
  `<image_2>` and so on.
- `target`: the edited result.
- `instruction`: the prompt.

The pipeline accepts edit prompts in two forms, and the `instruction` can use either:

- **Plain instruction:** the edit instruction alone.
- **Instruction with a caption of the result:** a structured JSON caption of the target image,
  with the instruction under `edit_instruction`. This is what
  [`bria-vgl`](https://github.com/Bria-AI/vgl)'s `modify()` writes, and what the model card's
  inference examples use.

The two forms give different text embeddings, so train in the form you will prompt with at
inference. Rows can mix the two forms.

#### Plain instruction

```json
{"sources": ["images/0001_src.jpg"], "target": "images/0001_tgt.jpg", "instruction": "Remove the lamp."}
```

#### Instruction with a caption of the result

The pipeline sends the instruction and the caption to the text encoder as separate fields. Write
the caption as a JSON object, or as that object serialised into a string. A JSON caption without
`edit_instruction` is an error.

```json
{"sources": ["images/0002_src.jpg"], "target": "images/0002_tgt.jpg", "instruction": {"short_description": "An empty oak desk by a window.", "objects": [], "edit_instruction": "Remove the lamp."}}
```

The caption can also sit under a `structured_caption` key, next to `edit_instruction` and nothing
else. The pipeline drops empty values and the `aesthetic_score` and `preference_score` keys from a
caption.

### Text-to-image

Put each image next to a caption file with the same name: `<name>.json` for a structured caption,
or `<name>.txt` for a plain sentence. Subfolders are searched too.

```
my_t2i/
├── bicycle.jpg
├── bicycle.json
└── street/
    ├── 0001.webp
    └── 0001.txt
```

- Images can be `png`, `jpg`, `jpeg` or `webp`.
- Images without a caption are skipped. If both caption files exist, the `.json` is used.
- The folder must not contain `edits.jsonl`, or the trainer switches to edit mode.

Prefer `.json` captions: the model was trained on structured captions, the same format you prompt
with at inference. [`bria-vgl`](https://github.com/Bria-AI/vgl)'s `create()` writes them. Leave
out keys you don't need. The trainer cleans and compacts the JSON as the pipeline does, and
handles `text_render` entries the same way too. For example, `bicycle.json`:

```json
{
  "short_description": "A red bicycle leaning against a white brick wall on a sunny morning.",
  "objects": [{"description": "a red city bicycle with a wicker basket", "location": "center"}],
  "background_setting": "a white brick wall with a green door",
  "lighting": {"conditions": "bright morning sun", "direction": "from the right"},
  "style_medium": "photograph"
}
```

## Train

```bash
python fine_tuning/train_lora.py --data my_edits/ --output-dir out/ \
    --resolution 1024 --rank 64 --lr 1e-4 --steps 1000 --grad-accum 8
```

The first run encodes every sample once and caches it in `<output-dir>/cache`. Checkpoints go to
`out/checkpoint-<step>/` and `out/final/`. The loss per step is in `out/loss.csv`. That loss uses a
random noise level per sample, so expect a noisy, flat-looking curve.

## Use the LoRA

```python
import torch
from diffusers import BriaFibo2Pipeline

pipe = BriaFibo2Pipeline.from_pretrained("briaai/Fibo-2.0-dev", dtype=torch.bfloat16).to("cuda")
pipe.transformer.load_lora_adapter("out/final", weight_name="pytorch_lora_weights.safetensors")
```

`weight_name` is required. Without it diffusers looks for a `.bin` file.

## Notes

- **Memory:** an edit LoRA at 1024, rank 128, grad accum 8 peaked at 31 GiB, at about 58 s per
  step on an L40S.
- **Defaults:** AdamW, LR 1e-4 constant, LoRA on all 415 linear layers, LoRA weights kept in fp32.
- **Adapter size:** a rank-128 adapter is about 2.1 GB in fp32.
- **Gradient checkpointing** is on by default. Turn it off with `--no-gradient-checkpointing`.