KaliYuga_BLIP+LoRA+Dreambooth_FineTuning.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/KaliYuga-ai/blip-lora-dreambooth-finetuning/blob/main/KaliYuga_BLIP%2BLoRA%2BDreambooth_FineTuning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# LoRA+DreamBooth Fine-tuning of Stable Diffusion Models (With BLIP Auto-Captioning)\n",
        "<div>\n",
        "<img src=\"https://images.squarespace-cdn.com/content/v1/6213c340453c3f502425776e/a432c21c-bb12-4f38-b5e2-1c12a3c403f6/Animated-Logo_1.gif\" width=\"150\"/>\n",
        "</div>\n",
        "\n",
        "\n",
        "[KaliYuga](https://twitter.com/KaliYuga_ai)'s simple fork of brian6091's [LoRA-Enabled Dreambooth notebook](https://github.com/brian6091/Dreambooth). \n",
        "\n",
        "In addition to some minor changes and rewording for clarity, **this fork adds a slightly modified version of BLIP dataset autocaptioning functionality** from [victorchall's EveryDream comapnion tools repo](https://github.com/victorchall/EveryDream) to brian6091's notebook.\n",
        "\n",
        "Once you've autocaptioned your datasets, you can use this same notebook to train Stable Diffusion models on those datasets using Dreambooth and/or Low-rank Adaptation (LoRA) approaches.\n",
        "\n",
        "I'm hoping to add BLIP2 functionality at a later date.\n",
        "\n",
        "Tested with [Stable Diffusion v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) and [Stable Diffusion v2-base](https://huggingface.co/stabilityai/stable-diffusion-2-base).\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "You can support victorchall's awesome EveryDream on \n",
        "[Patreon](https://www.patreon.com/everydream) or at [Kofi](https://ko-fi.com/everydream)!\n",
        "\n",
        "Click below to buy brian6091, whose notebook this is a fork of, a coffee!\n",
        "\n",
        "[<a href=\"https://www.buymeacoffee.com/jvsurfsqv\" target=\"_blank\"><img src=\"https://cdn.buymeacoffee.com/buttons/v2/default-yellow.png\" height=\"40px\" width=\"140px\" alt=\"Buy Me A Coffee\"></a>](https://www.buymeacoffee.com/jvsurfsqv)\n",
        "\n"
      ],
      "metadata": {
        "id": "EDwM5xLRggN5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## Mount Google Drive to access datasets, if initial model stored there, or you want to direct outputs there\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ],
      "metadata": {
        "id": "2wBnGW_v00va"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1: [OPTIONAL] BLIP Auto-Captioning\n",
        "\n",
        "**OPTIONAL Section**! \n",
        "\n",
        "If you don't want to use BLIP, or if your dataset is already labeled, you can skip this step!\n",
        "\n",
        "This section is taken (and modified slightly) from [victorchall](https://github.com/victorchall/EveryDream2trainer#docs)'s EveryDream 2 training notebook."
      ],
      "metadata": {
        "id": "Vm_Dbu7kpJri"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##1.1 Download Repo\n",
        "!git clone https://github.com/victorchall/EveryDream.git\n",
        "# Set working directory\n",
        "%cd EveryDream"
      ],
      "metadata": {
        "id": "Dx9WVBhTpPht"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##1.2 Install Requirements\n",
        "!pip install torch=='1.12.1+cu113' 'torchvision==0.13.1+cu113' --extra-index-url https://download.pytorch.org/whl/cu113\n",
        "!pip install pandas>='1.3.5'\n",
        "!git clone https://github.com/salesforce/BLIP scripts/BLIP\n",
        "!pip install timm\n",
        "!pip install fairscale=='0.4.4'\n",
        "!pip install transformers=='4.19.2'\n",
        "!pip install timm\n",
        "!pip install aiofiles"
      ],
      "metadata": {
        "id": "K04HRVgesXlN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 1.3 Upload your dataset to Google Drive (NOT to the Colab instance--doing this is very slow).\n",
        "Name it something you'll be able to remember/find easily. "
      ],
      "metadata": {
        "id": "PYXo-Tp6scbZ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "##1.4 Auto-Captioning\n",
        "\n",
        "*You cannot have commented lines between uncommented lines.  If you uncomment a line below, move it above any other commented lines.*\n",
        "\n",
        "*!python must remain the first line.*\n",
        "\n",
        "Default params should work fairly well.\n"
      ],
      "metadata": {
        "id": "CcziwfpeskwD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python scripts/auto_caption.py \\\n",
        "--img_dir /content/drive/MyDrive/YourDataset \\\n",
        "--out_dir /content/drive/MyDrive/output \\\n",
        "#--format mrwho \\\n",
        "#--min_length 34 \\\n",
        "#--q_factor 1.3 \\\n",
        "#--nucleus \\\n",
        "\n",
        "#IMPORTANT NOTE: replace \"[YourDataset]\" in the --img_dir line with your dataset folder name\n",
        "##ANOTHER NOTE: if you want to save over your original file names instead of making a new directory for your output files,\n",
        "##simply make your output path the same as your input path."
      ],
      "metadata": {
        "id": "rqr5sUIVskYF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 1.5\n",
        "Once your dataset is autocaptioned, download a bulk renaming app (I use [NameChanger](https://mrrsoftware.com/namechanger/) on Mac), download your labeled dataset and bulk-prepend all your file names with your rare token. After this, re-upload the dataset to Drive. This will ensure each filename acts as an instance prompt *which will include your rare token*."
      ],
      "metadata": {
        "id": "gKH1t_JVwuY1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "rilOGDYM2u0R"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2: Set Up Training Params"
      ],
      "metadata": {
        "id": "JQDQjocLwAup"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2.1: Install dependencies (takes about 1 minute)"
      ],
      "metadata": {
        "id": "LouTFfVYhRei"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "!cd /content/\n",
        "!git clone https://github.com/brian6091/Dreambooth --branch main --single-branch\n",
        "!pip install -r \"Dreambooth/requirements.txt\"\n",
        "!pip install -U --pre triton\n",
        "!pip install torchinfo\n",
        "\n",
        "!git clone https://github.com/brian6091/lora --branch v0.0.5 --single-branch\n",
        "!python -m pip install /content/lora/"
      ],
      "metadata": {
        "id": "PmsR_IPcvp7v"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title xformers\n",
        "#%%capture\n",
        "\n",
        "!nvidia-smi -L\n",
        "\n",
        "# Tested with Tesla T4 and A100 GPUs\n",
        "!pip install xformers==0.0.16rc425\n",
        "# May complain about some incompatibilities, which are resolved by upgrading the following:\n",
        "#!pip install -U --pre torchvision\n",
        "#!pip install -U --pre torchtext\n",
        "#!pip install -U --pre torchaudio"
      ],
      "metadata": {
        "id": "7tfenRTEQz_R"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2.2 Choose Models To Train From"
      ],
      "metadata": {
        "id": "8C3LrrpBvYn-"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "cellView": "form",
        "id": "YI6dHfQ8iqMg"
      },
      "outputs": [],
      "source": [
        "#@title ### 2.2.1: Name or path to initial model and VAE\n",
        "#@markdown Obligatory (e.g., runwayml/stable-diffusion-v1-5, stabilityai/stable-diffusion-2-base, or full path to model in diffusers format)\n",
        "MODEL_NAME_OR_PATH = \"runwayml/stable-diffusion-v1-5\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Optional (e.g., stabilityai/sd-vae-ft-mse), leaving empty will default to VAE packaged with the model\n",
        "VAE_NAME_OR_PATH = \"\" #@param {type:\"string\"}\n",
        "#if VAE_NAME_OR_PATH==\"\":\n",
        "#  VAE_NAME_OR_PATH = None\n",
        "\n",
        "#@markdown (Not yet implemented), leaving empty will default to text encoder packaged with the model\n",
        "TEXT_ENCODER_NAME_OR_PATH = \"\" #@param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ### 2.2.2 Hugging Face 🤗 credentials\n",
        "\n",
        "#@markdown If initiating training from official stable diffusion checkpoints (e.g., [stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5)), you must accept the license before using the model. You'll need a [🤗 Hugging Face](https://huggingface.co/) account to do so, after which you can [generate a login token](https://huggingface.co/settings/tokens) and paste it here.\n",
        "from huggingface_hub import login\n",
        "\n",
        "HUGGINGFACE_TOKEN = \"\" #@param {type:\"string\"}\n",
        "login(HUGGINGFACE_TOKEN)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "BGzA0N0C0pDB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2.3 Set up experiment parameters"
      ],
      "metadata": {
        "id": "fxP_d_n4mW2_"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "E1cJJ-P8jPhx",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title ## Training parameters\n",
        "\n",
        "import os\n",
        "from IPython.display import Markdown as md\n",
        "\n",
        "#@markdown Unique token for specific subject\n",
        "INSTANCE_TOKEN= \"raretoken\" #@param{type: 'string'}\n",
        "\n",
        "#@markdown Use image captions? Captions can be either the image filename, or a separate text file (that must be named identically to the image but w/ extension .txt). If a separate .txt file exists, filename is ignored.\n",
        "USE_IMAGE_CAPTIONS = True #@param {type:\"boolean\"}\n",
        "USE_IMAGE_CAPTIONS_FLAG = \"\"\n",
        "if USE_IMAGE_CAPTIONS:\n",
        "  USE_IMAGE_CAPTIONS_FLAG='--use_image_captions'\n",
        "\n",
        "#@markdown Path to instance images. Filenames are irrelevant, unless images are captioned *and* captions are not separate textfiles, in which case INSTANCE_TOKEN should appear in relevant filenames as part of the caption. This is what we did to our dataset in step 1.5.\n",
        "INSTANCE_DIR=\"/content/gdrive/MyDrive/[bliplabeleddatasetisbest]\" #@param{type: 'string'}\n",
        "\n",
        "RESOLUTION = 512 #@param{type: 'number'}\n",
        "\n",
        "TRAIN_BATCH_SIZE = 1 #@param{type: 'number'}\n",
        "\n",
        "GRADIENT_ACCUMULATION_STEPS = 1  #@param{type: 'number'}\n",
        "\n",
        "GRADIENT_CHECKPOINTING = True #@param {type:\"boolean\"}\n",
        "GRADIENT_CHECKPOINTING_FLAG=\"\"\n",
        "if GRADIENT_CHECKPOINTING:\n",
        "  GRADIENT_CHECKPOINTING_FLAG='--gradient_checkpointing'\n",
        "\n",
        "ENABLE_PRIOR_PRESERVATION = True #@param {type:\"boolean\"} \n",
        "ENABLE_PRIOR_PRESERVATION_FLAG=\"\"\n",
        "if ENABLE_PRIOR_PRESERVATION:\n",
        "  ENABLE_PRIOR_PRESERVATION_FLAG='--with_prior_preservation'\n",
        "\n",
        "#@markdown Prior loss weight. Note that if you set this to 0, but enable prior preservation and provide a CLASS_DIR, you can still monitor class loss.\n",
        "PRIOR_LOSS_WEIGHT = 1.0 #@param {type:\"number\"} \n",
        "\n",
        "#@markdown If using prior preservation, specify a path to class images. \n",
        "CLASS_DIR=\"/content/gdrive/MyDrive/[myregularizationimages]\" #@param{type: 'string'}\n",
        "if (CLASS_DIR !=\"\") and os.path.exists(str(CLASS_DIR)):\n",
        "  CLASS_DIR=CLASS_DIR\n",
        "elif (CLASS_DIR !=\"\") and not os.path.exists(str(CLASS_DIR)):\n",
        "  CLASS_DIR=input('\u001b[1;31mThe folder specified does not exist, use the colab file explorer to copy the path :')\n",
        "\n",
        "#@markdown Prompt for prior preservation class (e.g., 'person', 'a photo of a man', 'dog'). Used to generate regularization images. (The notebook this is a fork of says it's ignored if USE_IMAGE_CAPTIONS is checked, but this doesn't seem to be the case.)\n",
        "CLASS_PROMPT=\"[my class prompt]\" #@param {type:\"string\"}\n",
        "#@markdown Instance prompt, {SKS} will be automatically replaced by INSTANCE_TOKEN defined above. Ignored if USE_IMAGE_CAPTIONS checked--each individual caption with your token included will act as an instance prompt in that case.\n",
        "INSTANCE_PROMPT=\"a photo of {SKS} person\" #@param {type:\"string\"}\n",
        "INSTANCE_PROMPT=INSTANCE_PROMPT.replace(\"{SKS}\",INSTANCE_TOKEN)\n",
        "\n",
        "#@markdown Specify the number of class images used if prior preservation is enabled. If there are not enough images in CLASS_DIR (or CLASS_DIR is empty), additional images will be generated. A value of 1500 seems adequate for datasets of 500 or more, but may need to be lowered for smaller datasets.\n",
        "MIN_NUM_CLASS_IMAGES=1500 #@param{type: 'number'}\n",
        "\n",
        "#@markdown Batch size for generating class images \n",
        "SAMPLE_BATCH_SIZE = 1 #@param{type: 'number'}\n",
        "\n",
        "#@markdown Number of training iterations, e.g., # instance images * 100\n",
        "STEPS = 100000 #@param{type: 'number'}\n",
        "\n",
        "#@markdown Random number generator seed\n",
        "SEED = 1275017 #@param{type: 'number'}\n",
        "\n",
        "#@markdown Enable text encoder training? I leave this turned off, as I get much better results without it!\n",
        "TRAIN_TEXT_ENCODER = False #@param{type: 'boolean'}\n",
        "TRAIN_TEXT_ENCODER_FLAG=\"\"\n",
        "if TRAIN_TEXT_ENCODER:\n",
        "  TRAIN_TEXT_ENCODER_FLAG=\"--train_text_encoder\"\n",
        "\n",
        "#@markdown ## ADAM optimizer settings\n",
        "\n",
        "#@markdown Use 8-bit ADAM\n",
        "USE_8BIT_ADAM = True #@param {type:\"boolean\"} \n",
        "USE_8BIT_ADAM_FLAG=\"\"\n",
        "if USE_8BIT_ADAM:\n",
        "  USE_8BIT_ADAM_FLAG='--use_8bit_adam'\n",
        "\n",
        "#@markdown The exponential decay rate for the 1st moment estimates (the beta1 parameter for the Adam optimizer).\n",
        "ADAM_BETA1 = 0.9 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown The exponential decay rate for the 2nd moment estimates (the beta2 parameter for the Adam optimizer).\n",
        "ADAM_BETA2 = 0.999 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Weight decay magnitude for the Adam optimizer.\n",
        "ADAM_WEIGHT_DECAY = 1e-2 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Epsilon value for the Adam optimizer.\n",
        "ADAM_EPSILON = 1e-08 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown \"fp16\", \"bf16\", or \"no\" according to available VRAM. fl16 is a good option for lower VRAM.\n",
        "MIXED_PRECISION = \"no\" #@param{type: 'string'}\n",
        "\n",
        "#@markdown ## Learning rate parameters\n",
        "LR_SCHEDULE = \"cosine_with_restarts\" #@param [\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\", \"constant\", \"constant_with_warmup\"]\n",
        "LR = 3e-4 #@param{type: 'number'}\n",
        "#@markdown If training the text encoder, a different learning rate can be applied\n",
        "LR_TEXT_ENCODER = 1e-5 #@param{type: 'number'}\n",
        "LR_WARMUP_STEPS = 50 #@param{type: 'number'}\n",
        "#@markdown Applies only for cosine_with_restarts schedule\n",
        "LR_COSINE_NUM_CYCLES = 5 #@param{type: 'number'}"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## 2.4 (Experimental) [Data augmentation](https://journalofbigdata.springeropen.com/articles/10.1186/s40537-019-0197-0/)\n",
        "#@markdown Transformations to apply to images (both instance and class).\n",
        "#@markdown Useful to minimize the work of cropping and manually preparing images.\n",
        "#@markdown This may be useful for certain applications, such as training a style, where there may not be a specific subject in each image.\n",
        "#@markdown In this case, try not pre-cropping images, and instead enable random cropping, which presents to the network a randomly cropped (RESOLUTION X RESOLUTION) chunk of the original image selected for that iteration.\n",
        "#@markdown AUGMENT_MIN_RESOLUTION allows you to adjust how much of the image you will crop. So if you are training for RESOLUTION=512, setting AUGMENT_MIN_RESOLUTION will give you two crops (on average) for the shortest image dimension.\n",
        "\n",
        "#@markdown Resize image so that smallest dimension = AUGMENT_MIN_RESOLUTION (maintaining aspect ratio). Leave empty to skip.\n",
        "AUGMENT_MIN_RESOLUTION = None #@param{type: 'number'}\n",
        "AUGMENT_MIN_RESOLUTION_FLAG = \"\"\n",
        "if AUGMENT_MIN_RESOLUTION is not None:\n",
        "  AUGMENT_MIN_RESOLUTION = int(AUGMENT_MIN_RESOLUTION)\n",
        "  AUGMENT_MIN_RESOLUTION_FLAG = f\"--augment_min_resolution={AUGMENT_MIN_RESOLUTION}\"\n",
        "\n",
        "#@markdown If not enabled, defaults to center crop (which will do nothing if your images are already square at the RESOLUTION set above).\n",
        "AUGMENT_RANDOM_CROP = False #@param{type: 'boolean'}\n",
        "AUGMENT_CENTER_CROP_FLAG=\"--augment_center_crop\"\n",
        "if AUGMENT_RANDOM_CROP:\n",
        "  AUGMENT_CENTER_CROP_FLAG=\"\"\n",
        "\n",
        "#@markdown Randomly flip image horizontally. Not recommended if asymmetry is important (e.g., faces).\n",
        "AUGMENT_HFLIP = True #@param{type: 'boolean'}\n",
        "AUGMENT_HFLIP_FLAG=\"\"\n",
        "if AUGMENT_HFLIP:\n",
        "  AUGMENT_HFLIP_FLAG=\"--augment_hflip\""
      ],
      "metadata": {
        "id": "zLGpiF7xsLcb",
        "cellView": "form"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##2.5 (Experimental) other training parameters\n",
        "\n",
        "#@markdown ## [LoRA: Low-Rank Adaptation](https://arxiv.org/abs/2106.09685v2)\n",
        "#@markdown Uses [clonesimo's implementation](https://github.com/cloneofsimo/lora).\n",
        "\n",
        "#@markdown Read about brian6091's original notebook [here](https://github.com/cloneofsimo/lora/discussions/37).\n",
        "\n",
        "\n",
        "USE_LORA = True #@param{type: 'boolean'}\n",
        "USE_LORA_FLAG=\"\"\n",
        "if USE_LORA:\n",
        "  USE_LORA_FLAG=\"--use_lora\"\n",
        "\n",
        "#@markdown Rank of LoRA update matrix\n",
        "LORA_RANK = 4 #@param{type: 'number'}\n",
        "\n",
        "#@markdown ## [Drop text-conditioning to improve classifier-free guidance sampling](https://arxiv.org/abs/2207.12598)\n",
        "\n",
        "#@markdown Probability that image (applies to both instance and class images) will be selected for dropout (INSTANCE_PROMPT/CLASS_PROMPT will be replaced with UNCONDITIONAL_PROMPT)\n",
        "CONDITIONING_DROPOUT_PROB = 0.0 #@param{type: 'number'}\n",
        "#@markdown Defaults to an empty prompt. Unsure whether anything else would be useful.\n",
        "UNCONDITIONAL_PROMPT = \" \" #@param{type: 'string'}\n",
        "\n",
        "#@markdown ## Exponentially-weight moving average weights (unet only). Will not run on Tesla T4 (out of memory).\n",
        "USE_EMA = False #@param{type: 'boolean'}\n",
        "USE_EMA_FLAG=\"\"\n",
        "if USE_EMA:\n",
        "  USE_EMA_FLAG=\"--use_ema\"\n",
        "EMA_INV_GAMMA = 1.0 #@param{type: 'number'}\n",
        "EMA_POWER = 0.75 #@param{type: 'number'}\n",
        "EMA_MIN_VALUE = 0 #@param{type: 'number'}\n",
        "EMA_MAX_VALUE = 0.9999 #@param{type: 'number'}"
      ],
      "metadata": {
        "id": "LU7NC1Pkr47k",
        "cellView": "form"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##2.6: Where should outputs get saved?\n",
        "\n",
        "#@markdown Trained models (and intermediates) saved here\n",
        "OUTPUT_DIR=\"/content/gdrive/MyDrive/[myoutputdirectory]\" #@param{type: 'string'}\n",
        "\n",
        "#@markdown Training logs saved here\n",
        "LOGGING_DIR=\"/content/gdrive/MyDrive/[myoutputdirectory]/logs\" #@param{type: 'string'}\n",
        "\n",
        "if not os.path.exists(LOGGING_DIR):\n",
        "  !mkdir -p \"$LOGGING_DIR\"\n",
        "\n",
        "LOG_GPU = True #@param{type: 'boolean'}\n",
        "if LOG_GPU:\n",
        "  LOG_GPU_FLAG=\"--log_gpu\"\n",
        "else:\n",
        "  LOG_GPU_FLAG=\"\"\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Lji3GATOYIg_"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##2.7 Setup saving of intermediate models\n",
        "#@markdown To save intermediate checkpoints, set START_SAVING_FROM_STEP < STEPS\n",
        "\n",
        "#@markdown Number of steps between intermediate saves\n",
        "SAVE_CHECKPOINT_EVERY = 500 #@param{type: 'number'}\n",
        "if SAVE_CHECKPOINT_EVERY==None:\n",
        "  SAVE_CHECKPOINT_EVERY = STEPS+1\n",
        "\n",
        "START_SAVING_FROM_STEP=500 #@param{type: 'number'}\n",
        "if START_SAVING_FROM_STEP==None:\n",
        "  START_SAVING_FROM_STEP=STEPS\n",
        "\n",
        "#@markdown At each intermediate checkpoint, infer this many samples using SAVE_SAMPLE_PROMPT\n",
        "N_SAVE_SAMPLES=2 #@param{type: 'number'}\n",
        "\n",
        "#@markdown {SKS} is automatically replaced by INSTANCE_TOKEN. Give multiple prompts using // as a separator\n",
        "SAVE_SAMPLE_PROMPT= \"a painting of a woman in the style of {SKS}//a beautiful beach in the style of {SKS} \" #@param{type: 'string'}\n",
        "if SAVE_SAMPLE_PROMPT==\"\":\n",
        "  SAVE_SAMPLE_PROMPT=None\n",
        "else:\n",
        "  SAVE_SAMPLE_PROMPT=SAVE_SAMPLE_PROMPT.replace(\"{SKS}\",INSTANCE_TOKEN)\n",
        "\n",
        "#@markdown The negative prompt, on the other hand, applies to all SAVE_SAMPLE_PROMPTs\n",
        "SAVE_SAMPLE_NEGATIVE_PROMPT=\"border\" #@param{type: 'string'}"
      ],
      "metadata": {
        "cellView": "form",
        "id": "m9wXEuCnXn_0"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "8X0ohnP-2yrk"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 3: Train!"
      ],
      "metadata": {
        "id": "-sWFt9CCYkMO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## 3.1: (optional) Tensorboard visualization of loss and learning rate\n",
        "#@markdown Once the Tensorboard panel is launched (takes a good 10 seconds), click on the gear icon in upper right, and check Reload data. Then, after launching training in the next cell, click on TIME SERIES in upper left to see updates.\n",
        "#%load_ext tensorboard\n",
        "!rm -rf /content/logs\n",
        "%reload_ext tensorboard\n",
        "%tensorboard --logdir $LOGGING_DIR "
      ],
      "metadata": {
        "id": "WfL1DJnZYr-S",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## 3.2: Launch training\n",
        "!lsb_release -a | grep Description\n",
        "!pip freeze | grep diffusers\n",
        "!pip freeze | grep lora-diffusion\n",
        "!pip freeze | grep torchvision\n",
        "!pip freeze | grep transformers\n",
        "!pip freeze | grep xformers\n",
        "!accelerate env\n",
        "\n",
        "!accelerate launch \\\n",
        "    --mixed_precision=$MIXED_PRECISION \\\n",
        "    --num_machines=1 \\\n",
        "    --num_processes=1 \\\n",
        "    --dynamo_backend=\"no\" \\\n",
        "    /content/Dreambooth/train.py \\\n",
        "    $USE_LORA_FLAG \\\n",
        "    --lora_rank=$LORA_RANK \\\n",
        "    $TRAIN_TEXT_ENCODER_FLAG \\\n",
        "    --pretrained_model_name_or_path=$MODEL_NAME_OR_PATH \\\n",
        "    --pretrained_vae_name_or_path=$VAE_NAME_OR_PATH \\\n",
        "    --instance_data_dir=\"$INSTANCE_DIR\" \\\n",
        "    --class_data_dir=\"$CLASS_DIR\" \\\n",
        "    --output_dir=\"$OUTPUT_DIR\" \\\n",
        "    --logging_dir=\"$LOGGING_DIR\" \\\n",
        "    $LOG_GPU_FLAG \\\n",
        "    $ENABLE_PRIOR_PRESERVATION_FLAG \\\n",
        "    --prior_loss_weight=$PRIOR_LOSS_WEIGHT \\\n",
        "    --instance_prompt=\"$INSTANCE_PROMPT\" \\\n",
        "    --class_prompt=\"$CLASS_PROMPT\" \\\n",
        "    $USE_IMAGE_CAPTIONS_FLAG \\\n",
        "    --conditioning_dropout_prob=$CONDITIONING_DROPOUT_PROB \\\n",
        "    --unconditional_prompt=\"$UNCONDITIONAL_PROMPT\" \\\n",
        "    --seed=$SEED \\\n",
        "    --resolution=$RESOLUTION \\\n",
        "    --train_batch_size=$TRAIN_BATCH_SIZE \\\n",
        "    --gradient_accumulation_steps=$GRADIENT_ACCUMULATION_STEPS \\\n",
        "    $GRADIENT_CHECKPOINTING_FLAG \\\n",
        "    --mixed_precision=$MIXED_PRECISION \\\n",
        "    $USE_8BIT_ADAM_FLAG \\\n",
        "    --adam_beta1=$ADAM_BETA1 \\\n",
        "    --adam_beta2=$ADAM_BETA2 \\\n",
        "    --adam_weight_decay=$ADAM_WEIGHT_DECAY \\\n",
        "    --adam_epsilon=$ADAM_EPSILON \\\n",
        "    --learning_rate=$LR \\\n",
        "    --learning_rate_text=$LR_TEXT_ENCODER \\\n",
        "    --lr_scheduler=$LR_SCHEDULE \\\n",
        "    --lr_warmup_steps=$LR_WARMUP_STEPS \\\n",
        "    --lr_cosine_num_cycles=$LR_COSINE_NUM_CYCLES \\\n",
        "    $USE_EMA_FLAG \\\n",
        "    --ema_inv_gamma=$EMA_INV_GAMMA \\\n",
        "    --ema_power=$EMA_POWER \\\n",
        "    --ema_min_value=$EMA_MIN_VALUE \\\n",
        "    --ema_max_value=$EMA_MAX_VALUE \\\n",
        "    --max_train_steps=$STEPS \\\n",
        "    --num_class_images=$MIN_NUM_CLASS_IMAGES \\\n",
        "    --sample_batch_size=$SAMPLE_BATCH_SIZE \\\n",
        "    --save_min_steps=$START_SAVING_FROM_STEP \\\n",
        "    --save_interval=$SAVE_CHECKPOINT_EVERY \\\n",
        "    --n_save_sample=$N_SAVE_SAMPLES \\\n",
        "    --save_sample_prompt=\"$SAVE_SAMPLE_PROMPT\" \\\n",
        "    --save_sample_negative_prompt=\"$SAVE_SAMPLE_NEGATIVE_PROMPT\" \\\n",
        "    $AUGMENT_MIN_RESOLUTION_FLAG \\\n",
        "    $AUGMENT_CENTER_CROP_FLAG \\\n",
        "    $AUGMENT_HFLIP_FLAG"
      ],
      "metadata": {
        "id": "Zim-xlhbYlej"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "8hbIqh8221CX"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#4: Do inference with trained model(s)\n",
        "\n",
        "Cells in this section can be run to generate grids of images using the trained model(s). Useful for probing overtraining, concept bleeding, quality, etc."
      ],
      "metadata": {
        "id": "q9tudAPlex_2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##4.1: Some imports and utility functions\n",
        "import torch\n",
        "from diffusers import DiffusionPipeline, StableDiffusionPipeline, DPMSolverMultistepScheduler, AutoencoderKL\n",
        "from PIL import Image\n",
        "import os\n",
        "import json\n",
        "import random\n",
        "import string\n",
        "from lora_diffusion import monkeypatch_lora, tune_lora_scale\n",
        "\n",
        "device = \"cuda\"\n",
        "\n",
        "def image_grid(imgs, rows, cols):\n",
        "    assert len(imgs) == rows*cols\n",
        "    w, h = imgs[0].size\n",
        "    grid = Image.new('RGB', size=(cols*w, rows*h))\n",
        "    grid_w, grid_h = grid.size\n",
        "    for i, img in enumerate(imgs):\n",
        "        grid.paste(img, box=(i%cols*w, i//cols*h))\n",
        "    return grid\n",
        "\n",
        "def get_pipeline(model_name_or_path, \n",
        "                 vae_name_or_path=None, \n",
        "                 text_encoder_name_or_path=None,\n",
        "                 feature_extractor_name_or_path=None,\n",
        "                 revision=\"fp16\"):\n",
        "    #scheduler = DPMSolverMultistepScheduler.from_pretrained(model_name_or_path, subfolder=\"scheduler\")\n",
        "    scheduler = DPMSolverMultistepScheduler(\n",
        "        beta_start=0.00085,\n",
        "        beta_end=0.012,\n",
        "        beta_schedule=\"scaled_linear\",\n",
        "        num_train_timesteps=1000,\n",
        "        trained_betas=None,\n",
        "        prediction_type=\"epsilon\",\n",
        "        thresholding=False,\n",
        "        algorithm_type=\"dpmsolver++\",\n",
        "        solver_type=\"midpoint\",\n",
        "        lower_order_final=True,\n",
        "    )\n",
        "\n",
        "    pipe = DiffusionPipeline.from_pretrained(\n",
        "        model_name_or_path,\n",
        "        custom_pipeline=\"lpw_stable_diffusion\",\n",
        "        safety_checker=None,\n",
        "        revision=revision,\n",
        "        scheduler=scheduler,\n",
        "        vae=AutoencoderKL.from_pretrained(\n",
        "            vae_name_or_path or model_name_or_path,\n",
        "            subfolder=None if vae_name_or_path else \"vae\",\n",
        "            revision=None if vae_name_or_path else revision,\n",
        "            torch_dtype=torch.float16,\n",
        "        ),\n",
        "        feature_extractor=feature_extractor_name_or_path,\n",
        "        torch_dtype=torch.float16\n",
        "    ).to(\"cuda\")\n",
        "\n",
        "    #https://github.com/huggingface/diffusers/issues/1552\n",
        "    #pipe.enable_attention_slicing()\n",
        "    pipe.enable_xformers_memory_efficient_attention()\n",
        "    return pipe\n",
        "\n",
        "# Monkey patch LoRA pt files \n",
        "# Returns pipeline\n",
        "def get_lora_pipeline(model_dir, scale_unet=1.0, scale_text_encoder=1.0):\n",
        "    # Load untrained original model\n",
        "    pipe = get_pipeline(MODEL_NAME_OR_PATH, vae_name_or_path=VAE_NAME_OR_PATH)\n",
        "\n",
        "    print('Monkey patching unet pt file')\n",
        "    monkeypatch_lora(pipe.unet, torch.load(os.path.join(model_dir, \"lora_unet.pt\")))\n",
        "\n",
        "    print('Monkey patching text encoder pt file')\n",
        "    monkeypatch_lora(pipe.text_encoder, torch.load(os.path.join(model_dir, \"lora_text_encoder.pt\")), target_replace_module=[\"CLIPAttention\"])\n",
        "\n",
        "    tune_lora_scale(pipe.unet, scale_unet)\n",
        "    tune_lora_scale(pipe.text_encoder, scale_text_encoder)\n",
        "\n",
        "    return pipe\n",
        "\n",
        "def get_config(filename=None,\n",
        "               save_dir=None,\n",
        "               prompt=None, negative_prompt=None,\n",
        "               seeds=None,\n",
        "               num_samples=4,\n",
        "               width=512, height=512,\n",
        "               inference_steps=20,\n",
        "               guidance_scale=7.5,\n",
        "               ):\n",
        "    if filename==None:\n",
        "        num_prompts = len(prompt)\n",
        "        if seeds==None:\n",
        "            seeds = []\n",
        "            # fixed value seeds for easier comparision betwen subsequent runs/config files\n",
        "            for i in range(num_samples):\n",
        "                seeds.append(i * 1000000)\n",
        "        else:\n",
        "            num_samples = len(seeds)\n",
        "\n",
        "        tag = ''.join(random.choice(string.ascii_letters) for _ in range(8))\n",
        "        config = {\n",
        "            \"tag\": tag,\n",
        "            \"prompt\": prompt,\n",
        "            \"negative_prompt\": negative_prompt,\n",
        "            \"num_prompts\": num_prompts, \n",
        "            \"num_samples\": num_samples, \n",
        "            \"seeds\": seeds,\n",
        "            \"height\": height,\n",
        "            \"width\": width,\n",
        "            \"inference_steps\": inference_steps,\n",
        "            \"guidance_scale\": guidance_scale,\n",
        "        }\n",
        "\n",
        "        with open(os.path.join(save_dir, \"config_\"+tag+\".json\"), \"w\") as outfile:\n",
        "            json.dump(config, outfile)\n",
        "    else:\n",
        "        f = open(filename)\n",
        "        config = json.load(f)\n",
        "    \n",
        "    return config\n",
        "\n",
        "def get_images(pipe, sample_config, device=\"cuda\"):\n",
        "    generator = torch.Generator(\"cuda\")\n",
        "    with torch.autocast(device):\n",
        "        num_cfg = len(sample_config['guidance_scale'])\n",
        "        # Loop in order to use defined seed for each image in a batch\n",
        "        all_images = []\n",
        "        for i in range(sample_config['num_samples']):\n",
        "        #for _ in sample_config['num_samples']:\n",
        "            for cfg in sample_config['guidance_scale']:\n",
        "                # Manually generate latent\n",
        "                seed = sample_config['seeds'][i]\n",
        "                generator = generator.manual_seed(seed)\n",
        "                latent = torch.randn(\n",
        "                    (1, pipe.unet.in_channels, sample_config['height'] // 8, sample_config['width'] // 8),\n",
        "                    generator = generator,\n",
        "                    device = device\n",
        "                )\n",
        "                images = pipe(sample_config['prompt'],\n",
        "                    negative_prompt=sample_config['negative_prompt'],\n",
        "                    num_inference_steps=int(sample_config['inference_steps']),\n",
        "                    guidance_scale=cfg,\n",
        "                    latents=latent.repeat(sample_config['num_prompts'], 1, 1, 1),\n",
        "                ).images\n",
        "                all_images.extend(images)\n",
        "\n",
        "    grid = image_grid(all_images, rows=num_cfg*sample_config['num_samples'], cols=sample_config['num_prompts'])\n",
        "    return grid"
      ],
      "metadata": {
        "id": "-3nj_hjwGFCf",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##4.2: Specify which models to do inference with\n",
        "model_list = [os.path.join(OUTPUT_DIR,'1000'),\n",
        "              os.path.join(OUTPUT_DIR,'5000'),\n",
        "              os.path.join(OUTPUT_DIR,'10000'),\n",
        "              os.path.join(OUTPUT_DIR,'15000'),\n",
        "              os.path.join(OUTPUT_DIR,'20000'),\n",
        "              os.path.join(OUTPUT_DIR,'25000'),\n",
        "              os.path.join(OUTPUT_DIR,'30000'),\n",
        "              os.path.join(OUTPUT_DIR,'35000'), \n",
        "              os.path.join(OUTPUT_DIR,'45000'), \n",
        "              os.path.join(OUTPUT_DIR,'50000'),                                        \n",
        "]\n",
        "\n",
        "print(model_list)"
      ],
      "metadata": {
        "id": "w3Xv0Zks-fCp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## 4.3: Generate or load a configuration for inference\n",
        "\n",
        "config_name = None\n",
        "#config_name = os.path.join(OUTPUT_DIR, \"config_ZMasiqkP.json\")\n",
        "\n",
        "if config_name is None:\n",
        "    num_samples = 6\n",
        "    prompt = [\"photo of a cat\",\n",
        "              \"photo of a person\",\n",
        "              \"close-up studio portrait photo of Keanu Reeves, film, detail, studio lighting\",\n",
        "              \"{SKS}close-up studio portrait photo of a person, film, detail, studio lighting\",\n",
        "              \"{SKS} beautiful white (marble:1.1) bust of a person, highly detailed\",\n",
        "              \"{SKS} oil painting of a person on the beach\",\n",
        "    ]\n",
        "    negative_prompt = \"hands, nude, nudity, duplicate, frame, border\"\n",
        "    guidance_scale = [1.0, 3.0, 7.0, 15.0]\n",
        "\n",
        "    config = get_config(save_dir=OUTPUT_DIR,\n",
        "                        prompt=prompt, negative_prompt=negative_prompt,\n",
        "                        num_samples=num_samples,\n",
        "                        width=512, height=512, \n",
        "                        inference_steps=20, guidance_scale=guidance_scale\n",
        "                        )\n",
        "else:\n",
        "    config = get_config(filename=config_name)\n",
        "\n",
        "config['prompt'] = [sub.replace('{SKS}', INSTANCE_TOKEN) for sub in config['prompt']]\n",
        "print(config)"
      ],
      "metadata": {
        "id": "ub3iLhz0aBho"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## 4.4 Infer!\n",
        "\n",
        "LORA_SCALE_UNET = 1.0 #@param {type:\"slider\", min:0.0, max:2.0}\n",
        "LORA_SCALE_TENC = 1.0 #@param {type:\"slider\", min:0.0, max:2.0}\n",
        "\n",
        "for model in model_list:\n",
        "    print(model)\n",
        "    pipe = get_pipeline(model) if not USE_LORA else get_lora_pipeline(model, scale_unet=LORA_SCALE_UNET, scale_text_encoder=LORA_SCALE_TENC)\n",
        "    grid = get_images(pipe, config)\n",
        "    grid.save(os.path.join(OUTPUT_DIR, \"grid_\"+os.path.split(model)[1]+\"_\"+config['tag']+\".jpg\"), quality=90, optimize=True)\n",
        "    del pipe\n",
        "    if torch.cuda.is_available():\n",
        "        torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "ZMMHcQ55_PAb",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 4.5 (Optional) Generate grids for base model using same config\n",
        "model_name_or_path = MODEL_NAME_OR_PATH #'runwayml/stable-diffusion-v1-5'\n",
        "vae_name_or_path = VAE_NAME_OR_PATH #'stabilityai/sd-vae-ft-mse'\n",
        "pipe = get_pipeline(model_name_or_path, vae_name_or_path=vae_name_or_path)\n",
        "grid = get_images(pipe, config)\n",
        "grid.save(os.path.join(OUTPUT_DIR, \"grid_\"+os.path.split(model_name_or_path)[1]+\"_\"+config['tag']+\".jpg\"), quality=90, optimize=True)\n",
        "\n",
        "del pipe\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "Bq4VAh6EPhja",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "Odb5LDEb23xD"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 5: Convert to checkpoint (ckpt) format"
      ],
      "metadata": {
        "id": "W2WxeLuIqW2G"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!lora_add --path_1 runwayml/stable-diffusion-v1-5 --path_2 /content/gdrive/MyDrive/[youroutputdir]/[stepnumber]/lora_unet.pt --mode upl-ckpt-v2 --alpha 1.2 --output_path /content/gdrive/[outputdir]/[desiredckptname].ckpt\n"
      ],
      "metadata": {
        "id": "vUyA7tGPjIdb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Close Colab instance"
      ],
      "metadata": {
        "id": "_fE3xIr7lE2f"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import runtime\n",
        "runtime.unassign()"
      ],
      "metadata": {
        "id": "HYrrT17slGd5"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "premium"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}