Running IF with 🧨 diffusers on a Free Tier Google Colab
Alex Shonenkov, Daria Bakshandaeva, Misha Konstantinov, Will Berman, Patrick von Platen, Apolinário from multimodal AI art
TL;DR: We show how to run one of the most powerful open-source textto 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.Image compressed from official IF GitHub repo.IntroductionIF is a pixel-based text-to-image generation model and was released inlate April 2023 by DeepFloyd. Themodel architecture is strongly inspired by Google's closed-sourcedImagen.IF has two distinct advantages compared to existing text-to-image modelslike Stable Diffusion:The model operates directly in "pixel space" (i.e., onuncompressed images) instead of running the denoising process in thelatent space such as Stable Diffusion.The model is trained on outputs ofT5-XXL, a more powerfultext encoder than CLIP, used byStable Diffusion as the text encoder.As a result, IF is better at generating images with high-frequencydetails (e.g., human faces and hands) and is the first open-sourceimage generation model that can reliably generate images with text.The downside of operating in pixel space and using a more powerful textencoder 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, and1.2B parameters respectively. Compared to Stable Diffusion2.1's textencoder and UNet having just 400M and 900M parameters, respectively.Nevertheless, it is possible to run IF on consumer hardware if oneoptimizes the model for low-memory usage. We will show you can do thiswith 🧨 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 inpaintingcapabilities.💡 Note: We are trading gains in memory by gains inspeed here to make it possible to run IF in a free-tier Google Colab. Ifyou have access to high-end GPUs such as an A100, we recommend leavingall model components on GPU for maximum speed, as done in theofficial IF demo.💡 Note: Some of the larger images have been compressed to load faster in the blog format. When using the official model, they should be evenbetter quality!Let's dive in 🚀!IF's text generation capabilitiesTable of contentsAccepting the licenseOptimizing IF to run on memory constrained hardwareAvailable resourcesInstall dependenciesText-to-image generationImage variationInpaintingAccepting the licenseBefore you can use IF, you need to accept its usage conditions. To do so:Make sure to have a Hugging Face account and be logged inAccept the license on the model card of DeepFloyd/IF-I-XL-v1.0. Accepting the license on the stage I model card will auto accept for the other IF models.Make sure to login locally. Install huggingface_hubpip install huggingface_hub --upgraderun the login function in a Python shellfrom huggingface_hub import loginlogin()and enter your Hugging Face Hub access token.Optimizing IF to run on memory constrained hardwareState-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 justthe latest and greatest hardware.The deep learning community has created world class tools to runresource intensive models on consumer hardware:🤗 accelerate providesutilities for working with large models.bitsandbytes makes 8-bit quantization available to all PyTorch models.🤗 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 asimple API when optimizing large models.The free-tier Google Colab is both CPU RAM constrained (13 GB RAM) aswell as GPU VRAM constrained (15 GB RAM for T4), which makes running thewhole >10B IF model challenging!Let's map out the size of IF's model components in full float32precision:T5-XXL Text Encoder: 20GBStage 1 UNet: 17.2 GBStage 2 Super Resolution UNet: 2.5 GBStage 3 Super Resolution Model: 3.4 GBThere is no way we can run the model in float32 as the T5 and Stage 1UNet 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, butwe'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 usingbitsandbytes 8bit quantization, which allows saving the T5 checkpointwith as little as 8GB.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 forthemselves when needed.Diffusers supports modularly loading individual components i.e. we canload the text encoder without loading the UNet. This modular loadingwill ensure that we only load the component we need at a given step inthe pipeline to avoid exhausting the available CPU RAM and GPU VRAM.Let's give it a try 🚀Available resourcesThe free-tier Google Colab comes with around 13 GB CPU RAM:!grep MemTotal /proc/meminfoMemTotal: 13297192 kBAnd an NVIDIA T4 with 15 GB VRAM:!nvidia-smiSun Apr 23 23:14:19 2023 +-----------------------------------------------------------------------------+| NVIDIA-SMI 525.85.12 Driver Version: 525.85.12 CUDA Version: 12.0 ||-------------------------------+----------------------+----------------------+| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC || Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. || | | MIG M. ||===============================+======================+======================|| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 || N/A 72C P0 32W / 70W | 1335MiB / 15360MiB | 0% Default || | | N/A |+-------------------------------+----------------------+----------------------++-----------------------------------------------------------------------------+| Processes: || GPU GI CI PID Type Process name GPU Memory || ID ID Usage ||=============================================================================|+-----------------------------------------------------------------------------+Install dependenciesSome optimizations can require up-to-date versions of dependencies. Ifyou are having issues, please double check and upgrade versions.! pip install --upgrade \diffusers~=0.16 \transformers~=4.28 \safetensors~=0.3 \sentencepiece~=0.1 \accelerate~=0.18 \bitsandbytes~=0.38 \torch~=2.0 -q1. Text-to-image generationWe will walk step by step through text-to-image generation with IF usingDiffusers. We will explain briefly APIs and optimizations, but morein-depth explanations can be found in the official documentation forDiffusers,Transformers,Accelerate, andbitsandbytes.1.1 Load text encoderWe will load T5 using 8bit quantization. Transformers directly supportsbitsandbytesthrough the load_in_8bit flag.The flag variant="8bit" will download pre-quantized weights.We also use the device_map flag to allow transformers to offloadmodel layers to the CPU or disk. Transformers big modeling supportsarbitrary device maps, which can be used to separately load modelparameters directly to available devices. Passing "auto" willautomatically create a device map. See the transformersdocsfor more information.from transformers import T5EncoderModeltext_encoder = T5EncoderModel.from_pretrained("DeepFloyd/IF-I-XL-v1.0",subfolder="text_encoder", device_map="auto", load_in_8bit=True, variant="8bit")1.2 Create text embeddingsThe Diffusers API for accessing diffusion models is theDiffusionPipeline class and its subclasses. Each instance ofDiffusionPipeline is a fully self contained set of methods and modelsfor running diffusion networks. We can override the models it uses bypassing alternative instances as keyword arguments to from_pretrained.In this case, we pass None for the unet argument, so no UNet will beloaded. This allows us to run the text embedding portion of thediffusion process without loading the UNet into memory.from diffusers import DiffusionPipelinepipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", text_encoder=text_encoder, # pass the previously instantiated 8bit text encoderunet=None, device_map="auto")IF also comes with a super resolution pipeline. We will save the promptembeddings so we can later directly pass them to the superresolution pipeline. This will allow the super resolution pipeline to beloaded without a text encoder.Instead of an astronaut just riding ahorse, let's hand them asign as well!Let's define a fitting prompt:prompt = "a photograph of an astronaut riding a horse holding a sign that says Pixel's in space"and run it through the 8bit quantized T5 model:prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)1.3 Free memoryOnce the prompt embeddings have been created. We do not need the textencoder anymore. However, it is still in memory on the GPU. We need toremove it so that we can load the UNet.It's non-trivial to free PyTorch memory. We must garbage-collect thePython objects which point to the actual memory allocated on the GPU.First, use the Python keyword del to delete all Python objectsreferencing allocated GPU memorydel text_encoderdel pipeDeleting the python object is not enough to free the GPU memory.Garbage collection is when the actual GPU memory is freed.Additionally, we will call torch.cuda.empty_cache(). This methodisn't strictly necessary as the cached cuda memory will be immediatelyavailable for further allocations. Emptying the cache allows us toverify in the Colab UI that the memory is available.We'll use a helper function flush() to flush memory.import gcimport torchdef flush():gc.collect()torch.cuda.empty_cache()and run itflush()1.4 Stage 1: The main diffusion processWith our now available GPU memory, we can re-load theDiffusionPipeline with only the UNet to run the main diffusionprocess.The variant and torch_dtype flags are used by Diffusers to downloadand load the weights in 16 bit floating point format.pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16, device_map="auto")Often, we directly pass the text prompt to DiffusionPipeline.__call__.However, we previously computed our text embeddings which we can passinstead.IF also comes with a super resolution diffusion process. Settingoutput_type="pt" will return raw PyTorch tensors instead of a PILimage. This way, we can keep the PyTorch tensors on GPU and pass themdirectly to the stage 2 super resolution pipeline.Let's define a random generator and run the stage 1 diffusion process.generator = torch.Generator().manual_seed(1)image = pipe(prompt_embeds=prompt_embeds,negative_prompt_embeds=negative_embeds, output_type="pt",generator=generator,).imagesLet's manually convert the raw tensors to PIL and have a sneak peek atthe final result. The output of stage 1 is a 64x64 image.from diffusers.utils import pt_to_pilpil_image = pt_to_pil(image)pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)pil_image[0]And again, we remove the Python pointer and free CPU and GPU memory:del pipeflush()1.5 Stage 2: Super Resolution 64x64 to 256x256IF comes with a separate diffusion process for upscaling.We run each diffusion process with a separate pipeline.The super resolution pipeline can be loaded with a text encoder ifneeded. However, we will usually have pre-computed text embeddings fromthe first IF pipeline. If so, load the pipeline without the textencoder.Create the pipelinepipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-II-L-v1.0", text_encoder=None, # no use of text encoder => memory savings!variant="fp16", torch_dtype=torch.float16, device_map="auto")and run it, re-using the pre-computed text embeddingsimage = pipe(image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, output_type="pt",generator=generator,).imagesAgain we can inspect the intermediate results.pil_image = pt_to_pil(image)pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)pil_image[0]And again, we delete the Python pointer and free memorydel pipeflush()1.6 Stage 3: Super Resolution 256x256 to 1024x1024The second super resolution model for IF is the previously releaseStability AI's x4Upscaler.Let's create the pipeline and load it directly on GPU withdevice_map="auto".pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", torch_dtype=torch.float16, device_map="auto")🧨 diffusers makes independently developed diffusion models easilycomposable as pipelines can be chained together. Here we can just takethe previous PyTorch tensor output and pass it to the tage 3 pipeline asimage=image.💡 Note: The x4 Upscaler does not use T5 and has its own textencoder.Therefore, we cannot use the previously created prompt embeddings andinstead must pass the original prompt.pil_image = pipe(prompt, generator=generator, image=image).imagesUnlike the IF pipelines, the IF watermark will not be added by defaultto outputs from the Stable Diffusion x4 upscaler pipeline.We can instead manually apply the watermark.from diffusers.pipelines.deepfloyd_if import IFWatermarkerwatermarker = IFWatermarker.from_pretrained("DeepFloyd/IF-I-XL-v1.0", subfolder="watermarker")watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)View output imagepil_image[0]Et voila! A beautiful 1024x1024 image in a free-tier Google Colab.We have shown how 🧨 diffusers makes it easy to decompose and modularlyload resource-intensive diffusion models.💡 Note: We don't recommend using the above setup in production.8bit quantization, manual de-allocation of model weights, and diskoffloading all trade off memory for time (i.e., inference speed). Thiscan be especially noticable if the diffusion pipeline is re-used. Inproduction, we recommend using a 40GB A100 with all model componentsleft on the GPU. See the official IFdemo.2. Image variationThe same IF checkpoints can also be used for text guided image variationand inpainting. The core diffusion process is the same as text-to-imagegeneration except the initial noised image is created from the image tobe varied or inpainted.To run image variation, load the same checkpoints withIFImg2ImgPipeline.from_pretrained() andIFImg2ImgSuperResolution.from_pretrained().The APIs for memory optimization are all the same!Let's free the memory from the previous section.del pipeflush()For image variation, we start with an initial image that we want toadapt.For this section, we will adapt the famous "Slaps Roof of Car" meme.Let's download it from the internet.import requestsurl = "https://i.kym-cdn.com/entries/icons/original/000/026/561/car.jpg"response = requests.get(url)and load it into a PIL Imagefrom PIL import Imagefrom io import BytesIOoriginal_image = Image.open(BytesIO(response.content)).convert("RGB")original_image = original_image.resize((768, 512))original_imageThe image variation pipeline take both PIL images and raw tensors. Viewthe docstrings for more indepth documentation on expected inputs, here.2.1 Text EncoderImage variation is guided by text, so we can define a prompt and encodeit with T5's Text Encoder.Again we load the text encoder into 8bit precision.from transformers import T5EncoderModeltext_encoder = T5EncoderModel.from_pretrained("DeepFloyd/IF-I-XL-v1.0",subfolder="text_encoder", device_map="auto", load_in_8bit=True, variant="8bit")For image variation, we load the checkpoint withIFImg2ImgPipeline. When usingDiffusionPipeline.from_pretrained(...), checkpoints are loaded intotheir default pipeline. The default pipeline for the IF is thetext-to-image IFPipeline. When loading checkpointswith a non-default pipeline, the pipeline must be explicitly specified.from diffusers import IFImg2ImgPipelinepipe = IFImg2ImgPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", text_encoder=text_encoder, unet=None, device_map="auto")Let's turn our salesman into an anime character.prompt = "anime style"As before, we create the text embeddings with T5prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)and free GPU and CPU memory.First, remove the Python pointersdel text_encoderdel pipeand then free the memoryflush()2.2 Stage 1: The main diffusion processNext, we only load the stage 1 UNet weights into the pipeline object,just like we did in the previous section.pipe = IFImg2ImgPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16, device_map="auto")The image variation pipeline requires both the original image and theprompt embeddings.We can optionally use the strength argument to configure the amount ofvariation. strength directly controls the amount of noise added.Higher strength means more noise which means more variation.generator = torch.Generator().manual_seed(0)image = pipe(image=original_image,prompt_embeds=prompt_embeds,negative_prompt_embeds=negative_embeds, output_type="pt",generator=generator,).imagesLet's check the intermediate 64x64 again.pil_image = pt_to_pil(image)pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)pil_image[0]Looks good! We can free the memory and upscale the image again.del pipeflush()2.3 Stage 2: Super ResolutionFor super resolution, load the checkpoint withIFImg2ImgSuperResolutionPipeline and the same checkpoint as before.from diffusers import IFImg2ImgSuperResolutionPipelinepipe = IFImg2ImgSuperResolutionPipeline.from_pretrained("DeepFloyd/IF-II-L-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16, device_map="auto")💡 Note: The image variation super resolution pipeline requires thegenerated image as well as the original image.You can also use the Stable Diffusion x4 upscaler on this image. Feelfree to try it out using the code snippets in section 1.6.image = pipe(image=image,original_image=original_image,prompt_embeds=prompt_embeds,negative_prompt_embeds=negative_embeds, generator=generator,).images[0]imageNice! Let's free the memory and look at the final inpainting pipelines.del pipeflush()3. InpaintingThe IF inpainting pipeline is the same as the image variation, exceptonly a select area of the image is denoised.We specify the area to inpaint with an image mask.Let's show off IF's amazing "letter generation" capabilities. We canreplace this sign text with different slogan.First let's download the imageimport requestsurl = "https://i.imgflip.com/5j6x75.jpg"response = requests.get(url)and turn it into a PIL Imagefrom PIL import Imagefrom io import BytesIOoriginal_image = Image.open(BytesIO(response.content)).convert("RGB")original_image = original_image.resize((512, 768))original_imageWe will mask the sign so we can replace its text.For convenience, we have pre-generated the mask and loaded it into a HFdataset.Let's download it.from huggingface_hub import hf_hub_downloadmask_image = hf_hub_download("diffusers/docs-images", repo_type="dataset", filename="if/sign_man_mask.png")mask_image = Image.open(mask_image)mask_image💡 Note: You can create masks yourself by manually creating agreyscale image.from PIL import Imageimport numpy as npheight = 64width = 64example_mask = np.zeros((height, width), dtype=np.int8)# Set masked pixels to 255example_mask[20:30, 30:40] = 255# Make sure to create the image in mode 'L'# meaning single channel grayscaleexample_mask = Image.fromarray(example_mask, mode='L')example_maskNow we can start inpainting 🎨🖌3.1. Text EncoderAgain, we load the text encoder firstfrom transformers import T5EncoderModeltext_encoder = T5EncoderModel.from_pretrained("DeepFloyd/IF-I-XL-v1.0",subfolder="text_encoder", device_map="auto", load_in_8bit=True, variant="8bit")This time, we initialize the IFInpaintingPipeline in-painting pipelinewith the text encoder weights.from diffusers import IFInpaintingPipelinepipe = IFInpaintingPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", text_encoder=text_encoder, unet=None, device_map="auto")Alright, let's have the man advertise for more layers instead.prompt = 'the text, "just stack more layers"'Having defined the prompt, we can create the prompt embeddingsprompt_embeds, negative_embeds = pipe.encode_prompt(prompt)Just like before, we free the memorydel text_encoderdel pipeflush()3.2 Stage 1: The main diffusion processJust like before, we now load the stage 1 pipeline with only the UNet.pipe = IFInpaintingPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16, device_map="auto")Now, we need to pass the input image, the mask image, and the promptembeddings.image = pipe(image=original_image,mask_image=mask_image,prompt_embeds=prompt_embeds,negative_prompt_embeds=negative_embeds, output_type="pt",generator=generator,).imagesLet's take a look at the intermediate output.pil_image = pt_to_pil(image)pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)pil_image[0]Looks good! The text is pretty consistent!Let's free the memory so we can upscale the imagedel pipeflush()3.3 Stage 2: Super ResolutionFor super resolution, load the checkpoint withIFInpaintingSuperResolutionPipeline.from diffusers import IFInpaintingSuperResolutionPipelinepipe = IFInpaintingSuperResolutionPipeline.from_pretrained("DeepFloyd/IF-II-L-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16, device_map="auto")The inpainting super resolution pipeline requires the generated image,the original image, the mask image, and the prompt embeddings.Let's do a final denoising run.image = pipe(image=image,original_image=original_image,mask_image=mask_image,prompt_embeds=prompt_embeds,negative_prompt_embeds=negative_embeds, generator=generator,).images[0]imageNice, the model generated text without making a singlespelling error!ConclusionIF in 32-bit floating point precision uses 40 GB of weights in total. Weshowed how using only open source models and libraries, IF can be run ona free-tier Google Colab instance.The ML ecosystem benefits deeply from the sharing of open tools and openmodels. This notebook alone used models from DeepFloyd, StabilityAI, andGoogle. The libraries used -- Diffusers, Transformers, Accelerate, andbitsandbytes -- all benefit from countless contributors from differentorganizations.A massive thank you to the DeepFloyd team for the creation and opensourcing of IF, and for contributing to the democratization of goodmachine learning 🤗.