В общем, кому не хватало Midjourney у нас дома - добро пожаловать в @pixelmuse_bot.

Что бот умеет на данный момент:

  • Создавать картинки по текстовому запросу
  • Используется автоматический машинный перевод текста на английский, поэтому запрос может быть практически на любом языке
  • Можно выбрать из 5 моделей: реализм, аниме, киборги, sci-fi окружение, стикеры (по факту под капотом две модели + разные LoRa)
  • Можно загрузить свою фотографию для обработки img2img
  • Можно сделать апскейл результата x1.5. Для апскейла используется нейронка Real-ESRGAN
  • Даю 20 токенов для тестов (20 изображений) с лимитом 10 токенов в день

Dev-лог

Какое-то время ковырял Stable Diffusion в Автоматик111 и ComfyUI. В один прекрасный депрессивный день решил посмотреть, как работать с сетью напрямую в python. Сделать это можно с помощью библиотеке diffusers от Hugging Face. Написал пятистрочный скрипт и подумал почему бы ради прикола не приделать к нему телеграм бота, и тут понеслось.

Месяц я очень мало спал, забывал есть, ложился когда солнце уже вставало и накидал очень кривой, но рабочий прототип тг-бота, который умел получать команды от пользователя, пихать это в виде задачи в очередь rabbitmq. Из очереди её получал второй скрипт, который собственно генерил картинку и возвращал её назад для отправки в телегу.

SD1.5 считает блондинок шлюхами

Где-то тут еще столкнулся с проблемой, что встроенный в SD1.5 детектор nsfw работает очень плохо (а он мне был нужен) и явно предвзят к блондинкам. Если детектору показать две абсолютно одинаковых картинки на одной из которых светлая девушка - он с большой вероятностью определит её как шлюху nsfw. Пришлось его выкинуть и искать замену. Нашелся проект GantMan/nsfw_model. Он работает с библиотекой Keras, потратил какое-то время в попытках конвертировать модель для pytorch, но потом бросил и использовал как есть. В итоге эта модель показывает хорошие результаты, ложных срабатываний я не видел и в отличие от бинарного результата True/False от stable diffusion, она возвращает веса по нескольким категориям. Тут же научился блурить nsfw если у пользователя не отключен фильтр и накладывать водяной знак.

Еще после экспериментов сделал прикольный апскейл, когда после real-esrgan изображением попадает в sd1.5 img2img - на выходе получается очень хороший результат, а не просто мыло как после просто real-esrgan.

Показал бота паре знакомых и все сошлись на том, что это можно монетизировать.

Потом началась скучная часть переписывания бота без дополнительных потоков и глобальных переменных 🤦️ Переехал с Pika (библиотека для rabbitmq) на асинхронный Propan. Переписал работу с Redis, которая у меня используется как база. Сделал небольшую абстракцию для получения и сохранения объектов и стал использовать Pydantic для валидации объектов. Параллельно в налоговой поменял налоговый статус, заключил договор с платежной системой и добавил боту возможность покупок внутренней валюты для работы с нейронкой.

Изначально всё это работало с SD1.5, но пока писал, зарелизили SDXL1.0 и решил перевести проект на нее. И тут столкнулся с несколькими проблемами:

  1. sdxl значительно больше (6ГБ против 2ГБ) и поэтому отвалился функционал выбора из нескольких моделей
  2. sdxl почему-то падала с out-of-memory на моих 12ГБ VRAM на изображениях >= 960x960

В итоге первый вопрос решил тем, что предзагружаю несколько моделей в RAM, а не в видеопамять как раньше и затем выгружаю/загружаю в VRAM в зависимости от запроса пользователя. Временем выгрузки и загрузки можно пренебречь, происходит это быстро.

Второй вопрос решился заменой vae на madebyollin/sdxl-vae-fp16-fix. Пока искал как решить проблему, несколько раз встречал упоминание этой vae, но по описанию она решала проблему с черным декодированным изображением, а не OOM. Но в очередной подход к проблеме решил протестировать эту vae и о чудо, я могу генерировать изображения в нативном для SDXL разрешении 1024x1024 практически за то же время, что и 896x896. И еще можно тайлиг включить для меньшего потребления памяти. Делается как-то так:

vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)

pipe = StableDiffusionXLPipeline.from_single_file(
    model_path,
    vae=vae,
    use_safetensors=True,
    variant="fp16",
    torch_dtype=torch.float16)
pipe.enable_vae_tiling()

Что еще в планах

Апскейл Real-ESRGAN очень плохо обходится с делатями, всё становится гладким, особенно видно на коже и текстурах. Как писал выше, у меня был хороший вариант с r-esrgan + img2img, но от него пришлось отказаться при переходе от sd1.5 к sdxl т.к. она падала на больших изображениях. Возможно с фиксом vae это уже будет работать, надо тестировать.