fast-DreamBooth.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qEsNHTtVlbkV"
      },
      "source": [
        "# **fast-DreamBooth colab From https://github.com/TheLastBen/fast-stable-diffusion, if you face any issues, feel free to discuss them.** \n",
        "Keep your notebook updated for best experience. [Support](https://ko-fi.com/thelastben)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A4Bae3VP6UsE"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbKbx185zqlz"
      },
      "source": [
        "# Setting up the environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "QyvcqeiL65Tj"
      },
      "outputs": [],
      "source": [
        "#@markdown # Dependencies\n",
        "%%capture\n",
        "%cd /content/\n",
        "!git clone --branch updt https://github.com/TheLastBen/diffusers\n",
        "!pip install -q /content/diffusers\n",
        "!pip install -q accelerate==0.12.0\n",
        "!pip install -q OmegaConf\n",
        "!pip install -q wget\n",
        "!pip install -q open_clip_torch\n",
        "!pip install -q pillow==8.4\n",
        "!pip install -q torchsde\n",
        "!pip install -q pytorch_lightning\n",
        "!pip install -q huggingface_hub\n",
        "!pip install -U -q --no-cache-dir gdown\n",
        "!wget https://github.com/TheLastBen/fast-stable-diffusion/raw/main/Dreambooth/Deps\n",
        "!mv Deps Deps.7z\n",
        "!7z x Deps.7z\n",
        "!cp -r /content/usr/local/lib/python3.7/dist-packages /usr/local/lib/python3.7/\n",
        "!rm Deps.7z\n",
        "!rm -r /content/usr\n",
        "!sed -i 's@else prefix + \": \"@else prefix + \"\"@g' /usr/local/lib/python3.7/dist-packages/tqdm/std.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "1pld5ps87a1q"
      },
      "outputs": [],
      "source": [
        "#@markdown # xformers\n",
        "\n",
        "from subprocess import getoutput\n",
        "from IPython.display import HTML\n",
        "from IPython.display import clear_output\n",
        "import wget\n",
        "import time\n",
        "\n",
        "s = getoutput('nvidia-smi')\n",
        "if 'T4' in s:\n",
        "  gpu = 'T4'\n",
        "elif 'P100' in s:\n",
        "  gpu = 'P100'\n",
        "elif 'V100' in s:\n",
        "  gpu = 'V100'\n",
        "elif 'A100' in s:\n",
        "  gpu = 'A100'\n",
        "\n",
        "while True:\n",
        "    try: \n",
        "        gpu=='T4'or gpu=='P100'or gpu=='V100'or gpu=='A100'\n",
        "        break\n",
        "    except:\n",
        "        pass\n",
        "    print('\u001b[1;31mit seems that your GPU is not supported at the moment')\n",
        "    time.sleep(5)\n",
        "\n",
        "if (gpu=='T4'):\n",
        "  %pip install -q https://github.com/TheLastBen/fast-stable-diffusion/raw/main/precompiled/T4/xformers-0.0.13.dev0-py3-none-any.whl\n",
        "  \n",
        "elif (gpu=='P100'):\n",
        "  %pip install -q https://github.com/TheLastBen/fast-stable-diffusion/raw/main/precompiled/P100/xformers-0.0.13.dev0-py3-none-any.whl\n",
        "\n",
        "elif (gpu=='V100'):\n",
        "  %pip install -q https://github.com/TheLastBen/fast-stable-diffusion/raw/main/precompiled/V100/xformers-0.0.13.dev0-py3-none-any.whl\n",
        "\n",
        "elif (gpu=='A100'):\n",
        "  %pip install -q https://github.com/TheLastBen/fast-stable-diffusion/raw/main/precompiled/A100/xformers-0.0.13.dev0-py3-none-any.whl\n",
        "\n",
        "clear_output()\n",
        "print('\u001b[1;32mDONE !')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R3SsbIlxw66N"
      },
      "source": [
        "# Model Download"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "O3KHGKqyeJp9"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import time\n",
        "from IPython.display import clear_output\n",
        "from IPython.utils import capture\n",
        "\n",
        "#@markdown - Skip this cell if you are loading a previous session\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "Model_Version = \"1.5\" #@param [ \"1.5\", \"V2-512px\", \"V2-768px\"]\n",
        "\n",
        "#@markdown - Choose which version to finetune.\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "with capture.capture_output() as cap: \n",
        "  %cd /content/\n",
        "\n",
        "Huggingface_Token = \"\" #@param {type:\"string\"}\n",
        "token=Huggingface_Token\n",
        "\n",
        "#@markdown - Leave EMPTY if you're using the v2 model.\n",
        "#@markdown - Make sure you've accepted the terms in https://huggingface.co/runwayml/stable-diffusion-v1-5\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "Path_to_HuggingFace= \"\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown - Load and finetune a model from Hugging Face, use the format \"profile/model\" like : runwayml/stable-diffusion-v1-5\n",
        "\n",
        "#@markdown Or\n",
        "\n",
        "CKPT_Path = \"\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Or\n",
        "\n",
        "CKPT_Link = \"\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown - A CKPT direct link, huggingface CKPT link or a shared CKPT from gdrive.\n",
        "#@markdown ---\n",
        "\n",
        "Compatibility_Mode=\"\" #@param {type:\"boolean\"}\n",
        "#@markdown - Enable only if you're getting conversion errors.\n",
        "\n",
        "\n",
        "def downloadmodel():\n",
        "  token=Huggingface_Token\n",
        "  if token==\"\":\n",
        "      token=input(\"Insert your huggingface token :\")\n",
        "  if os.path.exists('/content/stable-diffusion-v1-5'):\n",
        "    !rm -r /content/stable-diffusion-v1-5\n",
        "  clear_output()\n",
        "\n",
        "  %cd /content/\n",
        "  clear_output()\n",
        "  !mkdir /content/stable-diffusion-v1-5\n",
        "  %cd /content/stable-diffusion-v1-5\n",
        "  !git init\n",
        "  !git lfs install --system --skip-repo\n",
        "  !git remote add -f origin  \"https://USER:{token}@huggingface.co/runwayml/stable-diffusion-v1-5\"\n",
        "  !git config core.sparsecheckout true\n",
        "  !echo -e \"scheduler\\ntext_encoder\\ntokenizer\\nunet\\nmodel_index.json\" > .git/info/sparse-checkout\n",
        "  !git pull origin main\n",
        "  if os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "    !git clone \"https://USER:{token}@huggingface.co/stabilityai/sd-vae-ft-mse\"\n",
        "    !mv /content/stable-diffusion-v1-5/sd-vae-ft-mse /content/stable-diffusion-v1-5/vae\n",
        "    !rm -r /content/stable-diffusion-v1-5/.git\n",
        "    %cd /content/stable-diffusion-v1-5\n",
        "    !rm model_index.json\n",
        "    time.sleep(1)    \n",
        "    wget.download('https://raw.githubusercontent.com/TheLastBen/fast-stable-diffusion/main/Dreambooth/model_index.json')\n",
        "    !sed -i 's@\"clip_sample\": false@@g' /content/stable-diffusion-v1-5/scheduler/scheduler_config.json\n",
        "    !sed -i 's@\"trained_betas\": null,@\"trained_betas\": null@g' /content/stable-diffusion-v1-5/scheduler/scheduler_config.json\n",
        "    !sed -i 's@\"sample_size\": 256,@\"sample_size\": 512,@g' /content/stable-diffusion-v1-5/vae/config.json  \n",
        "    %cd /content/    \n",
        "    clear_output()\n",
        "    print('\u001b[1;32mDONE !')\n",
        "  else:\n",
        "    while not os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "         print('\u001b[1;31mMake sure you accepted the terms in https://huggingface.co/runwayml/stable-diffusion-v1-5')\n",
        "         time.sleep(5)\n",
        "\n",
        "\n",
        "def newdownloadmodel():\n",
        "\n",
        "  %cd /content/\n",
        "  clear_output()\n",
        "  !mkdir /content/stable-diffusion-v2-768\n",
        "  %cd /content/stable-diffusion-v2-768\n",
        "  !git init\n",
        "  !git lfs install --system --skip-repo\n",
        "  !git remote add -f origin  \"https://USER:{token}@huggingface.co/stabilityai/stable-diffusion-2\"\n",
        "  !git config core.sparsecheckout true\n",
        "  !echo -e \"scheduler\\ntext_encoder\\ntokenizer\\nunet\\nvae\\nmodel_index.json\" > .git/info/sparse-checkout\n",
        "  !git pull origin main\n",
        "  clear_output()\n",
        "  print('\u001b[1;32mDONE !')\n",
        "\n",
        "\n",
        "def newdownloadmodelb():\n",
        "\n",
        "  %cd /content/\n",
        "  clear_output()\n",
        "  !mkdir /content/stable-diffusion-v2-512\n",
        "  %cd /content/stable-diffusion-v2-512\n",
        "  !git init\n",
        "  !git lfs install --system --skip-repo\n",
        "  !git remote add -f origin  \"https://USER:{token}@huggingface.co/stabilityai/stable-diffusion-2-base\"\n",
        "  !git config core.sparsecheckout true\n",
        "  !echo -e \"scheduler\\ntext_encoder\\ntokenizer\\nunet\\nvae\\nmodel_index.json\" > .git/info/sparse-checkout\n",
        "  !git pull origin main\n",
        "  clear_output()\n",
        "  print('\u001b[1;32mDONE !')\n",
        "    \n",
        "\n",
        "if Path_to_HuggingFace != \"\":\n",
        "  if os.path.exists('/content/stable-diffusion-custom'):\n",
        "    !rm -r /content/stable-diffusion-custom\n",
        "  clear_output()\n",
        "  %cd /content/\n",
        "  clear_output()\n",
        "  !mkdir /content/stable-diffusion-custom\n",
        "  %cd /content/stable-diffusion-custom\n",
        "  !git init\n",
        "  !git lfs install --system --skip-repo\n",
        "  !git remote add -f origin  \"https://USER:{token}@huggingface.co/{Path_to_HuggingFace}\"\n",
        "  !git config core.sparsecheckout true\n",
        "  !echo -e \"scheduler\\ntext_encoder\\ntokenizer\\nunet\\nmodel_index.json\" > .git/info/sparse-checkout\n",
        "  !git pull origin main\n",
        "  if os.path.exists('/content/stable-diffusion-custom/unet/diffusion_pytorch_model.bin'):\n",
        "    !git clone \"https://USER:{token}@huggingface.co/stabilityai/sd-vae-ft-mse\"\n",
        "    !mv /content/stable-diffusion-custom/sd-vae-ft-mse /content/stable-diffusion-custom/vae\n",
        "    !rm -r /content/stable-diffusion-custom/.git\n",
        "    %cd /content/stable-diffusion-custom\n",
        "    !rm model_index.json\n",
        "    time.sleep(1)\n",
        "    wget.download('https://raw.githubusercontent.com/TheLastBen/fast-stable-diffusion/main/Dreambooth/model_index.json')\n",
        "    !sed -i 's@\"clip_sample\": false@@g' /content/stable-diffusion-custom/scheduler/scheduler_config.json\n",
        "    !sed -i 's@\"trained_betas\": null,@\"trained_betas\": null@g' /content/stable-diffusion-custom/scheduler/scheduler_config.json\n",
        "    !sed -i 's@\"sample_size\": 256,@\"sample_size\": 512,@g' /content/stable-diffusion-custom/vae/config.json    \n",
        "    %cd /content/ \n",
        "    MODEL_NAME=\"/content/stable-diffusion-custom\"   \n",
        "    clear_output()\n",
        "    print('\u001b[1;32mDONE !')\n",
        "  else:\n",
        "    while not os.path.exists('/content/stable-diffusion-custom/unet/diffusion_pytorch_model.bin'):\n",
        "          print('\u001b[1;31mCheck the link you provided')\n",
        "          time.sleep(5)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "elif CKPT_Path !=\"\":\n",
        "  if os.path.exists('/content/stable-custom'):\n",
        "    !rm -r /content/stable-diffusion-custom\n",
        "  if os.path.exists(str(CKPT_Path)):\n",
        "    !mkdir /content/stable-diffusion-custom\n",
        "    with capture.capture_output() as cap:\n",
        "      if Compatibility_Mode:\n",
        "        !wget https://raw.githubusercontent.com/huggingface/diffusers/039958eae55ff0700cfb42a7e72739575ab341f1/scripts/convert_original_stable_diffusion_to_diffusers.py\n",
        "        !python /content/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path \"$CKPT_Path\" --dump_path /content/stable-diffusion-custom\n",
        "        !rm /content/convert_original_stable_diffusion_to_diffusers.py\n",
        "      else:           \n",
        "        !python /content/diffusers/scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path \"$CKPT_Path\" --dump_path /content/stable-diffusion-custom\n",
        "    if os.path.exists('/content/stable-diffusion-custom/unet/diffusion_pytorch_model.bin'):\n",
        "      !rm /content/v1-inference.yaml\n",
        "      clear_output()\n",
        "      MODEL_NAME=\"/content/stable-diffusion-custom\"\n",
        "      print('\u001b[1;32mDONE !')\n",
        "    else:\n",
        "      !rm /content/convert_original_stable_diffusion_to_diffusers.py\n",
        "      !rm /content/v1-inference.yaml\n",
        "      !rm -r /content/stable-diffusion-custom\n",
        "      while not os.path.exists('/content/stable-diffusion-custom/unet/diffusion_pytorch_model.bin'):\n",
        "        print('\u001b[1;31mConversion error, Insufficient RAM or corrupt CKPT, use a 4GB CKPT instead of 7GB')\n",
        "        time.sleep(5)\n",
        "  else:\n",
        "    while not os.path.exists(str(CKPT_Path)):\n",
        "       print('\u001b[1;31mWrong path, use the colab file explorer to copy the path')\n",
        "       time.sleep(5)\n",
        "  \n",
        "\n",
        "\n",
        "elif CKPT_Link !=\"\":   \n",
        "    if os.path.exists('/content/stable-diffusion-custom'):\n",
        "      !rm -r /content/stable-diffusion-custom   \n",
        "    !gdown --fuzzy -O model.ckpt $CKPT_Link\n",
        "    if os.path.exists('/content/model.ckpt'):\n",
        "      if os.path.getsize(\"/content/model.ckpt\") > 1810671599:\n",
        "        !mkdir /content/stable-diffusion-custom\n",
        "        with capture.capture_output() as cap: \n",
        "          if Compatibility_Mode:\n",
        "            !wget https://raw.githubusercontent.com/huggingface/diffusers/039958eae55ff0700cfb42a7e72739575ab341f1/scripts/convert_original_stable_diffusion_to_diffusers.py\n",
        "            !python /content/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path /content/model.ckpt --dump_path /content/stable-diffusion-custom\n",
        "            !rm /content/convert_original_stable_diffusion_to_diffusers.py            \n",
        "          else:           \n",
        "            !python /content/diffusers/scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path /content/model.ckpt --dump_path /content/stable-diffusion-custom\n",
        "        if os.path.exists('/content/stable-diffusion-custom/unet/diffusion_pytorch_model.bin'):\n",
        "          clear_output()\n",
        "          MODEL_NAME=\"/content/stable-diffusion-custom\"\n",
        "          print('\u001b[1;32mDONE !')\n",
        "          !rm /content/v1-inference.yaml\n",
        "          !rm /content/model.ckpt\n",
        "        else:\n",
        "          if os.path.exists('/content/v1-inference.yaml'):\n",
        "            !rm /content/v1-inference.yaml\n",
        "          !rm /content/convert_original_stable_diffusion_to_diffusers.py\n",
        "          !rm -r /content/stable-diffusion-custom\n",
        "          !rm /content/model.ckpt\n",
        "          while not os.path.exists('/content/stable-diffusion-custom/unet/diffusion_pytorch_model.bin'):\n",
        "            print('\u001b[1;31mConversion error, Insufficient RAM or corrupt CKPT, use a 4GB CKPT instead of 7GB')\n",
        "            time.sleep(5)\n",
        "      else:\n",
        "        while os.path.getsize('/content/model.ckpt') < 1810671599:\n",
        "           print('\u001b[1;31mWrong link, check that the link is valid')\n",
        "           time.sleep(5)\n",
        "    \n",
        "\n",
        "else:\n",
        "  if Model_Version==\"1.5\":\n",
        "    if not os.path.exists('/content/stable-diffusion-v1-5'):\n",
        "      downloadmodel()\n",
        "      MODEL_NAME=\"/content/stable-diffusion-v1-5\"\n",
        "    else:\n",
        "      MODEL_NAME=\"/content/stable-diffusion-v1-5\"\n",
        "      print(\"\u001b[1;32mThe v1.5 model already exist, using this model.\")\n",
        "  elif Model_Version==\"V2-512px\":\n",
        "    if not os.path.exists('/content/stable-diffusion-v2-512'):\n",
        "      newdownloadmodelb()\n",
        "      MODEL_NAME=\"/content/stable-diffusion-v2-512\"\n",
        "    else:\n",
        "      MODEL_NAME=\"/content/stable-diffusion-v2-512\"\n",
        "      print(\"\u001b[1;32mThe v2-512px model already exist, using this model.\")      \n",
        "  elif Model_Version==\"V2-768px\":\n",
        "    if not os.path.exists('/content/stable-diffusion-v2-768'):   \n",
        "      newdownloadmodel()\n",
        "      MODEL_NAME=\"/content/stable-diffusion-v2-768\"\n",
        "    else:\n",
        "      MODEL_NAME=\"/content/stable-diffusion-v2-768\"\n",
        "      print(\"\u001b[1;32mThe v2-768px model already exist, using this model.\")    "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0tN76Cj5P3RL"
      },
      "source": [
        "# Dreambooth"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "A1B299g-_VJo"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from IPython.display import clear_output\n",
        "from IPython.utils import capture\n",
        "from os import listdir\n",
        "from os.path import isfile\n",
        "import wget\n",
        "import time\n",
        "\n",
        "#@markdown #Create/Load a Session\n",
        "\n",
        "try:\n",
        "  MODEL_NAME\n",
        "  pass\n",
        "except:\n",
        "  MODEL_NAME=\"\"\n",
        "  \n",
        "PT=\"\"\n",
        "\n",
        "Captionned_instance_images = True\n",
        "\n",
        "Session_Name = \"\" #@param{type: 'string'}\n",
        "while Session_Name==\"\":\n",
        "  print('\u001b[1;31mInput the Session Name:') \n",
        "  Session_Name=input('')\n",
        "Session_Name=Session_Name.replace(\" \",\"_\")\n",
        "\n",
        "#@markdown - Enter the session name, it if it exists, it will load it, otherwise it'll create an new session.\n",
        "\n",
        "Session_Link_optional = \"\" #@param{type: 'string'}\n",
        "\n",
        "#@markdown - Import a session from another gdrive, the shared gdrive link must point to the specific session's folder that contains the trained CKPT, remove any intermediary CKPT if any.\n",
        "\n",
        "WORKSPACE='/content/gdrive/MyDrive/Fast-Dreambooth'\n",
        "\n",
        "if Session_Link_optional !=\"\":\n",
        "  print('\u001b[1;32mDownloading session...')\n",
        "with capture.capture_output() as cap:\n",
        "  %cd /content\n",
        "  if Session_Link_optional != \"\":\n",
        "    if not os.path.exists(str(WORKSPACE+'/Sessions')):\n",
        "      %mkdir -p $WORKSPACE'/Sessions'\n",
        "      time.sleep(1)\n",
        "    %cd $WORKSPACE'/Sessions'\n",
        "    !gdown --folder --remaining-ok -O $Session_Name  $Session_Link_optional\n",
        "    %cd $Session_Name\n",
        "    !rm -r instance_images\n",
        "    !rm -r Regularization_images\n",
        "    !unzip instance_images.zip\n",
        "    !mv *.ckpt $Session_Name\".ckpt\"\n",
        "    %cd /content\n",
        "\n",
        "\n",
        "INSTANCE_NAME=Session_Name\n",
        "OUTPUT_DIR=\"/content/models/\"+Session_Name\n",
        "SESSION_DIR=WORKSPACE+'/Sessions/'+Session_Name\n",
        "INSTANCE_DIR=SESSION_DIR+'/instance_images'\n",
        "MDLPTH=str(SESSION_DIR+\"/\"+Session_Name+'.ckpt')\n",
        "CLASS_DIR=SESSION_DIR+'/Regularization_images'\n",
        "\n",
        "Contains_faces = \"No\" #@param [\"No\", \"Female\", \"Male\", \"Both\"]\n",
        "\n",
        "#@markdown - If you're training on a subject with a face or a movie/style that contains faces. (experimental, still needs some tuning) \n",
        "\n",
        "def reg():\n",
        "  with capture.capture_output() as cap:\n",
        "    if Contains_faces!=\"No\":\n",
        "      if not os.path.exists(str(CLASS_DIR)):\n",
        "        %mkdir -p \"$CLASS_DIR\"\n",
        "      %cd $CLASS_DIR\n",
        "      !rm -r Women Men Mix\n",
        "      !wget -O Womenz 'https://github.com/TheLastBen/fast-stable-diffusion/raw/main/Dreambooth/Regularization/Women'\n",
        "      !wget -O Menz 'https://github.com/TheLastBen/fast-stable-diffusion/raw/main/Dreambooth/Regularization/Men'\n",
        "      !wget -O Mixz 'https://github.com/TheLastBen/fast-stable-diffusion/raw/main/Dreambooth/Regularization/Mix'\n",
        "      !unzip Menz\n",
        "      !unzip Womenz\n",
        "      !unzip Mixz\n",
        "      !rm Menz Womenz Mixz\n",
        "      !find . -name \"* *\" -type f | rename 's/ /_/g'\n",
        "      %cd /content               \n",
        "\n",
        "\n",
        "if os.path.exists(str(SESSION_DIR)):\n",
        "  if not os.path.exists(MDLPTH) and '.ckpt' in str([ckpt for ckpt in listdir(SESSION_DIR) if ckpt.split(\".\")[-1]==\"ckpt\"]):  \n",
        "    \n",
        "    def f(n):  \n",
        "      k=0\n",
        "      for i in [ckpt for ckpt in listdir(SESSION_DIR) if ckpt.split(\".\")[-1]==\"ckpt\"]:    \n",
        "        if k==n:    \n",
        "          !mv $SESSION_DIR/$i $MDLPTH\n",
        "        k=k+1\n",
        "\n",
        "    k=0\n",
        "    print('\u001b[1;33mNo final checkpoint model found, select which intermediary checkpoint to use (000 to skip):\\n\u001b[1;34m')\n",
        "\n",
        "    for i in [ckpt for ckpt in listdir(SESSION_DIR) if ckpt.split(\".\")[-1]==\"ckpt\"]:    \n",
        "      print(str(k)+'- '+i)\n",
        "      k=k+1\n",
        "    n=input()\n",
        "    while int(n)>k-1:\n",
        "      n=input()  \n",
        "    if n!=\"000\":\n",
        "      f(int(n))\n",
        "      print('\u001b[1;32mUsing the model '+ i+\" ...\")\n",
        "      time.sleep(2)\n",
        "    else:\n",
        "      print('\u001b[1;32mSkipping the intermediary checkpoints.')\n",
        "    del n\n",
        "\n",
        "\n",
        "if os.path.exists(str(SESSION_DIR)) and not os.path.exists(MDLPTH):\n",
        "  print('\u001b[1;32mLoading session with no previous model, using the original model or the custom downloaded model')\n",
        "  reg()\n",
        "  if MODEL_NAME==\"\":\n",
        "    print('\u001b[1;31mNo model found, use the \"Model Download\" cell to download a model.')\n",
        "  else:\n",
        "    print('\u001b[1;32mSession Loaded, proceed to uploading instance images')\n",
        "\n",
        "elif os.path.exists(MDLPTH):\n",
        "  print('\u001b[1;32mSession found, loading the trained model ...')\n",
        "  reg()\n",
        "  %mkdir -p \"$OUTPUT_DIR\"\n",
        "  !python /content/diffusers/scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path \"$MDLPTH\" --dump_path \"$OUTPUT_DIR\" --session_dir \"$SESSION_DIR\"\n",
        "  if os.path.exists(OUTPUT_DIR+'/unet/diffusion_pytorch_model.bin'):\n",
        "    resume=True    \n",
        "    !rm /content/v1-inference.yaml\n",
        "    clear_output()\n",
        "    print('\u001b[1;32mSession loaded.')\n",
        "  else:     \n",
        "    !rm /content/v1-inference.yaml\n",
        "    if not os.path.exists(OUTPUT_DIR+'/unet/diffusion_pytorch_model.bin'):\n",
        "      print('\u001b[1;31mConversion error, if the error persists, remove the CKPT file from the current session folder')\n",
        "\n",
        "\n",
        "elif not os.path.exists(str(SESSION_DIR)):\n",
        "    %mkdir -p \"$INSTANCE_DIR\"\n",
        "    print('\u001b[1;32mCreating session...')\n",
        "    reg()\n",
        "    if MODEL_NAME==\"\":\n",
        "      print('\u001b[1;31mNo model found, use the \"Model Download\" cell to download a model.')\n",
        "    else:\n",
        "      print('\u001b[1;32mSession created, proceed to uploading instance images')\n",
        "\n",
        "    \n",
        "if Contains_faces == \"Female\":\n",
        "  CLASS_DIR=CLASS_DIR+'/Women'\n",
        "if Contains_faces == \"Male\":\n",
        "  CLASS_DIR=CLASS_DIR+'/Men'\n",
        "if Contains_faces == \"Both\":\n",
        "  CLASS_DIR=CLASS_DIR+'/Mix'\n",
        "\n",
        "try:\n",
        "  Contain_f\n",
        "  del Contain_f\n",
        "except:\n",
        "  pass\n",
        "\n",
        "    #@markdown \n",
        "\n",
        "    #@markdown # The most importent step is to rename the instance pictures of each subject to a unique unknown identifier, example :\n",
        "    #@markdown - If you have 30 pictures of yourself, simply select them all and rename only one to the chosen identifier for example : phtmejhn, the files would be : phtmejhn (1).jpg, phtmejhn (2).png ....etc then upload them, do the same for other people or objects with a different identifier, and that's it.\n",
        "    #@markdown - Check out this example : https://i.imgur.com/d2lD3rz.jpeg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "LC4ukG60fgMy"
      },
      "outputs": [],
      "source": [
        "import shutil\n",
        "from google.colab import files\n",
        "from PIL import Image\n",
        "from tqdm import tqdm\n",
        "\n",
        "#@markdown #Instance Images\n",
        "#@markdown ----\n",
        "\n",
        "#@markdown\n",
        "#@markdown - Run the cell to Upload the instance pictures.\n",
        "\n",
        "Remove_existing_instance_images= True #@param{type: 'boolean'}\n",
        "#@markdown - Uncheck the box to keep the existing instance images.\n",
        "\n",
        "\n",
        "if Remove_existing_instance_images:\n",
        "  if os.path.exists(str(INSTANCE_DIR)):\n",
        "    !rm -r \"$INSTANCE_DIR\"\n",
        "\n",
        "if not os.path.exists(str(INSTANCE_DIR)):\n",
        "  %mkdir -p \"$INSTANCE_DIR\"\n",
        "\n",
        "IMAGES_FOLDER_OPTIONAL=\"\" #@param{type: 'string'}\n",
        "\n",
        "#@markdown - If you prefer to specify directly the folder of the pictures instead of uploading, this will add the pictures to the existing (if any) instance images. Leave EMPTY to upload.\n",
        "\n",
        "Crop_images= True #@param{type: 'boolean'}\n",
        "Crop_size = \"512\" #@param [\"512\", \"576\", \"640\", \"704\", \"768\", \"832\", \"896\", \"960\", \"1024\"]\n",
        "Crop_size=int(Crop_size)\n",
        "\n",
        "#@markdown - Unless you want to crop them manually in a precise way, you don't need to crop your instance images externally.\n",
        "\n",
        "while IMAGES_FOLDER_OPTIONAL !=\"\" and not os.path.exists(str(IMAGES_FOLDER_OPTIONAL)):\n",
        "  print('\u001b[1;31mThe image folder specified does not exist, use the colab file explorer to copy the path :')\n",
        "  IMAGES_FOLDER_OPTIONAL=input('')\n",
        "\n",
        "if IMAGES_FOLDER_OPTIONAL!=\"\":\n",
        "  if Crop_images:\n",
        "    for filename in tqdm(os.listdir(IMAGES_FOLDER_OPTIONAL), bar_format='  |{bar:15}| {n_fmt}/{total_fmt} Uploaded'):\n",
        "      extension = filename.split(\".\")[1]\n",
        "      identifier=filename.split(\".\")[0]\n",
        "      new_path_with_file = os.path.join(INSTANCE_DIR, filename)\n",
        "      file = Image.open(IMAGES_FOLDER_OPTIONAL+\"/\"+filename)\n",
        "      width, height = file.size\n",
        "      if file.size !=(Crop_size, Crop_size):      \n",
        "        side_length = min(width, height)\n",
        "        left = (width - side_length)/2\n",
        "        top = (height - side_length)/2\n",
        "        right = (width + side_length)/2\n",
        "        bottom = (height + side_length)/2\n",
        "        image = file.crop((left, top, right, bottom))\n",
        "        image = image.resize((Crop_size, Crop_size))\n",
        "        if (extension.upper() == \"JPG\"):\n",
        "            image.save(new_path_with_file, format=\"JPEG\", quality = 100)\n",
        "        else:\n",
        "            image.save(new_path_with_file, format=extension.upper())\n",
        "      else:\n",
        "        !cp \"$IMAGES_FOLDER_OPTIONAL/$filename\" \"$INSTANCE_DIR\"\n",
        "\n",
        "  else:\n",
        "    for filename in tqdm(os.listdir(IMAGES_FOLDER_OPTIONAL), bar_format='  |{bar:15}| {n_fmt}/{total_fmt} Uploaded'):\n",
        "      %cp -r \"$IMAGES_FOLDER_OPTIONAL/$filename\" \"$INSTANCE_DIR\"\n",
        " \n",
        "  print('\\n\u001b[1;32mDone, proceed to the training cell')\n",
        "\n",
        "\n",
        "elif IMAGES_FOLDER_OPTIONAL ==\"\":\n",
        "  uploaded = files.upload()\n",
        "  if Crop_images:\n",
        "    for filename in tqdm(uploaded.keys(), bar_format='  |{bar:15}| {n_fmt}/{total_fmt} Uploaded'):\n",
        "      shutil.move(filename, INSTANCE_DIR)\n",
        "      extension = filename.split(\".\")[1]\n",
        "      identifier=filename.split(\".\")[0]\n",
        "      new_path_with_file = os.path.join(INSTANCE_DIR, filename)\n",
        "      file = Image.open(new_path_with_file)\n",
        "      width, height = file.size\n",
        "      if file.size !=(Crop_size, Crop_size):        \n",
        "        side_length = min(width, height)\n",
        "        left = (width - side_length)/2\n",
        "        top = (height - side_length)/2\n",
        "        right = (width + side_length)/2\n",
        "        bottom = (height + side_length)/2\n",
        "        image = file.crop((left, top, right, bottom))\n",
        "        image = image.resize((Crop_size, Crop_size))\n",
        "        if (extension.upper() == \"JPG\"):\n",
        "            image.save(new_path_with_file, format=\"JPEG\", quality = 100)\n",
        "        else:\n",
        "            image.save(new_path_with_file, format=extension.upper())\n",
        "      clear_output()\n",
        "  else:\n",
        "    for filename in tqdm(uploaded.keys(), bar_format='  |{bar:15}| {n_fmt}/{total_fmt} Uploaded'):\n",
        "      shutil.move(filename, INSTANCE_DIR)\n",
        "      clear_output()\n",
        "\n",
        "  print('\\n\u001b[1;32mDone, proceed to the training cell')\n",
        "\n",
        "with capture.capture_output() as cap:\n",
        "  %cd \"$INSTANCE_DIR\"\n",
        "  !find . -name \"* *\" -type f | rename 's/ /-/g'\n",
        "  %cd /content\n",
        "  if os.path.exists(INSTANCE_DIR+\"/.ipynb_checkpoints\"):\n",
        "    %rm -r INSTANCE_DIR+\"/.ipynb_checkpoints\"    \n",
        "\n",
        "  %cd $SESSION_DIR\n",
        "  !rm instance_images.zip\n",
        "  !zip -r instance_images instance_images\n",
        "  %cd /content"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZnmQYfZilzY6"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "1-9QbkfAVYYU"
      },
      "outputs": [],
      "source": [
        "#@markdown ---\n",
        "#@markdown #Start DreamBooth\n",
        "#@markdown ---\n",
        "import os\n",
        "from subprocess import getoutput\n",
        "from IPython.display import HTML\n",
        "from IPython.display import clear_output\n",
        "from google.colab import runtime\n",
        "import time\n",
        "import random\n",
        "\n",
        "Resume_Training = False #@param {type:\"boolean\"}\n",
        "\n",
        "while not Resume_Training and MODEL_NAME==\"\":\n",
        "  print('\u001b[1;31mNo model found, use the \"Model Download\" cell to download a model.')\n",
        "  time.sleep(5)\n",
        "\n",
        "\n",
        "#@markdown  - If you're not satisfied with the result, check this box, run again the cell and it will continue training the current model.\n",
        "\n",
        "MODELT_NAME=MODEL_NAME\n",
        "\n",
        "Training_Steps=3000 #@param{type: 'number'}\n",
        "#@markdown - Total Steps = Number of Instance images * 200, if you use 30 images, use 6000 steps, if you're not satisfied with the result, resume training for another 500 steps, and so on ...\n",
        "\n",
        "Seed='' #@param{type: 'string'}\n",
        "\n",
        "#@markdown - Leave empty for a random seed.\n",
        "\n",
        "Resolution = \"512\" #@param [\"512\", \"576\", \"640\", \"704\", \"768\", \"832\", \"896\", \"960\", \"1024\"]\n",
        "Res=int(Resolution)\n",
        "\n",
        "#@markdown - Higher resolution = Higher quality, make sure the instance images are cropped to this selected size (or larger).\n",
        "\n",
        "fp16 = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown - Enable/disable half-precision, disabling it will double the training time and produce 4.7Gb checkpoints.\n",
        "\n",
        "#GC= \"\"\n",
        "#if Resolution!=\"512\":\n",
        "GC= \"--gradient_checkpointing\"\n",
        "\n",
        "if Seed =='' or Seed=='0':\n",
        "  Seed=random.randint(1, 999999)\n",
        "else:\n",
        "  Seed=int(Seed)\n",
        "\n",
        "if fp16:\n",
        "  prec=\"fp16\"\n",
        "else:\n",
        "  prec=\"no\"\n",
        "\n",
        "s = getoutput('nvidia-smi')\n",
        "if 'A100' in s:\n",
        "  precision=\"no\"\n",
        "  GC= \"\"\n",
        "else:\n",
        "  precision=prec\n",
        "\n",
        "try:\n",
        "   resume\n",
        "   if resume and not Resume_Training:\n",
        "     print('\u001b[1;31mOverwrite your previously trained model ?, answering \"yes\" will train a new model, answering \"no\" will resume the training of the previous model?  yes or no ?\u001b[0m')\n",
        "     while True:\n",
        "        ansres=input('')\n",
        "        if ansres=='no':\n",
        "          Resume_Training = True\n",
        "          del ansres\n",
        "          break\n",
        "        elif ansres=='yes':\n",
        "          Resume_Training = False\n",
        "          resume= False\n",
        "          break\n",
        "except:\n",
        "  pass\n",
        "\n",
        "if Resume_Training and os.path.exists(OUTPUT_DIR+'/unet/diffusion_pytorch_model.bin'):\n",
        "  MODELT_NAME=OUTPUT_DIR\n",
        "  print('\u001b[1;32mResuming Training...\u001b[0m')\n",
        "elif Resume_Training and not os.path.exists(OUTPUT_DIR+'/unet/diffusion_pytorch_model.bin'):\n",
        "  print('\u001b[1;31mPrevious model not found, training a new model...\u001b[0m')\n",
        "  MODELT_NAME=MODEL_NAME\n",
        "  while MODEL_NAME==\"\":\n",
        "    print('\u001b[1;31mNo model found, use the \"Model Download\" cell to download a model.')\n",
        "    time.sleep(5)\n",
        "\n",
        "#@markdown ---------------------------\n",
        "\n",
        "try:\n",
        "   Contain_f\n",
        "   pass\n",
        "except:\n",
        "   Contain_f=Contains_faces\n",
        "\n",
        "Enable_text_encoder_training= True #@param{type: 'boolean'}\n",
        "\n",
        "#@markdown - At least 10% of the total training steps are needed, it doesn't matter if they are at the beginning or in the middle or the end, in case you're training the model multiple times.\n",
        "#@markdown - For example you can devide 5%, 5%, 5% on 3 training runs on the model, or 0%, 0%, 15%, given that 15% will cover the total training steps count (15% of 200 steps is not enough).\n",
        "\n",
        "#@markdown - Enter the % of the total steps for which to train the text_encoder\n",
        "Train_text_encoder_for=100 #@param{type: 'number'}\n",
        "\n",
        "#@markdown - If you're training a style, keep it between 10-20%, if you're training on a person, set it between 50-70%, reduce it if you can't stylize the person/object.\n",
        "#@markdown - Higher % will give more weight to the instance, it gives stronger results at lower steps count, but harder to stylize.\n",
        "\n",
        "if Train_text_encoder_for>=100:\n",
        "  stptxt=Training_Steps\n",
        "elif Train_text_encoder_for==0:\n",
        "  Enable_text_encoder_training= False\n",
        "  stptxt=10\n",
        "else:\n",
        "  stptxt=int((Training_Steps*Train_text_encoder_for)/100)\n",
        "\n",
        "if not Enable_text_encoder_training:\n",
        "  Contains_faces=\"No\"\n",
        "else:\n",
        "   Contains_faces=Contain_f\n",
        "\n",
        "if Enable_text_encoder_training:\n",
        "  Textenc=\"--train_text_encoder\"\n",
        "else:\n",
        "  Textenc=\"\"\n",
        "\n",
        "#@markdown ---------------------------\n",
        "Save_Checkpoint_Every_n_Steps = False #@param {type:\"boolean\"}\n",
        "Save_Checkpoint_Every=500 #@param{type: 'number'}\n",
        "if Save_Checkpoint_Every==None:\n",
        "  Save_Checkpoint_Every=1\n",
        "#@markdown - Minimum 200 steps between each save.\n",
        "stp=0\n",
        "Start_saving_from_the_step=500 #@param{type: 'number'}\n",
        "if Start_saving_from_the_step==None:\n",
        "  Start_saving_from_the_step=0\n",
        "if (Start_saving_from_the_step < 200):\n",
        "  Start_saving_from_the_step=Save_Checkpoint_Every\n",
        "stpsv=Start_saving_from_the_step\n",
        "if Save_Checkpoint_Every_n_Steps:\n",
        "  stp=Save_Checkpoint_Every\n",
        "#@markdown - Start saving intermediary checkpoints from this step.\n",
        "\n",
        "Disconnect_after_training=False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown - Auto-disconnect from google colab after the training to avoid wasting compute units.\n",
        "\n",
        "Caption=''\n",
        "if Captionned_instance_images:\n",
        "  Caption='--image_captions_filename'\n",
        "\n",
        "\n",
        "def txtenc_train(Caption, stpsv, stp, MODELT_NAME, INSTANCE_DIR, CLASS_DIR, OUTPUT_DIR, PT, Seed, precision, Training_Steps):\n",
        "  print('\u001b[1;33mTraining the text encoder with regularization...\u001b[0m')\n",
        "  !accelerate launch /content/diffusers/examples/dreambooth/train_dreambooth.py \\\n",
        "    $Caption \\\n",
        "    --train_text_encoder \\\n",
        "    --dump_only_text_encoder \\\n",
        "    --pretrained_model_name_or_path=\"$MODELT_NAME\" \\\n",
        "    --instance_data_dir=\"$INSTANCE_DIR\" \\\n",
        "    --class_data_dir=\"$CLASS_DIR\" \\\n",
        "    --output_dir=\"$OUTPUT_DIR\" \\\n",
        "    --with_prior_preservation --prior_loss_weight=1.0 \\\n",
        "    --instance_prompt=\"$PT\"\\\n",
        "    --seed=$Seed \\\n",
        "    --resolution=$Res \\\n",
        "    --mixed_precision=$precision \\\n",
        "    --train_batch_size=1 \\\n",
        "    --gradient_accumulation_steps=1 --gradient_checkpointing \\\n",
        "    --use_8bit_adam \\\n",
        "    --learning_rate=2e-6 \\\n",
        "    --lr_scheduler=\"polynomial\" \\\n",
        "    --lr_warmup_steps=0 \\\n",
        "    --max_train_steps=$Training_Steps \\\n",
        "    --num_class_images=200\n",
        "\n",
        "def unet_train(Caption, SESSION_DIR, stpsv, stp, MODELT_NAME, INSTANCE_DIR, OUTPUT_DIR, PT, Seed, precision, Training_Steps):\n",
        "  clear_output()\n",
        "  print('\u001b[1;33mTraining the unet...\u001b[0m')\n",
        "  !accelerate launch /content/diffusers/examples/dreambooth/train_dreambooth.py \\\n",
        "    $Caption \\\n",
        "    --train_only_unet \\\n",
        "    --Session_dir=$SESSION_DIR \\\n",
        "    --save_starting_step=$stpsv \\\n",
        "    --save_n_steps=$stp \\\n",
        "    --pretrained_model_name_or_path=\"$MODELT_NAME\" \\\n",
        "    --instance_data_dir=\"$INSTANCE_DIR\" \\\n",
        "    --output_dir=\"$OUTPUT_DIR\" \\\n",
        "    --instance_prompt=\"$PT\" \\\n",
        "    --seed=$Seed \\\n",
        "    --resolution=$Res \\\n",
        "    --mixed_precision=$precision \\\n",
        "    --train_batch_size=1 \\\n",
        "    --gradient_accumulation_steps=1 $GC \\\n",
        "    --use_8bit_adam \\\n",
        "    --learning_rate=2e-6 \\\n",
        "    --lr_scheduler=\"polynomial\" \\\n",
        "    --lr_warmup_steps=0 \\\n",
        "    --max_train_steps=$Training_Steps\n",
        "\n",
        "if Contains_faces!=\"No\":\n",
        "  \n",
        "  txtenc_train(Caption, stpsv, stp, MODELT_NAME, INSTANCE_DIR, CLASS_DIR, OUTPUT_DIR, PT, Seed, precision, Training_Steps=stptxt)\n",
        "  unet_train(Caption, SESSION_DIR, stpsv, stp, MODELT_NAME, INSTANCE_DIR, OUTPUT_DIR, PT, Seed, precision, Training_Steps)\n",
        "\n",
        "else:\n",
        "  !accelerate launch /content/diffusers/examples/dreambooth/train_dreambooth.py \\\n",
        "    $Caption \\\n",
        "    $Textenc \\\n",
        "    --save_starting_step=$stpsv \\\n",
        "    --stop_text_encoder_training=$stptxt \\\n",
        "    --save_n_steps=$stp \\\n",
        "    --Session_dir=$SESSION_DIR \\\n",
        "    --pretrained_model_name_or_path=\"$MODELT_NAME\" \\\n",
        "    --instance_data_dir=\"$INSTANCE_DIR\" \\\n",
        "    --output_dir=\"$OUTPUT_DIR\" \\\n",
        "    --instance_prompt=\"$PT\" \\\n",
        "    --seed=$Seed \\\n",
        "    --resolution=$Res \\\n",
        "    --mixed_precision=$precision \\\n",
        "    --train_batch_size=1 \\\n",
        "    --gradient_accumulation_steps=1 $GC \\\n",
        "    --use_8bit_adam \\\n",
        "    --learning_rate=2e-6 \\\n",
        "    --lr_scheduler=\"polynomial\" \\\n",
        "    --lr_warmup_steps=0 \\\n",
        "    --max_train_steps=$Training_Steps\n",
        "\n",
        "\n",
        "if os.path.exists('/content/models/'+INSTANCE_NAME+'/unet/diffusion_pytorch_model.bin'):\n",
        "  print(\"Almost done ...\")\n",
        "  %cd /content    \n",
        "  !wget -O convertosd.py https://github.com/TheLastBen/fast-stable-diffusion/raw/main/Dreambooth/convertosd.py\n",
        "  clear_output()\n",
        "  if precision==\"no\":\n",
        "    !sed -i '226s@.*@@' /content/convertosd.py\n",
        "  !sed -i '201s@.*@    model_path = \"{OUTPUT_DIR}\"@' /content/convertosd.py\n",
        "  !sed -i '202s@.*@    checkpoint_path= \"{SESSION_DIR}/{Session_Name}.ckpt\"@' /content/convertosd.py\n",
        "  !python /content/convertosd.py\n",
        "  clear_output()\n",
        "  if os.path.exists(SESSION_DIR+\"/\"+INSTANCE_NAME+'.ckpt'):\n",
        "    if not os.path.exists(str(SESSION_DIR+'/tokenizer')):\n",
        "      !cp -R '/content/models/'$INSTANCE_NAME'/tokenizer' \"$SESSION_DIR\"\n",
        "    print(\"\u001b[1;32mDONE, the CKPT model is in your Gdrive in the sessions folder\")\n",
        "    if Disconnect_after_training :\n",
        "      runtime.unassign()    \n",
        "  else:\n",
        "    print(\"\u001b[1;31mSomething went wrong\")\n",
        "    \n",
        "else:\n",
        "  print(\"\u001b[1;31mSomething went wrong\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ehi1KKs-l-ZS"
      },
      "source": [
        "# Test The Trained Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iAZGngFcI8hq"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import time\n",
        "import sys\n",
        "import fileinput\n",
        "from IPython.display import clear_output\n",
        "from subprocess import getoutput\n",
        "from IPython.utils import capture\n",
        "\n",
        "\n",
        "Model_Version = \"1.5\" #@param [\"1.5\", \"V2-512\", \"V2-768\"]\n",
        "#@markdown  - Important! Choose the correct version and resolution of the model\n",
        "\n",
        "Update_repo = True #@param {type:\"boolean\"}\n",
        "\n",
        "Session__Name=\"\" #@param{type: 'string'}\n",
        "\n",
        "#@markdown - Leave empty if you want to use the current trained model.\n",
        "\n",
        "Use_Custom_Path = False #@param {type:\"boolean\"}\n",
        "\n",
        "try:\n",
        "  INSTANCE_NAME\n",
        "  INSTANCET=INSTANCE_NAME  \n",
        "except:\n",
        "  pass\n",
        "#@markdown - if checked, an input box will ask the full path to a desired model.\n",
        "\n",
        "if Session__Name!=\"\":\n",
        "  INSTANCET=Session__Name\n",
        "  INSTANCET=INSTANCET.replace(\" \",\"_\")\n",
        "\n",
        "if Use_Custom_Path:\n",
        "  try:\n",
        "    INSTANCET\n",
        "    del INSTANCET\n",
        "  except:\n",
        "    pass\n",
        "\n",
        "try:\n",
        "  INSTANCET\n",
        "  if Session__Name!=\"\":\n",
        "    path_to_trained_model='/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/'+Session__Name+\"/\"+Session__Name+'.ckpt'\n",
        "  else:\n",
        "    path_to_trained_model=SESSION_DIR+\"/\"+INSTANCET+'.ckpt'\n",
        "except:\n",
        "  print('\u001b[1;31mIt seems that you did not perform training during this session \u001b[1;32mor you chose to use a custom path,\\nprovide the full path to the model (including the name of the model):\\n')\n",
        "  path_to_trained_model=input()\n",
        "     \n",
        "while not os.path.exists(path_to_trained_model):\n",
        "   print(\"\u001b[1;31mThe model doesn't exist on you Gdrive, use the file explorer to get the path : \")\n",
        "   path_to_trained_model=input()\n",
        "\n",
        "         \n",
        "with capture.capture_output() as cap:\n",
        "    %cd /content/gdrive/MyDrive/\n",
        "    %mkdir sd\n",
        "    %cd sd\n",
        "    !git clone https://github.com/Stability-AI/stablediffusion\n",
        "    !git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui\n",
        "    %cd /content/gdrive/MyDrive/sd/stable-diffusion-webui/\n",
        "    !mkdir -p cache/{huggingface,torch}\n",
        "    %cd /content/\n",
        "    !ln -s /content/gdrive/MyDrive/sd/stable-diffusion-webui/cache/huggingface ../root/.cache/\n",
        "    !ln -s /content/gdrive/MyDrive/sd/stable-diffusion-webui/cache/torch ../root/.cache/\n",
        "\n",
        "if Update_repo:\n",
        "  !rm /content/gdrive/MyDrive/sd/stable-diffusion-webui/webui.sh  \n",
        "  !rm /content/gdrive/MyDrive/sd/stable-diffusion-webui/modules/paths.py\n",
        "  !rm /content/gdrive/MyDrive/sd/stable-diffusion-webui/webui.py \n",
        "  !rm /content/gdrive/MyDrive/sd/stable-diffusion-webui/modules/ui.py\n",
        "  !rm /content/gdrive/MyDrive/sd/stable-diffusion-webui/style.css\n",
        "  %cd /content/gdrive/MyDrive/sd/stable-diffusion-webui/\n",
        "  clear_output()\n",
        "  print('\u001b[1;32m')\n",
        "  !git pull\n",
        "\n",
        "\n",
        "with capture.capture_output() as cap:\n",
        "  \n",
        "  if not os.path.exists('/content/gdrive/MyDrive/sd/stablediffusion/src/k-diffusion/k_diffusion'):\n",
        "    !mkdir /content/gdrive/MyDrive/sd/stablediffusion/src\n",
        "    %cd /content/gdrive/MyDrive/sd/stablediffusion/src\n",
        "    !git clone https://github.com/CompVis/taming-transformers\n",
        "    !git clone https://github.com/openai/CLIP\n",
        "    !git clone https://github.com/salesforce/BLIP\n",
        "    !git clone https://github.com/sczhou/CodeFormer\n",
        "    !git clone https://github.com/crowsonkb/k-diffusion\n",
        "    !mv /content/gdrive/MyDrive/sd/stablediffusion/src/CLIP /content/gdrive/MyDrive/sd/stablediffusion/src/clip\n",
        "    !mv  /content/gdrive/MyDrive/sd/stablediffusion/src/BLIP /content/gdrive/MyDrive/sd/stablediffusion/src/blip    \n",
        "    !mv  /content/gdrive/MyDrive/sd/stablediffusion/src/CodeFormer /content/gdrive/MyDrive/sd/stablediffusion/src/codeformer        \n",
        "    !cp -r /content/gdrive/MyDrive/sd/stablediffusion/src/k-diffusion/k_diffusion /content/gdrive/MyDrive/sd/stable-diffusion-webui/    \n",
        "\n",
        "\n",
        "with capture.capture_output() as cap:\n",
        "  if not os.path.exists('/usr/local/lib/python3.7/dist-packages/facexlib-0.2.5.dist-info'):\n",
        "    %cd /content/\n",
        "    !pip install -q gradio==3.9\n",
        "    !pip install -q safetensors\n",
        "    !pip install -q torchsde    \n",
        "    !wget https://github.com/TheLastBen/fast-stable-diffusion/raw/main/Dependencies/Dependencies_AUT.1\n",
        "    !wget https://github.com/TheLastBen/fast-stable-diffusion/raw/main/Dependencies/Dependencies_AUT.2\n",
        "    %mv Dependencies_AUT.1 Dependencies_AUT.7z.001\n",
        "    %mv Dependencies_AUT.2 Dependencies_AUT.7z.002\n",
        "    !7z x Dependencies_AUT.7z.001\n",
        "    time.sleep(2)\n",
        "    !rm -r /content/usr/local/lib/python3.7/dist-packages/transformers\n",
        "    !rm -r /content/usr/local/lib/python3.7/dist-packages/transformers-4.19.2.dist-info\n",
        "    !rm -r /content/usr/local/lib/python3.7/dist-packages/diffusers\n",
        "    !rm -r /content/usr/local/lib/python3.7/dist-packages/diffusers-0.3.0.dist-info\n",
        "    !rm -r /content/usr/local/lib/python3.7/dist-packages/accelerate\n",
        "    !rm -r /content/usr/local/lib/python3.7/dist-packages/accelerate-0.12.0.dist-info    \n",
        "    !cp -r /content/usr/local/lib/python3.7/dist-packages /usr/local/lib/python3.7/\n",
        "    !rm -r /content/usr\n",
        "    !rm Dependencies_AUT.7z.001\n",
        "    !rm Dependencies_AUT.7z.002\n",
        "\n",
        "with capture.capture_output() as cap:    \n",
        "  %cd /content/gdrive/MyDrive/sd/stable-diffusion-webui/modules\n",
        "  !wget -O paths.py https://raw.githubusercontent.com/TheLastBen/fast-stable-diffusion/main/AUTOMATIC1111_files/paths.py\n",
        "\n",
        "with capture.capture_output() as cap:\n",
        "  if not os.path.exists('/tools/node/bin/lt'):\n",
        "    !npm install -g localtunnel\n",
        "\n",
        "with capture.capture_output() as cap: \n",
        "  %cd /content/gdrive/MyDrive/sd/stable-diffusion-webui/\n",
        "  time.sleep(1)\n",
        "  !wget -O webui.py https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui/master/webui.py\n",
        "  !sed -i 's@gpu_call).*@gpu_call) \\n        shared.demo.queue(concurrency_count=111500)@' /content/gdrive/MyDrive/sd/stable-diffusion-webui/webui.py\n",
        "  %cd /content/gdrive/MyDrive/sd/stable-diffusion-webui/modules/\n",
        "  !wget -O ui.py https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui/master/modules/ui.py\n",
        "  !sed -i 's@css = \"\".*@with open(os.path.join(script_path, \"style.css\"), \"r\", encoding=\"utf8\") as file:\\n        css = file.read()@' /content/gdrive/MyDrive/sd/stable-diffusion-webui/modules/ui.py  \n",
        "  %cd /content/gdrive/MyDrive/sd/stable-diffusion-webui\n",
        "  !wget -O style.css https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui/master/style.css\n",
        "  !sed -i 's@min-height: 4.*@min-height: 5.5em;@g' /content/gdrive/MyDrive/sd/stable-diffusion-webui/style.css  \n",
        "  %cd /content\n",
        "\n",
        "\n",
        "Use_Gradio_Server = False #@param {type:\"boolean\"}\n",
        "#@markdown  - Only if you have trouble connecting to the local server.\n",
        "\n",
        "\n",
        "share=''\n",
        "if Use_Gradio_Server:\n",
        "  share='--share'\n",
        "  for line in fileinput.input('/usr/local/lib/python3.7/dist-packages/gradio/blocks.py', inplace=True):\n",
        "    if line.strip().startswith('self.server_name ='):\n",
        "        line = '            self.server_name = server_name\\n'\n",
        "    if line.strip().startswith('self.server_port ='):\n",
        "        line = '            self.server_port = server_port\\n'\n",
        "    sys.stdout.write(line)\n",
        "  clear_output()\n",
        "  \n",
        "else:\n",
        "  share=''\n",
        "  !nohup lt --port 7860 > srv.txt 2>&1 &\n",
        "  time.sleep(2)\n",
        "  !grep -o 'https[^ ]*' /content/srv.txt >srvr.txt\n",
        "  time.sleep(2)\n",
        "  srv= getoutput('cat /content/srvr.txt')\n",
        "\n",
        "  for line in fileinput.input('/usr/local/lib/python3.7/dist-packages/gradio/blocks.py', inplace=True):\n",
        "    if line.strip().startswith('self.server_name ='):\n",
        "        line = f'            self.server_name = \"{srv[8:]}\"\\n'\n",
        "    if line.strip().startswith('self.server_port ='):\n",
        "        line = '            self.server_port = 443\\n'\n",
        "    if line.strip().startswith('self.protocol = \"https\"'):\n",
        "        line = '            self.protocol = \"https\"\\n'\n",
        "    if line.strip().startswith('if self.local_url.startswith(\"https\") or self.is_colab'):\n",
        "        line = ''    \n",
        "    if line.strip().startswith('else \"http\"'):\n",
        "        line = ''              \n",
        "    sys.stdout.write(line)\n",
        "    \n",
        "\n",
        "  !sed -i '13s@.*@    \"PUBLIC_SHARE_TRUE\": \"\u001b[32mConnected\",@' /usr/local/lib/python3.7/dist-packages/gradio/strings.py\n",
        "  \n",
        "  !rm /content/srv.txt\n",
        "  !rm /content/srvr.txt\n",
        "  clear_output()\n",
        "\n",
        "with capture.capture_output() as cap:\n",
        "  %cd /content/gdrive/MyDrive/sd/stablediffusion/\n",
        "\n",
        "if Model_Version == \"V2-768\":\n",
        "  configf=\"--config /content/gdrive/MyDrive/sd/stablediffusion/configs/stable-diffusion/v2-inference-v.yaml\"\n",
        "  NM==\"True\":\n",
        "elif Model_Version == \"V2-512\":\n",
        "  configf=\"--config /content/gdrive/MyDrive/sd/stablediffusion/configs/stable-diffusion/v2-inference.yaml\"\n",
        "  NM==\"True\":\n",
        "else:\n",
        "  configf=\"\"\n",
        "  NM==\"False\":\n",
        "\n",
        "if os.path.isfile(path_to_trained_model):\n",
        "  !python /content/gdrive/MyDrive/sd/stable-diffusion-webui/webui.py $share --disable-safe-unpickle --no-half-vae  --ckpt \"$path_to_trained_model\" $configf\n",
        "else:\n",
        "  !python /content/gdrive/MyDrive/sd/stable-diffusion-webui/webui.py $share --disable-safe-unpickle --no-half-vae  --ckpt-dir \"$path_to_trained_model\" $configf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d_mQ23XsOc5R"
      },
      "source": [
        "# Upload The Trained Model to Hugging Face "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NTqUIuhROdH4"
      },
      "outputs": [],
      "source": [
        "from slugify import slugify\n",
        "from huggingface_hub import HfApi, HfFolder, CommitOperationAdd\n",
        "from huggingface_hub import create_repo\n",
        "from IPython.display import display_markdown\n",
        "from IPython.display import clear_output\n",
        "from IPython.utils import capture\n",
        "from google.colab import files\n",
        "import shutil\n",
        "import time\n",
        "import os\n",
        "\n",
        "\n",
        "#@markdown ##Save it to your personal profile or collaborate to the public [library of concepts](https://huggingface.co/sd-dreambooth-library)\n",
        "#@markdown Leave `name_of_your_concept` blank if you want to name your concept the same as your Session Name\n",
        "name_of_your_concept = \"\" #@param {type:\"string\"}\n",
        "if(name_of_your_concept == \"\"):\n",
        "  name_of_your_concept = Session_Name\n",
        "name_of_your_concept=name_of_your_concept.replace(\" \",\"-\")  \n",
        "  \n",
        "Save_concept_to = \"My_Profile\" #@param [\"Public_Library\", \"My_Profile\"]\n",
        "\n",
        "#@markdown - [Create a write access token](https://huggingface.co/settings/tokens) , go to \"New token\" -> Role : Write. A regular read token won't work here.\n",
        "hf_token_write = \"\" #@param {type:\"string\"}\n",
        "if hf_token_write ==\"\":\n",
        "  print('\u001b[1;32mYour Hugging Face write access token : ')\n",
        "  hf_token_write=input()\n",
        "\n",
        "hf_token = hf_token_write\n",
        "\n",
        "api = HfApi()\n",
        "your_username = api.whoami(token=hf_token)[\"name\"]\n",
        "\n",
        "if(Save_concept_to == \"Public_Library\"):\n",
        "  repo_id = f\"sd-dreambooth-library/{slugify(name_of_your_concept)}\"\n",
        "  #Join the Concepts Library organization if you aren't part of it already\n",
        "  !curl -X POST -H 'Authorization: Bearer '$hf_token -H 'Content-Type: application/json' https://huggingface.co/organizations/sd-dreambooth-library/share/SSeOwppVCscfTEzFGQaqpfcjukVeNrKNHX\n",
        "else:\n",
        "  repo_id = f\"{your_username}/{slugify(name_of_your_concept)}\"\n",
        "output_dir = f'/content/models/'+INSTANCE_NAME\n",
        "\n",
        "def bar(prg):\n",
        "    br=\"\u001b[1;33mUploading to HuggingFace : \" '\u001b[0m|'+'█' * prg + ' ' * (25-prg)+'| ' +str(prg*4)+ \"%\"\n",
        "    return br\n",
        "\n",
        "print(\"\u001b[1;32mLoading...\")\n",
        "\n",
        "\n",
        "if NM==\"False\":\n",
        "  with capture.capture_output() as cap:\n",
        "    %cd $OUTPUT_DIR\n",
        "    !rm -r safety_checker feature_extractor .git\n",
        "    !rm model_index.json\n",
        "    !git init\n",
        "    !git lfs install --system --skip-repo\n",
        "    !git remote add -f origin  \"https://USER:{hf_token}@huggingface.co/runwayml/stable-diffusion-v1-5\"\n",
        "    !git config core.sparsecheckout true\n",
        "    !echo -e \"feature_extractor\\nsafety_checker\\nmodel_index.json\" > .git/info/sparse-checkout\n",
        "    !git pull origin main\n",
        "    !rm -r .git\n",
        "    %cd /content\n",
        "\n",
        "\n",
        "\n",
        "if os.path.exists('/content/sample_images'):\n",
        "  !rm -r /content/sample_images\n",
        "Samples=\"/content/sample_images\"\n",
        "!mkdir $Samples\n",
        "clear_output()\n",
        "print(\"\u001b[1;32mUpload Sample images of the model\")\n",
        "uploaded = files.upload()\n",
        "for filename in uploaded.keys():\n",
        "  shutil.move(filename, Samples)\n",
        "%cd $Samples\n",
        "!find . -name \"* *\" -type f | rename 's/ /_/g'\n",
        "%cd /content\n",
        "clear_output()\n",
        "\n",
        "print(bar(1))\n",
        "\n",
        "images_upload = os.listdir(Samples)\n",
        "image_string = \"\"\n",
        "instance_prompt_list = []\n",
        "previous_instance_prompt = ''\n",
        "for i, image in enumerate(images_upload):\n",
        "    image_string = f'''\n",
        "{image_string}![{i}](https://huggingface.co/{repo_id}/resolve/main/sample_images/{image})\n",
        "    '''\n",
        "    \n",
        "readme_text = f'''---\n",
        "license: creativeml-openrail-m\n",
        "tags:\n",
        "- text-to-image\n",
        "- stable-diffusion\n",
        "---\n",
        "### {name_of_your_concept} Dreambooth model trained by {api.whoami(token=hf_token)[\"name\"]} with [TheLastBen's fast-DreamBooth](https://colab.research.google.com/github/TheLastBen/fast-stable-diffusion/blob/main/fast-DreamBooth.ipynb) notebook\n",
        "\n",
        "\n",
        "Test the concept via A1111 Colab [fast-Colab-A1111](https://colab.research.google.com/github/TheLastBen/fast-stable-diffusion/blob/main/fast_stable_diffusion_AUTOMATIC1111.ipynb)\n",
        "Or you can run your new concept via `diffusers` [Colab Notebook for Inference](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/sd_dreambooth_inference.ipynb)\n",
        "\n",
        "Sample pictures of this concept:\n",
        "{image_string}\n",
        "'''\n",
        "#Save the readme to a file\n",
        "readme_file = open(\"README.md\", \"w\")\n",
        "readme_file.write(readme_text)\n",
        "readme_file.close()\n",
        "\n",
        "operations = [\n",
        "  CommitOperationAdd(path_in_repo=\"README.md\", path_or_fileobj=\"README.md\"),\n",
        "  CommitOperationAdd(path_in_repo=f\"{Session_Name}.ckpt\",path_or_fileobj=MDLPTH)\n",
        "\n",
        "]\n",
        "create_repo(repo_id,private=True, token=hf_token)\n",
        "\n",
        "api.create_commit(\n",
        "  repo_id=repo_id,\n",
        "  operations=operations,\n",
        "  commit_message=f\"Upload the concept {name_of_your_concept} embeds and token\",\n",
        "  token=hf_token\n",
        ")\n",
        "\n",
        "if NM==\"False\":\n",
        "  api.upload_folder(\n",
        "    folder_path=OUTPUT_DIR+\"/feature_extractor\",\n",
        "    path_in_repo=\"feature_extractor\",\n",
        "    repo_id=repo_id,\n",
        "    token=hf_token\n",
        "  )\n",
        "\n",
        "clear_output()\n",
        "print(bar(4))\n",
        "\n",
        "if NM==\"False\":\n",
        "  api.upload_folder(\n",
        "    folder_path=OUTPUT_DIR+\"/safety_checker\",\n",
        "    path_in_repo=\"safety_checker\",\n",
        "    repo_id=repo_id,\n",
        "    token=hf_token\n",
        "  )\n",
        "\n",
        "clear_output()\n",
        "print(bar(8))\n",
        "\n",
        "\n",
        "api.upload_folder(\n",
        "  folder_path=OUTPUT_DIR+\"/scheduler\",\n",
        "  path_in_repo=\"scheduler\",\n",
        "  repo_id=repo_id,\n",
        "  token=hf_token\n",
        ")\n",
        "\n",
        "clear_output()\n",
        "print(bar(9))\n",
        "\n",
        "api.upload_folder(\n",
        "  folder_path=OUTPUT_DIR+\"/text_encoder\",\n",
        "  path_in_repo=\"text_encoder\",\n",
        "  repo_id=repo_id,\n",
        "  token=hf_token\n",
        ")\n",
        "\n",
        "clear_output()\n",
        "print(bar(12))\n",
        "\n",
        "api.upload_folder(\n",
        "  folder_path=OUTPUT_DIR+\"/tokenizer\",\n",
        "  path_in_repo=\"tokenizer\",\n",
        "  repo_id=repo_id,\n",
        "  token=hf_token\n",
        ")\n",
        "\n",
        "clear_output()\n",
        "print(bar(13))\n",
        "\n",
        "api.upload_folder(\n",
        "  folder_path=OUTPUT_DIR+\"/unet\",\n",
        "  path_in_repo=\"unet\",\n",
        "  repo_id=repo_id,\n",
        "  token=hf_token\n",
        ")\n",
        "\n",
        "clear_output()\n",
        "print(bar(21))\n",
        "\n",
        "api.upload_folder(\n",
        "  folder_path=OUTPUT_DIR+\"/vae\",\n",
        "  path_in_repo=\"vae\",\n",
        "  repo_id=repo_id,\n",
        "  token=hf_token\n",
        ")\n",
        "\n",
        "clear_output()\n",
        "print(bar(23))\n",
        "\n",
        "api.upload_file(\n",
        "  path_or_fileobj=OUTPUT_DIR+\"/model_index.json\",\n",
        "  path_in_repo=\"model_index.json\",\n",
        "  repo_id=repo_id,\n",
        "  token=hf_token\n",
        ")\n",
        "\n",
        "clear_output()\n",
        "print(bar(24))\n",
        "\n",
        "api.upload_folder(\n",
        "  folder_path=Samples,\n",
        "  path_in_repo=\"sample_images\",\n",
        "  repo_id=repo_id,\n",
        "  token=hf_token\n",
        ")\n",
        "\n",
        "clear_output()\n",
        "print(bar(25))\n",
        "\n",
        "display_markdown(f'''## Your concept was saved successfully. [Click here to access it](https://huggingface.co/{repo_id})\n",
        "''', raw=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iVqNi8IDzA1Z"
      },
      "outputs": [],
      "source": [
        "#@markdown #Free Gdrive Space\n",
        "\n",
        "#@markdown Display the list of sessions from your gdrive and choose which ones to remove.\n",
        "\n",
        "import ipywidgets as widgets\n",
        "\n",
        "Sessions=os.listdir(\"/content/gdrive/MyDrive/Fast-Dreambooth/Sessions\")\n",
        "\n",
        "s = widgets.Select(\n",
        "    options=Sessions,\n",
        "    rows=5,\n",
        "    description='',\n",
        "    disabled=False\n",
        ")\n",
        "\n",
        "out=widgets.Output()\n",
        "\n",
        "d = widgets.Button(\n",
        "    description='Remove',\n",
        "    disabled=False,\n",
        "    button_style='warning',\n",
        "    tooltip='Removet the selected session',\n",
        "    icon='warning'\n",
        ")\n",
        "\n",
        "def rem(d):\n",
        "    with out:\n",
        "        if s.value is not None:\n",
        "            clear_output()\n",
        "            print(\"\u001b[1;33mTHE SESSION \u001b[1;31m\"+s.value+\" \u001b[1;33mHAS BEEN REMOVED FROM YOUR GDRIVE\")\n",
        "            !rm -r '/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/{s.value}'\n",
        "            s.options=os.listdir(\"/content/gdrive/MyDrive/Fast-Dreambooth/Sessions\")       \n",
        "        else:\n",
        "            d.close()\n",
        "            s.close()\n",
        "            clear_output()\n",
        "            print(\"\u001b[1;32mNOTHING TO REMOVE\")\n",
        "\n",
        "d.on_click(rem)\n",
        "if s.value is not None:\n",
        "    display(s,d,out)\n",
        "else:\n",
        "    print(\"\u001b[1;32mNOTHING TO REMOVE\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "bbKbx185zqlz",
        "AaLtXBbPleBr"
      ],
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}