576e8

README.md

@@ -2,50 +2,40 @@
 license: apache-2.0
 ---
 
-Краткая инструкция по установке
-Обновите систему и установите git-lfs:
-
-```
-apt update
-apt install git-lfs
-```
-Обновите pip и установите требуемые пакеты:
-
-```
-python -m pip install --upgrade pip
-pip install -r requirements.txt
-```
-Настройте git-credentials:
-
-```
-git config --global credential.helper store
-```
-Клонируйте репозиторий:
-
-```
-git clone https://huggingface.co/AiArtLab/sdxs
-cd sdxs/
-```
-Подготовьте датасет:
-
-```
-mkdir datasets
-cd datasets
-git clone https://huggingface.co/datasets/AiArtLab/imagenet-1kk
-```
-Выполните вход в сервисы:
-
-```
-huggingface-cli login
-wandb login
-```
-Настройте accelerate для многокарточного обучения:
-
-```
-accelerate config
-```
-Запустите обучение!
-
-```
-nohup accelerate launch train.py &
-```
+# SDXS: Быстрая, простая, многоязычная модель для генерации изображений
+
+**Компактная архитектура | Обучение на потребительских GPU | Свободная лицензия**
+
+## Предисловие
+Мы — **AI Art Lab**, небольшая команда энтузиастов с ограниченным бюджетом. Наша цель — создать компактную и быструю модель, которую можно обучать на пользовательских видеокартах (полный цикл обучения, не LoRA), мы выбрали U-Net за его способность эффективно работать с малыми датасетами и быстро обучаться даже на 16GB GPU (например, RTX 4080). Наш бюджет был ограничен несколькими тысячами долларов, что значительно меньше, чем например у SDXL (десятки миллионов).
+
+## Архитектура энкодера (текст и изображения)
+Мы исследовали различные энкодеры и пришли к выводу, что большие языковые модели / T5 XXL не нужны для качественной генерации. Нам был необходим энкодер, который хорошо понимает контекст запроса. Мы использовали мультиязычный энкодер **Mexma-SigLIP** с поддержкой 80 языков, который обрабатывает предложения, а не отдельные токены. Mexma принимает до 512 токенов, что создаёт большую матрицу и замедляет обучение. Поэтому мы использовали пулинг-слой, упрощающий вектор до 1x1152. Точнее мы пропускаем через линейную модель / текстовый проектор - чтобы добиться совместимости с эмбедингами SigLIP. Это позволило синхронизировать текстовые эмбединги с изображениями, что в перспективе даст единую мультимодальную модель. Данный функционал позволит смешивать в запросе ембединги изображений с текстовыми описаниями. Более того модель можно учить без описаний текстом вовсе, подавая только картинки. Что должно упростить тренировку модели например на видео - где задача разметки существенно затруднена, а также получить более конистентную и бесконечную видео генерацию - подавая на вход эмбединг предыдущих кадров с затуханием. Те в перспективе мы хотели бы расширить модель до 3д/видео генерации
+
+## Архитектура U-Net
+Мы выбрали плавную пирамиду каналов: [384, 576, 768, 960], с двумя слоями на блок, [4, 6, 8, 10] трансформеров и 1152/48=24 головы внимания. Эта архитектура обеспечивает наивысшую скорость обучения при размере модели около 2 миллиардов параметров (влазит в RTX 4080). Мы предполагаем, что благодаря большей "глубине" - качество будет выше, чем у SDXL, несмотря на меньший "размер". Модель может быть расширена до 4 миллиардов параметров (путем добавления слоя 1152), что приведёт к идеальной симметрии с размером эмбединга.
+
+## Архитектура VAE
+Мы выбрали нестандартный, 8x 16-канальный AuraDiffusion VAE, который сохраняет детали/текст/не искажает анатомию и цветовую палитру и при этом в нем отсутствует "дымка", характерная для SD3/Flux. Мы использовали быструю версию с FFN свёрткой, в которой наблюдаем незначительные повреждения текстур на тонких узорах (что видимо занижает рейтинг данного VAE). Апскейлеры, такие как ESRGAN, помогут справиться с этими артефактами. Вы всех остальных аспектах мы нашли данный VAE сильно недооцененным.
+
+## Процесс обучения
+### Оптимизатор
+Мы тестировали несколько оптимизаторов (adamw, laion, optimi-adamw, adafactor  и adamw-8bit) и выбрали **AdamW-8bit**. Наиболее плавную кривую с плавным затуханием градиентов продемонстрировал optimi-adamw. Наиболее хаотично ведет себя adamw-8bit, но он позволяет увеличить размер батча, что позволяет получить наивысшую скорость обучения на low price gpu (мы использовали 4*A6000/5*L40 для обучения).
+
+### Learning Rate
+Мы пришли к выводу что манипуляции с кривой затухания/прогревом дают эффект, но несущественный. Также выбор оптимального лернинг рейт кажется нам переоцененным. По нашим экспериментам адам прощает и позволяет выбрать лернинг рейт в широком диапазоне. Мы начинали с 1e-4 плавно понижая по мере обучения до 1e-6. Другими словами выбор корректной архитектуры модели имеет катастрафически большее значение чем манипуляции с гиперпараметрами. 
+
+### Датасет
+Мы обучали модель всего на ~1 миллионе изображений. 60 эпох на Imagenet в 256 (в целом бесполезный из-за низкого качества разметки), и 8 эпох в 576 на Midjourney/Nijijourney + реалистичные фото и аниме/арты. При разметке мы использовали человеческие промпты, промпты каптион емпорум, wd-tagger и moondream2. Мы по разному компоновали длину промпта и состав - от разметки только тегами до подробных описаний чтобы получить модель которая понимает различные стили промптирования. Размер датасета крайне мал - что приводит к тому что модель не знает огромного количества сущностей и будет неорректно рисовать например "гуся на велосипеде". Также в датасете было огромное количество вайфу и тому подобных генераций - так как нам интересно было исследовать насколько хорошо модель выучит анатомию человека, а не способность к рисованию бургеров. Подавляющее большинство описаний было на английском, но наши тесты на других языках показывают что модель мультиязычна.
+
+## Ограничения
+- Низкое покрытие концепциями из-за экстремально небольшого датасета.
+- Image2Image функционал требует дообучения. (Мы снизили долю siglip до 5% чтобы успеть "протренить" text2image)
+
+## Благодарности
+- **Стэн** — ключевой инвестор.
+- **Вадим** — материальная поддержка.
+- **Lovescape** & **Whargarbl** — сообщество тестировщиков.
+- **Caption Imperium** — датасеты.
+
+> "Мы верим, что будущее — за эффективными компактными моделями. Мы благодарим за безвомездные пожертвоания и очень надеемся на вашу поддержку в будующем. 

TRAIN.md

@@ -0,0 +1,51 @@
+---
+license: apache-2.0
+---
+
+Краткая инструкция по установке
+Обновите систему и установите git-lfs:
+
+```
+apt update
+apt install git-lfs
+```
+Обновите pip и установите требуемые пакеты:
+
+```
+python -m pip install --upgrade pip
+pip install -r requirements.txt
+```
+Настройте git-credentials:
+
+```
+git config --global credential.helper store
+```
+Клонируйте репозиторий:
+
+```
+git clone https://huggingface.co/AiArtLab/sdxs
+cd sdxs/
+```
+Подготовьте датасет:
+
+```
+mkdir datasets
+cd datasets
+git clone https://huggingface.co/datasets/AiArtLab/imagenet-1kk
+```
+Выполните вход в сервисы:
+
+```
+huggingface-cli login
+wandb login
+```
+Настройте accelerate для многокарточного обучения:
+
+```
+accelerate config
+```
+Запустите обучение!
+
+```
+nohup accelerate launch train.py &
+```

sdxs/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:74770515282e22ea2f330710e193ee7fb3c81f14895233a8bf2003501019940e
+oid sha256:c1521efedf07d86f8a40e6fe334128e0848faf9a60c92a7e7266f436246d34fe
 size 4529095968