🔥Consistency Decoder

Среди недавних релизов OpenAI больше всего в мое сердечко запал их новый декодер для LDM - Consistency Decoder.

Напомню, что Stable Diffusion генерирует изображения в Latent пространстве, которое затем декодируется VQ-VAE декодером (на картинке выше, кажется, его назвали GAN Decoder). Такой декодер довольно легковесен, но он не может точно восстанавливать мелки детали типа лиц людей на заднем плане. Поэтому парни из OpenAI решили натренировать еще одну диффузию вместо VAE декодера.

Но диффузия работает медленно и требует много шагов во время инференса... Тут на помощь пришла дистилляция из диффузии в Consistency Model, которая может неплохо работать за 1-2 шага, если руки не кривые.

В статье про Dalle-3 было вскользь упомянуто, что в новой архитектуре для перевода скрытого кода в RGB они как раз используют такой Consistency Decoder, который работает за 2 шага.

Consistency Decoder - это тоже Unet, и довольно большой, 620 M параметров (для сравнения SD 1.5 ~900M параметров). Но это все равно допустимая жертва ради улучшенного качества восстановления лиц, мелких деталей и регулярных линий.

Жаль, только, что кода тренировки нет, и остается только заниматься реверс-инженирингом и экспериментировать, чтобы понять как они натренировали этот декодер. Выложили только инференс и веса декодера, который совместим с базовым Stable Diffusion 1.x - 2x и может использоваться как drop-in replacement вместо стандартного декодера.

➡️ Код и веса

@ai_newz