DiffusionCLIP_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o38wi96CDr-0"
      },
      "source": [
        "論文<br>\n",
        "https://arxiv.org/abs/2110.02711<br>\n",
        "<br>\n",
        "GitHub<br>\n",
        "https://github.com/gwang-kim/diffusionclip<br>\n",
        "<br>\n",
        "<a href=\"https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/DiffusionCLIP_demo.ipynb\" target=\"_blank\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dl0geaqIDr-5"
      },
      "source": [
        "## 環境セットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I0fVT2GdDr-6"
      },
      "source": [
        "## GPU確認"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2PcII5fLdHnJ",
        "outputId": "2ffaee6a-5bce-4597-9e15-ccfcfe0329e8"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zshu3r5PDr-7"
      },
      "source": [
        "## GitHubからコード取得"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XA6vYWpA_ndq",
        "outputId": "3cb0dded-4f6d-4b89-caa9-e09b105c84ea"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "\n",
        "!git clone https://github.com/gwang-kim/DiffusionCLIP.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WaxAMTPz_yWu"
      },
      "source": [
        "## ライブラリのインストール"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DcmolIpv_xSX",
        "outputId": "d7ffca0b-1cfe-410f-9fc7-c8a90960c43e"
      },
      "outputs": [],
      "source": [
        "!pip install ftfy regex tqdm\n",
        "!pip install --upgrade gdown\n",
        "!pip install git+https://github.com/openai/CLIP.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1wcj95AhAi9L"
      },
      "source": [
        "## ライブラリのインポート"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mJgQRcIPAkwH",
        "outputId": "65aca97c-f485-4606-e74b-65959b9c4cd2"
      },
      "outputs": [],
      "source": [
        "%cd /content/DiffusionCLIP\n",
        "\n",
        "from diffusionclip import DiffusionCLIP\n",
        "from main import dict2namespace\n",
        "import argparse\n",
        "import yaml\n",
        "from PIL import Image\n",
        "import os\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(action='ignore')\n",
        "\n",
        "import torch\n",
        "device = 'cuda' if torch.cuda.is_available() else \"cpu\"\n",
        "print(\"using device is\", device)\n",
        "\n",
        "# モジュールの再読み込み\n",
        "%load_ext autoreload\n",
        "%autoreload 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "129s7MNwBA6x"
      },
      "source": [
        "# 学習済みモデルのダウンロード"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PzC4pPySBMst",
        "outputId": "16e1249d-718a-4937-9137-88b8db1c1cbd"
      },
      "outputs": [],
      "source": [
        "%cd /content/DiffusionCLIP\n",
        "!mkdir pretrained\n",
        "%cd pretrained\n",
        "\n",
        "# Finetune Human face model download\n",
        "if not os.path.exists('human_pixar_t601.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=1IoT7kZhtaoKf1uvhYhvyqzyG2MOJsqLe'\n",
        "if not os.path.exists('human_neanderthal_t601.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=1Uo0VI5kbATrQtckhEBKUPyRFNOcgwwne'\n",
        "if not os.path.exists('human_gogh_t601.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=1NXOL8oKTGLtpTsU_Vh5h0DmMeH7WG8rQ'\n",
        "if not os.path.exists('human_tanned_t201.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=1k6aDDOedRxhjFsJIA0dZLi2kKNvFkSYk'\n",
        "if not os.path.exists('human_male_t401.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=1n1GMVjVGxSwaQuWxoUGQ2pjV8Fhh72eh'\n",
        "if not os.path.exists('human_sketch_t601.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=1V9HDO8AEQzfWFypng72WQJRZTSQ272gb'\n",
        "if not os.path.exists('human_with_makeup_t301.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=1OL0mKK48wvaFaWGEs3GHsCwxxg7LexOh'\n",
        "if not os.path.exists('human_without_makeup_t301.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=157pTJBkXPoziGQdjy3SwdyeSpAjQiGRp'\n",
        "\n",
        "if not os.path.exists('512x512_diffusion.pt'):\n",
        "  !wget pretrained/ https://openaipublic.blob.core.windows.net/diffusion/jul-2021/512x512_diffusion.pt\n",
        "\n",
        "if not os.path.exists('imagenet_watercolor_t601.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=1l1vLwdL-6kC9jKcStASZ0KtX2OrmrSj6'\n",
        "if not os.path.exists('imagenet_pointillism_t601.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=1Am1Iii7jH986XQUuVaDs4v5s1h_acg0w'\n",
        "if not os.path.exists('imagenet_gogh_t601.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=1ZPeOvMpFStw8RXJga_0pWLJ7iIWEQIVY'\n",
        "if not os.path.exists('imagenet_cubism_t601.pth'):\n",
        "  !gdown 'https://drive.google.com/uc?id=1xEx4_MXvbvtSqLzn6z49RUnPDFoDv9Vm'\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "107DoyAeC93T"
      },
      "source": [
        "# テスト画像ダウンロード"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8KkFUbsiDAH_",
        "outputId": "857defc1-5658-43ad-cffd-150a0e979f6d"
      },
      "outputs": [],
      "source": [
        "%cd /content/DiffusionCLIP\n",
        "!mkdir test_imgs\n",
        "%cd test_imgs\n",
        "\n",
        "!wget https://www.pakutaso.com/shared/img/thumb/kys150922346900.jpg\n",
        "\n",
        "!wget https://www.pakutaso.com/shared/img/thumb/kawamurassIMGL3813_TP_V4.jpg\n",
        "# 画像の中心から512x512をcrop\n",
        "def crop_center(pil_img, crop_width, crop_height):\n",
        "    img_width, img_height = pil_img.size\n",
        "    return pil_img.crop(((img_width - crop_width) // 2,\n",
        "                         (img_height - crop_height) // 2,\n",
        "                         (img_width + crop_width) // 2,\n",
        "                         (img_height + crop_height) // 2))\n",
        "img = Image.open('kawamurassIMGL3813_TP_V4.jpg')\n",
        "im_crop = crop_center(img, 512, 512)\n",
        "im_crop.save('crop.jpg')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VP_vT_7WDyYc"
      },
      "source": [
        "# Human face manipulation (256x256)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "34epsqZeGM-J"
      },
      "outputs": [],
      "source": [
        "%cd /content/DiffusionCLIP\n",
        "\n",
        "model_dict = {\n",
        "    'Pixar':            \"pretrained/human_pixar_t601.pth\",\n",
        "    'Neanderthal':      \"pretrained/human_neanderthal_t601.pth\",\n",
        "    'Painting by Gogh': \"pretrained/human_gogh_t601.pth\",\n",
        "    'Tanned':           \"pretrained/human_tanned_t201.pth\",\n",
        "    'Female → Male':   \"pretrained/human_male_t401.pth\",\n",
        "    'Sketch':           \"pretrained/human_sketch_t601.pth\",\n",
        "    'With makeup':      \"pretrained/human_with_makeup_t301.pth\",\n",
        "    'Without makeup':   \"pretrained/human_without_makeup_t301.pth\",\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZrqgeW6DGJsS"
      },
      "source": [
        "## パラメータ設定"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 464
        },
        "id": "ATfUnoTlHVcY",
        "outputId": "89bca01a-22bd-49cb-ebc0-070fddfaa156"
      },
      "outputs": [],
      "source": [
        "%cd /content/DiffusionCLIP\n",
        "\n",
        "\n",
        "# @markdown 入力画像パス\n",
        "img_path = \"test_imgs/kys150922346900.jpg\" #@param {type:\"string\"}\n",
        "# @markdown 顔部分切り取り\n",
        "align_face = True #@param {type:\"boolean\"}\n",
        "\n",
        "# @markdown type\n",
        "edit_type = 'Sketch' #@param ['Pixar', 'Neanderthal','Sketch', 'Painting by Gogh', 'Tanned',  'With makeup', 'Without makeup', 'Female → Male']\n",
        "degree_of_change = 1 #@param {type:\"slider\", min:0.0, max:1.0, step:0.01}\n",
        "\n",
        "n_inv_step =  40#@param {type: \"integer\"}\n",
        "n_test_step = 6 #@param [6] \n",
        "\n",
        "\n",
        "model_path = model_dict[edit_type]\n",
        "t_0 = int(model_path.split('_t')[-1].replace('.pth',''))\n",
        "\n",
        "exp_dir = f\"runs/MANI_{img_path.split('/')[-1]}_align{align_face}\"\n",
        "os.makedirs(exp_dir, exist_ok=True)\n",
        "\n",
        "args_dic = {\n",
        "    'config': 'celeba.yml', \n",
        "    't_0': t_0, \n",
        "    'n_inv_step': int(n_inv_step), \n",
        "    'n_test_step': int(n_test_step),\n",
        "    'sample_type': 'ddim', \n",
        "    'eta': 0.0,\n",
        "    'bs_test': 1, \n",
        "    'model_path': model_path, \n",
        "    'img_path': img_path, \n",
        "    'deterministic_inv': 1, \n",
        "    'hybrid_noise': 0, \n",
        "    'n_iter': 1,  \n",
        "    'align_face': align_face, \n",
        "    'image_folder': exp_dir,\n",
        "    'model_ratio': degree_of_change,\n",
        "    'edit_attr': None, 'src_txts': None, 'trg_txts': None,\n",
        "    }\n",
        "args = dict2namespace(args_dic)\n",
        "\n",
        "with open(os.path.join('configs', args.config), 'r') as f:\n",
        "    config_dic = yaml.safe_load(f)\n",
        "config = dict2namespace(config_dic)\n",
        "config.device = device\n",
        "\n",
        "# Edit\n",
        "runner = DiffusionCLIP(args, config)\n",
        "runner.edit_one_image()\n",
        "\n",
        "# Result\n",
        "print()\n",
        "n_result = 1\n",
        "img = Image.open(os.path.join(exp_dir, '0_orig.png'))\n",
        "img = img.resize((int(img.width), int(img.height)))\n",
        "grid = Image.new(\"RGB\", (img.width*(n_result+1), img.height))\n",
        "grid.paste(img, (0, 0))\n",
        "for i in range(n_result):\n",
        "  img = Image.open(os.path.join(exp_dir, f\"3_gen_t{t_0}_it0_ninv{n_inv_step}_ngen{n_test_step}_mrat{degree_of_change}_{model_path.split('/')[-1].replace('.pth','')}.png\"))\n",
        "  img = img.resize((int(img.width), int(img.height)))\n",
        "  grid.paste(img, (int(img.height * (i+1)), 0))\n",
        "grid"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jf4FXwIiOMp7"
      },
      "source": [
        "# ImageNet Style Transfer (512x512)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yGs--AO0OO9G",
        "outputId": "f4a22c75-e21e-4d22-fe05-8402e26b527e"
      },
      "outputs": [],
      "source": [
        "%cd /content/DiffusionCLIP\n",
        "\n",
        "model_dict = {\n",
        "    'Watercolor art':            \"pretrained/imagenet_watercolor_t601.pth\",\n",
        "    'Pointillism art':      \"pretrained/imagenet_pointillism_t601.pth\",\n",
        "    'Painting by Gogh': \"pretrained/imagenet_gogh_t601.pth\",\n",
        "    'Cubism art':           \"pretrained/imagenet_cubism_t601.pth\",\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GS9ZAlAhOrhK"
      },
      "source": [
        "## パラメータ設定"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 612
        },
        "id": "AZCDAr5MOsu2",
        "outputId": "f50cf4b8-eb07-49e5-ae8f-e14ee4227241"
      },
      "outputs": [],
      "source": [
        "%cd /content/DiffusionCLIP\n",
        "\n",
        "\n",
        "# @markdown 入力画像パス\n",
        "img_path = \"test_imgs/crop.jpg\" #@param {type:\"string\"}\n",
        "# @markdown type\n",
        "edit_type = 'Watercolor art' #@param ['Watercolor art', 'Pointillism art','Painting by Gogh', 'Cubism art']\n",
        "degree_of_change = 1 #@param {type:\"slider\", min:0.0, max:1.0, step:0.01}\n",
        "\n",
        "n_inv_step =  40#@param {type: \"integer\"}\n",
        "n_test_step = 6 #@param [6]\n",
        "\n",
        "model_path = model_dict[edit_type]\n",
        "\n",
        "t_0 = int(model_path.split('_t')[-1].replace('.pth',''))\n",
        "\n",
        "exp_dir = f\"runs/MANI_{img_path.split('/')[-1]}\"\n",
        "os.makedirs(exp_dir, exist_ok=True)\n",
        "\n",
        "args_dic = {\n",
        "    'config': 'imagenet.yml', \n",
        "    't_0': t_0, \n",
        "    'n_inv_step': int(n_inv_step), \n",
        "    'n_test_step': int(n_test_step),\n",
        "    'sample_type': 'ddim', \n",
        "    'eta': 0.0,\n",
        "    'bs_test': 1, \n",
        "    'model_path': model_path, \n",
        "    'img_path': img_path, \n",
        "    'deterministic_inv': 1, \n",
        "    'hybrid_noise': 0, \n",
        "    'n_iter': 1,  \n",
        "    'align_face': 0,\n",
        "    'image_folder': exp_dir,\n",
        "    'model_ratio': degree_of_change,\n",
        "    'edit_attr': None, 'src_txts': None, 'trg_txts': None,\n",
        "    }\n",
        "args = dict2namespace(args_dic)\n",
        "\n",
        "with open(os.path.join('configs', args.config), 'r') as f:\n",
        "    config_dic = yaml.safe_load(f)\n",
        "config = dict2namespace(config_dic)\n",
        "config.device = device\n",
        "\n",
        "# Edit\n",
        "runner = DiffusionCLIP(args, config)\n",
        "runner.edit_one_image()\n",
        "\n",
        "# Result\n",
        "print()\n",
        "n_result = 1\n",
        "img = Image.open(os.path.join(exp_dir, '0_orig.png'))\n",
        "img = img.resize((int(img.width), int(img.height)))\n",
        "grid = Image.new(\"RGB\", (img.width*(n_result+1), img.height))\n",
        "grid.paste(img, (0, 0))\n",
        "for i in range(n_result):\n",
        "  img = Image.open(os.path.join(exp_dir, f\"3_gen_t{t_0}_it0_ninv{n_inv_step}_ngen{n_test_step}_mrat{degree_of_change}_{model_path.split('/')[-1].replace('.pth','')}.png\"))\n",
        "  img = img.resize((int(img.width), int(img.height)))\n",
        "  grid.paste(img, (int(img.height * (i+1)), 0))\n",
        "grid"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "DiffusionCLIP_demo.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}