Update app.py

app.py

@@ -52,8 +52,8 @@ model, diffusion = create_model_and_diffusion(**options)
 model.eval()
 if has_cuda:
     model.convert_to_fp16()
-model.to(device)
-model.load_state_dict(load_checkpoint('base', device))
+model.to(th.device('cpu'))
+model.load_state_dict(load_checkpoint('base', th.device('cpu'))
 print('total base parameters', sum(x.numel() for x in model.parameters()))
 
 # Create upsampler model.
@@ -64,8 +64,8 @@ model_up, diffusion_up = create_model_and_diffusion(**options_up)
 model_up.eval()
 if has_cuda:
     model_up.convert_to_fp16()
-model_up.to(device)
-model_up.load_state_dict(load_checkpoint('upsample', device))
+model_up.to(th.device('cpu'))
+model_up.load_state_dict(load_checkpoint('upsample', th.device('cpu')))
 print('total upsampler parameters', sum(x.numel() for x in model_up.parameters()))
 
 
@@ -77,6 +77,7 @@ def show_images(batch: th.Tensor):
 
 
 def compose_language_descriptions(prompt, guidance_scale, steps):
+    device = th.device('cpu')
     options['timestep_respacing'] = str(steps)
     _, diffusion = create_model_and_diffusion(**options)
 
@@ -238,12 +239,13 @@ clevr_model.eval()
 if has_cuda:
     clevr_model.convert_to_fp16()
 
-clevr_model.to(device)
-clevr_model.load_state_dict(th.load(download_model('clevr_pos'), device))
+clevr_model.to(th.device('cpu'))
+clevr_model.load_state_dict(th.load(download_model('clevr_pos'), th.device('cpu')))
 print('total clevr_pos parameters', sum(x.numel() for x in clevr_model.parameters()))
 
 
 def compose_clevr_objects(prompt, guidance_scale, steps):
+    device = th.device('cpu')
     coordinates = [[float(x.split(',')[0].strip()), float(x.split(',')[1].strip())]
                    for x in prompt.split('|')]
     coordinates += [[-1, -1]]  # add unconditional score label