Ra fp16 (step 212000)

.gitattributes

@@ -1,35 +1,9 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+docs/presentation.pdf filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,25 @@
+# python
+__pycache__/
+*.py[cod]
+
+# прогоны обучения, чекпойнты, сэмплы
+tmp/
+runs/
+*.pt
+nohup.out
+*.log
+
+# локальная постановка задачи (не для публикации)
+task.md
+
+# jupyter
+.ipynb_checkpoints/
+
+# кэши моделей и окружение
+.hf_home/
+.cache/
+text_encoder/.cache/
+.env
+
+# системное
+.DS_Store

docs/ARCH.md

@@ -0,0 +1,79 @@
+# Архитектура
+
+Три части, каждая со своей задачей. Всё, что ниже, проверено запуском.
+
+## 1. Тело: пирамидальный U-трансформер
+
+Принцип: **ширина растёт там, где токенов мало.** Стоимость слоя ~ `токены × C²`, поэтому на
+высоком разрешении держим узкие свёрточные блоки, а ширину и полное внимание — внизу.
+
+| уровень | токенов при 512² | ширина | блоки | механика |
+|---|---|---|---|---|
+| L0 ÷1 | 4096 (512²) / 1024 (256²) | **192** | 6 enc / 4 dec + 1 self-attn с каждой стороны | depthwise 3×3 + expand 4× (ConvNeXt-стиль), кросс-внимание к тексту в каждом блоке |
+| L1 ÷2 | 1024 / 256 | **512** | 4 enc / 3 dec | трансформер: self-attn + cross-attn к тексту + FFN |
+| L2 ÷4 | 256 / 64 | **1024** | 3 enc / 2 dec | то же, полное внимание дёшево |
+| L3 ÷8 | 64 / 16 | **1536** | 2 (боттлнек) | то же; широкий низ = глобальные ограничения (симметрия, композиция) |
+
+Переходы и связи:
+
+- **вниз**: mean-pool 2×2 + `Linear`, вверх: `Linear` + nearest 2×;
+- **скипы аддитивные** (SD-стиль): выходы `e0/e1/e2` прибавляются к декодеру того же уровня;
+  ширина на входе и выходе уровня совпадает, поэтому сложение без проекций;
+- **RoPE на каждом уровне**, фаза в **латентных пикселях** (L1 сдвинут вдвое, L2 вчетверо), поэтому
+  одна фаза означает одно расстояние независимо от уровня;
+- **модуляция**: общий вектор `t + pooled-текст` → `(shift, scale, gate)` на уровень; в conv-блоках
+  L0 пять чанков (три на conv-путь, два на вход кросс-внимания);
+- **zero-init**: выход conv-блоков и FFN обнулён, поэтому новые ветки не ломают старт.
+
+Размеры: тело **143.3M**. Шаг обучения 97 мс при bs8 на 256² (8.1 ГБ). Для сравнения, прежний
+плоский стек ширины 1024 на всех уровнях: 196 мс и hf 0.18 против 0.54 у пирамиды (архитектура
+удалена в `tmp/archive`).
+
+## 2. Текст
+
+- **Ettin-150M** (`text_encoder/`): ModernBERT, 768, 22 слоя, byte-level BPE, MIT, заморожен.
+- **SliceMixer** (`src/text.py`, 10.04M обучаемых): 7 срезов (выходы после глобальных слоёв
+  2, 5, 8, 11, 14, 17, 20) → пер-срезный LayerNorm (нормы растут 52 → 548, без нормировки ранние
+  слои не влияют) → 2 блока attention **по оси срезов** → `Linear(7→1)` → zero-init добавка к
+  среднему по срезам. Идея из `moss/adapter`, где 0.8B + 158M голова заменили 9B-энкодер.
+- Вход в тело: кросс-внимание на **всех** уровнях (общие ключи/значения на уровень) плюс
+  pooled-текст в модуляции.
+
+## 3. Хвост
+
+Декодер AE без изменений архитектуры: патчи 8×8, 6 блоков @512, adaLN с низкоранговыми дельтами,
+long skip из блока n−2, финальный `PixelShuffle + conv3×3` ×3. 23.87M.
+
+Обучается как обычный flow по пикселям: случайные `t`, случайный шум, цель — картинка, **условие —
+`z0 + (ошибка тела).detach()`**. Условие без градиента: рендерер видит реальную ошибку тела, но не
+может «научиться исправлять мусор» (именно это раньше увело его с 29.6 на 21.3 дБ).
+
+## 4. Сэмплинг
+
+```
+z ~ N(0,1)
+4–8 шагов Эйлера по телу:  x0 = тело(z, t, текст);  z ← z − (z − x0)/(t·N)
+CFG 1.5 (выше 2 при недообученной безусловной ветви ухудшает латент)
+1 + 3 шага декодера AE:    изображение
+```
+
+При 256² сэмпл стоит ~35 мс на 4800-й карте; при 512² — примерно вчетверо дороже.
+
+## 5. Бюджет
+
+| часть | параметры |
+|---|---|
+| тело | 143.3M |
+| выбор слоёв текста | 10.04M |
+| хвост | 23.87M |
+| **итого** | **177.3M** |
+
+## 6. Что зафиксировано и почему
+
+- **x0, а не velocity.** Смена параметризации тела не затрагивает декодер (это отдельный flow со
+  своей целью), но требует переучивания тела с нуля — поэтому оставлено как есть.
+- **Текст не экономим.** Кросс-внимание на L0 добавило 2.5M параметров и 9 мс, `switch` вырос
+  0.43 → 0.78 при том же уровне деталей.
+- **Никаких фиксированных шумов и клипов в обучении** — только стандартный flow-matching.
+- **Удалено как тупиковое:** плоский 1024-стек, резидуальная U-ветка, Mix-FFN поверх плоского
+  стека, velocity-переключение посреди обучения, LPIPS в лоссе, кропы одной картинки вместо набора.

docs/README.ru.md

@@ -0,0 +1,90 @@
+# Ra (UDiT) — пирамидальный U-трансформер для text-to-image
+
+Название модели — **Ra** (бог солнца, в честь пирамиды ширин 192 → 512 → 1024 → 1536);
+`UDiT` — имя архитектуры в коде.
+
+Одна модель из трёх частей: **тело** (пирамидальный U-трансформер, flow в латентном пространстве
+AE), **хвост** (декодер AE, flow по пикселям) и **текст** (замороженный Ettin-150M плюс обучаемый
+выбор слоёв). 177.3M параметров, обучение на 16 ГБ.
+
+```
+   текст: Ettin-150M (768, 22 слоя, заморожен)
+       └─ SliceMixer: 7 слоёв → attention по оси срезов → Linear(7→1)   [10.04M]
+
+   шум (32, h, w) ─► ТЕЛО, пирамида ширины 192 → 512 → 1024 → 1536                     [143.3M]
+        L0 ÷1  conv + текст   6 enc / 4 dec    4096 токенов при 512², 1024 при 256²
+        L1 ÷2  трансформер    4 / 3             1024 при 512², 256 при 256²
+        L2 ÷4  трансформер    3 / 2             256 / 64
+        L3 ÷8  трансформер    2                 64 / 16  ← широкий низ, полное внимание
+        скипы аддитивные, вниз mean-pool 2×+Linear, вверх Linear+nearest 2×
+   ─► латент z (32, h, w) ─► ХВОСТ: декодер AE (6 блоков @512, PixelShuffle ×3)        [23.87M]
+        flow по пикселям, условие z0 + (ошибка тела).detach()
+   ─► картинка
+```
+
+## Измерено (RTX 4080 16 ГБ, 256², bs8)
+
+| | |
+|---|---|
+| шаг обучения | **97 мс** (8.1 ГБ), 121 обновление на эпоху при 970 картинках |
+| параметры | тело 143.3M + текст 10.04M + хвост 23.87M = **177.3M** |
+| бабочки, 20k шагов | `mse_lat` 0.107 (low 0.022, high 0.187), `mse_pix` 0.0105, `switch` 0.78, `tail(z0)` 24.7 дБ |
+| детали (hf, Paper Kite) | **0.54** против AE 0.859 (63%); у прежней плоской 1024-архитектуры было 0.18 |
+| потолок AE (256², бабочки) | hf 0.859 при GT 0.913 |
+
+Подробности архитектуры — `docs/ARCH.md`, состояние и планы — `docs/STATUS.md`.
+
+## Состав
+
+```
+src/model.py               UDiT: тело-пирамида + хвост + sample()
+src/text.py                Ettin + SliceMixer (выбор слоёв)
+src/train.py               обучение на папке «картинка + .txt»
+src/make_dataset.py        сборка набора из плоской папки (бакеты кратны 64)
+src/make_dataset_people.py сборка набора по классам из allphoto_train (woman/man/cat)
+src/export.py              экспорт в transformer/ (safetensors + config.json)
+vae/                       AE f8c32-ae: энкодер 68.4M + декодер 23.9M, f8, 32 канала
+text_encoder/              Ettin-150M (768, 22 слоя, MIT)
+transformer/               текущие веса модели + config.json
+tmp/                       прогоны (samples, log.csv, last.pt); tmp/archive — старое
+docs/ARCH.md, docs/STATUS.md
+```
+
+## Запуск
+
+```bash
+# данные: бакеты кратны 64 (тело трижды делит сетку на 2)
+python3 src/make_dataset.py --src ~/dataset/butterflies/butterfly \
+        --out ~/dataset/butterflies_256x64 --long-side 256 --eval-n 30
+
+# обучение (256², реальный батч 8, ~16 минут на 10k шагов)
+python3 src/train.py --data ~/dataset/butterflies_256x64/train \
+        --eval ~/dataset/butterflies_256x64/eval --bs 8 --steps 20000 --every 500 --cfg 1.5
+
+# продолжить с чекпойнта
+python3 src/train.py ... --resume --out tmp/pyr2
+
+# экспорт весов
+python3 src/export.py --ckpt tmp/pyr2/last.pt --out transformer
+```
+
+Смотреть прогресс: `tmp/<run>/samples/latest.png` (панели по промптам), `tmp/<run>/log.csv`.
+
+## Правила, выведенные из замеров
+
+1. **Ширина растёт к низу** (192 / 512 / 1024 / 1536). Плоский стек ширины 1024 на всех уровнях
+   стоил 196 мс/шаг и дал hf 0.18; пирамида при 97 мс даёт 0.54.
+2. **Текст — на всех уровнях**, включая самый мелкий: кросс-внимание в каждом conv-блоке L0 и
+   отдельный self-attn блок с текстом на L0. Стоило +2.5M параметров и +9 мс, `switch` вырос
+   0.43 → 0.78.
+3. **Хвост учится на латенте `z0` плюс ошибка тела без градиента.** Обучение на предсказаниях тела
+   деградировало рендерер с 29.6 до 21.3 дБ; на чистом `z0` — рендерер не видит реальных ошибок.
+4. **Бакеты кратны 64**: три даунсэмпла требуют латента, кратного 8 (`h % 8 == 0`, `w % 8 == 0`).
+   Аспект — центр-кроп, не растяжение; в `allphoto_train` чёрные поля резать (`trim_black`).
+5. **Замер, а не впечатление.** В каждом логе: `mse_lat` (общий, low t<0.3, high t>0.7), `mse_pix`,
+   `switch` (один шум × два капшена), `tail(z0)` (сторож деградации хвоста), `skip` (не-конечные шаги).
+   Детали меряются долей высоких частот `hf(gen)/hf(AE.recon)`.
+6. **Цель — `x0`**, не velocity: обосновано только тем, что так наследуются веса декодера AE; смена
+   параметризации тела на декодер не влияет.
+7. **Никаких костылей в обучении** (фиксированный шум хвоста, клип латента, «mix» латентов):
+   всё, что влияет на результат, должно быть частью стандартного flow-matching.

docs/STATUS.md

@@ -0,0 +1,46 @@
+# Состояние и планы
+
+## Что работает (проверено)
+
+- **Пайплайн целиком**: текст → латент → пиксели, обучение и сэмплинг. Baseline-метрики в каждом
+  логе: `mse_lat` (общий, low t<0.3, high t>0.7), `mse_pix`, `switch`, `tail(z0)`, `skip`.
+- **Бабочки, 20 000 шагов** (`butterflies_256x64`, 970 train / 30 eval, 97 мс/шаг):
+
+| шаги | hf (доля от AE 0.859) | mse_lat | low | high | switch |
+|---|---|---|---|---|---|
+| 4 000 | 0.446 (52%) | 0.21 | 0.031 | 0.42 | 0.43 |
+| 10 000 | 0.436 (51%) | 0.164 | 0.024 | 0.322 | 0.68 |
+| 20 000 | **0.540 (63%)** | 0.107 | 0.022 | 0.187 | 0.78 |
+
+  Глазами: узнаваемая бабочка с жилками и рисунком, два сида дают две разные особи.
+  `tail(z0)` 24.7 дБ (потолок AE на 256² — 0.859 по hf), `skip` = 0.
+- **Потолок AE не мешает**: на 256² AE сохраняет 0.859 из 0.913 высоких частот GT.
+
+## Чего ещё не сделано
+
+1. **Плато на бабочках не найдено.** hf 0.436 → 0.540 за 10k шагов, лоссы падают; 30k шагов
+   покажут, где насыщение.
+2. **Люди не обучены на этой архитектуре.** Готов билдер `src/make_dataset_people.py`
+   (5749 фото в `~/dataset/allphoto_train`, классы woman/man/cat, бакеты кратны 64).
+   Это главный тест: там капшены длинные и детальными, то есть деталь *определена условием*.
+3. **512² не запускался** на пирамиде. Нужен набор с бакетами кратны 64 (512×384, 512×256, 384×512,
+   512×512) — `make_dataset.py --long-side 512 --multiple 64`. Ожидание: чешуйки/лица физически
+   разрешимы (бабочка 400 px вместо 200).
+4. **`high` = 0.187** — «выбор особи» на высоком шуме. Падает с обучением (0.42 → 0.19), поэтому
+   причина — недообучение, а не только бедность капшена.
+
+## Известные ограничения
+
+- Латент AE — **f8** (64×64 при 512²). Всё мельче 8×8 px додумывает декодер, а не тело.
+- Бакеты обязаны быть кратны **64** по каждой стороне (три даунсэмпла). При других размерах
+  даунсэмпл нецелый.
+- Сравнение с прежними архитектурами корректно только на одинаковых наборах и шагах: старые
+  чекпойнты в новую архитектуру не переносятся (уровни и ширины другие).
+
+## Порядок работ
+
+1. Дожать бабочек до плато (30k) — 16 минут на 10k шагов.
+2. Перейти на людей (5749) при 256²: 30k шагов ≈ 50 минут.
+3. 512² на людях — после того, как на 256² будут лица/руки.
+4. Ёмкость (ширина боттлнека 1536 → 2048, больше блоков L0) — только если данные и разрешение
+   уже не дают прироста.

docs/presentation.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e25cf92a5db610c2ff0ed13334eef2d22df88ae2d0378a3642afb4ace9fe65f6
+size 953712

docs/presentation.src.html

@@ -0,0 +1,618 @@
+<!DOCTYPE html>
+<html lang="ru" data-register="technical">
+<head>
+<meta charset="utf-8">
+<title>UDiT — пирамидальный U-трансформер</title>
+<style>
+  @page { size: 1280px 720px; margin: 0; }
+  :root {
+    --paper: #f6f2ea; --ink: #14120f; --muted: #6d675c; --rule: #d9d2c3;
+    --copper: #a8551f; --teal: #1f6b66; --sand: #efe8da; --white: #fffdf8;
+  }
+  * { box-sizing: border-box; }
+  html, body { margin: 0; padding: 0; background: var(--paper); }
+  body { font-family: "DejaVu Sans", sans-serif; color: var(--ink); font-size: 15px; line-height: 1.45; }
+  .slide {
+    width: 1280px; height: 720px; padding: 54px 64px 44px; position: relative;
+    background: var(--paper); page-break-after: always; overflow: hidden;
+  }
+  .slide:last-child { page-break-after: auto; }
+  .eyebrow { font-family: "DejaVu Sans Mono", monospace; font-size: 11px; letter-spacing: .16em;
+    text-transform: uppercase; color: var(--copper); margin: 0 0 10px; }
+  h1 { font-family: 'P052', 'URW Bookman', serif; font-size: 46px; line-height: 1.08; margin: 0 0 14px; font-weight: 600; }
+  h2 { font-family: 'P052', 'URW Bookman', serif; font-size: 28px; line-height: 1.18; margin: 2px 0 8px; font-weight: 500; }
+  h3 { font-size: 14px; margin: 0 0 4px; letter-spacing: .02em; }
+  p { margin: 0 0 8px; }
+  .lede { font-size: 17px; color: #2c2822; max-width: 62ch; }
+  .muted { color: var(--muted); }
+  .mono { font-family: "DejaVu Sans Mono", monospace; font-variant-numeric: tabular-nums; }
+  .num { font-family: "DejaVu Sans Mono", monospace; font-weight: 700; }
+  .rule { height: 1px; background: var(--rule); margin: 18px 0 22px; }
+  .foot { position: absolute; left: 64px; right: 64px; bottom: 20px; display: flex;
+    justify-content: space-between; font-family: "DejaVu Sans Mono", monospace; font-size: 10.5px; color: var(--muted); }
+  .grid2 { display: grid; grid-template-columns: 1fr 1fr; gap: 30px; }
+  .grid3 { display: grid; grid-template-columns: repeat(3, 1fr); gap: 22px; }
+  .grid4 { display: grid; grid-template-columns: repeat(4, 1fr); gap: 18px; }
+  .card { background: var(--white); border: 1px solid var(--rule); padding: 16px 18px; }
+  .kpi { display: flex; gap: 26px; flex-wrap: wrap; }
+  .kpi div { border-left: 2px solid var(--copper); padding-left: 12px; }
+  .kpi b { display: block; font-family: "DejaVu Sans Mono", monospace; font-size: 23px; font-weight: 700; }
+  .kpi span { font-size: 12.5px; color: var(--muted); }
+  table { border-collapse: collapse; width: 100%; font-size: 14px; }
+  th, td { text-align: left; padding: 7px 10px; border-bottom: 1px solid var(--rule); }
+  th { font-family: "DejaVu Sans Mono", monospace; font-size: 10.5px; letter-spacing: .1em;
+       text-transform: uppercase; color: var(--muted); font-weight: 400; }
+  td.n, th.n { text-align: right; font-family: "DejaVu Sans Mono", monospace; font-variant-numeric: tabular-nums; }
+  .tag { display: inline-block; font-family: "DejaVu Sans Mono", monospace; font-size: 11px;
+    padding: 1px 7px; border: 1px solid var(--rule); color: var(--muted); background: var(--sand); }
+  img.strip { width: 100%; display: block; border: 1px solid var(--rule); }
+  .caption { font-size: 12.5px; color: var(--muted); margin-top: 7px; }
+  .big { font-family: 'P052', serif; font-size: 21px; line-height: 1.3; }
+  ul { margin: 0; padding-left: 18px; }
+  li { margin-bottom: 5px; }
+  .accent { color: var(--copper); }
+  .teal { color: var(--teal); }
+  svg text { font-family: "DejaVu Sans", sans-serif; }
+  .mono-label { font-family: "DejaVu Sans Mono", monospace; font-size: 12px; }
+</style>
+</head>
+<body>
+
+<!-- 1 ---------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">text-to-image · 177.3M параметров · одна RTX 4080</p>
+  <h1>UDiT — пирамидальный U-трансформер</h1>
+  <p class="lede" style="max-width:88ch">Ширина растёт к низу: 192 → 512 → 1024 → 1536. Свёртки на мелком уровне,
+  внимание на грубом, текст на всех уровнях. Декодер готового автоэнкодера работает хвостом-рендерером,
+  поэтому тело считает 32 числа на клетку 8×8, а не пиксели.</p>
+
+  <div class="rule" style="margin:16px 0 14px"></div>
+  <div class="kpi">
+    <div><b>177.3M</b><span>параметров всего</span></div>
+    <div><b>97 мс</b><span>шаг обучения · bs8 · 256²</span></div>
+    <div><b>20.4</b><span>GFLOP forward у тела</span></div>
+    <div><b>0.647</b><span>детали (hf) против AE 0.859</span></div>
+    <div><b>8.1 ГБ</b><span>пик памяти из 16</span></div>
+  </div>
+
+  <img class="strip" src="img/strip.png" alt="оригинал, реконструкция AE и две генерации" style="margin-top:22px">
+  <p class="caption">слева направо: оригинал · реконструкция AE (потолок 0.859 hf) · две генерации модели на 30k шагов с разными сидами</p>
+
+  <div class="rule" style="margin:16px 0 10px"></div>
+  <div class="grid3">
+    <div><h3>тело · 143.3M</h3><p class="muted">пирамида 192/512/1024/1536, поток в латентном пространстве</p></div>
+    <div><h3>текст · 10.0M</h3><p class="muted">Ettin-150M заморожен, обучается только выбор слоёв</p></div>
+    <div><h3>хвост · 23.9M</h3><p class="muted">декодер AE, дообучается как flow по пикселям</p></div>
+  </div>
+
+  <div class="foot"><span>UDiT · архитектура</span><span>1 / 12</span></div>
+</section>
+
+<!-- 2 ---------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">три части</p>
+  <h2>Модель не считает пиксели — только 32 числа на клетку 8×8</h2>
+
+  <svg width="1150" height="250" viewBox="0 0 1150 250" style="margin-top:16px">
+    <defs>
+      <marker id="a1" markerWidth="10" markerHeight="10" refX="8" refY="5" orient="auto">
+        <path d="M0,0 L10,5 L0,10 z" fill="#14120f"/></marker>
+      <marker id="a2" markerWidth="10" markerHeight="10" refX="8" refY="5" orient="auto">
+        <path d="M0,0 L10,5 L0,10 z" fill="#1f6b66"/></marker>
+    </defs>
+
+    <rect x="0" y="60" width="140" height="92" fill="#fffdf8" stroke="#d9d2c3"/>
+    <text x="14" y="90" font-size="15">шум</text>
+    <text x="14" y="114" class="mono-label" fill="#6d675c">32 × h/8 × w/8</text>
+    <text x="14" y="134" class="mono-label" fill="#6d675c">при 256²: 32×32</text>
+
+    <rect x="196" y="24" width="330" height="164" fill="#f2ece0" stroke="#a8551f"/>
+    <text x="212" y="54" font-size="16" font-weight="bold">ТЕЛО · пирамида</text>
+    <text x="212" y="80" class="mono-label" fill="#6d675c">143.3M · 20.4 GFLOP</text>
+    <text x="212" y="102" class="mono-label" fill="#6d675c">4–8 шагов Эйлера · CFG 1.5</text>
+    <rect x="212" y="116" width="298" height="12" fill="#c98a5f"/><text x="212" y="146" class="mono-label" fill="#7c3d14">÷1 · 192</text>
+    <rect x="212" y="150" width="220" height="12" fill="#a8551f"/><text x="440" y="161" class="mono-label" fill="#7c3d14">÷2 · 512</text>
+    <rect x="212" y="168" width="140" height="12" fill="#7c3d14"/><text x="360" y="179" class="mono-label" fill="#7c3d14">÷4 · 1024 · ÷8 · 1536</text>
+
+    <rect x="582" y="60" width="150" height="92" fill="#fffdf8" stroke="#d9d2c3"/>
+    <text x="596" y="90" font-size="15">латент</text>
+    <text x="596" y="114" class="mono-label" fill="#6d675c">32 × h/8 × w/8</text>
+    <text x="596" y="134" class="mono-label" fill="#6d675c">x0-предсказание</text>
+
+    <rect x="788" y="24" width="230" height="164" fill="#e8f0ee" stroke="#1f6b66"/>
+    <text x="802" y="54" font-size="16" font-weight="bold">ХВОСТ · декодер AE</text>
+    <text x="802" y="80" class="mono-label" fill="#6d675c">23.9M · 4 шага flow</text>
+    <text x="802" y="102" class="mono-label" fill="#6d675c">PixelShuffle ×3 → 3 канала</text>
+    <text x="802" y="124" class="mono-label" fill="#6d675c">условие: z0 + ошибка тела</text>
+    <text x="802" y="146" class="mono-label" fill="#6d675c">обучен заранее, как декодер AE</text>
+    <text x="802" y="168" class="mono-label" fill="#6d675c">дообучается вместе с телом</text>
+
+    <rect x="1074" y="60" width="76" height="92" fill="#fffdf8" stroke="#d9d2c3"/>
+    <text x="1084" y="90" font-size="15">px</text>
+    <text x="1084" y="114" class="mono-label" fill="#6d675c">3×H×W</text>
+
+    <line x1="144" y1="106" x2="192" y2="106" stroke="#14120f" stroke-width="1.8" marker-end="url(#a1)"/>
+    <line x1="530" y1="106" x2="578" y2="106" stroke="#14120f" stroke-width="1.8" marker-end="url(#a1)"/>
+    <line x1="736" y1="106" x2="784" y2="106" stroke="#14120f" stroke-width="1.8" marker-end="url(#a1)"/>
+    <line x1="1022" y1="106" x2="1070" y2="106" stroke="#14120f" stroke-width="1.8" marker-end="url(#a1)"/>
+
+    <rect x="196" y="210" width="822" height="38" fill="#fffdf8" stroke="#1f6b66"/>
+    <text x="210" y="234" font-size="13.5">ТЕКСТ · Ettin-150M (768, заморожен) + SliceMixer 10.04M · кросс-внимание в теле на всех четырёх уровнях</text>
+    <line x1="300" y1="208" x2="300" y2="192" stroke="#1f6b66" stroke-width="1.8" marker-end="url(#a2)"/>
+  </svg>
+
+  <div class="grid3" style="margin-top:14px">
+    <div><h3>Латент вместо пикселей</h3><p class="muted">Один токен = 32 числа на клетку 8×8. При 256² это 1024 токена против 196 608 пикселей.</p></div>
+    <div><h3>Декодер уже обучен</h3><p class="muted">Хвост наследует декодер AE: телу не нужно учиться рисовать текстуру с нуля.</p></div>
+    <div><h3>Текст не экономим</h3><p class="muted">Кросс-внимание в каждом блоке до самого мелкого уровня: +2.5M, «switch» 0.43 → 0.78.</p></div>
+  </div>
+
+  <div class="foot"><span>UDiT · три части</span><span>2 / 12</span></div>
+</section>
+
+<!-- 3 ---------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">почему пирамида</p>
+  <h2>Плоский стек ширины 1024 платит за каждый токен по полной</h2>
+
+  <div class="grid2" style="margin-top:18px; gap:44px">
+    <div>
+      <svg width="530" height="212" viewBox="0 0 530 212">
+        <text x="0" y="12" font-size="13.5" font-weight="bold" fill="#a8551f">было: 12 блоков × ширина 1024</text>
+        <g fill="#e0b39a">
+          <rect x="0" y="26" width="360" height="9"/><rect x="0" y="39" width="360" height="9"/>
+          <rect x="0" y="52" width="360" height="9"/><rect x="0" y="65" width="360" height="9"/>
+          <rect x="0" y="78" width="360" height="9"/><rect x="0" y="91" width="360" height="9"/>
+          <rect x="0" y="104" width="360" height="9"/><rect x="0" y="117" width="360" height="9"/>
+          <rect x="0" y="130" width="360" height="9"/><rect x="0" y="143" width="360" height="9"/>
+        </g>
+        <line x1="372" y1="26" x2="372" y2="152" stroke="#a8551f"/>
+        <text x="380" y="70" class="mono-label" fill="#7c3d14">все блоки</text>
+        <text x="380" y="88" class="mono-label" fill="#7c3d14">на одном</text>
+        <text x="380" y="106" class="mono-label" fill="#7c3d14">разрешении</text>
+        <text x="0" y="176" class="mono-label" fill="#6d675c">токенов на блоке: 1024 · ширина: 1024 · уровней: 1</text>
+        <text x="0" y="196" class="mono-label" fill="#a8551f">232 GFLOP · 196 мс/шаг · hf 0.18</text>
+        <text x="0" y="212" class="mono-label" fill="#6d675c">нет уровня, где вся картинка — это 16 токенов</text>
+      </svg>
+    </div>
+    <div>
+      <svg width="530" height="212" viewBox="0 0 530 212">
+        <text x="0" y="12" font-size="13.5" font-weight="bold" fill="#1f6b66">стало: 192 → 512 → 1024 → 1536</text>
+        <rect x="0" y="26" width="400" height="26" fill="#c98a5f"/>
+        <text x="10" y="43" class="mono-label" fill="#fffdf8">÷1 · 1024 токена · ширина 192 · свёртки</text>
+        <rect x="60" y="58" width="310" height="26" fill="#b06f3f"/>
+        <text x="70" y="75" class="mono-label" fill="#fffdf8">÷2 · 256 токенов · 512</text>
+        <rect x="120" y="90" width="220" height="26" fill="#a8551f"/>
+        <text x="130" y="107" class="mono-label" fill="#fffdf8">÷4 · 64 · 1024</text>
+        <rect x="180" y="122" width="130" height="26" fill="#7c3d14"/>
+        <text x="190" y="139" class="mono-label" fill="#fffdf8">÷8 · 16 · 1536</text>
+        <text x="0" y="176" class="mono-label" fill="#6d675c">энкодер вверх, декодер вниз, аддитивный скип на каждом уровне</text>
+        <text x="0" y="196" class="mono-label" fill="#1f6b66">20.4 GFLOP · 97 мс/шаг · hf 0.54 → 0.65</text>
+        <text x="0" y="212" class="mono-label" fill="#6d675c">широкий низ: 16 токенов видят всю картинку целиком</text>
+      </svg>
+    </div>
+  </div>
+
+  <div class="rule" style="margin:6px 0 10px"></div>
+  <table>
+    <tr><th style="width:32%">метрика</th><th class="n" style="width:13%">плоский 1024</th><th class="n" style="width:12%">пирамида</th><th>что это значит</th></tr>
+    <tr><td>GFLOP на forward (тело)</td><td class="n">232</td><td class="n">20.4</td><td>в 11 раз меньше арифметики при том же числе параметров</td></tr>
+    <tr><td>шаг обучения, bs8 · 256²</td><td class="n">196 мс</td><td class="n">97 мс</td><td>вдвое больше шагов за то же время</td></tr>
+    <tr><td>детали: hf(gen)/hf(AE)</td><td class="n">21%</td><td class="n">51–75%</td><td>и это при 3–6 кратном меньшем числе шагов обучения</td></tr>
+  </table>
+
+  <div class="rule" style="margin:10px 0 12px"></div>
+  <p class="big">Стоимость слоя ≈ токены × ширина². Мы держим ширину там, где токенов мало, и снимаем
+  с мелкого уровня 28-кратный перерасход — эти же деньги уходят в глубину и разрешение.</p>
+
+  <div class="foot"><span>UDiT · сравнение</span><span>3 / 12</span></div>
+</section>
+
+<!-- 4 ---------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">��естница</p>
+  <h2>Ширина растёт там, где токенов мало — стоимость уровней сбалансирована</h2>
+
+  <div class="grid2" style="margin-top:12px; grid-template-columns: 1.25fr 1fr;">
+    <div>
+      <table>
+        <tr><th>уровень</th><th class="n">токенов 256²</th><th class="n">токенов 512²</th><th class="n">ширина</th><th>блоки</th><th>механика</th></tr>
+        <tr><td>L0 ÷1</td><td class="n">1024</td><td class="n">4096</td><td class="n">192</td><td>6 enc / 4 dec + self-attn</td><td>свёртки + текст</td></tr>
+        <tr><td>L1 ÷2</td><td class="n">256</td><td class="n">1024</td><td class="n">512</td><td>4 / 3</td><td>трансформер</td></tr>
+        <tr><td>L2 ÷4</td><td class="n">64</td><td class="n">256</td><td class="n">1024</td><td>3 / 2</td><td>трансформер</td></tr>
+        <tr><td>L3 ÷8</td><td class="n">16</td><td class="n">64</td><td class="n">1536</td><td>2 (боттлнек)</td><td>трансформер</td></tr>
+      </table>
+      <p class="caption">энкодер поднимается по лестнице, декодер спускается; на каждом уровне прибавляется скип</p>
+    </div>
+    <div>
+      <svg width="470" height="290" viewBox="0 0 470 290">
+        <text x="0" y="14" font-size="12" fill="#6d675c">стоимость уровня ≈ токены × ширина²</text>
+        <g>
+          <rect x="0" y="30" width="240" height="26" fill="#c98a5f"/><text x="250" y="48" class="mono-label">37.7 · L0</text>
+          <rect x="0" y="66" width="427" height="26" fill="#b06f3f"/><text x="437" y="84" class="mono-label">67.1 · L1</text>
+          <rect x="0" y="102" width="427" height="26" fill="#a8551f"/><text x="437" y="120" class="mono-label">67.1 · L2</text>
+          <rect x="0" y="138" width="240" height="26" fill="#7c3d14"/><text x="250" y="156" class="mono-label">37.7 · L3</text>
+        </g>
+        <line x1="0" y1="182" x2="430" y2="182" stroke="#d9d2c3"/>
+        <text x="0" y="204" font-size="12" fill="#a8551f">для сравнения: 1024 на всех уровнях</text>
+        <rect x="0" y="214" width="470" height="26" fill="#e7c8b2"/>
+        <text x="8" y="232" class="mono-label" fill="#7c3d14">1073 · L0 при ширине 1024 (в 28 раз дороже)</text>
+        <text x="0" y="272" class="mono-label" fill="#6d675c">параметры уезжают вниз, где токенов 16, а не 1024</text>
+      </svg>
+    </div>
+  </div>
+
+  <div class="foot"><span>UDiT · бюджет уровней</span><span>4 / 12</span></div>
+</section>
+
+<!-- 5 ---------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">роли уровней</p>
+  <h2>Каждый уровень отвечает за свой масштаб — как у художника</h2>
+
+  <div class="grid4" style="margin-top:18px">
+    <div class="card">
+      <svg width="100%" height="96" viewBox="0 0 220 96">
+        <rect x="18" y="14" width="184" height="68" fill="none" stroke="#d9d2c3"/>
+        <circle cx="110" cy="48" r="26" fill="#a8551f" opacity="0.22"/>
+        <circle cx="110" cy="48" r="12" fill="#a8551f" opacity="0.55"/>
+        <text x="70" y="92" class="mono-label" fill="#6d675c">16 токенов</text>
+      </svg>
+      <h3>L3 · «прищур»</h3>
+      <p class="muted">Вся картинка — 16 токенов. Полное внимание: все видят всех. Здесь решается композиция и симметрия.</p>
+    </div>
+    <div class="card">
+      <svg width="100%" height="96" viewBox="0 0 220 96">
+        <rect x="18" y="14" width="184" height="68" fill="none" stroke="#d9d2c3"/>
+        <path d="M60,70 Q110,10 160,70" fill="none" stroke="#1f6b66" stroke-width="3"/>
+        <text x="60" y="92" class="mono-label" fill="#6d675c">64 токена</text>
+      </svg>
+      <h3>L2 · крупные формы</h3>
+      <p class="muted">Тело, крылья, общая раскраска. Ширина 1024 — запас «смысла» на грубое описание.</p>
+    </div>
+    <div class="card">
+      <svg width="100%" height="96" viewBox="0 0 220 96">
+        <rect x="18" y="14" width="184" height="68" fill="none" stroke="#d9d2c3"/>
+        <path d="M40,60 C70,20 150,20 180,60" fill="none" stroke="#1f6b66" stroke-width="2"/>
+        <path d="M40,70 C70,34 150,34 180,70" fill="none" stroke="#1f6b66" stroke-width="2" opacity="0.6"/>
+        <text x="60" y="92" class="mono-label" fill="#6d675c">256 токенов</text>
+      </svg>
+      <h3>L1 · средние детали</h3>
+      <p class="muted">Форма крыла, крупный рисунок, где тёмное и где светлое. Ширина 512.</p>
+    </div>
+    <div class="card">
+      <svg width="100%" height="96" viewBox="0 0 220 96">
+        <rect x="18" y="14" width="184" height="68" fill="none" stroke="#d9d2c3"/>
+        <g stroke="#a8551f" stroke-width="1" opacity="0.8">
+          <line x1="30" y1="30" x2="46" y2="46"/><line x1="70" y1="26" x2="86" y2="42"/>
+          <line x1="110" y1="30" x2="126" y2="46"/><line x1="150" y1="26" x2="166" y2="42"/>
+          <line x1="40" y1="60" x2="56" y2="76"/><line x1="90" y1="58" x2="106" y2="74"/>
+          <line x1="140" y1="60" x2="156" y2="76"/>
+        </g>
+        <text x="56" y="92" class="mono-label" fill="#6d675c">1024 токена</text>
+      </svg>
+      <h3>L0 · текстура</h3>
+      <p class="muted">Жилки, кромки, крапинки. Здесь свёртка 3×3, а не внимание: текстура локальна, а свёртка в 20 раз дешевле.</p>
+    </div>
+  </div>
+
+  <div class="rule"></div>
+  <p class="big">Внимание на грубом уровне стоит копейки: 16 токенов × 16 = 256 пар. Внимание на мелком — 1024 × 1024 = миллион пар. Поэтому «прищур» мы делаем вниманием, а текстуру — свёрткой.</p>
+
+  <div class="foot"><span>UDiT · роли уровней</span><span>5 / 12</span></div>
+</section>
+
+<!-- 6 ---------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">путь и скипы</p>
+  <h2>Энкодер поднимается к смыслу, декодер возвращается к деталям — помня эскиз</h2>
+
+  <svg width="1150" height="300" viewBox="0 0 1150 300" style="margin-top:14px">
+    <defs>
+      <marker id="up" markerWidth="10" markerHeight="10" refX="8" refY="5" orient="auto">
+        <path d="M0,0 L10,5 L0,10 z" fill="#a8551f"/></marker>
+      <marker id="dn" markerWidth="10" markerHeight="10" refX="8" refY="5" orient="auto">
+        <path d="M0,0 L10,5 L0,10 z" fill="#1f6b66"/></marker>
+    </defs>
+
+    <g>
+      <rect x="0" y="196" width="210" height="66" fill="#fffdf8" stroke="#d9d2c3"/>
+      <text x="14" y="222" font-size="13.5">L0 · 1024 токена · 192</text>
+      <text x="14" y="244" class="mono-label" fill="#6d675c">6 блоков enc · 4 dec</text>
+
+      <rect x="272" y="140" width="210" height="66" fill="#fffdf8" stroke="#d9d2c3"/>
+      <text x="286" y="166" font-size="13.5">L1 · 256 · 512</text>
+      <text x="286" y="188" class="mono-label" fill="#6d675c">4 enc · 3 dec</text>
+
+      <rect x="544" y="84" width="210" height="66" fill="#fffdf8" stroke="#d9d2c3"/>
+      <text x="558" y="110" font-size="13.5">L2 · 64 · 1024</text>
+      <text x="558" y="132" class="mono-label" fill="#6d675c">3 enc · 2 dec</text>
+
+      <rect x="816" y="28" width="210" height="66" fill="#fffdf8" stroke="#d9d2c3"/>
+      <text x="830" y="54" font-size="13.5">L3 · 16 · 1536</text>
+      <text x="830" y="76" class="mono-label" fill="#6d675c">2 блока · полное внимание</text>
+    </g>
+
+    <g stroke="#a8551f" stroke-width="2" fill="none">
+      <line x1="212" y1="228" x2="266" y2="180" marker-end="url(#up)"/>
+      <line x1="484" y1="172" x2="538" y2="124" marker-end="url(#up)"/>
+      <line x1="756" y1="116" x2="810" y2="68" marker-end="url(#up)"/>
+    </g>
+    <text x="206" y="212" class="mono-label" fill="#a8551f">pool 2×</text>
+    <text x="478" y="156" class="mono-label" fill="#a8551f">pool 2×</text>
+    <text x="750" y="100" class="mono-label" fill="#a8551f">pool 2×</text>
+
+    <g stroke="#1f6b66" stroke-width="2" fill="none">
+      <line x1="866" y1="98" x2="792" y2="146" marker-end="url(#dn)"/>
+      <line x1="594" y1="154" x2="520" y2="202" marker-end="url(#dn)"/>
+      <line x1="322" y1="210" x2="248" y2="258" marker-end="url(#dn)"/>
+    </g>
+    <text x="880" y="120" class="mono-label" fill="#1f6b66">up 2×</text>
+    <text x="608" y="176" class="mono-label" fill="#1f6b66">up 2×</text>
+    <text x="336" y="232" class="mono-label" fill="#1f6b66">up 2×</text>
+
+    <rect x="0" y="286" width="1026" height="0" fill="none"/>
+    <text x="0" y="292" class="mono-label" fill="#a8551f">скип: на каждом уровне выход энкодера прибавляется ко входу декодера того же уровня (аддитивно, без проекций)</text>
+  </svg>
+
+  <div class="grid3" style="margin-top:10px">
+    <div><h3>Почему аддитивные</h3><p class="muted">Ширина входа и выхода уровня совпадает, проекции не нужны — так делают U-Net’ы SD.</p></div>
+    <div><h3>Почему pool, а не stride</h3><p class="muted">Среднее 2×2 плюс линейный слой: сжатие учится и не ломает фазу позиционного кодирования.</p></div>
+    <div><h3>Что это даёт</h3><p class="muted">Тонкая структура строится на своём уровне, а не вытягивается ��ниманием с одного масштаба.</p></div>
+  </div>
+
+  <div class="foot"><span>UDiT · путь и скипы</span><span>6 / 12</span></div>
+</section>
+
+<!-- 7 ---------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">текст</p>
+  <h2>Подпись не «берётся из последнего слоя» — модель сама выбирает, какие слои читать</h2>
+
+  <div class="grid2" style="margin-top:14px; grid-template-columns: 1.15fr 1fr;">
+    <div>
+      <svg width="560" height="252" viewBox="0 0 560 252">
+        <text x="0" y="12" font-size="11.5" fill="#6d675c">норма растёт с глубиной в 10 раз — без пер-слойной нормировки ранние слои не влияют</text>
+        <g class="mono-label" fill="#6d675c">
+          <text x="0" y="42">слой 2</text><text x="0" y="66">слой 5</text><text x="0" y="90">слой 8</text>
+          <text x="0" y="114">слой 11</text><text x="0" y="138">слой 14</text><text x="0" y="162">слой 17</text>
+          <text x="0" y="186">слой 20</text>
+        </g>
+        <g fill="#c98a5f">
+          <rect x="70" y="30" width="44" height="12"/><rect x="70" y="54" width="47" height="12"/>
+          <rect x="70" y="78" width="58" height="12"/><rect x="70" y="102" width="84" height="12"/>
+          <rect x="70" y="126" width="187" height="12"/><rect x="70" y="150" width="346" height="12"/>
+        </g>
+        <rect x="70" y="174" width="403" height="12" fill="#a8551f"/>
+        <text x="120" y="26" class="mono-label" fill="#7c3d14">52</text>
+        <text x="480" y="170" class="mono-label" fill="#7c3d14">548</text>
+        <rect x="0" y="206" width="560" height="24" fill="#e8f0ee" stroke="#1f6b66"/>
+        <text x="10" y="223" class="mono-label" fill="#1f6b66">7 срезов → LN на каждый → внимание по срезам → Linear(7→1) → + среднее</text>
+        <text x="0" y="248" class="mono-label" fill="#6d675c">«проектор слоёв»: какие слои читать, решает обучение, а не догадка</text>
+      </svg>
+      <p class="caption">SliceMixer · 10.04M обучаемых · идея из moss/adapter (0.8B + 158M голова вместо 9B-энкодера)</p>
+    </div>
+    <div>
+      <div class="card" style="margin-bottom:14px">
+        <h3>Два входа текста в каждый блок</h3>
+        <ul class="muted">
+          <li>кросс-внимание: каждый токен-клетка спрашивает «что в подписи про моё место?»</li>
+          <li>модуляция: общий вектор «t + усреднённый текст» настраивает shift/scale/gate фильтров</li>
+        </ul>
+      </div>
+      <table>
+        <tr><th>что сделали</th><th class="n">цена</th><th class="n">эффект</th></tr>
+        <tr><td>кросс-внимание и на L0</td><td class="n">+2.5M</td><td class="n">+9 мс</td></tr>
+        <tr><td>«switch» (один шум, два капшена)</td><td class="n">—</td><td class="n">0.43 → 0.78</td></tr>
+        <tr><td>детали (hf) не пострадали</td><td class="n">—</td><td class="n">0.44 → 0.65</td></tr>
+      </table>
+      <p class="caption">«switch» — насколько латент меняется от смены подписи при том же шуме</p>
+    </div>
+  </div>
+
+  <div class="foot"><span>UDiT · текст</span><span>7 / 12</span></div>
+</section>
+
+<!-- 8 ---------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">хвост и сэмплинг</p>
+  <h2>Тело ведёт латент, готовый декодер рисует пиксели</h2>
+
+  <svg width="1140" height="240" viewBox="0 0 1140 240" style="margin-top:12px">
+    <defs>
+      <marker id="a3" markerWidth="9" markerHeight="9" refX="7" refY="4.5" orient="auto">
+        <path d="M0,0 L9,4.5 L0,9 z" fill="#14120f"/></marker>
+    </defs>
+    <g>
+      <circle cx="40" cy="70" r="18" fill="#fffdf8" stroke="#d9d2c3"/><text x="34" y="75" class="mono-label">шум</text>
+      <line x1="60" y1="70" x2="120" y2="70" stroke="#14120f" marker-end="url(#a3)"/>
+      <rect x="126" y="44" width="180" height="52" fill="#f2ece0" stroke="#a8551f"/>
+      <text x="138" y="66" font-size="12">тело: x0-предсказание</text>
+      <text x="138" y="84" class="mono-label" fill="#6d675c">8 шагов · CFG 1.5</text>
+      <line x1="310" y1="70" x2="370" y2="70" stroke="#14120f" marker-end="url(#a3)"/>
+      <rect x="376" y="44" width="150" height="52" fill="#fffdf8" stroke="#d9d2c3"/>
+      <text x="388" y="66" font-size="12">латент z</text>
+      <text x="388" y="84" class="mono-label" fill="#6d675c">32 × h/8 × w/8</text>
+      <line x1="530" y1="70" x2="590" y2="70" stroke="#14120f" marker-end="url(#a3)"/>
+      <rect x="596" y="44" width="190" height="52" fill="#e8f0ee" stroke="#1f6b66"/>
+      <text x="608" y="66" font-size="12">хвост: декодер AE</text>
+      <text x="608" y="84" class="mono-label" fill="#6d675c">4 шага flow</text>
+      <line x1="790" y1="70" x2="850" y2="70" stroke="#14120f" marker-end="url(#a3)"/>
+      <rect x="856" y="44" width="130" height="52" fill="#fffdf8" stroke="#d9d2c3"/>
+      <text x="868" y="66" font-size="12">пиксели</text>
+      <text x="868" y="84" class="mono-label" fill="#6d675c">3 × H × W</text>
+    </g>
+  </svg>
+
+  <div class="rule" style="margin:6px 0 10px"></div>
+  <table style="table-layout:fixed">
+    <tr><th style="width:30%">на каком латенте учить хвост</th><th style="width:18%">что видит рендерер</th><th style="width:52%">что происходит</th></tr>
+    <tr><td>только z0 (чистый латент)</td><td>идеальный латент</td><td>на сэмплинге встречает ошибку тела, к которой не привык</td></tr>
+    <tr><td>только предсказание тела</td><td>мусор в начале обучения</td><td>учится «исправлять мусор» и разучивается рисовать: 29.6 → 21.3 дБ</td></tr>
+    <tr><td><b>z0 + (ошибка тела).detach()</b></td><td>реальную ошибку</td><td>видит промах, но градиент через него не идёт; 24.7–25.9 дБ стабильно</td></tr>
+  </table>
+  <p class="caption">Сэмпл: 8 шагов тела + 4 шага декодера ≈ 35 мс на 4080 при 256². CFG 1.5 — выше 2 при недоученной
+  безусловной ветви ухудшает латент (замерено 0.19 → 0.28). Метрика tail(z0) в каждом логе ловит деградацию рендерера.</p>
+
+  <div class="foot"><span>UDiT · хвост</span><span>8 / 12</span></div>
+</section>
+
+<!-- 9 ---------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">результат на узком домене</p>
+  <h2>Бабочки: детали растут с шагами и не насыщаются</h2>
+
+  <div class="grid2" style="margin-top:10px; grid-template-columns: 1fr 1.1fr;">
+    <div>
+      <table>
+        <tr><th>шаги</th><th class="n">hf / AE</th><th class="n">mse_lat</th><th class="n">high</th><th class="n">switch</th></tr>
+        <tr><td>4 000</td><td class="n">0.446 · 52%</td><td class="n">0.21</td><td class="n">0.42</td><td class="n">0.43</td></tr>
+        <tr><td>10 000</td><td class="n">0.436 · 51%</td><td class="n">0.164</td><td class="n">0.32</td><td class="n">0.68</td></tr>
+        <tr><td>20 000</td><td class="n">0.540 · 63%</td><td class="n">0.107</td><td class="n">0.19</td><td class="n">0.78</td></tr>
+        <tr><td><b>30 000</b></td><td class="n"><b>0.647 · 75%</b></td><td class="n"><b>0.072</b></td><td class="n"><b>0.114</b></td><td class="n">0.78</td></tr>
+      </table>
+      <p class="caption">hf — энергия высоких частот (лапласиан). AE-реконструкция того же вида: 0.859, оригинал 0.913.
+      «high» — ошибка на высоком шуме, то есть «какую особь рисовать»; она упала в 3.7 раза.</p>
+      <p class="big" style="margin-top:10px">Три сида дают три разные особи: 0.642 / 0.646 / 0.654 — модель не сколлапсировала в один мод.</p>
+    </div>
+    <div>
+      <img class="strip" src="img/samples.png" alt="Четыре вида бабочек, сгенерированные моделью">
+      <p class="caption">генерации на 30k шагов: Delias, Cepora, Troides, Catasticta — вид задаётся подписью из 16 токенов</p>
+    </div>
+  </div>
+
+  <div class="foot"><span>UDiT · результат</span><span>9 / 12</span></div>
+</section>
+
+<!-- 10 --------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">цена</p>
+  <h2>177.3M параметров, 97 мс на шаг, 8.1 ГБ — обучение на одной карте</h2>
+
+  <div class="grid3" style="margin-top:16px">
+    <div>
+      <h3>Параметры по частям</h3>
+      <svg width="330" height="170" viewBox="0 0 330 170">
+        <rect x="0" y="20" width="240" height="24" fill="#a8551f"/><text x="248" y="37" class="mono-label">143.3 M · тело</text>
+        <rect x="0" y="54" width="17" height="24" fill="#1f6b66"/><text x="25" y="71" class="mono-label">10.0 M · текст</text>
+        <rect x="0" y="88" width="40" height="24" fill="#6d675c"/><text x="48" y="105" class="mono-label">23.9 M · хвост</text>
+        <line x1="0" y1="132" x2="300" y2="132" stroke="#d9d2c3"/>
+        <text x="0" y="152" class="mono-label" fill="#6d675c">всего 177.3M · чекпойнт 709 МБ</text>
+      </svg>
+    </div>
+    <div>
+      <h3>Арифметика на forward</h3>
+      <svg width="330" height="170" viewBox="0 0 330 170">
+        <rect x="0" y="20" width="20" height="24" fill="#1f6b66"/><text x="28" y="37" class="mono-label">20.4 · пирамида</text>
+        <rect x="0" y="54" width="230" height="24" fill="#e7c8b2"/><text x="238" y="71" class="mono-label" fill="#7c3d14">232 · плоский 1024</text>
+        <text x="0" y="106" class="mono-label" fill="#6d675c">+ хвост ≈ 29 GFLOP</text>
+        <text x="0" y="126" class="mono-label" fill="#6d675c">итого ≈ 50 против ≈ 262</text>
+        <text x="0" y="152" class="mono-label" fill="#6d675c">та же карта, в 5 раз больше шагов</text>
+      </svg>
+    </div>
+    <div>
+      <h3>Скорость и память</h3>
+      <svg width="330" height="170" viewBox="0 0 330 170">
+        <rect x="0" y="20" width="110" height="24" fill="#a8551f"/><text x="118" y="37" class="mono-label">97 мс · пирамида</text>
+        <rect x="0" y="54" width="222" height="24" fill="#e7c8b2"/><text x="230" y="71" class="mono-label" fill="#7c3d14">196 мс · плоский</text>
+        <text x="0" y="106" class="mono-label" fill="#6d675c">8.1 ГБ из 16 (bs8, 256²)</text>
+        <text x="0" y="126" class="mono-label" fill="#6d675c">8-битный AdamW: −1 ГБ состояний</text>
+        <text x="0" y="152" class="mono-label" fill="#6d675c">эффективность ~1.6 TFLOPS — упор в память</text>
+      </svg>
+    </div>
+  </div>
+
+  <div class="rule"></div>
+  <div class="grid2">
+    <div class="card">
+      <h3>Почему хватает 16 ГБ</h3>
+      <p class="muted">Память ест не модель (2.5 ГБ весов и состояний), а активации: они ~ токены × ширина × блоки.
+      Пирамида держит широкие слои на 16–64 токенах, поэтому активации маленькие.</p>
+    </div>
+    <div class="card">
+      <h3>Где запас</h3>
+      <p class="muted">Шаг упирается в память, а не в счёт. Запас есть на батч 16 (14.8 ГБ с хвостом) и на 512²
+      с батчем 1–2 — то есть на разрешение, где деталь физически возможна.</p>
+    </div>
+  </div>
+
+  <div class="foot"><span>UDiT · цена</span><span>10 / 12</span></div>
+</section>
+
+<!-- 11 --------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">следующий прогон</p>
+  <h2>FLUX-Reason: 2 эпохи по 1.5M картинок, поток без кэша</h2>
+
+  <div class="grid2" style="margin-top:14px">
+    <div>
+      <table>
+        <tr><th>параметр</th><th>значение</th></tr>
+        <tr><td>датасет</td><td>LucasFang/FLUX-Reason-6M, часть Aesthetics-Part01 (415 шардов, ~250 ГБ)</td></tr>
+        <tr><td>разрешение / батч</td><td>256² · bs16</td></tr>
+        <tr><td>подпись</td><td>caption_entity → caption_composition (40–400 симв.), иначе caption_detail</td></tr>
+        <tr><td>бюджет</td><td>2 эпохи = 187 500 шагов ≈ 12.5 ч</td></tr>
+        <tr><td>оптимизатор</td><td>AdamW8bit, прогрев 500, cosine до 10%</td></tr>
+        <tr><td>дропауты</td><td>подпись 10% (CFG) · слова 5%</td></tr>
+        <tr><td>старт</td><td>веса бабочек (30k шагов)</td></tr>
+      </table>
+    </div>
+    <div>
+      <h3>Что уже видно (шаг 800)</h3>
+      <svg width="480" height="150" viewBox="0 0 480 150">
+        <text x="0" y="18" class="mono-label" fill="#6d675c">mse_lat общий и по зонам шума</text>
+        <line x1="0" y1="120" x2="470" y2="120" stroke="#d9d2c3"/>
+        <rect x="10" y="60" width="60" height="60" fill="#a8551f"/><text x="18" y="140" class="mono-label">0.33</text>
+        <rect x="110" y="113" width="60" height="7" fill="#1f6b66"/><text x="120" y="140" class="mono-label">0.037</text>
+        <rect x="210" y="120" width="60" height="0.5" fill="#6d675c"/>
+        <text x="330" y="70" class="mono-label" fill="#6d675c">low (t&lt;0.3): 0.037</text>
+        <text x="330" y="90" class="mono-label" fill="#6d675c">high (t&gt;0.7): 0.70</text>
+        <text x="330" y="110" class="mono-label" fill="#6d675c">241 мс/шаг · 66 картинок/с</text>
+        <text x="330" y="130" class="mono-label" fill="#6d675c">9.5 ГБ · skip 0</text>
+      </svg>
+      <p class="caption">«high» здесь в 6 раз выше, чем на бабочках: 1.5M разных сцен, «какую именно рисовать»
+      по подписи не угадать — это нормальный старт, а не поломка.</p>
+    </div>
+  </div>
+
+  <div class="rule"></div>
+  <div class="grid3">
+    <div><h3>Смотреть</h3><p class="muted">tmp/flux/log.csv · tmp/flux/samples/latest.png (лиса, стеклянный город, неон, импрессионист)</p></div>
+    <div><h3>Остановить / продолжить</h3><p class="muted">kill процесса · тот же запуск с --resume --out tmp/flux</p></div>
+    <div><h3>Честно про «2 эпохи»</h3><p class="muted">эпоха по Part01 = 1.5M картинок; все 4 части ≈ 1 ТБ, а на диске 352 ГБ</p></div>
+  </div>
+
+  <div class="foot"><span>UDiT · FLUX-Reason</span><span>11 / 12</span></div>
+</section>
+
+<!-- 12 --------------------------------------------------------------- -->
+<section class="slide">
+  <p class="eyebrow">правила</p>
+  <h2>Всё, что ниже — вывод из замеров, а не из вкуса</h2>
+
+  <div class="grid2" style="margin-top:16px; gap:26px 40px">
+    <div>
+      <p><span class="tag">1</span> <b>Ширина растёт к низу.</b> <span class="muted">Плоский 1024 на всех уровнях: 196 мс и hf 0.18. Пирамида: 97 мс и 0.54.</span></p>
+      <p><span class="tag">2</span> <b>Текст не экономим.</b> <span class="muted">Кросс-внимание на самом мелком уровне: +2.5M и +9 мс, «switch» 0.43 → 0.78.</span></p>
+      <p><span class="tag">3</span> <b>Хвост учится на z0 + ошибка тела.</b> <span class="muted">Только z0 — не видит реальных ошибок; только предсказания тела — деградация 29.6 → 21.3 дБ.</span></p>
+    </div>
+    <div>
+      <p><span class="tag">4</span> <b>Бакеты кратны 64.</b> <span class="muted">Три даунсэмпла требуют латента, кратного 8. Аспект — центр-кроп, не растяжение.</span></p>
+      <p><span class="tag">5</span> <b>Метрики вместо впечатлений.</b> <span class="muted">mse_lat (общий, low, high), mse_pix, switch, tail(z0), hf(gen)/hf(AE), skip.</span></p>
+      <p><span class="tag">6</span> <b>Никаких костылей.</b> <span class="muted">Фиксированный шум, клип латента, «mix» латентов — выброшены: они прячут то, что мы измеряем.</span></p>
+    </div>
+  </div>
+
+  <div class="rule"></div>
+  <div class="grid4">
+    <div><h3>Дальше</h3><p class="muted">дожать FLUX-Reason до 2 эпох по Part01</p></div>
+    <div><h3>Потом</h3><p class="muted">люди: 5749 фото, 256² — лица и кисти</p></div>
+    <div><h3>Затем</h3><p class="muted">512² на людях: деталь физически разрешима</p></div>
+    <div><h3>Только потом</h3><p class="muted">ёмкость: ширина боттлнека и число блоков L0</p></div>
+  </div>
+
+  <div class="foot"><span>UDiT · правила</span><span>12 / 12</span></div>
+</section>
+
+</body>
+</html>

src/export.py

@@ -0,0 +1,127 @@
+"""Экспорт обученного UDM в `transformer/`: safetensors + config.json.
+
+Три части в одном файле, по префиксам:
+  body.*        — тело (8 блоков из энкодера AE + 4 новых), поток в латентном пространстве
+  text_mixer.*  — обучаемый выбор слоёв Ettin (SliceMixer)
+  tail.*        — декодер AE (рендер латента в пиксели), дообученный как обычный flow
+
+  python3 src/export.py --ckpt tmp/run256/last.pt --out transformer
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from safetensors.torch import save_file
+
+ROOT = Path(__file__).resolve().parent.parent
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default=str(ROOT / "tmp" / "run256" / "last.pt"))
+    ap.add_argument("--out", default=str(ROOT / "transformer"))
+    ap.add_argument("--metrics", default="", help="дописать метрики в config.json (строка)")
+    ap.add_argument("--half", action="store_true", help="хранить веса в fp16 (вдвое меньше, для публикации)")
+    args = ap.parse_args()
+
+    ck = torch.load(args.ckpt, map_location="cpu")
+    tensors = {}
+    for name, sd in (("body", ck["body"]), ("text_mixer", ck["mixer"]), ("tail", ck.get("decoder") or {})):
+        for k, v in sd.items():
+            tensors[f"{name}.{k}"] = (v.half() if args.half else v).contiguous()
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    save_file(tensors, str(out / "diffusion_pytorch_model.safetensors"))
+
+    ae_cfg = json.loads((ROOT / "vae" / "config.json").read_text())
+    ta = ck.get("args") or {}          # чекпойнт помнит свои аргументы — не сочиняем заново
+    training = {
+        "dataset": ta.get("hf_dataset") or ta.get("data") or "~/dataset/butterflies_256x64",
+        "parts": ta.get("parts"),
+        "size": ta.get("size"),
+        "steps": int(ck["step"]),
+        "batch": ta.get("batch", 8),
+        "accum": ta.get("accum", 1),
+        "lr": ta.get("lr", 1e-4),
+        "lr_tail": ta.get("lr_tail", 3e-5),
+        "drop_text": ta.get("drop_text", 0.1),
+        "init": ta.get("init"),
+        "metrics": args.metrics,
+    }
+    cfg = {
+        "_class_name": "UDiT",
+        "name": "Ra",
+        "format": "udm/v1",
+        "step": int(ck["step"]),
+        "latent": {
+            "channels": ae_cfg["latent_channels"],
+            "spatial_factor": ae_cfg["patch_size"],
+            "mean": ae_cfg["latent_mean"],
+            "std": ae_cfg["latent_std"],
+            "source": "vae/ (Muinez/f8c32-ae)",
+        },
+        "text": {
+            "encoder": "text_encoder/ (jhu-clsp/ettin-encoder-150m, 768-dim, 22 слоя)",
+            "max_len": 192,
+            "slices": [2, 5, 8, 11, 14, 17, 20],
+            "slice_blocks": 2,
+            "trainable": "text_mixer.*",
+        },
+        "body": {
+            "kind": "пирамидальный U: ширина растёт к низу",
+            "levels": [
+                {"level": 0, "tokens": "÷1", "width": 192, "blocks": "conv 6 enc / 4 dec + 1 self-attn на сторону",
+                 "text": "кросс-внимание в каждом conv-блоке и в self-attn"},
+                {"level": 1, "tokens": "÷2", "width": 512, "blocks": "attn 4 enc / 3 dec", "text": "кросс-внимание"},
+                {"level": 2, "tokens": "÷4", "width": 1024, "blocks": "attn 3 enc / 2 dec", "text": "кросс-внимание"},
+                {"level": 3, "tokens": "÷8", "width": 1536, "blocks": "attn 2 (боттлнек)", "text": "кросс-внимание"},
+            ],
+            "skips": "аддитивные (SD-стиль): e0/e1/e2 прибавляются в декодере",
+            "down": "mean-pool 2× + Linear",
+            "up": "Linear + nearest 2×",
+            "rope": "2D RoPE, base 100, фаза в латентных пикселях на каждом уровне",
+            "parametrization": "x0 латента",
+        },
+        "tail": {
+            "source": "декодер AE, 6 блоков @512, PixelShuffle ×3",
+            "trained_as": "flow по пикселям, условие z0 + (ошибка тела).detach()",
+            "sample_steps": 4,
+        },
+        "params": {
+            "body": round(sum(v.numel() for k, v in tensors.items() if k.startswith("body.")) / 1e6, 2),
+            "text_mixer": round(sum(v.numel() for k, v in tensors.items()
+                                    if k.startswith("text_mixer.")) / 1e6, 2),
+            "tail": round(sum(v.numel() for k, v in tensors.items() if k.startswith("tail.")) / 1e6, 2),
+        },
+        "training": {k: v for k, v in training.items() if v is not None},
+        "storage_dtype": "float16" if args.half else "float32",
+    }
+    cfg["params"]["total"] = round(sum(cfg["params"].values()), 2)
+    (out / "config.json").write_text(json.dumps(cfg, indent=2, ensure_ascii=False))
+
+    print(f"сохранено: {out}/diffusion_pytorch_model.safetensors "
+          f"({sum(v.numel() for v in tensors.values())/1e6:.1f}M параметров, "
+          f"{(out / 'diffusion_pytorch_model.safetensors').stat().st_size/1e6:.0f} МБ)")
+    print(f"части: {cfg['params']}")
+
+
+def load_exported(model, enc, ae, path, device=None):
+    """Загрузить из `transformer/` (config.json + safetensors с префиксами body./text_mixer./tail.)."""
+    import json
+    from safetensors.torch import load_file
+    path = Path(path)
+    sd = load_file(str(path / "diffusion_pytorch_model.safetensors"), device=str(device) if device else None)
+    def sub(prefix):
+        return {k[len(prefix) + 1:]: v for k, v in sd.items() if k.startswith(prefix + ".")}
+    model.body.load_state_dict(sub("body"), strict=False)
+    enc.mixer.load_state_dict(sub("text_mixer"))
+    ae.decoder.load_state_dict(sub("tail"))
+    cfg_path = path / "config.json"
+    return json.loads(cfg_path.read_text()) if cfg_path.exists() else {}
+
+
+if __name__ == "__main__":
+    main()

src/make_dataset.py

@@ -0,0 +1,94 @@
+"""Сборка набора из папки «картинка + .txt»: бакеты (кратны --multiple) и центр-кроп.
+
+  python3 src/make_dataset.py --src ~/dataset/butterflies/butterfly \
+          --out ~/dataset/butterflies_256 --long-side 256 --eval-n 30
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import math
+import random
+from pathlib import Path
+
+from PIL import Image, ImageOps
+
+BUCKETS = {}
+
+
+def buckets(long_side: int, multiple: int = 16) -> dict:
+    L = long_side
+    def r16(x): return max(multiple, int(round(x / multiple)) * multiple)
+    return {
+        "1:1": (L, L),
+        "4:3": (L, r16(L * 3 / 4)),
+        "3:4": (r16(L * 3 / 4), L),
+        "3:2": (L, r16(L * 2 / 3)),
+        "2:3": (r16(L * 2 / 3), L),
+        "16:9": (L, r16(L * 9 / 16)),
+        "9:16": (r16(L * 9 / 16), L),
+    }
+
+
+def plan_size(w, h, max_upscale=1.4):
+    ratio = w / h
+    label, (tw, th) = min(BUCKETS.items(), key=lambda kv: abs(math.log(ratio / (kv[1][0] / kv[1][1]))))
+    if ratio > tw / th:
+        cw, ch = round(h * tw / th), h
+    else:
+        cw, ch = w, round(w * th / tw)
+    cw, ch = min(cw, w), min(ch, h)
+    if cw * max_upscale < tw or ch * max_upscale < th:
+        return None
+    return label, tw, th, cw, ch
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--src", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--long-side", type=int, default=256)
+    ap.add_argument("--eval-n", type=int, default=30)
+    ap.add_argument("--max-upscale", type=float, default=1.4)
+    ap.add_argument("--multiple", type=int, default=64,
+                    help="кратность сторон: три даунсэмпла тела требуют латент, кратный 8")
+    ap.add_argument("--seed", type=int, default=0)
+    args = ap.parse_args()
+
+    global BUCKETS
+    BUCKETS = buckets(args.long_side, args.multiple)
+    src, out = Path(args.src), Path(args.out)
+    for sub in ("train", "eval"):
+        (out / sub).mkdir(parents=True, exist_ok=True)
+
+    pairs = []
+    for img in sorted(list(src.glob("*.png")) + list(src.glob("*.jpg"))):
+        txt = img.with_suffix(".txt")
+        if txt.exists():
+            pairs.append((img, txt))
+    random.Random(args.seed).shuffle(pairs)
+    print(f"пар: {len(pairs)}")
+
+    rows, kept = [], 0
+    for i, (img, txt) in enumerate(pairs):
+        im = ImageOps.exif_transpose(Image.open(img)).convert("RGB")
+        plan = plan_size(im.width, im.height, args.max_upscale)
+        if plan is None:
+            continue
+        label, tw, th, cw, ch = plan
+        left, top = (im.width - cw) // 2, (im.height - ch) // 2
+        im = im.crop((left, top, left + cw, top + ch)).resize((tw, th), Image.LANCZOS)
+        split = "eval" if i < args.eval_n else "train"
+        name = f"{img.stem}"
+        im.save(out / split / f"{name}.jpg", quality=95)
+        (out / split / f"{name}.txt").write_text(txt.read_text().strip() + "\n")
+        rows.append({"split": split, "name": name, "w": tw, "h": th, "bucket": label})
+        kept += 1
+    with open(out / "index.csv", "w", newline="") as fh:
+        w = csv.DictWriter(fh, fieldnames=["split", "name", "w", "h", "bucket"])
+        w.writeheader(); w.writerows(rows)
+    print(f"готово: {kept} (train {kept-args.eval_n}, eval {args.eval_n}) -> {out}")
+
+
+if __name__ == "__main__":
+    main()

src/make_dataset_people.py

@@ -0,0 +1,174 @@
+"""Сборка набора по классам (woman/man/cat) из ~/dataset/allphoto_train — тест анатомии.
+
+Зачем: бабочки однородны и анатомию не показывают. Нужен набор, где модель обязана выучить
+тело: разные люди, позы, ракурсы, и капшены, которые называют, кто это и что делает.
+
+Классы: woman / man / cat — по N уникальных капшенов на класс (дедуп по первым 60 символам,
+иначе один photoshoot даёт десятки почти одинаковых кадров).
+
+  python3 src/make_dataset_people.py --n 60 --out ~/dataset/testset_anatomy
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import math
+import random
+import re
+import shutil
+from pathlib import Path
+
+from PIL import Image, ImageOps
+
+SRC = Path.home() / "dataset" / "allphoto_train"
+
+CLASSES = {
+    "woman": ("woman", "women", "girl"),
+    "man": ("man", "men", "boy"),
+    "cat": ("cat",),
+}
+# бакеты как в moss/src/make_dataset.py: обе стороны кратны 16 (латент /8, потом patch-2),
+# соотношение выбирается ближайшим по |log(ratio_источника / ratio_бакета)|, затем ЦЕНТР-КРОП —
+# растягивать нельзя, иначе лица и тела деформируются
+def buckets(long_side: int = 512, multiple: int = 64) -> dict:
+    """Бакеты: длинная сторона long_side, обе стороны кратны 16 (латент /8, затем /2)."""
+    L = long_side
+    def r16(x): return max(multiple, int(round(x / multiple)) * multiple)
+    return {
+        "1:1": (L, L),
+        "4:3": (L, r16(L * 3 / 4)),
+        "3:4": (r16(L * 3 / 4), L),
+        "3:2": (L, r16(L * 2 / 3)),
+        "2:3": (r16(L * 2 / 3), L),
+        "16:9": (L, r16(L * 9 / 16)),
+        "9:16": (r16(L * 9 / 16), L),
+    }
+BUCKETS = buckets(512)
+
+# для людей требуем, чтобы в кадре было тело, а не только лицо: бабочки-портреты анатомию не учат
+BODY_WORDS = ("standing", "walking", "sitting", "seated", "jeans", "shorts", "dress", "skirt",
+              "shirt", "legs", "arms", "hands", "wearing", "full body", "walks", "stands")
+
+
+def classify(cap: str) -> str | None:
+    low = cap.lower()
+    for cls, words in CLASSES.items():
+        if any(re.search(rf"\b{w}\b", low) for w in words):
+            # «man» внутри «woman» уже отсечён \b; но «woman» не должна становиться «man»
+            if cls == "man" and re.search(r"\bwoman\b", low):
+                return "woman"
+            return cls
+    return None
+
+
+def trim_black(im: Image.Image, thr: float = 12.0, var: float = 8.0) -> Image.Image:
+    """Срезает однородные чёрные поля по краям (в allphoto_train их полно: 71 из 171 кадра)."""
+    import numpy as np
+    a = np.asarray(im.convert("L"), dtype=np.float32)
+    def uniform_edge(line):        # чёрная и без структуры -> рамка, а не ночная сцена
+        return line.mean() < thr and line.std() < var
+    top, bot, left, right = 0, a.shape[0], 0, a.shape[1]
+    while top < bot - 8 and uniform_edge(a[top]): top += 1
+    while bot > top + 8 and uniform_edge(a[bot - 1]): bot -= 1
+    while left < right - 8 and uniform_edge(a[top:bot, left]): left += 1
+    while right > left + 8 and uniform_edge(a[top:bot, right - 1]): right -= 1
+    return im.crop((left, top, right, bot)) if (left or top or right != a.shape[1] or bot != a.shape[0]) else im
+
+
+def plan_size(w: int, h: int, max_upscale: float = 1.4):
+    """(w,h) -> (label, tw, th, cw, ch) или None, если источника не хватает (как в moss)."""
+    ratio = w / h
+    label, (tw, th) = min(BUCKETS.items(), key=lambda kv: abs(math.log(ratio / (kv[1][0] / kv[1][1]))))
+    if ratio > tw / th:
+        cw, ch = round(h * tw / th), h
+    else:
+        cw, ch = w, round(w * th / tw)
+    cw, ch = min(cw, w), min(ch, h)
+    if cw * max_upscale < tw or ch * max_upscale < th:
+        return None
+    return label, tw, th, cw, ch
+
+
+def load_bucketed(path: Path, max_upscale: float = 1.4):
+    """EXIF-поворот -> срез чёрных полей -> ближайший бакет -> центр-кроп -> ресайз."""
+    im = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
+    im = trim_black(im)
+    plan = plan_size(im.width, im.height, max_upscale)
+    if plan is None:
+        return None, None
+    label, tw, th, cw, ch = plan
+    left, top = (im.width - cw) // 2, (im.height - ch) // 2
+    return im.crop((left, top, left + cw, top + ch)).resize((tw, th), Image.LANCZOS), label
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--n", type=int, default=60, help="картинок на класс")
+    ap.add_argument("--eval-n", type=int, default=4, help="сколько на класс отложить для показа")
+    ap.add_argument("--max-upscale", type=float, default=1.4)
+    ap.add_argument("--long-side", type=int, default=512)
+    ap.add_argument("--multiple", type=int, default=64)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--out", default=str(Path.home() / "dataset" / "testset_anatomy"))
+    args = ap.parse_args()
+
+    global BUCKETS
+    BUCKETS = buckets(args.long_side, args.multiple)
+    out = Path(args.out)
+    for sub in ("train", "eval"):
+        (out / sub).mkdir(parents=True, exist_ok=True)
+
+    caps = sorted(SRC.glob("*.txt"))
+    rng = random.Random(args.seed)
+    rng.shuffle(caps)
+
+    seen_prefix: set[str] = set()
+    picked: dict[str, list[tuple[Path, str]]] = {c: [] for c in CLASSES}
+    need = args.n + args.eval_n
+    for f in caps:
+        cap = " ".join(f.read_text().split())
+        if len(cap) < 40:
+            continue
+        cls = classify(cap)
+        if cls is None or len(picked[cls]) >= need:
+            continue
+        low = cap.lower()
+        if cls == "cat" and re.search(r"\b(dog|puppy|kitten|rabbit)\b", low) and "cat" not in low:
+            continue
+        if cls in ("woman", "man") and not any(w in low for w in BODY_WORDS):
+            continue
+        key = cap[:60].lower()
+        if key in seen_prefix:
+            continue
+        img = f.with_suffix(".jpg")
+        if not img.exists():
+            continue
+        seen_prefix.add(key)
+        picked[cls].append((img, cap))
+        if all(len(v) >= need for v in picked.values()):
+            break
+
+    rows = []
+    for cls, items in picked.items():
+        for i, (img, cap) in enumerate(items):
+            split = "eval" if i < args.eval_n else "train"
+            im, bucket = load_bucketed(img)
+            if im is None:
+                continue
+            name = f"{cls}_{i:03d}"
+            im.save(out / split / f"{name}.jpg", quality=92)
+            (out / split / f"{name}.txt").write_text(cap + "\n")
+            rows.append({"split": split, "class": cls, "name": name, "w": im.width, "h": im.height,
+                         "bucket": bucket, "caption": cap})
+            print(f"  {split}/{name} {im.width}x{im.height} | {cap[:80]}")
+
+    with open(out / "index.csv", "w", newline="") as fh:
+        w = csv.DictWriter(fh, fieldnames=["split", "class", "name", "w", "h", "bucket", "caption"])
+        w.writeheader(); w.writerows(rows)
+    per = {c: sum(1 for r in rows if r["class"] == c and r["split"] == "train") for c in CLASSES}
+    print(f"\nитого: {sum(per.values())} train {per} + {sum(1 for r in rows if r['split']=='eval')} eval"
+          f" -> {out}")
+
+
+if __name__ == "__main__":
+    main()

src/model.py

@@ -0,0 +1,270 @@
+"""UDiT — пирамидальный U-трансформер: тело (латентный поток) + хвост (декодер AE).
+
+Идея: ширина растёт там, где токенов мало. Узко на высоком разрешении, широко внизу,
+как в U-Net'ах SD (320/640/1280 при 64/32/16), а не 1024 на всех уровнях.
+
+    512², латент f8 (64×36 = 2304 токена):
+        L0  ÷1   2304 tok  C=192   conv-блоки (depthwise 3×3 + expand 4×)    6 enc / 4 dec
+        L1  ÷2    576 tok  C=512   трансформер + cross-attn к тексту         4 enc / 3 dec
+        L2  ÷4    144 tok  C=1024  трансформер, полное внимание              3 enc / 2 dec
+        L3  ÷8     36 tok  C=1536  трансформер, полное внимание              2 (боттлнек)
+
+    Скипы — аддитивные (SD-стиль): после e0/e1/e2 признаки прибавляются в декодере.
+    Переходы: mean-pool 2× + Linear вниз, Linear + nearest 2× вверх.
+    Оценка: ~160M параметров и ~70 GFLOP против ~580 GFLOP у плоского 1024-стека.
+
+    Хвост — декодер AE (рендер латента в пиксели), текст — Ettin + SliceMixer (из `text.py`).
+"""
+from __future__ import annotations
+
+import math
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT / "vae"))
+
+from ae.model import AutoEncoder  # noqa: E402
+
+DEV = "cuda" if torch.cuda.is_available() else "cpu"
+_ROPE_CACHE: dict = {}
+
+
+def t_embed(t: torch.Tensor, dim: int = 256) -> torch.Tensor:
+    """Синусоидальный эмбеддинг времени (t в [0,1] умножается на 1000)."""
+    half = dim // 2
+    f = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device, dtype=torch.float32) / half)
+    a = (t.float()[:, None] * 1000.0) * f[None]
+    return torch.cat([a.sin(), a.cos()], -1)
+
+
+def body_rope(q, k, gh, gw, base: float = 100.0):
+    """2D RoPE: 4 четверти head_dim, две крутятся по x, две по y. Фаза — в латентных пикселях."""
+    key = (gh, gw, q.shape[-1], q.device, base)
+    tabs = _ROPE_CACHE.get(key)
+    if tabs is None:
+        dq = q.shape[-1] // 4
+        fr = torch.exp(-math.log(base) * torch.arange(dq, device=q.device, dtype=torch.float32) / dq)
+        yy = torch.arange(gh, device=q.device).repeat_interleave(gw).float()
+        xx = torch.arange(gw, device=q.device).repeat(gh).float()
+        ax, ay = xx[:, None] * fr, yy[:, None] * fr
+        tabs = (ax.cos()[None, None], ax.sin()[None, None], ay.cos()[None, None], ay.sin()[None, None])
+        _ROPE_CACHE[key] = tabs
+    cx, sx, cy, sy = tabs
+    dq = q.shape[-1] // 4
+
+    def rot(t):
+        a, b, c, e = t[..., :dq], t[..., dq:2 * dq], t[..., 2 * dq:3 * dq], t[..., 3 * dq:]
+        return torch.cat([a * cx - b * sx, a * sx + b * cx, c * cy - e * sy, c * sy + e * cy], dim=-1)
+
+    return rot(q), rot(k)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d: int, eps: float = 1e-6):
+        super().__init__()
+        self.d, self.eps = d, eps
+        self.scale = nn.Parameter(torch.zeros(d))
+
+    def forward(self, x):
+        return F.rms_norm(x.float(), (self.d,), weight=(self.scale.float() + 1.0), eps=self.eps).to(x.dtype)
+
+
+class ConvBlock(nn.Module):
+    """ConvNeXt-ish: RMSNorm → depthwise 3×3 → expand 4× → GELU → project, всё под adaLN-гейтом."""
+
+    def __init__(self, c: int, expand: int = 4, text_dim: int | None = None, heads: int = 6):
+        super().__init__()
+        self.c, self.hid, self.h, self.hd = c, c * expand, heads, c // heads
+        self.norm = RMSNorm(c)
+        self.norm_x = RMSNorm(c)
+        self.dw = nn.Conv2d(c, c, 3, padding=1, groups=c)
+        self.up = nn.Linear(c, self.hid)
+        self.down = nn.Linear(self.hid, c)
+        self.mod = nn.Parameter(torch.zeros(5, c))     # shift, scale, gate | shift_x, scale_x
+        self.xatt = text_dim is not None
+        if self.xatt:                                   # текст на мелком уровне: не экономим
+            self.xq, self.xo = nn.Linear(c, c, bias=False), nn.Linear(c, c, bias=False)
+        nn.init.zeros_(self.down.weight); nn.init.zeros_(self.down.bias)
+
+    def forward(self, x, vec, gh, gw, text=None, text_mask=None):
+        B = x.shape[0]
+        shift, scale, gate, shift_x, scale_x = (self.mod + vec).unbind(1)
+        h = self.norm(x) * (1 + scale[:, None]) + shift[:, None]
+        h = self.dw(h.transpose(1, 2).reshape(B, self.c, gh, gw)).flatten(2).transpose(1, 2)
+        x = x + gate[:, None] * self.down(F.gelu(self.up(h)))
+        if self.xatt and text is not None:
+            N = x.shape[1]
+            qq = self.xq(self.norm_x(x) * (1 + scale_x[:, None]) + shift_x[:, None])
+            qq = qq.reshape(B, N, self.h, self.hd).transpose(1, 2)
+            kk, vv = text
+            mask = text_mask[:, None, None, :] if text_mask is not None else None
+            xa = F.scaled_dot_product_attention(qq, kk, vv, attn_mask=mask)
+            x = x + self.xo(xa.transpose(1, 2).reshape(B, N, self.c))
+        return x
+
+
+class AttnBlock(nn.Module):
+    """Pre-norm трансформер: self-attention с 2D RoPE + (опц.) cross-attn к тексту + FFN."""
+
+    def __init__(self, c: int, heads: int, xatt_dim: int | None, mlp: float = 2.0):
+        super().__init__()
+        self.c, self.h, self.hd = c, heads, c // heads
+        self.n1, self.n2, self.n3 = RMSNorm(c), RMSNorm(c), RMSNorm(c)
+        self.qkv = nn.Linear(c, 3 * c, bias=False)
+        self.o = nn.Linear(c, c, bias=False)
+        self.up, self.down = nn.Linear(c, int(c * mlp)), nn.Linear(int(c * mlp), c)
+        self.mod = nn.Parameter(torch.zeros(6, c))
+        self.xatt = xatt_dim is not None
+        if self.xatt:
+            self.xq, self.xo = nn.Linear(c, c, bias=False), nn.Linear(c, c, bias=False)
+        nn.init.zeros_(self.down.weight); nn.init.zeros_(self.down.bias)
+
+    def forward(self, x, vec, gh, gw, text=None, text_mask=None):
+        B, N, C = x.shape
+        s1, g1, s2, g2, s3, g3 = (self.mod + vec).unbind(1)
+        h = self.n1(x) * (1 + s1[:, None])
+        q, k, v = self.qkv(h).reshape(B, N, 3, self.h, self.hd).permute(2, 0, 3, 1, 4)
+        q, k = body_rope(q, k, gh, gw)
+        a = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(B, N, C)
+        x = x + g1[:, None] * self.o(a)
+        if self.xatt and text is not None:
+            qq = self.xq(self.n2(x)).reshape(B, N, self.h, self.hd).transpose(1, 2)
+            kk, vv = text
+            mask = text_mask[:, None, None, :] if text_mask is not None else None
+            xa = F.scaled_dot_product_attention(qq, kk, vv, attn_mask=mask)
+            x = x + g2[:, None] * self.xo(xa.transpose(1, 2).reshape(B, N, C))
+        h = self.n3(x) * (1 + s3[:, None])
+        return x + g3[:, None] * self.down(F.gelu(self.up(h)))
+
+
+class Body(nn.Module):
+    """U-стек: ширина 192 → 512 → 1024 → 1536 с аддитивными скипами."""
+
+    def __init__(self, latent_ch: int = 32, text_dim: int = 768,
+                 widths: tuple = (192, 512, 1024, 1536), enc: tuple = (6, 4, 3, 2),
+                 dec: tuple = (4, 3, 2), xatt: bool = True):
+        super().__init__()
+        w0, w1, w2, w3 = widths
+        self.widths, self.enc, self.dec = widths, enc, dec
+        heads = {0: 6, 1: 8, 2: 8, 3: 8}
+        self.stem = nn.Linear(latent_ch, w0)
+        self.e0 = nn.ModuleList([ConvBlock(w0, text_dim=text_dim if xatt else None) for _ in range(enc[0])])
+        self.e0a = AttnBlock(w0, 6, text_dim if xatt else None)
+        self.d0a = AttnBlock(w0, 6, text_dim if xatt else None)
+        self.e1 = nn.ModuleList([AttnBlock(w1, heads[1], text_dim if xatt else None) for _ in range(enc[1])])
+        self.e2 = nn.ModuleList([AttnBlock(w2, heads[2], text_dim if xatt else None) for _ in range(enc[2])])
+        self.e3 = nn.ModuleList([AttnBlock(w3, heads[3], text_dim if xatt else None) for _ in range(enc[3])])
+        self.d2 = nn.ModuleList([AttnBlock(w2, heads[2], text_dim if xatt else None) for _ in range(dec[2])])
+        self.d1 = nn.ModuleList([AttnBlock(w1, heads[1], text_dim if xatt else None) for _ in range(dec[1])])
+        self.d0 = nn.ModuleList([ConvBlock(w0, text_dim=text_dim if xatt else None) for _ in range(dec[0])])
+        self.dn01, self.dn12, self.dn23 = nn.Linear(w0, w1), nn.Linear(w1, w2), nn.Linear(w2, w3)
+        self.up32, self.up21, self.up10 = nn.Linear(w3, w2), nn.Linear(w2, w1), nn.Linear(w1, w0)
+        self.tmlp = nn.Sequential(nn.Linear(256, 512), nn.SiLU(), nn.Linear(512, 512))
+        self.pmlp = nn.Linear(text_dim, 512)
+        self.ada = nn.ModuleDict({f"l{i}": nn.Sequential(nn.SiLU(), nn.Linear(512, (5 if i == 0 else 6) * w))
+                                  for i, w in enumerate(widths)})
+        self.ada["l0a"] = nn.Sequential(nn.SiLU(), nn.Linear(512, 6 * w0))   # для self-attn на L0
+        self.xk0 = nn.Linear(text_dim, 2 * w0)
+        self.xk1, self.xk2, self.xk3 = (nn.Linear(text_dim, 2 * w1), nn.Linear(text_dim, 2 * w2),
+                                        nn.Linear(text_dim, 2 * w3))
+        self.out_norm = RMSNorm(w0)
+        self.out = nn.Linear(w0, latent_ch)
+        nn.init.zeros_(self.out.weight); nn.init.zeros_(self.out.bias)
+
+    # ---- переходы
+    @staticmethod
+    def downsample(x, gh, gw, proj):
+        B = x.shape[0]
+        d = x.view(B, gh // 2, 2, gw // 2, 2, -1).mean(dim=(2, 4))
+        return proj(d.reshape(B, (gh // 2) * (gw // 2), -1))
+
+    @staticmethod
+    def upsample(x, gh, gw, proj):
+        """(B, gh*gw, C) -> (B, 4*gh*gw, C') в родительском row-major порядке."""
+        B = x.shape[0]
+        d = proj(x).view(B, gh, 1, gw, 1, -1).expand(B, gh, 2, gw, 2, proj.out_features)
+        return d.reshape(B, 4 * gh * gw, proj.out_features)
+
+    def _txt(self, proj, text, heads):
+        B, L, _ = text.shape
+        kk, vv = proj(text).chunk(2, dim=-1)
+        hd = kk.shape[-1] // heads
+        return (kk.reshape(B, L, heads, hd).transpose(1, 2), vv.reshape(B, L, heads, hd).transpose(1, 2))
+
+    def forward(self, z, t, text, text_mask, gh=None, gw=None):
+        """gh/gw принимаются для совместимости с интерфейсом UDM и игнорируются."""
+        B, C, gh0, gw0 = z.shape
+        cond = self.tmlp(t_embed(t, 256)) + self.pmlp(text.mean(1))
+        v = [self.ada[f"l{i}"](cond).view(B, 5 if i == 0 else 6, -1) for i in range(4)]
+        v0a = self.ada["l0a"](cond).view(B, 6, -1)
+        tk0 = self._txt(self.xk0, text, 6)
+        tk1 = self._txt(self.xk1, text, 8)
+        tk2 = self._txt(self.xk2, text, 8)
+        tk3 = self._txt(self.xk3, text, 8)
+
+        x = self.stem(z.flatten(2).transpose(1, 2))
+        for blk in self.e0:
+            x = blk(x, v[0], gh0, gw0, tk0, text_mask)
+        x = self.e0a(x, v0a, gh0, gw0, tk0, text_mask)
+        s0 = x
+        g1h, g1w = gh0 // 2, gw0 // 2
+        x = self.downsample(x, gh0, gw0, self.dn01)
+        for blk in self.e1:
+            x = blk(x, v[1], g1h, g1w, tk1, text_mask)
+        s1 = x
+        g2h, g2w = g1h // 2, g1w // 2
+        x = self.downsample(x, g1h, g1w, self.dn12)
+        for blk in self.e2:
+            x = blk(x, v[2], g2h, g2w, tk2, text_mask)
+        s2 = x
+        g3h, g3w = g2h // 2, g2w // 2
+        x = self.downsample(x, g2h, g2w, self.dn23)
+        for blk in self.e3:
+            x = blk(x, v[3], g3h, g3w, tk3, text_mask)
+        # декодер: скипы аддитивные
+        x = self.upsample(x, g3h, g3w, self.up32) + s2
+        for blk in self.d2:
+            x = blk(x, v[2], g2h, g2w, tk2, text_mask)
+        x = self.upsample(x, g2h, g2w, self.up21) + s1
+        for blk in self.d1:
+            x = blk(x, v[1], g1h, g1w, tk1, text_mask)
+        x = self.upsample(x, g1h, g1w, self.up10) + s0
+        x = self.d0a(x, v0a, gh0, gw0, tk0, text_mask)
+        for blk in self.d0:
+            x = blk(x, v[0], gh0, gw0, tk0, text_mask)
+        return self.out(self.out_norm(x)).transpose(1, 2).reshape(B, C, gh0, gw0)
+
+
+class UDiT(nn.Module):
+    """Пирамидальное тело + хвост-декодер AE; интерфейс как у UDM."""
+
+    def __init__(self, ae: AutoEncoder, text_dim: int = 768, **kw):
+        super().__init__()
+        self.ae = ae
+        self.body = Body(latent_ch=ae.config.latent_channels, text_dim=text_dim, **kw)
+
+    def to_raw(self, z):
+        return z * self.ae.latent_std + self.ae.latent_mean
+
+    @torch.no_grad()
+    def sample(self, text, text_mask, gh, gw, steps=4, tail_steps=4, generator=None,
+               param="x0", cfg=1.0, null_text=None, null_mask=None, clip_norm=0.0):
+        dev = self.ae.device
+        z = torch.randn(1, self.ae.config.latent_channels, gh, gw, device=dev, generator=generator)
+        with torch.autocast(dev.type, torch.bfloat16, enabled=dev.type == "cuda"):
+            for i in range(steps):
+                t = torch.full((1,), 1.0 - i / steps, device=dev)
+                o = self.body(z, t, text, text_mask).float()
+                if cfg > 1.0 and null_text is not None:
+                    ou = self.body(z, t, null_text, null_mask).float()
+                    o = ou + cfg * (o - ou)
+                v = o if param == "vel" else (z - o) / t.clamp(min=1e-2)[:, None, None, None]
+                z = z - v / steps
+        if clip_norm:
+            z = z.clamp(-clip_norm, clip_norm)
+        return self.ae.decode_raw(self.to_raw(z), steps=tail_steps, generator=generator)

src/publish.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+"""Сборка публикационного дерева: код + доки + веса в fp16.
+
+    python3 src/publish.py --out tmp/publish          # только собрать
+    python3 src/publish.py --out tmp/publish --push   # собрать и залить (REPO_ID)
+
+Локальные fp32-веса остаются на месте: они нужны обучению, а на Hub едет fp16.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+
+ROOT = Path(__file__).resolve().parent.parent
+REPO_ID = "AiArtLab/ra"
+
+COPY = ["README.md", ".gitattributes", ".gitignore", "docs", "src",
+        "vae/config.json", "vae/LICENSE", "vae/README.md", "vae/.gitattributes", "vae/example.py",
+        "vae/ae", "text_encoder/config.json", "text_encoder/tokenizer.json",
+        "text_encoder/tokenizer_config.json", "text_encoder/special_tokens_map.json",
+        "text_encoder/README.md", "text_encoder/.gitattributes",
+        "transformer/config.json"]
+
+
+def half_file(src: Path, dst: Path, key_prefix: str = ""):
+    tensors = load_file(str(src))
+    out, overflow = {}, 0
+    for k, v in tensors.items():
+        name = f"{key_prefix}{k}" if key_prefix else k
+        if v.is_floating_point():
+            if v.numel() and float(v.abs().max()) > 65504:
+                overflow += 1
+            v = v.half()
+        out[name] = v.contiguous()
+        assert torch.isfinite(out[name]).all(), f"{name}: не-конечные значения после fp16"
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    save_file(out, str(dst))
+    print(f"[fp16] {dst.relative_to(dst.parents[1])}: {len(out)} тензоров, "
+          f"{src.stat().st_size / 2**20:.0f} → {dst.stat().st_size / 2**20:.0f} МиБ"
+          + (f", ВНИМАНИЕ: {overflow} тензоров вне диапазона fp16" if overflow else ""), flush=True)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", default=str(ROOT / "tmp" / "publish"))
+    ap.add_argument("--repo", default=REPO_ID)
+    ap.add_argument("--push", action="store_true", help="залить собранное на Hub (публичный репозиторий)")
+    args = ap.parse_args()
+    out = Path(args.out)
+    if out.exists():
+        shutil.rmtree(out)
+    out.mkdir(parents=True)
+
+    for rel in COPY:
+        src = ROOT / rel
+        dst = out / rel
+        dst.parent.mkdir(parents=True, exist_ok=True)
+        if src.is_dir():
+            shutil.copytree(src, dst, ignore=shutil.ignore_patterns("__pycache__", ".ipynb_checkpoints"))
+        else:
+            shutil.copy2(src, dst)
+
+    half_file(ROOT / "transformer" / "diffusion_pytorch_model.safetensors",
+              out / "transformer" / "diffusion_pytorch_model.safetensors")
+    half_file(ROOT / "vae" / "model.safetensors", out / "vae" / "model.safetensors")
+
+    sd = torch.load(ROOT / "text_encoder" / "pytorch_model.bin", map_location="cpu", weights_only=True)
+    sd = {k: v for k, v in sd.items() if not k.startswith(("head.", "decoder."))}   # MLM-голова не нужна
+    save_file({k: v.half().contiguous() for k, v in sd.items()},
+              str(out / "text_encoder" / "model.safetensors"))
+    print(f"[fp16] text_encoder/model.safetensors: {len(sd)} тензоров "
+          f"({sum(v.numel() for v in sd.values()) / 1e6:.1f}M), "
+          f"{(out / 'text_encoder' / 'model.safetensors').stat().st_size / 2**20:.0f} МиБ", flush=True)
+
+    cfg_path = out / "transformer" / "config.json"
+    cfg = json.loads(cfg_path.read_text())
+    cfg["storage_dtype"] = "float16"
+    cfg_path.write_text(json.dumps(cfg, indent=2, ensure_ascii=False))
+
+    total = sum(f.stat().st_size for f in out.rglob("*") if f.is_file())
+    print(f"[готово] {out}: {total / 2**30:.2f} ГиБ", flush=True)
+
+    if args.push:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.create_repo(args.repo, repo_type="model", private=False, exist_ok=True)
+        api.upload_folder(folder_path=str(out), repo_id=args.repo, repo_type="model",
+                          commit_message=f"Ra fp16 (step {cfg['step']})")
+        print(f"[push] https://huggingface.co/{args.repo}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

src/sample.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+"""Сэмпл из transformer/ (или из .pt чекпойнта обучения).
+
+    python3 src/sample.py --prompt "A photograph of a red fox in fresh snow." --out fox.png
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path[:0] = [str(ROOT / "src"), str(ROOT / "vae")]
+
+from model import DEV, UDiT          # noqa: E402
+from text import EttinText           # noqa: E402
+from ae.model import AutoEncoder     # noqa: E402
+from export import load_exported     # noqa: E402
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default=str(ROOT / "transformer"), help="папка transformer/ или .pt")
+    ap.add_argument("--prompt", required=True)
+    ap.add_argument("--out", default=str(ROOT / "tmp" / "sample.png"))
+    ap.add_argument("--size", type=int, default=256, help="сторона кадра, кратна 64")
+    ap.add_argument("--steps", type=int, default=16, help="шагов тела")
+    ap.add_argument("--tail-steps", type=int, default=8, help="шагов декодера (хвоста)")
+    ap.add_argument("--cfg", type=float, default=2.5)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--n", type=int, default=1, help="сколько картинок подряд от --seed")
+    ap.add_argument("--text-len", type=int, default=192)
+    args = ap.parse_args()
+
+    ae = AutoEncoder.from_pretrained(str(ROOT / "vae"), device=DEV)
+    enc = EttinText(str(ROOT / "text_encoder"), DEV, max_len=args.text_len)
+    m = UDiT(ae, text_dim=enc.dim).to(DEV)
+
+    src = Path(args.model)
+    if src.is_dir():
+        step = load_exported(m, enc, ae, src, DEV).get("step")
+    else:
+        sd = torch.load(src, map_location=DEV, weights_only=False)
+        own = m.body.state_dict()
+        m.body.load_state_dict({k: v for k, v in sd["body"].items() if k in own}, strict=False)
+        enc.mixer.load_state_dict(sd["mixer"])
+        if sd.get("decoder") is not None:
+            ae.decoder.load_state_dict(sd["decoder"])
+        step = sd.get("step")
+    m.eval()
+    print(f"[model] {src}, шаг {step}", flush=True)
+
+    null_hs, null_mk = enc.slices([""], pad_to_len=args.text_len)
+    null_t = enc.mix(null_hs.float())
+    hs, mk = enc.slices([args.prompt], pad_to_len=args.text_len)
+    txt = enc.mix(hs.float())
+
+    out = Path(args.out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    gh = gw = args.size // 8
+    for i in range(args.n):
+        y = m.sample(txt, mk, gh, gw, steps=args.steps, tail_steps=args.tail_steps, cfg=args.cfg,
+                     null_text=null_t, null_mask=null_mk,
+                     generator=torch.Generator(DEV).manual_seed(args.seed + i))
+        im = Image.fromarray(((y[0].clamp(-1, 1) + 1) * 127.5).byte().permute(1, 2, 0).cpu().numpy())
+        path = out if args.n == 1 else out.with_name(f"{out.stem}_{args.seed + i}{out.suffix}")
+        im.save(path)
+        print(f"[ok] {path} {im.size[0]}x{im.size[1]}, seed {args.seed + i}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

src/text.py

@@ -0,0 +1,132 @@
+"""Текстовый путь UDM: замороженный Ettin + обучаемый выбор слоёв (по идее moss/adapter).
+
+Почему не один слой. Скрытые состояния растут по норме с глубиной (в Qwen3.5-9B — в 40 раз,
+13 у слоя 2 против 557 у слоя 27), поэтому «взять последний слой» теряет лексику и цифры,
+а «взять ранний» — семантику. В moss/adapter это решили иначе: несколько срезов по K
+нормируются ПОШТУЧНО, идут последовательностью длины K через 2 блока self-attention, затем
+`Linear(K→1)` схлопывает ось срезов, а результат zero-init'ом добавляется к основному пути.
+Тот же приём (плюс 0.8B вместо 9B) дал cos 0.924 к учителю.
+
+Здесь то же самое, но цель — не регрессия к чужому пространству, а обусловливание ДиТ:
+выход миксера идёт в `xkv` (кросс-внимание) и в `pmlp` (adaLN).
+
+Срезы Ettin-150M: 22 слоя, global attention каждые 3 → берём выходы после глобальных слоёв.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# выходы после глобальных слоёв ModernBERT (22 слоя, global_attn_every_n_layers=3)
+ETTIN_150M_SLICES = (2, 5, 8, 11, 14, 17, 20)
+
+
+class AttnBlock(nn.Module):
+    """Pre-norm self-attention + MLP (для последовательности длины K — оси срезов)."""
+
+    def __init__(self, d: int, heads: int = 8, ffn_mult: float = 2.0):
+        super().__init__()
+        self.h = heads
+        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
+        self.qkv = nn.Linear(d, 3 * d)
+        self.proj = nn.Linear(d, d)
+        self.fc1 = nn.Linear(d, int(d * ffn_mult))
+        self.fc2 = nn.Linear(int(d * ffn_mult), d)
+
+    def forward(self, x):
+        b, n, d = x.shape
+        q, k, v = self.qkv(self.n1(x)).reshape(b, n, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
+        a = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(b, n, d)
+        x = x + self.proj(a)
+        return x + self.fc2(F.gelu(self.fc1(self.n2(x))))
+
+
+class SliceMixer(nn.Module):
+    """Выбор слоя: пер-срезный LayerNorm → attention по оси срезов → Linear(K→1) → zero-init out.
+
+    blocks=0 оставляет только обученные веса слоёв (Linear(K→1) поверх нормированных срезов) —
+    дешёвый вариант «какой слой важнее», без внимания между срезами.
+    """
+
+    def __init__(self, n_slices: int, d: int, blocks: int = 2, heads: int = 8):
+        super().__init__()
+        self.n, self.d = n_slices, d
+        self.per = nn.LayerNorm(d, elementwise_affine=False)
+        self.blocks = nn.ModuleList([AttnBlock(d, heads) for _ in range(blocks)])
+        self.proj = nn.Linear(n_slices, 1, bias=False)
+        self.out = nn.Linear(d, d)
+        nn.init.zeros_(self.out.weight)
+        nn.init.zeros_(self.out.bias)
+
+    def forward(self, slices):                      # (B, L, K, D)
+        b, l, k, d = slices.shape
+        normed = self.per(slices)                   # пер-срезная нормировка ОБЯЗАТЕЛЬНА:
+        h = normed.reshape(b * l, k, d)             # сырые нормы срезов растут с глубиной (52…548)
+        for blk in self.blocks:
+            h = blk(h)
+        mixed = self.out(self.proj(h.permute(0, 2, 1)).squeeze(-1)).reshape(b, l, d)
+        return normed.mean(2) + mixed               # среднее по нормированным срезам + zero-init добавка
+
+    def layer_weights(self) -> torch.Tensor:
+        """Что микшер выучил: ненормированные веса слоёв (для логов)."""
+        return self.proj.weight.detach().squeeze()
+
+
+class EttinText(nn.Module):
+    """Замороженный энкодер Ettin-150M + обучаемый SliceMixer. Возвращает (B,L,768) и маску."""
+
+    def __init__(self, path: str, device, max_len: int = 192, slices=ETTIN_150M_SLICES,
+                 blocks: int = 2, dtype=torch.bfloat16):
+        super().__init__()
+        from transformers import AutoModel, AutoTokenizer
+        self.max_len, self.layer_idx = max_len, tuple(slices)
+        self.tok = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModel.from_pretrained(path, dtype=dtype).to(device).eval().requires_grad_(False)
+        self.dim = self.model.config.hidden_size
+        self.mixer = SliceMixer(len(self.layer_idx), self.dim, blocks=blocks).to(device)
+        self.device = torch.device(device)
+
+    def tokenize(self, caps: list[str], pad_to: int = 8, pad_to_len: int | None = None):
+        # padding="longest" (а не до max_len): короткие капшены стоят дешевле
+        t = self.tok(caps, padding="longest", truncation=True, max_length=self.max_len,
+                     return_tensors="pt")
+        ids, mask = t["input_ids"], t["attention_mask"].bool()
+        into = pad_to_len or ids.shape[1]
+        extra = into - ids.shape[1]
+        if extra > 0:
+            ids = F.pad(ids, (0, extra), value=self.tok.pad_token_id)
+            mask = F.pad(mask, (0, extra))
+        elif extra < 0:
+            ids, mask = ids[:, :into], mask[:, :into]
+        else:
+            extra = (-ids.shape[1]) % pad_to
+            if extra:
+                ids = F.pad(ids, (0, extra), value=self.tok.pad_token_id)
+                mask = F.pad(mask, (0, extra))
+        return ids.to(self.device), mask.to(self.device)
+
+    @torch.no_grad()
+    def slices(self, caps: list[str], pad_to_len: int | None = None):
+        """Срезы слоёв — считаются ОДИН раз на капшен (энкодер заморожен)."""
+        ids, mask = self.tokenize(caps, pad_to_len=pad_to_len)
+        out = self.model(input_ids=ids, attention_mask=mask, output_hidden_states=True)
+        return torch.stack([out.hidden_states[i] for i in self.layer_idx], dim=2), mask   # (B,L,K,D)
+
+    def mix(self, hs):
+        """Обучаемый путь: среднее по нормированным срезам + zero-init добавка микшера."""
+        return self.mixer(hs.float())
+
+    def forward(self, caps: list[str]):
+        hs, mask = self.slices(caps)
+        return self.mix(hs), mask, hs
+
+    def trainable(self):
+        return list(self.mixer.parameters())
+
+
+def empty_text(enc: EttinText):
+    """Пустой промпт = [CLS][SEP] — безусловный контекст для CFG."""
+    return enc([""])

src/train.py

@@ -0,0 +1,524 @@
+#!/usr/bin/env python3
+"""Обучение UDiT на папке «картинка + .txt». Квадратные кадры (--size), кэш латентов и текста.
+
+    python3 src/train.py --data ~/dataset/butterflies_256sq/train \
+            --eval ~/dataset/butterflies_256sq/eval --epochs 80 --batch 8
+
+Что внутри (перенесено из привычек agat/train.py):
+  * бюджет в эпохах (`--epochs`) или шагах (`--steps`), lr: прогрев → cosine → полка;
+  * AdamW8bit (bitsandbytes), если он есть, иначе обычный AdamW — печатается, что взято;
+  * дропаут подписи (`--drop-text`) для CFG и дропаут слов (`--drop-word`) внутри подписи;
+  * сэмплы на старте (шаг 0), каждые `--sample-every`; держим `--keep` последних чекпойнтов и
+    `--samples-keep` сеток; сохранение атомарное;
+  * `--probe`: подбирает максимальный батч, который влезает в память, и печатает его;
+  * EMA сглаживание печатаемого лосса (`--ema-decay`), `--seed` случайный по умолчанию (печатается);
+  * метрики: mse_lat (общий, low t<0.3, high t>0.7), mse_pix, switch, tail(z0), hf(gen)/hf(AE), skip;
+  * SIGUSR1 печатает стеки всех потоков (видно, ждёт процесс GPU или сеть, без py-spy).
+"""
+from __future__ import annotations
+
+import argparse
+import faulthandler
+import json
+import math
+import os
+import random
+import shutil
+import signal
+import time
+from pathlib import Path
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+signal.signal(signal.SIGINT, signal.default_int_handler)
+faulthandler.register(signal.SIGUSR1)
+
+import torch                      # noqa: E402
+import torch.nn.functional as F   # noqa: E402
+
+from model import DEV, UDiT        # noqa: E402
+from text import EttinText         # noqa: E402
+from ae.model import AutoEncoder, patchify   # noqa: E402
+
+ROOT = Path(__file__).resolve().parent.parent
+
+
+# ------------------------------------------------------------------ расписание
+def lr_at(step: int, base: float, min_lr: float, warmup: int, period: int, start: int = 0) -> float:
+    """Прогрев `warmup` шагов от `start`, затем cosine до `min_lr` за `period` шагов."""
+    rel = step - start
+    if warmup and rel <= warmup:
+        return base * min(1.0, rel / max(warmup, 1))
+    if period <= warmup:
+        return base
+    p = min(1.0, (rel - warmup) / (period - warmup))
+    return min_lr + (base - min_lr) * 0.5 * (1 + math.cos(math.pi * p))
+
+
+class HFImageStream:
+    """Поток parquet-шардов HF: декод/ресайз/кроп в фоновых потоках, обучение не ждёт диск.
+
+    Латент и текст НЕ кэшируются (каждая картинка встречается один раз), поэтому шаг дороже,
+    чем на папке: +3-4 мс на картинку у AE и +1.2 мс у Ettin.
+    """
+
+    def __init__(self, dataset, parts, size, fields, lo, hi, shuffle_buffer=1000, prefetch=48,
+                 workers=8, seed=0):
+        import queue
+        self.q, self.stop = queue.Queue(maxsize=prefetch), False
+        self.size, self.fields, self.lo, self.hi = size, fields, lo, hi
+        self.dataset, self.parts = dataset, parts
+        self.buffer, self.workers, self.seed = shuffle_buffer, workers, seed
+        import threading
+        threading.Thread(target=self._loop, daemon=True).start()
+
+    def _items(self):
+        """Шарды всех частей перемешиваются между собой (иначе эпоха идёт блоками по частям)."""
+        import random
+        from datasets import load_dataset
+        from huggingface_hub import list_repo_files
+        repo = list_repo_files(self.dataset, repo_type="dataset")
+        files = [f for f in repo
+                 if f.endswith(".parquet") and any(f.startswith(p + "/") for p in self.parts)]
+        if not files:                      # запасной путь, если листинг не отдал файлы
+            files = [f"hf://datasets/{self.dataset}/{p}/*.parquet" for p in self.parts]
+        else:
+            random.Random(self.seed).shuffle(files)
+            files = [f"hf://datasets/{self.dataset}/{f}" for f in files]
+        print(f"[stream] шардов: {len(files)} из частей {self.parts}", flush=True)
+        ds = load_dataset("parquet", data_files=files, split="train", streaming=True)
+        return ds.shuffle(buffer_size=self.buffer, seed=self.seed)
+
+    def _prepare(self, item):
+        from PIL import Image
+        import numpy as np
+        im = item["image"]
+        if im.mode != "RGB":
+            im = im.convert("RGB")
+        w, h = im.size
+        sc = self.size / min(w, h)
+        if abs(sc - 1.0) > 1e-3:
+            im = im.resize((max(self.size, round(w * sc)), max(self.size, round(h * sc))), Image.LANCZOS)
+            w, h = im.size
+        left, top = (w - self.size) // 2, (h - self.size) // 2
+        im = im.crop((left, top, left + self.size, top + self.size))
+        arr = torch.from_numpy(np.asarray(im, dtype=np.float32)).permute(2, 0, 1) / 127.5 - 1.0
+        return arr, self._caption(item)
+
+    def _caption(self, item):
+        clean = lambda x: (x or "").replace("\n", " ").strip()      # noqa: E731
+        for f in self.fields:
+            t = clean(item.get(f))
+            if self.lo <= len(t) <= self.hi:
+                return t
+        return clean(item.get("caption_detail"))
+
+    def _loop(self):
+        import collections
+        from concurrent.futures import ThreadPoolExecutor
+        with ThreadPoolExecutor(max_workers=self.workers) as pool:
+            while not self.stop:
+                try:
+                    it = iter(self._items())
+                    flying = collections.deque()
+                    while not self.stop:
+                        while len(flying) < self.workers:
+                            flying.append(pool.submit(self._prepare, next(it)))
+                        self.q.put(flying.popleft().result())
+                except Exception as e:
+                    print(f"[stream] сбой ({type(e).__name__}: {e}), переподключение через 5 с", flush=True)
+                    time.sleep(5)
+
+    def batch(self, n):
+        imgs, caps = [], []
+        while len(imgs) < n:
+            arr, cap = self.q.get()
+            imgs.append(arr); caps.append(cap)
+        return torch.stack(imgs), caps
+
+
+def hf_energy(t: torch.Tensor) -> float:
+    k = torch.tensor([[0.0, 1.0, 0.0], [1.0, -4.0, 1.0], [0.0, 1.0, 0.0]], device=t.device)
+    return F.conv2d(t.float().mean(1, keepdim=True), k.view(1, 1, 3, 3), padding=1).std().item()
+
+
+def psnr(a, b) -> float:
+    return 10 * math.log10(4 / max((a.clamp(-1, 1) - b).pow(2).mean().item(), 1e-12))
+
+
+def rotate(path: Path, keep: int, pattern: str):
+    """Оставить `keep` последних файлов, совпадающих с pattern (по номеру шага в имени)."""
+    files = sorted(path.glob(pattern))
+    for f in files[:-keep]:
+        f.unlink()
+
+
+def main():
+    ap = argparse.ArgumentParser(description="UDiT: обучение на папке пар <имя>.jpg + <имя>.txt")
+    ap.add_argument("--data", default=str(Path.home() / "dataset" / "butterflies_256sq" / "train"))
+    ap.add_argument("--eval", default=str(Path.home() / "dataset" / "butterflies_256sq" / "eval"))
+    ap.add_argument("--limit", type=int, default=0, help="взять только первые N картинок (папка)")
+    ap.add_argument("--hf-dataset", default="", help="например LucasFang/FLUX-Reason-6M (поток)")
+    ap.add_argument("--parts", default="Aesthetics-Part01,Aesthetics-Part02,Imaginative,Text")
+    ap.add_argument("--caption-fields", default="caption_entity,caption_composition",
+                    help="короткие поля подписи по порядку предпочтения (иначе caption_detail)")
+    ap.add_argument("--caption-len", default=(40, 400), nargs=2, type=int, metavar=("MIN", "MAX"))
+    ap.add_argument("--shuffle-buffer", type=int, default=1000)
+    ap.add_argument("--prefetch", type=int, default=48)
+    ap.add_argument("--workers", type=int, default=8)
+    ap.add_argument("--images-per-epoch", type=int, default=5_890_279,
+                    help="сколько картинок в эпохе потока (для пересчёта --epochs в шаги)")
+    ap.add_argument("--size", type=int, default=256, help="разрешение обучения (кадры квадратные)")
+    ap.add_argument("--text-len", type=int, default=192, help="предел длины подписи в токенах")
+    ap.add_argument("--epochs", type=float, default=None, help="бюджет в эпохах (приоритетнее --steps)")
+    ap.add_argument("--steps", type=int, default=20000, help="бюджет в шагах оптимизатора")
+    ap.add_argument("--batch", type=int, default=8)
+    ap.add_argument("--accum", type=int, default=1, help="накопление: эффективный батч = batch*accum")
+    ap.add_argument("--probe", action="store_true", help="подобрать максимальный батч под память и выйти")
+    ap.add_argument("--lr", type=float, default=1e-4)
+    ap.add_argument("--lr-warmup", type=int, default=500)
+    ap.add_argument("--lr-min", type=float, default=None, help="нижняя точка cosine (по умолчанию 10%% от --lr)")
+    ap.add_argument("--lr-period", type=int, default=0, help="к какому шагу cosine доходит до --lr-min")
+    ap.add_argument("--drop-text", type=float, default=0.1, help="доля строк с пустой подписью (CFG)")
+    ap.add_argument("--drop-word", type=float, default=0.05, help="доля токенов подписи, скрытых в строке")
+    ap.add_argument("--cfg", type=float, default=1.5, help="CFG при показе сэмплов")
+    ap.add_argument("--sample-steps", type=int, default=8, help="шагов тела при сэмплинге")
+    ap.add_argument("--tail-steps", type=int, default=4, help="шагов декодера AE")
+    ap.add_argument("--lr-tail", type=float, default=3e-5)
+    ap.add_argument("--tail", choices=["frozen", "train"], default="train")
+    ap.add_argument("--init", default=None, help="стартовые веса: чекпойнт или папка transformer/")
+    ap.add_argument("--resume", action="store_true", help="продолжить с <out>/last.pt")
+    ap.add_argument("--out", default=str(ROOT / "tmp" / "run"))
+    ap.add_argument("--log-every", type=int, default=50)
+    ap.add_argument("--sample-every", type=int, default=1000, help="0 = не сэмплировать")
+    ap.add_argument("--save-every", type=int, default=2000, help="0 = только в конце")
+    ap.add_argument("--keep", type=int, default=2, help="сколько чекпойнтов держать")
+    ap.add_argument("--samples-keep", type=int, default=4, help="сколько сеток сэмплов держать")
+    ap.add_argument("--ema-decay", type=float, default=0.995, help="сглаживание печатаемого лосса")
+    ap.add_argument("--seed", type=int, default=None)
+    args = ap.parse_args()
+
+    seed = args.seed if args.seed is not None else random.SystemRandom().randrange(2 ** 31)
+    torch.manual_seed(seed)
+    print(f"[seed] {seed}", flush=True)
+
+    out = Path(args.out); (out / "samples").mkdir(parents=True, exist_ok=True)
+    (out / "ckpt").mkdir(exist_ok=True)
+    log_path = out / "log.csv"
+    if not (args.resume and log_path.exists()):
+        log_path.write_text("step,mse_lat,low,high,mse_pix,switch,bright,contrast,lat_pc_min,hf_gen,tail_z0,sec\n")
+
+    from PIL import Image
+    import numpy as np
+
+    def load(path: Path):
+        im = Image.open(path).convert("RGB")
+        return torch.from_numpy(np.asarray(im, dtype=np.float32)).permute(2, 0, 1) / 127.5 - 1
+
+    # ---------------------------------------------------------------- данные
+    stream = None
+    if args.hf_dataset:
+        stream = HFImageStream(args.hf_dataset, args.parts.split(","), args.size,
+                               args.caption_fields.split(","), args.caption_len[0], args.caption_len[1],
+                               args.shuffle_buffer, args.prefetch, args.workers, seed)
+        n = args.images_per_epoch
+        per_epoch = max(1, n // (args.batch * args.accum))
+        steps = int(round(args.epochs * per_epoch)) if args.epochs else args.steps
+        print(f"[data] поток {args.hf_dataset} [{args.parts}], ~{n} картинок/эпоха, "
+              f"{per_epoch} шагов/эпоха → бюджет {steps} шагов "
+              f"(~{steps * args.batch / 62 / 3600:.1f} ч при 62 картинках/с)", flush=True)
+        ae = AutoEncoder.from_pretrained(str(ROOT / "vae"), device=DEV)
+        enc = EttinText(str(ROOT / "text_encoder"), DEV, max_len=args.text_len)
+        HS = MK = X = Z = None
+        HS_E, MK_E = enc.slices([""], pad_to_len=args.text_len)
+        HS_E = HS_E.to(torch.bfloat16)
+    else:
+        files = sorted(Path(args.data).glob("*.jpg")) + sorted(Path(args.data).glob("*.png"))
+        if args.limit:
+            files = files[: args.limit]
+        caps = [f.with_suffix(".txt").read_text().strip() for f in files]
+        ae = AutoEncoder.from_pretrained(str(ROOT / "vae"), device=DEV)
+        enc = EttinText(str(ROOT / "text_encoder"), DEV, max_len=args.text_len)
+        X = torch.stack([load(f) for f in files]).to(DEV)
+        with torch.no_grad():
+            Z = torch.cat([ae.encode(X[i:i + 8]) for i in range(0, len(X), 8)])
+        HS, MK = enc.slices(caps)
+        HS, MK = HS.to(torch.bfloat16), MK
+        HS_E, MK_E = enc.slices([""], pad_to_len=HS.shape[1])
+        HS_E = HS_E.to(torch.bfloat16)
+        n = len(files)
+        per_epoch = max(1, n // (args.batch * args.accum))
+        steps = int(round(args.epochs * per_epoch)) if args.epochs else args.steps
+        print(f"[data] {n} пар, латент {tuple(Z.shape[1:])}, текст {tuple(HS.shape[1:])}, "
+              f"эпоха = {per_epoch} шагов → бюджет {steps}", flush=True)
+
+    # ---------------------------------------------------------------- модель и оптимизатор
+    m = UDiT(ae, text_dim=enc.dim).to(DEV)
+    start = 0
+
+    def load_body(sd):
+        own = m.body.state_dict()
+        take = {k: v for k, v in sd.items() if k in own and own[k].shape == v.shape}
+        m.body.load_state_dict(take, strict=False)
+        return len(take)
+
+    if args.init:
+        p = Path(args.init)
+        if p.is_dir():
+            from export import load_exported
+            cfg = load_exported(m, enc, ae, p, DEV)
+            print(f"[init] папка {p}, шаг {cfg.get('step')}", flush=True)
+        else:
+            sd = torch.load(p, map_location=DEV)
+            n_take = load_body(sd["body"]); enc.mixer.load_state_dict(sd["mixer"])
+            if sd.get("decoder") is not None and args.tail == "train":
+                ae.decoder.load_state_dict(sd["decoder"])
+            start = int(sd.get("step", 0))
+            print(f"[init] {p}: перенесено {n_take} тензоров, шаг {start}", flush=True)
+    if args.resume and (out / "last.pt").exists():
+        sd = torch.load(out / "last.pt", map_location=DEV)
+        load_body(sd["body"]); enc.mixer.load_state_dict(sd["mixer"])
+        if sd.get("decoder") is not None and args.tail == "train":
+            ae.decoder.load_state_dict(sd["decoder"])
+        start = int(sd.get("step", 0))
+        print(f"[resume] {out / 'last.pt'} → шаг {start}", flush=True)
+
+    n_body = sum(p.numel() for p in m.body.parameters())
+    n_tail = sum(p.numel() for p in ae.decoder.parameters())
+    print(f"[model] тело {n_body/1e6:.1f}M + миксер {sum(p.numel() for p in enc.trainable())/1e6:.2f}M "
+          f"+ хвост {n_tail/1e6:.1f}M [{args.tail}]", flush=True)
+
+    groups = [{"params": list(m.body.parameters()), "lr": args.lr},
+              {"params": enc.trainable(), "lr": args.lr}]
+    if args.tail == "train":
+        ae.decoder.requires_grad_(True)
+        groups.append({"params": list(ae.decoder.parameters()), "lr": args.lr_tail})
+    params = [p for g in groups for p in g["params"]]
+    try:
+        import bitsandbytes as bnb
+        opt = bnb.optim.AdamW8bit(groups, betas=(0.9, 0.95), weight_decay=0.0)
+        print("[opt] AdamW8bit", flush=True)
+    except Exception as e:
+        opt = torch.optim.AdamW(groups, betas=(0.9, 0.95), weight_decay=0.0)
+        print(f"[opt] AdamW fp32 (8bit недоступен: {type(e).__name__})", flush=True)
+
+    lr_min = args.lr_min if args.lr_min is not None else args.lr * 0.1
+    period = args.lr_period or steps              # длина cosine в шагах, отсчёт от start
+    print(f"[lr] {args.lr:g}: прогрев {args.lr_warmup} от шага {start}, cosine до {lr_min:g} к шагу {start + period}", flush=True)
+
+    # ---------------------------------------------------------------- показ сэмплов
+    SAMPLE_PROMPTS = [
+        "A photograph of a red fox sleeping on a mossy log in a misty forest at dawn.",
+        "A city made of glass where rivers of light flow between the towers, cinematic lighting.",
+        'A neon sign reading "hello world" mounted on a brick wall at night in the rain.',
+        "An impressionist oil painting of a woman with a parasol standing by a lake.",
+    ]
+    show_txt, show_x = [], []
+    if stream is not None:
+        for cap in SAMPLE_PROMPTS:
+            hs, mk = enc.slices([cap])
+            show_txt.append((enc.mix(hs.float()), mk))
+    else:
+        ev_all = sorted(Path(args.eval).glob("*.jpg")) + sorted(Path(args.eval).glob("*.png"))
+        ev_cls = {}
+        for f in ev_all:
+            ev_cls.setdefault(f.stem.split("_")[0], f)
+        show = [ev_cls[k] for k in list(ev_cls)[:4]] or ev_all[:4]
+        for f in show:
+            hs, mk = enc.slices([f.with_suffix(".txt").read_text().strip()], pad_to_len=HS.shape[1])
+            show_txt.append((enc.mix(hs.float()), mk))
+            show_x.append(load(f)[None].to(DEV))
+
+    def sample_grid(step):
+        tiles, hfs, ys = [], [], []
+        for txt, mk in show_txt:
+            y = m.sample(txt, mk, lat_hw[0], lat_hw[1], steps=args.sample_steps,
+                         tail_steps=args.tail_steps, cfg=args.cfg,
+                         null_text=enc.mix(HS_E.float()), null_mask=MK_E,
+                         generator=torch.Generator(DEV).manual_seed(0))
+            tiles.append(Image.fromarray(((y[0].clamp(-1, 1) + 1) * 127.5).byte()
+                                         .permute(1, 2, 0).cpu().numpy()))
+            hfs.append(hf_energy(y)); ys.append(y)
+        if show_x:
+            ae_rec = ae.decode_raw(ae.encode(show_x[0]) * ae.latent_std + ae.latent_mean,
+                                   steps=args.tail_steps, generator=torch.Generator(DEV).manual_seed(0))
+            hf_ae = hf_energy(ae_rec)
+        else:
+            hf_ae = float("nan")
+        W, H = tiles[0].size
+        canvas = Image.new("RGB", (W * len(tiles), H), "white")
+        for i, t in enumerate(tiles):
+            canvas.paste(t, (i * W, 0))
+        canvas.save(out / "samples" / "latest.png")
+        canvas.save(out / "samples" / f"step_{step:07d}.png")
+        rotate(out / "samples", args.samples_keep, "step_*.png")
+        bright = float(torch.stack(ys).mean())
+        contrast = float(torch.stack(ys).std())
+        return hfs[0], hf_ae, canvas, bright, contrast
+
+    def save(step):
+        sd = {"body": {k: v.cpu() for k, v in m.body.state_dict().items()},
+              "mixer": {k: v.cpu() for k, v in enc.mixer.state_dict().items()},
+              "decoder": ({k: v.cpu() for k, v in ae.decoder.state_dict().items()}
+                          if args.tail == "train" else None),
+              "step": step, "args": vars(args), "seed": seed}
+        torch.save(sd, out / "last.pt")
+        torch.save(sd, out / "ckpt" / f"step_{step:07d}.pt")
+        rotate(out / "ckpt", args.keep, "step_*.pt")
+
+    # ---------------------------------------------------------------- цикл
+    lat_hw = (args.size // 8, args.size // 8)
+    pkx = m.ae._packing([lat_hw] * args.batch)
+    if args.probe:
+        best = args.batch
+        for bs in (args.batch, args.batch * 2, args.batch * 4, args.batch * 8):
+            try:
+                z = torch.randn(bs, 32, lat_hw[0], lat_hw[1], device=DEV)
+                t = torch.rand(bs, device=DEV)
+                if stream is not None:                 # в потоке текста нет — берём фиктивный
+                    hs_probe, mk_probe = enc.slices(["probe caption"] * bs)
+                else:
+                    hs_probe, mk_probe = HS[:bs], MK[:bs]
+                with torch.autocast(DEV, torch.bfloat16, enabled=DEV == "cuda"):
+                    m.body(z, t, enc.mix(hs_probe.float()), mk_probe).float().pow(2).mean().backward()
+                best = bs
+                m.zero_grad(set_to_none=True)
+            except torch.OutOfMemoryError:
+                m.zero_grad(set_to_none=True); torch.cuda.empty_cache(); break
+        print(f"[probe] максимальный батч ≈ {best} (пик {torch.cuda.max_memory_allocated()/1e9:.1f} ГБ)",
+              flush=True)
+        if stream is not None:
+            stream.stop = True
+        import os as _os
+        _os._exit(0)
+
+    if stream is not None:
+        def get_batch():
+            x, caps = stream.batch(args.batch)
+            x = x.to(DEV)
+            with torch.no_grad():
+                z = ae.encode(x)
+                hs, mk = enc.slices(caps)
+            return x, z, hs.to(torch.bfloat16), mk
+    else:
+        def get_batch():
+            idx = torch.randint(0, n, (args.batch,), device=DEV)
+            return X[idx], Z[idx], HS[idx], MK[idx]
+
+    hf_gen, hf_ae, _, bright0, contrast0 = sample_grid(start)
+    print(f"[sample] шаг {start}: яркость {bright0:+.3f}, контраст {contrast0:.3f}, hf {hf_gen:.3f}", flush=True)
+
+    t0 = time.perf_counter(); prev_el = t0
+    ema = {}
+    run = {"lat": 0.0, "low": 0.0, "high": 0.0, "pix": 0.0, "n_low": 0, "n_high": 0, "n": 0}
+    skips = 0
+    for step in range(start + 1, start + steps + 1):
+        opt.zero_grad(set_to_none=True)
+        ok = 0
+        for _ in range(args.accum):
+            x0, z0, hs, mk = get_batch()
+            # дропаут подписи (CFG) и дропаут слов внутри подписи
+            if args.drop_text:
+                drop = torch.rand(args.batch, device=DEV) < args.drop_text
+                if drop.any():                      # пустая подпись обрезается под длину батча
+                    L = hs.shape[1]
+                    hs = torch.where(drop[:, None, None, None], HS_E[:, :L].expand_as(hs), hs)
+                    mk = torch.where(drop[:, None], MK_E[:, :L].expand_as(mk), mk)
+            if args.drop_word:
+                word = (torch.rand(mk.shape, device=DEV) < args.drop_word) & mk
+                hs = hs.masked_fill(word[:, :, None, None], 0)
+                mk = mk & ~word
+            with torch.autocast(DEV, torch.bfloat16, enabled=DEV == "cuda"):
+                txt = enc.mix(hs.float())
+                tb = torch.rand(args.batch, device=DEV)
+                eb = torch.randn_like(z0)
+                zt = (1 - tb[:, None, None, None]) * z0 + tb[:, None, None, None] * eb
+                z_hat = m.body(zt, tb, txt, mk)
+                per_elem = (z_hat.float() - z0.float()).pow(2).mean(dim=(1, 2, 3))
+                low, high = tb < 0.3, tb > 0.7
+                loss_l = per_elem.mean()
+                raw = m.to_raw(z0 + (z_hat - z0).detach())
+                tx = torch.rand(args.batch, device=DEV)
+                ex = torch.randn_like(x0)
+                xt = (1 - tx[:, None, None, None]) * x0 + tx[:, None, None, None] * ex
+                xt_packed = pkx.pack([patchify(xt[i].to(ae.compute_dtype), 8) for i in range(args.batch)])
+                z_packed = pkx.pack([raw[i].flatten(1).transpose(0, 1) for i in range(args.batch)])
+                preds = ae.decoder(xt_packed, tx, z_packed, pkx)
+                pix_terms = torch.stack([(p.float() - x0[i]).pow(2).mean() for i, p in enumerate(preds)])
+                loss_p = pix_terms.mean()
+            total = loss_l + loss_p
+            if not torch.isfinite(total):
+                skips += 1
+                continue
+            (total / args.accum).backward()
+            ok += 1
+            run["lat"] += loss_l.item(); run["pix"] += loss_p.item()
+            if low.any():
+                run["low"] += per_elem[low].mean().item(); run["n_low"] += 1
+            if high.any():
+                run["high"] += per_elem[high].mean().item(); run["n_high"] += 1
+        if not ok:
+            print(f"  ! шаг {step}: не-конечный лосс, пропуск (всего {skips})", flush=True)
+            continue
+        run["n"] += 1
+        torch.nn.utils.clip_grad_norm_(params, 1.0)
+        lr = lr_at(step, args.lr, lr_min, args.lr_warmup, period, start)
+        for g, base in zip(opt.param_groups, (args.lr, args.lr, args.lr_tail)):
+            g["lr"] = lr * (base / args.lr)
+        opt.step()
+
+        if step % args.log_every == 0:
+            d = args.ema_decay
+            cur = {"lat": run["lat"] / run["n"], "pix": run["pix"] / run["n"],
+                   "low": run["low"] / max(run["n_low"], 1), "high": run["high"] / max(run["n_high"], 1)}
+            for k, v in cur.items():
+                ema[k] = v if k not in ema else d * ema[k] + (1 - d) * v
+            el = time.perf_counter() - t0; inc = el - prev_el; prev_el = el
+            print(f"  {step:6d} | mse_lat {cur['lat']:.4f} (ema {ema['lat']:.4f}; low {cur['low']:.3f} "
+                  f"high {cur['high']:.3f}) | mse_pix {cur['pix']:.4f} | lr {lr:.1e} | {el:6.0f}s "
+                  f"({inc / max(run['n'], 1) * 1e3:.0f} мс/шаг) | skip {skips}", flush=True)
+            run = {"lat": 0.0, "low": 0.0, "high": 0.0, "pix": 0.0, "n_low": 0, "n_high": 0, "n": 0}
+
+        if args.sample_every and step % args.sample_every == 0:
+            zs = []
+            lat_pc_min = float("nan")
+            for txt, mk in show_txt[:2]:
+                z = torch.randn(1, 32, lat_hw[0], lat_hw[1], device=DEV,
+                                generator=torch.Generator(DEV).manual_seed(7))
+                with torch.no_grad(), torch.autocast(DEV, torch.bfloat16, enabled=DEV == "cuda"):
+                    for i in range(args.sample_steps):
+                        t = torch.full((1,), 1.0 - i / args.sample_steps, device=DEV)
+                        o = m.body(z, t, txt, mk).float()
+                        ou = m.body(z, t, enc.mix(HS_E.float()), MK_E).float()
+                        o = ou + args.cfg * (o - ou)
+                        z = z - (z - o) / t.clamp(min=1e-2)[:, None, None, None] / args.sample_steps
+                zs.append(z)
+            switch = (((zs[0] - zs[1]).norm() / zs[0].norm()).item() if len(zs) > 1 else float("nan"))
+            if zs:
+                lat_pc_min = float(zs[0].std(dim=(0, 2, 3)).min())   # сторож сжатия дисперсии латента
+            hfg, hfa, canvas, bright, contrast = sample_grid(step)
+            if show_x:
+                z_ref = ae.encode(show_x[0])
+                tail_psnr = psnr(ae.decode_raw(m.to_raw(z_ref), steps=args.tail_steps,
+                                               generator=torch.Generator(DEV).manual_seed(0)), show_x[0])
+            else:
+                tail_psnr = float("nan")
+            el = time.perf_counter() - t0
+            print(f"  ↳ switch {switch:.3f} | яркость {bright:+.3f} контраст {contrast:.3f} | "
+                  f"латент: std {lat_pc_min:.2f} (у данных ~0.8–1.3) | hf {hfg:.3f} | tail(z0) {tail_psnr:.2f}",
+                  flush=True)
+            with open(log_path, "a") as fh:
+                fh.write(f"{step},{ema.get('lat', 0):.5f},{ema.get('low', 0):.5f},{ema.get('high', 0):.5f},"
+                         f"{ema.get('pix', 0):.5f},{switch:.4f},{bright:.4f},{contrast:.4f},"
+                         f"{lat_pc_min:.4f},{hfg:.4f},{tail_psnr:.2f},{el:.1f}\n")
+        if args.save_every and step % args.save_every == 0:
+            save(step)
+    save(start + steps)
+    if stream is not None:
+        stream.stop = True          # фоновые потоки не должны переживать интерпретатор
+    print("[done]", flush=True)
+    import os as _os
+    _os._exit(0)                    # обойти шумную финализацию streaming-генераторов
+
+
+if __name__ == "__main__":
+    main()

text_encoder/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

text_encoder/README.md

@@ -0,0 +1,971 @@
+---
+language:
+- en
+license: mit
+pipeline_tags:
+- fill-mask
+- feature-extraction
+library_name: transformers
+tags:
+- modernbert
+- ettin
+- encoder
+- text-embeddings
+- retrieval
+- classification
+---
+
+# Ettin: an Open Suite of Paired Encoders and Decoders
+
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Paper](https://img.shields.io/badge/Paper-Arxiv-red)](https://arxiv.org/abs/2507.11412)
+[![Models](https://img.shields.io/badge/🤗%20Hugging%20Face-12%20Models-blue)](https://huggingface.co/jhu-clsp)
+[![Data](https://img.shields.io/badge/🤗%20Training%20Data-2T%20Tokens-green)](https://huggingface.co/datasets/jhu-clsp)
+[![GitHub](https://img.shields.io/badge/GitHub-Code-black)](https://github.com/jhu-clsp/ettin-encoder-vs-decoder)
+
+> 🎯 **TL;DR**: State-of-the-art paired encoder and decoder models (17M-1B params) trained identically for fair comparison with open data. Encoders beat ModernBERT. Decoders beat Llama 3.2/SmolLM2.
+
+📄 [Paper](https://arxiv.org/abs/2507.11412) | 🚀 [GitHub Repository](https://github.com/jhu-clsp/ettin-encoder-vs-decoder)
+
+This model is part of the Ettin suite - the first collection of paired encoder-only and decoder-only models trained with identical data, architecture, and training recipes. Ettin enables fair comparisons between encoder and decoder architectures across multiple scales, providing state-of-the-art performance for open-data models in their respective size categories.
+
+## Table of Contents
+- [Performance Highlights](#performance-highlights)
+- [Quick Start](#quick-start)
+- [Model Description](#model-description)
+- [Training Data](#training-data)
+- [Model Family](#model-family)
+  - [Encoder Models](#encoder-models)
+  - [Decoder Models](#decoder-models)
+  - [Cross-Objective Models](#cross-objective-models)
+- [Accessing Training Checkpoints](#accessing-training-checkpoints)
+- [Research Applications](#research-applications)
+- [Training Details](#training-details)
+- [Model Architecture](#model-architecture)
+- [Usage Examples](#usage-examples)
+- [Fine-tuning Examples](#fine-tuning-examples)
+- [Citation](#citation)
+
+## 📊 Performance Highlights
+
+### Encoder Tasks (vs. ModernBERT)
+- **GLUE Average**: 88.9 vs 88.4 (Base), 90.8 vs 90.4 (Large)
+- **MTEB v2 English Retrieval**: 45.7 vs 43.9 (Base), 48.4 vs 47.0 (Large)
+- **Code Search and Long Context**: Superior performance on CodeSearchNet and MLDR
+
+### Decoder Tasks (vs. SmolLM2 & Llama 3.2)
+- **Average Score**: 46.2 vs 45.2 (SmolLM2-135M)
+- **1B Model**: 59.0 vs 56.6 (Llama 3.2-1B)
+- **Generative Tasks**: Competitive across all model sizes
+
+### Key Finding
+**Architecture-specific advantages persist**: A 400M encoder outperforms a 1B decoder on classification tasks, while a 400M decoder outperforms a 1B encoder on generation tasks.
+
+## 🚀 Quick Start
+
+### Installation
+```bash
+pip install torch>=1.9.0
+# until the new pip release, install from main to use decoders (transformers>=4.54.X will contain it)
+# encoders work with transformers>=4.48.0
+pip install git+https://github.com/huggingface/transformers.git
+```
+
+### 30-Second Examples
+
+**Encoder for Classification/Embeddings:**
+```python
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-encoder-150m")
+model = AutoModel.from_pretrained("jhu-clsp/ettin-encoder-150m")
+```
+
+**Decoder for Text Generation:**
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-decoder-150m")
+model = AutoModelForCausalLM.from_pretrained("jhu-clsp/ettin-decoder-150m")
+```
+
+## Model Description
+
+Ettin models are designed to provide a foundation for comparing encoder-only and decoder-only architectures. Unlike previous comparisons that were limited by different training data, architectures, and recipes, Ettin models use:
+
+1. **Identical training data** - Same high-quality mixture across all models
+2. **Open Training Data** - Data is available now with batch-level training data for each of the 250+ checkpoints
+3. **Matched architectures** - Only differing in attention patterns (bidirectional vs causal) and training objectives (MLM vs CLM)
+4. **Consistent training recipe** - Three-phase training with 2T tokens
+5. **Multiple scales** - From 17M to 1B parameters
+
+This approach allows for true apples-to-apples comparisons between encoder and decoder models, revealing the inherent strengths of each architecture.
+
+## Training Data
+
+The training data is publicly available and split across different phases:
+
+- **Pre-training Data**: [jhu-clsp/ettin-pretraining-data](https://huggingface.co/datasets/jhu-clsp/ettin-pretraining-data) - 1.7T tokens of diverse data mixture
+- **Mid-training/Extension Data**: [jhu-clsp/ettin-extension-data](https://huggingface.co/datasets/jhu-clsp/ettin-extension-data) - 250B tokens of higher-quality filtered data
+- **Decay Phase Data**: [jhu-clsp/ettin-decay-data](https://huggingface.co/datasets/jhu-clsp/ettin-decay-data) - 100B tokens of premium data sources
+- **Training Data Order**: [jhu-clsp/ettin-data-order](https://huggingface.co/datasets/jhu-clsp/ettin-data-order) - Batch-level training order (columns: input_ids, step)
+
+## Model Family
+
+### Encoder Models
+
+| Size | Model | Parameters | Best For | Download |
+|:-----|:------|:-----------|:---------|:---------|
+| XXS | [ettin-encoder-17m](https://huggingface.co/jhu-clsp/ettin-encoder-17m) | 17M | Mobile/Edge devices | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-17m) |
+| XS | [ettin-encoder-32m](https://huggingface.co/jhu-clsp/ettin-encoder-32m) | 32M | Fast inference | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-32m) |
+| Small | [ettin-encoder-68m](https://huggingface.co/jhu-clsp/ettin-encoder-68m) | 68M | Balanced performance | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-68m) |
+| Base | [ettin-encoder-150m](https://huggingface.co/jhu-clsp/ettin-encoder-150m) | 150M | Standard use cases | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-150m) |
+| Large | [ettin-encoder-400m](https://huggingface.co/jhu-clsp/ettin-encoder-400m) | 400M | High accuracy needs | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-400m) |
+| XL | [ettin-encoder-1b](https://huggingface.co/jhu-clsp/ettin-encoder-1b) | 1B | Best performance | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-1b) |
+
+### Decoder Models
+
+| Size | Model | Parameters | Best For | Download |
+|:-----|:------|:-----------|:---------|:---------|
+| XXS | [ettin-decoder-17m](https://huggingface.co/jhu-clsp/ettin-decoder-17m) | 17M | Lightweight generation | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-17m) |
+| XS | [ettin-decoder-32m](https://huggingface.co/jhu-clsp/ettin-decoder-32m) | 32M | Quick prototyping | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-32m) |
+| Small | [ettin-decoder-68m](https://huggingface.co/jhu-clsp/ettin-decoder-68m) | 68M | Efficient generation | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-68m) |
+| Base | [ettin-decoder-150m](https://huggingface.co/jhu-clsp/ettin-decoder-150m) | 150M | Standard generation | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-150m) |
+| Large | [ettin-decoder-400m](https://huggingface.co/jhu-clsp/ettin-decoder-400m) | 400M | Quality generation | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-400m) |
+| XL | [ettin-decoder-1b](https://huggingface.co/jhu-clsp/ettin-decoder-1b) | 1B | Best generation | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-1b) |
+
+### Cross-Objective Models
+
+These models demonstrate what happens when you continue training encoders as decoders (and vice versa). **Important**: Load these models using the architecture they were *converted to*, not their original architecture.
+
+#### Encoders Trained from Decoders (Decoder → MLM)
+**Load as encoders** using `AutoModel` or `AutoModelForMaskedLM`:
+
+| Size | Model | Parameters | Description | Download |
+|:-----|:------|:-----------|:------------|:---------|
+| XXS | [ettin-encoder-from-decoder-17m](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-17m) | 17M | Decoder → MLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-17m) |
+| XS | [ettin-encoder-from-decoder-32m](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-32m) | 32M | Decoder → MLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-32m) |
+| Small | [ettin-encoder-from-decoder-68m](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-68m) | 68M | Decoder → MLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-68m) |
+| Base | [ettin-encoder-from-decoder-150m](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-150m) | 150M | Decoder → MLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-150m) |
+| Large | [ettin-encoder-from-decoder-400m](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-400m) | 400M | Decoder → MLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-400m) |
+| XL | [ettin-encoder-from-decoder-1b](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-1b) | 1B | Decoder → MLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-encoder-from-decoder-1b) |
+
+#### Decoders Trained from Encoders (Encoder → CLM)
+**Load as decoders** using `AutoModelForCausalLM`:
+
+| Size | Model | Parameters | Description | Download |
+|:-----|:------|:-----------|:------------|:---------|
+| XXS | [ettin-decoder-from-encoder-17m](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-17m) | 17M | Encoder → CLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-17m) |
+| XS | [ettin-decoder-from-encoder-32m](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-32m) | 32M | Encoder → CLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-32m) |
+| Small | [ettin-decoder-from-encoder-68m](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-68m) | 68M | Encoder → CLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-68m) |
+| Base | [ettin-decoder-from-encoder-150m](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-150m) | 150M | Encoder → CLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-150m) |
+| Large | [ettin-decoder-from-encoder-400m](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-400m) | 400M | Encoder → CLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-400m) |
+| XL | [ettin-decoder-from-encoder-1b](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-1b) | 1B | Encoder → CLM continued training | [![Download](https://img.shields.io/badge/🤗-Download-blue)](https://huggingface.co/jhu-clsp/ettin-decoder-from-encoder-1b) |
+
+**Example Usage for Cross-Objective Models:**
+```python
+# Encoder-from-decoder: Load as encoder
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-encoder-from-decoder-150m")
+model = AutoModel.from_pretrained("jhu-clsp/ettin-encoder-from-decoder-150m")
+
+# Decoder-from-encoder: Load as decoder  
+from transformers import AutoTokenizer, AutoModelForCausalLM
+tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-decoder-from-encoder-150m")
+model = AutoModelForCausalLM.from_pretrained("jhu-clsp/ettin-decoder-from-encoder-150m")
+```
+
+## Accessing Training Checkpoints
+
+Beyond the final models listed above, we provide access to intermediate training checkpoints for research and analysis purposes. These checkpoints allow you to study model behavior and performance throughout the training process. You can get the checkpoints either in HF format or raw for continued pre-training (e.g. Composer format).
+
+#### Raw Checkpoints
+All raw training checkpoints are available in the [jhu-clsp/ettin-checkpoints](https://huggingface.co/datasets/jhu-clsp/ettin-checkpoints) dataset.
+
+#### HuggingFace Format Checkpoints
+Each model repository contains multiple tagged versions representing different training stages:
+
+- **`step{number}`** - Pretraining phase checkpoints (e.g., `step599525`, `step596528`)
+- **`ext{number}`** - Extension/mid-training phase checkpoints (e.g., `ext1000`, `ext2000`) 
+- **`decay{number}`** - Decay phase checkpoints (e.g., `decay100`, `decay500`)
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Load a specific pretraining checkpoint
+model = AutoModelForCausalLM.from_pretrained(
+    "jhu-clsp/ettin-decoder-400m", 
+    revision="step590532"  # Specific checkpoint tag
+)
+
+# Load an extension phase checkpoint
+model = AutoModelForCausalLM.from_pretrained(
+    "jhu-clsp/ettin-decoder-400m", 
+    revision="ext1000"
+)
+
+# Load a decay phase checkpoint  
+model = AutoModelForCausalLM.from_pretrained(
+    "jhu-clsp/ettin-decoder-400m", 
+    revision="decay100"
+)
+```
+
+This checkpoint availability enables detailed analysis of training dynamics, loss curves, and capability emergence across the complete 2T token training process.
+
+
+## 🔬 Research Applications
+
+### What Makes Ettin Unique
+
+Ettin provides the first **controlled comparison** of encoder vs. decoder architectures:
+
+- **Identical Training Data**: Same 2T token mixture across all models
+- **Matched Architectures**: Only attention patterns and objectives differ  
+- **Open Everything**: Training data, model weights, and batch-level training order
+- **Multiple Scales**: Fair comparison from 17M to 1B parameters
+- **250+ Checkpoints**: Complete training trajectory analysis
+
+### Use Cases for Researchers
+
+- **Architecture Studies**: Compare encoder vs decoder capabilities fairly
+- **Training Dynamics**: Analyze 250+ checkpoints with batch-level data ordering  
+- **Scaling Laws**: Study how architectural advantages change with scale
+- **Transfer Learning**: Investigate cross-objective training effectiveness
+- **Replication Studies**: First open replication of ModernBERT training recipe
+
+### Reproducibility
+
+All training artifacts are publicly available:
+- Training data with exact batch ordering
+- Model checkpoints every 8.5B tokens
+- Complete hyperparameter configurations
+- Training code and evaluation scripts
+
+## Training Details
+
+**Data:** High-quality mixture including DCLM, Dolma v1.7, scientific papers, code, and curated sources totaling 2T+ tokens
+
+**Architecture:** Transformer with RoPE, GLU activations, and prenorm layers
+
+**Training Phases:**
+- **Pre-training**: 1.7T tokens with diverse data mixture
+- **Mid-training**: 250B tokens with higher-quality filtered data and context extension to 8K
+- **Decay phase**: 100B tokens with premium data sources
+
+**Key Features:**
+- Context length: Up to 8K tokens
+- Vocabulary: 50,368 tokens (ModernBERT tokenizer)
+- Deep but efficient architectures following MobileLLM principles
+
+## Model Architecture
+
+| Parameter | 17M | 32M | 68M | 150M | 400M | 1B |
+|:----------|:----|:----|:----|:-----|:-----|:---|
+| Layers | 7 | 10 | 19 | 22 | 28 | 28 |
+| Hidden Size | 256 | 384 | 512 | 768 | 1024 | 1792 |
+| Intermediate Size | 384 | 576 | 768 | 1152 | 2624 | 3840 |
+| Attention Heads | 4 | 6 | 8 | 12 | 16 | 28 |
+
+
+
+## Usage Examples
+
+### Encoder: Masked Language Modeling
+<details>
+<summary>Click to expand **encoder** usage examples</summary>
+
+```python
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+import torch
+
+# Load MLM model
+tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-encoder-150m")
+model = AutoModelForMaskedLM.from_pretrained("jhu-clsp/ettin-encoder-150m")
+
+def predict_masked_token(text):
+    inputs = tokenizer(text, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # Get predictions for [MASK] tokens
+    mask_indices = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)
+    predictions = outputs.logits[mask_indices]
+    
+    # Get top 5 predictions
+    top_tokens = torch.topk(predictions, 5, dim=-1)
+    return [tokenizer.decode(token) for token in top_tokens.indices[0]]
+
+# Example
+masked_text = "The capital of France is [MASK]."
+predictions = predict_masked_token(masked_text)
+print(f"Predictions: {predictions}")
+```
+
+</details>
+
+### Decoder: Text Generation
+
+<details>
+<summary>Click to expand **decoder text generation**</summary>
+  
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+# Load model and tokenizer  
+tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-decoder-150m")
+model = AutoModelForCausalLM.from_pretrained("jhu-clsp/ettin-decoder-150m")
+
+# Set pad token if needed
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+def generate_text(prompt, max_length=100, temperature=0.7):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs.input_ids,
+            max_length=max_length,
+            temperature=temperature,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            num_return_sequences=1
+        )
+    
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Example usage
+prompt = "The future of artificial intelligence is"
+generated = generate_text(prompt)
+print(generated)
+```
+
+</details>
+
+
+## Fine-tuning Examples
+
+### Encoders
+<details><summary>Click to see how to finetune this into a dense embedding model using Sentence Transformers</summary> 
+
+```python
+import argparse
+
+from datasets import load_dataset
+from sentence_transformers import (
+    SentenceTransformer,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.evaluation import TripletEvaluator
+from sentence_transformers.losses import CachedMultipleNegativesRankingLoss
+from sentence_transformers.training_args import BatchSamplers
+
+def main():
+    # parse the lr & model name
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--lr", type=float, default=8e-5)
+    parser.add_argument("--model_name", type=str, default="jhu-clsp/ettin-encoder-150m")
+    args = parser.parse_args()
+    lr = args.lr
+    model_name = args.model_name
+    model_shortname = model_name.split("/")[-1]
+
+    # 1. Load a model to finetune
+    model = SentenceTransformer(model_name)
+
+    # 2. Load a dataset to finetune on
+    dataset = load_dataset(
+        "sentence-transformers/msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1",
+        "triplet-hard",
+        split="train",
+    )
+    dataset_dict = dataset.train_test_split(test_size=1_000, seed=12)
+    train_dataset = dataset_dict["train"].select(range(1_250_000))
+    eval_dataset = dataset_dict["test"]
+
+    # 3. Define a loss function
+    loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=16)  # Increase mini_batch_size if you have enough VRAM
+
+    run_name = f"{model_shortname}-DPR-{lr}"
+    # 4. (Optional) Specify training arguments
+    args = SentenceTransformerTrainingArguments(
+        # Required parameter:
+        output_dir=f"output/{model_shortname}/{run_name}",
+        # Optional training parameters:
+        num_train_epochs=1,
+        per_device_train_batch_size=512,
+        per_device_eval_batch_size=512,
+        warmup_ratio=0.05,
+        fp16=False,  # Set to False if GPU can't handle FP16
+        bf16=True,  # Set to True if GPU supports BF16
+        batch_sampler=BatchSamplers.NO_DUPLICATES,  # (Cached)MultipleNegativesRankingLoss benefits from no duplicates
+        learning_rate=lr,
+        # Optional tracking/debugging parameters:
+        save_strategy="steps",
+        save_steps=500,
+        save_total_limit=2,
+        logging_steps=500,
+        run_name=run_name,  # Used in `wandb`, `tensorboard`, `neptune`, etc. if installed
+    )
+
+    # 5. (Optional) Create an evaluator & evaluate the base model
+    dev_evaluator = TripletEvaluator(
+        anchors=eval_dataset["query"],
+        positives=eval_dataset["positive"],
+        negatives=eval_dataset["negative"],
+        name="msmarco-co-condenser-dev",
+    )
+    dev_evaluator(model)
+
+    # 6. Create a trainer & train
+    trainer = SentenceTransformerTrainer(
+        model=model,
+        args=args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        loss=loss,
+        evaluator=dev_evaluator,
+    )
+    trainer.train()
+
+    # 7. (Optional) Evaluate the trained model on the evaluator after training
+    dev_evaluator(model)
+
+    # 8. Save the model
+    model.save_pretrained(f"output/{model_shortname}/{run_name}/final")
+
+    # 9. (Optional) Push it to the Hugging Face Hub
+    model.push_to_hub(run_name, private=False)
+
+if __name__ == "__main__":
+    main()
+```
+</details>
+
+
+<details><summary>Click to see how to finetune this into a multi-vector embedding model with PyLate</summary>
+
+```python
+from datasets import load_dataset
+from pylate import losses, models, utils
+from sentence_transformers import (
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+
+def main():
+    # Load the datasets required for knowledge distillation (train, queries, documents)
+    train = load_dataset(
+        path="lightonai/ms-marco-en-bge",
+        name="train",
+    )
+
+    queries = load_dataset(
+        path="lightonai/ms-marco-en-bge",
+        name="queries",
+    )
+
+    documents = load_dataset(
+        path="lightonai/ms-marco-en-bge",
+        name="documents",
+    )
+
+    # Set the transformation to load the documents/queries texts using the corresponding ids on the fly
+    train.set_transform(
+        utils.KDProcessing(queries=queries, documents=documents).transform,
+    )
+
+    # Define the base model, training parameters, and output directory
+    num_train_epochs = 1
+    lr = 8e-5
+    batch_size = 16
+    accum_steps = 1
+    model_name = "jhu-clsp/ettin-encoder-150m"
+    model_shortname = model_name.split("/")[-1]
+
+    # Set the run name for logging and output directory
+    run_name = f"{model_shortname}-colbert-KD-{lr}"
+    output_dir = f"output/{model_shortname}/{run_name}"
+
+    # Initialize the ColBERT model from the base model
+    model = models.ColBERT(model_name_or_path=model_name)
+
+    # Configure the training arguments (e.g., epochs, batch size, learning rate)
+    args = SentenceTransformerTrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=num_train_epochs,
+        per_device_train_batch_size=batch_size,
+        fp16=False,  # Set to False if you get an error that your GPU can't run on FP16
+        bf16=True,  # Set to True if you have a GPU that supports BF16
+        run_name=run_name,
+        logging_steps=10,
+        learning_rate=lr,
+        gradient_accumulation_steps=accum_steps,
+        warmup_ratio=0.05,
+    )
+
+    # Use the Distillation loss function for training
+    train_loss = losses.Distillation(model=model)
+
+    # Initialize the trainer
+    trainer = SentenceTransformerTrainer(
+        model=model,
+        args=args,
+        train_dataset=train,
+        loss=train_loss,
+        data_collator=utils.ColBERTCollator(tokenize_fn=model.tokenize),
+    )
+
+    # Start the training process
+    trainer.train()
+
+    model.save_pretrained(f"{output_dir}/final")
+
+if __name__ == "__main__":
+    main()
+
+```
+</details>
+
+<details><summary>Click to see how to finetune this into a sparse retrieval model using Sentence Transformers</summary>
+
+```python
+import logging
+
+from datasets import load_dataset
+
+from sentence_transformers import (
+    SparseEncoder,
+    SparseEncoderModelCardData,
+    SparseEncoderTrainer,
+    SparseEncoderTrainingArguments,
+)
+from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
+from sentence_transformers.sparse_encoder.losses import SparseMultipleNegativesRankingLoss, SpladeLoss
+from sentence_transformers.training_args import BatchSamplers
+
+logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)
+
+# 1. Load a model to finetune with 2. (Optional) model card data
+model = SparseEncoder(
+    "jhu-clsp/ettin-encoder-150m",
+    model_card_data=SparseEncoderModelCardData(
+        language="en",
+        license="apache-2.0",
+    )
+)
+
+# 3. Load a dataset to finetune on
+full_dataset = load_dataset("sentence-transformers/natural-questions", split="train").select(range(100_000))
+dataset_dict = full_dataset.train_test_split(test_size=1_000, seed=12)
+train_dataset = dataset_dict["train"]
+eval_dataset = dataset_dict["test"]
+
+# 4. Define a loss function
+loss = SpladeLoss(
+    model=model,
+    loss=SparseMultipleNegativesRankingLoss(model=model),
+    query_regularizer_weight=5e-5,
+    document_regularizer_weight=3e-5,
+)
+
+# 5. (Optional) Specify training arguments
+run_name = "splade-distilbert-base-uncased-nq"
+args = SparseEncoderTrainingArguments(
+    # Required parameter:
+    output_dir=f"models/{run_name}",
+    # Optional training parameters:
+    num_train_epochs=1,
+    per_device_train_batch_size=16,
+    per_device_eval_batch_size=16,
+    learning_rate=2e-5,
+    warmup_ratio=0.1,
+    fp16=True,  # Set to False if you get an error that your GPU can't run on FP16
+    bf16=False,  # Set to True if you have a GPU that supports BF16
+    batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch
+    # Optional tracking/debugging parameters:
+    eval_strategy="steps",
+    eval_steps=1000,
+    save_strategy="steps",
+    save_steps=1000,
+    save_total_limit=2,
+    logging_steps=200,
+    run_name=run_name,  # Will be used in W&B if `wandb` is installed
+)
+
+# 6. (Optional) Create an evaluator & evaluate the base model
+dev_evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"], batch_size=16)
+
+# 7. Create a trainer & train
+trainer = SparseEncoderTrainer(
+    model=model,
+    args=args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    loss=loss,
+    evaluator=dev_evaluator,
+)
+trainer.train()
+
+# 8. Evaluate the model performance again after training
+dev_evaluator(model)
+
+# 9. Save the trained model
+model.save_pretrained(f"models/{run_name}/final")
+
+# 10. (Optional) Push it to the Hugging Face Hub
+model.push_to_hub(run_name)
+
+```
+</details>
+
+<details><summary>Click to see how to finetune this into a reranker model using Sentence Transformers</summary>
+
+```python
+import logging
+import traceback
+
+import torch
+from datasets import load_dataset
+
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.cross_encoder import (
+    CrossEncoder,
+    CrossEncoderModelCardData,
+    CrossEncoderTrainer,
+    CrossEncoderTrainingArguments,
+)
+from sentence_transformers.cross_encoder.evaluation import (
+    CrossEncoderNanoBEIREvaluator,
+    CrossEncoderRerankingEvaluator,
+)
+from sentence_transformers.cross_encoder.losses import BinaryCrossEntropyLoss
+from sentence_transformers.evaluation import SequentialEvaluator
+from sentence_transformers.util import mine_hard_negatives
+
+# Set the log level to INFO to get more information
+logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)
+
+
+def main():
+    model_name = "jhu-clsp/ettin-encoder-150m"
+
+    train_batch_size = 64
+    num_epochs = 1
+    num_hard_negatives = 5  # How many hard negatives should be mined for each question-answer pair
+
+    # 1a. Load a model to finetune with 1b. (Optional) model card data
+    model = CrossEncoder(
+        model_name,
+        model_card_data=CrossEncoderModelCardData(
+            language="en",
+            license="apache-2.0",
+        ),
+    )
+    print("Model max length:", model.max_length)
+    print("Model num labels:", model.num_labels)
+
+    # 2a. Load the GooAQ dataset: https://huggingface.co/datasets/sentence-transformers/gooaq
+    logging.info("Read the gooaq training dataset")
+    full_dataset = load_dataset("sentence-transformers/gooaq", split="train").select(range(100_000))
+    dataset_dict = full_dataset.train_test_split(test_size=1_000, seed=12)
+    train_dataset = dataset_dict["train"]
+    eval_dataset = dataset_dict["test"]
+    logging.info(train_dataset)
+    logging.info(eval_dataset)
+
+    # 2b. Modify our training dataset to include hard negatives using a very efficient embedding model
+    embedding_model = SentenceTransformer("sentence-transformers/static-retrieval-mrl-en-v1", device="cpu")
+    hard_train_dataset = mine_hard_negatives(
+        train_dataset,
+        embedding_model,
+        num_negatives=num_hard_negatives,  # How many negatives per question-answer pair
+        margin=0,  # Similarity between query and negative samples should be x lower than query-positive similarity
+        range_min=0,  # Skip the x most similar samples
+        range_max=100,  # Consider only the x most similar samples
+        sampling_strategy="top",  # Sample the top negatives from the range
+        batch_size=4096,  # Use a batch size of 4096 for the embedding model
+        output_format="labeled-pair",  # The output format is (query, passage, label), as required by BinaryCrossEntropyLoss
+        use_faiss=True,
+    )
+    logging.info(hard_train_dataset)
+
+    # 2c. (Optionally) Save the hard training dataset to disk
+    # hard_train_dataset.save_to_disk("gooaq-hard-train")
+    # Load again with:
+    # hard_train_dataset = load_from_disk("gooaq-hard-train")
+
+    # 3. Define our training loss.
+    # pos_weight is recommended to be set as the ratio between positives to negatives, a.k.a. `num_hard_negatives`
+    loss = BinaryCrossEntropyLoss(model=model, pos_weight=torch.tensor(num_hard_negatives))
+
+    # 4a. Define evaluators. We use the CrossEncoderNanoBEIREvaluator, which is a light-weight evaluator for English reranking
+    nano_beir_evaluator = CrossEncoderNanoBEIREvaluator(
+        dataset_names=["msmarco", "nfcorpus", "nq"],
+        batch_size=train_batch_size,
+    )
+
+    # 4b. Define a reranking evaluator by mining hard negatives given query-answer pairs
+    # We include the positive answer in the list of negatives, so the evaluator can use the performance of the
+    # embedding model as a baseline.
+    hard_eval_dataset = mine_hard_negatives(
+        eval_dataset,
+        embedding_model,
+        corpus=full_dataset["answer"],  # Use the full dataset as the corpus
+        num_negatives=30,  # How many documents to rerank
+        batch_size=4096,
+        include_positives=True,
+        output_format="n-tuple",
+        use_faiss=True,
+    )
+    logging.info(hard_eval_dataset)
+    reranking_evaluator = CrossEncoderRerankingEvaluator(
+        samples=[
+            {
+                "query": sample["question"],
+                "positive": [sample["answer"]],
+                "documents": [sample[column_name] for column_name in hard_eval_dataset.column_names[2:]],
+            }
+            for sample in hard_eval_dataset
+        ],
+        batch_size=train_batch_size,
+        name="gooaq-dev",
+        # Realistic setting: only rerank the positives that the retriever found
+        # Set to True to rerank *all* positives
+        always_rerank_positives=False,
+    )
+
+    # 4c. Combine the evaluators & run the base model on them
+    evaluator = SequentialEvaluator([reranking_evaluator, nano_beir_evaluator])
+    evaluator(model)
+
+    # 5. Define the training arguments
+    short_model_name = model_name if "/" not in model_name else model_name.split("/")[-1]
+    run_name = f"reranker-{short_model_name}-gooaq-bce"
+    args = CrossEncoderTrainingArguments(
+        # Required parameter:
+        output_dir=f"models/{run_name}",
+        # Optional training parameters:
+        num_train_epochs=num_epochs,
+        per_device_train_batch_size=train_batch_size,
+        per_device_eval_batch_size=train_batch_size,
+        learning_rate=2e-5,
+        warmup_ratio=0.1,
+        fp16=False,  # Set to False if you get an error that your GPU can't run on FP16
+        bf16=True,  # Set to True if you have a GPU that supports BF16
+        dataloader_num_workers=4,
+        load_best_model_at_end=True,
+        metric_for_best_model="eval_gooaq-dev_ndcg@10",
+        # Optional tracking/debugging parameters:
+        eval_strategy="steps",
+        eval_steps=1000,
+        save_strategy="steps",
+        save_steps=1000,
+        save_total_limit=2,
+        logging_steps=200,
+        logging_first_step=True,
+        run_name=run_name,  # Will be used in W&B if `wandb` is installed
+        seed=12,
+    )
+
+    # 6. Create the trainer & start training
+    trainer = CrossEncoderTrainer(
+        model=model,
+        args=args,
+        train_dataset=hard_train_dataset,
+        loss=loss,
+        evaluator=evaluator,
+    )
+    trainer.train()
+
+    # 7. Evaluate the final model, useful to include these in the model card
+    evaluator(model)
+
+    # 8. Save the final model
+    final_output_dir = f"models/{run_name}/final"
+    model.save_pretrained(final_output_dir)
+
+    # 9. (Optional) save the model to the Hugging Face Hub!
+    # It is recommended to run `huggingface-cli login` to log into your Hugging Face account first
+    try:
+        model.push_to_hub(run_name)
+    except Exception:
+        logging.error(
+            f"Error uploading model to the Hugging Face Hub:
+{traceback.format_exc()}To upload it manually, you can run "
+            f"`huggingface-cli login`, followed by loading the model using `model = CrossEncoder({final_output_dir!r})` "
+            f"and saving it using `model.push_to_hub('{run_name}')`."
+        )
+
+
+if __name__ == "__main__":
+    main()
+
+```
+</details>
+
+### Decoders
+
+<details>
+<summary>Click to expand decoder training code</summary>
+
+# Full training
+```bash
+python trl/scripts/sft.py \
+    --model_name_or_path jhu-clsp/ettin-decoder-17m \
+    --dataset_name trl-lib/Capybara \
+    --learning_rate 2.0e-5 \
+    --num_train_epochs 1 \
+    --packing \
+    --per_device_train_batch_size 2 \
+    --gradient_accumulation_steps 8 \
+    --gradient_checkpointing \
+    --eos_token '<|im_end|>' \
+    --eval_strategy steps \
+    --eval_steps 100 \
+    --output_dir ettin-decoder-17m \
+    --push_to_hub
+```
+
+# LoRA
+```bash
+python trl/scripts/sft.py \
+    --model_name_or_path jhu-clsp/ettin-decoder-17m \
+    --dataset_name trl-lib/Capybara \
+    --learning_rate 2.0e-4 \
+    --num_train_epochs 1 \
+    --packing \
+    --per_device_train_batch_size 2 \
+    --gradient_accumulation_steps 8 \
+    --gradient_checkpointing \
+    --eos_token '<|im_end|>' \
+    --eval_strategy steps \
+    --eval_steps 100 \
+    --use_peft \
+    --lora_r 32 \
+    --lora_alpha 16 \
+    --output_dir ettin-decoder-17m \
+    --push_to_hub
+```
+
+with `sft.py`:
+```python
+import argparse
+
+from datasets import load_dataset
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES
+
+from trl import (
+    ModelConfig,
+    ScriptArguments,
+    SFTConfig,
+    SFTTrainer,
+    TrlParser,
+    clone_chat_template,
+    get_kbit_device_map,
+    get_peft_config,
+    get_quantization_config,
+)
+
+
+def main(script_args, training_args, model_args):
+    ################
+    # Model init kwargs & Tokenizer
+    ################
+    quantization_config = get_quantization_config(model_args)
+    model_kwargs = dict(
+        revision=model_args.model_revision,
+        trust_remote_code=model_args.trust_remote_code,
+        attn_implementation=model_args.attn_implementation,
+        torch_dtype=model_args.torch_dtype,
+        use_cache=False if training_args.gradient_checkpointing else True,
+        device_map=get_kbit_device_map() if quantization_config is not None else None,
+        quantization_config=quantization_config,
+    )
+
+    # Create model
+    config = AutoConfig.from_pretrained(model_args.model_name_or_path)
+    valid_image_text_architectures = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values()
+
+    if config.architectures and any(arch in valid_image_text_architectures for arch in config.architectures):
+        from transformers import AutoModelForImageTextToText
+
+        model_kwargs.pop("use_cache", None)  # Image models do not support cache
+        model = AutoModelForImageTextToText.from_pretrained(model_args.model_name_or_path, **model_kwargs)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, **model_kwargs)
+
+    # Create tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code, use_fast=True
+    )
+
+    # Set default chat template if needed
+    if tokenizer.chat_template is None:
+        # TODO: source should be passed as an argument
+        model, tokenizer = clone_chat_template(model, tokenizer, "Qwen/Qwen3-0.6B")
+
+    ################
+    # Dataset
+    ################
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
+
+    ################
+    # Training
+    ################
+    trainer = SFTTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset[script_args.dataset_train_split],
+        eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None,
+        processing_class=tokenizer,
+        peft_config=get_peft_config(model_args),
+    )
+
+    trainer.train()
+
+    # Save and push to hub
+    trainer.save_model(training_args.output_dir)
+    if training_args.push_to_hub:
+        trainer.push_to_hub(dataset_name=script_args.dataset_name)
+
+
+def make_parser(subparsers: argparse._SubParsersAction = None):
+    dataclass_types = (ScriptArguments, SFTConfig, ModelConfig)
+    if subparsers is not None:
+        parser = subparsers.add_parser("sft", help="Run the SFT training script", dataclass_types=dataclass_types)
+    else:
+        parser = TrlParser(dataclass_types)
+    return parser
+
+
+if __name__ == "__main__":
+    parser = make_parser()
+    # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments.
+    # To ensure that their parsing does not interfere with the script arguments, parse the arguments with
+    # `return_remaining_strings=True`, then ignore the remaining strings.
+    script_args, training_args, model_args, _ = parser.parse_args_and_config(return_remaining_strings=True)
+    main(script_args, training_args, model_args)
+
+```
+</details>
+
+## Citation
+
+If you use Ettin models in your research, please cite our work:
+
+```bibtex
+@misc{weller2025seqvsseqopen,
+      title={Seq vs Seq: An Open Suite of Paired Encoders and Decoders}, 
+      author={Orion Weller and Kathryn Ricci and Marc Marone and Antoine Chaffin and Dawn Lawrie and Benjamin Van Durme},
+      year={2025},
+      eprint={2507.11412},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2507.11412}, 
+}
+```

text_encoder/config.json

@@ -0,0 +1,46 @@
+{
+  "_name_or_path": "ettin-encoder-150m",
+  "architectures": [
+    "ModernBertForMaskedLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 50281,
+  "classifier_activation": "gelu",
+  "classifier_bias": false,
+  "classifier_dropout": 0.0,
+  "classifier_pooling": "mean",
+  "cls_token_id": 50281,
+  "decoder_bias": true,
+  "deterministic_flash_attn": false,
+  "embedding_dropout": 0.0,
+  "eos_token_id": 50282,
+  "global_attn_every_n_layers": 3,
+  "global_rope_theta": 160000.0,
+  "gradient_checkpointing": false,
+  "hidden_activation": "gelu",
+  "hidden_size": 768,
+  "initializer_cutoff_factor": 2.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 1152,
+  "layer_norm_eps": 1e-05,
+  "local_attention": 128,
+  "local_rope_theta": 160000.0,
+  "max_position_embeddings": 7999,
+  "mlp_bias": false,
+  "mlp_dropout": 0.0,
+  "model_type": "modernbert",
+  "norm_bias": false,
+  "norm_eps": 1e-05,
+  "num_attention_heads": 12,
+  "num_hidden_layers": 22,
+  "pad_token_id": 50283,
+  "position_embedding_type": "sans_pos",
+  "sep_token_id": 50282,
+  "tie_word_embeddings": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.47.0.dev0",
+  "vocab_size": 50368,
+  "is_causal": false,
+  "causal_mask": false
+}

text_encoder/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f95f652c607616a076c9e05c82f00f7372ce332de764dd016d230be6ab0d5864
+size 298042336

text_encoder/special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

text_encoder/tokenizer_config.json

@@ -0,0 +1,944 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "|||IP_ADDRESS|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "1": {
+      "content": "<|padding|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50254": {
+      "content": "                        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50255": {
+      "content": "                       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50256": {
+      "content": "                      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50257": {
+      "content": "                     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50258": {
+      "content": "                    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50259": {
+      "content": "                   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50260": {
+      "content": "                  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50261": {
+      "content": "                 ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50262": {
+      "content": "                ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50263": {
+      "content": "               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50264": {
+      "content": "              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50265": {
+      "content": "             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50266": {
+      "content": "            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50267": {
+      "content": "           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50268": {
+      "content": "          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50269": {
+      "content": "         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50270": {
+      "content": "        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50271": {
+      "content": "       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50272": {
+      "content": "      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50273": {
+      "content": "     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50274": {
+      "content": "    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50275": {
+      "content": "   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50276": {
+      "content": "  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50277": {
+      "content": "|||EMAIL_ADDRESS|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50278": {
+      "content": "|||PHONE_NUMBER|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50279": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50280": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50281": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50282": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50283": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50284": {
+      "content": "[MASK]",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50285": {
+      "content": "[unused0]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50286": {
+      "content": "[unused1]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50287": {
+      "content": "[unused2]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50288": {
+      "content": "[unused3]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50289": {
+      "content": "[unused4]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50290": {
+      "content": "[unused5]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50291": {
+      "content": "[unused6]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50292": {
+      "content": "[unused7]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50293": {
+      "content": "[unused8]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50294": {
+      "content": "[unused9]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50295": {
+      "content": "[unused10]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50296": {
+      "content": "[unused11]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50297": {
+      "content": "[unused12]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50298": {
+      "content": "[unused13]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50299": {
+      "content": "[unused14]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50300": {
+      "content": "[unused15]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50301": {
+      "content": "[unused16]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50302": {
+      "content": "[unused17]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50303": {
+      "content": "[unused18]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50304": {
+      "content": "[unused19]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50305": {
+      "content": "[unused20]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50306": {
+      "content": "[unused21]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50307": {
+      "content": "[unused22]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50308": {
+      "content": "[unused23]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50309": {
+      "content": "[unused24]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50310": {
+      "content": "[unused25]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50311": {
+      "content": "[unused26]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50312": {
+      "content": "[unused27]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50313": {
+      "content": "[unused28]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50314": {
+      "content": "[unused29]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50315": {
+      "content": "[unused30]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50316": {
+      "content": "[unused31]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50317": {
+      "content": "[unused32]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50318": {
+      "content": "[unused33]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50319": {
+      "content": "[unused34]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50320": {
+      "content": "[unused35]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50321": {
+      "content": "[unused36]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50322": {
+      "content": "[unused37]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50323": {
+      "content": "[unused38]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50324": {
+      "content": "[unused39]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50325": {
+      "content": "[unused40]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50326": {
+      "content": "[unused41]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50327": {
+      "content": "[unused42]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50328": {
+      "content": "[unused43]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50329": {
+      "content": "[unused44]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50330": {
+      "content": "[unused45]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50331": {
+      "content": "[unused46]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50332": {
+      "content": "[unused47]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50333": {
+      "content": "[unused48]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50334": {
+      "content": "[unused49]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50335": {
+      "content": "[unused50]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50336": {
+      "content": "[unused51]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50337": {
+      "content": "[unused52]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50338": {
+      "content": "[unused53]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50339": {
+      "content": "[unused54]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50340": {
+      "content": "[unused55]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50341": {
+      "content": "[unused56]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50342": {
+      "content": "[unused57]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50343": {
+      "content": "[unused58]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50344": {
+      "content": "[unused59]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50345": {
+      "content": "[unused60]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50346": {
+      "content": "[unused61]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50347": {
+      "content": "[unused62]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50348": {
+      "content": "[unused63]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50349": {
+      "content": "[unused64]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50350": {
+      "content": "[unused65]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50351": {
+      "content": "[unused66]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50352": {
+      "content": "[unused67]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50353": {
+      "content": "[unused68]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50354": {
+      "content": "[unused69]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50355": {
+      "content": "[unused70]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50356": {
+      "content": "[unused71]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50357": {
+      "content": "[unused72]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50358": {
+      "content": "[unused73]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50359": {
+      "content": "[unused74]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50360": {
+      "content": "[unused75]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50361": {
+      "content": "[unused76]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50362": {
+      "content": "[unused77]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50363": {
+      "content": "[unused78]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50364": {
+      "content": "[unused79]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50365": {
+      "content": "[unused80]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50366": {
+      "content": "[unused81]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50367": {
+      "content": "[unused82]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "model_max_length": 8192,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "unk_token": "[UNK]"
+}

transformer/config.json

@@ -0,0 +1,157 @@
+{
+  "_class_name": "UDiT",
+  "name": "Ra",
+  "format": "udm/v1",
+  "step": 212000,
+  "latent": {
+    "channels": 32,
+    "spatial_factor": 8,
+    "mean": [
+      -4.137261,
+      2.725146,
+      -1.435803,
+      -2.187023,
+      -4.161275,
+      -1.58918,
+      -2.527523,
+      3.002302,
+      2.793705,
+      -2.091501,
+      -2.632896,
+      -3.59589,
+      -3.069669,
+      -4.568713,
+      3.26668,
+      -0.897648,
+      2.072541,
+      5.451158,
+      1.747871,
+      -2.113284,
+      2.03444,
+      -1.379274,
+      0.66369,
+      2.712183,
+      -0.111181,
+      6.054838,
+      1.235881,
+      -2.02777,
+      3.663243,
+      -2.454632,
+      1.141335,
+      1.786136
+    ],
+    "std": [
+      1.030738,
+      0.951682,
+      1.079687,
+      0.957871,
+      1.198266,
+      1.102266,
+      1.213605,
+      0.947544,
+      1.085991,
+      1.093223,
+      0.9027,
+      1.010205,
+      0.920466,
+      1.044107,
+      1.023179,
+      0.992978,
+      0.991627,
+      1.073585,
+      1.114629,
+      1.056403,
+      1.093012,
+      1.015109,
+      1.007069,
+      0.976026,
+      0.982324,
+      1.093805,
+      1.036086,
+      1.005727,
+      1.104078,
+      1.042298,
+      1.095522,
+      0.960255
+    ],
+    "source": "vae/ (Muinez/f8c32-ae)"
+  },
+  "text": {
+    "encoder": "text_encoder/ (jhu-clsp/ettin-encoder-150m, 768-dim, 22 слоя)",
+    "max_len": 192,
+    "slices": [
+      2,
+      5,
+      8,
+      11,
+      14,
+      17,
+      20
+    ],
+    "slice_blocks": 2,
+    "trainable": "text_mixer.*"
+  },
+  "body": {
+    "kind": "пирамидальный U: ширина растёт к низу",
+    "levels": [
+      {
+        "level": 0,
+        "tokens": "÷1",
+        "width": 192,
+        "blocks": "conv 6 enc / 4 dec + 1 self-attn на сторону",
+        "text": "кросс-внимание в каждом conv-блоке и в self-attn"
+      },
+      {
+        "level": 1,
+        "tokens": "÷2",
+        "width": 512,
+        "blocks": "attn 4 enc / 3 dec",
+        "text": "кросс-внимание"
+      },
+      {
+        "level": 2,
+        "tokens": "÷4",
+        "width": 1024,
+        "blocks": "attn 3 enc / 2 dec",
+        "text": "кросс-внимание"
+      },
+      {
+        "level": 3,
+        "tokens": "÷8",
+        "width": 1536,
+        "blocks": "attn 2 (боттлнек)",
+        "text": "кросс-внимание"
+      }
+    ],
+    "skips": "аддитивные (SD-стиль): e0/e1/e2 прибавляются в декодере",
+    "down": "mean-pool 2× + Linear",
+    "up": "Linear + nearest 2×",
+    "rope": "2D RoPE, base 100, фаза в латентных пикселях на каждом уровне",
+    "parametrization": "x0 латента"
+  },
+  "tail": {
+    "source": "декодер AE, 6 блоков @512, PixelShuffle ×3",
+    "trained_as": "flow по пикселям, условие z0 + (ошибка тела).detach()",
+    "sample_steps": 4
+  },
+  "params": {
+    "body": 143.33,
+    "text_mixer": 10.04,
+    "tail": 23.87,
+    "total": 177.24
+  },
+  "training": {
+    "dataset": "LucasFang/FLUX-Reason-6M",
+    "parts": "Aesthetics-Part01,Aesthetics-Part02,Imaginative,Text",
+    "size": 256,
+    "steps": 212000,
+    "batch": 16,
+    "accum": 1,
+    "lr": 0.0001,
+    "lr_tail": 3e-05,
+    "drop_text": 0.1,
+    "init": "tmp/flux2/ckpt/step_0192000.pt",
+    "metrics": "FLUX-Reason Part01 on top of the 20k butterfly run; continued on Parts 01+02+Imaginative+Text from step 192000, mse_lat 0.244 (ema 0.245), low 0.026, high 0.57, mse_pix 0.019"
+  },
+  "storage_dtype": "float16"
+}

transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f288233710ddf86c9d50099d59fa5d1f19438ff9a55b418bcc27ff2a23405658
+size 354543188

vae/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

vae/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

vae/README.md

@@ -0,0 +1,104 @@
+---
+license: apache-2.0
+---
+
+# Image Autoencoder (8x, 32 channels, flow-matching decoder)
+
+Inference-only package for an image autoencoder intended as a latent space for diffusion models.
+
+- **Compression:** 8x spatially, 32 latent channels — an `H x W` RGB image becomes a `32 x H/8 x W/8` latent.
+- **Encoder:** deterministic transformer with local window attention (2D RoPE), works at any resolution
+  whose sides are multiples of 8.
+- **Decoder:** conditional flow-matching model (x0-prediction). It starts from Gaussian noise and is
+  integrated in `N` steps (default **4**) conditioned on the latent.
+- **Normalization built in:** `encode()` returns latents already normalized per channel
+  (`(z - mean) / std`, roughly zero mean / unit variance); `decode()` undoes it. Use
+  `encode_raw()` / `decode_raw()` for unnormalized latents.
+
+## Install
+
+```bash
+pip install torch safetensors huggingface_hub numpy pillow
+# optional: triton (CUDA) for the fused window-attention kernel; otherwise PyTorch SDPA is used
+```
+
+The model code lives in the `ae/` folder of this repo:
+
+```bash
+huggingface-cli download Muinez/f8c32-ae --local-dir f8c32-ae
+export PYTHONPATH=$PWD/f8c32-ae:$PYTHONPATH
+```
+
+## Usage
+
+```python
+import torch
+from ae import AutoEncoder
+
+model = AutoEncoder.from_pretrained("Muinez/f8c32-ae", device="cuda")   # bf16 compute on CUDA
+
+x = ...                                   # (B, 3, H, W) float in [-1, 1], H and W multiples of 8
+z = model.encode(x)                       # (B, 32, H/8, W/8), normalized latents
+y = model.decode(z, steps=4)              # (B, 3, H, W), approximately in [-1, 1]
+y = y.clamp(-1, 1)
+
+# reproducible decoding
+g = torch.Generator("cuda").manual_seed(0)
+y = model.decode(z, steps=4, generator=g)
+
+# raw (unnormalized) latents
+z_raw = model.encode_raw(x)
+y = model.decode_raw(z_raw, steps=4)
+```
+
+### Batch of different resolutions / aspect ratios
+
+`encode` / `decode` (and the `_raw` variants) also take a **list** of tensors of different
+sizes and return a list. The whole list goes through the model in one forward pass: each
+image keeps its own size (no resizing or cropping), is cut into its own grid of 8x8 patch
+tokens, and the grids are packed into one sequence batch padded only up to the largest image
+area. Attention never crosses into other images or padding, so results match encoding each
+image on its own. Only the decoder's small convolutional upsampling head runs per group of
+equal-shape images.
+
+```python
+from PIL import Image
+import numpy as np
+
+def load(path, max_side=1024):
+    im = Image.open(path).convert("RGB")
+    s = min(1.0, max_side / max(im.size))
+    w, h = (int(im.width * s) // 8 * 8, int(im.height * s) // 8 * 8)   # sides: multiples of 8
+    im = im.resize((w, h), Image.LANCZOS)
+    return torch.from_numpy(np.asarray(im)).permute(2, 0, 1)            # uint8 (3, H, W)
+
+images = [load("portrait.jpg"), load("landscape.png"), load("square.webp")]   # any sizes
+latents = model.encode(images)              # list of (32, H_i/8, W_i/8), normalized
+recons = model.decode(latents, steps=4)     # list of (3, H_i, W_i)
+```
+
+Images may be uint8 `[0, 255]` or float `[-1, 1]`.
+
+See `example.py` for a complete round trip (`python example.py input.png recon.png`).
+
+### Notes
+
+- `steps` trades speed for detail: 1 step is fastest and slightly softer, 4 is the default.
+  Because the decoder samples from noise, different seeds give slightly different fine detail.
+- `dtype` in `from_pretrained` sets the compute dtype (default: bfloat16 on CUDA, float32 on CPU).
+  Encoder weights are cast to it; the decoder keeps float32 weights and runs under autocast.
+- On CUDA with Triton installed, a fused window-attention kernel is used for any image sizes;
+  otherwise (CPU, no Triton) an equivalent masked SDPA path is used. Force SDPA with
+  `import ae.window_attention as wa; wa.BACKEND = "sdpa"`.
+
+## Files
+
+- `config.json` — architecture and latent normalization statistics
+- `model.safetensors` — encoder and decoder weights (float32)
+- `ae/` — model code (`AutoEncoder`, `AEConfig`)
+- `example.py` — round-trip example
+- `LICENSE` — Apache License 2.0
+
+## License
+
+Apache License 2.0 — see `LICENSE`.

vae/ae/__init__.py

@@ -0,0 +1,3 @@
+from .model import AEConfig, AutoEncoder, Encoder, FlowDecoder, Packing
+
+__all__ = ["AEConfig", "AutoEncoder", "Encoder", "FlowDecoder", "Packing"]

vae/ae/model.py

@@ -0,0 +1,469 @@
+"""Image autoencoder: deterministic transformer encoder + flow-matching decoder.
+
+* Encoder: 8x8 patches -> windowed-attention transformer blocks -> 32-channel latent
+  on a grid 8x smaller than the image.
+* Decoder: a conditional flow-matching model. Given the latent, a noisy image ``x_t`` and
+  the time ``t``, it predicts the clean image ``x_0``. Images are produced by integrating
+  from pure noise (t=1) to the image (t=0) in a small number of steps (default 4).
+
+Everything except the decoder's upsampling head works on token sequences, so a batch may
+mix images of any sizes and aspect ratios: each image is cut into its own grid of patch
+tokens, the grids are packed into one padded sequence batch (:class:`Packing`), and
+attention is restricted to each image's own grid. The convolutional upsampling head runs
+per group of equal-shape images.
+
+All architecture hyper-parameters come from :class:`AEConfig`.
+"""
+import json
+import math
+import os
+from dataclasses import asdict, dataclass
+from typing import List, Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .window_attention import _tile_width, packed_window_attention
+
+
+@dataclass
+class AEConfig:
+    in_channels: int = 3
+    patch_size: int = 8               # spatial compression factor
+    latent_channels: int = 32
+    # encoder
+    enc_width: int = 1024
+    enc_depth: int = 8
+    enc_mlp_ratio: float = 2.0
+    # decoder
+    dec_width: int = 512
+    dec_depth: int = 6
+    dec_mlp_ratio: float = 4.0
+    dec_adaln_rank: int = 128         # rank of the per-block low-rank modulation deltas
+    dec_skip_blocks: int = 2          # long skip: output of block (n-1) feeds the last n blocks
+    dec_head_min_channels: int = 64
+    time_embed_dim: int = 256
+    # shared
+    head_dim: int = 64
+    window_radius: int = 8            # attention window: |dy| <= r and |dx| <= r tokens
+    rope_base: float = 100.0
+    # per-channel statistics of raw latents (used by encode()/decode())
+    latent_mean: Optional[List[float]] = None
+    latent_std: Optional[List[float]] = None
+
+    @classmethod
+    def from_json(cls, path):
+        with open(path) as f:
+            d = json.load(f)
+        known = {k: v for k, v in d.items() if k in cls.__dataclass_fields__}
+        return cls(**known)
+
+    def to_json(self, path):
+        with open(path, "w") as f:
+            json.dump(asdict(self), f, indent=2)
+
+
+# --------------------------------------------------------------------------- packing
+
+class Packing:
+    """Layout of a batch of token grids of sizes ``hw = [(h, w), ...]``.
+
+    Sample ``b`` is stored row-major in its own padded ``gyp x gxp`` grid at the start of a
+    sequence of length ``N`` (the largest padded area); its real tokens are the top-left
+    ``h x w`` corner. ``pos[b]`` are the sequence indices of its real tokens.
+    """
+
+    def __init__(self, hw: Sequence[Tuple[int, int]], radius: int, rope_base: float,
+                 head_dim: int, device):
+        self.hw = [(int(h), int(w)) for h, w in hw]
+        WY, WX = 8, _tile_width(radius)
+        B = len(self.hw)
+        geo = [(math.ceil(h / WY) * WY, math.ceil(w / WX) * WX) for h, w in self.hw]
+        self.N = max(a * b for a, b in geo)
+        self.max_tiles = max((a // WY) * (b // WX) for a, b in geo)
+        self.dims = torch.tensor([[a, b, h, w] for (a, b), (h, w) in zip(geo, self.hw)],
+                                 dtype=torch.int32, device=device)
+        self.valid = torch.zeros(B, self.N, 1, device=device)
+        pos_y = torch.zeros(B, self.N, device=device)
+        pos_x = torch.zeros(B, self.N, device=device)
+        self.pos = []
+        for b, ((_, gxp), (h, w)) in enumerate(zip(geo, self.hw)):
+            yy = torch.arange(h, device=device).repeat_interleave(w)
+            xx = torch.arange(w, device=device).repeat(h)
+            p = yy * gxp + xx
+            self.pos.append(p)
+            self.valid[b, p] = 1
+            pos_y[b, p], pos_x[b, p] = yy.float(), xx.float()
+        # 2D axial RoPE tables: the head dim is split into four quarters,
+        # two rotate with the x coordinate, two with the y coordinate.
+        dq = head_dim // 4
+        fr = torch.exp(-math.log(rope_base) * torch.arange(dq, device=device, dtype=torch.float32) / dq)
+        ax, ay = pos_x[..., None] * fr, pos_y[..., None] * fr                  # (B, N, dq)
+        self._rope = [t[:, None] for t in (ax.cos(), ax.sin(), ay.cos(), ay.sin())]
+        self._dq = dq
+        self.radius = radius
+
+    def pack(self, seqs):
+        """list of (n_b, C) -> (B, N, C), zero padded."""
+        out = seqs[0].new_zeros(len(seqs), self.N, seqs[0].shape[-1])
+        for b, s in enumerate(seqs):
+            out[b, self.pos[b]] = s
+        return out
+
+    def unpack(self, x):
+        """(B, N, C) -> list of (C, h, w)."""
+        return [x[b, self.pos[b]].transpose(0, 1).reshape(-1, h, w) for b, (h, w) in enumerate(self.hw)]
+
+    def rope(self, q, k):
+        dq = self._dq
+        cx, sx, cy, sy = (c.to(q.dtype) for c in self._rope)
+
+        def rot(t):
+            a, b, c, e = t[..., :dq], t[..., dq:2 * dq], t[..., 2 * dq:3 * dq], t[..., 3 * dq:]
+            return torch.cat([a * cx - b * sx, a * sx + b * cx, c * cy - e * sy, c * sy + e * cy], dim=-1)
+
+        return rot(q), rot(k)
+
+    def attention(self, q, k, v):
+        return packed_window_attention(q, k, v, self.radius, self.dims, self.max_tiles)
+
+
+def patchify(img, p):
+    """(C, H, W) -> (H/p * W/p, C*p*p) row-major patch tokens."""
+    C, H, W = img.shape
+    return img.reshape(C, H // p, p, W // p, p).permute(1, 3, 0, 2, 4).reshape(-1, C * p * p)
+
+
+# --------------------------------------------------------------------------- encoder
+
+class RMSNorm(nn.Module):
+    def __init__(self, ch, eps=1e-8):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(ch))
+
+    def forward(self, x):
+        return x * x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() * self.weight
+
+
+class Mlp(nn.Module):
+    def __init__(self, i, h, o):
+        super().__init__()
+        self.m = nn.Sequential(nn.Linear(i, h), nn.GELU(), nn.Linear(h, o))
+
+    def forward(self, x):
+        return self.m(x)
+
+
+class PlainMlp(nn.Module):
+    def __init__(self, d, ratio):
+        super().__init__()
+        h = int(d * ratio)
+        self.up = nn.Linear(d, h)
+        self.down = nn.Linear(h, d)
+        self.act = nn.GELU(approximate="tanh")
+
+    def forward(self, x):
+        return self.down(self.act(self.up(x)))
+
+
+class EncoderBlock(nn.Module):
+    """Pre-norm transformer block with 2D-RoPE windowed self-attention and an MLP."""
+
+    def __init__(self, ch, head_dim, mlp_ratio):
+        super().__init__()
+        self.h = max(1, ch // head_dim)
+        self.norm1 = RMSNorm(ch)
+        self.qkv = nn.Linear(ch, ch * 3, bias=False)
+        self.proj = nn.Linear(ch, ch, bias=False)
+        self.norm2 = RMSNorm(ch)
+        self.mlp = PlainMlp(ch, mlp_ratio)
+
+    def forward(self, x, pk: Packing):
+        B, N, C = x.shape
+        q, k, v = self.qkv(self.norm1(x)).reshape(B, N, 3, self.h, C // self.h).permute(2, 0, 3, 1, 4)
+        q, k = pk.rope(q, k)
+        a = pk.attention(q, k, v).transpose(1, 2).reshape(B, N, C).to(x.dtype)
+        x = x + self.proj(a)
+        return x + self.mlp(self.norm2(x))
+
+
+class Encoder(nn.Module):
+    """patch tokens (B, N, 3*p*p) -> raw latent tokens (B, N, latent_channels)."""
+
+    def __init__(self, cfg: AEConfig):
+        super().__init__()
+        w = cfg.enc_width
+        self.initial = nn.Linear(cfg.in_channels * cfg.patch_size ** 2, w)
+        self.blocks = nn.ModuleList(EncoderBlock(w, cfg.head_dim, cfg.enc_mlp_ratio)
+                                    for _ in range(cfg.enc_depth))
+        self.out = Mlp(w, w, cfg.latent_channels)
+
+    def forward(self, x, pk: Packing):
+        valid = pk.valid.to(x.dtype)
+        x = self.initial(x)
+        for b in self.blocks:
+            x = b(x, pk) * valid
+        return self.out(x)
+
+
+# --------------------------------------------------------------------------- decoder
+
+class DecoderBlock(nn.Module):
+    """Transformer block with windowed 2D-RoPE attention, modulated by the timestep:
+    the modulation vector carries (scale, gate) for the attention and the MLP sub-layers."""
+
+    def __init__(self, ch, head_dim, mlp_ratio):
+        super().__init__()
+        self.h = max(1, ch // head_dim)
+        self.norm1 = nn.RMSNorm(ch)
+        self.qkv = nn.Linear(ch, 3 * ch)
+        self.proj = nn.Linear(ch, ch)
+        self.norm2 = nn.RMSNorm(ch)
+        hid = int(ch * mlp_ratio)
+        self.fc1 = nn.Linear(ch, hid)
+        self.fc2 = nn.Linear(hid, ch)
+
+    def forward(self, x, mod, pk: Packing):
+        B, N, C = x.shape
+        s1, g1, s2, g2 = mod[:, None, :].chunk(4, dim=-1)
+        h = self.norm1(x) * (1 + s1)
+        q, k, v = self.qkv(h).reshape(B, N, 3, self.h, C // self.h).permute(2, 0, 3, 1, 4)
+        q, k = pk.rope(q, k)
+        a = pk.attention(q, k, v).transpose(1, 2).reshape(B, N, C).to(h.dtype)
+        x = x + g1 * self.proj(a)
+        h = self.norm2(x) * (1 + s2)
+        return x + g2 * self.fc2(F.gelu(self.fc1(h)))
+
+
+class FlowDecoder(nn.Module):
+    """x0-prediction network: (x_t, t, latent) -> predicted clean image.
+
+    The transformer works on packed tokens; the upsampling head (PixelShuffle + 3x3 convs)
+    runs per group of equal-shape samples.
+    """
+
+    def __init__(self, cfg: AEConfig):
+        super().__init__()
+        D, p = cfg.dec_width, cfg.patch_size
+        self.depth = cfg.dec_depth
+        self.n_skip = cfg.dec_skip_blocks
+        self.time_embed_dim = cfg.time_embed_dim
+        self.patchify = nn.Linear(cfg.in_channels * p * p, D)
+        self.latent_up = nn.Linear(cfg.latent_channels, D)
+        self.fuse_in = nn.Linear(2 * D, D)
+        self.tmlp = nn.Sequential(nn.Linear(cfg.time_embed_dim, D), nn.SiLU(), nn.Linear(D, D))
+        self.adaln_act = nn.SiLU()
+        mw = 4 * D
+        self.adaln_base = nn.Linear(D, mw)
+        self.d_down = nn.ModuleList(nn.Linear(D, cfg.dec_adaln_rank, bias=False) for _ in range(self.depth))
+        self.d_up = nn.ModuleList(nn.Linear(cfg.dec_adaln_rank, mw, bias=False) for _ in range(self.depth))
+        self.blocks = nn.ModuleList(DecoderBlock(D, cfg.head_dim, cfg.dec_mlp_ratio)
+                                    for _ in range(self.depth))
+        self.fuse_skip = nn.Linear(2 * D, D)
+        # progressive upsampling head: [PixelShuffle(2) -> conv3x3 -> GELU] x log2(p)
+        n = int(round(math.log2(p)))
+        ch, stages = D, []
+        for _ in range(n):
+            nxt = max(ch // 2, cfg.dec_head_min_channels)
+            stages.append(nn.Sequential(nn.PixelShuffle(2), nn.Conv2d(ch // 4, nxt, 3, padding=1), nn.GELU()))
+            ch = nxt
+        self.up_stages = nn.ModuleList(stages)
+        self.up_final = nn.Conv2d(ch, cfg.in_channels, 3, padding=1)
+
+    def time_embed(self, t):
+        half = self.time_embed_dim // 2
+        freqs = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device, dtype=torch.float32) / half)
+        a = t.float()[:, None] * 1000.0 * freqs[None]
+        return self.tmlp(torch.cat([a.sin(), a.cos()], dim=-1).to(self.tmlp[0].weight.dtype))
+
+    def head(self, feats):
+        """feats: (B, D, h, w) -> (B, 3, h*p, w*p)."""
+        for st in self.up_stages:
+            feats = st(feats)
+        return self.up_final(feats)
+
+    def forward(self, xt, t, z, pk: Packing):
+        """xt: packed noisy patch tokens (B, N, 3*p*p); t: (B,); z: packed latents (B, N, C).
+        Returns a list of per-sample x0 predictions (3, H_b, W_b)."""
+        valid = pk.valid
+        cond = self.adaln_act(self.time_embed(t))
+        base = self.adaln_base(cond)
+        x = self.fuse_in(torch.cat([self.patchify(xt), self.latent_up(z)], dim=-1))
+        valid = valid.to(x.dtype)
+        x = x * valid
+        outs = []
+        for i, b in enumerate(self.blocks):
+            if i == self.depth - self.n_skip:
+                x = self.fuse_skip(torch.cat([outs[self.n_skip - 1], x], dim=-1))
+            x = b(x, base + self.d_up[i](self.d_down[i](cond)), pk) * valid
+            outs.append(x)
+        feats = pk.unpack(x)
+        preds = [None] * len(feats)
+        groups = {}
+        for i, f in enumerate(feats):
+            groups.setdefault(tuple(f.shape), []).append(i)
+        for idx in groups.values():
+            y = self.head(torch.stack([feats[i] for i in idx]))
+            for j, i in enumerate(idx):
+                preds[i] = y[j]
+        return preds
+
+
+# --------------------------------------------------------------------------- wrapper
+
+def _as_list(x):
+    """Tensor (B, C, H, W) or list of (C, H, W) -> (list, was_tensor)."""
+    if torch.is_tensor(x):
+        assert x.dim() == 4, f"expected (B, C, H, W), got {tuple(x.shape)}"
+        return list(x.unbind(0)), True
+    return list(x), False
+
+
+class AutoEncoder(nn.Module):
+    """Encoder + flow-matching decoder with built-in latent normalization.
+
+    ``encode`` returns latents normalized per channel ((z - mean) / std), ready to be used
+    as a diffusion target; ``decode`` expects such latents. ``encode_raw`` / ``decode_raw``
+    skip the normalization.
+
+    All four methods take either a batch tensor ``(B, C, H, W)`` (returns a tensor) or a
+    list of ``(C, H_i, W_i)`` tensors of different sizes (returns a list), processed in one
+    forward pass.
+    """
+
+    def __init__(self, cfg: AEConfig):
+        super().__init__()
+        self.config = cfg
+        self.encoder = Encoder(cfg)
+        self.decoder = FlowDecoder(cfg)
+        C = cfg.latent_channels
+        mean = torch.tensor(cfg.latent_mean, dtype=torch.float32) if cfg.latent_mean else torch.zeros(C)
+        std = torch.tensor(cfg.latent_std, dtype=torch.float32) if cfg.latent_std else torch.ones(C)
+        self.register_buffer("latent_mean", mean.view(C, 1, 1), persistent=False)
+        self.register_buffer("latent_std", std.view(C, 1, 1), persistent=False)
+        # compute dtype: encoder weights are stored in it; the decoder keeps fp32 weights and
+        # runs under autocast with this dtype (when it is not float32).
+        self.compute_dtype = torch.float32
+
+    # ---------------------------------------------------------------- loading
+    @classmethod
+    def from_pretrained(cls, path_or_repo, device=None, dtype=None, **hub_kwargs):
+        """Load from a local directory or a Hugging Face repo id containing
+        ``config.json`` and ``model.safetensors``.
+
+        dtype: compute dtype (default: bfloat16 on CUDA, float32 on CPU).
+        """
+        from safetensors.torch import load_file
+
+        if os.path.isdir(path_or_repo):
+            cfg_path = os.path.join(path_or_repo, "config.json")
+            w_path = os.path.join(path_or_repo, "model.safetensors")
+        else:
+            from huggingface_hub import hf_hub_download
+            cfg_path = hf_hub_download(path_or_repo, "config.json", **hub_kwargs)
+            w_path = hf_hub_download(path_or_repo, "model.safetensors", **hub_kwargs)
+        device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        if dtype is None:
+            dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+        model = cls(AEConfig.from_json(cfg_path))
+        model.load_state_dict(load_file(w_path), strict=True)
+        model.to(device).eval().requires_grad_(False)
+        model.set_dtype(dtype)
+        return model
+
+    def set_dtype(self, dtype):
+        self.compute_dtype = dtype
+        self.encoder.to(dtype)
+        self.decoder.float()
+        return self
+
+    @property
+    def device(self):
+        return self.latent_mean.device
+
+    def _autocast(self):
+        return torch.autocast(self.device.type, dtype=self.compute_dtype,
+                              enabled=self.compute_dtype != torch.float32)
+
+    def _packing(self, hw):
+        c = self.config
+        return Packing(hw, c.window_radius, c.rope_base, c.head_dim, self.device)
+
+    # ---------------------------------------------------------------- encoder side
+    @torch.no_grad()
+    def encode_raw(self, images):
+        """images: (B, 3, H, W) or a list of (3, H_i, W_i); float in [-1, 1] or uint8 in
+        [0, 255]; sides multiples of the patch size.
+        Returns raw (unnormalized) float32 latents: (B, C, H/p, W/p) or a list of (C, H_i/p, W_i/p)."""
+        p = self.config.patch_size
+        imgs, was_tensor = _as_list(images)
+        for im in imgs:
+            assert im.dim() == 3 and im.shape[-2] % p == 0 and im.shape[-1] % p == 0, \
+                f"each image must be (3, H, W) with H, W multiples of {p}, got {tuple(im.shape)}"
+        pk = self._packing([(im.shape[-2] // p, im.shape[-1] // p) for im in imgs])
+        toks = []
+        for im in imgs:
+            im = im.to(self.device)
+            im = im.float() / 127.5 - 1.0 if im.dtype == torch.uint8 else im
+            toks.append(patchify(im.to(self.compute_dtype), p))
+        z = [t.float() for t in pk.unpack(self.encoder(pk.pack(toks), pk))]
+        return torch.stack(z) if was_tensor else z
+
+    @torch.no_grad()
+    def encode(self, images):
+        """Images -> normalized latents (zero mean / unit std per channel)."""
+        z = self.encode_raw(images)
+        if torch.is_tensor(z):
+            return (z - self.latent_mean) / self.latent_std
+        return [(t - self.latent_mean) / self.latent_std for t in z]
+
+    # ---------------------------------------------------------------- decoder side
+    @torch.no_grad()
+    def decode_raw(self, latents, steps=4, generator=None, noise=None):
+        """Raw latents -> images (approximately in [-1, 1], not clamped).
+
+        latents: (B, C, h, w) or a list of (C, h_i, w_i).
+        steps: number of flow-matching integration steps (default 4; 1 gives a fast,
+               softer result, more steps add a bit of detail).
+        generator: optional torch.Generator for the initial noise (on the model device).
+        noise: optional explicit initial noise, (B, 3, H, W) or a list of (3, H_i, W_i).
+        """
+        p, C = self.config.patch_size, self.config.in_channels
+        zs, was_tensor = _as_list(latents)
+        zs = [z.to(self.device, torch.float32) for z in zs]
+        pk = self._packing([tuple(z.shape[-2:]) for z in zs])
+        zp = pk.pack([z.flatten(1).transpose(0, 1) for z in zs])
+        if noise is not None:
+            xs = [n.to(self.device, torch.float32) for n in _as_list(noise)[0]]
+        elif was_tensor:
+            B, _, h, w = latents.shape
+            xs = list(torch.randn((B, C, h * p, w * p), generator=generator, device=self.device).unbind(0))
+        else:
+            xs = [torch.randn((C, z.shape[1] * p, z.shape[2] * p), generator=generator, device=self.device)
+                  for z in zs]
+        ts = torch.linspace(1.0, 0.0, steps + 1, device=self.device)
+        B = len(zs)
+        for i in range(steps):
+            with self._autocast():
+                xt = pk.pack([patchify(x.to(self.compute_dtype), p) for x in xs])
+                preds = self.decoder(xt, ts[i].expand(B), zp, pk)
+            # x0 prediction -> implied noise -> move along the straight path to ts[i+1]
+            xs = [ts[i + 1] * (x - (1 - ts[i]) * pr.float()) / ts[i] + (1 - ts[i + 1]) * pr.float()
+                  for x, pr in zip(xs, preds)]
+        return torch.stack(xs) if was_tensor else xs
+
+    @torch.no_grad()
+    def decode(self, latents, steps=4, generator=None, noise=None):
+        """Normalized latents (as returned by ``encode``) -> images in about [-1, 1]."""
+        if torch.is_tensor(latents):
+            z = latents.to(self.device, torch.float32) * self.latent_std + self.latent_mean
+        else:
+            z = [t.to(self.device, torch.float32) * self.latent_std + self.latent_mean for t in latents]
+        return self.decode_raw(z, steps=steps, generator=generator, noise=noise)
+
+    @torch.no_grad()
+    def forward(self, images, steps=4, generator=None):
+        """Round trip: images -> latents -> images."""
+        return self.decode(self.encode(images), steps=steps, generator=generator)

vae/ae/window_attention.py

@@ -0,0 +1,129 @@
+"""Local 2D window attention.
+
+Every token attends to all tokens within a square window of Chebyshev radius ``r``
+(i.e. ``|dy| <= r and |dx| <= r``) on its 2D grid.
+
+Two numerically equivalent implementations are provided:
+
+* a fused Triton kernel (CUDA only, forward pass only);
+* PyTorch ``scaled_dot_product_attention`` with a boolean window mask (CPU, or when
+  Triton is unavailable).
+
+Both work on a packed batch: every sample is its own token grid of any size.
+"""
+import torch
+import torch.nn.functional as F
+
+try:  # Triton is optional; without it the SDPA path is used.
+    import triton
+    import triton.language as tl
+
+    _HAS_TRITON = True
+except Exception:  # pragma: no cover
+    _HAS_TRITON = False
+
+_LOG2E = 1.4426950408889634
+
+# Global switch: "auto" (kernel when possible), "sdpa" (always masked SDPA).
+BACKEND = "auto"
+
+
+if _HAS_TRITON:
+
+    @triton.jit
+    def _window_attn_fwd(Q, K, V, O, DIMS, NH: tl.constexpr, NMAX, WY: tl.constexpr,
+                         WX: tl.constexpr, R: tl.constexpr, D: tl.constexpr,
+                         BR: tl.constexpr, qk_scale):
+        # One program = one WY x WX tile of queries for one (batch, head).
+        # Keys are streamed from the tile extended by R on every side ("halo").
+        t, bh = tl.program_id(0), tl.program_id(1)
+        b = bh // NH
+        gyp = tl.load(DIMS + b * 4 + 0)
+        gxp = tl.load(DIMS + b * 4 + 1)
+        gyr = tl.load(DIMS + b * 4 + 2)
+        gxr = tl.load(DIMS + b * 4 + 3)
+        ntx = gxp // WX
+        if t >= (gyp // WY) * ntx:
+            return
+        ty, tx = t // ntx, t % ntx
+        base = bh * NMAX * D
+        SX: tl.constexpr = WX + 2 * R
+        SY: tl.constexpr = WY + 2 * R
+
+        qi = tl.arange(0, WY * WX)
+        qyr, qxr = qi // WX, qi % WX
+        qy, qx = ty * WY + qyr, tx * WX + qxr
+        d = tl.arange(0, D)
+        qo = base + (qy * gxp + qx)[:, None] * D + d[None, :]
+        q = tl.load(Q + qo)
+        okq = (qy < gyr) & (qx < gxr)
+
+        kxr = tl.arange(0, BR * SX) % SX
+        okx = okq[:, None] & (tl.abs(qxr[:, None] - (kxr[None, :] - R)) <= R)
+        kxg = tx * WX - R + kxr
+        okg = (kxg >= 0) & (kxg < gxr)
+
+        m = tl.full((WY * WX,), -1e30, tl.float32)
+        l = tl.zeros((WY * WX,), tl.float32)
+        acc = tl.zeros((WY * WX, D), tl.float32)
+
+        for r0 in range(0, SY, BR):
+            kyr = r0 + tl.arange(0, BR * SX) // SX
+            kyg = ty * WY - R + kyr
+            ok = okg & (kyg >= 0) & (kyg < gyr) & (kyr < SY)
+            off = base + (kyg * gxp + kxg) * D
+            k = tl.load(K + off[:, None] + d[None, :], mask=ok[:, None], other=0.0)
+            v = tl.load(V + off[:, None] + d[None, :], mask=ok[:, None], other=0.0)
+
+            s = tl.dot(q, tl.trans(k)) * qk_scale
+            good = okx & (tl.abs(qyr[:, None] - (kyr[None, :] - R)) <= R) & ok[None, :]
+            s = tl.where(good, s, -1e30)
+
+            mn = tl.maximum(m, tl.max(s, 1))
+            p = tl.where(good, tl.math.exp2(s - mn[:, None]), 0.0)
+            a = tl.math.exp2(m - mn)
+            l = l * a + tl.sum(p, 1)
+            acc = acc * a[:, None] + tl.dot(p.to(v.dtype), v).to(tl.float32)
+            m = mn
+
+        lc = tl.maximum(l, 1e-30)
+        tl.store(O + qo, (acc / lc[:, None]).to(O.dtype.element_ty))
+
+
+def _tile_width(r):
+    # The kernel needs a power-of-two key strip width (WX + 2R).
+    for wx in (16, 8):
+        sx = wx + 2 * r
+        if sx & (sx - 1) == 0:
+            return wx
+    raise ValueError(f"window radius {r}: no supported tile width")
+
+
+@torch._dynamo.disable
+def packed_window_attention(q, k, v, r, dims, max_tiles):
+    """Windowed attention for a packed batch of differently shaped token grids.
+
+    Sample ``b`` is stored row-major in a padded ``gyp x gxp`` grid at the start of the
+    sequence, with its real ``gyr x gxr`` tokens in the top-left corner; the rest is padding.
+    dims: (B, 4) int32 ``[gyp, gxp, gyr, gxr]`` (gyp a multiple of 8, gxp of the tile width).
+    q, k, v: (B, heads, Nmax, head_dim). Padding outputs are zero (kernel) or unused (SDPA).
+    """
+    B, nh, N, D = q.shape
+    if BACKEND == "auto" and _HAS_TRITON and q.is_cuda and D in (16, 32, 64, 128):
+        q, k, v = (t.to(torch.bfloat16).contiguous() for t in (q, k, v))
+        o = torch.zeros_like(q)
+        _window_attn_fwd[(max_tiles, B * nh)](
+            q, k, v, o, dims.to(q.device, torch.int32).contiguous(), NH=nh, NMAX=N, WY=8,
+            WX=_tile_width(r), R=r, D=D, BR=4, qk_scale=D ** -0.5 * _LOG2E, num_warps=8, num_stages=3)
+        return o
+    # reference path: per-sample block mask over the packed sequence
+    n = torch.arange(N, device=q.device)
+    d = dims.to(q.device).long()
+    yy, xx = n[None] // d[:, 1:2], n[None] % d[:, 1:2]                      # (B, N)
+    valid = (yy < d[:, 2:3]) & (xx < d[:, 3:4])
+    win = ((yy[:, :, None] - yy[:, None]).abs() <= r) & ((xx[:, :, None] - xx[:, None]).abs() <= r)
+    m = win & valid[:, None, :]
+    m = m | (~valid[:, :, None] & torch.eye(N, dtype=torch.bool, device=q.device))
+    return F.scaled_dot_product_attention(q.contiguous(), k.contiguous(), v.contiguous(),
+                                          attn_mask=m[:, None])
+

vae/config.json

@@ -0,0 +1,86 @@
+{
+  "in_channels": 3,
+  "patch_size": 8,
+  "latent_channels": 32,
+  "enc_width": 1024,
+  "enc_depth": 8,
+  "enc_mlp_ratio": 2.0,
+  "dec_width": 512,
+  "dec_depth": 6,
+  "dec_mlp_ratio": 4.0,
+  "dec_adaln_rank": 128,
+  "dec_skip_blocks": 2,
+  "dec_head_min_channels": 64,
+  "time_embed_dim": 256,
+  "head_dim": 64,
+  "window_radius": 8,
+  "rope_base": 100.0,
+  "latent_mean": [
+    -4.137261,
+    2.725146,
+    -1.435803,
+    -2.187023,
+    -4.161275,
+    -1.58918,
+    -2.527523,
+    3.002302,
+    2.793705,
+    -2.091501,
+    -2.632896,
+    -3.59589,
+    -3.069669,
+    -4.568713,
+    3.26668,
+    -0.897648,
+    2.072541,
+    5.451158,
+    1.747871,
+    -2.113284,
+    2.03444,
+    -1.379274,
+    0.66369,
+    2.712183,
+    -0.111181,
+    6.054838,
+    1.235881,
+    -2.02777,
+    3.663243,
+    -2.454632,
+    1.141335,
+    1.786136
+  ],
+  "latent_std": [
+    1.030738,
+    0.951682,
+    1.079687,
+    0.957871,
+    1.198266,
+    1.102266,
+    1.213605,
+    0.947544,
+    1.085991,
+    1.093223,
+    0.9027,
+    1.010205,
+    0.920466,
+    1.044107,
+    1.023179,
+    0.992978,
+    0.991627,
+    1.073585,
+    1.114629,
+    1.056403,
+    1.093012,
+    1.015109,
+    1.007069,
+    0.976026,
+    0.982324,
+    1.093805,
+    1.036086,
+    1.005727,
+    1.104078,
+    1.042298,
+    1.095522,
+    0.960255
+  ]
+}

vae/example.py

@@ -0,0 +1,29 @@
+"""Minimal round trip: image -> latents -> image."""
+import sys
+
+import numpy as np
+import torch
+from PIL import Image
+
+from ae import AutoEncoder
+
+src = sys.argv[1] if len(sys.argv) > 1 else "input.png"
+dst = sys.argv[2] if len(sys.argv) > 2 else "recon.png"
+
+model = AutoEncoder.from_pretrained("Muinez/f8c32-ae")   # HF repo id or local dir
+
+img = Image.open(src).convert("RGB")
+w, h = (img.width // 8) * 8, (img.height // 8) * 8  # sides must be multiples of 8
+img = img.crop((0, 0, w, h))
+x = torch.from_numpy(np.asarray(img, dtype=np.float32)).permute(2, 0, 1)[None] / 127.5 - 1
+
+z = model.encode(x)                                 # (1, 32, h/8, w/8), normalized
+print("latent", tuple(z.shape), f"mean {z.mean():.3f} std {z.std():.3f}")
+
+gen = torch.Generator(model.device).manual_seed(0)
+y = model.decode(z, steps=4, generator=gen).clamp(-1, 1)
+
+mse = (y.cpu() - x).pow(2).mean()
+print(f"PSNR {10 * torch.log10(4 / mse):.2f} dB")
+Image.fromarray(((y[0] + 1) * 127.5).round().byte().permute(1, 2, 0).cpu().numpy()).save(dst)
+print("saved", dst)

vae/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a484ad128d0837630dd16248af9b2265d2721d7434dc696bf7451fcbf1809665
+size 184623110