|
| 1 | +## LoRA fine-tuning Flux.1 Dev with quantization |
| 2 | + |
| 3 | +> [!NOTE] |
| 4 | +> This example is educational in nature and fixes some arguments to keep things simple. It should act as a reference to build things further. |
| 5 | +
|
| 6 | +This example shows how to fine-tune [Flux.1 Dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) with LoRA and quantization. We show this by using the [`Norod78/Yarn-art-style`](https://huggingface.co/datasets/Norod78/Yarn-art-style) dataset. Steps below summarize the workflow: |
| 7 | + |
| 8 | +* We precompute the text embeddings in `compute_embeddings.py` and serialize them into a parquet file. |
| 9 | +* `train_dreambooth_lora_flux_miniature.py` takes care of training: |
| 10 | + * Since we already precomputed the text embeddings, we don't load the text encoders. |
| 11 | + * We load the VAE and use it to precompute the image latents and we then delete it. |
| 12 | + * Load the Flux transformer, quantize it with the [NF4 datatype](https://arxiv.org/abs/2305.14314) through `bitsandbytes`, prepare it for 4bit training. |
| 13 | + * Add LoRA adapter layers to it and then ensure they are kept in FP32 precision. |
| 14 | + * Train! |
| 15 | + |
| 16 | +To run training in a memory-optimized manner, we additionally use: |
| 17 | + |
| 18 | +* 8Bit Adam |
| 19 | +* Gradient checkpointing |
| 20 | + |
| 21 | +We have tested the scripts on a 24GB 4090. It works on a free-tier Colab Notebook, too, but it's extremely slow. |
| 22 | + |
| 23 | +## Training |
| 24 | + |
| 25 | +Ensure you have installed the required libraries: |
| 26 | + |
| 27 | +```bash |
| 28 | +pip install -U transformers accelerate bitsandbytes peft datasets |
| 29 | +pip install git+https://github.com/huggingface/diffusers -U |
| 30 | +``` |
| 31 | + |
| 32 | +Now, compute the text embeddings: |
| 33 | + |
| 34 | +```bash |
| 35 | +python compute_embeddings.py |
| 36 | +``` |
| 37 | + |
| 38 | +It should create a file named `embeddings.parquet`. We're then ready to launch training. First, authenticate so that you can access the Flux.1 Dev model: |
| 39 | + |
| 40 | +```bash |
| 41 | +huggingface-cli |
| 42 | +``` |
| 43 | + |
| 44 | +Then launch: |
| 45 | + |
| 46 | +```bash |
| 47 | +accelerate launch --config_file=accelerate.yaml \ |
| 48 | + train_dreambooth_lora_flux_miniature.py \ |
| 49 | + --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \ |
| 50 | + --data_df_path="embeddings.parquet" \ |
| 51 | + --output_dir="yarn_art_lora_flux_nf4" \ |
| 52 | + --mixed_precision="fp16" \ |
| 53 | + --use_8bit_adam \ |
| 54 | + --weighting_scheme="none" \ |
| 55 | + --resolution=1024 \ |
| 56 | + --train_batch_size=1 \ |
| 57 | + --repeats=1 \ |
| 58 | + --learning_rate=1e-4 \ |
| 59 | + --guidance_scale=1 \ |
| 60 | + --report_to="wandb" \ |
| 61 | + --gradient_accumulation_steps=4 \ |
| 62 | + --gradient_checkpointing \ |
| 63 | + --lr_scheduler="constant" \ |
| 64 | + --lr_warmup_steps=0 \ |
| 65 | + --cache_latents \ |
| 66 | + --rank=4 \ |
| 67 | + --max_train_steps=700 \ |
| 68 | + --seed="0" |
| 69 | +``` |
| 70 | + |
| 71 | +We can direcly pass a quantized checkpoint path, too: |
| 72 | + |
| 73 | +```diff |
| 74 | ++ --quantized_model_path="hf-internal-testing/flux.1-dev-nf4-pkg" |
| 75 | +``` |
| 76 | + |
| 77 | +Depending on the machine, training time will vary but for our case, it was 1.5 hours. It maybe possible to speed this up by using `torch.bfloat16`. |
| 78 | + |
| 79 | +We support training with the DeepSpeed Zero2 optimizer, too. To use it, first install DeepSpeed: |
| 80 | + |
| 81 | +```bash |
| 82 | +pip install -Uq deepspeed |
| 83 | +``` |
| 84 | + |
| 85 | +And then launch: |
| 86 | + |
| 87 | +```bash |
| 88 | +accelerate launch --config_file=ds2.yaml \ |
| 89 | + train_dreambooth_lora_flux_miniature.py \ |
| 90 | + --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \ |
| 91 | + --data_df_path="embeddings.parquet" \ |
| 92 | + --output_dir="yarn_art_lora_flux_nf4" \ |
| 93 | + --mixed_precision="no" \ |
| 94 | + --use_8bit_adam \ |
| 95 | + --weighting_scheme="none" \ |
| 96 | + --resolution=1024 \ |
| 97 | + --train_batch_size=1 \ |
| 98 | + --repeats=1 \ |
| 99 | + --learning_rate=1e-4 \ |
| 100 | + --guidance_scale=1 \ |
| 101 | + --report_to="wandb" \ |
| 102 | + --gradient_accumulation_steps=4 \ |
| 103 | + --gradient_checkpointing \ |
| 104 | + --lr_scheduler="constant" \ |
| 105 | + --lr_warmup_steps=0 \ |
| 106 | + --cache_latents \ |
| 107 | + --rank=4 \ |
| 108 | + --max_train_steps=700 \ |
| 109 | + --seed="0" |
| 110 | +``` |
| 111 | + |
| 112 | +## Inference |
| 113 | + |
| 114 | +When loading the LoRA params (that were obtained on a quantized base model) and merging them into the base model, it is recommended to first dequantize the base model, merge the LoRA params into it, and then quantize the model again. This is because merging into 4bit quantized models can lead to some rounding errors. Below, we provide an end-to-end example: |
| 115 | + |
| 116 | +1. First, load the original model and merge the LoRA params into it: |
| 117 | + |
| 118 | +```py |
| 119 | +from diffusers import FluxPipeline |
| 120 | +import torch |
| 121 | + |
| 122 | +ckpt_id = "black-forest-labs/FLUX.1-dev" |
| 123 | +pipeline = FluxPipeline.from_pretrained( |
| 124 | + ckpt_id, text_encoder=None, text_encoder_2=None, torch_dtype=torch.float16 |
| 125 | +) |
| 126 | +pipeline.load_lora_weights("yarn_art_lora_flux_nf4", weight_name="pytorch_lora_weights.safetensors") |
| 127 | +pipeline.fuse_lora() |
| 128 | +pipeline.unload_lora_weights() |
| 129 | + |
| 130 | +pipeline.transformer.save_pretrained("fused_transformer") |
| 131 | +``` |
| 132 | + |
| 133 | +2. Quantize the model and run inference |
| 134 | + |
| 135 | +```py |
| 136 | +from diffusers import AutoPipelineForText2Image, FluxTransformer2DModel, BitsAndBytesConfig |
| 137 | +import torch |
| 138 | + |
| 139 | +ckpt_id = "black-forest-labs/FLUX.1-dev" |
| 140 | +bnb_4bit_compute_dtype = torch.float16 |
| 141 | +nf4_config = BitsAndBytesConfig( |
| 142 | + load_in_4bit=True, |
| 143 | + bnb_4bit_quant_type="nf4", |
| 144 | + bnb_4bit_compute_dtype=bnb_4bit_compute_dtype, |
| 145 | +) |
| 146 | +transformer = FluxTransformer2DModel.from_pretrained( |
| 147 | + "fused_transformer", |
| 148 | + quantization_config=nf4_config, |
| 149 | + torch_dtype=bnb_4bit_compute_dtype, |
| 150 | +) |
| 151 | +pipeline = AutoPipelineForText2Image.from_pretrained( |
| 152 | + ckpt_id, transformer=transformer, torch_dtype=bnb_4bit_compute_dtype |
| 153 | +) |
| 154 | +pipeline.enable_model_cpu_offload() |
| 155 | + |
| 156 | +image = pipeline( |
| 157 | + "a puppy in a pond, yarn art style", num_inference_steps=28, guidance_scale=3.5, height=768 |
| 158 | +).images[0] |
| 159 | +image.save("yarn_merged.png") |
| 160 | +``` |
| 161 | + |
| 162 | +| Dequantize, merge, quantize | Merging directly into quantized model | |
| 163 | +|-------|-------| |
| 164 | +|  |  | |
| 165 | + |
| 166 | +As we can notice the first column result follows the style more closely. |
0 commit comments