Создаем воспоминания. Осваиваем FLUX, LoRA и ComfyUI

2025-01-03 в 15:02, admin, рубрики: computer vision, flux, машинное обучение, нейросети, Стругацкие

Разбираюсь на праздниках с дообучением моделей для генерации изображений. Было интересно, насколько сложно дообучить модель для генерации изображений по тексту FLUX в домашних условиях, сколько нужно обучающих данных и как затем генерировать качественные фотографии и иллюстрации.

Чтобы через время не забыть про особенности процесса и как-то его зафиксировать, решил поделиться наработками. Под катом подробности и еще немного фотографий АБССС.

Введение

Если интересно как модели работают, то информация на вход такая — есть нейросетевые модели для одного типа данных (текст, изображения, звук и т.д.), есть мультимодальные (в данном случае текст-картинка).

Такие модели можно обучать с нуля и это дорого, нужен кластер с GPU (видеокарты) и много данных. После обучения получится некоторая фундаментальная модель, которая может делать много обычных задач. В домене текст-картинка бывают открытые модели, типа Stable Diffusion, Kandinsky и FLUX, бывают закрытые, типа DALL-E.

Открытую модель можно дообучать разными способами. Один из самых экономных — дообучение методом LoRA (Low-Rank Adaptation), так обучается только малая часть параметров, эту часть можно хранить отдельно, а при запуске модели (это еще называется инференсом) LoRA «приклеивается» к основной модели. Данных тут тоже понадобится мало. С них и начнем.

Данные

Объем: 20-25 фотографий с лицом.

Самый простой способ начать эксперименты — это сделать датасет из своих фото. Просто скиньте их себе, например, в телеграм или по почте и откройте на компьютере.

На компьютере просматривайте их и делайте околоквадратные скриншоты любой удобной утилитой типа Greenshot или встроенной в Windows. По-моему, такая штука ещё идет Яндекс браузером. Жмёте PrtSc, выделяете область, готово.

Так делаете 20-30 фотографий, просматриваете, удаляете слишком отличающиеся или некачественные. Откладываете в отдельную папку. Датасет готов.

Особенности: Для личностей типа Стругацких или Бродского, качественных фотографий крайне мало, но много и не надо. То, что найдете, по возможности оцветните и улушите качество (уберите шум и т.д.). Чем лучше данные на входе, тем стабильней и качественней будет выход.

Обучаем модель

Спортивный подход: обучаем локально при наличии достаточно мощной видеокарты (16-24Gb памяти).

Время работы: моя 3090 пыхтит и делает «бррр» в течение 40-60 минут (зависимости от каличества фото и эпох — циклов обучения по датасету).

Для локального обучения используем Flux Gym. Этот репозиторий является оберткой для запуска скриптов из друго проекта — Kohya Scripts.

Процесс установки очень простой (описан в readme) — копируем репозиторий, устанавливаем зависимости, готово. После запуска открываем страничку в браузере и видим интерфейс.

Пишем название для будущей модели, задаем триггер — то, на что модель должна будет среагировать, чтобы начать генерировать картинки на основе обучающих данных. Это должно быть что-то необычное, чтобы не накладываться на уже знакомые модели понятия. Например b0r1s_strug. Можно и фразу.

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

При выборе есть три модели, выбирайте flux-dev или bdsqlsz/flux1-dev2pro-single (это тот же flux-dev, но дообученный на качественных картинках, как заявляют авторы, огромных различий на паре запусках я не заметил). В первый раз модели выкачаются из интернета и сохранятся на диске, вес порядка 23 Гб.

Выберите ограничение по памяти, остальные параметры пока оставьте по-умолчанию и меняйте в зависимости от времени обучения и результата.

Загрузите ваши фото. Везде автоматически проставится триггер. Можно оставить так и ничего больше не делать. Если вы загружаете какие-то концептуальные фото с различными сценами, то лучше добавить описания. Это можно сделать автоматически, нажав на кнопку "Add captions" и они автоматически проставятся скачанной моделью, подредактируйте их.

Справа будет виден скрипт, который будет запускаться. Все, можно начать обучение.

Особенности: О чем авторы Flyx Gym не пишут, но что можно улучшить. Объем полученной LoRA модели будет около 35Мб. Это мало, что в целом нормально. Но я заметил, что если обучать LoRA онлайн платными сервисами (об этом дальше), то при экспорте она весит около 130Мб.

Дело тут в том, что LoRA можно конфигурировать. По сути, это матрицы определенного размера, которые добавляются к различным слоям исходной модели. Если почитать скрипты, то становится понятно, что этот размер можно задавать через дополнительные параметры. Либо через network_dim, либо более тонко. Методом научного тыка я подобрал так:

--network_args "img_attn_dim=8" "img_mlp_dim=16" "txt_attn_dim=8" "txt_mlp_dim=8" "img_mod_dim=8" "txt_mod_dim=8" "single_dim=16" "single_mod_dim=8"

В итоге LoRA занимает около 110Мб и работает получше. Больших исследований я не проводил, так все это очень времязатратно. Здесь вам задел для будущих оптимизаций.

Добавляем эту строку к скрипту запуска и запускаем через консоль. Ах да, если указать в интерфейсе опцию генерации семплов по мере обучения, то можно будет смотреть как модель обучается. Но по наблюдениям, эти фото (даже последние) гораздо хуже того, что будет при запуске кончной модели.

После некоторого времени в папке outputs/<name> видим нашу модель и промежуточные чекпоинты в формате .safetensors. Это то, что нам надо.

Если нет видеокарты: попробуйте поискать сервера с видеокартой в аренду (типа vast.ai), настроив все скрипты один раз и арендовав что-нибудь типа A100, можно уложиться минут в 15 обучения, получится не очень дорого. Есть вариант ещё проще.

Неспортивный подход: обучаем онлайн. Стоит около 5-7 долларов, модель можно скачать и запускать локально. Самый простой метод, качество получается хорошее, хотя с настройками уже не поиграешься.

Заходим на fal.ai и регистрируемся, раньше там давали 5 долларов бесплатно, как раз хватило бы на один запуск, сейчас, по-моему, такого нет. Пополняем баланс и идем в flux-lora-portrait-trainer.

Загружаем ваши фото, жмем на запуск, ждем минут 5-7, готово. Последующий инференс стоит около 3 рублей за генерацию. Все работает быстро и удобно.

Модель можно экспортировать на Hugging Face и скачать. Весит, как было сказано, около 130Мб.

Запускаем модель

Пожалуй, самая интересная часть. Для запуска моделей будем использовать ComfyUI. Я долгое время его избегал, так как на скриншотах он всегда выглядит круто, но не очень интуитивно:

Зато примеров его использования пруд пруди, в данном случае это и плюс и минус ComfyUI. Избегал я его зря, немного покопавшись в нём, инструмент оказался максимально гибким и мощным.

Для чего: если у вас есть пул картинок и какая-то многоступенчатая задача, например, вы хотите оцветнять изображения, находить на каждом из них лицо, вырезать его и увеличивать при помощи нейросетей, то нужно один раз сделать такой пайплайн (а таких примеров наверняка уже куча), то это тот инструмент, который нужен.

Установка ComfyUI: проще всего установить портативную версию, скачав её из репозитория. Собственно, надо только распаковать архив и запустить. Интерфейс будет доступен в браузере.

Вся работа идет визуально, через блоки. Каждый блок умеет делать одну вещь: загружать модель, сохранять картинку, сравнивать что-то и т.д. Есть огромное количество кастомных блоков, которые надо устанавливать отдельно. Блоки объединяются в workflow.

Сразу установите расширение ComfyUI-Manager, оно позволит вам на лету устанавливать недостающие модули.

В папке с Comfy есть другие папочки для моделей, входных и выходных картинок. Подложим наши модели туда. Все необходимые модели уже автоматически скачались при запуске fluxgym. Нам надо их скопировать.

fluxgym/models/unet/flux1-dev.sft → comfyui/models/unet
fluxgym/models/vae/ae.sft → comfyui/models/vae
fluxgym/models/clip/clip_l.safetensors → comfyui/models/clip
fluxgym/models/clip/t5xxl_fp16.safetensors → comfyui/models/clip
fluxgym/outputs/<lora_name>/<lora_name>.safetensors → comfyui/models/loras

Workflow: необходимый пайплайн я собрал. В нем есть загрузка исходной модели, применение вашей LoRA, генерация картинки по промпту и масштабирование результата другйо нейросетью. Берем workflow отсюда. Это json файл, сохраните его в папке с ComfyUI.

Открывайте этот пайплайн через интерфейс программы (Workflow → Open). Убедитесь, что все скопированные модели выбраны:

После этого напишите промпт, используя выбранное на этапе обучения триггерное слово. Нажмите Queue.

Советы: одним из ключевых параметров тут является cfg (Classifier-Free Guidance) в компоненте KSampler, который производит семплирование («проявляет» картинку из начального шума в сторону вектора, несущего данные о вашей текстовой подсказке). Чем он выше, тем больше модель добавляет деталей «от себя» жертвуя при этом фотореалистичностью. Покрутите его от 0.8 в обе стороны и посмотрите на результат.

Также влияет доля, с которой LoRA модель добавляется к основной модели, вид начального семплирования и ваш промпт. Тут у вас простор для экспериментов. Также мы при обучении не меняли начальные параметры — число повторений каждой картинки и число эпох обучения. Это тоже одна из осей для экспериментов.

Особенности: в зависимости от качества картинок модель может быть более или менее стабильной. Чем хуже данные, тем больше придется делать генераций, чтобы получить что-то стоящее. Впрочем даже по немногочисленным изображениям Стругацких удалось сделать отличные фото, лучше которых в сети не найти.

Также модель может иногда менять пол (!) человека. Поэтому стоит добавлять к промпту man или woman. Также улучшают качество стандартные дописки типа «professional shot», «high resolution» и т.д.

Напоследок

Персонажи в данной статье выбранны неспроста — братьев Стругацких я просто люблю читать, да и вчера, 2 января, был День научной фантастики (приуроченный к дню рождения Азимова).

Владимира Сурдина же и Алексея Семихатова вы наверняка знаете как замечательных популяризаторов науки (надеюсь, они простят мне использование их фотографий, если им попадется эта статья).