Running DeepFloyd IF on Free‑Tier Google Colab with Diffusers

TL;DR

DeepFloyd’s IF model can be executed on a free‑tier Google Colab by quantizing the T5‑XXL text encoder to 8‑bit, loading model components modularly with 🤗 Diffusers, and offloading weights to CPU or disk, enabling full text‑to‑image, image‑variation, and inpainting pipelines despite the notebook’s 13 GB RAM and 15 GB GPU limits.

Introduction – IF’s capabilities and constraints

DeepFloyd released IF in late April 2023 as a pixel‑space text‑to‑image diffusion model inspired by Google’s Imagen. Compared with Stable Diffusion, IF:

  • Operates directly on uncompressed images, preserving high‑frequency details such as faces and hands.
  • Uses the powerful T5‑XXL encoder instead of CLIP, improving text fidelity and making it the first open‑source model that reliably renders readable text.

The trade‑off is size: T5‑XXL (4.5 B parameters), IF‑I UNet (4.3 B), and IF‑II upscaler UNet (1.2 B) together exceed 10 B parameters, far larger than Stable Diffusion 2.1’s ~1.3 B total. Running the full model in float32 would require >40 GB of GPU memory, which is unavailable on a free Colab T4 (15 GB VRAM) or its 13 GB CPU RAM.

Optimizing for memory‑constrained hardware

Key tools

  • 🤗 Accelerate – utilities for large‑model device placement.
  • bitsandbytes – 8‑bit quantization for PyTorch models.
  • 🤗 safetensors – fast, safe checkpoint loading.
  • 🤗 Diffusers – high‑level diffusion pipelines that integrate the above.

Memory‑saving strategy

  1. Quantize T5‑XXL to 8‑bit – reduces the encoder checkpoint from ~20 GB (fp32) to ~8 GB (8‑bit safetensors).
  2. Load components lazily – Diffusers allows loading only the text encoder or only the UNet at a given time, preventing simultaneous memory peaks.
  3. Use fp16 for UNet weights – Stage 1 and Stage 2 UNets fit in GPU memory when loaded as half‑precision.
  4. Explicitly free PyTorch objects – del the model objects, run gc.collect() and torch.cuda.empty_cache() after each stage.

Step‑by‑step: Text‑to‑Image generation

1. Install up‑to‑date dependencies

pip install --upgrade \
  diffusers~=0.16 \
  transformers~=4.28 \
  safetensors~=0.3 \
  sentencepiece~=0.1 \
  accelerate~=0.18 \
  bitsandbytes~=0.38 \
  torch~=2.0 -q

2. Load the 8‑bit T5 encoder

from transformers import T5EncoderModel
text_encoder = T5EncoderModel.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="text_encoder",
    device_map="auto",
    load_in_8bit=True,
    variant="8bit",
)

3. Create prompt embeddings without loading the UNet

from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    text_encoder=text_encoder,
    unet=None,
    device_map="auto",
)
prompt = "a photograph of an astronaut riding a horse holding a sign that says Pixel's in space"
prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)

4. Free the encoder to make room for the UNet

import gc, torch

def flush():
    gc.collect()
    torch.cuda.empty_cache()

del text_encoder, pipe
flush()

5. Stage 1 diffusion (64×64) – load only the UNet

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    text_encoder=None,
    variant="fp16",
    torch_dtype=torch.float16,
    device_map="auto",
)
generator = torch.Generator().manual_seed(1)
image = pipe(
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds,
    output_type="pt",
    generator=generator,
).images

The result is a 64 × 64 tensor that can be visualized with pt_to_pil.

6. Stage 2 super‑resolution (64→256) – IF‑II pipeline

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-II-L-v1.0",
    text_encoder=None,
    variant="fp16",
    torch_dtype=torch.float16,
    device_map="auto",
)
image = pipe(
    image=image,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds,
    output_type="pt",
    generator=generator,
).images

7. Stage 3 super‑resolution (256→1024) – StabilityAI x4 upscaler

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-x4-upscaler",
    torch_dtype=torch.float16,
    device_map="auto",
)
final_image = pipe(prompt, image=image, generator=generator).images[0]

Apply the IF watermark manually if desired:

from diffusers.pipelines.deepfloyd_if import IFWatermarker
watermarker = IFWatermarker.from_pretrained("DeepFloyd/IF-I-XL-v1.0", subfolder="watermarker")
watermarker.apply_watermark(final_image, pipe.unet.config.sample_size)

The pipeline now yields a 1024 × 1024 image entirely within a free Colab session.

Image variation and inpainting – reusing the same checkpoints

The IF checkpoints also support IFImg2ImgPipeline (variation) and IFInpaintingPipeline. The workflow mirrors the text‑to‑image steps:

  1. Load the 8‑bit T5 encoder, encode a variation prompt, then free the encoder.
  2. Load the Stage 1 UNet with unet=None for the variation pipeline, pass the original image and strength to control how much noise is added.
  3. Upscale with IFImg2ImgSuperResolutionPipeline (or the Stable Diffusion upscaler).
  4. For inpainting, additionally provide a binary mask image; the same modular loading and memory‑freeing pattern applies.

All three pipelines share the same memory‑optimisation tricks: 8‑bit text encoder, fp16 UNet, device‑map auto placement, and explicit garbage collection.

Practical considerations and performance trade‑offs

  • Speed vs. memory – 8‑bit quantization and repeated loading/unloading increase inference latency. For production, a GPU with ≥40 GB VRAM (e.g., A100) should keep all components on‑device for maximal throughput.
  • Quality – The official IF demo (hosted on Hugging Face Spaces) runs the full‑precision model and yields slightly higher fidelity, especially for large images.
  • License – Users must accept the DeepFloyd IF license on the model card before loading any checkpoint.

Conclusion – Democratizing large diffusion models

By combining open‑source checkpoints (DeepFloyd IF, StabilityAI upscaler) with community‑driven libraries (Diffusers, Transformers, Accelerate, bitsandbytes), Hugging Face demonstrates that a >10 B‑parameter diffusion model can be run on a free Google Colab instance. This showcases the power of modular pipelines and quantization to broaden access to state‑of‑the‑art generative AI.

Sources