Update Reconstructions.ipynb

vd_pipe needed to be float16
MedARC-AI · Oct 21, 2023 · 4d02ab3 · 4d02ab3
1 parent 59d95dd
commit 4d02ab3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/Reconstructions.ipynb b/src/Reconstructions.ipynb
@@ -278,12 +278,12 @@
     "from diffusers import VersatileDiffusionDualGuidedPipeline, UniPCMultistepScheduler\n",
     "from diffusers.models import DualTransformer2DModel\n",
     "try:\n",
-    "    vd_pipe =  VersatileDiffusionDualGuidedPipeline.from_pretrained(vd_cache_dir).to(device)\n",
+    "    vd_pipe =  VersatileDiffusionDualGuidedPipeline.from_pretrained(vd_cache_dir).to(device).to(torch.float16)\n",
     "except:\n",
     "    print(\"Downloading Versatile Diffusion to\", vd_cache_dir)\n",
     "    vd_pipe =  VersatileDiffusionDualGuidedPipeline.from_pretrained(\n",
     "            \"shi-labs/versatile-diffusion\",\n",
-    "            cache_dir = vd_cache_dir).to(device)\n",
+    "            cache_dir = vd_cache_dir).to(device).to(torch.float16)\n",
     "vd_pipe.image_unet.eval()\n",
     "vd_pipe.vae.eval()\n",
     "vd_pipe.image_unet.requires_grad_(False)\n",