community-efforts/image_preferences/05_fine_tune_flux_lora.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tune Flux LoRA on the image preferences dataset\n",
    "\n",
    "Note, we will not use preferences from the dev set for this fine-tuning. We will only use the chosen images for an Supervised fine-tuning phase. Additionally, we recommend using a A100 GPU (4$/hour on Hugging Face) for this fine-tuning because of the memory requirements. The fine-tuning script will take about 4 hours to complete for a single epoch.\n",
    "\n",
    "## Install dependencies\n",
    "\n",
    "We first make sure we have the latest version of diffusers installed. This is a development version of diffusers, so we need to install it from source. Additionally, we install the other dependencies that are required for the fine-tuning script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/huggingface/diffusers\n",
    "!pip install -e diffusers/.\n",
    "!pip install datasets sentencepiece protobuf accelerate peft wandb torchvision prodigyopt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Logins and config\n",
    "\n",
    "We will use Weights & Biases to log the training process. Additionally, we log in to Hugging Face to push the finetuned model to the Hub."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!huggingface-cli login --token \"hf_xxx\"\n",
    "!wandb login \"xxx\"\n",
    "!accelerate config default"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine-tune the model\n",
    "\n",
    "Lastly, we fine-tune the Flux LoRA on the chosen images from the image preferences dataset. We heavily inspired from the [Dreambooth fine-tuning script](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux.md) and modified it to work for our use case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!accelerate launch diffusers/examples/dreambooth/train_dreambooth_lora_flux.py \\\n",
    "    --pretrained_model_name_or_path \"black-forest-labs/FLUX.1-dev\" \\\n",
    "    --dataset_name \"data-is-better-together/open-image-preferences-v1-binarized\" \\\n",
    "    --hub_model_id \"davidberenstein1957/open-image-preferences-v1-flux-dev-lora\" \\\n",
    "    --push_to_hub \\\n",
    "    --output_dir \"open-image-preferences-v1-flux-dev-lora\" \\\n",
    "    --image_column \"chosen\" \\\n",
    "    --caption_column \"prompt\" \\\n",
    "    --mixed_precision=\"bf16\" \\\n",
    "    --resolution=1024 \\\n",
    "    --train_batch_size=1 \\\n",
    "    --repeats=1 \\\n",
    "    --report_to=\"wandb\"\\\n",
    "    --gradient_accumulation_steps=1 \\\n",
    "    --gradient_checkpointing \\\n",
    "    --learning_rate=1.0 \\\n",
    "    --text_encoder_lr=1.0 \\\n",
    "    --optimizer=\"prodigy\"\\\n",
    "    --lr_scheduler=\"constant\" \\\n",
    "    --lr_warmup_steps=0 \\\n",
    "    --rank=8 \\\n",
    "    --checkpointing_steps=2000 \\\n",
    "    --seed=\"0\" "
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}