Supra2-IMG: DiT-модель text-to-image на 100M параметров за 10 часов

Официальный баннер Supra2-IMG из репозитория модели.

Что это такое

Supra2-IMG: крошечный DiT, 104.1M параметров, генерирующий изображения 256x256 из 128 токенов контекста Flan-T5. Текстовое условие формируется замороженным энкодером Flan-T5-Base, а латентные представления декодируются через SD-VAE-FT-MSE, поэтому оба этих компонента берутся из своих репозиториев, а не поставляются в checkpoint.

SupraLabs: независимая лаборатория, стоящая за языковыми моделями Supra2 на 100M параметров, и Supra2-IMG применяет тот же подход «обучай малым, обучай с нуля» к генерации изображений. Модель выпущена как чистый код PyTorch плюс веса: inference.py, config.json и model_final_ema.pt.

Обучение

Запуск намеренно скромный, в чём и смысл релиза:

• Датасет: 10 эпох по полному датасету LucasFang/FLUX-Reason-6M, 5.6M изображений после подготовки.
• Выбор подписей: для каждого изображения берётся первая доступная подпись в цепочке caption_composition, затем caption_entity, caption_text, caption_style и caption_imaginative, что сохраняет наиболее качественную аннотацию для каждого образца.
• Оборудование: одна Nvidia H100 SXM 80GB на RunPod, 9 часов включая подготовку данных, на диске 2.5TB.

Поскольку вся модель занимает несколько сотен мегабайт, это практичная база для изучения предобучения в малом масштабе и для экспериментов по файнтюнингу, которые были бы слишком дорогими в масштабе 8B.

Примеры

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

Пример вывода 256x256, полученный поставляемым скриптом вывода.

Доступность

Поддержки ComfyUI нет: Supra2-IMG это обычный релиз PyTorch с собственной архитектурой SupraDiT, поэтому это не pipeline diffusers, и ни шаблон, ни нода ComfyUI не опубликованы. Для запуска нужен поставляемый скрипт плюс два внешних компонента:

wget https://huggingface.co/SupraLabs/Supra2-IMG/resolve/main/inference.py
python inference.py --prompt "a sea jellyfish floating in the pitch-black ocean depths" \
  --seed 0 --cfg 3.0 --steps 50 --n 1 --out jellyfish.png

Рекомендуемые настройки выборки: seed 0, CFG 3.0 и 50 шагов.

• Веса и код: SupraLabs/Supra2-IMG
• Текстовый энкодер: google/flan-t5-base
• VAE: stabilityai/sd-vae-ft-mse
• Лаборатория: supra-labs.com