Kandinsky 2.1

Возможности

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

Архитектура Kandinsky 2.1 базируется на латентной диффузии с добавлением Image prior модели для повышения качества генераций. Image prior модель была разработана и обучена для маппинга текстов и изображений, в основе которой лежит трансформерная архитектура.

После обучения Image prior модели была обучена основная модель Text2Image, в архитектуре которой Image prior модель DiffusionMapping использовалась для маппинга входного текстового описания в визуальный эмбеддинг CLIP. Далее он применялся в обучении и инференсе диффузионной модели. Конечный блок — декодер, который из латентного представления позволяет получить финальное синтезированное изображение.

В отличие от предыдущих реализаций, в модели Kandinsky 2.1 используется новый декодер MoVQ вместо VQGAN, что позволяет значительно повысить качество генераций.

Среди нескольких возможных способов использования генеративной модели мы особое внимание уделено возможности смешивания изображений. Для этого в обученную диффузионную модель подается два визуальных эмбеддинга CLIP. И далее декодер восстанавливает “смешанное” изображение.

Архитектура Kandinsky 2.1: обучение Image prior (слева), генерация изображения по текстовому описанию (в центре), смешивание изображений (справа).

Преимущества

Модель содержит 3,3 млрд параметров, поэтому ее можно развернуть на обычной GPU.

Сценарии использования

1. Генерация картинки по тексту

Эйнштейн в космосе, окруженный логарифмической схемой

2. Смешивание изображений

3. Вариации изображения

Инструкции по использованию

1. Text-to-Image

from omegaconf import OmegaConf
from kandinsky2.kandinsky2_1_model import Kandinsky2_1

config = OmegaConf.load('inference.yaml')
unclip_weights = 'unclip.ckpt'
prior_weights = 'prior.ckpt'
device = 'cuda'

model = Kandinsky2_1(
	config,
	unclip_weights,
	prior_weights,
	device=device
)
images = model.generate_text2img(
	'woman with barocco flowers',
	batch_size=1,
	h=768, w=768
)

2. Image Fusion

from omegaconf import OmegaConf
from PIL import Image
from kandinsky2.kandinsky2_1_model import Kandinsky2_1

config = OmegaConf.load('inference.yaml')
unclip_weights = 'unclip.ckpt'
prior_weights = 'prior.ckpt'
device = 'cuda'

model = Kandinsky2_1(
	config,
	unclip_weights,
	prior_weights,
	device=device
)

path_to_image1 = 'image1.jpg'
path_to_image2 = 'image2.jpg'
images = model.mix_images(
	[Image.open(path_to_image1), Image.open(path_to_image2)],
	[0.5, 0.5],
	batch_size=1,
	h=768, w=768
)

3. Image Variations

from omegaconf import OmegaConf
from PIL import Image
from kandinsky2.kandinsky2_1_model import Kandinsky2_1

config = OmegaConf.load('inference.yaml')
unclip_weights = 'unclip.ckpt'
prior_weights = 'prior.ckpt'
device = 'cuda'

model = Kandinsky2_1(
	config,
	unclip_weights,
	prior_weights,
	device=device
)

path_to_image = 'image1.jpg'
image = Image.open(path_to_image)
images = model.mix_images(
	[image, image],
	[1.0, 0.0],
	batch_size=1,
	h=768, w=768
)

Ccылки

GitHub