mrfakename's picture
Mirror briaai/Fibo-2.0-dev
b3c38d3 verified
|
Raw History Blame Contribute Delete
5.32 kB

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. Until it is merged, install the branch the model card uses:

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:

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'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

{"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.

{"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'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:

{
  "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

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

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.