Как обучить нейросеть рисовать в стиле любого художника
Василиса Премудрая лабает на модуляре под над мухоморами. Stable Иван Билибин, 2023
Прошло 3 года с момента когда я обучал StyleGAN на панельках и мне стало интересно что там сейчас с генерацией картинок. А там — ого — можно дообучить целый stable diffusion на любом стиле любого художника! Как? А вот щас расскажу
Подготавливаем данные
Для успешного файнтюнинга нам понадобится (всего!) около 10–20 картинок. Да-да, если раньше надо было собирать датасеты из десятков тысяч картинок то сейчас можно отделаться просто десятком. Умные люди из гугла придумали для этого метод DreamBooth, а другие умные люди из майкрософта — метод LoRa, чтобы уменьшить размерность обучаемых весов и ускорить файнтюнинг. В сумме эти 2 метода позволяют за ~10 минут на хорошей GPU дообучить stable diffusion на стиле художника по совсем небольшому датасету.
Обучаем модель
Можно обучать модель локально, если позволяют ресурсы, как это сделать описано например вот здесь, в официальной документации библиотеки diffusers. А для тех у кого под рукой нет мощной GPU или просто лень заморачиваться с настройкой окружения и скриптов есть например Replicate. Я решил воспользоваться этим сервисом и заплатил около … 0.5$ за целый файнтюнинг. Подробная инструкция по файнтюнингу на replicate есть вот здесь, я лишь акцентирую внимание на некоторых нюансах:
Нюанс первый. По-умолчанию пайплайн использует BLIP для генерации подписей к картинкам. Это хорошо работает на лицах, собаках, машинах, в общем на каких-то понятных и предсказуемых образах. А вот если хочется обучить модель на необычном стиле, да ещё и так чтобы она правильно реагировала на определённые слова то этот вариант не подойдёт. Пример — работы художника Всеволода Иванова, в которых любой кто хоть раз смотрел Рен-ТВ увидит гиперборею и атлантиду, но для BLIP-а это будут просто рисунки людей и мамонтов. В общем, чтобы слова в промптах имели смысл лучше сделать все подписи руками, и положить получившуюся csv-шку в архив вместе с картинками
Нюанс второй. Если вам непременно нужны какие-то специфичные образы в генерациях, не стесняйтесь добавлять в датасет картинки реальных объектов, пусть выбивающиеся из общего стиля. Подпись желательно тоже делать с одной стороны уникальной, с другой стороны близкой по смыслу к тому что хочется получить. Пример — если добавить в датасет картинку подмосковной электрички с подписью electric train то процент удачных генераций электрички в стиле художника получится меньше чем если при обучении сказать что это sobaka electric train
Нюанс третий. В ноутбуке из инструкции конфиг написан для файнтюнинга на лицах, чтобы учить стиль надо скопировать правильный конфиг из раздела «Fine-tuning a style» инструкции. Если используются подписи из csv, то можно оставить caption_prefix пустым и учить сетку сразу на нескольких художниках, чтобы например потом смешивать их стили
Запускаем генерацию локально
После того как модель обучена, её можно скачать и запускать генерацию примерно на любом кирпиче — главное чтобы памяти хватило. Для replicate способ скачать модель оказался не совсем очевидным и описан он тут в предпоследнем пункте инструкции. Надеюсь, скачать модель через их API с сайта всё же можно, просто я не понял как
Бекенд replicate использует репозиторий cog-sdxl для генерации, так что помимо pip install torch и diffusers надо будет скачать и положить в локальную папку репо cog-sdxl
Также важно обратить внимание на параметр lora_scale — это как бы сила с которой полученные веса LoRa применяются к основной модели
import torch
from diffusers import DiffusionPipeline
from cog_sdxl.dataset_and_utils import TokenEmbeddingsHandler
pipe = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
).to("cuda")
model_dir = "папка с моделью"
pipe.load_lora_weights(f"{model_dir}", weight_name="lora.safetensors")
pipe.fuse_lora(lora_scale=0.6)
text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
tokenizers = [pipe.tokenizer, pipe.tokenizer_2]
embhandler = TokenEmbeddingsHandler(text_encoders, tokenizers)
embhandler.load_embeddings(f"{model_dir}/embeddings.pti")
images = pipe(
'пишем промпт',
# negative_prompt='если надо пишем negative промпт',
cross_attention_kwargs={"scale": 0.8},
).images
images[0].save(f"картинка.jpg")
Результаты
русы против ящеров
Первым делом попробовал на том самом Всеволоде Иванове, иллюстрировавшем Рен-ТВ нулевых. До того что надо самому писать подписи к картинкам тогда ещё не додумался, поэтому промпты подбирались методом тыка и пирамиды Ведической Гипербореи оказалась на самом деле градирнями
ин стайл оф Фоменко э пэйнтинг оф геометрик абстракшнс
Следующим был Анатолий Фоменко, мультфильм «перевал» которого я однажды в детстве увидел и так и не смог развидеть. Подписать картины руками было сложно, не всегда понятно как вообще описать то что на них изображено, но сетка отлично выучила органическо-геометрические абстракции, суровые лица советских мужчин и цветовую гамму
Russian Rave in Forest 2006 (MEDNOE OZERO)
Ну и пока лучшее что получилось сделать — русский рейв в лесу в стиле Билибина. Так получилось что я неравнодушен к музыкально-синтезаторной тематике и, обучив сетку, решил попробовать сгенерировать Василису с синтами. А сетка внезапно взяла и не забыла что это и как оно выглядит!
Редактируем результаты
Примеры осознанного редактирования
Как обычно не обошлось без нашего музыкального проекта — захотелось сгенерить арты к альбому. Задача тут была немного сложнее чем просто что-то нарисовать — надо было получить осмысленные образы в определённом стиле (за основу взяли работы художницы электри4ка), да ещё и согласующиеся с атмосферой треков. И вот хочешь ты огромного кота на электричке, а сетка рисует усреднённый паровоз. Тут на помощь приходит инпентинг и масочное редактирование. Делаем маску, белое перерисовываем, чёрное оставляем, сужаем промпт до интересующего объекта, и — ура — кот едет на той самой электричке. Ну и лапки тоже пофиксить можно. Сделать это можно либо в веб-интерфейсе replicate либо локально, с помощью аргумента mask_image в методе pipe
Так себе результаты
На самом деле их было гораздо больше, в основном это привычные 10 пальцев и 5 хвостов
Слева направо: ЪУЪ, Salvatore Ganacci, котёнок который обязательно выживет
Секунда самопиара
Если вам вдруг понравились такие приколы, то заходите в телегу, их у меня ещё будет