File size: 4,185 Bytes
de9c421 3dd752f de9c421 3dd752f de9c421 3dd752f de9c421 3dd752f de9c421 0f8003c de9c421 4ad8e7e de9c421 4ad8e7e de9c421 4ad8e7e de9c421 3dd752f de9c421 3dd752f de9c421 3dd752f de9c421 3dd752f de9c421 3dd752f de9c421 | 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 | # 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 карт/с);
кэширование латентов ускорило бы вдвое — пока не сделано.
|