import torch
from diffusers import DiffusionPipeline
# switch to "mps" for apple devices
pipe = DiffusionPipeline.from_pretrained("AiArtLab/ra", dtype=torch.bfloat16, device_map="cuda")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]YAML Metadata Warning:empty or missing yaml metadata in repo card
Check out the documentation for more information.
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_lat0.323 → 0.237 за 4 тысячи шагов (10 минут), силуэт бабочки на 512×512 уже читается. - Латентный режим дешевле: 114 мс/шаг при батче 16 против ~240 мс с хвостом.
- Энкодер этого VAE стоит ~172 мс на батч 16, поэтому на потоке данных шаг ~347 мс (46 карт/с); кэширование латентов ускорило бы вдвое — пока не сделано.
- Downloads last month
- -