EveryDream2trainer/Train_Colab.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/nawnie/EveryDream2trainer/blob/main/Train_Colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "blaLMSbkPHhG"
      },
      "source": [
        "EveryDream2 Colab Edition"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WsYIcz9HY9lx",
        "cellView": "form"
      },
      "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",
      "source": [
        "#@title verify python version\n",
        "!python --version"
      ],
      "metadata": {
        "id": "f2cdMtCt9Wb6",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Optional connect Gdrive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "\n",
        "!mkdir /content/drive/MyDrive/logs/ckpt"
      ],
      "metadata": {
        "id": "d1di4EC6ygw1",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hAuBbtSvGpau",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@markdown # Install Dependencies\n",
        "from IPython.display import clear_output\n",
        "from subprocess import getoutput\n",
        "s = getoutput('nvidia-smi')\n",
        "!pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 --extra-index-url \"https://download.pytorch.org/whl/cu116\"\n",
        "!pip install transformers==4.25.1\n",
        "!pip install diffusers[torch]==0.10.2\n",
        "!pip install pynvml==11.4.1\n",
        "!pip install bitsandbytes==0.35.0\n",
        "!pip install ftfy==6.1.1\n",
        "!pip install aiohttp==3.8.3\n",
        "!pip install tensorboard>=2.11.0\n",
        "!pip install protobuf==3.20.1\n",
        "!pip install wandb==0.13.6\n",
        "!pip install pyre-extensions==0.0.23\n",
        "if \"A100\" in s:\n",
        "  !pip install https://huggingface.co/industriaditat/xformers_precompiles/blob/main/A100_13dev/xformers-0.0.13.dev0-py3-none-any.whl\n",
        "else:\n",
        "  !pip install https://huggingface.co/industriaditat/xformers_precompiles/resolve/main/T4_13dev/xformers-0.0.13.dev0-py3-none-any.whl\n",
        "!pip install pytorch-lightning==1.6.5\n",
        "!pip install OmegaConf==2.2.3\n",
        "!pip install numpy==1.23.5\n",
        "!pip install colorama\n",
        "!pip install 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": {
        "id": "unaffeqGP_0A",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Get A Base Model\n",
        "#@markdown in this space you can type any url however i provide 2 options in the drop down\n",
        "from IPython.display import clear_output\n",
        "!mkdir input\n",
        "%cd /content/EveryDream2trainer\n",
        "!python utils/get_yamls.py\n",
        "!rm train.py\n",
        "!wget \"https://github.com/nawnie/EveryDream2trainer/blob/e605cc59fb1f84f4f8e9290d88f7ede7b4c4fe0c/train_colab.py\"\n",
        "!mkdir Ckpt_cache\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",
        "\n",
        "clear_output()\n",
        "print(\"DONE!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "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"
      ],
      "metadata": {
        "id": "nEzuEYH0536C"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "#@title Convert to diffuers currently only 1.5 based models\n",
        "#@markdown * This Path should be to an ckpt you want to use the one we just downloaded will be inside EveryDream2trainer\n",
        "\n",
        "Base_path = \"/content/EveryDream2trainer/sd_v1-5_vae.ckpt\" #@param {type:\"string\"}\n",
        "#@markdown * the name that you will use when selecting this model in the future traning sessons.\n",
        "save_name = \"SD_1.5\" #@param{type:\"string\"}\n",
        "\n",
        "!python utils/convert_original_stable_diffusion_to_diffusers.py --scheduler_type ddim  \\\n",
        "--original_config_file v1-inference.yaml \\\n",
        "--image_size 512 \\\n",
        "--checkpoint_path \"$Base_path\" \\\n",
        "--prediction_type epsilon \\\n",
        "--upcast_attn False \\\n",
        "--dump_path \"$save_name\"\n"
      ],
      "metadata": {
        "id": "tPvQSo6ScF2c",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Resume from a diffuser model saved to your Gdrive\n",
        "#@markdown * if you have preveiously saved diffusers on your drive you can slect them here\n",
        "Resume_Model = \"\" #@param{type:\"string\"} \n",
        "save_name= Resume_Model"
      ],
      "metadata": {
        "id": "bLpcvpGJB4Gu",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "JXVu-W2lCjwX"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j9pEI69WXS9w",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title \n",
        "#@markdown # Run Everydream 2\n",
        "\n",
        "Save_to_Gdrive = True #@param{type:\"boolean\"}\n",
        "#@markdown * Use resume to contnue training you just ran\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\n",
        "Gradient_checkpointing = True #@param{type:\"boolean\"}\n",
        "#@markdown * Xformers saves ram and offers a great speed up\n",
        "Disable_Xformers = False #@param{type:\"boolean\"}\n",
        "#@markdown * 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 stop the text encoder to attempt to reduce overfitting when resuming an unfinished model\n",
        "Disable_text_Encoder= False #@param{type:\"boolean\"}\n",
        "Project_Name = \"\" #@param{type: 'string'}\n",
        "Max_Epochs = 100 #@param {type:\"slider\", min:0, max:200, step:10}\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. See Advanced Tweaking for more info. Once you have started, the learning rate is a good first knob to turn as you move into more advanced tweaking.\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 * The constant scheduler is the default and keeps your LR set to the value you set in the command line. That's really it for constant! I recommend sticking with it until you are comfortable with general training.\n",
        "\n",
        "Schedule = \"polynomial\" #@param [\"constant\", \"polynomial\", \"linear\", \"cosine\"] {allow-input: true}\n",
        "Resolution = 512#@param{type:\"integer\"} \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 * While very small batch sizes can impact performance negatively, at some point larger sizes have little impact on overall speed as well, so shooting for the moon is not always advisable. Changing batch size may also impact what learning rate you use, with typically larger batch_size requiring a slightly higher learning rate.\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 * 1 or 5 steps will take the same vram as a batch of 1\n",
        "#@markdown * in colab free teir you can expect the fastest proformance from a batch of 4 and a 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 fast Doesnt always mean better nor does slow experement\n",
        "\n",
        "Gradient_steps = 2 #@param{type: 'number'}\n",
        "Dataset_Location = \"/content/drive/MyDrive/Zmanda\" #@param {type:\"string\"}\n",
        "dataset = Dataset_Location\n",
        "model = save_name\n",
        "\n",
        "Save_every_N_epoch = 25#@param{type:\"integer\"}\n",
        "\n",
        "#@markdown You can set your own sample prompts by adding them, one line at a time, to sample_prompts.txt.\n",
        "\n",
        "Steps_between_samples = 30#@param{type:\"integer\"}\n",
        "\n",
        "Drive=\"\"\n",
        "if Save_to_Gdrive:\n",
        "  Drive = \"--logdir /content/drive/MyDrive/logs --save_ckpt_dir /content/drive/MyDrive/logs/ckpt\"\n",
        "\n",
        "if Max_Epochs==0:\n",
        "  Max_Epoch=1\n",
        "\n",
        "if resume:\n",
        "  model = \"findlast\"\n",
        "\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",
        "\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 2 \\\n",
        "  --cond_dropout 0.00 \\\n",
        "  --data_root \"$dataset\" \\\n",
        "  --flip_p 0.01 \\\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",
      "source": [
        "#@title Adv Setup use the .Json to setup your paramaters\n",
        "!python train.py --config train.json"
      ],
      "metadata": {
        "id": "Iuoa_1B9jRGU"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}