DreamBooth.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Last update：2022/11/12"
      ],
      "metadata": {
        "id": "FFuIIvdaSGo0"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XU7NuMAA2drw",
        "outputId": "6c673e34-26d6-4d60-b35b-05383aa029e6",
        "cellView": "form"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Tesla T4, 15109 MiB, 15109 MiB\n"
          ]
        }
      ],
      "source": [
        "#@markdown Check type of GPU and VRAM available.\n",
        "!nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv,noheader"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BzM7j0ZSc_9c"
      },
      "source": [
        "https://github.com/ShivamShrirao/diffusers/tree/main/examples/dreambooth"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wnTMyW41cC1E"
      },
      "source": [
        "## Install Requirements"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!wget -q https://github.com/ShivamShrirao/diffusers/raw/main/examples/dreambooth/train_dreambooth.py\n",
        "%pip install -qq git+https://github.com/ShivamShrirao/diffusers\n",
        "%pip install -q -U --pre triton\n",
        "%pip install -q accelerate==0.12.0 transformers ftfy bitsandbytes gradio natsort"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RXJP21VrFN6K",
        "outputId": "2086ca94-7f65-4390-bfb2-1477fa31e632"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[K     |████████████████████████████████| 163 kB 36.8 MB/s \n",
            "\u001b[?25h  Building wheel for diffusers (PEP 517) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[K     |████████████████████████████████| 18.7 MB 32.4 MB/s \n",
            "\u001b[K     |████████████████████████████████| 143 kB 38.4 MB/s \n",
            "\u001b[K     |████████████████████████████████| 5.5 MB 58.3 MB/s \n",
            "\u001b[K     |████████████████████████████████| 53 kB 1.6 MB/s \n",
            "\u001b[K     |████████████████████████████████| 62.5 MB 1.2 MB/s \n",
            "\u001b[K     |████████████████████████████████| 11.6 MB 55.4 MB/s \n",
            "\u001b[K     |████████████████████████████████| 7.6 MB 45.1 MB/s \n",
            "\u001b[K     |████████████████████████████████| 55 kB 3.1 MB/s \n",
            "\u001b[K     |████████████████████████████████| 56 kB 4.9 MB/s \n",
            "\u001b[K     |████████████████████████████████| 54 kB 555 kB/s \n",
            "\u001b[K     |████████████████████████████████| 213 kB 66.6 MB/s \n",
            "\u001b[K     |████████████████████████████████| 84 kB 3.8 MB/s \n",
            "\u001b[K     |████████████████████████████████| 84 kB 3.3 MB/s \n",
            "\u001b[K     |████████████████████████████████| 272 kB 73.3 MB/s \n",
            "\u001b[K     |████████████████████████████████| 2.3 MB 51.0 MB/s \n",
            "\u001b[K     |████████████████████████████████| 106 kB 74.6 MB/s \n",
            "\u001b[K     |████████████████████████████████| 63 kB 1.2 MB/s \n",
            "\u001b[K     |████████████████████████████████| 80 kB 8.1 MB/s \n",
            "\u001b[K     |████████████████████████████████| 68 kB 6.9 MB/s \n",
            "\u001b[K     |████████████████████████████████| 46 kB 4.7 MB/s \n",
            "\u001b[K     |████████████████████████████████| 4.1 MB 58.1 MB/s \n",
            "\u001b[K     |████████████████████████████████| 593 kB 74.1 MB/s \n",
            "\u001b[K     |████████████████████████████████| 856 kB 50.1 MB/s \n",
            "\u001b[?25h  Building wheel for ffmpy (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for python-multipart (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hd207Y7E-LZP"
      },
      "source": [
        "### Mount Google Drive."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ihkGY7yMZHDD",
        "outputId": "2b9604be-d2a9-4296-dbc9-f263794f58b4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vwBJIHWV8X5b"
      },
      "source": [
        "### Load Model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o5te03bQ86PQ"
      },
      "source": [
        "Before running this cell, you need to copy the folder to the Google Drive.  \n",
        "https://drive.google.com/drive/folders/1vOaFitkJ1MOig6VD-ubnQkg239Dzpn15?usp=sharing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "pJdrvDCWUhLq"
      },
      "outputs": [],
      "source": [
        "!cp -r \"/content/drive/MyDrive/NovelAI_SD_Model\" \"/content\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DTxPP7lJ-mSo"
      },
      "source": [
        "If you want to use other model from HuggingFace, uncomment the cell below(delete the \"# \") and run the cell below, filling in the token as instructed"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "y4lqqWT_uxD2"
      },
      "outputs": [],
      "source": [
        "# from huggingface_hub import notebook_login\n",
        "# !git config --global credential.helper store\n",
        "# notebook_login()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XfTlc8Mqb8iH"
      },
      "source": [
        "### Install xformers from precompiled wheel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "n6dcjPnnaiCn",
        "outputId": "f8023e73-03b7-4ea0-e475-dc356dec503b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[K     |████████████████████████████████| 102.9 MB 103 kB/s \n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "# %pip install xformers\n",
        "%pip install xformers==0.0.16rc425\n",
        "\n",
        "# If precompiled wheels don't work, install it with the following command. It will take around 40 minutes to compile.\n",
        "# %pip install git+https://github.com/facebookresearch/xformers@1d31a3a#egg=xformers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G0NV324ZcL9L"
      },
      "source": [
        "## Settings and run"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Name/Path of the initial model.\n",
        "MODEL_NAME = \"/content/NovelAI_SD_Model\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Path for images of the concept for training.\n",
        "INSTANCE_DIR = \"/content/data/sks\" #@param {type:\"string\"}\n",
        "!mkdir -p $INSTANCE_DIR\n",
        "\n",
        "#@markdown A general name for class like dog for dog images.\n",
        "CLASS_NAME = \"girl\" #@param {type:\"string\"}\n",
        "CLASS_DIR = f\"/content/data/{CLASS_NAME}\"\n",
        "\n",
        "#@markdown Enter the directory name to save model at.\n",
        "OUTPUT_DIR = \"/content/stable_diffusion_weights/sks\" #@param {type:\"string\"}\n",
        "\n",
        "print(f\"[*] Weights will be saved at {OUTPUT_DIR}\")\n",
        "\n",
        "!mkdir -p $OUTPUT_DIR\n",
        "\n",
        "#@markdown sks is a rare identifier, feel free to replace it."
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wjEHqo76FWCA",
        "outputId": "3561880b-6b0b-4c7d-8760-5e505c6c23bb"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[*] Weights will be saved at /content/stable_diffusion_weights/sks\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# You can also add multiple concepts here. Try tweaking `--max_train_steps` accordingly.\n",
        "\n",
        "concepts_list = [\n",
        "    {\n",
        "        \"instance_prompt\":\"photo of sks \" +CLASS_NAME,\n",
        "        \"class_prompt\":\"photo of a \" + CLASS_NAME,\n",
        "        \"instance_data_dir\":INSTANCE_DIR,\n",
        "        \"class_data_dir\":CLASS_DIR\n",
        "    },\n",
        "]\n",
        "# `class_data_dir` contains regularization images\n",
        "import json\n",
        "import os\n",
        "for c in concepts_list:\n",
        "    os.makedirs(c[\"instance_data_dir\"], exist_ok=True)\n",
        "\n",
        "with open(\"concepts_list.json\", \"w\") as f:\n",
        "    json.dump(concepts_list, f, indent=4)"
      ],
      "metadata": {
        "id": "bvIa165_KSLp"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 292
        },
        "id": "fe-GgtnUVO_e",
        "outputId": "422e188c-9bf7-412d-ec7b-3ec7de492406"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "     <input type=\"file\" id=\"files-08f6091d-0006-4ede-a611-45a88eed75e0\" name=\"files[]\" multiple disabled\n",
              "        style=\"border:none\" />\n",
              "     <output id=\"result-08f6091d-0006-4ede-a611-45a88eed75e0\">\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> "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Saving 6.jpg to 6.jpg\n",
            "Saving 8.png to 8.png\n",
            "Saving 10.png to 10.png\n",
            "Saving 11.png to 11.png\n",
            "Saving 13.png to 13.png\n",
            "Saving 14.png to 14.png\n",
            "Saving 15.png to 15.png\n"
          ]
        }
      ],
      "source": [
        "#@markdown Upload your images by running this cell.\n",
        "\n",
        "#@markdown OR\n",
        "\n",
        "#@markdown You can use the file manager on the left panel to upload (drag and drop) to INSTANCE_DIR (it uploads faster)\n",
        "\n",
        "import os\n",
        "from google.colab import files\n",
        "import shutil\n",
        "\n",
        "uploaded = files.upload()\n",
        "for filename in uploaded.keys():\n",
        "    dst_path = os.path.join(INSTANCE_DIR, filename)\n",
        "    shutil.move(filename, dst_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qn5ILIyDJIcX"
      },
      "source": [
        "# Start Training\n",
        "\n",
        "Use the table below to choose the best flags based on your memory and speed requirements. Tested on Tesla T4 GPU.\n",
        "\n",
        "\n",
        "| `fp16` | `train_batch_size` | `gradient_accumulation_steps` | `gradient_checkpointing` | `use_8bit_adam` | GB VRAM usage | Speed (it/s) |\n",
        "| ---- | ------------------ | ----------------------------- | ----------------------- | --------------- | ---------- | ------------ |\n",
        "| fp16 | 1                  | 1                             | TRUE                    | TRUE            | 9.92       | 0.93         |\n",
        "| no   | 1                  | 1                             | TRUE                    | TRUE            | 10.08      | 0.42         |\n",
        "| fp16 | 2                  | 1                             | TRUE                    | TRUE            | 10.4       | 0.66         |\n",
        "| fp16 | 1                  | 1                             | FALSE                   | TRUE            | 11.17      | 1.14         |\n",
        "| no   | 1                  | 1                             | FALSE                   | TRUE            | 11.17      | 0.49         |\n",
        "| fp16 | 1                  | 2                             | TRUE                    | TRUE            | 11.56      | 1            |\n",
        "| fp16 | 2                  | 1                             | FALSE                   | TRUE            | 13.67      | 0.82         |\n",
        "| fp16 | 1                  | 2                             | FALSE                   | TRUE            | 13.7       | 0.83          |\n",
        "| fp16 | 1                  | 1                             | TRUE                    | FALSE           | 15.79      | 0.77         |\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Add `--gradient_checkpointing` flag for around 9.92 GB VRAM usage.\n",
        "\n",
        "remove `--use_8bit_adam` flag for full precision. Requires 15.79 GB with `--gradient_checkpointing` else 17.8 GB.\n",
        "\n",
        "remove `--train_text_encoder` flag to reduce memory usage further, degrades output quality."
      ],
      "metadata": {
        "id": "8QvNtCtRF7HH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!accelerate launch train_dreambooth.py \\\n",
        "  --pretrained_model_name_or_path=$MODEL_NAME \\\n",
        "  --output_dir=$OUTPUT_DIR \\\n",
        "  --revision=\"fp16\" \\\n",
        "  --with_prior_preservation --prior_loss_weight=1.0 \\\n",
        "  --seed=1337 \\\n",
        "  --resolution=512 \\\n",
        "  --train_batch_size=1 \\\n",
        "  --train_text_encoder \\\n",
        "  --mixed_precision=\"fp16\" \\\n",
        "  --use_8bit_adam \\\n",
        "  --gradient_accumulation_steps=1 \\\n",
        "  --learning_rate=2e-6 \\\n",
        "  --lr_scheduler=\"constant\" \\\n",
        "  --lr_warmup_steps=0 \\\n",
        "  --num_class_images=50 \\\n",
        "  --sample_batch_size=4 \\\n",
        "  --max_train_steps=2000 \\\n",
        "  --save_interval=10000 \\\n",
        "  --save_sample_prompt=\"photo of sks {CLASS_NAME}\" \\\n",
        "  --concepts_list=\"concepts_list.json\"\n",
        "\n",
        "# Reduce the `--save_interval` to lower than `--max_train_steps` to save weights from intermediate steps.\n",
        "# `--save_sample_prompt` can be same as `--instance_prompt` to generate intermediate samples (saved along with weights in samples directory)."
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sJ5qObOWMKKH",
        "outputId": "869a43b9-a780-4346-eaef-d029b5423b77"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The following values were not passed to `accelerate launch` and had defaults used instead:\n",
            "\t`--num_processes` was set to a value of `1`\n",
            "\t`--num_machines` was set to a value of `1`\n",
            "\t`--mixed_precision` was set to a value of `'no'`\n",
            "\t`--num_cpu_threads_per_process` was set to `1` to improve out-of-box performance\n",
            "To avoid this warning pass in values for each of the problematic parameters or run `accelerate config`.\n",
            "You have passed `None` for safety_checker to disable its functionality in <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'>. Note that this might lead to problems when using <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> and is not recommended.\n",
            "Generating class images: 100% 13/13 [07:10<00:00, 33.15s/it]\n",
            "Caching latents: 100% 50/50 [00:10<00:00,  4.70it/s]\n",
            "Steps: 100% 2000/2000 [31:47<00:00,  1.05it/s, loss=0.24, lr=1e-6]You have passed `None` for safety_checker to disable its functionality in <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'>. Note that this might lead to problems when using <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> and is not recommended.\n",
            "\n",
            "Generating samples:   0% 0/4 [00:00<?, ?it/s]\u001b[A\n",
            "Generating samples:  25% 1/4 [00:12<00:38, 12.76s/it]\u001b[A\n",
            "Generating samples:  50% 2/4 [00:21<00:21, 10.58s/it]\u001b[A\n",
            "Generating samples:  75% 3/4 [00:30<00:09,  9.92s/it]\u001b[A\n",
            "Generating samples: 100% 4/4 [00:40<00:00, 10.01s/it]\n",
            "[*] Weights saved at /content/stable_diffusion_weights/sks/2000\n",
            "Steps: 100% 2000/2000 [32:36<00:00,  1.02it/s, loss=0.24, lr=1e-6]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5V8wgU0HN-Kq"
      },
      "source": [
        "## Convert weights to ckpt to use in web UIs like AUTOMATIC1111."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "baL22PHzOLeP"
      },
      "outputs": [],
      "source": [
        "#@markdown Download script\n",
        "!wget -q https://github.com/ShivamShrirao/diffusers/raw/main/scripts/convert_diffusers_to_original_stable_diffusion.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "89Az5NUxOWdy",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c4c22fe9-9800-48ad-f72a-73a1f43210ee"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Reshaping encoder.mid.attn_1.q.weight for SD format\n",
            "Reshaping encoder.mid.attn_1.k.weight for SD format\n",
            "Reshaping encoder.mid.attn_1.v.weight for SD format\n",
            "Reshaping encoder.mid.attn_1.proj_out.weight for SD format\n",
            "Reshaping decoder.mid.attn_1.q.weight for SD format\n",
            "Reshaping decoder.mid.attn_1.k.weight for SD format\n",
            "Reshaping decoder.mid.attn_1.v.weight for SD format\n",
            "Reshaping decoder.mid.attn_1.proj_out.weight for SD format\n",
            "[*] Converted ckpt saved at /content/model.ckpt\n"
          ]
        }
      ],
      "source": [
        "#@markdown Run conversion.\n",
        "\n",
        "#@markdown Path for save ckpt model.\n",
        "\n",
        "ckpt_path = \"/content/model.ckpt\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown The train steps of the model you want to save.\n",
        "\n",
        "#@markdown Normally fill in the max_train_steps.\n",
        "\n",
        "steps = \"2000\" #@param {type:\"string\"}\n",
        "\n",
        "model_path = OUTPUT_DIR + \"/\" + str(steps)\n",
        "\n",
        "half_arg = \"\"\n",
        "#@markdown  Whether to convert to fp16, takes half the space (2GB), might loose some quality.\n",
        "fp16 = False #@param {type: \"boolean\"}\n",
        "if fp16:\n",
        "    half_arg = \"--half\"\n",
        "!python convert_diffusers_to_original_stable_diffusion.py --model_path $model_path --checkpoint_path $ckpt_path $half_arg\n",
        "print(f\"[*] Converted ckpt saved at {ckpt_path}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AfqXUgnPACqX"
      },
      "source": [
        "### Save Model to Google Drive."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "05x2PKiJ2mS8"
      },
      "outputs": [],
      "source": [
        "!cp \"/content/model.ckpt\" \"/content/drive/MyDrive/model.ckpt\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ToNG4fd_dTbF"
      },
      "source": [
        "## Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "gW15FjffdTID"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch import autocast\n",
        "from diffusers import StableDiffusionPipeline\n",
        "from IPython.display import display\n",
        "\n",
        "# If you want to use previously trained model saved in gdrive, replace this with the full path of model in gdrive\n",
        "model_path = OUTPUT_DIR + \"/\" + str(steps) \n",
        "\n",
        "pipe = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16).to(\"cuda\")\n",
        "g_cuda = None"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oIzkltjpVO_f",
        "outputId": "096d9c39-91b3-423e-aec4-256afb1c2e6b"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x7f82ab3cea10>"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ],
      "source": [
        "#@markdown Can set random seed here for reproducibility.\n",
        "g_cuda = torch.Generator(device='cuda')\n",
        "seed = 1234 #@param {type:\"number\"}\n",
        "g_cuda.manual_seed(seed)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "K6xoHWSsbcS3"
      },
      "outputs": [],
      "source": [
        "#@title Run for generating images.\n",
        "\n",
        "prompt = \"photo of sks girl, digital painting,(best quality), ((masterpiece)), (highres), original, extremely detailed wallpaper, 1girl,beach,watercolor (medium),starry sky ,light blue hair,light blue eyes,(an extremely delicate and beautiful),loli,ahoge,long_hair,\" #@param {type:\"string\"}\n",
        "negative_prompt = \"multiple breasts, (mutated hands and fingers:1.5 ), (long body :1.3), (mutation, poorly drawn :1.2) , black-white, bad anatomy, liquid body, liquid tongue, disfigured, malformed, mutated, anatomical nonsense, text font ui, error, malformed hands, long neck, blurred, lowers, lowres, bad anatomy, bad proportions, bad shadow, uncoordinated body, unnatural body, fused breasts, bad breasts, huge breasts, poorly drawn breasts, extra breasts, liquid breasts, heavy breasts, missing breasts, huge haunch, huge thighs, huge calf, bad hands, fused hand, missing hand, disappearing arms, disappearing thigh, disappearing calf, disappearing legs, fused ears, bad ears, poorly drawn ears, extra ears, liquid ears, heavy ears, missing ears, fused animal ears, bad animal ears, poorly drawn animal ears, extra animal ears, liquid animal ears, heavy animal ears, missing animal ears, text, ui, error, missing fingers, missing limb, fused fingers, one hand with more than 5 fingers, one hand with less than 5 fingers, one hand with more than 5 digit, one hand with less than 5 digit, extra digit, fewer digits, fused digit, missing digit, bad digit, liquid digit, colorful tongue, black tongue, cropped, watermark, username, blurry, JPEG artifacts, signature, 3D, 3D game, 3D game scene, 3D character, malformed feet, extra feet, bad feet, poorly drawn feet, fused feet, missing feet, extra shoes, bad shoes, fused shoes, more than two shoes, poorly drawn shoes, bad gloves, poorly drawn gloves, fused gloves, bad cum, poorly drawn cum, fused cum, bad hairs, poorly drawn hairs, fused hairs, big muscles, ugly, bad face, fused face, poorly drawn face, cloned face, big face, long face, bad eyes, fused eyes poorly drawn eyes, extra eyes, malformed limbs, more than 2 nipples, missing nipples, different nipples, fused nipples, bad nipples, poorly drawn nipples, black nipples, colorful nipples, gross proportions. short arm, (((missing arms))), missing thighs, missing calf, missing legs, mutation, duplicate, morbid, mutilated, poorly drawn hands, more than 1 left hand, more than 1 right hand, deformed, (blurry), disfigured, missing legs, extra arms, extra thighs, more than 2 thighs, extra calf, fused calf, extra legs, bad knee, extra knee, more than 2 legs, bad tails, bad mouth, fused mouth, poorly drawn mouth, bad tongue, tongue within mouth, too long tongue, black tongue, big mouth, cracked mouth, bad mouth, dirty face, dirty teeth, dirty pantie, fused pantie, poorly drawn pantie, fused cloth, poorly drawn cloth, bad pantie, yellow teeth, thick lips, bad cameltoe, colorful cameltoe, bad asshole, poorly drawn asshole, fused asshole, missing asshole, bad anus, bad pussy, bad crotch, bad crotch seam, fused anus, fused pussy, fused anus, fused crotch, poorly drawn crotch, fused seam, poorly drawn anus, poorly drawn pussy, poorly drawn crotch, poorly drawn crotch seam, bad thigh gap, missing thigh gap, fused thigh gap, liquid thigh gap, poorly drawn thigh gap, poorly drawn anus, bad collarbone, fused collarbone, missing collarbone, liquid collarbone, strong girl, obesity, worst quality, low quality, normal quality, liquid tentacles, bad tentacles, poorly drawn tentacles, split tentacles, fused tentacles, missing clit, bad clit, fused clit, colorful clit, black clit, liquid clit, QR code, bar code, censored, safety panties, safety knickers, beard, furry ,pony, pubic hair, mosaic, excrement, faeces, shit, futa, testis,nsfw\" #@param {type:\"string\"}\n",
        "num_samples = 4 #@param {type:\"number\"}\n",
        "guidance_scale = 7.5 #@param {type:\"number\"}\n",
        "num_inference_steps = 50 #@param {type:\"number\"}\n",
        "height = 512 #@param {type:\"number\"}\n",
        "width = 512 #@param {type:\"number\"}\n",
        "\n",
        "with autocast(\"cuda\"), torch.inference_mode():\n",
        "    images = pipe(\n",
        "        prompt,\n",
        "        height=height,\n",
        "        width=width,\n",
        "        negative_prompt=negative_prompt,\n",
        "        num_images_per_prompt=num_samples,\n",
        "        num_inference_steps=num_inference_steps,\n",
        "        guidance_scale=guidance_scale,\n",
        "        generator=g_cuda\n",
        "    ).images\n",
        "\n",
        "for img in images:\n",
        "    display(img)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3.10.7 64-bit",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.7"
    },
    "vscode": {
      "interpreter": {
        "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
      }
    },
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}