fine_tuning_openai_diffusion_model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "fine-tuning openai diffusion model.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "avkq78LOjVhV",
        "3Ol_0nghwwGC",
        "1NZ2Yi2CxITo",
        "CiMreX-_n6Kz"
      ],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Sxela/guided-diffusion-sxela/blob/main/fine_tuning_openai_diffusion_model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "A simple colab to fine-tune openai diffusion models.\n",
        "\n",
        "\n",
        "Feel free to ask questions in this post's comments: https://www.patreon.com/posts/66246423\n",
        "\n",
        "by [Alex Spirin](https://twitter.com/devdef)\n",
        "\n",
        "![visitors](https://visitor-badge.glitch.me/badge?page_id=sxela_finetune_openai_colab)"
      ],
      "metadata": {
        "id": "3hBAjQO1kiEW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup (run once per session)\n",
        "\n",
        "This mounts your google drive for easier storage"
      ],
      "metadata": {
        "id": "ufaUo7olwoF0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "id": "EtMv2MEzSzjN",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "29947e74-992a-42a0-d21c-78a138be6014"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This downloads the training code and installs it"
      ],
      "metadata": {
        "id": "Eg3mlCMIe1B6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content\n",
        "!git clone https://github.com/Sxela/guided-diffusion-sxela\n",
        "%cd /content/guided-diffusion-sxela\n",
        "!pip install -e ."
      ],
      "metadata": {
        "id": "yXvOPC8PfnUG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Train (tune) BEDROOM model :D\n",
        "Needs 16gb GPU RAM\n",
        "\n",
        "Works in colab pro and on kaggle "
      ],
      "metadata": {
        "id": "JZ8BNzApp_Xk"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download a pre-trained LSUN BEDROOM model that we will be tuning on our dataset"
      ],
      "metadata": {
        "id": "RmI7jtj5fzJJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://openaipublic.blob.core.windows.net/diffusion/march-2021/lsun_uncond_100M_1200K_bs128.pt -P /content/"
      ],
      "metadata": {
        "id": "h-fL3fb8wpxZ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6880b17d-e2c7-43ef-90d6-9c74b71d6e2f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "28.pt                25%[====>               ] 110.18M  15.9MB/s    eta 28s    ^C\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Tune"
      ],
      "metadata": {
        "id": "OV2gIxZhw2me"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "For gigachads. \n",
        "We're going to do what's called a pro-gamer move (or not): tune a small model, trained on bedrooms, on our own dataset. Just because we can and it's much faster than training from scratch. \n",
        "\n",
        "Don't forget to change the paths:\n",
        "You need to change DATASET_PATH to point to your dataset images folder, and CHECKPOINT_PATH - to point to a folder you'd like to save progress to. \n",
        "\n",
        "For, example here /content/drive/MyDrive/deep_learning/guided-diffusion-sxela/ - this path points to a location, where all the training checkpoints will be saved\n",
        "\n",
        "and /content/YourDatasetHere/ - this path points to your dataset, i.e. a folder with images (no captions needed)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "We will be using this model together with CLIP inside DiscoDiffusion, so we can train less, stop early and let CLIP do the heavy lifting.\n",
        "\n",
        "This will run almost forever, but you should start checking your results at around ~50k iterations. Good results begin to appear at 100-200k iterations, depending on your dataset.\n",
        "\n",
        "Validating means opening your CHECKPOINT_PATH folder, taking the ema_0.9999_(some number of steps).pt file with the highest number (the latest one), going to this version of DiscoDiffusion here \n",
        "https://github.com/Sxela/DiscoDiffusion-Warp/blob/main/Disco_Diffusion_v5_2_Warp_custom_model.ipynb and setting this: diffusion-model - custom, custom_path - path to that ema file from the previous step (if you saved it on google drive - then just point it there), and set width_height to 256x256, then run DD as usual\n"
      ],
      "metadata": {
        "id": "WqBBkqPjqESf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "MODEL_FLAGS=\"--image_size 256 --num_channels 128 --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16\"\n",
        "DIFFUSION_FLAGS=\"--diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --use_scale_shift_norm False\"\n",
        "TRAIN_FLAGS=\"--lr 2e-5 --batch_size 4 --save_interval 2000 --log_interval 50 --resume_checkpoint /content/lsun_uncond_100M_1200K_bs128.pt\"\n",
        "DATASET_PATH=\"/content/YourDatasetHere/\" #change to point to your dataset path \n",
        "OUTPUT_PATH=\"/content/drive/MyDrive/deep_learning/guided-diffusion-sxela/\" #models will be saved here, change to your drive folder or else\n",
        "%cd /content/guided-diffusion-sxela\n",
        "!python scripts/image_train.py --data_dir $DATASET_PATH $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS --logdir $OUTPUT_PATH\n",
        "\n",
        "#if you are using vanilla openai repo, then you will ned to run this: \n",
        "#!OPENAI_LOGDIR=$OUTPUT_PATH python scripts/image_train.py --data_dir $DATASET_PATH $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS"
      ],
      "metadata": {
        "id": "apH5i0hTqz1y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Sampling\n",
        "The best way to sample your model in real-life conditions is to plug it into DiscoDiffusion. \n",
        "\n",
        "\n",
        "Grab your latest ema checkpoint, open this colab here - https://github.com/Sxela/DiscoDiffusion-Warp/blob/main/Disco_Diffusion_v5_2_Warp_custom_model.ipynb\n",
        "\n",
        "and change model settings > custom model path to your ema checkpoint's location, as described in the previous cell."
      ],
      "metadata": {
        "id": "udICgtfHEiQn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can still sample using vanilla openai code, just plug your checkpoint in the cell below\n",
        "\n",
        "Don't forget to change all the paths"
      ],
      "metadata": {
        "id": "57cMKNlWF1VY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "checkpoint_path = 'input some checkpoint path here' #use ema checkpoint\n",
        "OUTPUT_PATH=\"/content/drive/MyDrive/deep_learning/guided-diffusion-sxela/\"\n",
        "!python scripts/image_sample.py --num_samples 1 --model_path $checkpoint_path $MODEL_FLAGS $DIFFUSION_FLAGS --timestep_respacing ddim100 --logdir $OUTPUT_PATH\n",
        "\n",
        "#if you are using vanilla openai repo, then you will ned to run this: \n",
        "#!OPENAI_LOGDIR=/content/drive/MyDrive/deep_learning/guided-diffusion-sxela/samples/  python scripts/image_sample.py --num_samples 1 --model_path $checkpoint_path $MODEL_FLAGS $DIFFUSION_FLAGS --timestep_respacing ddim100"
      ],
      "metadata": {
        "id": "O-RCVDtuGArz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import PIL\n",
        "\n",
        "sample_path = 'some sample path'\n",
        "im = np.load(sample_path)\n",
        "PIL.Image.fromarray(im.f.arr_0[0])"
      ],
      "metadata": {
        "id": "nFPy3r8AGEW8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Train (tune) 256x256 vanilla DD model\n",
        "Only if you have a beefy GPU with more than 16gb RAM\n",
        "\n",
        "For lvl 50 AI bosses, \n",
        "Will not fit into colab pro, only in colab pro+ with A100 gpu\n"
      ],
      "metadata": {
        "id": "avkq78LOjVhV"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download a pre-trained openai 256x256 model (the one used in DiscoDiffusion) that we will be tuning on our dataset"
      ],
      "metadata": {
        "id": "sk2sVBAxwurI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#download model checkpoint\n",
        "!wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt -P /content/\n",
        "#if you wish to tune the 512x512 finetuned model from DD, you need to download it and change image size and checkpoint path later here: \n",
        "#!wget https://huggingface.co/lowlevelware/512x512_diffusion_unconditional_ImageNet/resolve/main/512x512_diffusion_uncond_finetune_008100.pt"
      ],
      "metadata": {
        "id": "Jjf4ZopAwwyo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Tune"
      ],
      "metadata": {
        "id": "3Ol_0nghwwGC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Don't forget to change the paths:\n",
        "You need to change DATASET_PATH to point to your dataset images folder, and CHECKPOINT_PATH - to point to a folder you'd like to save progress to. \n",
        "\n",
        "For, example here /content/drive/MyDrive/deep_learning/guided-diffusion-sxela/ - this path points to a location, where all the training checkpoints will be saved\n",
        "\n",
        "and /content/YourDatasetHere/ - this path points to your dataset, i.e. a folder with images (no captions needed)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "We will be using this model together with CLIP inside DiscoDiffusion, so we can train less, stop early and let CLIP do the heavy lifting.\n",
        "\n",
        "This will run almost forever, but you should start checking your results at around ~50k iterations. Good results begin to appear at 100-200k iterations, depending on your dataset.\n",
        "\n",
        "Validating means opening your CHECKPOINT_PATH folder, taking the ema_0.9999_(some number of steps).pt file with the highest number (the latest one), going to this version of DiscoDiffusion here \n",
        "https://github.com/Sxela/DiscoDiffusion-Warp/blob/main/Disco_Diffusion_v5_2_Warp_custom_model.ipynb and setting this: diffusion-model - custom, custom_path - path to that ema file from the previous step (if you saved it on google drive - then just point it there), \n",
        "\n",
        "you'll need to set custom model settings to this: \n",
        "\n",
        "    model_config.update({\n",
        "        'attention_resolutions': '32, 16, 8',\n",
        "        'class_cond': False,\n",
        "        'diffusion_steps': diffusion_steps,\n",
        "        'rescale_timesteps': True,\n",
        "        'timestep_respacing': timestep_respacing,\n",
        "        'image_size': 256,\n",
        "        'learn_sigma': True,\n",
        "        'noise_schedule': 'linear',\n",
        "        'num_channels': 256,\n",
        "        'num_head_channels': 64,\n",
        "        'num_res_blocks': 2,\n",
        "        'resblock_updown': True,\n",
        "        'use_checkpoint': use_checkpoint,\n",
        "        'use_fp16': True,\n",
        "        'use_scale_shift_norm': True,\n",
        "    })"
      ],
      "metadata": {
        "id": "1_xc7GvAwgNU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "MODEL_FLAGS=\"--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64  --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True\"\n",
        "TRAIN_FLAGS=\"--lr 2e-5 --batch_size 4 --save_interval 1000 --log_interval 50 --resume_checkpoint /content/256x256_diffusion_uncond.pt\"  \n",
        "DATASET_PATH=\"/content/YourDatasetHere/\" #change to point to your dataset path \n",
        "OUTPUT_PATH=\"/content/drive/MyDrive/deep_learning/guided-diffusion/\"\n",
        "%cd /content/guided-diffusion\n",
        "!python scripts/image_train.py --data_dir $DATASET_PATH $MODEL_FLAGS $TRAIN_FLAGS --logdir $OUTPUT_PATH\n",
        "\n",
        "#if you are using vanilla openai repo, then you will ned to run this:\n",
        "# !OPENAI_LOGDIR=$OUTPUT_PATH python scripts/image_train.py --data_dir $DATASET_PATH $MODEL_FLAGS $TRAIN_FLAGS "
      ],
      "metadata": {
        "id": "fJtcF4C_jDjz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Sample from model"
      ],
      "metadata": {
        "id": "AHbxCkynj2h0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Sampling\n",
        "The best way to sample your model in real-life conditions is to plug it into DiscoDiffusion. \n",
        "\n",
        "\n",
        "Grab your latest ema checkpoint, open this colab here - https://github.com/Sxela/DiscoDiffusion-Warp/blob/main/Disco_Diffusion_v5_2_Warp_custom_model.ipynb\n",
        "\n",
        "and change settings like described in the previous cell"
      ],
      "metadata": {
        "id": "1NZ2Yi2CxITo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can still sample using vanilla openai code, just plug your checkpoint in the cell below\n",
        "\n",
        "Don't forget to change all the paths"
      ],
      "metadata": {
        "id": "_fD1dA5vxRDb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "checkpoint_path = 'input some checkpoint path here' #use ema checkpoint\n",
        "OUTPUT_PATH=\"/content/drive/MyDrive/deep_learning/guided-diffusion-sxela/\"\n",
        "!python scripts/image_sample.py --num_samples 1 --model_path $checkpoint_path $MODEL_FLAGS --timestep_respacing ddim100 --logdir $OUTPUT_PATH\n",
        "\n",
        "#if you are using vanilla openai repo, then you will ned to run this: \n",
        "#!OPENAI_LOGDIR=/content/drive/MyDrive/deep_learning/guided-diffusion-sxela/samples/  python scripts/image_sample.py --num_samples 1 --model_path $checkpoint_path $MODEL_FLAGS --timestep_respacing ddim100"
      ],
      "metadata": {
        "id": "tAZ1CwLkj11s"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Show results"
      ],
      "metadata": {
        "id": "l3cMMZLKkatO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import PIL\n",
        "\n",
        "sample_path = 'some sample path'\n",
        "im = np.load(sample_path)\n",
        "PIL.Image.fromarray(im.f.arr_0[0])"
      ],
      "metadata": {
        "id": "_WGeIjHhkbnr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Train from scratch (smaller model than vanilla DD, but larger than LSUN)\n",
        "For lvl 1 AI crooks like me, should fit into colab pro"
      ],
      "metadata": {
        "id": "CiMreX-_n6Kz"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Train a smaller model that will fit definitely into colab pro."
      ],
      "metadata": {
        "id": "MLR2sbXSoNdB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Don't forget to change the paths:\n",
        "You need to change DATASET_PATH to point to your dataset images folder, and CHECKPOINT_PATH - to point to a folder you'd like to save progress to. \n",
        "\n",
        "For, example here /content/drive/MyDrive/deep_learning/guided-diffusion-sxela/ - this path points to a location, where all the training checkpoints will be saved\n",
        "\n",
        "and /content/YourDatasetHere/ - this path points to your dataset, i.e. a folder with images (no captions needed)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "We will be using this model together with CLIP inside DiscoDiffusion, so we can train less, stop early and let CLIP do the heavy lifting.\n",
        "\n",
        "This will run almost forever, but you should start checking your results at around ~50k iterations. Good results begin to appear at 100-200k iterations, depending on your dataset.\n",
        "\n",
        "Validating means opening your CHECKPOINT_PATH folder, taking the ema_0.9999_(some number of steps).pt file with the highest number (the latest one), going to this version of DiscoDiffusion here \n",
        "https://github.com/Sxela/DiscoDiffusion-Warp/blob/main/Disco_Diffusion_v5_2_Warp_custom_model.ipynb and setting this: diffusion-model - custom, custom_path - path to that ema file from the previous step (if you saved it on google drive - then just point it there), \n",
        "\n",
        "you'll need to set custom model settings to this: \n",
        "\n",
        "    model_config.update({\n",
        "        'attention_resolutions': '32, 16, 8',\n",
        "        'class_cond': False,\n",
        "        'diffusion_steps': diffusion_steps,\n",
        "        'rescale_timesteps': True,\n",
        "        'timestep_respacing': timestep_respacing,\n",
        "        'image_size': 256,\n",
        "        'learn_sigma': True,\n",
        "        'noise_schedule': 'linear',\n",
        "        'num_channels': 128,\n",
        "        'num_heads': 4,\n",
        "        'num_res_blocks': 2,\n",
        "        'resblock_updown': True,\n",
        "        'use_checkpoint': use_checkpoint,\n",
        "        'use_fp16': True,\n",
        "        'use_scale_shift_norm': True,\n",
        "    })"
      ],
      "metadata": {
        "id": "3p6ThbjFxtBm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "MODEL_FLAGS=\"--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 128 --num_heads 4  --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True\"\n",
        "TRAIN_FLAGS=\"--lr 2e-5 --batch_size 4 --save_interval 1000 --log_interval 50\"\n",
        "DATASET_PATH=\"/content/YourDatasetHere/\" #change to point to your dataset path \n",
        "OUTPUT_PATH=\"/content/drive/MyDrive/deep_learning/guided-diffusion-sxela/\"\n",
        "%cd /content/guided-diffusion-sxela\n",
        "!python scripts/image_train.py --data_dir $DATASET_PATH $MODEL_FLAGS $TRAIN_FLAGS --logdir $OUTPUT_PATH\n",
        "\n",
        "#if you are using vanilla openai repo, then you will ned to run this:\n",
        "# !OPENAI_LOGDIR=$OUTPUT_PATH python scripts/image_train.py --data_dir $DATASET_PATH $MODEL_FLAGS $TRAIN_FLAGS "
      ],
      "metadata": {
        "id": "UfH7XSbKn7ib"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Sampling\n",
        "The best way to sample your model in real-life conditions is to plug it into DiscoDiffusion. \n",
        "\n",
        "\n",
        "Grab your latest ema checkpoint, open this colab here - https://github.com/Sxela/DiscoDiffusion-Warp/blob/main/Disco_Diffusion_v5_2_Warp_custom_model.ipynb\n",
        "\n",
        "and change settings like described in the previous cell"
      ],
      "metadata": {
        "id": "8seIEPF9pF7Q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "checkpoint_path = 'input some checkpoint path here' #use ema checkpoint\n",
        "OUTPUT_PATH=\"/content/drive/MyDrive/deep_learning/guided-diffusion-sxela/\"\n",
        "!python scripts/image_sample.py --num_samples 1 --model_path $checkpoint_path $MODEL_FLAGS  --timestep_respacing ddim100 --logdir $OUTPUT_PATH\n",
        "\n",
        "#if you are using vanilla openai repo, then you will ned to run this: \n",
        "#!OPENAI_LOGDIR=/content/drive/MyDrive/deep_learning/guided-diffusion-sxela/samples/  python scripts/image_sample.py --num_samples 1 --model_path $checkpoint_path $MODEL_FLAGS  --timestep_respacing ddim100"
      ],
      "metadata": {
        "id": "GaHnukcKpEX-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Show results"
      ],
      "metadata": {
        "id": "3mfZb81vpIK_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import PIL\n",
        "\n",
        "sample_path = 'some sample path'\n",
        "im = np.load(sample_path)\n",
        "PIL.Image.fromarray(im.f.arr_0[0])"
      ],
      "metadata": {
        "id": "QCwCF0NhpHy6"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}