VQGAN+CLIP_Video_with_Optical_Flow.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "VQGAN+CLIP Video with Optical Flow",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [
        "XYU-20jFsaW8"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Setup"
      ],
      "metadata": {
        "id": "XYU-20jFsaW8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Prepare workspace\n",
        "! rm -r -f .config sample_data\n",
        "! git clone --recursive https://github.com/robobeebop/VQGAN-CLIP-Video.git ."
      ],
      "metadata": {
        "cellView": "form",
        "id": "cjDu-Hbtmrwb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Mount Drive\n",
        "\n",
        "from google.colab import output\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "0mEllCWgnrf2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Install Required Libraries\n",
        "! pip install lpips ftfy omegaconf einops pytorch-lightning\n",
        "! pip3 install transformers"
      ],
      "metadata": {
        "cellView": "form",
        "id": "lhwh9J2klrkf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# RUN"
      ],
      "metadata": {
        "id": "0A_d725Ssdm3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Select a model\n",
        "from dream import Dream\n",
        "from os.path import exists\n",
        "!mkdir checkpoints\n",
        "ckpt_dir = \"/content/checkpoints/\"\n",
        "\n",
        "vqgan_model = \"imagenet_1024\" #@param [\"imagenet_1024\", \"imagenet_16384\", \"coco\", \"sflckr\"]\n",
        "\n",
        "vqgan_options = {\n",
        "  \"imagenet_1024\": [\"1-7QlixzWxZAO8ktGFqvxrZ_JzapzI5hH\", \"1-8mSOBsutfkE95piiGf4ZuVX0zkAwzkn\"],\n",
        "  \"imagenet_16384\": [\"1_1q5zxEBx17AyTALEhGqhSsS7tyCJ4fe\", \"1-0D4pbu7NHrvWzTfbw4hiA1Sno75Z2_C\"],\n",
        "  \"coco\": [\"1-9gq1a4yGOKC3rDw-X9NBe5_JVcKcLPG\", \"1-CPBZXsCgCv-Z6Uy4Sf4lKeyqG_C5i-Y\"],\n",
        "  \"sflckr\": [\"1iIgSRV4H6og3l2myXPRE043ULoPlqn8w\", \"1-1vMpPmB6QZhGzriXG9iI6WeFZLl7VP2\"],\n",
        "}\n",
        "\n",
        "yaml, ckpt = ckpt_dir + \"%s.yaml\"%vqgan_model, ckpt_dir + \"%s.ckpt\"%vqgan_model\n",
        "\n",
        "if not exists(ckpt) or not exists(yaml):\n",
        "  yaml_id = vqgan_options[vqgan_model][0]\n",
        "  ckpt_id = vqgan_options[vqgan_model][1]\n",
        "\n",
        "  !gdown --id \"$yaml_id\" -O \"$yaml\"\n",
        "  !gdown --id \"$ckpt_id\" -O \"$ckpt\"\n",
        "\n",
        "\n",
        "\n",
        "dream = Dream()\n",
        "dream.cook([yaml, ckpt])"
      ],
      "metadata": {
        "cellView": "form",
        "id": "gzz5pv6-nXLk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the cell above once when it's initial googlecolab run or you want to change the VQGAN model"
      ],
      "metadata": {
        "id": "TRXTBau3bWdf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Dream  { display-mode: \"form\" }\n",
        "#@markdown You can see output frames in /content/output folder.\n",
        "\n",
        "from dream import cv2\n",
        "from dream import save_img\n",
        "from dream import reduce_res\n",
        "from dream import np\n",
        "from dream import get_opflow_image\n",
        "from dream import PIL\n",
        "from dream import trange\n",
        "from dream import glob\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown Describe how deepdream should look like. Each scene ( | ... | ) can be weighted: \"deepdream:3 | dogs:-1\"\n",
        "text_prompts = \"trending on artstation | lovecraftian horror | deepdream | vibrant colors | 4k | made by Edvard Munch\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown Video paths.\n",
        "vid_path = '/content/samples/sample_policeman.mp4' #@param {type:\"string\"}\n",
        "output_vid_path = '/content/samples/deepdreamed_policeman.mp4' #@param {type:\"string\"}\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown Play around with these settings, finding optimal settings may vary video to video. Set both to 0 for more chaotic experience.\n",
        "frame_weight = 10#@param {type:\"number\"}\n",
        "previous_frame_weight =  0.1#@param {type:\"number\"}\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown Usual VQGAN+CLIP settings\n",
        "step_size = 0.15 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "iter_n =  5#@param {type:\"number\"}\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown Dream more intense on first frame of the video. \n",
        "do_wait_first_frame = True #@param {type:\"boolean\"}\n",
        "wait_step_size = 0.15 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "wait_iter_n =  15#@param {type:\"number\"}\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown Weights of how previous deepdreamed frame should effect current frame.\n",
        "blendflow = 0.6#@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "blendstatic =  0.6#@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "#@markdown ---\n",
        "#@markdown Video resolution and fps\n",
        "w = 1920#@param {type:\"number\"}\n",
        "h = 1080#@param {type:\"number\"}\n",
        "fps =  24#@param {type:\"number\"}\n",
        "#@markdown ---\n",
        "#@markdown Make a test run.\n",
        "is_test = False #@param {type:\"boolean\"}\n",
        "test_finish_at = 24#@param {type:\"number\"}\n",
        "#@markdown ---\n",
        "#@markdown Get all frames from video, can be set to False if video didn't change. \n",
        "video_to_frames = True #@param {type:\"boolean\"}\n",
        "\n",
        "!mkdir input\n",
        "!mkdir output\n",
        "!rm -r -f ./output/*.jpg\n",
        "\n",
        "\n",
        "if(video_to_frames):\n",
        "    !rm -r -f ./input/*.jpg\n",
        "    vidcap = cv2.VideoCapture(vid_path)\n",
        "    success,image = vidcap.read()\n",
        "    index = 1\n",
        "    while success:\n",
        "        cv2.imwrite(\"./input/%04d.jpg\" % index, image)\n",
        "        success, image = vidcap.read()\n",
        "        index += 1\n",
        "\n",
        "x, y = reduce_res((w, h))\n",
        "img_arr = sorted(glob('input/*.jpg'))\n",
        "\n",
        "np_img = np.float32(PIL.Image.open(img_arr[0]))\n",
        "np_img = cv2.resize(np_img, dsize=(x, y), interpolation=cv2.INTER_CUBIC)\n",
        "h, w, c = np_img.shape\n",
        "\n",
        "frame = None\n",
        "\n",
        "if do_wait_first_frame:\n",
        "    frame = dream.deepdream(np_img, text_prompts, [x, y], iter_n=wait_iter_n, step_size=step_size, init_weight=frame_weight)\n",
        "else:\n",
        "    frame = dream.deepdream(np_img, text_prompts, [x, y], iter_n=iter_n, step_size=step_size, init_weight=frame_weight)\n",
        "\n",
        "frame = cv2.resize(frame, dsize=(x, y), interpolation=cv2.INTER_CUBIC)\n",
        "save_img(frame, 'output/%04d.jpg'%0)\n",
        "\n",
        "img_range = trange(len(img_arr[:test_finish_at]), desc=\"Dreaming\") if is_test else trange(len(img_arr), desc=\"Dreaming\")\n",
        "prev_frame = None\n",
        "for i in img_range:  \n",
        "    if previous_frame_weight != 0:\n",
        "        prev_frame = np.copy(frame)\n",
        "\n",
        "    img = img_arr[i]\n",
        "    np_prev_img = np_img\n",
        "    np_img = np.float32(PIL.Image.open(img))\n",
        "    np_img = cv2.resize(np_img, dsize=(x, y), interpolation=cv2.INTER_CUBIC)\n",
        "    frame = cv2.resize(frame, dsize=(x, y), interpolation=cv2.INTER_CUBIC)\n",
        "    \n",
        "    frame_flow_masked, background_masked = get_opflow_image(np_prev_img, frame, np_img, blendflow, blendstatic)\n",
        "    frame = frame_flow_masked + background_masked\n",
        "    frame = dream.deepdream(frame, text_prompts, [x, y], iter_n=iter_n, init_weight=frame_weight, step_size=step_size, image_prompts=prev_frame, image_prompt_weight=previous_frame_weight)\n",
        "    \n",
        "    save_img(frame, 'output/%04d.jpg'%i)\n",
        "\n",
        "\n",
        "mp4_fourcc = cv2.VideoWriter_fourcc(*'MP4V')\n",
        "out = cv2.VideoWriter(output_vid_path, mp4_fourcc, fps, (w, h))\n",
        "filelist = sorted(glob('output/*.jpg'))\n",
        "\n",
        "for i in trange(len(filelist), desc=\"Generating Video\"):\n",
        "    img = cv2.imread(filelist[i])\n",
        "    img = cv2.resize(img, dsize=(w, h), interpolation=cv2.INTER_CUBIC)\n",
        "    out.write(img)\n",
        "out.release()"
      ],
      "metadata": {
        "id": "O6GwujqCowd9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Download Generated Video.\n",
        "# @markdown (doesn't work on Safari)\n",
        "from google.colab import files\n",
        "\n",
        "files.download(output_vid_path)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "PXZmbS0IeKUn"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}