| # 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. |
|
|
