Update app.py

app.py

@@ -34,14 +34,14 @@ from diffusers import StableDiffusionPipeline
 # On a GPU, it should be under a minute.
 
 has_cuda = False
-gpu = th.device('cpu' if not th.cuda.is_available() else 'cuda')
+device = 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(gpu)
+).to(cpu)
 
 # Create base model.
 timestep_respacing = 100  # @param{type: 'number'}
@@ -77,7 +77,6 @@ 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)
 
@@ -245,7 +244,6 @@ print('total clevr_pos parameters', sum(x.numel() for x in clevr_model.parameter
 
 
 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
@@ -302,10 +300,22 @@ def stable_diffusion_compose(prompt, scale, steps):
 def compose(prompt, version, guidance_scale, steps):
     with th.no_grad():
         if version == 'GLIDE':
+            clevr_model.to(cpu)
+            pipe.to(cpu)
+            model.to(device)
+            model_up.to(device)
             return compose_language_descriptions(prompt, guidance_scale, steps)
         elif version == 'Stable_Diffusion_1v_4':
+            model.to(cpu)
+            model_up.to(cpu)
+            clevr_model.to(cpu)
+            pipe.to(device)
             return stable_diffusion_compose(prompt, guidance_scale, steps)
         else:
+            model.to(cpu)
+            model_up.to(cpu)
+            pipe.to(cpu)
+            clevr_model.to(device)
             return compose_clevr_objects(prompt, guidance_scale, steps)