File size: 6,440 Bytes
b4f5f26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | # 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
```
## Запуск
```bash
# данные: бакеты кратны 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`.
## Правила, выведенные из замеров
1. **Ширина растёт к низу** (192 / 512 / 1024 / 1536). Плоский стек ширины 1024 на всех уровнях
стоил 196 мс/шаг и дал hf 0.18; пирамида при 97 мс даёт 0.54.
2. **Текст — на всех уровнях**, включая самый мелкий: кросс-внимание в каждом conv-блоке L0 и
отдельный self-attn блок с текстом на L0. Стоило +2.5M параметров и +9 мс, `switch` вырос
0.43 → 0.78.
3. **Хвост учится на латенте `z0` плюс ошибка тела без градиента.** Обучение на предсказаниях тела
деградировало рендерер с 29.6 до 21.3 дБ; на чистом `z0` — рендерер не видит реальных ошибок.
4. **Бакеты кратны 64**: три даунсэмпла требуют латента, кратного 8 (`h % 8 == 0`, `w % 8 == 0`).
Аспект — центр-кроп, не растяжение; в `allphoto_train` чёрные поля резать (`trim_black`).
5. **Замер, а не впечатление.** В каждом логе: `mse_lat` (общий, low t<0.3, high t>0.7), `mse_pix`,
`switch` (один шум × два капшена), `tail(z0)` (сторож деградации хвоста), `skip` (не-конечные шаги).
Детали меряются долей высоких частот `hf(gen)/hf(AE.recon)`.
6. **Цель — `x0`**, не velocity: обосновано только тем, что так наследуются веса декодера AE; смена
параметризации тела на декодер не влияет.
7. **Никаких костылей в обучении** (фиксированный шум хвоста, клип латента, «mix» латентов):
всё, что влияет на результат, должно быть частью стандартного flow-matching.
|