Ra — sdxs-VAE branch (experimental)

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

	основная ветка (`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)

Сэмплинг

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