File size: 665 Bytes
e402ae6 27dfa17 e402ae6 27dfa17 e402ae6 27dfa17 e402ae6 27dfa17 e402ae6 27dfa17 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
#!/usr/bin/env python3
import torch
import gc
from diffusers import DiffusionPipeline
shape = (30_000, 30_000)
input = torch.randn(shape, device="cuda")
def clear_memory(model):
model.to('cpu')
gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
torch.clear_autocast_cache()
for _ids in ["runwayml/stable-diffusion-v1-5", "CompVis/stable-diffusion-v1-4", "runwayml/stable-diffusion-v1-5", "CompVis/stable-diffusion-v1-4", "runwayml/stable-diffusion-v1-5"]:
pipe = DiffusionPipeline.from_pretrained(_ids, use_safetensors=True).to("cuda")
pipe("hey", num_inference_steps=1)
print("finished...")
clear_memory(pipe)
|