From 7c5e365f74cbe31f09029f531948f9073be62a12 Mon Sep 17 00:00:00 2001 From: Nate Baer Date: Sat, 12 Mar 2022 23:48:33 -0800 Subject: [PATCH] Check if using secondary model before loading it --- Disco_Diffusion.ipynb | 5 ++--- disco.py | 5 ++--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/Disco_Diffusion.ipynb b/Disco_Diffusion.ipynb index 881e482..9a6ff1d 100644 --- a/Disco_Diffusion.ipynb +++ b/Disco_Diffusion.ipynb @@ -2321,15 +2321,14 @@ " 'use_scale_shift_norm': True,\n", " })\n", "\n", - "secondary_model_ver = 2\n", "model_default = model_config['image_size']\n", "\n", "\n", "\n", - "if secondary_model_ver == 2:\n", + "if use_secondary_model:\n", " secondary_model = SecondaryDiffusionImageNet2()\n", " secondary_model.load_state_dict(torch.load(f'{model_path}/secondary_model_imagenet_2.pth', map_location='cpu'))\n", - "secondary_model.eval().requires_grad_(False).to(device)\n", + " secondary_model.eval().requires_grad_(False).to(device)\n", "\n", "clip_models = []\n", "if ViTB32 is True: clip_models.append(clip.load('ViT-B/32', jit=False)[0].eval().requires_grad_(False).to(device)) \n", diff --git a/disco.py b/disco.py index e06ea5e..58ad5b5 100644 --- a/disco.py +++ b/disco.py @@ -2199,15 +2199,14 @@ elif diffusion_model == '256x256_diffusion_uncond': 'use_scale_shift_norm': True, }) -secondary_model_ver = 2 model_default = model_config['image_size'] -if secondary_model_ver == 2: +if use_secondary_model: secondary_model = SecondaryDiffusionImageNet2() secondary_model.load_state_dict(torch.load(f'{model_path}/secondary_model_imagenet_2.pth', map_location='cpu')) -secondary_model.eval().requires_grad_(False).to(device) + secondary_model.eval().requires_grad_(False).to(device) clip_models = [] if ViTB32 is True: clip_models.append(clip.load('ViT-B/32', jit=False)[0].eval().requires_grad_(False).to(device))