Update app.py

app.py

@@ -33,15 +33,15 @@ from diffusers import StableDiffusionPipeline
 # On CPU, generating one sample may take on the order of 20 minutes.
 # On a GPU, it should be under a minute.
 
-has_cuda = th.cuda.is_available()
-device = th.device('cpu' if not has_cuda else 'cuda')
-print(device)
+has_cuda = False
+gpu = th.device('cpu' if not th.cuda.is_available() else 'cuda')
+cpu = th.device('cpu')
 
 # iniatilize stable diffusion model
 pipe = StableDiffusionPipeline.from_pretrained(
     "CompVis/stable-diffusion-v1-4",
     use_auth_token='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'
-).to(device)
+).to(gpu)
 
 # Create base model.
 timestep_respacing = 100  # @param{type: 'number'}
@@ -52,8 +52,8 @@ model, diffusion = create_model_and_diffusion(**options)
 model.eval()
 if has_cuda:
     model.convert_to_fp16()
-model.to(th.device('cpu'))
-model.load_state_dict(load_checkpoint('base', th.device('cpu'))
+model.to(cpu)
+model.load_state_dict(load_checkpoint('base', 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(th.device('cpu'))
-model_up.load_state_dict(load_checkpoint('upsample', th.device('cpu')))
+model_up.to(cpu)
+model_up.load_state_dict(load_checkpoint('upsample', cpu))
 print('total upsampler parameters', sum(x.numel() for x in model_up.parameters()))
 
 
@@ -294,7 +294,7 @@ def compose_clevr_objects(prompt, guidance_scale, steps):
 
 
 def stable_diffusion_compose(prompt, scale, steps):
-    with autocast('cpu' if not has_cuda else 'cuda'):
+    with autocast('cpu' if not th.cuda.is_available() else 'cuda'):
         image = pipe(prompt, guidance_scale=scale, num_inference_steps=steps)["sample"][0]
         return image