README.md

metadata

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

⚡ FlashDiffusion: FlashPixart ⚡

Flash Diffusion is a diffusion distillation method proposed in ADD ARXIV by Clément Chadebec, Onur Tasar and Benjamin Aubin. This model is a 66.5M LoRA distilled version of Pixart-α model that is able to generate 1024x1024 images in 4 steps. See our live demo.

How to use?

The model can be used using the StableDiffusionPipeline from diffusers library directly. It can allow reducing the number of required sampling steps to 2-4 steps.

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]

Training Details

The model was trained for 40k iterations on 4 H100 GPUs. Please refer to the paper for further parameters details.

License

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