agat

Генератор картинок 512×512 по текстовой подписи: 223.4M параметров, замороженный текстовый энкодер Ettin-68M, свой латентный VAE 4x. Это Agate Preview 003 (MIT), доработанный и переобученный на LucasFang/FLUX-Reason-6M с нашим латентом, целью и тренером.

Как работает

подпись -> Ettin-68M -> 256 токенов (паддинг скрыт маской)
                          │
шумный латент 128×128 ────┼─► МЫШЛЕНИЕ 43.1M ─► план 16×16×640 ─┐
                          │                                    │
                          └─► РЕНДЕРЕР 149M, U-Net на 128×128 ───┴─► скорость (шум − данные)
                              288 каналов на 128, 832 на 64, 1664 на 32

• Текст входит в рендерер на трёх масштабах: joint-внимание над [латент, план, текст] на 32×32, cross-attention адаптер на 64×64 и адаптер на 128×128 — там, где решается деталь в 4 px.
• План — интерфейс, а не признаки: одна карта 16×16, превращённая в shift/scale/gate для каждого блока, плюс 4 канала свёртки 3×3 по самому латенту для локальной подстройки.
• Flat, без второй сети: U-Net считает прямо на латенте 128×128, никакого апсемплинга; 256 px — тот же код на латенте 64×64.
• Скорость шум − данные, t — уровень шума (1 = шум, 0 = данные), обучение на logit-normal t.

Установка и запуск

./install.sh                                        # torch (cu130) + requirements.txt
python3 test_pipeline.py                            # 25 проверок, exit 1 при падении
python3 pipeline.py gen --prompt "a red fox on a mossy log" --out fox.png    # 512 px
python3 pipeline.py gen --prompt "a red fox" --size 256 --out fox256.png     # 256 px

Из питона:

from pipeline import load_components, sample_images
model, text, vae, meta = load_components(device="cuda")      # transformer/model.safetensors
imgs = sample_images(model, text, vae, meta, ["a red fox on a mossy log"],
                     steps=50, cfg=3.0, seed=0, size=512)

Архитектура читается из config.json рядом с весами, окружение задавать не нужно.

Обучение

setsid nohup ./train.sh > logs/watch.log 2>&1 &        # resume с transformer/, 512 px, батч 4, 250 000 шагов

Примеры

Раскладка

agat_pipeline.py   обёртка модели (наши конвенции), TextEncoder, константы, config_dict
model/             генератор: thinker.py (Agate 003, наши правки помечены PATCH:), layers.py,
                   joint.py (joint-внимание и cross-attention адаптеры), text_encoder.py (Ettin-68M)
pipeline.py        CLI: prepare раскладки и gen; load_components/sample_images
train.py           обучение: поток FLUX-Reason / папка пар / одна картинка; лоссы; DDP
train.sh           прогон под сторожем: все параметры — дефолты train.py
test_pipeline.py   25 проверок пайплайна, без фреймворков
watch_train.sh     сторож: перезапуск при падении и зависании
transformer/       чекпойнт + config.json (чекпойнт описывает сам себя)
vae/  text_encoder/  scheduler/  media/
docs/TECHNICAL.md  подробности: геометрия, конвенция t, лоссы, DDP, замеры
logs/  samples/  runs/   генерируется, в .gitignore