Kandinsky 2.0

Возможности

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

Архитектура Kandinsky 2.0 относится к классу диффузионных моделей с обуславливанием. Основная идея данного подхода заключается в постепенном восстановлении картинки из шума, путём движения в сторону максимизации правдоподобия на реальных картинках. Степень этого движения предсказывается нейронной сетью UNet на каждом обратном шаге диффузии. Обуславливание на текст реализован за счёт добавления слоёв перекрёстного внимания в архитектуру сети UNet. Отметим, что обучение UNet происходило не напрямую в пространстве изображений, а в латентном пространстве сети KL-AE.

Архитектура Kandinsky 2.0

Процесс обучения Kandinsky 2.0 состоял из трёх частей:

1. Предобучение на данных с разрешением картинок 256x256
2. Обучение на данных с разрешением картинок 512x512
3. Обучение на сильно отфильтрованных данных (100M) с разрешением 512x512

Для решения задач inpainting и outpainting обучались дополнительно со случайным маскированием областей изображений.

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

Модель содержит 2 млрд. параметров, поэтому её можно развернуть на обычной GPU. Высокой скорости инференса модели удалось также добиться за счёт того, что необходимые для генерации картинки текстовые эмбеддинги могут быть посчитаны отдельно на CPU.

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

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

Железный человек on the Moon 背景中的烟花

2. Inpainting
3. Outpainting

Зелёная планета в космосе

4. Картина по наброску

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

1. Text2Img

from natalle.natalle_model import Natalle

config_name = 'inference.yaml'
path2weights = 'natalle.pt'
device = 'cuda'

model = Natalle(config_name, path2weights, device, task_type='text2img')
images = model.generate_text2img('caption', batch_size=2, h=512, w=512)

2. Img2Img

from natalle.natalle_model import Natalle
from PIL import Image

config_name = 'inference.yaml'
path2weights = 'natalle.pt'
device = 'cuda'

model = Natalle(config_name, path2weights, device, task_type='img2img')
init_image = Image.open('image.jpg')
images = model.generate_img2img('caption', init_image, strength=0.6)

3. Inpainting and Outpainting

from natalle.natalle_model import Natalle
from PIL import Image
import numpy as np

config_name = 'inference.yaml'
path2weights = 'natalle_inpainting.pt'
device = 'cuda'

model = Natalle(config_name, path2weights, device, task_type='inpainting')
init_image = Image.open('image.jpg')
mask = np.ones((512, 512), dtype=np.float32)
mask[100:] =  0
images = model.generate_inpainting(
'caption', init_image, mask, guidance_scale=7, num_steps=50
)

Ccылки

GitHub