fix

app.py

@@ -51,15 +51,15 @@ def predict(dict, prompt="", negative_prompt="", guidance_scale=7.5, steps=20, s
     
     pipe.vae.to('cuda')
     pipe.unet.to('cuda')
+
     output = pipe(prompt = prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps), strength=strength)
+
     pipe.vae.to('cpu')
     pipe.unet.to('cpu')
-
     gc.collect()
     torch.cuda.empty_cache()
 
     comparing_unet.to('cuda')
-    comparing_vae.to('cuda')
     comparing_controlnet.to('cuda')
 
     image = TF.to_tensor(dict["image"].convert("RGB").resize((1024, 1024)))
@@ -77,12 +77,19 @@ def predict(dict, prompt="", negative_prompt="", guidance_scale=7.5, steps=20, s
         prompts=prompt, negative_prompts=negative_prompt, images=image, guidance_scale=guidance_scale, sigmas=sigmas, timesteps=timesteps,
         text_encoder_one=pipe.text_encoder, text_encoder_two=pipe.text_encoder_2, unet=comparing_unet, controlnet=comparing_controlnet
     )
-    out = comparing_vae.output_tensor_to_pil(comparing_vae.decode(out))
 
     comparing_unet.to('cpu')
-    comparing_vae.to('cpu')
     comparing_controlnet.to('cpu')
 
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    comparing_vae.to('cuda')
+
+    out = comparing_vae.output_tensor_to_pil(comparing_vae.decode(out))
+
+    comparing_vae.to('cpu')
+
     gc.collect()
     torch.cuda.empty_cache()