Update app.py

app.py

@@ -18,7 +18,7 @@ from diffusers import AutoencoderKL, StableDiffusionXLPipeline, EulerAncestralDi
 from typing import Tuple
 #from transformers import AutoTokenizer, AutoModelForCausalLM
 import paramiko
-
+import gc
 #os.system("chmod +x ./cusparselt.sh")
 #os.system("./cusparselt.sh")
 #os.system("chmod +x ./cudnn.sh")
@@ -204,8 +204,8 @@ def generate(
         options["use_resolution_binning"] = True
 
     images = []
-    with torch.no_grad():
-      for i in range(0, num_images, BATCH_SIZE):
+    #with torch.no_grad():
+    for i in range(0, num_images, BATCH_SIZE):
         batch_options = options.copy()
         batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
         if "negative_prompt" in batch_options:
@@ -215,6 +215,8 @@ def generate(
     images[0].save(sd_image_path,optimize=False,compress_level=0)
     upload_to_ftp(sd_image_path)    
     image_paths = [save_image(img) for img in images]
+    torch.cuda.empty_cache()
+    gc.collect()
     return image_paths, seed
 
 def generate_cpu(
@@ -238,7 +240,7 @@ def generate_cpu(
     pipe.to("cpu")
 
     seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator(device='cpu').manual_seed(seed)
+    generator = torch.Generator(device='cuda').manual_seed(seed)
 
     prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)