Running IF with 🧨 diffusers on a Free Tier Google Colab
shonenkov, Gugutse, ZeroShot-AI, williamberman, patrickvonplaten, multimodalart
<a target="_blank" href="https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/deepfloyd_if_free_tier_google_colab.ipynb">
    <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
</a>


**TL;DR**: We show how to run one of the most powerful open-source text
to image models **IF** on a free-tier Google Colab with 🧨 diffusers.

You can also explore the capabilities of the model directly in the [Hugging Face Space](https://huggingface.co/spaces/DeepFloyd/IF).

<p align="center">
    <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/if/nabla.jpg" alt="if-collage"><br>
    <em>Image compressed from official <a href="https://github.com/deep-floyd/IF/blob/release/pics/nabla.jpg">IF GitHub repo</a>.</em>
</p>

## Introduction

IF is a pixel-based text-to-image generation model and was [released in
late April 2023 by DeepFloyd](https://github.com/deep-floyd/IF). The
model architecture is strongly inspired by [Google's closed-sourced
Imagen](https://imagen.research.google/).

IF has two distinct advantages compared to existing text-to-image models
like Stable Diffusion:

- The model operates directly in "pixel space" (*i.e.,* on
uncompressed images) instead of running the denoising process in the
latent space such as [Stable Diffusion](http://hf.co/blog/stable_diffusion).
- The model is trained on outputs of
[T5-XXL](https://huggingface.co/google/t5-v1_1-xxl), a more powerful
text encoder than [CLIP](https://openai.com/research/clip), used by
Stable Diffusion as the text encoder.

As a result, IF is better at generating images with high-frequency
details (*e.g.,* human faces and hands) and is the first open-source
image generation model that can reliably generate images with text.

The downside of operating in pixel space and using a more powerful text
encoder is that IF has a significantly higher amount of parameters. T5,
IF\'s text-to-image UNet, and IF\'s upscaler UNet have 4.5B, 4.3B, and
1.2B parameters respectively. Compared to [Stable Diffusion
2.1](https://huggingface.co/stabilityai/stable-diffusion-2-1)\'s text
encoder and UNet having just 400M and 900M parameters, respectively.

Nevertheless, it is possible to run IF on consumer hardware if one
optimizes the model for low-memory usage. We will show you can do this
with 🧨 diffusers in this blog post.

In 1.), we explain how to use IF for text-to-image generation, and in 2.)
and 3.), we go over IF's image variation and image inpainting
capabilities.

💡 **Note**: We are trading gains in memory by gains in
speed here to make it possible to run IF in a free-tier Google Colab. If
you have access to high-end GPUs such as an A100, we recommend leaving
all model components on GPU for maximum speed, as done in the
[official IF demo](https://huggingface.co/spaces/DeepFloyd/IF).

💡 **Note**: Some of the larger images have been compressed to load faster 
in the blog format. When using the official model, they should be even
better quality!

Let\'s dive in 🚀!

<p align="center">
    <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/if/meme.png"><br>
    <em>IF's text generation capabilities</em>
</p>

## Table of contents

* [Accepting the license](#accepting-the-license)
* [Optimizing IF to run on memory constrained hardware](#optimizing-if-to-run-on-memory-constrained-hardware)
* [Available resources](#available-resources)
* [Install dependencies](#install-dependencies)
* [Text-to-image generation](#1-text-to-image-generation)
* [Image variation](#2-image-variation)
* [Inpainting](#3-inpainting)

## Accepting the license

Before you can use IF, you need to accept its usage conditions. To do so:

- 1. Make sure to have a [Hugging Face account](https://huggingface.co/join) and be logged in
- 2. Accept the license on the model card of [DeepFloyd/IF-I-XL-v1.0](https://huggingface.co/DeepFloyd/IF-I-XL-v1.0). Accepting the license on the stage I model card will auto accept for the other IF models.
- 3. Make sure to login locally. Install `huggingface_hub`

```sh
pip install huggingface_hub --upgrade
```

run the login function in a Python shell

```py
from huggingface_hub import login

login()
```

and enter your [Hugging Face Hub access token](https://huggingface.co/docs/hub/security-tokens#what-are-user-access-tokens).

## Optimizing IF to run on memory constrained hardware

State-of-the-art ML should not just be in the hands of an elite few.
Democratizing ML means making models available to run on more than just
the latest and greatest hardware.

The deep learning community has created world class tools to run
resource intensive models on consumer hardware:

- [🤗 accelerate](https://github.com/huggingface/accelerate) provides
utilities for working with [large models](https://huggingface.co/docs/accelerate/usage_guides/big_modeling).
- [bitsandbytes](https://github.com/TimDettmers/bitsandbytes) makes [8-bit quantization](https://github.com/TimDettmers/bitsandbytes#features) available to all PyTorch models.
- [🤗 safetensors](https://github.com/huggingface/safetensors) not only ensures that save code is executed but also significantly speeds up the loading time of large models.

Diffusers seamlessly integrates the above libraries to allow for a
simple API when optimizing large models.

The free-tier Google Colab is both CPU RAM constrained (13 GB RAM) as
well as GPU VRAM constrained (15 GB RAM for T4), which makes running the
whole >10B IF model challenging!

Let\'s map out the size of IF\'s model components in full float32
precision:

- [T5-XXL Text Encoder](https://huggingface.co/DeepFloyd/IF-I-XL-v1.0/tree/main/text_encoder): 20GB
- [Stage 1 UNet](https://huggingface.co/DeepFloyd/IF-I-XL-v1.0/tree/main/unet): 17.2 GB
- [Stage 2 Super Resolution UNet](https://huggingface.co/DeepFloyd/IF-II-L-v1.0/blob/main/pytorch_model.bin): 2.5 GB
- [Stage 3 Super Resolution Model](https://huggingface.co/stabilityai/stable-diffusion-x4-upscaler): 3.4 GB

There is no way we can run the model in float32 as the T5 and Stage 1
UNet weights are each larger than the available CPU RAM.

In float16, the component sizes are 11GB, 8.6GB, and 1.25GB for T5,
Stage1 and Stage2 UNets, respectively, which is doable for the GPU, but
we're still running into CPU memory overflow errors when loading the T5
(some CPU is occupied by other processes).

Therefore, we lower the precision of T5 even more by using
`bitsandbytes` 8bit quantization, which allows saving the T5 checkpoint
with as little as [8
GB](https://huggingface.co/DeepFloyd/IF-I-XL-v1.0/blob/main/text_encoder/model.8bit.safetensors).

Now that each component fits individually into both CPU and GPU memory,
we need to make sure that components have all the CPU and GPU memory for
themselves when needed.

Diffusers supports modularly loading individual components i.e. we can
load the text encoder without loading the UNet. This modular loading
will ensure that we only load the component we need at a given step in
the pipeline to avoid exhausting the available CPU RAM and GPU VRAM.

Let\'s give it a try 🚀

![t2i_64](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/if/t2i_64.png)

## Available resources

The free-tier Google Colab comes with around 13 GB CPU RAM:

``` python
!grep MemTotal /proc/meminfo
```

```bash
MemTotal:       13297192 kB
```

And an NVIDIA T4 with 15 GB VRAM:

``` python
!nvidia-smi
```

```bash
Sun Apr 23 23:14:19 2023       
+