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