Update app.py

app.py

@@ -10,6 +10,7 @@ import urllib
 import time
 import os
 import datetime
+import gc
 
 from models.transformer_sd3 import SD3Transformer2DModel
 #from diffusers import StableDiffusion3Pipeline
@@ -125,6 +126,9 @@ def infer(
         nb_token=64, 
     )
     upscaler_2.to(torch.device('cpu'))
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
     torch.set_float32_matmul_precision("highest")
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
@@ -183,6 +187,9 @@ def infer(
         sd_image.save(rv_path,optimize=False,compress_level=0)
         upload_to_ftp(rv_path)
         upscaler_2.to(torch.device('cuda'))
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
         with torch.no_grad():
             upscale2 = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
         print('-- got upscaled image --')