ByteDance

/

SDXL-Lightning

like 1.92k

Follow

ByteDance 708

Will provide safetensors. Coming soon

Hi @PeterL1n I converted the existing weights into the Diffusers format that doesn't require hf_hub_download and torch.load. Added a section in the model card to demonstrate.

I gotta say I am not a fan of this structure... And now it is kinda late to move .pth files into folders with .safetensors because some HF spaces already use them.

Here is what I propose.
We could upload safetensors to root directory following the same naming.

Change loading logic to this.

import torch
from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

base = "stabilityai/stable-diffusion-xl-base-1.0"
repo = "ByteDance/SDXL-Lightning"
ckpt = "sdxl_lightning_4step_lora.safetensors"

unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")

The benefits are

1. We don't need to store model config.json
2. The naming is just way better. I prefer to have more general naming instead of diffusers specific naming.

Anything I do not foresee?

I already updated the doc on how to load unet pth using this new method to avoid downloading original SDXL.

Updated to safetensors. Closing this PR.