tools / new_loading.py
patrickvonplaten's picture
all
6c27fdd
raw
history blame
725 Bytes
#!/usr/bin/env python3
import time
import torch
import safetensors.torch
from diffusers import StableDiffusionXLPipeline
pipe = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
use_safetensors=True,
local_files_only=True,
)
pipe = pipe.to("cuda")
# !wget https://civitai.com/api/download/models/135931 -O loras/pixel-art-xl.safetensors
lora_weights = safetensors.torch.load_file(
"pixel-art-xl.safetensors", device="cpu"
)
for _ in range(5):
t0 = time.perf_counter()
pipe.load_lora_weights(lora_weights.copy())
pipe.unload_lora_weights()
print("Load + unload cycle took: ", time.perf_counter() - t0)