Ra (UDiT) — пирамидальный U-трансформер для text-to-image
Название модели — Ra (бог солнца, в честь пирамиды ширин 192 → 512 → 1024 → 1536);
UDiT — имя архитектуры в коде.
Одна модель из трёх частей: тело (пирамидальный U-трансформер, flow в латентном пространстве AE), хвост (декодер AE, flow по пикселям) и текст (замороженный Ettin-150M плюс обучаемый выбор слоёв). 177.3M параметров, обучение на 16 ГБ.
текст: Ettin-150M (768, 22 слоя, заморожен)
└─ SliceMixer: 7 слоёв → attention по оси срезов → Linear(7→1) [10.04M]
шум (32, h, w) ─► ТЕЛО, пирамида ширины 192 → 512 → 1024 → 1536 [143.3M]
L0 ÷1 conv + текст 6 enc / 4 dec 4096 токенов при 512², 1024 при 256²
L1 ÷2 трансформер 4 / 3 1024 при 512², 256 при 256²
L2 ÷4 трансформер 3 / 2 256 / 64
L3 ÷8 трансформер 2 64 / 16 ← широкий низ, полное внимание
скипы аддитивные, вниз mean-pool 2×+Linear, вверх Linear+nearest 2×
─► латент z (32, h, w) ─► ХВОСТ: декодер AE (6 блоков @512, PixelShuffle ×3) [23.87M]
flow по пикселям, условие z0 + (ошибка тела).detach()
─► картинка
Измерено (RTX 4080 16 ГБ, 256², bs8)
| шаг обучения | 97 мс (8.1 ГБ), 121 обновление на эпоху при 970 картинках |
| параметры | тело 143.3M + текст 10.04M + хвост 23.87M = 177.3M |
| бабочки, 20k шагов | mse_lat 0.107 (low 0.022, high 0.187), mse_pix 0.0105, switch 0.78, tail(z0) 24.7 дБ |
| детали (hf, Paper Kite) | 0.54 против AE 0.859 (63%); у прежней плоской 1024-архитектуры было 0.18 |
| потолок AE (256², бабочки) | hf 0.859 при GT 0.913 |
Подробности архитектуры — docs/ARCH.md, состояние и планы — docs/STATUS.md.
Состав
src/model.py UDiT: тело-пирамида + хвост + sample()
src/text.py Ettin + SliceMixer (выбор слоёв)
src/train.py обучение на папке «картинка + .txt»
src/make_dataset.py сборка набора из плоской папки (бакеты кратны 64)
src/make_dataset_people.py сборка набора по классам из allphoto_train (woman/man/cat)
src/export.py экспорт в transformer/ (safetensors + config.json)
vae/ AE f8c32-ae: энкодер 68.4M + декодер 23.9M, f8, 32 канала
text_encoder/ Ettin-150M (768, 22 слоя, MIT)
transformer/ текущие веса модели + config.json
tmp/ прогоны (samples, log.csv, last.pt); tmp/archive — старое
docs/ARCH.md, docs/STATUS.md
Запуск
# данные: бакеты кратны 64 (тело трижды делит сетку на 2)
python3 src/make_dataset.py --src ~/dataset/butterflies/butterfly \
--out ~/dataset/butterflies_256x64 --long-side 256 --eval-n 30
# обучение (256², реальный батч 8, ~16 минут на 10k шагов)
python3 src/train.py --data ~/dataset/butterflies_256x64/train \
--eval ~/dataset/butterflies_256x64/eval --bs 8 --steps 20000 --every 500 --cfg 1.5
# продолжить с чекпойнта
python3 src/train.py ... --resume --out tmp/pyr2
# экспорт весов
python3 src/export.py --ckpt tmp/pyr2/last.pt --out transformer
Смотреть прогресс: tmp/<run>/samples/latest.png (панели по промптам), tmp/<run>/log.csv.
Правила, выведенные из замеров
- Ширина растёт к низу (192 / 512 / 1024 / 1536). Плоский стек ширины 1024 на всех уровнях стоил 196 мс/шаг и дал hf 0.18; пирамида при 97 мс даёт 0.54.
- Текст — на всех уровнях, включая самый мелкий: кросс-внимание в каждом conv-блоке L0 и
отдельный self-attn блок с текстом на L0. Стоило +2.5M параметров и +9 мс,
switchвырос 0.43 → 0.78. - Хвост учится на латенте
z0плюс ошибка тела без градиента. Обучение на предсказаниях тела деградировало рендерер с 29.6 до 21.3 дБ; на чистомz0— рендерер не видит реальных ошибок. - Бакеты кратны 64: три даунсэмпла требуют латента, кратного 8 (
h % 8 == 0,w % 8 == 0). Аспект — центр-кроп, не растяжение; вallphoto_trainчёрные поля резать (trim_black). - Замер, а не впечатление. В каждом логе:
mse_lat(общий, low t<0.3, high t>0.7),mse_pix,switch(один шум × два капшена),tail(z0)(сторож деградации хвоста),skip(не-конечные шаги). Детали меряются долей высоких частотhf(gen)/hf(AE.recon). - Цель —
x0, не velocity: обосновано только тем, что так наследуются веса декодера AE; смена параметризации тела на декодер не влияет. - Никаких костылей в обучении (фиксированный шум хвоста, клип латента, «mix» латентов): всё, что влияет на результат, должно быть частью стандартного flow-matching.