diffusers/tools

1#!/usr/bin/env python32import torch3import gc4from diffusers import DiffusionPipeline5 6shape = (30_000, 30_000)7 8input = torch.randn(shape, device="cuda")9 10 11def clear_memory(model):12    model.to('cpu')13    gc.collect()14    torch.cuda.empty_cache()15    torch.cuda.ipc_collect()16    torch.clear_autocast_cache()17 18for _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"]:19    pipe = DiffusionPipeline.from_pretrained(_ids, use_safetensors=True).to("cuda")20    pipe("hey", num_inference_steps=1)21    print("finished...")22 23    clear_memory(pipe)24