Як працюють моделі для генерації зображень. Stable Diffusion, Midjourney та інші

💡 Усі статті, обговорення, новини про AI — в одному місці. Приєднуйтесь до AI спільноти!

Привіт, спільното! Мене звати Паша, я понад шість років працюю у сфері Data Science. Наразі я частина команди Universe Group — продуктової IT-компанії, що спеціалізується на розробці мобільних і вебзастосунків з використанням штучного інтелекту.

За цей час я встиг попрацювати з найрізноманітнішими проєктами: від класичних регресійних моделей для банківського сектору до повноцінних систем у сферах рекомендацій, обробки природної мови (NLP, LLM) та комп’ютерного зору (CV).

У цьому матеріалі я детально розберу, як працює модель для генерації зображень Stable Diffusion: які ідеї лежать в основі її архітектури, як вона навчається та чому генерує саме такі результати. Ділюсь практичним досвідом — це буде корисно Data Science-інженерам, технічним командам, а також CEO, CTO та всім, хто вже використовує або планує інтегрувати цей інструмент у свої продукти.

Архітектура

Всі сучасні моделі для генерації зображень працюють на одному принципі — принципі дифузії, саме тому їх часто називаються дифузними моделями. Сьогодні ми розберемо архітектуру Stable Diffusion — основну нейромережу для генерації зображень зараз. Але за однаковим принципом (з мінімальними змінами) працюють всі дифузні моделі, про які ви чули: Midjourney, DALL·E 2, etc.

Архітектура моделі Stable Diffusion складається з 3 різних нейромереж:

• Autoencoder (VAE).
• Text-encoder, e.g. CLIP’s Text Encoder.
• U-Net.

Autoencoder (VAE)

Розберемо для чого потрібний варіаційний автоенкодер (VAE). Автоенкодер складається з двох основних компонентів: енкодера та декодера. Задача енкодера закодувати зображення, а задача декодера розкодувати назад. Ось і вся суть🙂

Архітектура Autoencoder

Autoencoder навчають наступним чином: беруть велику кількість доступних зображень і подають їх на вхід цій моделі, при цьому на виході модель намагається відтворити ці самі зображення. Тобто усе, що робить модель, — це отримує зображення і намагається його ідеально відтворити. Проте є важливий нюанс: на вході та виході модель працює із зображеннями розміром 512×512×3, але в середині архітектури створюється спеціальне «вузьке місце» (bottleneck), наприклад, розміром 64×64×4.

Через це нейромережа не може передавати значення пікселів безпосередньо, оскільки в неї просто немає достатньо нейронів для такого обсягу даних. В результаті модель навчається поміщати в одне значення нейрона кілька важливих характеристик зображення, таких як колір, об’єкти чи персонажі на зображенні, стиль зображення та інші параметри. Таким чином вона передає не окремі значення пікселів, а більш високорівневі абстракції. Ці стислі представлення зображень називаються латентним представленням, тобто «прихованим».

Це виглядає приблизно ось так:

На вхід подається зображення собачки, а на виході нейромережа намагається його відтворити максимально схоже з початковим зображенням. Ми бачимо, що зображення неідеально відповідає вхідному, але в цьому і суть навчання: енкодер вчиться максимально якісно кодувати інформацію в латентні ознаки, а декодер вчиться максимально якісно відновлювати зображення з цього латентного простору.

Ще один спосіб пояснити принцип роботи автоенкодера — через аналогію. Уявіть, що вам показали зображення собаки розміром 512×512×3 і попросили передати це зображення іншій людині. Найточніший спосіб — це скопіювати значення кожного пікселя й передати їх разом із координатами, щоб отримувач міг повністю відтворити зображення. Але тепер уявіть, що у вас є можливість передати лише 10 характеристик. У такому випадку ви намагатиметесь у ці 10 характеристик вмістити якомога більше суттєвої інформації: не окремі пікселі, а більш узагальнені ознаки — наприклад, «фотографія», «собака», «чорно-біла», «сидить», «дивиться в камеру», «на фоні плитка», «порода: спанієль» тощо.

Це і є суть вузького місця (bottleneck) в автокодері: змусити модель передавати не повні піксельні дані, а стислі, змістовні характеристики зображення. У нейромережі ці характеристики не виражаються словами, як у нашій аналогії, а кодуються у вигляді числових векторів. І саме це перетворення — це і задача моделі Autoencoder.

Що довше нейромережа навчається, то краще вона відтворює початкове зображення. Але головна мета не в цьому. Нам важливо, щоб модель навчилася максимально точно закодовувати зображення у свій латентний простір (це робить частина encoder) і так само точно відновлювати його з цього простору (а це decoder). Це критично важливо, щоб ми могли працювати з величезною кількістю зображень і складними архітектурами в майбутньому, оскільки перетворення зображення з форми 512×512×3 до латенту з формою 64×64×4 зменшує обсяг даних приблизно у 48 разів. Це суттєво знижує вимоги до оперативної пам’яті та обчислювальних ресурсів порівняно з генерацією зображень у піксельному просторі. Завдяки цьому можна ефективно генерувати зображення розміром 512×512 навіть на споживчих GPU, таких як 16GB у Google Colab (colab.research.google.com).

Отже, підсумуємо: перший ключовий компонент — варіаційний автоенкодер (VAE). Він використовується для перетворення зображення у компактне латентне представлення. Це дозволяє моделі працювати не з повнорозмірними піксельними даними, а з їх стислою, змістовною версією. Такий підхід зменшує обсяг оброблюваних даних, а відповідно і потребу в обчислювальних ресурсах, приблизно у 48 разів.

Text-encoder, e.g. CLIP’s Text Encoder

Текстові енкодери відповідають за обробку текстових запитів і перетворення їх у числові вектори, які слугують основою для генерації зображень. Коли користувач вводить текстовий запит — наприклад, «чарівний пейзаж на заході сонця» — це висловлювання природною мовою необхідно трансформувати у числовий формат, який може обробити нейронна мережа.

Для цього в Stable Diffusion використовується модель кодування тексту — CLIP (Contrastive Language-Image Pre-Training). Модель CLIP розроблена в OpenAI та призначена для одночасного розуміння як тексту, так і зображень. Модель формує спільний простір ознак (embedding space) для обох модальностей, що дозволяє порівнювати текстові та візуальні вектори на основі їхньої семантичної схожості. Під час генерації зображення за текстовим запитом CLIP гарантує, що текстове представлення максимально відповідає відповідним візуальним характеристикам. Для цього текст і зображення обробляються окремими нейронними мережами, виходи яких об’єднуються у спільному латентному просторі — саме це забезпечує точну й узгоджену генерацію.

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

Навчання CLIP (Contrastive Language-Image Pre-Training) є процесом, спрямованим на створення спільного латентного простору для текстових і візуальних даних, що дозволяє моделі одночасно розуміти зображення та текст. Ось як відбувається навчання CLIP.

Дані для навчання:

CLIP навчається на великій кількості пар текст-зображення. Наприклад, одна пара може бути: зображення кота і текстовий опис «чорний кіт, що сидить на дереві». Це дозволяє моделі асоціювати текстові описи з конкретними візуальними характеристиками зображень.

Роздільне кодування тексту і зображення:

— Текстовий енкодер: для тексту використовується модель трансформера, яка перетворює текстовий опис у векторне представлення (текстовий вектор). Цей вектор містить інформацію про семантику тексту.

— Візуальний енкодер: зображення обробляються через архітектуру, подібну до CNN (Convolutional Neural Network), наприклад, ResNet або Vision Transformer. Вона перетворює зображення на вектор, який містить його основні ознаки.

Контрастивне навчання:

Модель використовує контрастивне навчання (contrastive learning), що означає, що її мета — мінімізувати відстань між векторами зображення та тексту, які належать до однієї пари (позитивні пари), і максимізувати відстань між векторами зображень і текстів, які не пов’язані (негативні пари). Іншими словами, CLIP вчиться, щоб правильний текст і відповідне зображення були ближчими в латентному просторі, а неправильні пари — віддаленими.

Використання позитивних і негативних прикладів:

Під час навчання для кожного текстового запиту модель генерує векторне представлення зображень і перевіряє його подібність з іншими зображеннями, а також з іншими текстами. Якщо текст і зображення правильні, їхні вектори будуть дуже схожими, а якщо пара неправильна, модель має розрізняти їх.

Отже, після навчання модель CLIP може генерувати вектори як для зображень, так і для текстів, що дозволяє знаходити візуально схожі зображення за текстовим запитом або навпаки. Цей спільний латентний простір робить CLIP потужним інструментом для задач, де важлива взаємодія між зображенням і текстом, наприклад, пошук за зображеннями, класифікація та генерація зображень з тексту.

У моделях для генерації зображень текстовий енкодер має завдання перетворити запит користувача (prompt) на числовий вектор (embedding), який модель може інтерпретувати, щоб зрозуміти, яке зображення потрібно згенерувати.

U-net

Третій важливий компонент архітектури Stable Diffusion — це U-Net. U-Net — це згорткова нейронна мережа, але вона також містить частини енкодера та декодера. Основна ідея якої полягає в тому, щоб в процесі генерації зображення модель могла зберігати важливу інформацію про структуру зображення, а також здатність поступово уточнювати деталі на всіх етапах генерації.

Цю модель навчають таким чином, щоб вона могла поступово відновлювати зображення з повного шуму. Для тренування використовують набір фотографій, до яких поетапно додається випадковий шум (аж до повної нерозпізнаваності).

Завдання моделі полягає в тому, щоб поетапно зменшувати рівень шуму на зображенні. Тобто на кожному кроці U-net намагається зменшити шум і наблизити зображення до його оригінального вигляду. Навчаючись на зображеннях з різними рівнями шуму, модель здатна, отримавши повністю зашумлене зображення, відновити його без шуму.

У контексті дифузійних моделей U-Net використовується для відновлення зображень, зменшуючи шум на кожному кроці дифузії. Під час навчання мережа вчиться передбачати конкретний шаблон шуму, який має бути доданий на кожному етапі для отримання найбільш реалістичних результатів. Цей процес дозволяє моделі генерувати нові зображення, починаючи з чистого випадкового шуму та поступово його уточнюючи до отримання бажаного зображення.

Завдяки своїй здатності точно відновлювати зображення та передбачати деталі на кожному етапі, U-Net є основним компонентом у генеративних моделях. У дифузійних моделях, таких як Stable Diffusion, вона відіграє важливу роль під час етапу видалення шуму, поступово «очищаючи» шум для створення чіткого, реалістичного зображення, яке точно відповідає наданому текстовому опису. Її гнучкість і висока точність роблять U-Net однією з найбільш ефективних архітектур для таких завдань.

Об’єднуємо все разом

А тепер, використовуючи усі 3 нейромережі — зберемо модель Stable Diffusion.

Процес генерації зображення починається з того, що користувач вводить текстовий запит — опис того, що саме він хоче побачити. Цей текст потрапляє до моделі Text Encoder, яка перетворює його у числовий вектор (ембеддинг), зрозумілий для нейромережі.

Після цього генерується абсолютно випадковий шум у латентному просторі розміром 64×64. З цим шумом модель працює так, ніби це дуже зашумлене зображення, і починає поетапно його «очищати» за допомогою нейромережі U-Net. Але оскільки ми не маємо жодного початкового зображення, лише шум, то моделі необхідно пояснити, що саме ми хочемо отримати. І тут вступає в дію ембеддинг з Text Encoder — на кожному кроці дифузії він подається разом із шумом до U-Net, допомагаючи моделі орієнтуватися на вказаний запит.

Таким чином, на кожному етапі зменшення шуму модель не просто «очищає» зображення, а робить це з урахуванням змісту промпта. Коли U-Net завершує N кроків дифузії, отримане латентне зображення передається декодеру з варіаційного автоенкодера (VAE), який перетворює його у фінальне зображення в піксельному просторі — те, яке ми вже можемо побачити й оцінити 🙂

Ініціалізація випадкового шуму — ключовий етап у дифузійних моделях. Саме завдяки цьому, навіть при однаковому текстовому запиті, модель може щоразу генерувати абсолютно різні зображення. Це досягається тим, що латентний простір розміром 64×64 кожного разу заповнюється новим випадковим шумом. У результаті маємо астрономічну кількість варіацій — понад 18 квінтильйонів (2⁶⁴ = 18446744073709600000) можливих зображень для одного й того ж промпту. Саме така стохастичність робить ці моделі настільки гнучкими, креативними та захопливими для використання!

Ось як працюють сучасні дифузійні моделі генерації зображень

Ми розібрали три ключові компоненти моделі Stable Diffusion: варіаційний автоенкодер (VAE), текстовий енкодер та архітектуру U-Net. Кожен з них виконує свою унікальну функцію, а разом вони перетворюють звичайний текстовий опис у реалістичне, деталізоване зображення.

Ця архітектура не лише забезпечує високу якість генерації, а й оптимізована для роботи з меншими обчислювальними ресурсами. Завдяки цьому створювати зображення стало можливо навіть на звичайних споживчих відеокартах. Саме така ефективність і гнучкість зробили Stable Diffusion та інші дифузійні моделі настільки популярними серед широкого кола користувачів.

Сподіваюся, цей матеріал допоміг вам краще зрозуміти, як працюють сучасні моделі генерації зображень. Якщо маєте запитання чи хочете поділитися думками — пишіть у коментарях.