Clear CUDA cache

app.py

@@ -22,6 +22,7 @@ def load_model(image_size=256):
 
 torch.set_grad_enabled(False)
 device = "cuda" if torch.cuda.is_available() else "cpu"
+find_model(f"DiT-XL-2-512x512.pt")
 model = load_model(image_size=256)
 vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device)
 current_image_size = 256
@@ -34,8 +35,8 @@ def generate(image_size, vae_model, class_label, cfg_scale, num_sampling_steps,
     if image_size != current_image_size:
         global model
         del model
-        # if device == "cuda":
-        #     torch.cuda.empty_cache()
+        if device == "cuda":
+            torch.cuda.empty_cache()
         model = load_model(image_size=image_size)
     current_image_size = image_size