Upload usage.ipynb

usage.ipynb

@@ -0,0 +1,116 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import io\n",
+    "import os, sys\n",
+    "import requests\n",
+    "import PIL\n",
+    "\n",
+    "import torch\n",
+    "import torchvision.transforms as T\n",
+    "import torchvision.transforms.functional as TF\n",
+    "\n",
+    "from dall_e          import map_pixels, unmap_pixels, load_model\n",
+    "from IPython.display import display, display_markdown\n",
+    "\n",
+    "target_image_size = 256\n",
+    "\n",
+    "def download_image(url):\n",
+    "    resp = requests.get(url)\n",
+    "    resp.raise_for_status()\n",
+    "    return PIL.Image.open(io.BytesIO(resp.content))\n",
+    "\n",
+    "def preprocess(img):\n",
+    "    s = min(img.size)\n",
+    "    \n",
+    "    if s < target_image_size:\n",
+    "        raise ValueError(f'min dim for image {s} < {target_image_size}')\n",
+    "        \n",
+    "    r = target_image_size / s\n",
+    "    s = (round(r * img.size[1]), round(r * img.size[0]))\n",
+    "    img = TF.resize(img, s, interpolation=PIL.Image.LANCZOS)\n",
+    "    img = TF.center_crop(img, output_size=2 * [target_image_size])\n",
+    "    img = torch.unsqueeze(T.ToTensor()(img), 0)\n",
+    "    return map_pixels(img)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# This can be changed to a GPU, e.g. 'cuda:0'.\n",
+    "dev = torch.device('cpu')\n",
+    "\n",
+    "# For faster load times, download these files locally and use the local paths instead.\n",
+    "enc = load_model(\"https://cdn.openai.com/dall-e/encoder.pkl\", dev)\n",
+    "dec = load_model(\"https://cdn.openai.com/dall-e/decoder.pkl\", dev)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x = preprocess(download_image('https://assets.bwbx.io/images/users/iqjWHBFdfxIU/iKIWgaiJUtss/v2/1000x-1.jpg'))\n",
+    "display_markdown('Original image:')\n",
+    "display(T.ToPILImage(mode='RGB')(x[0]))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch.nn.functional as F\n",
+    "\n",
+    "z_logits = enc(x)\n",
+    "z = torch.argmax(z_logits, axis=1)\n",
+    "z = F.one_hot(z, num_classes=enc.vocab_size).permute(0, 3, 1, 2).float()\n",
+    "\n",
+    "x_stats = dec(z).float()\n",
+    "x_rec = unmap_pixels(torch.sigmoid(x_stats[:, :3]))\n",
+    "x_rec = T.ToPILImage(mode='RGB')(x_rec[0])\n",
+    "\n",
+    "display_markdown('Reconstructed image:')\n",
+    "display(x_rec)"
+   ]
+  },
+  {
+   "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.9.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}