Update app.py

app.py

@@ -59,19 +59,19 @@ print('loading models')
 ae_config = OmegaConf.load(ae_config_path)
 ae_model = ldm.models.autoencoder.AutoencoderKL(**ae_config.model.params)
 ae_model.eval().requires_grad_(False).to(device)
-ae_model.load_state_dict(torch.load(ae_model_path), weights_only=True)
+ae_model.load_state_dict(torch.load(ae_model_path, weights_only=True))
 n_ch, side_y, side_x = 4, 32, 32
 
 # diffusion model
 model = train.DiffusionModel(192, [1,1,2,2], autoencoder_scale=torch.tensor(4.3084))
-model.load_state_dict(torch.load(checkpoint, map_location='cpu'), weights_only=True)
+model.load_state_dict(torch.load(checkpoint, map_location='cpu', weights_only=True))
 model = model.to(device).eval().requires_grad_(False)
 
 # CLOOB
 cloob_config = pretrained.get_config('cloob_laion_400m_vit_b_16_16_epochs')
 cloob = model_pt.get_pt_model(cloob_config)
 checkpoint = pretrained.download_checkpoint(cloob_config)
-cloob.load_state_dict(model_pt.get_pt_params(cloob_config, checkpoint), weights_only=True)
+cloob.load_state_dict(model_pt.get_pt_params(cloob_config, checkpoint))
 cloob.eval().requires_grad_(False).to(device)