StyleGAN3+Inversion.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "StyleGAN3+Inversion",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyM2y4Cw4X6i/UzzkwyFnjbz",
      "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/ouhenio/stylegan3-projector/blob/add-video-download-to-notebook/StyleGAN3%2BInversion.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "hvhFSELI8nYi"
      },
      "outputs": [],
      "source": [
        "#@markdown #**Setup libraries and functions** 🏗️🛠️\n",
        "\n",
        "!pip install --upgrade torch==1.9.1+cu111 torchvision==0.10.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html\n",
        "#!pip install --upgrade https://download.pytorch.org/whl/nightly/cu111/torch-1.11.0.dev20211012%2Bcu111-cp37-cp37m-linux_x86_64.whl https://download.pytorch.org/whl/nightly/cu111/torchvision-0.12.0.dev20211012%2Bcu111-cp37-cp37m-linux_x86_64.whl\n",
        "!git clone https://github.com/NVlabs/stylegan3\n",
        "!pip install einops ninja\n",
        "!pip install wget \n",
        "\n",
        "import sys\n",
        "sys.path.append('./stylegan3')\n",
        "\n",
        "import io\n",
        "import os, time\n",
        "import pickle\n",
        "import shutil\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import requests\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.transforms.functional as TF\n",
        "import copy\n",
        "import imageio\n",
        "import unicodedata\n",
        "import re\n",
        "from PIL import Image\n",
        "from tqdm.notebook import tqdm\n",
        "from torchvision.transforms import Compose, Resize, ToTensor, Normalize\n",
        "from IPython.display import display\n",
        "from einops import rearrange\n",
        "from google.colab import files\n",
        "from time import perf_counter\n",
        "from stylegan3.dnnlib.util import open_url\n",
        "import wget\n",
        "\n",
        "\n",
        "device = torch.device('cuda:0')\n",
        "\n",
        "# Load VGG16 feature detector.\n",
        "url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/vgg16.pt'\n",
        "with open_url(url) as f:\n",
        "    vgg16 = torch.jit.load(f).eval().to(device)\n",
        "print('Using device:', device, file=sys.stderr)\n",
        "\n",
        "def fetch_model(url_or_path):\n",
        "    if \"drive.google\" in url_or_path:\n",
        "      if \"18MOpwTMJsl_Z17q-wQVnaRLCUFZYSNkj\" in url_or_path: \n",
        "        basename = \"wikiart-1024-stylegan3-t-17.2Mimg.pkl\"\n",
        "      elif \"14UGDDOusZ9TMb-pOrF0PAjMGVWLSAii1\" in url_or_path:\n",
        "        basename = \"lhq-256-stylegan3-t-25Mimg.pkl\"\n",
        "    else:\n",
        "        basename = os.path.basename(url_or_path)\n",
        "    if os.path.exists(basename):\n",
        "        return basename\n",
        "    else:\n",
        "        if \"drive.google\" not in url_or_path:\n",
        "          !wget -c '{url_or_path}'\n",
        "        else:\n",
        "          path_id = url_or_path.split(\"id=\")[-1]\n",
        "          !gdown --id '{path_id}'\n",
        "        return basename\n",
        "\n",
        "def download_img(img_url):\n",
        "    try:\n",
        "        return wget.download(img_url,out=\"input.jpg\")\n",
        "    except:\n",
        "        return\n",
        "\n",
        "# Projector\n",
        "\n",
        "def project(\n",
        "    G,\n",
        "    target: torch.Tensor, # [C,H,W] and dynamic range [0,255], W & H must match G output resolution\n",
        "    *,\n",
        "    num_steps                  = 1000,\n",
        "    w_avg_samples              = -1,\n",
        "    initial_learning_rate      = 0.1,\n",
        "    initial_noise_factor       = 0.05,\n",
        "    lr_rampdown_length         = 0.25,\n",
        "    lr_rampup_length           = 0.05,\n",
        "    noise_ramp_length          = 0.75,\n",
        "    regularize_noise_weight    = 1e5,\n",
        "    verbose                    = False,\n",
        "    device: torch.device\n",
        "):\n",
        "\n",
        "    assert target.shape == (G.img_channels, G.img_resolution, G.img_resolution)\n",
        "\n",
        "    def logprint(*args):\n",
        "        if verbose:\n",
        "            print(*args)\n",
        "\n",
        "    G = copy.deepcopy(G).eval().requires_grad_(False).to(device) # type: ignore\n",
        "\n",
        "    # Compute w stats.\n",
        "    if w_avg_samples > 0:\n",
        "      logprint(f'Computing W midpoint and stddev using {w_avg_samples} samples...')\n",
        "      z_samples = np.random.RandomState(123).randn(w_avg_samples, G.z_dim)\n",
        "    else:\n",
        "      seed = np.random.randint(0, 2**32 - 1)\n",
        "      z_samples = np.random.RandomState(seed).randn(1, G.z_dim)\n",
        "    w_samples = G.mapping(torch.from_numpy(z_samples).to(device), None)  # [N, L, C]\n",
        "    w_samples = w_samples[:, :1, :].cpu().numpy().astype(np.float32)       # [N, 1, C]\n",
        "    w_avg = np.mean(w_samples, axis=0, keepdims=True)      # [1, 1, C]\n",
        "    w_std = (np.sum((w_samples - w_avg) ** 2) / w_avg_samples) ** 0.5\n",
        "\n",
        "    # Setup noise inputs.\n",
        "    noise_bufs = { name: buf for (name, buf) in G.synthesis.named_buffers() if 'noise_const' in name }\n",
        "\n",
        "    # Features for target image.\n",
        "    target_images = target.unsqueeze(0).to(device).to(torch.float32)\n",
        "    if target_images.shape[2] > 256:\n",
        "        target_images = F.interpolate(target_images, size=(256, 256), mode='area')\n",
        "    target_features = vgg16(target_images, resize_images=False, return_lpips=True)\n",
        "\n",
        "    w_opt = torch.tensor(w_avg, dtype=torch.float32, device=device, requires_grad=True) # pylint: disable=not-callable\n",
        "    w_out = torch.zeros([num_steps] + list(w_opt.shape[1:]), dtype=torch.float32, device=device)\n",
        "    optimizer = torch.optim.Adam([w_opt] + list(noise_bufs.values()), betas=(0.9, 0.999), lr=initial_learning_rate)\n",
        "\n",
        "    # Init noise.\n",
        "    for buf in noise_bufs.values():\n",
        "        buf[:] = torch.randn_like(buf)\n",
        "        buf.requires_grad = True\n",
        "\n",
        "    for step in range(num_steps):\n",
        "        # Learning rate schedule.\n",
        "        t = step / num_steps\n",
        "        w_noise_scale = w_std * initial_noise_factor * max(0.0, 1.0 - t / noise_ramp_length) ** 2\n",
        "        lr_ramp = min(1.0, (1.0 - t) / lr_rampdown_length)\n",
        "        lr_ramp = 0.5 - 0.5 * np.cos(lr_ramp * np.pi)\n",
        "        lr_ramp = lr_ramp * min(1.0, t / lr_rampup_length)\n",
        "        lr = initial_learning_rate * lr_ramp\n",
        "        for param_group in optimizer.param_groups:\n",
        "            param_group['lr'] = lr\n",
        "\n",
        "        # Synth images from opt_w.\n",
        "        w_noise = torch.randn_like(w_opt) * w_noise_scale\n",
        "        ws = (w_opt + w_noise).repeat([1, G.mapping.num_ws, 1])\n",
        "        synth_images = G.synthesis(ws, noise_mode='const')\n",
        "\n",
        "        # Downsample image to 256x256 if it's larger than that. VGG was built for 224x224 images.\n",
        "        synth_images = (synth_images + 1) * (255/2)\n",
        "        if synth_images.shape[2] > 256:\n",
        "            synth_images = F.interpolate(synth_images, size=(256, 256), mode='area')\n",
        "\n",
        "        # Features for synth images.\n",
        "        synth_features = vgg16(synth_images, resize_images=False, return_lpips=True)\n",
        "        dist = (target_features - synth_features).square().sum()\n",
        "\n",
        "        # Noise regularization.\n",
        "        reg_loss = 0.0\n",
        "        for v in noise_bufs.values():\n",
        "            noise = v[None,None,:,:] # must be [1,1,H,W] for F.avg_pool2d()\n",
        "            while True:\n",
        "                reg_loss += (noise*torch.roll(noise, shifts=1, dims=3)).mean()**2\n",
        "                reg_loss += (noise*torch.roll(noise, shifts=1, dims=2)).mean()**2\n",
        "                if noise.shape[2] <= 8:\n",
        "                    break\n",
        "                noise = F.avg_pool2d(noise, kernel_size=2)\n",
        "        loss = dist + reg_loss * regularize_noise_weight\n",
        "\n",
        "        # Step\n",
        "        optimizer.zero_grad(set_to_none=True)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        logprint(f'step {step+1:>4d}/{num_steps}: dist {dist:<4.2f} loss {float(loss):<5.2f}')\n",
        "\n",
        "        # Save projected W for each optimization step.\n",
        "        w_out[step] = w_opt.detach()[0]\n",
        "\n",
        "        # Normalize noise.\n",
        "        with torch.no_grad():\n",
        "            for buf in noise_bufs.values():\n",
        "                buf -= buf.mean()\n",
        "                buf *= buf.square().mean().rsqrt()\n",
        "\n",
        "    return w_out.repeat([1, G.mapping.num_ws, 1])\n",
        "\n",
        "def get_perceptual_loss(synth_image, target_features):\n",
        "    # Downsample image to 256x256 if it's larger than that. VGG was built for 224x224 images.\n",
        "    synth_image = (synth_image + 1) * (255/2)\n",
        "    if synth_image.shape[2] > 256:\n",
        "        synth_image = F.interpolate(synth_image, size=(256, 256), mode='area')\n",
        "\n",
        "    # Features for synth images.\n",
        "    synth_features = vgg16(synth_image, resize_images=False, return_lpips=True)\n",
        "    return (target_features - synth_features).square().sum()\n",
        "\n",
        "def get_target_features(target):\n",
        "    target_images = target.unsqueeze(0).to(device).to(torch.float32)\n",
        "    if target_images.shape[2] > 256:\n",
        "        target_images = F.interpolate(target_images, size=(256, 256), mode='area')\n",
        "    return vgg16(target_images, resize_images=False, return_lpips=True)\n",
        "\n",
        "def slugify(value, allow_unicode=False):\n",
        "    \"\"\"\n",
        "    Taken from https://github.com/django/django/blob/master/django/utils/text.py\n",
        "    Convert to ASCII if 'allow_unicode' is False. Convert spaces or repeated\n",
        "    dashes to single dashes. Remove characters that aren't alphanumerics,\n",
        "    underscores, or hyphens. Convert to lowercase. Also strip leading and\n",
        "    trailing whitespace, dashes, and underscores.\n",
        "    \"\"\"\n",
        "    value = str(value)\n",
        "    if allow_unicode:\n",
        "        value = unicodedata.normalize('NFKC', value)\n",
        "    else:\n",
        "        value = unicodedata.normalize('NFKD', value).encode('ascii', 'ignore').decode('ascii')\n",
        "    value = re.sub(r'[^\\w\\s-]', '', value.lower())\n",
        "    return re.sub(r'[-\\s]+', '-', value).strip('-_')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown #**Model selection** 🎭\n",
        "\n",
        "\n",
        "#@markdown There are 4 pre-trained options to play with:\n",
        "#@markdown - FFHQ: Trained with human faces.\n",
        "#@markdown - MetFaces: Trained with paintings/portraits of human faces.\n",
        "#@markdown - AFHQv2: Trained with animal faces.\n",
        "#@markdown - Cosplay: Trained by [l4rz](https://twitter.com/l4rz) with cosplayer's faces.\n",
        "#@markdown - Wikiart: Trained by [Justin Pinkney](https://www.justinpinkney.com/) with the Wikiart 1024 dataset.\n",
        "#@markdown - Landscapes: Trained by [Justin Pinkney](https://www.justinpinkney.com/) with the LHQ dataset.\n",
        "\n",
        "\n",
        "#@markdown **Run this cell again if you change the model**.\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "base_url = \"https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/\"\n",
        "\n",
        "Model = 'FFHQ' #@param [\"FFHQ\", \"MetFaces\", \"AFHQv2\", \"cosplay\", \"Wikiart\", \"Landscapes\"]\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "model_name = {\n",
        "    \"FFHQ\": base_url + \"stylegan3-t-ffhqu-1024x1024.pkl\",\n",
        "    \"MetFaces\": base_url + \"stylegan3-r-metfacesu-1024x1024.pkl\",\n",
        "    \"AFHQv2\": base_url + \"stylegan3-t-afhqv2-512x512.pkl\",\n",
        "    \"cosplay\": \"https://l4rz.net/cosplayface-snapshot-stylegan3t-008000.pkl\",\n",
        "    \"Wikiart\": \"https://archive.org/download/wikiart-1024-stylegan3-t-17.2Mimg/wikiart-1024-stylegan3-t-17.2Mimg.pkl\",\n",
        "    \"Landscapes\": \"https://archive.org/download/lhq-256-stylegan3-t-25Mimg/lhq-256-stylegan3-t-25Mimg.pkl\"\n",
        "}\n",
        "\n",
        "network_url = model_name[Model]\n",
        "\n",
        "with open(fetch_model(network_url), 'rb') as fp:\n",
        "  G = pickle.load(fp)['G_ema'].to(device)\n",
        "\n",
        "zs = torch.randn([10000, G.mapping.z_dim], device=device)\n",
        "w_stds = G.mapping(zs, None).std(0)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "XUDCwVS6LOUc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown #**Parameters** ✍️\n",
        "#@markdown ---\n",
        "#@markdown `target_image_filename`: URL or an uploaded file\n",
        "\n",
        "target_image_filename = \"https://static.wikia.nocookie.net/fiveworldadventures/images/5/5f/Monkey_D._Luffy.png\"#@param {type:\"string\"}\n",
        "steps = 500#@param {type:\"number\"}\n",
        "\n",
        "if target_image_filename.lower().startswith(\"http\"):\n",
        "  target_image_filename = download_img(target_image_filename)\n",
        "target_pil = Image.open(target_image_filename).convert('RGB')\n",
        "w, h = target_pil.size\n",
        "s = min(w, h)\n",
        "target_pil = target_pil.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))\n",
        "target_pil = target_pil.resize((G.img_resolution, G.img_resolution), Image.LANCZOS)\n",
        "target_uint8 = np.array(target_pil, dtype=np.uint8)\n",
        "target_tensor = torch.tensor(target_uint8.transpose([2, 0, 1]), device=device)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "sCuEficELTvn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown #**Run the inversion** 🚀\n",
        "\n",
        "tf = Compose([\n",
        "  Resize(224),\n",
        "  lambda x: torch.clamp((x+1)/2,min=0,max=1),\n",
        "])\n",
        "\n",
        "seed = 14\n",
        "\n",
        "def run(timestring, projection_target):\n",
        "  torch.manual_seed(seed)\n",
        "\n",
        "  target_features = get_target_features(projection_target)\n",
        "\n",
        "  with torch.no_grad():\n",
        "    qs = []\n",
        "    losses = []\n",
        "    for _ in range(8):\n",
        "      q = (G.mapping(torch.randn([4,G.mapping.z_dim], device=device), None, truncation_psi=0.7) - G.mapping.w_avg) / w_stds\n",
        "      images = G.synthesis(q * w_stds + G.mapping.w_avg)\n",
        "      loss = get_perceptual_loss(images, target_features)\n",
        "      i = torch.argmin(loss)\n",
        "      qs.append(q[i])\n",
        "      losses.append(loss)\n",
        "    qs = torch.stack(qs)\n",
        "    losses = torch.stack(losses)\n",
        "    i = torch.argmin(losses)\n",
        "    q = qs[i].unsqueeze(0).requires_grad_()\n",
        "\n",
        "  # Sampling loop\n",
        "  q_ema = q\n",
        "  opt = torch.optim.AdamW([q], lr=0.03, betas=(0.0,0.999))\n",
        "  loop = tqdm(range(steps))\n",
        "  for i in loop:\n",
        "    opt.zero_grad()\n",
        "    w = q * w_stds\n",
        "    image = G.synthesis(w + G.mapping.w_avg, noise_mode='const')\n",
        "    loss = get_perceptual_loss(image, target_features)\n",
        "    loss.backward()\n",
        "    opt.step()\n",
        "    loop.set_postfix(loss=loss.item(), q_magnitude=q.std().item())\n",
        "\n",
        "    q_ema = q_ema * 0.9 + q * 0.1\n",
        "    image = G.synthesis(q_ema * w_stds + G.mapping.w_avg, noise_mode='const')\n",
        "\n",
        "    if i % 10 == 0:\n",
        "      display(TF.to_pil_image(tf(image)[0]))\n",
        "      print(f\"image {i}/{steps} | loss: {loss}\")\n",
        "    pil_image = TF.to_pil_image(image[0].add(1).div(2).clamp(0,1))\n",
        "    os.makedirs(f'samples/{timestring}', exist_ok=True)\n",
        "    pil_image.save(f'samples/{timestring}/{i:04}.jpg')\n",
        "\n",
        "try:\n",
        "  timestring = time.strftime('%Y%m%d%H%M%S')\n",
        "  run(timestring, target_tensor)\n",
        "except KeyboardInterrupt:\n",
        "  pass"
      ],
      "metadata": {
        "cellView": "form",
        "id": "H7qjLteTHnpo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown #**Generate video** 🎥\n",
        "\n",
        "#@markdown You can edit frame rate and stuff by double-clicking this tab.\n",
        "\n",
        "frames = os.listdir(f\"samples/{timestring}\")\n",
        "frames = len(list(filter(lambda filename: filename.endswith(\".jpg\"), frames))) #Get number of jpg generated\n",
        "\n",
        "init_frame = 1 #This is the frame where the video will start\n",
        "last_frame = frames #You can change i to the number of the last frame you want to generate. It will raise an error if that number of frames does not exist.\n",
        "\n",
        "min_fps = 10\n",
        "max_fps = 60\n",
        "\n",
        "total_frames = last_frame-init_frame\n",
        "\n",
        "#Desired video time in seconds\n",
        "video_length = 14 #@param {type:\"number\"}\n",
        "#Video filename\n",
        "video_name = \"\" #@param {type:\"string\"}\n",
        "video_name = slugify(video_name)\n",
        "\n",
        "if not video_name:\n",
        "  video_name = \"video\"\n",
        "# frames = []\n",
        "# tqdm.write('Generating video...')\n",
        "# for i in range(init_frame,last_frame): #\n",
        "#     filename = f\"samples/{timestring}/{i:04}.jpg\"\n",
        "#     frames.append(Image.open(filename))\n",
        "\n",
        "fps = np.clip(total_frames/video_length,min_fps,max_fps)\n",
        "\n",
        "print(\"Generating video...\")\n",
        "!ffmpeg -r {fps} -i samples/{timestring}/%04d.jpg -c:v libx264 -vf fps={fps} -pix_fmt yuv420p samples/{video_name}.mp4 -frames:v {total_frames}\n",
        "\n",
        "# from subprocess import Popen, PIPE\n",
        "# p = Popen(['ffmpeg', '-y', '-f', 'image2pipe', '-vcodec', 'png', '-r', str(fps), '-i', '-', '-vcodec', 'libx264', '-r', str(fps), '-pix_fmt', 'yuv420p', '-crf', '17', '-preset', 'veryslow', f'samples/{video_name}.mp4'], stdin=PIPE)\n",
        "# for im in tqdm(frames):\n",
        "#     im.save(p.stdin, 'PNG')\n",
        "# p.stdin.close()\n",
        "\n",
        "print(\"The video is ready\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "1Z__XJ9HuaU5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown #**Download video** 📀\n",
        "\n",
        "#Video filename\n",
        "\n",
        "files.download(f\"samples/{video_name}.mp4\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "2FocdVISvgCX"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}