| # 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`. |
|
|
