device_map support to pipelines (#6857)
* get device <-> component mapping when using multiple gpus. * condition the device_map bits. * relax condition * device_map progress. * device_map enhancement * some cleaning up and debugging * Apply suggestions from code review Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com> * incorporate suggestions from PR. * remove multi-gpu condition for now. * guard check the component -> device mapping * fix: device_memory variable * dispatching transformers model to have force_hooks=True * better guarding for transformers device_map * introduce support balanced_low_memory and balanced_ultra_low_memory. * remove device_map patch. * fix: intermediate variable scoping. * fix: condition in cpu offload. * fix: flax class restrictions. * remove modifications from cpu_offload and model_offload * incorporate changes. * add a simple forward pass test * add: torch_device in get_inputs() * add: tests * remove print * safe-guard to(), model offloading and cpu offloading when balanced is used as a device_map. * style * remove . * safeguard device_map with more checks and remove invalid device_mapping strategues. * make a class attribute and adjust tests accordingly. * fix device_map check * fix test * adjust comment * fix: device_map attribute * fix: dispatching. * max_memory test for pipeline * version guard the tests * fix guard. * address review feedback. * reset_device_map method. * add: test for reset_hf_device_map * fix a couple things. * add reset_device_map() in the error message. * add tests for checking reset_device_map doesn't have unintended consequences. * fix reset_device_map and offloading tests. * create _get_final_device_map utility. * hf_device_map -> _hf_device_map * add documentation * add notes suggested by Marc. * styling. * Apply suggestions from code review Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * move updates within gpu condition. * other docs related things * note on ignore a device not specified in . * provide a suggestion if device mapping errors out. * fix: typo. * _hf_device_map -> hf_device_map * Empty-Commit * add: example hf_device_map. --------- Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: Pedro Cuenca <pedro@huggingface.co>
π€ Diffusers is the go-to library for state-of-the-art pretrained diffusion models for generating images, audio, and even 3D structures of molecules. Whether you're looking for a simple inference solution or training your own diffusion models, π€ Diffusers is a modular toolbox that supports both. Our library is designed with a focus on usability over performance, simple over easy, and customizability over abstractions.
π€ Diffusers offers three core components:
- State-of-the-art diffusion pipelines that can be run in inference with just a few lines of code.
- Interchangeable noise schedulers for different diffusion speeds and output quality.
- Pretrained models that can be used as building blocks, and combined with schedulers, for creating your own end-to-end diffusion systems.
Installation
We recommend installing π€ Diffusers in a virtual environment from PyPI or Conda. For more details about installing PyTorch and Flax, please refer to their official documentation.
PyTorch
With pip (official package):
pip install --upgrade diffusers[torch]
With conda (maintained by the community):
conda install -c conda-forge diffusers
Flax
With pip (official package):
pip install --upgrade diffusers[flax]
Apple Silicon (M1/M2) support
Please refer to the How to use Stable Diffusion in Apple Silicon guide.
Quickstart
Generating outputs is super easy with π€ Diffusers. To generate an image from text, use the from_pretrained method to load any pretrained diffusion model (browse the Hub for 22000+ checkpoints):
from diffusers import DiffusionPipeline
import torch
pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipeline.to("cuda")
pipeline("An image of a squirrel in Picasso style").images[0]
You can also dig into the models and schedulers toolbox to build your own diffusion system:
from diffusers import DDPMScheduler, UNet2DModel
from PIL import Image
import torch
scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
model = UNet2DModel.from_pretrained("google/ddpm-cat-256").to("cuda")
scheduler.set_timesteps(50)
sample_size = model.config.sample_size
noise = torch.randn((1, 3, sample_size, sample_size), device="cuda")
input = noise
for t in scheduler.timesteps:
with torch.no_grad():
noisy_residual = model(input, t).sample
prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
input = prev_noisy_sample
image = (input / 2 + 0.5).clamp(0, 1)
image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
image = Image.fromarray((image * 255).round().astype("uint8"))
image
Check out the Quickstart to launch your diffusion journey today!
How to navigate the documentation
| Documentation | What can I learn? |
|---|---|
| Tutorial | A basic crash course for learning how to use the library's most important features like using models and schedulers to build your own diffusion system, and training your own diffusion model. |
| Loading | Guides for how to load and configure all the components (pipelines, models, and schedulers) of the library, as well as how to use different schedulers. |
| Pipelines for inference | Guides for how to use pipelines for different inference tasks, batched generation, controlling generated outputs and randomness, and how to contribute a pipeline to the library. |
| Optimization | Guides for how to optimize your diffusion model to run faster and consume less memory. |
| Training | Guides for how to train a diffusion model for different tasks with different training techniques. |
Contribution
We β€οΈ contributions from the open-source community! If you want to contribute to this library, please check out our Contribution guide. You can look out for issues you'd like to tackle to contribute to the library.
- See Good first issues for general opportunities to contribute
- See New model/pipeline to contribute exciting new diffusion models / diffusion pipelines
- See New scheduler
Also, say π in our public Discord channel . We discuss the hottest trends about diffusion models, help each other with contributions, personal projects or just hang out β.
Popular Tasks & Pipelines
| Task | Pipeline | π€ Hub |
|---|---|---|
| Unconditional Image Generation | DDPM | google/ddpm-ema-church-256 |
| Text-to-Image | Stable Diffusion Text-to-Image | runwayml/stable-diffusion-v1-5 |
| Text-to-Image | unCLIP | kakaobrain/karlo-v1-alpha |
| Text-to-Image | DeepFloyd IF | DeepFloyd/IF-I-XL-v1.0 |
| Text-to-Image | Kandinsky | kandinsky-community/kandinsky-2-2-decoder |
| Text-guided Image-to-Image | ControlNet | lllyasviel/sd-controlnet-canny |
| Text-guided Image-to-Image | InstructPix2Pix | timbrooks/instruct-pix2pix |
| Text-guided Image-to-Image | Stable Diffusion Image-to-Image | runwayml/stable-diffusion-v1-5 |
| Text-guided Image Inpainting | Stable Diffusion Inpainting | runwayml/stable-diffusion-inpainting |
| Image Variation | Stable Diffusion Image Variation | lambdalabs/sd-image-variations-diffusers |
| Super Resolution | Stable Diffusion Upscale | stabilityai/stable-diffusion-x4-upscaler |
| Super Resolution | Stable Diffusion Latent Upscale | stabilityai/sd-x2-latent-upscaler |
Popular libraries using 𧨠Diffusers
- https://github.com/microsoft/TaskMatrix
- https://github.com/invoke-ai/InvokeAI
- https://github.com/apple/ml-stable-diffusion
- https://github.com/Sanster/lama-cleaner
- https://github.com/IDEA-Research/Grounded-Segment-Anything
- https://github.com/ashawkey/stable-dreamfusion
- https://github.com/deep-floyd/IF
- https://github.com/bentoml/BentoML
- https://github.com/bmaltais/kohya_ss
- +9000 other amazing GitHub repositories πͺ
Thank you for using us β€οΈ.
Credits
This library concretizes previous work by many different authors and would not have been possible without their great research and implementations. We'd like to thank, in particular, the following implementations which have helped us in our development and without which the API could not have been as polished today:
- @CompVis' latent diffusion models library, available here
- @hojonathanho original DDPM implementation, available here as well as the extremely useful translation into PyTorch by @pesser, available here
- @ermongroup's DDIM implementation, available here
- @yang-song's Score-VE and Score-VP implementations, available here
We also want to thank @heejkoo for the very helpful overview of papers, code and resources on diffusion models, available here as well as @crowsonkb and @rromb for useful discussions and insights.
Citation
@misc{von-platen-etal-2022-diffusers,
author = {Patrick von Platen and Suraj Patil and Anton Lozhkov and Pedro Cuenca and Nathan Lambert and Kashif Rasul and Mishig Davaadorj and Dhruv Nair and Sayak Paul and William Berman and Yiyi Xu and Steven Liu and Thomas Wolf},
title = {Diffusers: State-of-the-art diffusion models},
year = {2022},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/huggingface/diffusers}}
}
