Add option to run with CPU

Disco_Diffusion.ipynb

@@ -459,10 +459,13 @@
         "id": "InstallDeps"
       },
       "source": [
-        "#@title ### 1.3 Install and import dependencies\n",
+        "#@title ### 1.3 Install, import dependencies and set up runtime devices\n",
         "\n",
         "import pathlib, shutil, os, sys\n",
         "\n",
+        "#@markdown Check this if you want to use CPU\n",
+        "useCPU = False #@param {type:\"boolean\"}\n",
+        "\n",
         "if not is_colab:\n",
         "  # If running locally, there's a good chance your env will need this in order to not crash upon np.matmul() or similar operations.\n",
         "  os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'\n",
@@ -597,9 +600,10 @@
         "print('Using device:', DEVICE)\n",
         "device = DEVICE # At least one of the modules expects this name..\n",
         "\n",
-        "if torch.cuda.get_device_capability(DEVICE) == (8,0): ## A100 fix thanks to Emad\n",
+        "if not useCPU:\n",
-        "  print('Disabling CUDNN for A100 gpu', file=sys.stderr)\n",
+        "  if torch.cuda.get_device_capability(DEVICE) == (8,0): ## A100 fix thanks to Emad\n",
-        "  torch.backends.cudnn.enabled = False"
+        "    print('Disabling CUDNN for A100 gpu', file=sys.stderr)\n",
+        "    torch.backends.cudnn.enabled = False"
       ],
       "outputs": [],
       "execution_count": null
@@ -1864,7 +1868,7 @@
         "        'num_res_blocks': 2,\n",
         "        'resblock_updown': True,\n",
         "        'use_checkpoint': use_checkpoint,\n",
-        "        'use_fp16': True,\n",
+        "        'use_fp16': not useCPU,\n",
         "        'use_scale_shift_norm': True,\n",
         "    })\n",
         "elif diffusion_model == '256x256_diffusion_uncond':\n",
@@ -1882,7 +1886,7 @@
         "        'num_res_blocks': 2,\n",
         "        'resblock_updown': True,\n",
         "        'use_checkpoint': use_checkpoint,\n",
-        "        'use_fp16': True,\n",
+        "        'use_fp16': not useCPU,\n",
         "        'use_scale_shift_norm': True,\n",
         "    })\n",
         "\n",

disco.py

@@ -432,10 +432,13 @@ else:
 # !!    "cellView": "form",
 # !!    "id": "InstallDeps"
 # !! }}
-#@title ### 1.3 Install and import dependencies
+#@title ### 1.3 Install, import dependencies and set up runtime devices
 
 import pathlib, shutil, os, sys
 
+# Check this if you want to use CPU
+useCPU = False
+
 if not is_colab:
   # If running locally, there's a good chance your env will need this in order to not crash upon np.matmul() or similar operations.
   os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'
@@ -570,9 +573,10 @@ DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 print('Using device:', DEVICE)
 device = DEVICE # At least one of the modules expects this name..
 
-if torch.cuda.get_device_capability(DEVICE) == (8,0): ## A100 fix thanks to Emad
+if not useCPU:
-  print('Disabling CUDNN for A100 gpu', file=sys.stderr)
+  if torch.cuda.get_device_capability(DEVICE) == (8,0): ## A100 fix thanks to Emad
-  torch.backends.cudnn.enabled = False
+    print('Disabling CUDNN for A100 gpu', file=sys.stderr)
+    torch.backends.cudnn.enabled = False
 
 # %%
 # !! {"metadata": {
@@ -1817,7 +1821,7 @@ if diffusion_model == '512x512_diffusion_uncond_finetune_008100':
         'num_res_blocks': 2,
         'resblock_updown': True,
         'use_checkpoint': use_checkpoint,
-        'use_fp16': True,
+        'use_fp16': not useCPU,
         'use_scale_shift_norm': True,
     })
 elif diffusion_model == '256x256_diffusion_uncond':
@@ -1835,7 +1839,7 @@ elif diffusion_model == '256x256_diffusion_uncond':
         'num_res_blocks': 2,
         'resblock_updown': True,
         'use_checkpoint': use_checkpoint,
-        'use_fp16': True,
+        'use_fp16': not useCPU,
         'use_scale_shift_norm': True,
     })