В общем, кому не хватало 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 и решил перевести проект на нее. И тут столкнулся с несколькими проблемами:
- sdxl значительно больше (6ГБ против 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 это уже будет работать, надо тестировать.