Концепция
Моей задачей было обучить нейросеть генерации изображений космоса. Тема была выбрана мной из-за сильной тяги к звездам и всему внеземному. Изображения хотелось видеть яркие, не лишенные фантастичности. Такими, каким я представляю себе дальний космос.
Исходные изображения
Для обучения модели я собирала фотографии с открытой лицензией и находящиеся в открытом доступе, которые собрала в один архив. Все фотографии соответствуют теме космоса, звезд и планет, тому, что я хотела бы видеть на получившихся изображениях.
Исходные изображения. Пример.
Исходные изображения. Пример.
В рамках проекта я реализовала полный цикл обучения адаптера LoRA (Low-Rank Adaptation) для модели Stable Diffusion v1.5, используя 20 собранных изображений.
На этом этапе я установила необходимые библиотеки. bitsandbytes позволил мне использовать 8-битную оптимизацию памяти, что критически важно при работе с ограниченными ресурсами GPU. Библиотеки accelerate и peft обеспечили эффективное обучение с использованием метода LoRA. Актуальную версию diffusers я установила напрямую из репозитория Hugging Face, чтобы получить все последние обновления.
Я создала рабочую директорию /content/Universe, где будут храниться все обучающие данные. Используя функцию files.upload () из Google Colab, я загрузила 20 подготовленных изображений космической тематики. Для удобства я временно переключилась в созданную папку, чтобы файлы сохранились по правильному пути, после чего вернулась в корневую директорию.
Я реализовала автоматический поиск всех изображений в папке, поддерживая форматы JPEG, PNG и WEBP. Все найденные файлы я загрузила в память с помощью библиотеки PIL. Для визуального контроля качества датасета я написала функцию image_grid, которая позволяет отобразить набор изображений в виде сетки — это помогло мне убедиться, что все файлы загружены корректно и имеют подходящее качество.
Для автоматической генерации описаний я использовала предобученную модель BLIP (Bootstrapping Language-Image Pre-training). Я загрузила модель и процессор из репозитория Salesforce, переведя модель в формат float16 для экономии видеопамяти. Функция caption_image принимает изображение и возвращает текстовое описание. Чтобы усилить стилистическую направленность на космическую тематику, я добавила к каждому описанию пользовательский префикс.
Я создала файл metadata.jsonl в формате JSON Lines, где каждая строка содержит пару «имя_файла»: «текстовое_описание». Для каждого изображения я сгенерировала описание с помощью BLIP, объединила его с заданным префиксом и записала в файл. Этот формат является стандартным для обучения моделей в экосистеме Hugging Face.
Перед началом обучения я явно указала использование GPU с индексом 0, чтобы избежать конфликтов при наличии нескольких устройств. Я также добавила диагностический вывод, который позволил мне убедиться, что GPU доступен и имеет достаточный объем памяти для обучения. Это важный шаг, так как обучение моделей диффузии требует значительных вычислительных ресурсов.
Для демонстрации полученных результатов я собрала первые пять сгенерированных изображений и отобразила их в виде сетки с помощью библиотеки matplotlib. Это позволило мне визуально оценить качество дообучения модели и убедиться, что LoRA адаптер успешно усвоил стилистические особенности обучающего датасета.
Результаты обучения
Четыре наиболее точных изображения, передающих те ощущения, которые мне хотелось вложить в работу.
Заключение
В ходе работы я реализовала полный цикл адаптации модели Stable Diffusion под кастомный стиль с использованием LoRA. Ключевыми результатами стали:
Успешное обучение на ограниченном датасете из 20 изображений.
Эффективное использование ресурсов GPU благодаря LoRA, 8-битной оптимизации и mixed precision.
Автоматизация создания аннотаций через интеграцию с BLIP.
Получение качественных изображений космической тематики с использованием стилевого токена.
Данный подход может быть масштабирован на большее количество изображений и адаптирован для различных стилей и предметных областей.
Для помощи написания промптов использовался искусствнный интеллект.
