Check if using secondary model before loading it

Disco_Diffusion.ipynb

@@ -2321,12 +2321,11 @@
         "        '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",

disco.py

@@ -2199,12 +2199,11 @@ 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)