generate.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import data\n",
    "from utils import viz\n",
    "import models\n",
    "import torch\n",
    "\n",
    "model = models.HumanGAN.load_from_checkpoint(\n",
    "    \"checkpoints/ours/ours.ckpt\"\n",
    ")\n",
    "generator = model.generator_ema.eval().cuda().requires_grad_(False)\n",
    "discriminator = model.discriminator.eval().cuda().requires_grad_(False)\n",
    "\n",
    "dataset = data.HumansDataset(\n",
    "    \"humans_in_context\", deterministic=True, split=\"test\", num_frames=1, spacing=1,\n",
    ")\n",
    "\n",
    "seed = 0\n",
    "num_poses = 16\n",
    "num_scenes = 16\n",
    "scenes_per_row = 4\n",
    "\n",
    "print(f\"Next seed: {seed + num_poses * (num_scenes + 1)}\")\n",
    "\n",
    "for pose_idx in range(num_poses):\n",
    "    pose_seed = seed + pose_idx * (num_scenes + 1)\n",
    "    \n",
    "    _, keypoints = dataset.sample(seed=pose_seed)\n",
    "    keypoints = keypoints[None].cuda()\n",
    "    \n",
    "    skeletons = viz.visualize_skeletons(keypoints, dataset)\n",
    "    images = []\n",
    "    scene_seeds = []\n",
    "    \n",
    "    for scene_idx in range(num_scenes):\n",
    "        scene_seed = pose_seed + scene_idx + 1\n",
    "        scene_seeds.append(scene_seed)\n",
    "        \n",
    "        styles, frame_keypoints = viz.sample_styles(\n",
    "            keypoints, generator, seed=scene_seed, samples=1, strength=0.25\n",
    "        )\n",
    "        image = generator.synthesis_network(styles, frame_keypoints)\n",
    "        scene = generator.synthesis_network(styles, 0 * frame_keypoints)\n",
    "        \n",
    "        images.append(scene.cpu())\n",
    "        images.append(image.cpu())\n",
    "    \n",
    "    images = torch.cat(images)\n",
    "    viz.display_grid(images.cpu(), nrow=2 * scenes_per_row)\n",
    "    scene_seeds = \"\\n\".join([str(scene_seeds[i:i+scenes_per_row]) for i in range(0, num_scenes, scenes_per_row)])\n",
    "    print(f\"{pose_seed}\\n{scene_seeds}\")\n",
    "import data\n",
    "from utils import viz\n",
    "import models\n",
    "import torch\n",
    "\n",
    "model = models.HumanGAN.load_from_checkpoint(\n",
    "    \"PATH_TO_PRETRAINED/1x2sxwu9_step=00899999.ckpt\"\n",
    ")\n",
    "generator = model.generator_ema.eval().cuda().requires_grad_(False)\n",
    "discriminator = model.discriminator.eval().cuda().requires_grad_(False)\n",
    "\n",
    "dataset = data.HumansDataset(\n",
    "    \"PATH_TO_DATASET\", deterministic=True, split=\"test\", num_frames=1, spacing=1,\n",
    ")\n",
    "\n",
    "seed = 0\n",
    "num_poses = 16\n",
    "num_scenes = 16\n",
    "scenes_per_row = 4\n",
    "\n",
    "print(f\"Next seed: {seed + num_poses * (num_scenes + 1)}\")\n",
    "\n",
    "for pose_idx in range(num_poses):\n",
    "    pose_seed = seed + pose_idx * (num_scenes + 1)\n",
    "    \n",
    "    _, keypoints = dataset.sample(seed=pose_seed)\n",
    "    keypoints = keypoints[None].cuda()\n",
    "    \n",
    "    skeletons = viz.visualize_skeletons(keypoints, dataset)\n",
    "    images = []\n",
    "    scene_seeds = []\n",
    "    \n",
    "    for scene_idx in range(num_scenes):\n",
    "        scene_seed = pose_seed + scene_idx + 1\n",
    "        scene_seeds.append(scene_seed)\n",
    "        \n",
    "        styles, frame_keypoints = viz.sample_styles(\n",
    "            keypoints, generator, seed=scene_seed, samples=1, strength=0.25\n",
    "        )\n",
    "        image = generator.synthesis_network(styles, frame_keypoints)\n",
    "        scene = generator.synthesis_network(styles, 0 * frame_keypoints)\n",
    "        \n",
    "        images.append(scene.cpu())\n",
    "        images.append(image.cpu())\n",
    "    \n",
    "    images = torch.cat(images)\n",
    "    viz.display_grid(images.cpu(), nrow=2 * scenes_per_row)\n",
    "    scene_seeds = \"\\n\".join([str(scene_seeds[i:i+scenes_per_row]) for i in range(0, num_scenes, scenes_per_row)])\n",
    "    print(f\"{pose_seed}\\n{scene_seeds}\")\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}