finetune_sdxl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Description\n",
    "This notebook lets you finetune the [SDXL](https://huggingface.co/papers/2307.01952) model on a set of dozens of images related to one topic/theme. It uses DreamBooth and LoRA training methods. You can get good results with less than 20 images. \n",
    "\n",
    "DreamBooth is a training technique that updates the entire diffusion model by training on just a few images of a common subject or style. It works by associating a special prompt \"_instance_prompt_\" with the images. The keyword in this prompt is what will trigger the tuned weights during the inference.\n",
    "\n",
    "LoRA is a training technique that significantly reduces the number of trainable parameters. It works by inserting a smaller number of new weights into the model and only these are trained. This produces smaller weights after the finetuning.\n",
    "\n",
    "#### Experiment\n",
    "The model was finetuned on 15 images of moroccan cities of size 512x512, with instance prompt = \"a moroccan city\".\n",
    "\n",
    "The finetuning takes ~2 hours on a TESLA P100 GPU with a batch size of 1 and a learning rate of 1e-4.\n",
    "\n",
    "Some results can be seen below:\n",
    "\n",
    "![image.png](docs/result_1.jpg)\n",
    "![image.png](docs/result_2.jpg)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cloning into 'diffusers'...\n",
      "Updating files:  48% (617/1261)\n",
      "Updating files:  49% (618/1261)\n",
      "Updating files:  50% (631/1261)\n",
      "Updating files:  51% (644/1261)\n",
      "Updating files:  52% (656/1261)\n",
      "Updating files:  53% (669/1261)\n",
      "Updating files:  54% (681/1261)\n",
      "Updating files:  55% (694/1261)\n",
      "Updating files:  56% (707/1261)\n",
      "Updating files:  57% (719/1261)\n",
      "Updating files:  58% (732/1261)\n",
      "Updating files:  59% (744/1261)\n",
      "Updating files:  60% (757/1261)\n",
      "Updating files:  61% (770/1261)\n",
      "Updating files:  62% (782/1261)\n",
      "Updating files:  63% (795/1261)\n",
      "Updating files:  64% (808/1261)\n",
      "Updating files:  65% (820/1261)\n",
      "Updating files:  66% (833/1261)\n",
      "Updating files:  67% (845/1261)\n",
      "Updating files:  68% (858/1261)\n",
      "Updating files:  69% (871/1261)\n",
      "Updating files:  70% (883/1261)\n",
      "Updating files:  71% (896/1261)\n",
      "Updating files:  72% (908/1261)\n",
      "Updating files:  73% (921/1261)\n",
      "Updating files:  74% (934/1261)\n",
      "Updating files:  75% (946/1261)\n",
      "Updating files:  76% (959/1261)\n",
      "Updating files:  77% (971/1261)\n",
      "Updating files:  78% (984/1261)\n",
      "Updating files:  79% (997/1261)\n",
      "Updating files:  80% (1009/1261)\n",
      "Updating files:  81% (1022/1261)\n",
      "Updating files:  82% (1035/1261)\n",
      "Updating files:  83% (1047/1261)\n",
      "Updating files:  84% (1060/1261)\n",
      "Updating files:  85% (1072/1261)\n",
      "Updating files:  86% (1085/1261)\n",
      "Updating files:  87% (1098/1261)\n",
      "Updating files:  88% (1110/1261)\n",
      "Updating files:  89% (1123/1261)\n",
      "Updating files:  90% (1135/1261)\n",
      "Updating files:  91% (1148/1261)\n",
      "Updating files:  92% (1161/1261)\n",
      "Updating files:  93% (1173/1261)\n",
      "Updating files:  94% (1186/1261)\n",
      "Updating files:  95% (1198/1261)\n",
      "Updating files:  96% (1211/1261)\n",
      "Updating files:  97% (1224/1261)\n",
      "Updating files:  98% (1236/1261)\n",
      "Updating files:  99% (1249/1261)\n",
      "Updating files: 100% (1261/1261)\n",
      "Updating files: 100% (1261/1261), done.\n"
     ]
    }
   ],
   "source": [
    "!git clone https://github.com/huggingface/diffusers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -q /diffusers/.\n",
    "! pip install -q -r /diffusers/examples/dreambooth/requirements.txt\n",
    "! pip install -q bitsandbytes>=0.40.0\n",
    "! pip install -q xformers>=0.0.20\n",
    "! pip install -q numpy>= 1.22.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from diffusers import DiffusionPipeline, AutoencoderKL\n",
    "from accelerate.utils import write_basic_config\n",
    "from huggingface_hub import whoami, upload_folder, create_repo, snapshot_download\n",
    "\n",
    "write_basic_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters to set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify your HF token if you want to push the model to a HF repo\n",
    "HF_TOKEN = \"\"\n",
    "save_to_hf = False\n",
    "\n",
    "# the directory where the images are stored or will be stored if you choose to download them from HF hub (see below)\n",
    "img_data_dir = \"data\" \n",
    "\n",
    "# the directory where the lora weights will be stored\n",
    "model_dir = \"model\" \n",
    "\n",
    "# the prompt that contains the common theme keyword of your images\n",
    "instance_prompt = \"a photo of a moroccan city\" \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(model_dir, exist_ok=True)\n",
    "os.makedirs('output/pretrained/', exist_ok=True)\n",
    "os.makedirs('output/finetuned/', exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prompt the pre-trained model\n",
    "Uncomment only if you need to test the pretrained sdxl. \n",
    "\n",
    "If you're running this notebook on < 14GB VRAM then you cannot launch both this section and the finetuning script. In this case, restart your session. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"moroccan city\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = DiffusionPipeline.from_pretrained(\"stabilityai/stable-diffusion-xl-base-1.0\").to(\"cuda\")\n",
    "\n",
    "num_inference_steps = 50 \n",
    "guidance_scale = 10\n",
    "image = pipeline(prompt,\n",
    "                 num_inference_steps=num_inference_steps,\n",
    "                 guidance_scale=guidance_scale).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finetuning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download dataset from HF (optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# os.makedirs(img_data_dir, exist_ok=True)\n",
    "# snapshot_download(\n",
    "#     \"imomayiz/morocco-img\",\n",
    "#     local_dir=img_data_dir,\n",
    "#     repo_type=\"dataset\",\n",
    "#     ignore_patterns=\".gitattributes\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finetune\n",
    "Upload your images to _data_dir_ or download an img dataset from HF by executing the cell above.\n",
    " \n",
    "Set the training parameters depending on the available VRAM.\n",
    "\n",
    "The weights will be saved under _model_dir_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env bash\n",
    "! accelerate launch /diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py \\\n",
    "  --pretrained_model_name_or_path=\"stabilityai/stable-diffusion-xl-base-1.0\" \\\n",
    "  --pretrained_vae_model_name_or_path=\"madebyollin/sdxl-vae-fp16-fix\" \\\n",
    "  --instance_data_dir=img_data_dir \\\n",
    "  --output_dir=model_dir \\\n",
    "  --instance_prompt=instance_prompt \\\n",
    "  --resolution=1024 \\\n",
    "  --train_batch_size=1 \\\n",
    "  --gradient_accumulation_steps=3 \\\n",
    "  --gradient_checkpointing \\\n",
    "  --learning_rate=1e-4 \\\n",
    "  --snr_gamma=5.0 \\\n",
    "  --lr_scheduler=\"constant\" \\\n",
    "  --lr_warmup_steps=0 \\\n",
    "  --mixed_precision=\"fp16\" \\\n",
    "  --use_8bit_adam \\\n",
    "  --max_train_steps=500 \\\n",
    "  --checkpointing_steps=717 \\\n",
    "  --seed=\"0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the finetuned model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae = AutoencoderKL.from_pretrained(\"madebyollin/sdxl-vae-fp16-fix\", torch_dtype=torch.float16)\n",
    "pipe = DiffusionPipeline.from_pretrained(\n",
    "    \"stabilityai/stable-diffusion-xl-base-1.0\",\n",
    "    vae=vae,\n",
    "    torch_dtype=torch.float16,\n",
    "    variant=\"fp16\",\n",
    "    use_safetensors=True\n",
    ")\n",
    "pipe.load_lora_weights(model_dir)\n",
    "_ = pipe.to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"a photo of a modern moroccan city\" # @param\n",
    "\n",
    "image = pipe(prompt=prompt, num_inference_steps=20).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save model to HF (optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_model_card(\n",
    "    repo_id: str,\n",
    "    images=None,\n",
    "    base_model=str,\n",
    "    train_text_encoder=False,\n",
    "    instance_prompt=str,\n",
    "    validation_prompt=str,\n",
    "    repo_folder=None,\n",
    "    vae_path=None,\n",
    "):\n",
    "    img_str = \"widget:\\n\" if images else \"\"\n",
    "    for i, image in enumerate(images):\n",
    "        image.save(os.path.join(repo_folder, f\"image_{i}.png\"))\n",
    "        img_str += f\"\"\"\n",
    "        - text: '{validation_prompt if validation_prompt else ' ' }'\n",
    "          output:\n",
    "            url:\n",
    "                \"image_{i}.png\"\n",
    "        \"\"\"\n",
    "\n",
    "    yaml = f\"\"\"\n",
    "---\n",
    "tags:\n",
    "- stable-diffusion-xl\n",
    "- stable-diffusion-xl-diffusers\n",
    "- text-to-image\n",
    "- diffusers\n",
    "- lora\n",
    "- template:sd-lora\n",
    "{img_str}\n",
    "base_model: {base_model}\n",
    "instance_prompt: {instance_prompt}\n",
    "license: openrail++\n",
    "---\n",
    "    \"\"\"\n",
    "\n",
    "    model_card = f\"\"\"\n",
    "# SDXL LoRA DreamBooth - {repo_id}\n",
    "\n",
    "<Gallery />\n",
    "\n",
    "## Model description\n",
    "\n",
    "These are {repo_id} LoRA adaption weights for {base_model}.\n",
    "\n",
    "The weights were trained  using [DreamBooth](https://dreambooth.github.io/).\n",
    "\n",
    "LoRA for the text encoder was enabled: {train_text_encoder}.\n",
    "\n",
    "Special VAE used for training: {vae_path}.\n",
    "\n",
    "## Trigger words\n",
    "\n",
    "You should use {instance_prompt} to trigger the image generation.\n",
    "\n",
    "## Download model\n",
    "\n",
    "Weights for this model are available in Safetensors format.\n",
    "\n",
    "[Download]({repo_id}/tree/main) them in the Files & versions tab.\n",
    "\n",
    "\"\"\"\n",
    "    with open(os.path.join(repo_folder, \"README.md\"), \"w\") as f:\n",
    "        f.write(yaml + model_card)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if save_to_hf:\n",
    "    username = whoami(token=HF_TOKEN)[\"name\"]\n",
    "    repo_id = f\"{username}/sdxl_lora\"\n",
    "    repo_id = create_repo(repo_id, exist_ok=True, token=HF_TOKEN).repo_id\n",
    "    \n",
    "    save_model_card(\n",
    "        repo_id = repo_id,\n",
    "        images=[],\n",
    "        base_model=\"stabilityai/stable-diffusion-xl-base-1.0\",\n",
    "        train_text_encoder=False,\n",
    "        instance_prompt=instance_prompt,\n",
    "        validation_prompt=None,\n",
    "        repo_folder=\"\",\n",
    "        vae_path=\"madebyollin/sdxl-vae-fp16-fix\",\n",
    "    )\n",
    "\n",
    "    upload_folder(\n",
    "        token=HF_TOKEN,\n",
    "        repo_id=repo_id,\n",
    "        folder_path=model_dir,\n",
    "        commit_message=\"End of training\",\n",
    "        ignore_patterns=[\"step_*\", \"epoch_*\"],\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}