Перейти к содержанию
Публикация AiManual

Запуск IF на бесплатном Google Colab: полное руководство

Запустите DeepFloyd IF на бесплатном Google Colab: 8-битное квантование T5, модульная загрузка компонентов и все три этапа генерации до 1024x1024. Готовый код д

Коротко

Что будет в материале

  1. 01

    Введение: IF - мощная модель генерации изображений

  2. 02

    Подготовка среды в Google Colab

  3. 03

    Оптимизация памяти: ключ к успеху

  4. 04

    Пошаговый запуск трех этапов генерации

Модель IF от DeepFloyd суммарно содержит более 10 миллиардов параметров и требует около 50 ГБ памяти при стандартной загрузке в float32. Бесплатный Colab даёт 13 ГБ CPU RAM и 15 ГБ VRAM. Запуск возможен, если разбить модель на модули и квантовать текстовый энкодер. В этом руководстве показан рабочий пайплайн: 8-битное квантование T5, последовательная загрузка компонентов, освобождение памяти между этапами и финальная генерация изображения 1024x1024.

Вы получите готовый код для text-to-image, image variation и inpainting. Также разберём, как сохранять предварительно вычисленные эмбеддинги и повторно использовать их без повторного кодирования текста.

Введение: IF - мощная модель генерации изображений

DeepFloyd IF использует каскадную архитектуру из трёх компонентов. Первый - текстовый энкодер T5-XXL с 4,5 миллиарда параметров, который преобразует промпт в эмбеддинги. Второй - UNet, генерирующий изображение 64x64 пикселя. Третий - два модуля суперразрешения, которые последовательно увеличивают картинку до 256x256 и затем до 1024x1024.

Качество генерации IF сопоставимо с топовыми диффузионными моделями, а в задачах точного следования тексту модель часто показывает лучшие результаты. Плата за это - высокие требования к памяти. Полная загрузка всех компонентов в float32 занимает свыше 40 ГБ, что недоступно на бесплатном Colab. Решение: квантование и модульная работа с компонентами.

Если вы уже работали с библиотекой Diffusers, материал будет знакомым. Для погружения в эволюцию этой библиотеки и её роль в поддержке DeepFloyd IF рекомендую разбор первого года Diffusers.

Подготовка среды в Google Colab

Начните с проверки доступных ресурсов. Бесплатный тариф Colab обычно предоставляет GPU Tesla T4 с 15 ГБ VRAM и около 13 ГБ системной памяти. Этого достаточно для запуска IF при правильной оптимизации.

Установка необходимых библиотек

Установите четыре библиотеки: diffusers для пайплайнов, transformers для загрузки T5, accelerate для управления устройствами и bitsandbytes для 8-битного квантования.

!pip install diffusers transformers accelerate bitsandbytes

После установки проверьте версии. Разные версии diffusers могут иметь отличия в API, поэтому фиксируйте совместимые версии, если пайплайн уже отлажен.

Проверка ресурсов Colab

Выполните команды для вывода информации о GPU и оперативной памяти:

!nvidia-smi
!free -h

В выводе nvidia-smi ищите строку с Tesla T4 и объёмом памяти 15109MiB. Если GPU не назначен, включите его через меню Runtime → Change runtime type → GPU.

Оптимизация памяти: ключ к успеху

Стандартный запуск IF через DiffusionPipeline загружает все компоненты одновременно. На Colab это приводит к ошибке Out of Memory. Нужны две техники: квантование T5 и модульная загрузка.

8-битное квантование текстового энкодера T5

T5-XXL в float32 занимает около 18 ГБ. Квантование в 8 бит через bitsandbytes сокращает потребление до 5-6 ГБ. Загрузка выполняется с параметром load_in_8bit:

from transformers import T5EncoderModel, BitsAndBytesConfig

quant_config = BitsAndBytesConfig(load_in_8bit=True)
text_encoder = T5EncoderModel.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="text_encoder",
    quantization_config=quant_config,
    device_map="auto"
)

Квантование незначительно влияет на качество эмбеддингов, но радикально снижает требования к VRAM. Это стандартный компромисс для запуска больших моделей на ограниченном железе.

Модульная загрузка и освобождение памяти

Суть подхода: в каждый момент времени в памяти находится только один компонент. Сначала загружается T5, вычисляются эмбеддинги, затем T5 удаляется. Потом загружается UNet, генерируется изображение 64x64, UNet удаляется. Далее по очереди загружаются модули суперразрешения.

Освобождение памяти выполняется через del и torch.cuda.empty_cache():

import torch

del text_encoder
torch.cuda.empty_cache()

Этот цикл повторяется после каждого этапа. Без явного вызова empty_cache фрагментация памяти может привести к сбоям даже при формально достаточном объёме.

Пошаговый запуск трех этапов генерации

Пайплайн IF состоит из трёх последовательных шагов. Каждый шаг использует свой компонент модели и требует отдельной загрузки.

Этап 1: Создание текстовых эмбеддингов

Текстовый промпт кодируется в эмбеддинги с помощью T5. Эти эмбеддинги будут использованы на этапе генерации базового изображения.

from transformers import T5Tokenizer

tokenizer = T5Tokenizer.from_pretrained("DeepFloyd/IF-I-XL-v1.0", subfolder="tokenizer")
prompt = "a photorealistic image of a red fox in a snowy forest"
text_inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)

with torch.no_grad():
    text_embeddings = text_encoder(**text_inputs.to("cuda")).last_hidden_state

Сохраните эмбеддинги в переменную и освободите T5. Эмбеддинги занимают немного памяти, их можно хранить в RAM или на диске.

Этап 2: Генерация изображения 64x64

Загрузите UNet и диффузионный планировщик. Передайте текстовые эмбеддинги и запустите генерацию.

from diffusers import DDPMScheduler, IFPipeline

scheduler = DDPMScheduler.from_pretrained("DeepFloyd/IF-I-XL-v1.0", subfolder="scheduler")
unet = UNet2DConditionModel.from_pretrained("DeepFloyd/IF-I-XL-v1.0", subfolder="unet").to("cuda")

pipe = IFPipeline(
    text_encoder=None,
    tokenizer=tokenizer,
    unet=unet,
    scheduler=scheduler
)

image_64 = pipe(
    prompt_embeds=text_embeddings,
    output_type="pt",
    num_inference_steps=50
).images

На выходе получается тензор изображения 64x64. UNet занимает около 4 ГБ в float32, что укладывается в лимиты T4.

Этап 3: Суперразрешение до 1024x1024

Сначала изображение увеличивается до 256x256 с помощью первой модели суперразрешения, затем до 1024x1024 с помощью второй.

from diffusers import IFSuperResolutionPipeline

# Удаляем UNet, загружаем первый апскейлер
del unet
torch.cuda.empty_cache()

sr_pipe_256 = IFSuperResolutionPipeline.from_pretrained(
    "DeepFloyd/IF-II-L-v1.0",
    text_encoder=None,
    tokenizer=tokenizer
).to("cuda")

image_256 = sr_pipe_256(
    image=image_64,
    prompt_embeds=text_embeddings,
    output_type="pt",
    num_inference_steps=50
).images

del sr_pipe_256
torch.cuda.empty_cache()

# Загружаем второй апскейлер
sr_pipe_1024 = IFSuperResolutionPipeline.from_pretrained(
    "DeepFloyd/IF-III-L-v1.0",
    text_encoder=None,
    tokenizer=tokenizer
).to("cuda")

image_1024 = sr_pipe_1024(
    image=image_256,
    prompt_embeds=text_embeddings,
    output_type="pt",
    num_inference_steps=75
).images

Финальное изображение имеет разрешение 1024x1024. Каждый апскейлер загружается и удаляется по очереди, чтобы не превысить лимит VRAM.

Использование предварительно вычисленных эмбеддингов

Кодирование текста через T5-XXL занимает 20-40 секунд. При повторных запусках с тем же промптом это время можно сэкономить. Сохраните эмбеддинги на диск:

torch.save(text_embeddings.cpu(), "text_embeddings.pt")

При следующем запуске загрузите их напрямую, пропуская этап T5:

text_embeddings = torch.load("text_embeddings.pt").to("cuda")

Это особенно полезно при отладке пайплайна или генерации множества изображений с одинаковым промптом.

Дополнительные возможности: image variation и inpainting

IF поддерживает три режима работы. Text-to-image мы разобрали. Два других - вариации изображения и инпейнтинг - используют ту же архитектуру, но с другими входными данными.

Image variation

Для генерации вариаций существующего изображения вместо текстовых эмбеддингов используются эмбеддинги изображения. Они извлекаются через отдельный энкодер изображений:

from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection

image_encoder = CLIPVisionModelWithProjection.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="image_encoder"
).to("cuda")

image_processor = CLIPImageProcessor.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="image_processor"
)

image_inputs = image_processor(source_image, return_tensors="pt")
with torch.no_grad():
    image_embeddings = image_encoder(**image_inputs.to("cuda")).image_embeds

Дальше пайплайн аналогичен text-to-image, но вместо prompt_embeds передаются image_embeds.

Inpainting

Инпейнтинг позволяет редактировать часть изображения по текстовому описанию. Нужны исходное изображение, маска и промпт. Пайплайн комбинирует текстовые эмбеддинги с замаскированным изображением:

from diffusers import IFInpaintingPipeline

inpaint_pipe = IFInpaintingPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    text_encoder=None,
    tokenizer=tokenizer,
    unet=unet,
    scheduler=scheduler
)

result = inpaint_pipe(
    prompt_embeds=text_embeddings,
    image=source_image,
    mask_image=mask,
    output_type="pt",
    num_inference_steps=50
).images

Маска должна быть бинарной: белые пиксели обозначают область для редактирования, чёрные - сохраняемую область.

Для понимания общих принципов оптимизации инференса диффузионных моделей полезен разбор ускорения Stable Diffusion на CPU Intel. Техники управления памятью и батчами применимы и к IF.

Заключение: итоги и перспективы

Запуск IF на бесплатном Colab - решаемая задача. Ключевые шаги: квантование T5 до 8 бит, модульная загрузка компонентов, освобождение памяти между этапами и использование предварительно вычисленных эмбеддингов. В результате модель с 10+ млрд параметров работает на Tesla T4 с 15 ГБ VRAM.

Ограничения: генерация занимает 3-5 минут на изображение, бесплатный тариф Colab может отключать сессию при длительной неактивности, а качество 8-битного квантования T5 незначительно уступает float32. Для продакшена рассмотрите платные GPU или оптимизированные пайплайны.

Если вы сравниваете разные генеративные модели для своих задач, обратите внимание на обзор Gemini 3.6 Flash с бенчмарками и примерами кода.

Подписаться на канал