Создание более интеллектуальных моделей преобразования текста в изображение с помощью ControlNet

Управляйте моделями диффузии, добавляя дополнительные условия

Авторы: Саи Бхаргав, Пратмеш, Кавья, Алиша, Амритангшу

В последние годы модели преобразования текста в изображение добились впечатляющих успехов в создании высококачественных изображений из текстовых описаний. Эти модели учатся преобразовывать естественный язык в пространство пикселей, что позволяет им генерировать реалистичные изображения, соответствующие входному описанию. Тем не менее, еще есть возможности для улучшения с точки зрения контроля над генерируемыми изображениями. Добавление условного контроля к моделям распространения текста в изображение может позволить пользователям более точно направлять создание изображений, что приводит к более точным и адаптированным результатам. В связи с этим все более популярным становится использование эскизов, набросков, поз как средств условного контроля. В этом блоге мы рассмотрим концепцию условного управления в моделях диффузии текста в изображение, в частности, как различные методы могут повысить точность генерации изображений, а также преимущества использования таких методов.

В этом блоге мы рассмотрим концепцию объединения ControlNet со стабильной диффузией и ее потенциал для улучшения преобразования текста в изображение.

Что такое стабильная диффузия?

Стабильная диффузия — это вариант диффузионной модели, генеративной модели, которая учится генерировать высококачественные изображения путем моделирования распределения промежуточных источников шума. Модель диффузии представляет собой последовательный процесс, который начинается с фиксированного распределения шума и применяет последовательность преобразований, каждое из которых изменяет распределение шума для создания нового распределения, которое лучше соответствует целевому распределению. Окончательное распределение используется для создания выходного изображения. Было показано, что модель диффузии эффективна для создания высококачественных изображений, но ее обучение может быть сложным из-за необходимости оценивать сложные условные распределения.

Распространение — это процесс, происходящий внутри розового компонента «создатель информации об изображении». Имея вложения токенов, которые представляют входной текст, и случайный начальный массив информации об изображении (их также называют скрытыми), процесс создает массив информации, который использует декодер изображения. чтобы нарисовать окончательный образ.

Эта процедура проводится поэтапно. Каждый шаг добавляет больше релевантных данных. Мы можем проверить массив случайных латентных данных и посмотреть, как он преобразуется в визуальный шум, чтобы получить представление о процессе. В этом случае визуальный осмотр влечет за собой прогон его через декодер изображений. Распространение происходит в несколько этапов, каждый из которых работает с входным массивом латентных данных и создает другой массив латентных значений, который больше напоминает входной текст, а также всю визуальную информацию, полученную моделью из всех изображений, на которых она обучалась. Набор этих скрытых данных можно визуализировать, чтобы увидеть, какая информация добавляется на каждом этапе.

Использование стабильной диффузии на python с использованием конвейера диффузоров

#!pip install -q diffusers==0.14.0 transformers xformers git+https://github.com/huggingface/accelerate.git
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
import torch

repo_id = "stabilityai/stable-diffusion-2-base"
pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.float16, revision="fp16")

pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

prompt = "Panda in space"
image = pipe(prompt, num_inference_steps=25).images[0]
image.save("astronaut.png")

Что такое ControlNet?

ControlNet — это структура нейронной сети, которая позволяет явно управлять предварительно обученными моделями большой диффузии, обеспечивая поддержку дополнительных входных условий. ControlNet предназначен для комплексного изучения конкретных условий задачи, что делает процесс обучения надежным, даже если набор обучающих данных невелик и содержит менее 50 000 образцов. Эта возможность делает ControlNet мощным инструментом для преобразования текста в изображение, где управляющие переменные могут использоваться для указания желаемых атрибутов сгенерированных изображений, таких как форма, цвет или текстура.

Как это работает

ControlNet работает, копируя веса блоков нейронной сети в две копии: «заблокированную» копию и «обучаемую» копию. «Заблокированная» копия сохраняет исходную модель, а «обучаемая» копия используется для изучения заданных условий.

Во время обучения ControlNet использует технику, называемую «нулевой сверткой», которая представляет собой свертку 1x1, в которой как вес, так и смещение инициализируются нулями. Перед обучением все нулевые свертки выводят нули, и ControlNet не вызывает никаких искажений исходной модели. «Обучаемая» копия обучается заданным условиям, в то время как «заблокированная» копия остается неизменной.

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

Использование конвейера StableDiffusion ControlNet

мы используем Pipeline для генерации текста в изображение, используя Stable Diffusion с руководством ControlNet

Доступны предварительно обученные модели:

• Управление с помощью обнаружения краев: lllyasviel/sd-controlnet-canny: обучено обнаружению краев, используемому для управления краями изображения.
• Управление с помощью определения позы: lllyasviel/sd-controlnet_openpose: обучение с использованием изображения кости OpenPose, используемого для управления позами человека.
• Управление с помощью обнаружения каракулей:lllyasviel/sd-controlnet_scribble: обучено человеческим каракулям, используемым для управления контурами изображения.
• lllyasviel/sd-controlnet-depth: обучение выполнено с оценкой глубины Midas, используемой для управления глубиной изображения.
• lllyasviel/sd-controlnet-hed: Обучено обнаружению краев HED (мягкие края), используемому для управления мягкими краями изображения.
• lllyasviel/sd-controlnet-mlsd: обучен обнаружению линий M-LSD, используемому для управления прямыми линиями на изображении.
• lllyasviel/sd-controlnet-normal: обучен с картой нормалей, используемой для управления нормалями изображения.
• lllyasviel/sd-controlnet_seg: обучен семантической сегментации, используется для управления семантической сегментацией изображения.

Шаг 1. Настройка среды и конвейера

Во-первых, мы должны настроить нашу среду и импортировать стабильную диффузию из обнимающего лица, здесь мы используем стабильную диффузию-v1–5.

Код устанавливает некоторые пакеты с помощью pip, а затем импортирует модули и классы из этих пакетов.

Первая строка устанавливает следующие пакеты:

• diffusers==0.14.0: пакет для реализации моделей распространения в PyTorch.
• transformers: пакет для задач обработки естественного языка (NLP), таких как классификация текста и генерация языка.
• xformers: Пакет для моделей преобразователей, которые представляют собой тип архитектуры нейронной сети, используемый в задачах НЛП.
• git+https://github.com/huggingface/accelerate.git: пакет для оптимизации обучения моделей глубокого обучения.

!pip install -q diffusers==0.14.0 transformers xformers git+https://github.com/huggingface/accelerate.git
!pip install -q opencv-contrib-python
!pip install -q controlnet_aux

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
import torch

controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
)

Шаг 2. Загрузка нашего тестового изображения в данные

from diffusers import StableDiffusionControlNetPipeline
from diffusers.utils import load_image

image = load_image(
    "/content/daniel-craig-007.jpg-303a730.png"
)
image

Шаг 3. Использование алгоритма обнаружения границ Canny

Этот код применяет алгоритм обнаружения границ Кэнни к изображению, преобразуя его в изображение в градациях серого с тремя каналами, устанавливая нижнее и верхнее пороговые значения и используя функцию cv2.Canny() для создания бинарного изображения с белыми пикселями, представляющими края. Полученное изображение затем преобразуется обратно в объект Image с помощью функции PIL Image.fromarray().

import cv2
from PIL import Image
import numpy as np

image = np.array(image)

low_threshold = 100
high_threshold = 200

image = cv2.Canny(image, low_threshold, high_threshold)
image = image[:, :, None]
image = np.concatenate([image, image, image], axis=2)
canny_image = Image.fromarray(image)
canny_image

Шаг 4. Использование конвейера для создания выходных данных на основе входных данных и SD

Результатом объекта pipe является набор из четырех изображений, которые генерируются на основе подсказок и предоставленных параметров управления. Подсказки представлены в виде списка из четырех строк, каждая из которых представляет собой имя знаменитости, за которым следует сообщение, указывающее желаемое качество изображения. Подсказки указывают на стремление к «наилучшему качеству» и «чрезвычайно детализированным» изображениям. Параметр negative_prompt установлен для указания желаемого качества изображения в негативном ключе, запрашивая изображения «монохромное», «низкое разрешение», «плохая анатомия», «наихудшее качество» и «низкое качество». Параметр generator используется для управления случайностью генерируемых изображений, а параметр num_inference_steps задает количество итераций, используемых для генерации каждого изображения.

def image_grid(imgs, rows, cols):
    assert len(imgs) == rows * cols

    w, h = imgs[0].size
    grid = Image.new("RGB", size=(cols * w, rows * h))
    grid_w, grid_h = grid.size

    for i, img in enumerate(imgs):
        grid.paste(img, box=(i % cols * w, i // cols * h))
    return grid

prompt = ", best quality, extremely detailed"
prompt = [t + prompt for t in ["Tom Cruise", "Donald Trump", "rihanna", "taylor swift"]]
generator = [torch.Generator(device="cpu").manual_seed(2) for i in range(len(prompt))]

output = pipe(
    prompt,
    canny_image,
    negative_prompt=["monochrome, lowres, bad anatomy, worst quality, low quality"] * len(prompt),
    generator=generator,
    num_inference_steps=20,
)

image_grid(output.images, 2, 2)

Шаг 5. Использование OpenposeDetector для определения позы на изображении

Этот код загружает предварительно обученную модель OpenposeDetector из концентратора моделей Hugging Face с помощью метода from_pretrained() и применяет ее к изображению для обнаружения человеческих поз. Он также загружает предварительно обученные ControlNetModel и StableDiffusionControlNetPipeline из пакета diffusers, устанавливает некоторые параметры управления для конвейера и включает некоторые оптимизации памяти и производительности.

from controlnet_aux import OpenposeDetector
model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
poses = model(image)
#image_grid(poses, 2, 2)
poses
controlnet = ControlNetModel.from_pretrained(
    "fusing/stable-diffusion-v1-5-controlnet-openpose", torch_dtype=torch.float16
)
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    model_id,
    controlnet=controlnet,
    torch_dtype=torch.float16,
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()
pipe.enable_xformers_memory_efficient_attention()

Обучение модели Объединение ControlNet со стабильной диффузией

Мы хотели посмотреть, как мы можем обучить модель ControlNet на наборе данных. Итак, для достижения нашей цели по управлению SD мы взяли на себя задачу закрасить кружки цветами. Наш подход будет состоять в том, чтобы использовать простую задачу, чтобы продемонстрировать, как ControlNet может изучать специфические для задачи условия сквозным способом. Мы будем использовать подсказку, которая описывает нашу цель и включает «Контрольное изображение (исходное изображение)».

Стабильная диффузия уже обучена на миллиардах изображений и уже знает, что такое цвета и формы. Однако он не понимает значения «Контрольного изображения (исходного изображения)» в нашей подсказке. Добавив ControlNet, мы можем научить SD понимать это условие ввода и закрашивать круг указанным цветом.

Набор данных

Набор данных fill50k содержит набор из 50 000 пар изображений кругов, которые хранятся в двух папках: «исходная» и «целевая». Папка «исходник» содержит изображения кругов только с линиями, а папка «цель» содержит изображения кругов, залитых определенным цветом.

В дополнение к изображениям набор данных включает файл «prompt.json», в котором содержится информация о каждом изображении. Каждая подсказка в виде описания круга и цвета фона. Например, в подсказке может быть написано «синий кружок на желтом фоне».

Чтобы использовать этот набор данных с PyTorch, вам нужно написать скрипт, который считывает изображения и соответствующие им подсказки из папки «fill50k» и подготавливает их к обучению. Это будет включать загрузку изображений в тензоры PyTorch, создание пакетов данных и выполнение любой необходимой предварительной обработки, такой как нормализация или увеличение данных.

В целом, набор данных fill50k предназначен для обучения моделей машинного обучения заполнению кругов определенными цветами на основе подсказок, представленных в наборе данных. Это полезный ресурс для исследователей и разработчиков, работающих над созданием изображений и связанными с ними задачами с использованием методов глубокого обучения.

Загрузка набора данных

import json
import cv2
import numpy as np

from torch.utils.data import Dataset


class MyDataset(Dataset):
    def __init__(self):
        self.data = []
        with open('./training/fill50k/prompt.json', 'rt') as f:
            for line in f:
                self.data.append(json.loads(line))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]

        source_filename = item['source']
        target_filename = item['target']
        prompt = item['prompt']

        source = cv2.imread('./training/fill50k/' + source_filename)
        target = cv2.imread('./training/fill50k/' + target_filename)

        # Do not forget that OpenCV read images in BGR order.
        source = cv2.cvtColor(source, cv2.COLOR_BGR2RGB)
        target = cv2.cvtColor(target, cv2.COLOR_BGR2RGB)

        # Normalize source images to [0, 1].
        source = source.astype(np.float32) / 255.0

        # Normalize target images to [-1, 1].
        target = (target.astype(np.float32) / 127.5) - 1.0

        return dict(jpg=target, txt=prompt, hint=source)

Какую модель SD использовать

Обратите внимание, что все веса внутри ControlNet также копируются из SD, поэтому ни один слой не обучается с нуля, и вы все еще выполняете точную настройку всей модели.

Обучение модели

Этот скрипт Python использует PyTorch Lightning для обучения пользовательской модели машинного обучения на пользовательском наборе данных. Он загружает предварительно обученную модель, устанавливает такие конфигурации, как размер пакета, частота регистратора, скорость обучения и многое другое. Затем он создает DataLoader для загрузки данных из пользовательского набора данных, создает регистратор изображений для записи хода обучения и настраивает PyTorch Lightning Trainer для управления процессом обучения. Наконец, он обучает модель с помощью Trainer и DataLoader и регистрирует прогресс с помощью ImageLogger.

import pytorch_lightning as pl
from torch.utils.data import DataLoader
from tutorial_dataset import MyDataset
from cldm.logger import ImageLogger
from cldm.model import create_model, load_state_dict


# Configs
resume_path = './models/control_sd15_ini.ckpt'
batch_size = 4
logger_freq = 300
learning_rate = 1e-5
sd_locked = True
only_mid_control = False


# First use cpu to load models. Pytorch Lightning will automatically move it to GPUs.
model = create_model('./models/cldm_v15.yaml').cpu()
model.load_state_dict(load_state_dict(resume_path, location='cpu'))
model.learning_rate = learning_rate
model.sd_locked = sd_locked
model.only_mid_control = only_mid_control


# Misc
dataset = MyDataset()
dataloader = DataLoader(dataset, num_workers=0, batch_size=batch_size, shuffle=True)
logger = ImageLogger(batch_frequency=logger_freq)
trainer = pl.Trainer(gpus=1, precision=32, callbacks=[logger])


# Train!
trainer.fit(model, dataloader)

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

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

Обучение

Подготовка данных и модели:

• Предварительно обученные модели могут быть полезны для самых разных задач и легко интегрируются в проекты.
• Доступны различные предварительно обученные модели для управления различными аспектами изображения, такими как края, глубина и позы человека.
• Библиотеки с открытым исходным кодом, такие как PyTorch, можно использовать для обучения и развертывания моделей машинного обучения.

Создание изображения:

• Стабильная диффузия — это мощный метод создания высококачественных изображений на основе управляющих входных данных.
• Использование диффузионной модели с управляющей сетью может обеспечить более точный контроль над сгенерированным изображением.
• Можно использовать несколько подсказок для создания разных изображений на основе одного и того же управляющего входа.

Обработка изображений:

• Методы обработки изображений, такие как обнаружение краев Канни и OpenPose, могут использоваться для извлечения конкретной информации из изображений, которая затем может использоваться в качестве управляющих входных данных для создания изображений.
• Различные типы методов обработки изображений подходят для различных типов входных данных управления, таких как края, позы и контуры.

Если вы заинтересованы, наша кодовая база доступна ниже:

Источники

Мы использовали следующие ресурсы в нашем анализе и базе кода для справки.