Update README.md

README.md

@@ -1,3 +1,70 @@
----
-license: cc-by-nc-4.0
----
+---
+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 **26.4M** LoRA distilled version of Pixart-α model that is able to generate 1024x1024 images in **4 steps**.
+
+
+<p align="center">
+   <img style="width:700px;" src="images/hf_grid.png">
+</p>
+
+# 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**.
+
+```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="images/raccoon.png">
+</p>
+
+# 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.