MP_v_diffusion_images512x512.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/metaphorz/v-diffusion-pytorch/blob/master/MP_v_diffusion_images512x512.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kfsbjhr0GjMi"
      },
      "source": [
        "# Multi-Parameter v-diffusion-pytorch\n",
        "### This is Katherine Crowson's v-diffusion-pytorch repository with some additions\n",
        "### P. Fishwick, 01/17/22\n",
        "1. Center installation in Google Drive by mounting it.\n",
        "2. First run: the installation - store in Drive\n",
        "3. First run: the pretrained model - store in Drive\n",
        "4. A loop over a range of skip_timestep (st) and seed values (bottom of notebook) -- this allows you to \"batch\" run the notebook to create multiple images. There is a --batch_size option in cfg_sample.py and this is NOT used.\n",
        "5. A collage option -- makes rectangular collages e.g. 3x3, 6x6, 4x7, 2x4\n",
        "6. Storing resulting images in a zip file and then downloading\n",
        "\n",
        "This notebook creates a \"multiparam-v-diffusion-pytorch\" directory within Google Drive."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cFfU7JXtB_GV"
      },
      "outputs": [],
      "source": [
        "# Load Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n5MnJSsDaU8p"
      },
      "outputs": [],
      "source": [
        "#\n",
        "# discover the assigned GPU\n",
        "#\n",
        "!nvidia-smi -L"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u3alYfr8aYdr"
      },
      "outputs": [],
      "source": [
        "# install requirements\n",
        "!pip install ftfy regex requests tqdm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_rndN6sXYsRg"
      },
      "outputs": [],
      "source": [
        " # obtain v-diffusion-pytorch\n",
        "# run this one time\n",
        "root = '/content/drive/MyDrive/multiparam-v-diffusion-pytorch'\n",
        "import os\n",
        "if (not os.path.isdir(root)):\n",
        "  !git clone --recursive https://github.com/metaphorz/v-diffusion-pytorch {root}\n",
        "else:\n",
        "  print(root+' exists')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "igMHf8BdcBmz"
      },
      "outputs": [],
      "source": [
        "#\n",
        "# Download the diffusion model\n",
        "# The model list (checkpoints) can be found https://github.com/crowsonkb/v-diffusion-pytorch#model-checkpoints\n",
        "# SHA-256: 4fc95ee1b3205a3f7422a07746383776e1dbc367eaf06a5b658ad351e77b7bda\n",
        "#\n",
        "# Save pretrained model in Google Drive for first run (v-diffusion-pytorch/ should not pre-exist!)\n",
        "import os\n",
        "if (not os.path.isdir(root+'/checkpoints')):\n",
        "  !mkdir {root}/results\n",
        "  !mkdir {root}/checkpoints\n",
        "  !curl -L https://v-diffusion.s3.us-west-2.amazonaws.com/cc12m_1_cfg.pth > {root}/checkpoints/cc12m_1_cfg.pth\n",
        "else:\n",
        "  print('pretrained model already downloaded')\n",
        "# From original \n",
        "# !curl -L https://v-diffusion.s3.us-west-2.amazonaws.com/cc12m_1_cfg.pth > {root}/checkpoints/cc12m_1_cfg.pth"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sQhDMRNac4RV"
      },
      "outputs": [],
      "source": [
        "%cd {root}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BFWNudcQZFKM"
      },
      "outputs": [],
      "source": [
        "# Run inference loop given the init image for each st and seed value\n",
        "# Edit the st and range values\n",
        "import numpy as np\n",
        "#\n",
        "# Specify number of skip_timesteps\n",
        "# in between st_start and st_stop\n",
        "#\n",
        "st_num = 1\n",
        "st_start = 0.85\n",
        "st_stop = 0.90\n",
        "st_step = (st_stop - st_start)/st_num\n",
        "st_list = np.arange(st_start, st_stop, st_step)\n",
        "st_values = list(st_list)\n",
        "print('# of st_values ', len(st_values))\n",
        "#\n",
        "# Specify number of seeds\n",
        "# starting at seed_start\n",
        "#\n",
        "seed_num = 9\n",
        "seed_start = 0\n",
        "seed_stop = seed_start + seed_num\n",
        "seed_step = 1\n",
        "seed_list = np.arange(seed_start, seed_stop, seed_step)\n",
        "seed_values = list(seed_list)\n",
        "print('# of seeds ', seed_num)\n",
        "print('# of generated images ', st_num*seed_num)\n",
        "print('Image Generation ...')\n",
        "for st_num in st_values:\n",
        "  for seed_num2 in seed_values:\n",
        "    # enter your own init image\n",
        "    !python cfg_sample.py --init /content/drive/MyDrive/images/abandoned2.jpg -st $st_num --size 512 512 --seed $seed_num2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0sp3wdUSf3fi"
      },
      "outputs": [],
      "source": [
        "# save all images in a zip\n",
        "!zip {root}/images.zip {root}/results/*.png"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W3W81jhlk7G4"
      },
      "outputs": [],
      "source": [
        "# download images.zip\n",
        "# sometimes this gets stuck, so it is commented out\n",
        "# instead, click on the images.zip file in Drive under multiparam-v-diffusion-pytorch\n",
        "# and download the zip file\n",
        "#from google.colab import files\n",
        "#files.download(root+'/images.zip')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iBW4Mn2GO-8N"
      },
      "outputs": [],
      "source": [
        "# Optional: Google Drive results/ cleaning\n",
        "# clean your results folder between notebook runs\n",
        "# once images.zip has downloaded\n",
        "#!rm {root}/images.zip\n",
        "#!rm {root}/results/*.png"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-E4OHTZhRo-J"
      },
      "outputs": [],
      "source": [
        "%cd {root}/results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HLLmTSu-RVcQ"
      },
      "outputs": [],
      "source": [
        "# Generate a collage of all images\n",
        "# and store in multiparam-v-diffusion-pytorch/\n",
        "# Code adapted from: https://holypython.com/python-pil-tutorial/creating-photo-collages/\n",
        "\n",
        "from PIL import Image, ImageDraw\n",
        "import glob\n",
        "lst = sorted(glob.glob(\"*.png\"))\n",
        "print('check list length ',len(lst))\n",
        "# to make the collage, you need num_x * num_y .png files\n",
        "# need not be square\n",
        "# 9 images to make a 3x3 collage\n",
        "# 36 images to make a 6x6 collage\n",
        "# 500 pixels for each side of square component images\n",
        "num_x = 3\n",
        "num_y = 3\n",
        "print('generating a ',num_x,' * ',num_y,' collage')\n",
        "# choose this to modify the size of each square in the collage\n",
        "# side_square = 200 # for smaller collage image\n",
        "side_square = 512 # for the largest collage image\n",
        "\n",
        "\n",
        "collage = Image.new(\"RGBA\", (side_square*num_x,side_square*num_y), color=(255,255,255,255))\n",
        "\n",
        "file_num = 0\n",
        "for j in range(0,num_y*side_square,side_square):\n",
        "    for i in range(0,num_x*side_square,side_square):\n",
        "        image_file = str(lst[file_num])\n",
        "        photo = Image.open(image_file).convert(\"RGBA\")\n",
        "        photo = photo.resize((side_square,side_square))          \n",
        "        collage.paste(photo, (i,j))\n",
        "        file_num += 1\n",
        "\n",
        "display(collage)\n",
        "collage.save(root+\"/collage.png\")\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "name": "MP-v-diffusion-images512x512.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}