Лягушки в смешных шляпках обучают Stable Diffusion

С детства я люблю лягушек и окружаю себя ими, будь то игрушки, заставка на телефон или домашний питомец.

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

Для обучения нейросети я собрала датасет из 60-и фотографий лягушек и жаб в разнообразных шляпках: от ковбойских шляп до цветочков на голове.

Источник изображений: pinterest.com

Итак, датасет создан, пора приступать к обучению!

Работа с кодом велась в среде разработки Google Colab.

Сначала я импортировала библиотеки, нужные для обучения Stable Diffusion.

Теперь нужно импортировать подготовленный датасет, чтобы Google Collab смог «видеть» наших прекрасных лягушат. Для этого я указала пусть на гугл диск с собранными ранее изобрабражениями.

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

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

Итоговый префикс: «a photo of a cute TOK frog in a funny hat». Дальше нейросеть сама описывает, что видит на картинке.

Теперь у нас есть текстовые описания изображений! Они сохраняются в папку с датасетом в формате .jsonl.

Следующим этапом заходим в аккаунт на Hugging Face и создаём токен. Так у системы Google Colab будет доступ к аккаунту, чтобы сохранить туда будущую модель обучения.

Наконец мы приступаем к самому обучению!

Размер для изображения был задан в 512×512 пикселей, с шагом в 500 и чек-поинтом на 250.

Обучение заняло чуть меньше 38 минут.

Обучение завершено! Теперь нужно сохранить модель обучения на Hugging Face.

Можно начинать генерировать лягушек в смешных шляпках! Я делала это уже на странице модели на Hugging Face.