fast-DreamBooth.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# **DreamBooth colab From https://github.com/TheLastBen/fast-stable-diffusion, if you face any issues, feel free to discuss them.** \n",
        "Keep you notebook updated for best experience.\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 git+https://github.com/TheLastBen/diffusers\n",
        "!pip install -q accelerate==0.12.0\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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "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 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 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 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 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 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": [
        "#@markdown # Downloading the model\n",
        "import os\n",
        "import time\n",
        "from IPython.display import clear_output\n",
        "\n",
        "Huggingface_Token = \"\" #@param {type:\"string\"}\n",
        "#@markdown ---\n",
        "#@markdown (Make sure you accepted the terms in https://huggingface.co/CompVis/stable-diffusion-v1-4)\n",
        "token=Huggingface_Token\n",
        "if token == \"\" and not os.path.exists('/content/stable-diffusion-v1-4'):\n",
        "  token=input(\"Insert your huggingface token :\")\n",
        "  %cd /content/\n",
        "  clear_output()\n",
        "  !git init\n",
        "  !git lfs install --system --skip-repo\n",
        "  !git clone \"https://USER:{token}@huggingface.co/CompVis/stable-diffusion-v1-4\"\n",
        "  if os.path.exists('/content/stable-diffusion-v1-4/unet/diffusion_pytorch_model.bin'):\n",
        "    clear_output()\n",
        "    print('\u001b[1;32mDONE !')\n",
        "  else:\n",
        "    while not os.path.exists('/content/stable-diffusion-v1-4'):\n",
        "         print('\u001b[1;31mMake sure you accepted the terms in https://huggingface.co/CompVis/stable-diffusion-v1-4')\n",
        "         time.sleep(5)\n",
        "         \n",
        "elif not os.path.exists('/content/stable-diffusion-v1-4'):\n",
        "  %cd /content/\n",
        "  clear_output()\n",
        "  !git init\n",
        "  !git lfs install --system --skip-repo\n",
        "  !git clone \"https://USER:{token}@huggingface.co/CompVis/stable-diffusion-v1-4\"\n",
        "  if os.path.exists('/content/stable-diffusion-v1-4/unet/diffusion_pytorch_model.bin'):\n",
        "    clear_output()\n",
        "    print('\u001b[1;32mDONE !')\n",
        "  else:\n",
        "    while not os.path.exists('/content/stable-diffusion-v1-4'):\n",
        "         print('\u001b[1;31mMake sure you accepted the terms in https://huggingface.co/CompVis/stable-diffusion-v1-4')\n",
        "         time.sleep(5)\n",
        "else:\n",
        "  print(\"\u001b[1;32mModel already exists\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0tN76Cj5P3RL"
      },
      "source": [
        "# Dreambooth"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 38
        },
        "id": "1pH1oP-7yBZm",
        "outputId": "9780a9ab-60e2-4bdb-f5bc-3e15b8619366"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "     <input type=\"file\" id=\"files-572594a5-b1b0-4cd0-838c-85eb7bc4e89f\" name=\"files[]\" multiple disabled\n",
              "        style=\"border:none\" />\n",
              "     <output id=\"result-572594a5-b1b0-4cd0-838c-85eb7bc4e89f\">\n",
              "      Upload widget is only available when the cell has been executed in the\n",
              "      current browser session. Please rerun this cell to enable.\n",
              "      </output>\n",
              "      <script>// Copyright 2017 Google LLC\n",
              "//\n",
              "// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
              "// you may not use this file except in compliance with the License.\n",
              "// You may obtain a copy of the License at\n",
              "//\n",
              "//      http://www.apache.org/licenses/LICENSE-2.0\n",
              "//\n",
              "// Unless required by applicable law or agreed to in writing, software\n",
              "// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
              "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
              "// See the License for the specific language governing permissions and\n",
              "// limitations under the License.\n",
              "\n",
              "/**\n",
              " * @fileoverview Helpers for google.colab Python module.\n",
              " */\n",
              "(function(scope) {\n",
              "function span(text, styleAttributes = {}) {\n",
              "  const element = document.createElement('span');\n",
              "  element.textContent = text;\n",
              "  for (const key of Object.keys(styleAttributes)) {\n",
              "    element.style[key] = styleAttributes[key];\n",
              "  }\n",
              "  return element;\n",
              "}\n",
              "\n",
              "// Max number of bytes which will be uploaded at a time.\n",
              "const MAX_PAYLOAD_SIZE = 100 * 1024;\n",
              "\n",
              "function _uploadFiles(inputId, outputId) {\n",
              "  const steps = uploadFilesStep(inputId, outputId);\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  // Cache steps on the outputElement to make it available for the next call\n",
              "  // to uploadFilesContinue from Python.\n",
              "  outputElement.steps = steps;\n",
              "\n",
              "  return _uploadFilesContinue(outputId);\n",
              "}\n",
              "\n",
              "// This is roughly an async generator (not supported in the browser yet),\n",
              "// where there are multiple asynchronous steps and the Python side is going\n",
              "// to poll for completion of each step.\n",
              "// This uses a Promise to block the python side on completion of each step,\n",
              "// then passes the result of the previous step as the input to the next step.\n",
              "function _uploadFilesContinue(outputId) {\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  const steps = outputElement.steps;\n",
              "\n",
              "  const next = steps.next(outputElement.lastPromiseValue);\n",
              "  return Promise.resolve(next.value.promise).then((value) => {\n",
              "    // Cache the last promise value to make it available to the next\n",
              "    // step of the generator.\n",
              "    outputElement.lastPromiseValue = value;\n",
              "    return next.value.response;\n",
              "  });\n",
              "}\n",
              "\n",
              "/**\n",
              " * Generator function which is called between each async step of the upload\n",
              " * process.\n",
              " * @param {string} inputId Element ID of the input file picker element.\n",
              " * @param {string} outputId Element ID of the output display.\n",
              " * @return {!Iterable<!Object>} Iterable of next steps.\n",
              " */\n",
              "function* uploadFilesStep(inputId, outputId) {\n",
              "  const inputElement = document.getElementById(inputId);\n",
              "  inputElement.disabled = false;\n",
              "\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  outputElement.innerHTML = '';\n",
              "\n",
              "  const pickedPromise = new Promise((resolve) => {\n",
              "    inputElement.addEventListener('change', (e) => {\n",
              "      resolve(e.target.files);\n",
              "    });\n",
              "  });\n",
              "\n",
              "  const cancel = document.createElement('button');\n",
              "  inputElement.parentElement.appendChild(cancel);\n",
              "  cancel.textContent = 'Cancel upload';\n",
              "  const cancelPromise = new Promise((resolve) => {\n",
              "    cancel.onclick = () => {\n",
              "      resolve(null);\n",
              "    };\n",
              "  });\n",
              "\n",
              "  // Wait for the user to pick the files.\n",
              "  const files = yield {\n",
              "    promise: Promise.race([pickedPromise, cancelPromise]),\n",
              "    response: {\n",
              "      action: 'starting',\n",
              "    }\n",
              "  };\n",
              "\n",
              "  cancel.remove();\n",
              "\n",
              "  // Disable the input element since further picks are not allowed.\n",
              "  inputElement.disabled = true;\n",
              "\n",
              "  if (!files) {\n",
              "    return {\n",
              "      response: {\n",
              "        action: 'complete',\n",
              "      }\n",
              "    };\n",
              "  }\n",
              "\n",
              "  for (const file of files) {\n",
              "    const li = document.createElement('li');\n",
              "    li.append(span(file.name, {fontWeight: 'bold'}));\n",
              "    li.append(span(\n",
              "        `(${file.type || 'n/a'}) - ${file.size} bytes, ` +\n",
              "        `last modified: ${\n",
              "            file.lastModifiedDate ? file.lastModifiedDate.toLocaleDateString() :\n",
              "                                    'n/a'} - `));\n",
              "    const percent = span('0% done');\n",
              "    li.appendChild(percent);\n",
              "\n",
              "    outputElement.appendChild(li);\n",
              "\n",
              "    const fileDataPromise = new Promise((resolve) => {\n",
              "      const reader = new FileReader();\n",
              "      reader.onload = (e) => {\n",
              "        resolve(e.target.result);\n",
              "      };\n",
              "      reader.readAsArrayBuffer(file);\n",
              "    });\n",
              "    // Wait for the data to be ready.\n",
              "    let fileData = yield {\n",
              "      promise: fileDataPromise,\n",
              "      response: {\n",
              "        action: 'continue',\n",
              "      }\n",
              "    };\n",
              "\n",
              "    // Use a chunked sending to avoid message size limits. See b/62115660.\n",
              "    let position = 0;\n",
              "    do {\n",
              "      const length = Math.min(fileData.byteLength - position, MAX_PAYLOAD_SIZE);\n",
              "      const chunk = new Uint8Array(fileData, position, length);\n",
              "      position += length;\n",
              "\n",
              "      const base64 = btoa(String.fromCharCode.apply(null, chunk));\n",
              "      yield {\n",
              "        response: {\n",
              "          action: 'append',\n",
              "          file: file.name,\n",
              "          data: base64,\n",
              "        },\n",
              "      };\n",
              "\n",
              "      let percentDone = fileData.byteLength === 0 ?\n",
              "          100 :\n",
              "          Math.round((position / fileData.byteLength) * 100);\n",
              "      percent.textContent = `${percentDone}% done`;\n",
              "\n",
              "    } while (position < fileData.byteLength);\n",
              "  }\n",
              "\n",
              "  // All done.\n",
              "  yield {\n",
              "    response: {\n",
              "      action: 'complete',\n",
              "    }\n",
              "  };\n",
              "}\n",
              "\n",
              "scope.google = scope.google || {};\n",
              "scope.google.colab = scope.google.colab || {};\n",
              "scope.google.colab._files = {\n",
              "  _uploadFiles,\n",
              "  _uploadFilesContinue,\n",
              "};\n",
              "})(self);\n",
              "</script> "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import os\n",
        "import shutil\n",
        "from google.colab import files\n",
        "from IPython.display import clear_output\n",
        "from IPython.utils import capture\n",
        "#@markdown #Setting up\n",
        "#@markdown ---\n",
        "\n",
        "Training_Subject = \"Character\" #@param [\"Character\", \"Object\", \"Style\", \"Artist\", \"Movie\", \"TV Show\"] \n",
        "\n",
        "With_Prior_Preservation = \"Yes\" #@param [\"Yes\", \"No\"] \n",
        "#@markdown - (Only for characters or objects) With the prior reservation method the results are more coherent and better, you will have to either upload around 200 pictures of the class you're training (dog, person, car, house ...), but more is better, if you don't upload them, it will automatically generate them at the cost of quality and time. For excellent results, manually filter around 400 images of the class and 50 pictures of the instance (a photo of yourself or your dog ...), try 2000 steps then tune it from there.\n",
        "\n",
        "MODEL_NAME=\"/content/stable-diffusion-v1-4\"\n",
        "\n",
        "\n",
        "SUBJECT_TYPE = \"\" #@param{type: 'string'}\n",
        "while SUBJECT_TYPE==\"\":\n",
        "   SUBJECT_TYPE=input('Input the subject type:')\n",
        "\n",
        "#@markdown - If you're training on a character or an object, the subject type would be : Man, Woman, Shirt, Car, Dog, Baby ...etc\n",
        "#@markdown - If you're training on a Style, the subject type would be : impressionist, brutalist, abstract, use \"beautiful\" for a general style...etc\n",
        "#@markdown - If you're training on a Movie/Show, the subject type would be : Action, Drama, Science-fiction, Comedy ...etc\n",
        "#@markdown - If you're training on an Artist, the subject type would be : Painting, sketch, drawing, photography, art ...etc\n",
        "\n",
        "\n",
        "INSTANCE_NAME= \"\" #@param{type: 'string'}\n",
        "while INSTANCE_NAME==\"\":\n",
        "   INSTANCE_NAME=input('Input the instance name (identifier) :')\n",
        "\n",
        "#@markdown - The instance is an identifier, choose a unique identifier unknown by stable diffusion. \n",
        "\n",
        "INSTANCE_DIR_OPTIONAL=\"\" #@param{type: 'string'}\n",
        "INSTANCE_DIR=INSTANCE_DIR_OPTIONAL\n",
        "while INSTANCE_DIR_OPTIONAL!=\"\" and not os.path.exists(str(INSTANCE_DIR)):\n",
        "    INSTANCE_DIR=input('\u001b[1;31mThe instance folder specified does not exist, use the colab file explorer to copy the path :')\n",
        "\n",
        "#@markdown - If the number of instance pictures is large, it is preferable to specify directly the folder instead of uploading, leave EMPTY to upload.\n",
        "\n",
        "CLASS_DIR=\"/content/data/\"+ SUBJECT_TYPE\n",
        "Number_of_subject_images=200#@param{type: 'number'}\n",
        "while Number_of_subject_images==None:\n",
        "     Number_of_subject_images=input('Input the number of subject images :')\n",
        "SUBJECT_IMAGES=Number_of_subject_images\n",
        "\n",
        "if Training_Subject==\"Character\" or Training_Subject==\"Object\":\n",
        "  PT=\"photo of \"+INSTANCE_NAME+\" \"+SUBJECT_TYPE\n",
        "  CPT=\"a photo of a \"+SUBJECT_TYPE+\", ultra detailed, 8k\"\n",
        "elif Training_Subject==\"Style\":\n",
        "  With_Prior_Preservation = \"No\"\n",
        "  PT=\"in the \"+SUBJECT_TYPE+\" style of \"+INSTANCE_NAME\n",
        "elif Training_Subject==\"Artist\":\n",
        "  With_Prior_Preservation = \"No\"\n",
        "  PT=SUBJECT_TYPE+\" By \"+INSTANCE_NAME\n",
        "elif Training_Subject==\"Movie\": \n",
        "  With_Prior_Preservation = \"No\"\n",
        "  PT=\"from the \"+SUBJECT_TYPE+\" movie \"+ INSTANCE_NAME\n",
        "elif Training_Subject==\"TV Show\": \n",
        "  With_Prior_Preservation = \"No\"\n",
        "  PT=\"from the \"+SUBJECT_TYPE+\" tv show \"+ INSTANCE_NAME   \n",
        "  \n",
        "OUTPUT_DIR=\"/content/models/\"+ INSTANCE_NAME\n",
        "\n",
        "if INSTANCE_DIR_OPTIONAL==\"\":\n",
        "  INSTANCE_DIR=\"/content/data/\"+INSTANCE_NAME\n",
        "  !mkdir -p \"$INSTANCE_DIR\"\n",
        "  uploaded = files.upload()\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",
        "print('\u001b[1;32mOK')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 38
        },
        "id": "LVCHYyYTIc61",
        "outputId": "179cb686-15d0-45a1-b84c-5d59611513bc"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "     <input type=\"file\" id=\"files-260fabff-59c2-4ac2-96e1-bc726973f55d\" name=\"files[]\" multiple disabled\n",
              "        style=\"border:none\" />\n",
              "     <output id=\"result-260fabff-59c2-4ac2-96e1-bc726973f55d\">\n",
              "      Upload widget is only available when the cell has been executed in the\n",
              "      current browser session. Please rerun this cell to enable.\n",
              "      </output>\n",
              "      <script>// Copyright 2017 Google LLC\n",
              "//\n",
              "// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
              "// you may not use this file except in compliance with the License.\n",
              "// You may obtain a copy of the License at\n",
              "//\n",
              "//      http://www.apache.org/licenses/LICENSE-2.0\n",
              "//\n",
              "// Unless required by applicable law or agreed to in writing, software\n",
              "// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
              "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
              "// See the License for the specific language governing permissions and\n",
              "// limitations under the License.\n",
              "\n",
              "/**\n",
              " * @fileoverview Helpers for google.colab Python module.\n",
              " */\n",
              "(function(scope) {\n",
              "function span(text, styleAttributes = {}) {\n",
              "  const element = document.createElement('span');\n",
              "  element.textContent = text;\n",
              "  for (const key of Object.keys(styleAttributes)) {\n",
              "    element.style[key] = styleAttributes[key];\n",
              "  }\n",
              "  return element;\n",
              "}\n",
              "\n",
              "// Max number of bytes which will be uploaded at a time.\n",
              "const MAX_PAYLOAD_SIZE = 100 * 1024;\n",
              "\n",
              "function _uploadFiles(inputId, outputId) {\n",
              "  const steps = uploadFilesStep(inputId, outputId);\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  // Cache steps on the outputElement to make it available for the next call\n",
              "  // to uploadFilesContinue from Python.\n",
              "  outputElement.steps = steps;\n",
              "\n",
              "  return _uploadFilesContinue(outputId);\n",
              "}\n",
              "\n",
              "// This is roughly an async generator (not supported in the browser yet),\n",
              "// where there are multiple asynchronous steps and the Python side is going\n",
              "// to poll for completion of each step.\n",
              "// This uses a Promise to block the python side on completion of each step,\n",
              "// then passes the result of the previous step as the input to the next step.\n",
              "function _uploadFilesContinue(outputId) {\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  const steps = outputElement.steps;\n",
              "\n",
              "  const next = steps.next(outputElement.lastPromiseValue);\n",
              "  return Promise.resolve(next.value.promise).then((value) => {\n",
              "    // Cache the last promise value to make it available to the next\n",
              "    // step of the generator.\n",
              "    outputElement.lastPromiseValue = value;\n",
              "    return next.value.response;\n",
              "  });\n",
              "}\n",
              "\n",
              "/**\n",
              " * Generator function which is called between each async step of the upload\n",
              " * process.\n",
              " * @param {string} inputId Element ID of the input file picker element.\n",
              " * @param {string} outputId Element ID of the output display.\n",
              " * @return {!Iterable<!Object>} Iterable of next steps.\n",
              " */\n",
              "function* uploadFilesStep(inputId, outputId) {\n",
              "  const inputElement = document.getElementById(inputId);\n",
              "  inputElement.disabled = false;\n",
              "\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  outputElement.innerHTML = '';\n",
              "\n",
              "  const pickedPromise = new Promise((resolve) => {\n",
              "    inputElement.addEventListener('change', (e) => {\n",
              "      resolve(e.target.files);\n",
              "    });\n",
              "  });\n",
              "\n",
              "  const cancel = document.createElement('button');\n",
              "  inputElement.parentElement.appendChild(cancel);\n",
              "  cancel.textContent = 'Cancel upload';\n",
              "  const cancelPromise = new Promise((resolve) => {\n",
              "    cancel.onclick = () => {\n",
              "      resolve(null);\n",
              "    };\n",
              "  });\n",
              "\n",
              "  // Wait for the user to pick the files.\n",
              "  const files = yield {\n",
              "    promise: Promise.race([pickedPromise, cancelPromise]),\n",
              "    response: {\n",
              "      action: 'starting',\n",
              "    }\n",
              "  };\n",
              "\n",
              "  cancel.remove();\n",
              "\n",
              "  // Disable the input element since further picks are not allowed.\n",
              "  inputElement.disabled = true;\n",
              "\n",
              "  if (!files) {\n",
              "    return {\n",
              "      response: {\n",
              "        action: 'complete',\n",
              "      }\n",
              "    };\n",
              "  }\n",
              "\n",
              "  for (const file of files) {\n",
              "    const li = document.createElement('li');\n",
              "    li.append(span(file.name, {fontWeight: 'bold'}));\n",
              "    li.append(span(\n",
              "        `(${file.type || 'n/a'}) - ${file.size} bytes, ` +\n",
              "        `last modified: ${\n",
              "            file.lastModifiedDate ? file.lastModifiedDate.toLocaleDateString() :\n",
              "                                    'n/a'} - `));\n",
              "    const percent = span('0% done');\n",
              "    li.appendChild(percent);\n",
              "\n",
              "    outputElement.appendChild(li);\n",
              "\n",
              "    const fileDataPromise = new Promise((resolve) => {\n",
              "      const reader = new FileReader();\n",
              "      reader.onload = (e) => {\n",
              "        resolve(e.target.result);\n",
              "      };\n",
              "      reader.readAsArrayBuffer(file);\n",
              "    });\n",
              "    // Wait for the data to be ready.\n",
              "    let fileData = yield {\n",
              "      promise: fileDataPromise,\n",
              "      response: {\n",
              "        action: 'continue',\n",
              "      }\n",
              "    };\n",
              "\n",
              "    // Use a chunked sending to avoid message size limits. See b/62115660.\n",
              "    let position = 0;\n",
              "    do {\n",
              "      const length = Math.min(fileData.byteLength - position, MAX_PAYLOAD_SIZE);\n",
              "      const chunk = new Uint8Array(fileData, position, length);\n",
              "      position += length;\n",
              "\n",
              "      const base64 = btoa(String.fromCharCode.apply(null, chunk));\n",
              "      yield {\n",
              "        response: {\n",
              "          action: 'append',\n",
              "          file: file.name,\n",
              "          data: base64,\n",
              "        },\n",
              "      };\n",
              "\n",
              "      let percentDone = fileData.byteLength === 0 ?\n",
              "          100 :\n",
              "          Math.round((position / fileData.byteLength) * 100);\n",
              "      percent.textContent = `${percentDone}% done`;\n",
              "\n",
              "    } while (position < fileData.byteLength);\n",
              "  }\n",
              "\n",
              "  // All done.\n",
              "  yield {\n",
              "    response: {\n",
              "      action: 'complete',\n",
              "    }\n",
              "  };\n",
              "}\n",
              "\n",
              "scope.google = scope.google || {};\n",
              "scope.google.colab = scope.google.colab || {};\n",
              "scope.google.colab._files = {\n",
              "  _uploadFiles,\n",
              "  _uploadFilesContinue,\n",
              "};\n",
              "})(self);\n",
              "</script> "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "#@markdown ##[Optional] Upload or choose a folder of the class pictures (pictures of dogs if you're training on a dog), 200 is good, more is better, if you upload less than Number_of_subject_images, it will automatically generate the rest.\n",
        "\n",
        "import os\n",
        "from google.colab import files\n",
        "from IPython.display import clear_output\n",
        "from IPython.utils import capture\n",
        "import shutil\n",
        "\n",
        "if (With_Prior_Preservation=='No'):\n",
        "  print(\"\u001b[1;32mThis training method/subject doesn't require class images\")\n",
        "\n",
        "else:\n",
        "  CLASS_DIR=\"\" #@param{type: 'string'}\n",
        "  if (CLASS_DIR !=\"\") and os.path.exists(str(CLASS_DIR)):\n",
        "    CLASS_DIR=CLASS_DIR\n",
        "  elif (CLASS_DIR !=\"\") and not os.path.exists(str(CLASS_DIR)):\n",
        "    CLASS_DIR=input('\u001b[1;31mThe folder specified does not exist, use the colab file explorer to copy the path :')\n",
        "  elif (CLASS_DIR ==\"\"):\n",
        "    CLASS_DIR=\"/content/data/\"+ SUBJECT_TYPE\n",
        "    !mkdir -p \"data/$SUBJECT_TYPE\"\n",
        "    uploaded = files.upload()\n",
        "    for filename in uploaded.keys():\n",
        "      shutil.move(filename, CLASS_DIR)\n",
        "      clear_output()\n",
        "\n",
        "with capture.capture_output() as cap:\n",
        "  if os.path.exists(str(CLASS_DIR)):\n",
        "    %cd \"$CLASS_DIR\"\n",
        "    !find . -name \"* *\" -type f | rename 's/ /_/g'\n",
        "    %cd /content\n",
        "\n",
        "\n",
        "#@markdown - To save time, if you specify a CLASS_DIR which is a folder that containes class images (eg: 200 pics of a dog), dreambooth will use this folder. \n",
        "#@markdown -Leave it empty if you want to upload\n",
        "\n",
        "#@markdown - Skip the cell if you want it to generate class (subject) images.  "
      ]
    },
    {
      "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",
        "\n",
        "fp16 = True #@param {type:\"boolean\"}\n",
        "if fp16:\n",
        "  prec=\"fp16\"\n",
        "else:\n",
        "  prec=\"no\"\n",
        "\n",
        "#@markdown  - fp16 or half precision meaning slightly lower quality but double the speed.\n",
        "s = getoutput('nvidia-smi')\n",
        "if 'A100' in s:\n",
        "  precision=\"no\"\n",
        "else:\n",
        "  precision=prec\n",
        "\n",
        "Training_Steps=\"1500\" #@param{type: 'string'}\n",
        "Seed=45576 #@param{type: 'number'}\n",
        "\n",
        "if With_Prior_Preservation=='No':\n",
        "  #@markdown ####More steps, better results, but longer training time\n",
        "  !accelerate launch /content/diffusers/examples/dreambooth/train_dreambooth.py \\\n",
        "    --train_text_encoder \\\n",
        "    --pretrained_model_name_or_path=\"$MODEL_NAME\" \\\n",
        "    --instance_data_dir=\"$INSTANCE_DIR\" \\\n",
        "    --output_dir=\"$OUTPUT_DIR\" \\\n",
        "    --instance_prompt=\"$PT\" \\\n",
        "    --seed=$Seed \\\n",
        "    --resolution=512 \\\n",
        "    --mixed_precision=$precision \\\n",
        "    --train_batch_size=1 \\\n",
        "    --gradient_accumulation_steps=1 \\\n",
        "    --use_8bit_adam \\\n",
        "    --learning_rate=2e-6 \\\n",
        "    --lr_scheduler=\"constant\" \\\n",
        "    --lr_warmup_steps=0 \\\n",
        "    --max_train_steps=$Training_Steps \n",
        "\n",
        "else:\n",
        "\n",
        "  !accelerate launch /content/diffusers/examples/dreambooth/train_dreambooth.py \\\n",
        "    --train_text_encoder \\\n",
        "    --pretrained_model_name_or_path=\"$MODEL_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",
        "    --class_prompt=\"$CPT\" \\\n",
        "    --seed=$Seed \\\n",
        "    --resolution=512 \\\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=\"constant\" \\\n",
        "    --lr_warmup_steps=0 \\\n",
        "    --max_train_steps=$Training_Steps \\\n",
        "    --num_class_images=$SUBJECT_IMAGES\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= \"/content/gdrive/MyDrive/{INSTANCE_NAME}.ckpt\"@' /content/convertosd.py\n",
        "  !python /content/convertosd.py\n",
        "  clear_output()\n",
        "  if os.path.exists('/content/gdrive/MyDrive/'+INSTANCE_NAME+'.ckpt'):\n",
        "    print(\"\u001b[1;32mDONE, the CKPT model is in your Gdrive\")\n",
        "  else:\n",
        "    print(\"\u001b[1;31mSomething went wrong\")\n",
        "else:\n",
        "  print(\"\u001b[1;31mSomething went wrong\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iAZGngFcI8hq"
      },
      "outputs": [],
      "source": [
        "#@markdown # Test the trained model\n",
        "import os\n",
        "from IPython.display import clear_output\n",
        "from subprocess import getoutput\n",
        "from IPython.utils import capture\n",
        "import time\n",
        "\n",
        "INSTANCE__NAME=\"\" #@param{type: 'string'}\n",
        "\n",
        "#@markdown - Leave empty if you want to use the current trained model\n",
        "\n",
        "if INSTANCE__NAME!=\"\":\n",
        "  INSTANCE_NAME=INSTANCE__NAME\n",
        "\n",
        "Use_Custom_Path = False #@param {type:\"boolean\"}\n",
        "\n",
        "try:\n",
        "  INSTANCE_NAME\n",
        "  if Use_Custom_Path:\n",
        "    del INSTANCE_NAME\n",
        "except:\n",
        "  pass\n",
        "#@markdown - if checked, an input box will ask the full path to a desired model\n",
        "\n",
        "try:\n",
        "  INSTANCE_NAME\n",
        "  path_to_trained_model='/content/gdrive/MyDrive/'+INSTANCE_NAME+'.ckpt'\n",
        "except:\n",
        "  path_to_trained_model=input(\"\u001b[1;31mIt seems that you did not perform training during this session or you chose to use a custom path,\\nprovide the full path to the model (including the name of the model):\\n\")\n",
        "     \n",
        "while not os.path.exists(path_to_trained_model):\n",
        "   path_to_trained_model=input(\"\u001b[1;31mThe model doesn't exist on you Gdrive, use the file explorer to get the path : \")\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",
        "\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/gradio-3.4b3.dist-info'):\n",
        "    %cd /content/\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",
        "    %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        demo.queue(concurrency_count=111500)@' /content/gdrive/MyDrive/sd/stable-diffusion-webui/webui.py\n",
        "Use_Gradio_Server = False #@param {type:\"boolean\"}\n",
        "#@markdown  - Only if you have trouble connecting to the local server\n",
        "\n",
        "with capture.capture_output() as cap: \n",
        "  %cd /content\n",
        "\n",
        "share=''\n",
        "if Use_Gradio_Server:\n",
        "  share='--share'\n",
        "  !sed -i '1037s@.*@            self.server_name = server_name@' /usr/local/lib/python3.7/dist-packages/gradio/blocks.py\n",
        "  !sed -i '1039s@.*@            self.server_port = server_port@' /usr/local/lib/python3.7/dist-packages/gradio/blocks.py  \n",
        "  !sed -i '1043s@.*@            self.protocol = \"https\" if self.local_url.startswith(\"https\") else \"http\"@' /usr/local/lib/python3.7/dist-packages/gradio/blocks.py  \n",
        "\n",
        "else:\n",
        "  share=''\n",
        "\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",
        "  !sed -i '1037s@.*@            self.server_name = \"{srv[8:]}\"@' /usr/local/lib/python3.7/dist-packages/gradio/blocks.py\n",
        "  !sed -i '1039s@.*@            self.server_port = 443@' /usr/local/lib/python3.7/dist-packages/gradio/blocks.py\n",
        "  !sed -i '1043s@.*@            self.protocol = \"https\"@' /usr/local/lib/python3.7/dist-packages/gradio/blocks.py  \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",
        "\n",
        "\n",
        "with capture.capture_output() as cap:\n",
        "  %cd /content/gdrive/MyDrive/sd/stable-diffusion/\n",
        "\n",
        "!python /content/gdrive/MyDrive/sd/stable-diffusion-webui/webui.py $share --disable-safe-unpickle --ckpt \"$path_to_trained_model\""
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "bbKbx185zqlz",
        "P8aNBqn9JviD"
      ],
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}