| |
| 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) |
|
|
