Text-to-Image
Diffusers
Safetensors
lora
File size: 2,791 Bytes
2139044
 
 
 
 
 
 
 
 
dd0b758
2139044
 
bbe7d5c
 
eee6df9
02cd9ac
2139044
 
 
1166844
2139044
 
 
 
127083a
2139044
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
201d281
2139044
 
 
bbe7d5c
2139044
aa9a9bb
 
 
2139044
eee6df9
 
 
 
 
 
 
 
 
 
 
 
 
 
2139044
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
---
license: cc-by-nc-4.0
library_name: diffusers
base_model: PixArt-alpha/PixArt-XL-2-1024-MS
tags:
- lora
- text-to-image
inference: False
---
# ⚡ Flash Diffusion: FlashPixart ⚡


Flash Diffusion is a diffusion distillation method proposed in [Flash Diffusion: Accelerating Any Conditional
Diffusion Model for Few Steps Image Generation](http://arxiv.org/abs/2406.02347) *by Clément Chadebec, Onur Tasar, Eyal Benaroche, and Benjamin Aubin.*
This model is a **66.5M** LoRA distilled version of [Pixart-α](https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS) model that is able to generate 1024x1024 images in **4 steps**.
See our [live demo](https://huggingface.co/spaces/jasperai/FlashPixart) and [official Github repo](https://github.com/gojasper/flash-diffusion).


<p align="center">
   <img style="width:700px;" src="assets/flash_pixart.jpg">
</p>

# How to use?

The model can be used using the `PixArtAlphaPipeline` from `diffusers` library directly. It can allow reducing the number of required sampling steps to **4 steps**.

```python
import torch
from diffusers import PixArtAlphaPipeline, Transformer2DModel, LCMScheduler
from peft import PeftModel

# Load LoRA
transformer = Transformer2DModel.from_pretrained(
  "PixArt-alpha/PixArt-XL-2-1024-MS",
  subfolder="transformer",
  torch_dtype=torch.float16
)
transformer = PeftModel.from_pretrained(
  transformer,
  "jasperai/flash-pixart"
)

# Pipeline
pipe = PixArtAlphaPipeline.from_pretrained(
  "PixArt-alpha/PixArt-XL-2-1024-MS",
  transformer=transformer,
  torch_dtype=torch.float16
)

# Scheduler
pipe.scheduler = LCMScheduler.from_pretrained(
  "PixArt-alpha/PixArt-XL-2-1024-MS",
  subfolder="scheduler",
  timestep_spacing="trailing",
)

pipe.to("cuda")

prompt = "A raccoon reading a book in a lush forest."

image = pipe(prompt, num_inference_steps=4, guidance_scale=0).images[0]
```
<p align="center">
   <img style="width:400px;" src="assets/raccoon.png">
</p>

# Training Details
The model was trained for 40k iterations on 4 H100 GPUs (representing approximately 188 hours of training). Please refer to the [paper](http://arxiv.org/abs/2406.02347) for further parameters details. 

**Metrics on COCO 2014 validation (Table 4)**
  - FID-10k: 29.30 (4 NFE)
  - CLIP Score: 0.303 (4 NFE)

## Citation
If you find this work useful or use it in your research, please consider citing us

```bibtex
@misc{chadebec2024flash,
      title={Flash Diffusion: Accelerating Any Conditional Diffusion Model for Few Steps Image Generation}, 
      author={Clement Chadebec and Onur Tasar and Eyal Benaroche and Benjamin Aubin},
      year={2024},
      eprint={2406.02347},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}
```

## License
This model is released under the the Creative Commons BY-NC license.