Update app.py

app.py

@@ -218,6 +218,8 @@ def generate_30(
     pipe.text_encoder_2=text_encoder_2
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
     if latent_file is not None:  # Check if a latent file is provided
         sd_image_a = Image.open(latent_file.name).convert('RGB')
         sd_image_a.resize((height,width), Image.LANCZOS)
@@ -269,6 +271,8 @@ def generate_30(
         upload_to_ftp(filename)
         uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
         torch.set_float32_matmul_precision("medium")
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
         with torch.no_grad():
             upscale = upscaler(sd_image, tiling=True, tile_width=256, tile_height=256)
         downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
@@ -309,25 +313,28 @@ def generate_60(
     pipe.text_encoder_2=text_encoder_2
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
     if latent_file is not None:  # Check if a latent file is provided
-        sd_image_a = Image.open(latent_file.name)
+        sd_image_a = Image.open(latent_file.name).convert('RGB')
+        sd_image_a.resize((height,width), Image.LANCZOS)
         if latent_file_2 is not None:  # Check if a latent file is provided
-            sd_image_b = Image.open(latent_file_2.name)
+            sd_image_b = Image.open(latent_file_2.name).convert('RGB')
             sd_image_b.resize((height,width), Image.LANCZOS)
         else:
             sd_image_b = None
         if latent_file_3 is not None:  # Check if a latent file is provided
-            sd_image_c = Image.open(latent_file_3.name)
+            sd_image_c = Image.open(latent_file_3.name).convert('RGB')
             sd_image_c.resize((height,width), Image.LANCZOS)
         else:
             sd_image_c = None
         if latent_file_4 is not None:  # Check if a latent file is provided
-            sd_image_d = Image.open(latent_file_4.name)
+            sd_image_d = Image.open(latent_file_4.name).convert('RGB')
             sd_image_d.resize((height,width), Image.LANCZOS)
         else:
             sd_image_d = None
         if latent_file_5 is not None:  # Check if a latent file is provided
-            sd_image_e = Image.open(latent_file_5.name)
+            sd_image_e = Image.open(latent_file_5.name).convert('RGB')
             sd_image_e.resize((height,width), Image.LANCZOS)
         else:
             sd_image_e = None
@@ -359,6 +366,8 @@ def generate_60(
         upload_to_ftp(filename)
         uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
         torch.set_float32_matmul_precision("medium")
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
         with torch.no_grad():
             upscale = upscaler(sd_image, tiling=True, tile_width=256, tile_height=256)
         downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
@@ -399,25 +408,28 @@ def generate_90(
     pipe.text_encoder_2=text_encoder_2
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
     if latent_file is not None:  # Check if a latent file is provided
-        sd_image_a = Image.open(latent_file.name)
+        sd_image_a = Image.open(latent_file.name).convert('RGB')
+        sd_image_a.resize((height,width), Image.LANCZOS)
         if latent_file_2 is not None:  # Check if a latent file is provided
-            sd_image_b = Image.open(latent_file_2.name)
+            sd_image_b = Image.open(latent_file_2.name).convert('RGB')
             sd_image_b.resize((height,width), Image.LANCZOS)
         else:
             sd_image_b = None
         if latent_file_3 is not None:  # Check if a latent file is provided
-            sd_image_c = Image.open(latent_file_3.name)
+            sd_image_c = Image.open(latent_file_3.name).convert('RGB')
             sd_image_c.resize((height,width), Image.LANCZOS)
         else:
             sd_image_c = None
         if latent_file_4 is not None:  # Check if a latent file is provided
-            sd_image_d = Image.open(latent_file_4.name)
+            sd_image_d = Image.open(latent_file_4.name).convert('RGB')
             sd_image_d.resize((height,width), Image.LANCZOS)
         else:
             sd_image_d = None
         if latent_file_5 is not None:  # Check if a latent file is provided
-            sd_image_e = Image.open(latent_file_5.name)
+            sd_image_e = Image.open(latent_file_5.name).convert('RGB')
             sd_image_e.resize((height,width), Image.LANCZOS)
         else:
             sd_image_e = None
@@ -425,7 +437,6 @@ def generate_90(
         filename= f'rv_IP_{timestamp}.png'
         print("-- using image file --")
         print('-- generating image --')
-        #with torch.no_grad():
         sd_image = ip_model.generate(
                 pil_image_1=sd_image_a,
                 pil_image_2=sd_image_b,
@@ -450,6 +461,8 @@ def generate_90(
         upload_to_ftp(filename)
         uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
         torch.set_float32_matmul_precision("medium")
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
         with torch.no_grad():
             upscale = upscaler(sd_image, tiling=True, tile_width=256, tile_height=256)
         downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)