|
--- |
|
license: openrail++ |
|
tags: |
|
- text-to-image |
|
- stable-diffusion |
|
library_name: diffusers |
|
--- |
|
|
|
# SDXL-Lightning |
|
|
|
![Intro Image](images/intro.jpg) |
|
|
|
SDXL-Lightning is a lightning fast text-to-image generation model. It can generate high-quality 1024px images under a few steps. For more information, please refer to our paper: [SDXL-Lightning: Progressive Adversarial Diffusion Distillation](https://huggingface.co./ByteDance/SDXL-Lightning/resolve/main/sdxl_lightning_report.pdf). The models are released for research purposes only. |
|
|
|
Our models are distilled from [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co./stabilityai/stable-diffusion-xl-base-1.0). This repository contains checkpoints for 1-step, 2-step, 4-step, and 8-step distilled models. The generation quality of our 2-step, 4-step, and 8-step model is amazing. Our 1-step model is more experimental. |
|
|
|
We provide both full UNet and LoRA checkpoints. The full UNet models have the best quality while the LoRA models can be applied to other base models. |
|
|
|
|
|
## Diffusers Usage |
|
|
|
Please always use the correct checkpoint for the corresponding inference steps. |
|
|
|
### 2-Step, 4-Step, 8-Step UNet |
|
|
|
```python |
|
import torch |
|
from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler |
|
from huggingface_hub import hf_hub_download |
|
|
|
base = "stabilityai/stable-diffusion-xl-base-1.0" |
|
repo = "ByteDance/SDXL-Lightning" |
|
ckpt = "sdxl_lightning_4step_unet.pth" # Use the correct ckpt for your step setting! |
|
|
|
# Load model. |
|
pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda") |
|
pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, ckpt), map_location="cuda")) |
|
|
|
# Ensure sampler uses "trailing" timesteps. |
|
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") |
|
|
|
# Ensure using the same inference steps as the loaded model and CFG set to 0. |
|
pipe("A girl smiling", num_inference_steps=4, guidance_scale=0).images[0].save("output.png") |
|
``` |
|
|
|
### 2-Step, 4-Step, 8-Step LoRA |
|
|
|
```python |
|
import torch |
|
from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler |
|
from huggingface_hub import hf_hub_download |
|
|
|
base = "stabilityai/stable-diffusion-xl-base-1.0" |
|
repo = "ByteDance/SDXL-Lightning" |
|
ckpt = "sdxl_lightning_4step_lora.pth" # Use the correct ckpt for your step setting! |
|
|
|
# Load model. |
|
pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda") |
|
pipe.load_lora_weights(hf_hub_download(repo, ckpt)) |
|
pipe.fuse_lora() |
|
|
|
# Ensure sampler uses "trailing" timesteps. |
|
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") |
|
|
|
# Ensure using the same inference steps as the loaded model and CFG set to 0. |
|
pipe("A girl smiling", num_inference_steps=4, guidance_scale=0).images[0].save("output.png") |
|
``` |
|
|
|
### 1-Step UNet |
|
|
|
The 1-step model uses "sample" prediction instead of "epsilon" prediction! The scheduler needs to be configured correctly. |
|
|
|
```python |
|
import torch |
|
from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler |
|
from huggingface_hub import hf_hub_download |
|
|
|
base = "stabilityai/stable-diffusion-xl-base-1.0" |
|
repo = "ByteDance/SDXL-Lightning" |
|
ckpt = "sdxl_lightning_1step_unet_x0.pth" # Use the correct ckpt for your step setting! |
|
|
|
# Load model. |
|
pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda") |
|
pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, ckpt), map_location="cuda")) |
|
|
|
# Ensure sampler uses "trailing" timesteps and "sample" prediction type. |
|
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample") |
|
|
|
# Ensure using the same inference steps as the loaded model and CFG set to 0. |
|
pipe("A girl smiling", num_inference_steps=1, guidance_scale=0).images[0].save("output.png") |
|
``` |
|
|
|
|
|
## ComfyUI Usage |
|
|
|
Please always use the correct checkpoint for the corresponding inference steps. |
|
Please use Euler sampler with sgm_uniform scheduler. |
|
|
|
### 2-Step, 4-Step, 8-Step UNet |
|
|
|
1. Download the UNet checkpoint to `/ComfyUI/models/unet`. |
|
2. Download our [ComfyUI UNet workflow](comfyui/sdxl_lightning_unet.json). |
|
|
|
![SDXL-Lightning ComfyUI UNet Workflow](images/comfyui_unet.png) |
|
|
|
### 2-Step, 4-Step, 8-Step LoRA |
|
|
|
1. Download the LoRA checkpoint to `/ComfyUI/models/loras` |
|
2. Download our [ComfyUI LoRA workflow](comfyui/sdxl_lightning_lora.json). |
|
|
|
![SDXL-Lightning ComfyUI UNet Workflow](images/comfyui_lora.png) |
|
|
|
### 1-Step UNet |
|
|
|
ComfyUI does not support changing model formulation to x0-prediction, so it is not usable in ComfyUI yet. Hopefully ComfyUI gets updated soon. |