| # 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 карт/с); |
| кэширование латентов ускорило бы вдвое — пока не сделано. |
|
|
