Update app.py

app.py

@@ -29,7 +29,7 @@ pipe_i2i.enable_xformers_memory_efficient_attention()
 
 upscaler = StableDiffusionLatentUpscalePipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16)
 upscaler.enable_xformers_memory_efficient_attention()
-upscaler.to("cuda")
+
 
 if torch.cuda.is_available():
   pipe = pipe.to("cuda")
@@ -137,6 +137,7 @@ def auto_prompt_correction(prompt_ui,neg_prompt_ui,cool_japan_type_ui,disable_au
     return prompt,neg_prompt
     
 def txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator,superreso=False):
+    global pipe, upscaler
     if(superreso):
       low_res_latents = pipe(
         prompt,
@@ -147,7 +148,8 @@ def txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator,sup
         height = height,
         output_type="latent",
         generator = generator).images
-
+      pipe=pipe.to("cpu")
+      upscaler=upscaler.to("cuda")
       result = upscaler(
         prompt=prompt,
         negative_prompt = neg_prompt,
@@ -156,6 +158,9 @@ def txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator,sup
         guidance_scale=guidance,
         generator=generator,
       )
+      result=result.to("cpu")
+      pipe=pipe.to("cuda")
+      upscaler=upscaler.to("cpu")
     else:  
       result = pipe(
         prompt,
@@ -183,7 +188,8 @@ def img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height
         #height = height,
         output_type="latent",
         generator = generator).images
-
+      pipe=pipe.to("cpu")
+      upscaler=upscaler.to("cuda")
       result = upscaler(
         prompt=prompt,
         negative_prompt = neg_prompt,
@@ -192,6 +198,9 @@ def img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height
         guidance_scale=guidance,
         generator=generator,
       )
+      result=result.to("cpu")
+      pipe=pipe.to("cuda")
+      upscaler=upscaler.to("cpu")
     else:  
       result = pipe_i2i(
         prompt,