{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "#@title 🤗 AutoTrain DreamBooth\n",
        "# @markdown In order to use this colab\n",
        "# @markdown - upload images to a folder named `images/`\n",
        "# @markdown - choose a project name if you wish\n",
        "# @markdown - change model if you wish, you can also select sd2/2.1 or sd1.5\n",
        "# @markdown - update prompt and remember it. choose keywords that don't usually appear in dictionaries\n",
        "# @markdown - add huggingface information (token and repo_id) if you wish to push trained model to huggingface hub\n",
        "# @markdown - update hyperparameters if you wish\n",
        "# @markdown - click `Runtime > Run all` or run each cell individually\n",
        "\n",
        "import os\n",
        "!pip install -U autotrain-advanced > install_logs.txt\n",
        "!autotrain setup > setup_logs.txt"
      ],
      "metadata": {
        "cellView": "code",
        "id": "9iClNdQayIv5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# New Section"
      ],
      "metadata": {
        "id": "BqqPQXPhRgU2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DL29zZgjRtll",
        "outputId": "ced02a72-8dda-48c1-c111-a682dd5e2881"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "A2-_lkBS1WKA",
        "cellView": "code"
      },
      "outputs": [],
      "source": [
        "#@markdown ---\n",
        "#@markdown #### Project Config\n",
        "project_name = 'my_dreambooth_project' # @param {type:\"string\"}\n",
        "model_name = 'stabilityai/stable-diffusion-xl-base-1.0' # @param [\"stabilityai/stable-diffusion-xl-base-1.0\", \"runwayml/stable-diffusion-v1-5\", \"stabilityai/stable-diffusion-2-1\", \"stabilityai/stable-diffusion-2-1-base\"]\n",
        "prompt = 'photo of a nkl person' # @param {type: \"string\"}\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown #### Push to Hub?\n",
        "#@markdown Use these only if you want to push your trained model to a private repo in your Hugging Face Account\n",
        "#@markdown If you dont use these, the model will be saved in Google Colab and you are required to download it manually.\n",
        "#@markdown Please enter your Hugging Face write token. The trained model will be saved to your Hugging Face account.\n",
        "#@markdown You can find your token here: https://huggingface.co/settings/tokens\n",
        "push_to_hub = True # @param [\"False\", \"True\"] {type:\"raw\"}\n",
        "hf_token = \"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\" #@param {type:\"string\"}\n",
        "repo_id = \"NEXAS/stable_diff_custom\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown #### Hyperparameters\n",
        "learning_rate = 1e-4 # @param {type:\"number\"}\n",
        "num_steps = 500 #@param {type:\"number\"}\n",
        "batch_size = 1 # @param {type:\"slider\", min:1, max:32, step:1}\n",
        "gradient_accumulation = 4 # @param {type:\"slider\", min:1, max:32, step:1}\n",
        "resolution = 1024 # @param {type:\"slider\", min:128, max:1024, step:128}\n",
        "use_8bit_adam = True # @param [\"False\", \"True\"] {type:\"raw\"}\n",
        "use_xformers = True # @param [\"False\", \"True\"] {type:\"raw\"}\n",
        "use_fp16 = True # @param [\"False\", \"True\"] {type:\"raw\"}\n",
        "train_text_encoder = False # @param [\"False\", \"True\"] {type:\"raw\"}\n",
        "gradient_checkpointing = True # @param [\"False\", \"True\"] {type:\"raw\"}\n",
        "os.environ[\"PROJECT_NAME\"] = project_name\n",
        "os.environ[\"MODEL_NAME\"] = model_name\n",
        "os.environ[\"PROMPT\"] = prompt\n",
        "os.environ[\"PUSH_TO_HUB\"] = str(push_to_hub)\n",
        "os.environ[\"HF_TOKEN\"] = hf_token\n",
        "os.environ[\"REPO_ID\"] = repo_id\n",
        "os.environ[\"LEARNING_RATE\"] = str(learning_rate)\n",
        "os.environ[\"NUM_STEPS\"] = str(num_steps)\n",
        "os.environ[\"BATCH_SIZE\"] = str(batch_size)\n",
        "os.environ[\"GRADIENT_ACCUMULATION\"] = str(gradient_accumulation)\n",
        "os.environ[\"RESOLUTION\"] = str(resolution)\n",
        "os.environ[\"USE_8BIT_ADAM\"] = str(use_8bit_adam)\n",
        "os.environ[\"USE_XFORMERS\"] = str(use_xformers)\n",
        "os.environ[\"USE_FP16\"] = str(use_fp16)\n",
        "os.environ[\"TRAIN_TEXT_ENCODER\"] = str(train_text_encoder)\n",
        "os.environ[\"GRADIENT_CHECKPOINTING\"] = str(gradient_checkpointing)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "g3cd_ED_yXXt"
      },
      "outputs": [],
      "source": [
        "!autotrain dreambooth \\\n",
        "--model ${MODEL_NAME} \\\n",
        "--output ${PROJECT_NAME} \\\n",
        "--image-path images/ \\\n",
        "--prompt \"${PROMPT}\" \\\n",
        "--resolution ${RESOLUTION} \\\n",
        "--batch-size ${BATCH_SIZE} \\\n",
        "--num-steps ${NUM_STEPS} \\\n",
        "--gradient-accumulation ${GRADIENT_ACCUMULATION} \\\n",
        "--lr ${LEARNING_RATE} \\\n",
        "$( [[ \"$USE_FP16\" == \"True\" ]] && echo \"--fp16\" ) \\\n",
        "$( [[ \"$USE_XFORMERS\" == \"True\" ]] && echo \"--xformers\" ) \\\n",
        "$( [[ \"$TRAIN_TEXT_ENCODER\" == \"True\" ]] && echo \"--train-text-encoder\" ) \\\n",
        "$( [[ \"$USE_8BIT_ADAM\" == \"True\" ]] && echo \"--use-8bit-adam\" ) \\\n",
        "$( [[ \"$GRADIENT_CHECKPOINTING\" == \"True\" ]] && echo \"--gradient-checkpointing\" ) \\\n",
        "$( [[ \"$PUSH_TO_HUB\" == \"True\" ]] && echo \"--push-to-hub --hub-token ${HF_TOKEN} --hub-model-id ${REPO_ID}\" )"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Inference\n",
        "# this is the inference code that you can use after you have trained your model\n",
        "# Unhide code below and change prj_path to your repo or local path (e.g. my_dreambooth_project)\n",
        "#\n",
        "#\n",
        "#\n",
        "from diffusers import DiffusionPipeline\n",
        "import torch\n",
        "\n",
        "prj_path = \"/content/my_dreambooth_project\"\n",
        "model = \"stabilityai/stable-diffusion-xl-base-1.0\"\n",
        "pipe = DiffusionPipeline.from_pretrained(\n",
        "    model,\n",
        "    torch_dtype=torch.float16,\n",
        ")\n",
        "pipe.to(\"cuda\")\n",
        "pipe.load_lora_weights(prj_path, weight_name=\"pytorch_lora_weights.safetensors\")\n",
        "\n",
        "prompt = \"potrait Photo of a nkl person, in anime style\"\n",
        "\n",
        "for seed in range(5):\n",
        "    generator = torch.Generator(\"cuda\").manual_seed(seed)\n",
        "    image = pipe(prompt=prompt, generator=generator,num_inference_steps=25).images[0]\n",
        "    image.save(f\"images/{seed}.png\")"
      ],
      "metadata": {
        "id": "L2Zn_1Knlmgs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Inference\n",
        "# this is the inference code that you can use after you have trained your model\n",
        "# Unhide code below and change prj_path to your repo or local path (e.g. my_dreambooth_project)\n",
        "#\n",
        "#\n",
        "#\n",
        "from diffusers import DiffusionPipeline,StableDiffusionXLImg2ImgPipeline\n",
        "import torch\n",
        "\n",
        "prj_path = \"NEXAS/stable_diff_custom\"\n",
        "model = \"stabilityai/stable-diffusion-xl-base-1.0\"\n",
        "pipe = DiffusionPipeline.from_pretrained(\n",
        "     model,\n",
        "     torch_dtype=torch.float16,\n",
        "     )\n",
        "pipe.to(\"cuda\")\n",
        "pipe.load_lora_weights(prj_path, weight_name=\"pytorch_lora_weights.safetensors\")\n",
        "\n",
        "refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(\n",
        "     \"stabilityai/stable-diffusion-xl-refiner-1.0\",\n",
        "     torch_dtype=torch.float16,\n",
        " )\n",
        "refiner.to(\"cuda\")\n",
        "\n",
        "prompt = \"photo of a nkl person,in a black suit 4k\"\n",
        "\n",
        "seed = 42\n",
        "generator = torch.Generator(\"cuda\").manual_seed(seed)\n",
        "image = pipe(prompt=prompt, generator=generator).images[0]\n",
        "image = refiner(prompt=prompt, generator=generator, image=image).images[0]\n",
        "image.save(f\"generated_image.png\")"
      ],
      "metadata": {
        "id": "M8i_ae_obcGe"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}