wpeebles commited on
Commit
4fbe107
1 Parent(s): 2dc4acc

Clear CUDA cache

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -22,6 +22,7 @@ def load_model(image_size=256):
22
 
23
  torch.set_grad_enabled(False)
24
  device = "cuda" if torch.cuda.is_available() else "cpu"
 
25
  model = load_model(image_size=256)
26
  vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device)
27
  current_image_size = 256
@@ -34,8 +35,8 @@ def generate(image_size, vae_model, class_label, cfg_scale, num_sampling_steps,
34
  if image_size != current_image_size:
35
  global model
36
  del model
37
- # if device == "cuda":
38
- # torch.cuda.empty_cache()
39
  model = load_model(image_size=image_size)
40
  current_image_size = image_size
41
 
 
22
 
23
  torch.set_grad_enabled(False)
24
  device = "cuda" if torch.cuda.is_available() else "cpu"
25
+ find_model(f"DiT-XL-2-512x512.pt")
26
  model = load_model(image_size=256)
27
  vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device)
28
  current_image_size = 256
 
35
  if image_size != current_image_size:
36
  global model
37
  del model
38
+ if device == "cuda":
39
+ torch.cuda.empty_cache()
40
  model = load_model(image_size=image_size)
41
  current_image_size = image_size
42