# Ra — sdxs-VAE branch (experimental)

Это **отдельная ветка** модели Ra: то же пирамидальное тело, но латентное пространство берётся
не из нашего `f8c32` автоэнкодера, а из **`AsymmetricAutoencoderKL`** (`sdxs-micro/vae`):
энкодер f8, декодер f16, поэтому латент 32×32 декодируется в **512×512** — при тех же 1024 токенах
вдвое большее разрешение на выходе.

| | основная ветка ([`AiArtLab/Ra`](https://huggingface.co/AiArtLab/Ra)) | эта ветка |
|---|---|---|
| VAE | `vae/` — наш f8c32 (AE + обученный поточный хвост) | `vae_sdxs/` — AsymmetricAutoencoderKL |
| тело | 272.5M (тело 238.5M + миксер 10.0M + хвост 23.9M) | 248.6M (тело 238.5M + миксер 10.0M), хвоста нет |
| где учится | латент + пиксельный лосс через хвост | **только латент**, декодирует VAE |
| латент | 32×32 (f8, 256² на входе) | 32×32 (f8, 256² на входе) |
| выход | 256×256 | **512×512** (декодер ×2) |

Хвост (поточный пиксельный декодер) в этой схеме не нужен: это классическая латентная диффузия —
модель предсказывает латент, картинку делает VAE.

## Что лежит тут

```
transformer_sdxs/   Safetensors + config.json — тело этой ветки (fp16)
vae_sdxs/           AsymmetricAutoencoderKL (fp16): config.json + diffusion_pytorch_model.safetensors
src/                код (тот же, что в основной ветке, плюс адаптер VAE и поток TFRecord)
run_sdxs.sh         запуск дообучения одной командой
requirements.txt    зависимости (для этого VAE нужен diffusers >= 0.37)
```

## Сэмплинг

```bash
python3 src/sample.py --vae sdxs --model transformer_sdxs \
    --prompt "A butterfly with white and brown patterned wings." --out butterfly.png
```
`--vae sdxs` обязателен: без него латент будет декодироваться нашим VAE и получится мусор.

## Обучение

```bash
bash run_sdxs.sh /path/to/tfrecord_shards 10      # 10 эпох по шардам i1 (PNG 256² + подписи)
```
Скрипт сам подставляет `--vae sdxs`, `--init transformer_sdxs`, латентный режим без хвоста,
`--batch 16`, lr `1e-4 → 1e-5`, сэмплы каждые 2000 шагов. Переопределяется через env:
`INIT`, `STEP`, `OUT`, `BATCH`, `LOG`.

Данные — TFRecord-шарды из [`zlab-princeton/i1-midjourneyv6-tfrecord`](https://huggingface.co/datasets/zlab-princeton/i1-midjourneyv6-tfrecord)
(PNG 256×256 + 5 длинных подписей на картинку); поток читает их без TensorFlow и берёт случайную
из двух самых коротких подписей.

Веса в `transformer_sdxs/` — шаг **5000** (экспорт из чекпойнта sdxs-ветки;
номер шага в config не пишется, он передаётся параметром `--step`, по умолчанию 5000).

## Что проверено

- Переход с основной ветки на этот VAE работает: тот же чекпойнт, дообучение на бабочках —
  `mse_lat` 0.323 → 0.237 за 4 тысячи шагов (10 минут), силуэт бабочки на 512×512 уже читается.
- Латентный режим дешевле: 114 мс/шаг при батче 16 против ~240 мс с хвостом.
- Энкодер этого VAE стоит ~172 мс на батч 16, поэтому на потоке данных шаг ~347 мс (46 карт/с);
  кэширование латентов ускорило бы вдвое — пока не сделано.