From e7f7a2255d1423ba718df1fdecacf7e4619687cc Mon Sep 17 00:00:00 2001
From: Aleksey Morozov <36787333+amrzv@users.noreply.github.com>
Date: Sun, 26 Jun 2022 17:16:51 +0300
Subject: [PATCH] Add image downloading in case colab env

---
 notebooks/inpaint.ipynb | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/notebooks/inpaint.ipynb b/notebooks/inpaint.ipynb
index 1191336..961b1d7 100644
--- a/notebooks/inpaint.ipynb
+++ b/notebooks/inpaint.ipynb
@@ -47,6 +47,17 @@
     "device = th.device('cpu' if not has_cuda else 'cuda')"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "if 'COLAB_GPU' in os.environ:  # Downloading image in case running in colab env\n",
+    "  !wget https://raw.githubusercontent.com/openai/glide-text2im/main/notebooks/grass.png"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -57,7 +68,7 @@
     "options = model_and_diffusion_defaults()\n",
     "options['inpaint'] = True\n",
     "options['use_fp16'] = has_cuda\n",
-    "options['timestep_respacing'] = '100' # use 100 diffusion steps for fast sampling\n",
+    "options['timestep_respacing'] = '100'  # use 100 diffusion steps for fast sampling\n",
     "model, diffusion = create_model_and_diffusion(**options)\n",
     "model.eval()\n",
     "if has_cuda:\n",
@@ -77,7 +88,7 @@
     "options_up = model_and_diffusion_defaults_upsampler()\n",
     "options_up['inpaint'] = True\n",
     "options_up['use_fp16'] = has_cuda\n",
-    "options_up['timestep_respacing'] = 'fast27' # use 27 diffusion steps for very fast sampling\n",
+    "options_up['timestep_respacing'] = 'fast27'  # use 27 diffusion steps for very fast sampling\n",
     "model_up, diffusion_up = create_model_and_diffusion(**options_up)\n",
     "model_up.eval()\n",
     "if has_cuda:\n",
@@ -95,10 +106,11 @@
    "source": [
     "def show_images(batch: th.Tensor):\n",
     "    \"\"\" Display a batch of images inline. \"\"\"\n",
-    "    scaled = ((batch + 1)*127.5).round().clamp(0,255).to(th.uint8).cpu()\n",
+    "    scaled = ((batch + 1)*127.5).round().clamp(0, 255).to(th.uint8).cpu()\n",
     "    reshaped = scaled.permute(2, 0, 3, 1).reshape([batch.shape[2], -1, 3])\n",
     "    display(Image.fromarray(reshaped.numpy()))\n",
     "\n",
+    "\n",
     "def read_image(path: str, size: int = 256) -> Tuple[th.Tensor, th.Tensor]:\n",
     "    pil_img = Image.open(path).convert('RGB')\n",
     "    pil_img = pil_img.resize((size, size), resample=Image.BICUBIC)\n",
@@ -173,6 +185,7 @@
     "    inpaint_mask=source_mask_64.repeat(full_batch_size, 1, 1, 1).to(device),\n",
     ")\n",
     "\n",
+    "\n",
     "# Create an classifier-free guidance sampling function\n",
     "def model_fn(x_t, ts, **kwargs):\n",
     "    half = x_t[: len(x_t) // 2]\n",
@@ -184,6 +197,7 @@
     "    eps = th.cat([half_eps, half_eps], dim=0)\n",
     "    return th.cat([eps, rest], dim=1)\n",
     "\n",
+    "\n",
     "def denoised_fn(x_start):\n",
     "    # Force the model to have the exact right x_start predictions\n",
     "    # for the part of the image which is known.\n",
@@ -192,6 +206,7 @@
     "        + model_kwargs['inpaint_image'] * model_kwargs['inpaint_mask']\n",
     "    )\n",
     "\n",
+    "\n",
     "# Sample from the base model.\n",
     "model.del_cache()\n",
     "samples = diffusion.p_sample_loop(\n",
@@ -245,6 +260,7 @@
     "    inpaint_mask=source_mask_256.repeat(batch_size, 1, 1, 1).to(device),\n",
     ")\n",
     "\n",
+    "\n",
     "def denoised_fn(x_start):\n",
     "    # Force the model to have the exact right x_start predictions\n",
     "    # for the part of the image which is known.\n",
@@ -253,6 +269,7 @@
     "        + model_kwargs['inpaint_image'] * model_kwargs['inpaint_mask']\n",
     "    )\n",
     "\n",
+    "\n",
     "# Sample from the base model.\n",
     "model_up.del_cache()\n",
     "up_shape = (batch_size, 3, options_up[\"image_size\"], options_up[\"image_size\"])\n",