Simple Diffusion — компактная модель генерации изображений

Simple Diffusion — компактная модель генерации изображений

Мы представляем Simple Diffusion (sdxs-1b) — компактную, быструю и открытую диффузионную модель, которую можно обучать и запускать на обычных видеокартах, сохранив высокое качество генерации. Модель выпускается как альфа-версия под лицензией Apache-2.0 вместе с открытым кодом для подготовки данных и обучения. Доступна на Hugging Face.

Как создавалась модель

Работа над Simple Diffusion началась в декабре 2024 года в рамках AiArtLab — небольшой группы энтузиастов, экспериментирующих с генеративными моделями. Изначально мы изучали линейные трансформеры на базе Sana, но после неудачных попыток дообучения решили разрабатывать архитектуру с нуля.

С февраля 2025 года один из авторов перешёл к модификации UNet-архитектуры, вдохновлённой Stable Diffusion 1.5, а второй сосредоточился на DiT и разработке быстрых автоенкодеров. За полгода было протестировано множество архитектур, но большинство из них не справлялись с обобщением анатомии или детализацией.

С августа по сентябрь 2025 года мы сосредоточились на улучшении VAE. Эксперименты привели к созданию SOTA-версии VAE (на тот момент) и набору улучшений для различных автоенкодеров.

В декабре 2025 года была обучена первая версия sdxs-0.8 — модель с 0,8 млрд параметров, использующая UNet в стиле SD1.5, Long CLIP и 16-канальный Simple VAE с flow matching таргетом.

Хотя модель неплохо справлялась с аниме-стилями, мы стремились к более высокому качеству и фотореалистичности. В январе–феврале 2026 года протестировали ключевые идеи из SDXL, но они либо ухудшали результат, либо замедляли обучение без заметного прироста качества.

Разработка VAE и апскейлинг

В этот период мы вернулись к изучению Flux.2 и обнаружили, что его VAE — 128-канальный, что оказалось избыточным для небольших моделей. В итоге мы адаптировали его под 32 канала, создав асимметричный VAE, который был дообучен за два дня на одной GPU.

Обучение VAE напрямую приводило к размытию деталей, поэтому мы использовали 5 параллельных целевых функций и специальный алгоритм нормализации.

Качество VAE

Наш 32-канальный VAE показывает рекордные значения PSNR и LPIPS среди 16× моделей, значительно превосходя Flux1 VAE даже при сравнении с 8× сжатием. Это стало возможным благодаря удачной основе — Flux.2.

Позже мы перешли на Qwen3.5 в качестве текстового энкодера. Его скорость обучения оказалась значительно выше, чем у немультимодальных моделей, а также добавила мультимодальность и поддержку нескольких языков.

Данные и предобработка

Для обучения использовали выборку из 1–2 миллионов изображений из открытых датасетов, включая Midjourney и Nijourney. Большинство — иллюстрации и рисунки, с небольшой долей фотографий.

Подписи представлены как в формате danbooru (теги), так и в виде описаний до 250 токенов. Изображения масштабировались от 768 до 1408 пикселей (с шагом 64) для поддержки разных пропорций, хотя обучение велось на вдвое меньшем разрешении.

Финальная архитектура

SDXS-1B состоит из трёх компонентов:

  • UNet — 1,6 млрд параметров. Архитектура ближе к Stable Diffusion 1.5, чем к SDXL. Небольшое и равномерное количество блоков обеспечивает сбалансированное внимание к анатомии и деталям.
  • Текстовый энкодер Qwen3.5-2B. Мы протестировали CLIP, LongCLIP, SigLIP и другие. Qwen3.5 показал качество, близкое к LongCLIP, но добавил мультимодальность и возможность рефайнинга. Эмбеддинги берутся с предпоследнего слоя для лучшего понимания структуры.
  • Асимметричный VAE — 32 канала в латентном пространстве, энкодер 8×, декодер 16×. Это означает, что изображение сжимается в 8 раз, а при декодировании восстанавливается вдвое крупнее. Такой подход работает как встроенный 2× апскейлер.

Ключевая особенность: встроенный апскейлер

Апскейлер увеличивает изображение в 2 раза без добавления новых деталей, сохраняя стиль и устраняя пикселизацию. Это особенно полезно при генерации в диапазоне 512–768 пикселей, где инференс значительно быстрее, чем в 1024, но качество остаётся высоким.

Такой подход применим не только к генерации, но и к увеличению медицинских снимков или копий картин, где важна точность, а не креативность.

Уточнение промптов и режимы ввода

Модель поддерживает улучшение коротких или теговых запросов с помощью встроенного рефайнера. Также можно подавать на вход изображение (image-to-image) или даже звук — все эти режимы находятся в экспериментальной стадии.

Процесс обучения и гиперпараметры

Обучение проходило поэтапно: сначала — базовые формы, затем композиция, и наконец — детали (лица, текстуры). Использовался AdamW8bit с базовым LR 4e-5 и минимальным 4e-6. Изменение learning rate по сценариям (0.5, 1.0, 3.0, 5.0) не дало стабильного улучшения.

Эксперименты с Muon привели к деградации анатомии. В целом, роль гиперпараметров оказалась переоценена — модель либо обучается, либо нет.

Обучение проводилось на 1–8 GPU RTX 5090 и заняло около 2–3 месяцев. Для устойчивости к шуму применялся dropout 10% на выходе текстового энкодера, а длина текста ограничивалась 250 токенами.

Весь код открыт: в репозитории доступны скрипты для подготовки данных и монолитный train.py. Подробные инструкции — в README.

Использование модели

SDXS-1B доступна через библиотеку Diffusers от Hugging Face. В пайплайн можно интегрировать улучшение промптов и апскейлинг.

Среднее качество на случайных промптах ниже из-за ограниченного датасета.

Ограничения

Модель остаётся в альфа-версии. Из-за небольшого и несбалансированного датасета она лучше справляется с иллюстрациями, чем с фотографиями. Рекомендуется добавлять «photo» в negative prompt, чтобы избежать попыток фотореализма.

Лицензия — Apache-2.0, разрешает свободное использование в коммерческих и исследовательских целях.

Перспективы

Обучение временно приостановлено. Мы публикуем модель «как есть», надеясь на вклад сообщества.

Следующие шаги: расширение датасета, завершение обучения, разработка Turbo LoRA, ControlNet и переход к генерации видео.

Модель уже доступна на Hugging Face. Приветствуются предложения о сотрудничестве — контакты указаны на странице модели.

Читать оригинал