readme: {
    "nbformat": 4,
    "nbformat_minor": 0,
    "metadata": {
      "accelerator": "GPU",
      "colab": {
        "name": "animegan_v2_for_videos.ipynb",
        "provenance": [],
        "collapsed_sections": [],
        "authorship_tag": "ABX9TyP/bydrfrVmE0CzRt9JBw+x",
        "include_colab_link": true
      },
      "kernelspec": {
        "display_name": "Python 3",
        "name": "python3"
      },
      "language_info": {
        "name": "python"
      }
    },
    "cells": [
      {
        "cell_type": "markdown",
        "metadata": {
          "id": "view-in-github",
          "colab_type": "text"
        },
        "source": [
          "<a href=\"https://colab.research.google.com/github/nateraw/animegan-v2-for-videos/blob/main/animegan_v2_for_videos.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
        ]
      },
      {
        "cell_type": "code",
        "metadata": {
          "id": "dufmM-T1Helt"
        },
        "source": [
          "%%capture\n",
          "! pip install gradio encoded-video"
        ],
        "execution_count": null,
        "outputs": []
      },
      {
        "cell_type": "code",
        "metadata": {
          "id": "9CY3n8A0Lvdi"
        },
        "source": [
          "import gc\n",
          "import math\n",
          "import tempfile\n",
          "from PIL import Image\n",
          "from io import BytesIO\n",
          "\n",
          "import torch\n",
          "import gradio as gr\n",
          "import numpy as np\n",
          "from encoded_video import EncodedVideo, write_video\n",
          "from torchvision.transforms.functional import to_tensor, center_crop"
        ],
        "execution_count": null,
        "outputs": []
      },
      {
        "cell_type": "code",
        "metadata": {
          "id": "YxdCnrTzLw5V"
        },
        "source": [
          "model = torch.hub.load(\n",
          "    \"AK391/animegan2-pytorch:main\",\n",
          "    \"generator\",\n",
          "    pretrained=True,\n",
          "    device=\"cuda\",\n",
          "    progress=True,\n",
          ")"
        ],
        "execution_count": null,
        "outputs": []
      },
      {
        "cell_type": "code",
        "metadata": {
          "id": "TYAyXUP1UeOd"
        },
        "source": [
          "! curl https://upload.wikimedia.org/wikipedia/commons/transcoded/2/29/2017-01-07_President_Obama%27s_Weekly_Address.webm/2017-01-07_President_Obama%27s_Weekly_Address.webm.360p.vp9.webm -o obama.webm"
        ],
        "execution_count": null,
        "outputs": []
      },
      {
        "cell_type": "code",
        "metadata": {
          "id": "TxT45Nlc88tD"
        },
        "source": [
          "def face2paint(model: torch.nn.Module, img: Image.Image, size: int = 512, device: str = 'cuda'):\n",
          "    w, h = img.size\n",
          "    s = min(w, h)\n",
          "    img = img.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))\n",
          "    img = img.resize((size, size), Image.LANCZOS)\n",
          "\n",
          "    with torch.no_grad():\n",
          "        input = to_tensor(img).unsqueeze(0) * 2 - 1\n",
          "        output = model(input.to(device)).cpu()[0]\n",
          "\n",
          "        output = (output * 0.5 + 0.5).clip(0, 1) * 255.\n",
          "\n",
          "    return output\n",
          "\n",
          "# This function is taken from pytorchvideo!\n",
          "def uniform_temporal_subsample(x: torch.Tensor, num_samples: int, temporal_dim: int = -3) -> torch.Tensor:\n",
          "    \"\"\"\n",
          "    Uniformly subsamples num_samples indices from the temporal dimension of the video.\n",
          "    When num_samples is larger than the size of temporal dimension of the video, it\n",
          "    will sample frames based on nearest neighbor interpolation.\n",
          "    Args:\n",
          "        x (torch.Tensor): A video tensor with dimension larger than one with torch\n",
          "            tensor type includes int, long, float, complex, etc.\n",
          "        num_samples (int): The number of equispaced samples to be selected\n",
          "        temporal_dim (int): dimension of temporal to perform temporal subsample.\n",
          "    Returns:\n",
          "        An x-like Tensor with subsampled temporal dimension.\n",
          "    \"\"\"\n",
          "    t = x.shape[temporal_dim]\n",
          "    assert num_samples > 0 and t > 0\n",
          "    # Sample by nearest neighbor interpolation if num_samples > t.\n",
          "    indices = torch.linspace(0, t - 1, num_samples)\n",
          "    indices = torch.clamp(indices, 0, t - 1).long()\n",
          "    return torch.index_select(x, temporal_dim, indices)\n",
          "\n",
          "\n",
          "def short_side_scale(\n",
          "    x: torch.Tensor,\n",
          "    size: int,\n",
          "    interpolation: str = \"bilinear\",\n",
          ") -> torch.Tensor:\n",
          "    \"\"\"\n",
          "    Determines the shorter spatial dim of the video (i.e. width or height) and scales\n",
          "    it to the given size. To maintain aspect ratio, the longer side is then s