notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import os\n",
    "from typing import NamedTuple\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision\n",
    "from torch.utils import data\n",
    "from PIL import Image\n",
    "\n",
    "import model\n",
    "from tensor_transforms import convert_to_coord_format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FFHQ256Arguments(NamedTuple):\n",
    "    \"\"\"CIPSskip for FFHQ-256\"\"\"\n",
    "    Generator = 'CIPSskip'\n",
    "    output_dir = 'ffhq256_g_ema.pt'\n",
    "    out_path = 'checkpoint'\n",
    "    size = 256\n",
    "    coords_size = 256\n",
    "    fc_dim = 512 \n",
    "    latent = 512\n",
    "    style_dim = 512\n",
    "    n_mlp = 8\n",
    "    activation = None\n",
    "    channel_multiplier = 2\n",
    "    ckpt = os.path.join(out_path, output_dir)\n",
    "    coords_integer_values = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FFHQ1024Arguments(NamedTuple):\n",
    "    \"\"\"CIPSskip-progressive for FFHQ-1024\"\"\"\n",
    "    Generator = 'CIPSskip'\n",
    "    output_dir = 'ffhq1024_g_ema.pt'\n",
    "    out_path = 'checkpoint'\n",
    "    size = 256\n",
    "    coords_size = 1024\n",
    "    fc_dim = 512 \n",
    "    latent = 512\n",
    "    style_dim = 512\n",
    "    n_mlp = 8\n",
    "    activation = None\n",
    "    channel_multiplier = 2\n",
    "    ckpt = os.path.join(out_path, output_dir)\n",
    "    coords_integer_values = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Churches256Arguments(NamedTuple):\n",
    "    \"\"\"CIPSskip for LSUN-Churches-256\"\"\"\n",
    "    Generator = 'CIPSskip'\n",
    "    output_dir = 'churches_g_ema.pt'\n",
    "    out_path = 'checkpoint'\n",
    "    size = 256\n",
    "    coords_size = 256\n",
    "    fc_dim = 512 \n",
    "    latent = 512\n",
    "    style_dim = 512\n",
    "    n_mlp = 8\n",
    "    activation = None\n",
    "    channel_multiplier = 2\n",
    "    ckpt = os.path.join(out_path, output_dir)\n",
    "    coords_integer_values = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Lanscapes256Arguments(NamedTuple):\n",
    "    \"\"\"CIPSres for Landscapes-256\"\"\"\n",
    "    Generator = 'CIPSres'\n",
    "    output_dir = 'landscapes_g_ema.pt'\n",
    "    out_path = 'checkpoint'\n",
    "    size = 256\n",
    "    coords_size = 256\n",
    "    fc_dim = 512 \n",
    "    latent = 512\n",
    "    style_dim = 512\n",
    "    n_mlp = 8\n",
    "    activation = None\n",
    "    channel_multiplier = 2\n",
    "    ckpt = os.path.join(out_path, output_dir)\n",
    "    coords_integer_values = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = Lanscapes256Arguments()\n",
    "device = 'cuda'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Generator = getattr(model, args.Generator)\n",
    "g_ema = Generator(size=args.size, hidden_size=args.fc_dim, style_dim=args.latent, n_mlp=args.n_mlp,\n",
    "                  activation=args.activation, channel_multiplier=args.channel_multiplier,\n",
    "                  ).to(device)\n",
    "g_ema.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt = torch.load(path)\n",
    "g_ema.load_state_dict(ckpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image(tensor, nrow=8, padding=2,\n",
    "               normalize=False, range=None, scale_each=False, pad_value=0):\n",
    "    \n",
    "    grid = torchvision.utils.make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,\n",
    "                     normalize=normalize, range=range, scale_each=scale_each)\n",
    "    ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()\n",
    "    im = Image.fromarray(ndarr)\n",
    "    return im"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finding mean for truncation trick "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_sample = 1\n",
    "sample_z = torch.randn(n_sample, args.latent, device=device)\n",
    "converted_full = convert_to_coord_format(sample_z.size(0), args.coords_size, args.coords_size, device,\n",
    "                                         integer_values=args.coords_integer_values)\n",
    "\n",
    "latents = []\n",
    "samples = []\n",
    "with torch.no_grad():\n",
    "    for _ in range(100):\n",
    "        sample_z = torch.randn(n_sample, args.latent, device=device)\n",
    "        sample, latent = g_ema(converted_full, [sample_z], return_latents=True)\n",
    "        latents.append(latent.cpu())\n",
    "        samples.append(sample.cpu())\n",
    "\n",
    "samples = torch.cat(samples, 0)\n",
    "latents = torch.cat(latents, 0)\n",
    "\n",
    "truncation_latent = latents.mean(0).cuda()\n",
    "\n",
    "print('truncation_latent', truncation_latent.shape)\n",
    "assert len(truncation_latent.shape)==1 and truncation_latent.size(0) == 512, 'smt wrong'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sampling with truncation trick "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_sample = 4\n",
    "sample_z = torch.randn(n_sample, args.latent, device=device)\n",
    "converted_full = convert_to_coord_format(sample_z.size(0), args.coords_size, args.coords_size, device,\n",
    "                                         integer_values=args.coords_integer_values)\n",
    "\n",
    "with torch.no_grad():\n",
    "    style = g_ema.style(sample_z)\n",
    "    sample, _ = g_ema(converted_full, [style], \n",
    "                      truncation=0.6,\n",
    "                      truncation_latent=truncation_latent,\n",
    "                      input_is_latent=True,)\n",
    "    \n",
    "im = get_image(sample,                        \n",
    "                nrow=int(n_sample ** 0.5),\n",
    "                normalize=True,\n",
    "                range=(-1, 1),)\n",
    "\n",
    "im"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}