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 https://github.com/TheLastBen/diffusers\n",
        "!pip install -q git+https://github.com/TheLastBen/diffusers\n",
        "!pip install -q accelerate==0.12.0\n",
        "!pip install -q OmegaConf\n",
        "!pip install -q wget\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": [
        "# Downloading the model"
      ]
    },
    {
      "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",
        "with capture.capture_output() as cap: \n",
        "  %cd /content/\n",
        "\n",
        "Huggingface_Token = \"\" #@param {type:\"string\"}\n",
        "token=Huggingface_Token\n",
        "\n",
        "#@markdown - Download the original v1.5 model.\n",
        "\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",
        "Compatiblity_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",
        "def downloadmodel_hf():\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/{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-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;31mCheck the link you provided')\n",
        "         time.sleep(5)\n",
        "\n",
        "if Path_to_HuggingFace != \"\":\n",
        "  downloadmodel_hf()\n",
        "\n",
        "elif CKPT_Path !=\"\":\n",
        "  if os.path.exists('/content/stable-diffusion-v1-5'):\n",
        "    !rm -r /content/stable-diffusion-v1-5\n",
        "  if os.path.exists(str(CKPT_Path)):\n",
        "    !mkdir /content/stable-diffusion-v1-5\n",
        "    with capture.capture_output() as cap:\n",
        "      if Compatiblity_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-v1-5\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-v1-5        \n",
        "    if os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "      !rm /content/v1-inference.yaml\n",
        "      clear_output()\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-v1-5\n",
        "      while not os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "        print('\u001b[1;31mConversion error, Insufficient RAM or corrupt CKPT, use a 4.7GB 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",
        "elif CKPT_Link !=\"\":   \n",
        "    if os.path.exists('/content/stable-diffusion-v1-5'):\n",
        "      !rm -r /content/stable-diffusion-v1-5     \n",
        "    !gdown --fuzzy $CKPT_Link -O model.ckpt    \n",
        "    if os.path.exists('/content/model.ckpt'):\n",
        "      if os.path.getsize(\"/content/model.ckpt\") > 1810671599:\n",
        "        !mkdir /content/stable-diffusion-v1-5\n",
        "        with capture.capture_output() as cap: \n",
        "          if Compatiblity_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-v1-5\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-v1-5\n",
        "        if os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "          clear_output()\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-v1-5\n",
        "          !rm /content/model.ckpt\n",
        "          while not os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "            print('\u001b[1;31mConversion error, Insufficient RAM or corrupt CKPT, use a 4.7GB 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",
        "else:\n",
        "  downloadmodel()"
      ]
    },
    {
      "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",
        "import wget\n",
        "import time\n",
        "\n",
        "#@markdown #Create/Load a Session\n",
        "\n",
        "def fdownloadmodel():\n",
        "  token=input(\"Insert your huggingface token :\")\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",
        "MODEL_NAME=\"/content/stable-diffusion-v1-5\"\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",
        "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",
        "#@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",
        "if os.path.exists(str(SESSION_DIR)) and not os.path.exists(str(SESSION_DIR+\"/\"+Session_Name+'.ckpt')):\n",
        "  print('\u001b[1;32mLoading session with no previous model, using the original model or the custom downloaded model')\n",
        "  reg()\n",
        "  if not os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "    if os.path.exists('/content/stable-diffusion-v1-5'):\n",
        "      !rm -r '/content/stable-diffusion-v1-5'    \n",
        "    fdownloadmodel()\n",
        "  if not os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "    print('\u001b[1;31mError downloading the model, make sure you have accepted the terms at https://huggingface.co/runwayml/stable-diffusion-v1-5')\n",
        "  else:\n",
        "    print('\u001b[1;32mSession Loaded, proceed to uploading instance images')\n",
        "\n",
        "elif os.path.exists(str(SESSION_DIR+\"/\"+Session_Name+'.ckpt')):\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 not os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "      if os.path.exists('/content/stable-diffusion-v1-5'):\n",
        "        !rm -r '/content/stable-diffusion-v1-5'\n",
        "      fdownloadmodel()\n",
        "    if os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "      print('\u001b[1;32mSession created, proceed to uploading instance images')\n",
        "    else:\n",
        "      print('\u001b[1;31mError downloading the model, make sure you have accepted the terms at https://huggingface.co/runwayml/stable-diffusion-v1-5')  \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",
        "    #@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",
        "\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",
        "  with capture.capture_output() as cap:\n",
        "    if Crop_images:\n",
        "      for filename in os.listdir(IMAGES_FOLDER_OPTIONAL):\n",
        "        extension = filename.split(\".\")[1]\n",
        "        identifier=filename.split(\".\")[0]\n",
        "        new_path_with_file = os.path.join(IMAGES_FOLDER_OPTIONAL, filename)\n",
        "        file = Image.open(new_path_with_file)\n",
        "        width, height = file.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",
        "        %cp -r \"$IMAGES_FOLDER_OPTIONAL/.\" \"$INSTANCE_DIR\"\n",
        "        clear_output()      \n",
        "    else:\n",
        "      %cp -r \"$IMAGES_FOLDER_OPTIONAL/.\" \"$INSTANCE_DIR\"\n",
        "\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",
        "  print('\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 uploaded.keys():\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",
        "      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 uploaded.keys():\n",
        "      shutil.move(filename, INSTANCE_DIR)\n",
        "      clear_output()\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",
        "  print('\u001b[1;32mDone, proceed to the training cell')\n",
        "\n",
        "with capture.capture_output() as cap:\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",
        "import random\n",
        "\n",
        "Resume_Training = False #@param {type:\"boolean\"}\n",
        "\n",
        "if not Resume_Training and not os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "  if os.path.exists('/content/stable-diffusion-v1-5'):\n",
        "    !rm -r '/content/stable-diffusion-v1-5'\n",
        "  print('\u001b[1;31mOriginal model not found, downloading....\u001b[0m')\n",
        "  fdownloadmodel()\n",
        "  if os.path.exists('/content/stable-diffusion-v1-5/unet/diffusion_pytorch_model.bin'):\n",
        "     print('\u001b[1;32mModel downloaded, proceeding to training...')\n",
        "  else:\n",
        "     print('\u001b[1;31mError downloading the model, make sure you have accepted the terms at https://huggingface.co/runwayml/stable-diffusion-v1-5')  \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), if you're getting memory issues, check the box below (slower speed but memory effecient) :\n",
        "\n",
        "Reduce_memory_usage = False #@param {type:\"boolean\"}\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 Reduce_memory_usage:\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",
        "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",
        "\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 - Keep the % low for better style transfer, more training steps will be necessary for good results.\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",
        "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",
        "  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",
        "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/CompVis/stable-diffusion\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",
        "  if not os.path.exists('/content/gdrive/MyDrive/sd/stable-diffusion/src/k-diffusion/k_diffusion'):\n",
        "    !mkdir /content/gdrive/MyDrive/sd/stable-diffusion/src\n",
        "    %cd /content/gdrive/MyDrive/sd/stable-diffusion/src\n",
        "    !git clone https://github.com/CompVis/taming-transformers\n",
        "    !git clone https://github.com/openai/CLIP\n",
        "    !mv /content/gdrive/MyDrive/sd/stable-diffusion/src/CLIP /content/gdrive/MyDrive/sd/stable-diffusion/src/clip\n",
        "    !git clone https://github.com/TencentARC/GFPGAN\n",
        "    !mv  /content/gdrive/MyDrive/sd/stable-diffusion/src/GFPGAN/gfpgan /content/gdrive/MyDrive/sd/stable-diffusion-webui\n",
        "    !git clone https://github.com/salesforce/BLIP\n",
        "    !mv  /content/gdrive/MyDrive/sd/stable-diffusion/src/BLIP /content/gdrive/MyDrive/sd/stable-diffusion/src/blip\n",
        "    !git clone https://github.com/sczhou/CodeFormer\n",
        "    !mv  /content/gdrive/MyDrive/sd/stable-diffusion/src/CodeFormer /content/gdrive/MyDrive/sd/stable-diffusion/src/codeformer\n",
        "    !git clone https://github.com/xinntao/Real-ESRGAN\n",
        "    !mv  /content/gdrive/MyDrive/sd/stable-diffusion/src/Real-ESRGAN/ /content/gdrive/MyDrive/sd/stable-diffusion/src/realesrgan\n",
        "    !git clone https://github.com/crowsonkb/k-diffusion.git\n",
        "    !cp -r /content/gdrive/MyDrive/sd/stable-diffusion/src/k-diffusion/k_diffusion /content/gdrive/MyDrive/sd/stable-diffusion-webui\n",
        "    !git clone https://github.com/Hafiidz/latent-diffusion\n",
        "    !cp -r  /content/gdrive/MyDrive/sd/stable-diffusion/ldm /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",
        "    !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",
        "    %cd /content/gdrive/MyDrive/sd/stable-diffusion-webui/ldm/modules\n",
        "    !wget -O attention.py https://raw.githubusercontent.com/TheLastBen/fast-stable-diffusion/main/precompiled/attention.py    \n",
        "\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",
        "  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        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/stable-diffusion/\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 --ckpt \"$path_to_trained_model\"\n",
        "else:\n",
        "  !python /content/gdrive/MyDrive/sd/stable-diffusion-webui/webui.py $share --disable-safe-unpickle --ckpt-dir \"$path_to_trained_model\""
      ]
    },
    {
      "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",
        "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",
        "    instance_prompt = image.split(\"_\")[0]\n",
        "    if(instance_prompt != previous_instance_prompt):\n",
        "      title_instance_prompt_string = instance_prompt\n",
        "      instance_prompt_list.append(instance_prompt)\n",
        "    else:\n",
        "      title_instance_prompt_string = ''\n",
        "    previous_instance_prompt = instance_prompt\n",
        "    image_string = f'''\n",
        "{title_instance_prompt_string}\n",
        "{image_string}![{instance_prompt} {i}](https://huggingface.co/{repo_id}/resolve/main/concept_images/{image})\n",
        "    '''\n",
        "    \n",
        "readme_text = f'''---\n",
        "license: creativeml-openrail-m\n",
        "tags:\n",
        "- text-to-image\n",
        "---\n",
        "### {name_of_your_concept} on Stable Diffusion via Dreambooth trained on the [fast-DreamBooth.ipynb by TheLastBen](https://colab.research.google.com/github/TheLastBen/fast-stable-diffusion/blob/main/fast-DreamBooth.ipynb) notebook\n",
        "#### Model by {api.whoami(token=hf_token)[\"name\"]}\n",
        "This your the Stable Diffusion model fine-tuned the {name_of_your_concept} concept taught to Stable Diffusion with Dreambooth.\n",
        "It can be used by modifying the `instance_prompt(s)`: **{', '.join(instance_prompt_list)}**\n",
        "\n",
        "You can also train your own concepts and upload them to the library by using [the fast-DremaBooth.ipynb by TheLastBen](https://colab.research.google.com/github/TheLastBen/fast-stable-diffusion/blob/main/fast-DreamBooth.ipynb).\n",
        "You can run your new 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), [Spaces with the Public Concepts loaded](https://huggingface.co/spaces/sd-dreambooth-library/stable-diffusion-dreambooth-concepts)\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",
        "#Save the token identifier to a file\n",
        "text_file = open(\"token_identifier.txt\", \"w\")\n",
        "text_file.write(', '.join(instance_prompt_list))\n",
        "text_file.close()\n",
        "\n",
        "operations = [\n",
        "  CommitOperationAdd(path_in_repo=\"token_identifier.txt\", path_or_fileobj=\"token_identifier.txt\"),\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",
        "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",
        "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",
        "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=\"concept_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)"
      ]
    }
  ],
  "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
}