Flash Pixart

⚡ Flash Diffusion: FlashPixart ⚡

Flash Diffusion is a diffusion distillation method presented in the paper Flash Diffusion: Accelerating Any Conditional Diffusion Model for Few Steps Image Generation by Clément Chadebec, Onur Tasar, Eyal Benaroche, and Benjamin Aubin from Jasper Research. This model is a 66.5M LoRA distilled version of the Pixart-α model, capable of generating 1024x1024 images in just 4 steps. Check out our live demo and the official Github repo.

🚀 Quick Start

The model can be utilized directly via the PixArtAlphaPipeline from the diffusers library, which enables reducing the necessary sampling steps to 4 steps.

💻 Usage Examples

🔍 Basic Usage

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]

🔧 Technical Details

The model was trained for 40k iterations on 4 H100 GPUs, which is approximately 188 hours of training. For more detailed parameters, please refer to the paper.

📊 Metrics on COCO 2014 validation (Table 4)

📚 Documentation

📖 Citation

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

@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 Creative Commons BY - NC license.