From c9821db89236178e9a90b188a44bd38b89c10c07 Mon Sep 17 00:00:00 2001 From: OodavidsinoO Date: Thu, 5 May 2022 10:12:06 +0800 Subject: [PATCH] Add option to run with CPU --- Disco_Diffusion.ipynb | 16 ++++++++++------ disco.py | 16 ++++++++++------ 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/Disco_Diffusion.ipynb b/Disco_Diffusion.ipynb index 6e77387..17493be 100644 --- a/Disco_Diffusion.ipynb +++ b/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", - " print('Disabling CUDNN for A100 gpu', file=sys.stderr)\n", - " torch.backends.cudnn.enabled = False" + "if not useCPU:\n", + " if torch.cuda.get_device_capability(DEVICE) == (8,0): ## A100 fix thanks to Emad\n", + " 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", diff --git a/disco.py b/disco.py index e6db74d..254c57a 100644 --- a/disco.py +++ b/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 - print('Disabling CUDNN for A100 gpu', file=sys.stderr) - torch.backends.cudnn.enabled = False +if not useCPU: + if torch.cuda.get_device_capability(DEVICE) == (8,0): ## A100 fix thanks to Emad + 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, })