AutoTrain_Dreambooth.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "id": "JvMRbVLEJlZT",
        "outputId": "aa97ec7b-af8c-4072-8885-2aa36a1a3cdf"
      },
      "outputs": [],
      "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) 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",
        "#@markdown - report issues / feature requests here: https://github.com/huggingface/autotrain-advanced/issues\n",
        "\n",
        "import os\n",
        "!pip install -U autotrain-advanced > install_logs.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "A2-_lkBS1WKA"
      },
      "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 sks dog' # @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 = False # @param [\"False\", \"True\"] {type:\"raw\"}\n",
        "hf_token = \"hf_XXX\" #@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 = False # @param [\"False\", \"True\"] {type:\"raw\"}\n",
        "use_xformers = False # @param [\"False\", \"True\"] {type:\"raw\"}\n",
        "mixed_precision = \"fp16\" # @param [\"fp16\", \"bf16\", \"none\"] {type:\"raw\"}\n",
        "train_text_encoder = False # @param [\"False\", \"True\"] {type:\"raw\"}\n",
        "disable_gradient_checkpointing = False # @param [\"False\", \"True\"] {type:\"raw\"}\n",
        "\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[\"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[\"MIXED_PRECISION\"] = str(mixed_precision)\n",
        "os.environ[\"TRAIN_TEXT_ENCODER\"] = str(train_text_encoder)\n",
        "os.environ[\"DISABLE_GRADIENT_CHECKPOINTING\"] = str(disable_gradient_checkpointing)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "id": "g3cd_ED_yXXt",
        "outputId": "f8e88e16-07c4-463c-9f6c-f0470a0288a2"
      },
      "outputs": [],
      "source": [
        "!autotrain dreambooth \\\n",
        "--model ${MODEL_NAME} \\\n",
        "--project-name ${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",
        "--mixed-precision ${MIXED_PRECISION} \\\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",
        "$( [[ \"$DISABLE_GRADIENT_CHECKPOINTING\" == \"True\" ]] && echo \"--disable_gradient-checkpointing\" ) \\\n",
        "$( [[ \"$PUSH_TO_HUB\" == \"True\" ]] && echo \"--push-to-hub --token ${HF_TOKEN}\" )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_LvIS7-7PcLT"
      },
      "outputs": [],
      "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 = \"username/repo_name\"\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 sks dog in a bucket\"\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": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}