support HRA

examples/hra_dreambooth/README.md

@@ -0,0 +1,112 @@
+<!--Copyright 2023 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+
+-->
+
+# DreamBooth fine-tuning with HRA
+
+This guide demonstrates how to use Householder reflection adaptation (HRA) method, to fine-tune Dreambooth with `stabilityai/stable-diffusion-2-1` model.
+
+HRA provides a new perspective connecting LoRA to OFT and achieves encouraging performance in various downstream tasks.
+HRA adapts a pre-trained model by multiplying each frozen weight matrix with a chain of r learnable Householder reflections (HRs).
+HRA can be interpreted as either an OFT adapter or an adaptive LoRA. 
+Consequently, it harnesses the advantages of both strategies, reducing parameters and computation costs while penalizing the loss of pre-training knowledge
+For further details on BOFT, please consult the [original HRA paper](https://arxiv.org/abs/2405.17484).
+
+In this guide we provide a Dreambooth fine-tuning script that is available in [PEFT's GitHub repo examples](https://github.com/huggingface/peft/tree/main/examples/hra_dreambooth). This implementation is adapted from [peft's boft_dreambooth](https://github.com/huggingface/peft/tree/main/examples/boft_dreambooth). 
+
+You can try it out and fine-tune on your custom images.
+
+## Set up your environment
+
+Start by cloning the PEFT repository:
+
+```bash
+git clone --recursive https://github.com/huggingface/peft
+```
+
+Navigate to the directory containing the training scripts for fine-tuning Dreambooth with HRA:
+
+```bash
+cd peft/examples/hra_dreambooth
+```
+
+Set up your environment: install PEFT, and all the required libraries. At the time of writing this guide we recommend installing PEFT from source. The following environment setup should work on A100 and H100:
+
+```bash
+conda create --name peft python=3.10
+conda activate peft
+conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 pytorch-cuda=11.8 -c pytorch -c nvidia
+conda install xformers -c xformers
+pip install -r requirements.txt
+pip install git+https://github.com/huggingface/peft
+```
+
+## Download the data
+
+Download [dreambooth](https://github.com/google/dreambooth) dataset by running this script.
+
+```bash
+bash download_dreambooth.sh
+```
+
+After downloading the data, your directory structure should look like this:
+
+```
+hra_dreambooth
+├── dreambooth
+│   └── dataset
+│       ├── backpack
+│       └── backpack_dog
+│       ...
+```
+
+You can also put your custom images into `hra_dreambooth/dreambooth/dataset`.
+
+## Fine-tune Dreambooth with HRA
+
+```bash
+./train_dreambooth.sh $prompt_idx $class_idx
+```
+
+where the `$prompt_idx` corresponds to different prompts ranging from 0 to 24 and the `$class_idx` corresponds to different subjects ranging from 0 to 29.
+
+Launch the training script with `accelerate` and pass hyperparameters, as well as LoRa-specific arguments to it such as:
+
+- `use_hra`: Enables HRA in the training script.
+- `hra_r`: the number of HRs (i.e., r) across different layers, expressed in `int`. 
+As r increases, the number of trainable parameters increases, which generally leads to improved performance.
+However, this also results in higher memory consumption and longer computation times. 
+Therefore, r is usually set to 8.
+**Note**, please set r to an even number to avoid potential issues during initialization.
+- `hra_apply_GS`: Applys Gram-Schmidt orthogonalization. Default is `false`.
+- `hra_bias`: specify if the `bias` paramteres should be traind. Can be `none`, `all` or `hra_only`.
+
+If you are running this script on Windows, you may need to set the `--num_dataloader_workers` to 0.
+
+To learn more about DreamBooth fine-tuning with prior-preserving loss, check out the [Diffusers documentation](https://huggingface.co/docs/diffusers/training/dreambooth#finetuning-with-priorpreserving-loss).
+
+**Note**, after finishing the fine-tuning stage, the script will output the following information:
+
+```bash
+BASE_MODEL_NAME = XXX
+ADAPTER_MODEL_PATH = XXX
+PROMPT = XXX
+```
+
+Please fill in the information in the `hra_dreambooth_inference.ipynb`.
+
+## Generate images with the fine-tuned model
+
+To generate images with the fine-tuned model, simply run the jupyter notebook `hra_dreambooth_inference.ipynb` for visualization with `jupyter notebook` under `./examples/hra_dreambooth`.

examples/hra_dreambooth/download_dreambooth.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+echo -e "\nDownloading dreambooth dataset..."
+git clone https://github.com/google/dreambooth.git

examples/hra_dreambooth/hra_dreambooth_inference.ipynb

@@ -0,0 +1,92 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "acd7b15e",
+   "metadata": {},
+   "source": [
+    "# Dreambooth with HRA\n",
+    "This Notebook assumes that you already ran the train_dreambooth.py script to create your own adapter."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "acab479f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from diffusers import DiffusionPipeline\n",
+    "from diffusers.utils import check_min_version, get_logger\n",
+    "from peft import PeftModel\n",
+    "\n",
+    "# Will error if the minimal version of diffusers is not installed. Remove at your own risks.\n",
+    "check_min_version(\"0.10.0.dev0\")\n",
+    "\n",
+    "logger = get_logger(__name__)\n",
+    "\n",
+    "# Please fill in the parameters that are output by running train_dreambooth.sh.\n",
+    "BASE_MODEL_NAME = \"stabilityai/stable-diffusion-2-1-base\"\n",
+    "ADAPTER_MODEL_PATH = \"logs/backpack-0\"\n",
+    "PROMPT = \"a qwe backpack in the jungle\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pipe = DiffusionPipeline.from_pretrained(\n",
+    "    BASE_MODEL_NAME,\n",
+    ")\n",
+    "pipe.to(\"cuda\")\n",
+    "pipe.unet = PeftModel.from_pretrained(pipe.unet, ADAPTER_MODEL_PATH + \"/unet\", adapter_name=\"default\")\n",
+    "pipe.text_encoder = PeftModel.from_pretrained(\n",
+    "    pipe.text_encoder, ADAPTER_MODEL_PATH + \"/text_encoder\", adapter_name=\"default\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = pipe(\n",
+    "    PROMPT,\n",
+    "    num_inference_steps=50,\n",
+    "    height=512,\n",
+    "    width=512,\n",
+    ").images[0]\n",
+    "image"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.14"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

examples/hra_dreambooth/requirements.txt

@@ -0,0 +1,13 @@
+transformers==4.36.2
+accelerate==0.25.0
+evaluate
+tqdm
+datasets==2.16.1
+diffusers==0.17.1
+Pillow
+huggingface_hub
+safetensors
+nb_conda_kernels
+ipykernel
+ipywidgets
+wandb==0.16.1