EveryDream2trainer/Train_Colab.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/victorchall/EveryDream2trainer/blob/main/Train_Colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "blaLMSbkPHhG"
      },
      "source": [
        "# EveryDream2 Colab Edition\n",
        "\n",
        "Check out documentation here: https://github.com/victorchall/EveryDream2trainer#docs\n",
        "\n",
        "And join the discord: https://discord.gg/uheqxU6sXN"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "WsYIcz9HY9lx"
      },
      "outputs": [],
      "source": [
        "#@title # Install python 3.10 \n",
        "import os\n",
        "from IPython.display import clear_output\n",
        "!wget https://github.com/korakot/kora/releases/download/v0.10/py310.sh\n",
        "!bash ./py310.sh -b -f -p /usr/local\n",
        "!python -m ipykernel install --name \"py310\" --user\n",
        "clear_output()\n",
        "os.kill(os.getpid(), 9)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "f2cdMtCt9Wb6"
      },
      "outputs": [],
      "source": [
        "#@title Verify python version, should be 3.10.something\n",
        "!python --version"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "d1di4EC6ygw1"
      },
      "outputs": [],
      "source": [
        "#@title Optional connect Gdrive\n",
        "#@markdown # but strongly recommended\n",
        "#@markdown This will let you put all your training data and checkpoints directly on your drive.  Much faster/easier to continue later, less setup time.\n",
        "\n",
        "#@markdown Creates /content/drive/MyDrive/everydreamlogs/ckpt\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "!mkdir -p /content/drive/MyDrive/everydreamlogs/ckpt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "hAuBbtSvGpau"
      },
      "outputs": [],
      "source": [
        "#@markdown # Install Dependencies\n",
        "#@markdown This will take a couple minutes, be patient and watch the output for \"DONE!\"\n",
        "from IPython.display import clear_output\n",
        "from subprocess import getoutput\n",
        "s = getoutput('nvidia-smi')\n",
        "!pip install -q torch==1.12.1+cu116 torchvision==0.13.1+cu116 --extra-index-url \"https://download.pytorch.org/whl/cu116\"\n",
        "!pip install -q transformers==4.25.1\n",
        "!pip install -q diffusers[torch]==0.10.2\n",
        "!pip install -q pynvml==11.4.1\n",
        "!pip install -q bitsandbytes==0.35.0\n",
        "!pip install -q ftfy==6.1.1\n",
        "!pip install -q aiohttp==3.8.3\n",
        "!pip install -q tensorboard>=2.11.0\n",
        "!pip install -q protobuf==3.20.1\n",
        "!pip install -q wandb==0.13.6\n",
        "!pip install -q pyre-extensions==0.0.23\n",
        "if \"A100\" in s:\n",
        "  !pip install -q https://huggingface.co/industriaditat/xformers_precompiles/blob/main/A100_13dev/xformers-0.0.13.dev0-py3-none-any.whl\n",
        "else:\n",
        "  !pip install -q https://huggingface.co/industriaditat/xformers_precompiles/resolve/main/T4_13dev/xformers-0.0.13.dev0-py3-none-any.whl\n",
        "!pip install -q pytorch-lightning==1.6.5\n",
        "!pip install -q OmegaConf==2.2.3\n",
        "!pip install -q numpy==1.23.5\n",
        "!pip install -q colorama\n",
        "!pip install -q keyboard\n",
        "clear_output()\n",
        "!git clone https://github.com/victorchall/EveryDream2trainer.git\n",
        "%cd /content/EveryDream2trainer\n",
        "!wget \"https://raw.githubusercontent.com/nawnie/EveryDream2trainer/main/train_colab.py\"\n",
        "clear_output()\n",
        "print(\"DONE!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "unaffeqGP_0A"
      },
      "outputs": [],
      "source": [
        "#@title Get A Base Model\n",
        "#@markdown Choose SD1.5 or Waifu Diffusion 1.3 from the dropdown, or paste your own URL in the box\n",
        "\n",
        "#@markdown If you already did this once with Gdrive connected, you can skip this step as the cached copy is on your gdrive\n",
        "from IPython.display import clear_output\n",
        "!mkdir input\n",
        "%cd /content/EveryDream2trainer\n",
        "!python utils/get_yamls.py\n",
        "MODEL_URL = \"https://huggingface.co/panopstor/EveryDream/resolve/main/sd_v1-5_vae.ckpt\" #@param [\"https://huggingface.co/panopstor/EveryDream/resolve/main/sd_v1-5_vae.ckpt\", \"https://huggingface.co/hakurei/waifu-diffusion-v1-3/resolve/main/wd-v1-3-float16.ckpt\"] {allow-input: true}\n",
        "print(\"Downloading \")\n",
        "!wget $MODEL_URL\n",
        "\n",
        "%cd /content/EveryDream2trainer\n",
        "\n",
        "clear_output()\n",
        "print(\"DONE!\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "nEzuEYH0536C"
      },
      "source": [
        "In order to train, you need a base model on which to train. This is a one-time setup to configure base models when you want to use a particular base.  \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tPvQSo6ScF2c"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "#@title Setup conversion\n",
        "\n",
        "#@markdown **If you already did this once with Gdrive connected, you can skip this step as the cached copy is on your gdrive.**  \n",
        "# \n",
        "# If you are not sure, look in your Gdrive for `everydreamlogs/ckpt` and see if you have a folder with the `save_name` below.\n",
        "\n",
        "#@markdown Pick the `model_type` in the dropdown. This is the model type that you are converting and you downloaded above.  This is important as it will determine the model architecture and the correct settings to use.\n",
        "\n",
        "#@markdown * `SD1x` is all SD1.x based models *(SD1.4, SD1.5, Waifu Diffusion 1.3, etc)*\n",
        "\n",
        "#@markdown * `SD2_512_base` is the SD2 512 base model\n",
        "\n",
        "#@markdown * `SD21` is all SD2 768 models.  *(ex. SD2.1 768, or trained models based on that)*\n",
        "\n",
        "#@markdown If you are not sure, double check the model author's page or ask for help on [Discord](https://discord.gg/uheqxU6sXN).\n",
        "model_type = \"SD1x\" #@param [\"SD1x\", \"SD2_512_base\", \"SD21\"]\n",
        "\n",
        "#@markdown This is the temporary ckpt file that was downloaded above. If you downloaded a different model, you can change this.  *Hint: look at your file manager in the EveryDream2trainer folder for .ckpt files*.\n",
        "base_path = \"/content/EveryDream2trainer/sd_v1-5_vae.ckpt\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown The name that you will use when selecting this model in the future training sessons.\n",
        "save_name = \"SD15\" #@param{type:\"string\"}\n",
        "\n",
        "#@markdown If you are using Gdrive, this will save the converted model to your Gdrive for future use so you can skip downloading and converting the model.\n",
        "cache_to_gdrive = True #@param{type:\"boolean\"}\n",
        "\n",
        "if cache_to_gdrive:\n",
        "    save_name = os.path.join(\"/content/drive/MyDrive/everydreamlogs/ckpt\", save_name)\n",
        "\n",
        "img_size = 512\n",
        "upscale_attention = False\n",
        "if model_type == \"SD1x\":\n",
        "    inference_yaml = \"v1-inference.yaml\"\n",
        "elif model_type == \"SD2_512_base\":\n",
        "    upscale_attention = True\n",
        "    inference_yaml = \"v2-inference.yaml\"\n",
        "elif model_type == \"SD21\":\n",
        "    upscale_attention = True\n",
        "    inference_yaml = \"v2-inference-v.yaml\"\n",
        "    img_size = 768\n",
        "\n",
        "print(base_path)\n",
        "print(inference_yaml)\n",
        "\n",
        "!python utils/convert_original_stable_diffusion_to_diffusers.py --scheduler_type ddim  \\\n",
        "--original_config_file {inference_yaml} \\\n",
        "--image_size {img_size} \\\n",
        "--checkpoint_path {base_path} \\\n",
        "--prediction_type epsilon \\\n",
        "--upcast_attn False \\\n",
        "--dump_path {save_name}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "bLpcvpGJB4Gu"
      },
      "outputs": [],
      "source": [
        "#@title Pick your base model from a diffusers model saved to your Gdrive (converted above)\n",
        "\n",
        "#@markdown Do not skip this cell.\n",
        "\n",
        "#@markdown * If you have preveiously saved diffusers on your drive you can select it here\n",
        "\n",
        "#@markdown ex. */content/drive/MyDrive/everydreamlogs/myproject_202208/ckpts/interrupted-gs023*\n",
        "\n",
        "#@markdown The default for SD1.5 converted above would be */content/drive/MyDrive/everydreamlogs/ckpt/SD15*\n",
        "Resume_Model = \"/content/drive/MyDrive/everydreamlogs/ckpt/SD15\" #@param{type:\"string\"} \n",
        "save_name = Resume_Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JXVu-W2lCjwX"
      },
      "source": [
        "For a more indepth Explanation of each of these paramaters check out /content/EveryDream2trainer/doc.\n",
        "\n",
        "\n",
        "After youve tried a few models you will find /content/EveryDream2trainer/doc/ATWEAKING.md to be extremly helpful."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "j9pEI69WXS9w"
      },
      "outputs": [],
      "source": [
        "#@title \n",
        "#@markdown # Run Everydream 2\n",
        "#@markdown If you want to use a .json config or upload your own, skip this cell and run the cell below instead\n",
        "\n",
        "#@markdown * Save logs and output ckpts to Gdrive (strongly suggested)\n",
        "Save_to_Gdrive = True #@param{type:\"boolean\"}\n",
        "#@markdown * Use resume to contnue training you just ran, will also find latest diffusers log in your Gdrive to continue.\n",
        "resume = False #@param{type:\"boolean\"}\n",
        "#@markdown * Checkpointing Saves Vram to allow larger batch sizes minor slow down on a single batch size but will can allow room for a higher traning resolution (suggested on Colab Free tier, turn off for A100)\n",
        "Gradient_checkpointing = True #@param{type:\"boolean\"}\n",
        "Disable_Xformers = False\n",
        "#@markdown * Tag shuffling, mainly for booru training. Best to just read this if interested in shufflng tags /content/EveryDream2trainer/doc/SHUFFLING_TAGS.md\n",
        "shuffle_tags = False #@param{type:\"boolean\"}\n",
        "#@markdown * You can turn off the text encoder training (generally not suggested)\n",
        "Disable_text_Encoder= False #@param{type:\"boolean\"}\n",
        "#@markdown * Name your project so you can find it in your logs\n",
        "Project_Name = \"my_project\" #@param{type: 'string'}\n",
        "\n",
        "#@markdown * The learning rate affects how much \"training\" is done on the model per training step. It is a very careful balance to select a value that will learn your data and not wreck the model. \n",
        "#@markdown Leave this default unless you are very comfortable with training and know what you are doing.\n",
        "\n",
        "Learning_Rate = 1e-6 #@param{type: 'number'}\n",
        "\n",
        "#@markdown * A learning rate scheduler can change your learning rate as training progresses.\n",
        "\n",
        "#@markdown I recommend sticking with constant until you are comfortable with general training. \n",
        "\n",
        "Schedule = \"constant\" #@param [\"constant\", \"polynomial\", \"linear\", \"cosine\"] {allow-input: true}\n",
        "\n",
        "#@markdown * Resolution to train at (recommend 512).  Higher resolution will require lower batch size (below).\n",
        "Resolution = 512 #@param {type:\"slider\", min:256, max:768, step:64}\n",
        "\n",
        "#@markdown * Batch size is also another \"hyperparamter\" of itself and there are tradeoffs. It may not always be best to use the highest batch size possible. Once of the primary reasons to change it is if you get \"CUDA out of memory\" errors where lowering the value may help.\n",
        "\n",
        "#@markdown * Batch size impacts VRAM use.  4 should work on SD1.x models and 3 for SD2.x models at 512 resolution.  Lower this if you get CUDA out of memory errors.\n",
        "\n",
        "Batch_Size = 4 #@param{type: 'number'}\n",
        "\n",
        "#@markdown * Gradient accumulation is sort of like a virtual batch size increase use this to increase batch size with out increasing vram usage\n",
        "#@markdown * Increasing this will not have much impact on VRAM use.\n",
        "#@markdown * In colab free teir you can expect the fastest proformance from a batch of 4 and a gradient step of 2 giving us a total batch size of 8 at 512 resolution \n",
        "#@markdown * Due to bucketng you may need to decresse batch size to 3\n",
        "#@markdown * Remember more gradient accumulation (or batch size) doesn't automatically mean better\n",
        "\n",
        "Gradient_steps = 1 #@param{type:\"slider\", min:1, max:10, step:1}\n",
        "\n",
        "#@markdown * Location on your Gdrive where your training images are.\n",
        "Dataset_Location = \"/content/drive/MyDrive/training_samples\" #@param {type:\"string\"}\n",
        "dataset = Dataset_Location\n",
        "model = save_name\n",
        "\n",
        "#@markdown * Max Epochs to train for, this defines how many total times all your training data is used.\n",
        "\n",
        "Max_Epochs = 100 #@param {type:\"slider\", min:0, max:200, step:5}\n",
        "\n",
        "#@markdown * How often to save checkpoints.\n",
        "Save_every_N_epoch = 20 #@param{type:\"integer\"}\n",
        "\n",
        "#@markdown * Test sample generation steps, how often to generate samples during training.\n",
        "\n",
        "#@markdown You can set your own sample prompts by adding them, one line at a time, to `/content/EveryDream2trainer/sample_prompts.txt`.  If left empty, it will use the captions from your training images.\n",
        "\n",
        "#@markdown Use the steps_between_samples to set how often the samples are generated.\n",
        "Steps_between_samples = 300 #@param{type:\"integer\"}\n",
        "\n",
        "#@markdown * That's it!  Run the cell!\n",
        "\n",
        "Drive=\"\"\n",
        "if Save_to_Gdrive:\n",
        "  Drive = \"--logdir /content/drive/MyDrive/everydreamlogs --save_ckpt_dir /content/drive/MyDrive/everydreamlogs/ckpt\"\n",
        "\n",
        "if Max_Epochs==0:\n",
        "  Max_Epoch=1\n",
        "\n",
        "if resume:\n",
        "  model = \"findlast\"\n",
        "\n",
        "Gradient = \"\"\n",
        "if Gradient_checkpointing:\n",
        "  Gradient = \"--gradient_checkpointing \"\n",
        "if \"A100\" in s:\n",
        "  Gradient = \"\"\n",
        "\n",
        "DX = \"\"  \n",
        "if Disable_Xformers:\n",
        "  DX = \"--disable_xformers \"\n",
        "\n",
        "shuffle = \"\"\n",
        "if shuffle_tags:\n",
        "  shuffle = \"--shuffle_tags \"\n",
        "\n",
        "textencode = \"\"\n",
        "if Disable_text_Encoder:\n",
        "  textencode = \"--disable_textenc_training Train_text \"\n",
        "\n",
        "!python train_colab.py --resume_ckpt \"$model\" \\\n",
        "  $textencode \\\n",
        "  $Gradient \\\n",
        "  $shuffle \\\n",
        "  $Drive \\\n",
        "  $DX \\\n",
        "  --amp \\\n",
        "  --batch_size $Batch_Size \\\n",
        "  --grad_accum $Gradient_steps \\\n",
        "  --cond_dropout 0.00 \\\n",
        "  --data_root \"$dataset\" \\\n",
        "  --flip_p 0.00 \\\n",
        "  --lr $Learning_Rate \\\n",
        "  --lr_decay_steps 0 \\\n",
        "  --lr_scheduler \"$Schedule\" \\\n",
        "  --lr_warmup_steps 0 \\\n",
        "  --max_epochs $Max_Epochs \\\n",
        "  --project_name \"$Project_Name\" \\\n",
        "  --resolution $Resolution \\\n",
        "  --sample_prompts \"sample_prompts.txt\" \\\n",
        "  --sample_steps $Steps_between_samples \\\n",
        "  --save_every_n_epoch $Save_every_N_epoch \\\n",
        "  --seed 555 \\\n",
        "  --shuffle_tags \\\n",
        "  --useadam8bit \\\n",
        "  --notebook\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Iuoa_1B9jRGU"
      },
      "outputs": [],
      "source": [
        "#@title Alternate startup script\n",
        "#@markdown * Edit train.json to setup your paramaters\n",
        "#@markdown * Edit chain0.json to make use of chaining\n",
        "#@markdown * make sure to check each confguration you will need 1 Json per chain length 3 are provided\n",
        "\n",
        "\n",
        "%cd /content/EveryDream2trainer\n",
        "Chain_Length=0 #@param{type:\"integer\"}\n",
        "l = Chain_Length \n",
        "I=0 #repeat counter\n",
        "if l == None or l == 0:\n",
        "  l=1\n",
        "while l > 0:\n",
        "  !python train_colab.py --config chain{I}.json\n",
        "  l -= 1\n",
        "  I =+ 1"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "include_colab_link": true,
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "venv",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.5 (tags/v3.10.5:f377153, Jun  6 2022, 16:14:13) [MSC v.1929 64 bit (AMD64)]"
    },
    "vscode": {
      "interpreter": {
        "hash": "e602395b73d27e246c3f66de86a1ed4dc1e5a85e8356fd1a2f027b9d2f1f8162"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}